source/src/state/graph_state.py

"""
Graph State với conversation cache support - LangGraph compatible
"""
from typing import Optional, Dict, Any, List
from typing_extensions import TypedDict

class TransportationState(TypedDict):
    # Input/Output
    user_message: str
    ai_response: str
    
    # Conversation cache  
    conversation_cache: List[Dict[str, Any]]
    
    # Function tracking
    function_calls_made: List[Dict[str, Any]]
    
    # Error handling
    error_message: Optional[str]
    
    # Status
    current_step: str

def create_initial_state(user_message: str = "") -> TransportationState:
    """Create initial state with defaults"""
    return TransportationState(
        user_message=user_message,
        ai_response="",
        conversation_cache=[],
        function_calls_made=[],
        error_message=None,
        current_step="input"
    )

def add_function_call(state: TransportationState, function_name: str, arguments: Dict[str, Any], result: Dict[str, Any]) -> TransportationState:
    """Add function call to state"""
    state["function_calls_made"].append({
        "function_name": function_name,
        "arguments": arguments,
        "result": result
    })
    return state

def set_error(state: TransportationState, error: str) -> TransportationState:
    """Set error in state"""
    state["error_message"] = error
    return state