139 lines
3.5 KiB
Python
139 lines
3.5 KiB
Python
"""Synchronization store tools."""
|
|
|
|
from typing import Dict, Any, List
|
|
from langchain_core.tools import tool
|
|
|
|
|
|
def _get_registry():
|
|
"""Get the global registry instance."""
|
|
from . import _registry
|
|
return _registry
|
|
|
|
|
|
@tool
|
|
def list_sync_stores() -> List[str]:
|
|
"""List all available synchronization stores.
|
|
|
|
Returns:
|
|
List of store names that can be read/written
|
|
"""
|
|
registry = _get_registry()
|
|
if not registry:
|
|
return []
|
|
return list(registry.entries.keys())
|
|
|
|
|
|
@tool
|
|
def read_sync_state(store_name: str) -> Dict[str, Any]:
|
|
"""Read the current state of a synchronization store.
|
|
|
|
Args:
|
|
store_name: Name of the store to read (e.g., "TraderState", "StrategyState")
|
|
|
|
Returns:
|
|
Dictionary containing the current state of the store
|
|
|
|
Raises:
|
|
ValueError: If store_name doesn't exist
|
|
"""
|
|
registry = _get_registry()
|
|
if not registry:
|
|
raise ValueError("SyncRegistry not initialized")
|
|
|
|
entry = registry.entries.get(store_name)
|
|
if not entry:
|
|
available = list(registry.entries.keys())
|
|
raise ValueError(f"Store '{store_name}' not found. Available stores: {available}")
|
|
|
|
return entry.model.model_dump(mode="json")
|
|
|
|
|
|
@tool
|
|
async def write_sync_state(store_name: str, updates: Dict[str, Any]) -> Dict[str, str]:
|
|
"""Update the state of a synchronization store.
|
|
|
|
This will apply the updates to the store and trigger synchronization
|
|
with all connected clients.
|
|
|
|
Args:
|
|
store_name: Name of the store to update
|
|
updates: Dictionary of field updates (field_name: new_value)
|
|
|
|
Returns:
|
|
Dictionary with status and updated fields
|
|
|
|
Raises:
|
|
ValueError: If store_name doesn't exist or updates are invalid
|
|
"""
|
|
registry = _get_registry()
|
|
if not registry:
|
|
raise ValueError("SyncRegistry not initialized")
|
|
|
|
entry = registry.entries.get(store_name)
|
|
if not entry:
|
|
available = list(registry.entries.keys())
|
|
raise ValueError(f"Store '{store_name}' not found. Available stores: {available}")
|
|
|
|
try:
|
|
# Get current state
|
|
current_state = entry.model.model_dump(mode="json")
|
|
|
|
# Apply updates
|
|
new_state = {**current_state, **updates}
|
|
|
|
# Update the model
|
|
registry._update_model(entry.model, new_state)
|
|
|
|
# Trigger sync
|
|
await registry.push_all()
|
|
|
|
return {
|
|
"status": "success",
|
|
"store": store_name,
|
|
"updated_fields": list(updates.keys())
|
|
}
|
|
|
|
except Exception as e:
|
|
raise ValueError(f"Failed to update store '{store_name}': {str(e)}")
|
|
|
|
|
|
@tool
|
|
def get_store_schema(store_name: str) -> Dict[str, Any]:
|
|
"""Get the schema/structure of a synchronization store.
|
|
|
|
This shows what fields are available and their types.
|
|
|
|
Args:
|
|
store_name: Name of the store
|
|
|
|
Returns:
|
|
Dictionary describing the store's schema
|
|
|
|
Raises:
|
|
ValueError: If store_name doesn't exist
|
|
"""
|
|
registry = _get_registry()
|
|
if not registry:
|
|
raise ValueError("SyncRegistry not initialized")
|
|
|
|
entry = registry.entries.get(store_name)
|
|
if not entry:
|
|
available = list(registry.entries.keys())
|
|
raise ValueError(f"Store '{store_name}' not found. Available stores: {available}")
|
|
|
|
# Get model schema
|
|
schema = entry.model.model_json_schema()
|
|
|
|
return {
|
|
"store_name": store_name,
|
|
"schema": schema
|
|
}
|
|
|
|
|
|
SYNC_TOOLS = [
|
|
list_sync_stores,
|
|
read_sync_state,
|
|
write_sync_state,
|
|
get_store_schema
|
|
]
|