app/llm/state.py

"""
State definitions for the LangGraph multi-agent workflow.
"""
from typing import TypedDict, List, Dict, Optional, Annotated
import operator


class AgentState(TypedDict):
    """State for the multi-agent RAG system."""
    
    # User input
    query: str
    user_id: Optional[str]
    policy_ids: Optional[List[str]]  # NEW: Selected policy IDs
    
    # Chat history
    chat_history: Annotated[List[Dict], operator.add]
    
    # Routing decision
    agent_type: Optional[str]
    routing_reasoning: Optional[str]
    
    # Retrieved context (for RAG)
    context_chunks: Optional[List[Dict]]
    
    # Search results (for Search agent)
    search_results: Optional[List[Dict]]
    
    # Final response
    answer: Optional[str]
    sources: Optional[List]
    policy_names: Optional[List[str]]  # Policy document titles used in answer
    
    # Metadata
    metadata: Optional[Dict]
    error: Optional[str]