systemr-cli 1.0.0__py3-none-any.whl

neo/orchestrator.py

@@ -0,0 +1,288 @@
+"""Multi-agent orchestrator — parallel sub-agents for complex tasks.
+
+Mirrors Claude Code's Agent tool: the main agent spawns focused
+sub-agents that run in parallel via asyncio.gather, each with a
+focused system prompt and subset of tools.
+
+Seven pre-defined agent types cover the full trading day loop:
+SETUP → RESEARCH → PLAN → EXECUTE → DOCUMENT → ANALYZE
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from dataclasses import dataclass, field
+from decimal import Decimal
+from typing import Any, Callable
+
+import structlog
+
+from neo.streaming import ChatRequest, chat_blocking
+
+logger = structlog.get_logger(module="orchestrator")
+
+
+@dataclass
+class SubAgent:
+    """A focused sub-agent with a specific task and system prompt.
+
+    Attributes:
+        name: Human-readable agent name (shown in status display).
+        prompt: The focused task prompt sent to the LLM.
+    """
+
+    name: str
+    prompt: str
+
+    def to_request(
+        self,
+        model: str | None = None,
+        profile: str | None = None,
+        rules: str | None = None,
+    ) -> ChatRequest:
+        """Build a ChatRequest for this sub-agent.
+
+        Args:
+            model: Optional model override.
+            profile: PROFILE.md content for context.
+            rules: RULES.md content for rule enforcement.
+
+        Returns:
+            A ChatRequest configured for this agent's task.
+        """
+        return ChatRequest(
+            user_input=self.prompt,
+            model=model,
+            profile=profile,
+            rules=rules,
+            research_mode=True,
+        )
+
+
+@dataclass
+class AgentResult:
+    """Result from a sub-agent execution.
+
+    Attributes:
+        agent_name: Which agent produced this result.
+        success: Whether execution completed without error.
+        content: The agent's response text.
+        error: Error message if success is False.
+        duration_seconds: How long the agent took.
+        credits_used: Credits consumed by this agent (Decimal for precision).
+        tools_called: Number of tool calls made.
+    """
+
+    agent_name: str
+    success: bool
+    content: str
+    error: str | None = None
+    duration_seconds: float = 0.0
+    credits_used: Decimal = field(default_factory=lambda: Decimal("0"))
+    tools_called: int = 0
+
+
+# ── Pre-defined agent roster ────────────────────────────────────────
+
+PORTFOLIO_SCANNER = SubAgent(
+    name="Portfolio Scanner",
+    prompt=(
+        "Analyze my current portfolio state. Show all open positions "
+        "with current P&L, total account value, buying power, and "
+        "overall exposure percentage. Be concise — use a table format."
+    ),
+)
+
+MARKET_ANALYST = SubAgent(
+    name="Market Analyst",
+    prompt=(
+        "Give me a brief market overview for today. Cover: major index "
+        "futures (S&P, Nasdaq, Dow), VIX level, any significant overnight "
+        "moves, and key economic events on today's calendar. "
+        "Be concise — bullet points only."
+    ),
+)
+
+RISK_CALCULATOR = SubAgent(
+    name="Risk Calculator",
+    prompt=(
+        "Calculate my current risk exposure. Show: total portfolio risk "
+        "as % of account, risk per position, remaining risk budget for "
+        "today based on my daily loss limit, and how many more trades I "
+        "can take at my standard position size. Be concise."
+    ),
+)
+
+WATCHLIST_SCANNER = SubAgent(
+    name="Watchlist Scanner",
+    prompt=(
+        "Scan my watchlist for actionable setups today. For each ticker "
+        "showing a setup, identify: the pattern (breakout, pullback, "
+        "support bounce), entry zone, stop level, and estimated R:R. "
+        "Only show tickers with clear setups."
+    ),
+)
+
+STRATEGY_ANALYST = SubAgent(
+    name="Strategy Analyst",
+    prompt=(
+        "Review today's trading session. Analyze: total P&L, number of "
+        "trades, win rate, average R-multiple, risk budget used vs allowed, "
+        "any rule violations, and one key lesson. End with a discipline "
+        "score out of 10."
+    ),
+)
+
+JOURNAL_WRITER = SubAgent(
+    name="Journal Writer",
+    prompt=(
+        "Write a trade journal entry for today. Include: market conditions, "
+        "trades taken (with entry/exit/R), what went well, what could "
+        "improve, and emotional state. Keep it honest."
+    ),
+)
+
+TRADE_PLANNER = SubAgent(
+    name="Trade Planner",
+    prompt=(
+        "Based on my risk budget, current positions, and market conditions, "
+        "suggest up to 3 trade ideas for today. For each: ticker, direction, "
+        "entry, stop, target, position size, and risk amount. Ensure total "
+        "planned risk stays within my daily limit."
+    ),
+)
+
+
+# ── Orchestration ───────────────────────────────────────────────────
+
+async def run_agents(
+    agents: list[SubAgent],
+    access_token: str,
+    model: str | None = None,
+    profile: str | None = None,
+    rules: str | None = None,
+    on_start: Callable[[str], None] | None = None,
+    on_complete: Callable[[str, bool], None] | None = None,
+) -> list[AgentResult]:
+    """Run multiple sub-agents in parallel via asyncio.gather.
+
+    Each agent gets its own Bedrock API call. Results are collected
+    and returned in the same order as the input agents.
+
+    Args:
+        agents: List of sub-agents to execute.
+        access_token: Bearer token for API auth.
+        model: Optional model override for all agents.
+        profile: PROFILE.md content.
+        rules: RULES.md content.
+        on_start: Optional callback when an agent starts.
+        on_complete: Optional callback when an agent finishes.
+
+    Returns:
+        List of AgentResult objects.
+    """
+    async def _run_one(agent: SubAgent) -> AgentResult:
+        if on_start:
+            try:
+                on_start(agent.name)
+            except Exception:
+                pass
+
+        start_time = time.monotonic()
+        try:
+            request = agent.to_request(
+                model=model, profile=profile, rules=rules,
+            )
+            result = await chat_blocking(request, access_token)
+            duration = time.monotonic() - start_time
+            content = result.get(
+                "response", result.get("content", result.get("message", "")),
+            )
+            # Parse cost tracking from response
+            credits_raw = result.get("credits_used")
+            credits = Decimal(str(credits_raw)) if credits_raw is not None else Decimal("0")
+            tools = int(result.get("tools_called", 0))
+
+            if on_complete:
+                try:
+                    on_complete(agent.name, True)
+                except Exception:
+                    pass
+            logger.info(
+                "agent_completed", agent=agent.name, success=True,
+                duration=round(duration, 2), credits=str(credits), tools=tools,
+            )
+            return AgentResult(
+                agent_name=agent.name, success=True, content=content,
+                duration_seconds=duration, credits_used=credits, tools_called=tools,
+            )
+        except Exception as exc:
+            duration = time.monotonic() - start_time
+            if on_complete:
+                try:
+                    on_complete(agent.name, False)
+                except Exception:
+                    pass
+            logger.warning("agent_failed", agent=agent.name, error=str(exc), duration=round(duration, 2))
+            return AgentResult(
+                agent_name=agent.name, success=False, content="", error=str(exc),
+                duration_seconds=duration,
+            )
+
+    results = await asyncio.gather(*[_run_one(a) for a in agents])
+    return list(results)
+
+
+async def run_morning_briefing(
+    access_token: str,
+    model: str | None = None,
+    profile: str | None = None,
+    rules: str | None = None,
+    on_start: Callable[[str], None] | None = None,
+    on_complete: Callable[[str, bool], None] | None = None,
+) -> list[AgentResult]:
+    """Run the morning briefing — 4 parallel agents.
+
+    Agents: Portfolio Scanner, Market Analyst, Risk Calculator, Watchlist Scanner.
+    """
+    return await run_agents(
+        agents=[PORTFOLIO_SCANNER, MARKET_ANALYST, RISK_CALCULATOR, WATCHLIST_SCANNER],
+        access_token=access_token, model=model,
+        profile=profile, rules=rules,
+        on_start=on_start, on_complete=on_complete,
+    )
+
+
+async def run_eod_review(
+    access_token: str,
+    model: str | None = None,
+    profile: str | None = None,
+    rules: str | None = None,
+    on_start: Callable[[str], None] | None = None,
+    on_complete: Callable[[str, bool], None] | None = None,
+) -> list[AgentResult]:
+    """Run end-of-day review — Strategy Analyst + Journal Writer."""
+    return await run_agents(
+        agents=[STRATEGY_ANALYST, JOURNAL_WRITER],
+        access_token=access_token, model=model,
+        profile=profile, rules=rules,
+        on_start=on_start, on_complete=on_complete,
+    )
+
+
+async def run_trade_plan(
+    access_token: str,
+    model: str | None = None,
+    profile: str | None = None,
+    rules: str | None = None,
+    on_start: Callable[[str], None] | None = None,
+    on_complete: Callable[[str, bool], None] | None = None,
+) -> list[AgentResult]:
+    """Run trade planning — Market Analyst + Risk Calculator + Trade Planner."""
+    return await run_agents(
+        agents=[MARKET_ANALYST, RISK_CALCULATOR, TRADE_PLANNER],
+        access_token=access_token, model=model,
+        profile=profile, rules=rules,
+        on_start=on_start, on_complete=on_complete,
+    )