Spaces:
Sleeping
Sleeping
File size: 1,724 Bytes
a2cbcac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | from typing import TypedDict, Optional, Dict, Any, List
from datetime import datetime
class AgentState(TypedDict):
"""
Unified state structure for all PrimoAgent agents.
Used by Data Collection, Technical Analysis, News Intelligence, and Portfolio Manager.
"""
# Core workflow
session_id: str
symbols: List[str]
current_step: str
analysis_date: str # Date for analysis in YYYY-MM-DD format
# Agent results (simple dictionaries)
data_collection_results: Optional[Dict[str, Any]]
technical_analysis_results: Optional[Dict[str, Any]]
news_intelligence_results: Optional[Dict[str, Any]]
portfolio_manager_results: Optional[Dict[str, Any]]
# Simple error handling
error: Optional[str]
def create_initial_state(session_id: str, symbols: List[str], analysis_date: Optional[str] = None) -> AgentState:
"""Create initial workflow state."""
# If no date provided, use today's date
if analysis_date is None:
analysis_date = datetime.now().strftime("%Y-%m-%d")
return AgentState(
session_id=session_id,
symbols=symbols,
analysis_date=analysis_date,
current_step="initialized",
data_collection_results=None,
technical_analysis_results=None,
news_intelligence_results=None,
portfolio_manager_results=None,
error=None
)
def update_step(state: AgentState, step: str) -> AgentState:
"""Update current step."""
state["current_step"] = step
return state
def set_error(state: AgentState, error_message: str) -> AgentState:
"""Set error in state."""
state["error"] = error_message
state["current_step"] = "error"
return state |