File size: 1,404 Bytes
8ff817c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
"""

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