62 lines
2.4 KiB
Python
62 lines
2.4 KiB
Python
import logging
|
|
from collections import defaultdict
|
|
from typing import Generator
|
|
from api.store.runs_store import RunsStore
|
|
from api.models.run import RunConfig, RunStatus
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
try:
|
|
from tradingagents.graph.trading_graph import TradingAgentsGraph
|
|
from tradingagents.default_config import DEFAULT_CONFIG
|
|
except ImportError:
|
|
TradingAgentsGraph = None # type: ignore
|
|
DEFAULT_CONFIG = {}
|
|
|
|
|
|
class RunService:
|
|
def __init__(self, store: RunsStore):
|
|
self._store = store
|
|
|
|
def stream_events(self, run_id: str) -> Generator[dict, None, None]:
|
|
run = self._store.get(run_id)
|
|
if not run or not run.config:
|
|
yield {"event": "run:error", "data": {"message": "Run not found"}}
|
|
return
|
|
|
|
self._store.update_status(run_id, RunStatus.RUNNING)
|
|
config = run.config
|
|
|
|
ta_config = DEFAULT_CONFIG.copy()
|
|
ta_config["llm_provider"] = config.llm_provider
|
|
ta_config["deep_think_llm"] = config.deep_think_llm
|
|
ta_config["quick_think_llm"] = config.quick_think_llm
|
|
ta_config["max_debate_rounds"] = config.max_debate_rounds
|
|
ta_config["max_risk_discuss_rounds"] = config.max_risk_discuss_rounds
|
|
|
|
try:
|
|
ta = TradingAgentsGraph(
|
|
debug=False,
|
|
config=ta_config,
|
|
selected_analysts=config.enabled_analysts or
|
|
["market", "news", "fundamentals", "social"],
|
|
)
|
|
|
|
turn_counts: defaultdict[str, int] = defaultdict(int)
|
|
|
|
for step_key, report in ta.stream_propagate(config.ticker, config.date):
|
|
turn = turn_counts[step_key]
|
|
yield {"event": "agent:start", "data": {"step": step_key, "turn": turn}}
|
|
yield {"event": "agent:complete", "data": {"step": step_key, "turn": turn, "report": report}}
|
|
self._store.add_report(run_id, f"{step_key}:{turn}", report)
|
|
turn_counts[step_key] += 1
|
|
|
|
decision = ta._last_decision or "HOLD"
|
|
self._store.update_decision(run_id, decision)
|
|
self._store.update_status(run_id, RunStatus.COMPLETE)
|
|
yield {"event": "run:complete", "data": {"decision": decision, "run_id": run_id}}
|
|
|
|
except Exception as e:
|
|
self._store.set_error(run_id, str(e))
|
|
yield {"event": "run:error", "data": {"message": str(e)}}
|