File size: 2,363 Bytes
ac6a4ef |
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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
"""Define the state structures for the agent supervisor."""
from __future__ import annotations
from typing import Dict, List, Literal, Optional, Sequence, Any
from langchain_core.messages import AnyMessage
from langgraph.graph import MessagesState, add_messages
from typing_extensions import TypedDict, Annotated
# --- Constants and shared definitions ---------------------------------------
# Define worker types (specialized agents that perform tasks)
WORKERS = ["researcher", "coder"]
# Define all member types (including control nodes)
MEMBERS = WORKERS + ["planner", "critic", "supervisor"]
# Define status/routing options
VERDICTS = ["CORRECT", "RETRY"]
ROUTING = ["FINISH"] + WORKERS
OPTIONS = ROUTING + VERDICTS
# --- Router for supervisor decisions ---------------------------------------
class Router(TypedDict):
"""Determines which worker to route to next or if the task is complete.
The supervisor returns this structure to navigate the workflow.
Valid values are defined in the ROUTING list.
"""
next: Literal[*ROUTING]
# --- Plan structure for the Planner node -----------------------------------
class PlanStep(TypedDict):
"""A single step in the plan created by the Planner."""
worker: Literal[*WORKERS]
instruction: str
class Plan(TypedDict):
"""The complete plan produced by the Planner node."""
steps: List[PlanStep]
# --- Critic verdict structure ----------------------------------------------
class CriticVerdict(TypedDict):
"""The verdict from the Critic on whether the answer is satisfactory."""
verdict: Literal[*VERDICTS]
reason: Optional[str]
# --- State for the agent supervisor ----------------------------------------
class State(MessagesState):
"""State for the agent supervisor workflow.
Extends MessagesState which provides message history tracking.
Adds fields to track routing information, plan, and critic verdict.
"""
next: str
plan: Optional[Plan] = None
current_step_index: Optional[int] = None
draft_answer: Optional[str] = None
critic_verdict: Optional[CriticVerdict] = None
context: Dict[str, Any] = {} # Shared context accessible to all agents
worker_results: Dict[str, List[str]] = {} # Store results from each worker
|