src/utils/state.py

"""
State schema for DS-STAR multi-agent system.
Defines the centralized state structure shared across all agents.
"""

import operator
from typing import Annotated, List, TypedDict

from langchain_core.messages import AIMessage, HumanMessage


class PlanStep(TypedDict):
    """Individual plan step with number and description"""

    step_number: int
    description: str


class DSStarState(TypedDict):
    """
    Centralized state for DS-STAR pipeline.

    This state is passed between all agents in the graph.
    Uses reducer pattern for message accumulation.
    """

    # User input
    query: str

    # Data file descriptions (from Stage 1 - Analyzer)
    data_descriptions: dict[str, str]

    # Current plan - list of completed steps
    # NO REDUCER - we need full control for backtracking
    plan: List[PlanStep]

    # Code and execution
    current_code: str
    execution_result: str

    # Verification and routing
    is_sufficient: bool
    router_decision: str  # "Add Step" or "Step N"

    # Iteration tracking
    iteration: int
    max_iterations: int

    # Messages for agent communication (accumulated with reducer)
    messages: Annotated[List[HumanMessage | AIMessage], operator.add]

    # Next node routing (for internal control flow)
    next: str

    # LLM instance (shared across all agents)
    llm: object