115 lines
4.8 KiB
Python
115 lines
4.8 KiB
Python
from typing import Annotated, Any, Mapping, Optional
|
|
from typing_extensions import NotRequired, TypedDict
|
|
from langgraph.graph import MessagesState
|
|
|
|
|
|
RESEARCH_PROVENANCE_FIELDS = (
|
|
"research_status",
|
|
"research_mode",
|
|
"timed_out_nodes",
|
|
"degraded_reason",
|
|
"covered_dimensions",
|
|
"manager_confidence",
|
|
)
|
|
|
|
|
|
def extract_research_provenance(
|
|
debate_state: Mapping[str, Any] | None,
|
|
) -> dict[str, Any] | None:
|
|
if not isinstance(debate_state, Mapping):
|
|
return None
|
|
metadata = {
|
|
key: debate_state.get(key)
|
|
for key in RESEARCH_PROVENANCE_FIELDS
|
|
if key in debate_state
|
|
}
|
|
return metadata or None
|
|
|
|
|
|
# Researcher team state
|
|
class InvestDebateState(TypedDict, total=False):
|
|
bull_history: Annotated[
|
|
str, "Bullish Conversation history"
|
|
] # Bullish Conversation history
|
|
bear_history: Annotated[
|
|
str, "Bearish Conversation history"
|
|
] # Bullish Conversation history
|
|
history: Annotated[str, "Conversation history"] # Conversation history
|
|
current_response: Annotated[str, "Latest response"] # Last response
|
|
judge_decision: Annotated[str, "Final judge decision"] # Last response
|
|
count: Annotated[int, "Length of the current conversation"] # Conversation length
|
|
research_status: NotRequired[Annotated[str, "Research stage status: full/degraded/failed"]]
|
|
research_mode: NotRequired[Annotated[str, "Research mode: debate/degraded_synthesis"]]
|
|
timed_out_nodes: NotRequired[Annotated[list[str], "Research nodes that timed out"]]
|
|
degraded_reason: NotRequired[Annotated[Optional[str], "Research degradation reason"]]
|
|
covered_dimensions: NotRequired[Annotated[list[str], "Research dimensions covered so far"]]
|
|
manager_confidence: NotRequired[Annotated[Optional[float], "Research manager confidence"]]
|
|
|
|
|
|
# Risk management team state
|
|
class RiskDebateState(TypedDict):
|
|
aggressive_history: Annotated[
|
|
str, "Aggressive Agent's Conversation history"
|
|
] # Conversation history
|
|
conservative_history: Annotated[
|
|
str, "Conservative Agent's Conversation history"
|
|
] # Conversation history
|
|
neutral_history: Annotated[
|
|
str, "Neutral Agent's Conversation history"
|
|
] # Conversation history
|
|
history: Annotated[str, "Conversation history"] # Conversation history
|
|
latest_speaker: Annotated[str, "Analyst that spoke last"]
|
|
current_aggressive_response: Annotated[
|
|
str, "Latest response by the aggressive analyst"
|
|
] # Last response
|
|
current_conservative_response: Annotated[
|
|
str, "Latest response by the conservative analyst"
|
|
] # Last response
|
|
current_neutral_response: Annotated[
|
|
str, "Latest response by the neutral analyst"
|
|
] # Last response
|
|
judge_decision: Annotated[str, "Judge's decision"]
|
|
count: Annotated[int, "Length of the current conversation"] # Conversation length
|
|
|
|
|
|
class AgentState(MessagesState):
|
|
company_of_interest: Annotated[str, "Company that we are interested in trading"]
|
|
trade_date: Annotated[str, "What date we are trading at"]
|
|
portfolio_context: Annotated[str, "Optional portfolio/account context for this analysis"]
|
|
peer_context: Annotated[str, "Optional same-theme or peer ranking context for this analysis"]
|
|
peer_context_mode: Annotated[str, "Mode describing whether peer_context is same-theme normalized or only a book snapshot"]
|
|
|
|
sender: Annotated[str, "Agent that sent this message"]
|
|
|
|
# research step
|
|
market_report: Annotated[str, "Report from the Market Analyst"]
|
|
sentiment_report: Annotated[str, "Report from the Social Media Analyst"]
|
|
news_report: Annotated[
|
|
str, "Report from the News Researcher of current world affairs"
|
|
]
|
|
fundamentals_report: Annotated[str, "Report from the Fundamentals Researcher"]
|
|
|
|
# researcher team discussion step
|
|
investment_debate_state: Annotated[
|
|
InvestDebateState, "Current state of the debate on if to invest or not"
|
|
]
|
|
investment_plan: Annotated[str, "Plan generated by the Analyst"]
|
|
investment_plan_structured: Annotated[
|
|
Mapping[str, Any], "Structured metadata extracted from the research-manager decision"
|
|
]
|
|
|
|
trader_investment_plan: Annotated[str, "Plan generated by the Trader"]
|
|
trader_investment_plan_structured: Annotated[
|
|
Mapping[str, Any], "Structured metadata extracted from the trader decision"
|
|
]
|
|
|
|
# risk management team discussion step
|
|
risk_debate_state: Annotated[
|
|
RiskDebateState, "Current state of the debate on evaluating risk"
|
|
]
|
|
final_trade_decision: Annotated[str, "Final decision made by the Risk Analysts"]
|
|
final_trade_decision_report: Annotated[str, "Human-readable final decision report"]
|
|
final_trade_decision_structured: Annotated[
|
|
Mapping[str, Any], "Structured metadata extracted from the portfolio-manager decision"
|
|
]
|