149 lines
5.6 KiB
Python
149 lines
5.6 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
from tradingagents.services.account import AccountService, AccountSnapshot
|
|
from tradingagents.services.auto_trade import AutoTradeResult, AutoTradeService, TickerDecision
|
|
from tradingagents.services.autopilot_events import AutopilotEventQueue, HypothesisEvent, _utcnow
|
|
from tradingagents.services.hypothesis_store import HypothesisRecord, HypothesisStore, PlanStepRecord
|
|
|
|
|
|
@dataclass
|
|
class ProcessedEvent:
|
|
event: HypothesisEvent
|
|
status: str
|
|
message: str
|
|
|
|
|
|
class AutopilotWorker:
|
|
"""Processes hypothesis events and runs focused reevaluations."""
|
|
|
|
def __init__(
|
|
self,
|
|
results_root: Path,
|
|
auto_trader: AutoTradeService,
|
|
account_service: AccountService,
|
|
) -> None:
|
|
self.results_root = Path(results_root)
|
|
self.store = HypothesisStore(self.results_root / "hypotheses")
|
|
self.queue = AutopilotEventQueue(self.results_root / "autopilot")
|
|
self.auto_trader = auto_trader
|
|
self.account_service = account_service
|
|
|
|
def enqueue_event(self, hypothesis_id: str, event_type: str, payload: Optional[Dict[str, Any]] = None) -> HypothesisEvent:
|
|
event = HypothesisEvent(
|
|
id=f"evt_{hypothesis_id}_{event_type}_{_utcnow()}",
|
|
hypothesis_id=hypothesis_id,
|
|
event_type=event_type,
|
|
payload=payload or {},
|
|
created_at=_utcnow(),
|
|
)
|
|
self.queue.enqueue(event)
|
|
return event
|
|
|
|
def list_events(self) -> List[HypothesisEvent]:
|
|
return self.queue.list()
|
|
|
|
def process_all(self) -> List[ProcessedEvent]:
|
|
events = self.queue.dequeue_all()
|
|
processed: List[ProcessedEvent] = []
|
|
for event in events:
|
|
status, message = self._handle_event(event)
|
|
processed.append(ProcessedEvent(event=event, status=status, message=message))
|
|
return processed
|
|
|
|
def _handle_event(self, event: HypothesisEvent) -> Tuple[str, str]:
|
|
record = self.store.get(event.hypothesis_id)
|
|
if not record:
|
|
return ("skipped", f"Hypothesis {event.hypothesis_id} not found")
|
|
|
|
step = record.next_open_step()
|
|
if not step:
|
|
record.status = "completed"
|
|
record.updated_at = _utcnow()
|
|
self.store.upsert(record)
|
|
return ("completed", "Hypothesis already completed; marked as completed")
|
|
|
|
step.status = "done"
|
|
step.metadata.setdefault("events", []).append(
|
|
{
|
|
"type": event.event_type,
|
|
"payload": event.payload,
|
|
"timestamp": event.created_at,
|
|
}
|
|
)
|
|
record.updated_at = _utcnow()
|
|
if record.next_open_step() is None:
|
|
record.status = "completed"
|
|
self.store.upsert(record)
|
|
|
|
reevaluation_msg = self._reevaluate(record, event)
|
|
return ("updated", f"Marked step '{step.description}' as done. {reevaluation_msg}")
|
|
|
|
def _reevaluate(self, record: HypothesisRecord, event: HypothesisEvent) -> str:
|
|
try:
|
|
snapshot = self.account_service.refresh()
|
|
except Exception as exc:
|
|
return f"Failed to refresh account snapshot: {exc}"
|
|
|
|
try:
|
|
result = self.auto_trader.run(
|
|
snapshot,
|
|
focus_override=[record.ticker],
|
|
allow_market_closed=True,
|
|
)
|
|
except Exception as exc:
|
|
return f"Auto-trade reevaluation failed: {exc}"
|
|
|
|
decision = self._extract_decision(result, record.ticker)
|
|
if not decision:
|
|
return "No decision returned for ticker"
|
|
|
|
self._apply_decision(record, decision)
|
|
record.updated_at = _utcnow()
|
|
self.store.upsert(record)
|
|
return f"Reevaluated with action {record.action}"
|
|
|
|
def _extract_decision(self, result: AutoTradeResult, ticker: str) -> Optional[TickerDecision]:
|
|
for decision in result.decisions:
|
|
if decision.ticker.upper() == ticker.upper():
|
|
return decision
|
|
return None
|
|
|
|
def _apply_decision(self, record: HypothesisRecord, decision: TickerDecision) -> None:
|
|
record.action = (decision.final_decision or decision.immediate_action or record.action).upper()
|
|
record.priority = decision.priority
|
|
record.notes = decision.final_notes or decision.sequential_plan.notes or record.notes
|
|
record.plan = self._build_plan_from_decision(record.ticker, decision)
|
|
record.triggers = decision.action_queue or record.triggers
|
|
record.status = "monitoring"
|
|
if getattr(decision, "strategy", None):
|
|
record.strategy = decision.strategy.to_dict()
|
|
|
|
def _build_plan_from_decision(self, ticker: str, decision: TickerDecision) -> List[PlanStepRecord]:
|
|
steps: List[PlanStepRecord] = []
|
|
actions = decision.sequential_plan.actions or []
|
|
for idx, action in enumerate(actions, 1):
|
|
steps.append(
|
|
PlanStepRecord(
|
|
id=f"{ticker.lower()}_{idx}",
|
|
description=str(action),
|
|
status="pending",
|
|
metadata={
|
|
"next_decision": decision.sequential_plan.next_decision,
|
|
"source": "autopilot",
|
|
},
|
|
)
|
|
)
|
|
if not steps:
|
|
steps.append(
|
|
PlanStepRecord(
|
|
id=f"{ticker.lower()}_plan",
|
|
description=f"Monitor hypothesis for {ticker}",
|
|
status="pending",
|
|
)
|
|
)
|
|
return steps
|