agent_states/agent_state.py

#!/usr/bin/env python3
"""
Agent state management for the LangGraph cyber-legal assistant
"""

from typing import TypedDict, List, Dict, Any, Optional
from datetime import datetime


class AgentState(TypedDict):
    """
    State definition for the LangGraph agent workflow
    """
    # User interaction
    user_query: str
    user_id: Optional[str]
    conversation_history: List[Dict[str, str]]
    intermediate_steps: List[Dict[str, Any]]
    system_prompt: Optional[str]
    
    # Context processing
    relevant_documents: List[str]
    
    # Metadata
    query_timestamp: str
    processing_time: Optional[float]
    jurisdiction: Optional[str]