apps/api/app/graph/state.py

"""Typed, reducer-merged shared state. State design *is* the architecture.

The data agent and the news agent run in the same LangGraph superstep. Both write
into the same state object concurrently, so every field either belongs to exactly
one writer or carries an explicit reducer that merges concurrent updates without
losing either side.

Fields written by **both** workers — ``errors``, ``tool_calls``, ``attempts``,
``token_spend``, ``agents_completed`` — all have reducers. That is what makes the
parallel fan-out race-free: there is no read-modify-write anywhere, only
commutative merges applied by the graph runtime.
"""

from __future__ import annotations

import operator
from collections.abc import Mapping, Sequence
from typing import Annotated, Any, TypedDict

from app.models.brief import Brief
from app.models.market import Fundamentals, Metrics, NewsFeed, PriceHistory, RiskEvent, Sentiment
from app.models.run import HumanDecision, RunError, RunMode, RunStatus, TokenSpend

#: Hard cap so a pathological run cannot grow state without bound.
MAX_TOOL_CALL_RECORDS = 2000
MAX_ERRORS = 500


# --------------------------------------------------------------------------- #
# Reducers                                                                     #
# --------------------------------------------------------------------------- #


def merge_mapping[VT](
    left: Mapping[str, VT] | None, right: Mapping[str, VT] | None
) -> dict[str, VT]:
    """Last-writer-wins per key. Disjoint keys from parallel agents both survive."""
    if not left:
        return dict(right or {})
    if not right:
        return dict(left)
    merged = dict(left)
    merged.update(right)
    return merged


def merge_counters(
    left: Mapping[str, int] | None, right: Mapping[str, int] | None
) -> dict[str, int]:
    """Per-key integer addition — used for per-ticker attempt counts."""
    merged: dict[str, int] = dict(left or {})
    for key, value in (right or {}).items():
        merged[key] = merged.get(key, 0) + value
    return merged


def append_errors(
    left: Sequence[RunError] | None, right: Sequence[RunError] | None
) -> list[RunError]:
    """Append while de-duplicating identical (stage, ticker, message) triples."""
    merged: list[RunError] = list(left or [])
    seen = {error.key for error in merged}
    for error in right or []:
        if error.key in seen:
            continue
        seen.add(error.key)
        merged.append(error)
    return merged[-MAX_ERRORS:]


def append_records(
    left: Sequence[dict[str, Any]] | None, right: Sequence[dict[str, Any]] | None
) -> list[dict[str, Any]]:
    """Plain append for telemetry rows, bounded to keep state small."""
    merged = list(left or []) + list(right or [])
    return merged[-MAX_TOOL_CALL_RECORDS:]


def append_risk_events(
    left: Sequence[RiskEvent] | None, right: Sequence[RiskEvent] | None
) -> list[RiskEvent]:
    """Append de-duplicating on (ticker, headline)."""
    merged: list[RiskEvent] = list(left or [])
    seen = {(event.ticker, event.headline) for event in merged}
    for event in right or []:
        key = (event.ticker, event.headline)
        if key in seen:
            continue
        seen.add(key)
        merged.append(event)
    return merged


def append_unique_strings(left: Sequence[str] | None, right: Sequence[str] | None) -> list[str]:
    """Ordered set union."""
    merged = list(left or [])
    seen = set(merged)
    for value in right or []:
        if value not in seen:
            seen.add(value)
            merged.append(value)
    return merged


def merge_spend(left: TokenSpend | None, right: TokenSpend | None) -> TokenSpend:
    """Field-wise addition of token spend from concurrent model calls."""
    if left is None:
        return right or TokenSpend()
    if right is None:
        return left
    return left.plus(right)


def take_last[VT](left: VT | None, right: VT | None) -> VT | None:
    """Last non-None write wins. For fields only one node ever sets."""
    return right if right is not None else left


# --------------------------------------------------------------------------- #
# State                                                                        #
# --------------------------------------------------------------------------- #


class RunState(TypedDict, total=False):
    """Shared state for one brief."""

    # --- immutable run identity -------------------------------------------
    run_id: str
    watchlist_key: str
    tickers: list[str]
    mode: RunMode
    days: int
    news_limit: int
    session_date: str
    max_regenerations: int

    # --- reducer-merged worker output --------------------------------------
    prices: Annotated[dict[str, PriceHistory], merge_mapping]
    fundamentals: Annotated[dict[str, Fundamentals], merge_mapping]
    metrics: Annotated[dict[str, Metrics], merge_mapping]
    news: Annotated[dict[str, NewsFeed], merge_mapping]
    sentiment: Annotated[dict[str, Sentiment], merge_mapping]
    risk_events: Annotated[list[RiskEvent], append_risk_events]

    # --- reducer-merged cross-cutting --------------------------------------
    errors: Annotated[list[RunError], append_errors]
    tool_calls: Annotated[list[dict[str, Any]], append_records]
    attempts: Annotated[dict[str, int], merge_counters]
    agents_completed: Annotated[list[str], append_unique_strings]
    iterations: Annotated[int, operator.add]
    token_spend: Annotated[TokenSpend, merge_spend]

    # --- single-writer fields ----------------------------------------------
    plan: dict[str, Any]
    brief: Brief | None
    verification: dict[str, Any] | None
    regenerations: int
    status: RunStatus
    decision: HumanDecision | None
    delivery: dict[str, Any] | None
    abort_reason: str | None


def initial_state(
    *,
    run_id: str,
    watchlist_key: str,
    tickers: Sequence[str],
    session_date: str,
    mode: RunMode = RunMode.STANDARD,
    days: int = 120,
    news_limit: int = 6,
    max_regenerations: int = 1,
) -> RunState:
    """A fully-populated starting state — every channel initialised."""
    return RunState(
        run_id=run_id,
        watchlist_key=watchlist_key,
        tickers=list(tickers),
        mode=mode,
        days=days,
        news_limit=news_limit,
        session_date=session_date,
        max_regenerations=max_regenerations,
        prices={},
        fundamentals={},
        metrics={},
        news={},
        sentiment={},
        risk_events=[],
        errors=[],
        tool_calls=[],
        attempts={},
        agents_completed=[],
        iterations=0,
        token_spend=TokenSpend(),
        plan={},
        brief=None,
        verification=None,
        regenerations=0,
        status=RunStatus.RUNNING,
        decision=None,
        delivery=None,
        abort_reason=None,
    )


# --------------------------------------------------------------------------- #
# Completeness helpers — used by the supervisor to decide when to route on      #
# --------------------------------------------------------------------------- #


def has_market_data(state: RunState, ticker: str) -> bool:
    metrics = state.get("metrics", {}).get(ticker)
    return metrics is not None and metrics.ok


def has_news_data(state: RunState, ticker: str) -> bool:
    return ticker in state.get("sentiment", {})


def market_attempted(state: RunState, ticker: str) -> bool:
    return f"market:{ticker}" in state.get("attempts", {})


def news_attempted(state: RunState, ticker: str) -> bool:
    return f"news:{ticker}" in state.get("attempts", {})


def attempts_for(state: RunState, kind: str, ticker: str) -> int:
    return state.get("attempts", {}).get(f"{kind}:{ticker}", 0)


def pending_market_tickers(state: RunState, max_attempts: int = 2) -> list[str]:
    """Tickers still missing market data that have retries left."""
    return [
        ticker
        for ticker in state.get("tickers", [])
        if not has_market_data(state, ticker)
        and attempts_for(state, "market", ticker) < max_attempts
    ]


def pending_news_tickers(state: RunState, max_attempts: int = 2) -> list[str]:
    """Tickers still missing news/sentiment that have retries left."""
    return [
        ticker
        for ticker in state.get("tickers", [])
        if not has_news_data(state, ticker) and attempts_for(state, "news", ticker) < max_attempts
    ]


def completeness(state: RunState) -> dict[str, float]:
    """Per-ticker completeness in [0, 1] — drives the UI progress bars."""
    result: dict[str, float] = {}
    for ticker in state.get("tickers", []):
        score = 0.0
        if has_market_data(state, ticker):
            score += 0.6
        if has_news_data(state, ticker):
            score += 0.4
        result[ticker] = round(score, 2)
    return result


def usable_tickers(state: RunState) -> list[str]:
    """Tickers with enough data to appear as a full snapshot row."""
    return [t for t in state.get("tickers", []) if has_market_data(state, t)]


def is_partial(state: RunState) -> bool:
    """True when any ticker failed to produce market data."""
    return len(usable_tickers(state)) < len(state.get("tickers", []))