ailine/backend/app/graph/state.py

"""
主图状态定义 - React 模式增强版
Main Graph State Definition - React Mode Enhanced
"""

from enum import Enum, auto
from typing import Optional, Dict, Any, Annotated, Sequence, TypedDict, List
from dataclasses import dataclass, field
from langgraph.graph import add_messages
from langchain_core.messages import BaseMessage


# ========== 兼容旧代码的类型 ==========
class MessagesState(TypedDict):
    """旧的MessagesState类型（保留兼容性）"""
    messages: Annotated[Sequence[BaseMessage], add_messages]


class GraphContext(TypedDict):
    """旧的GraphContext类型（保留兼容性）"""
    llm_calls: int
    memory_context: str
    system_prompt: str


# ========== 新的类型 ==========
class CurrentAction(Enum):
    """主图当前操作类型"""
    NONE = auto()
    GENERAL_CHAT = auto()
    NEWS_ANALYSIS = auto()
    DICTIONARY = auto()
    CONTACT = auto()


class ErrorSeverity(Enum):
    """错误严重程度"""
    INFO = auto()     # 信息级别，继续执行
    WARNING = auto()  # 警告级别，可以重试
    ERROR = auto()    # 错误级别，需要处理
    FATAL = auto()    # 致命错误，终止执行


@dataclass
class ErrorRecord:
    """错误记录"""
    error_type: str
    error_message: str
    severity: ErrorSeverity = ErrorSeverity.ERROR
    source: str = ""  # 来源：哪个节点/子图/工具
    timestamp: str = ""
    retry_count: int = 0  # 已重试次数
    max_retries: int = 3  # 最大重试次数
    context: Dict[str, Any] = field(default_factory=dict)  # 错误上下文


@dataclass
class MainGraphState:
    """
    主图状态 - React 循环推理版本
    
    包含：
    1. 旧代码的MessagesState兼容性字段
    2. React 推理控制字段
    3. 循环和错误处理
    4. 子图结果占位
    5. 用户信息
    """
    # ========== 兼容性字段（保留旧的MessagesState） ==========
    messages: Annotated[Sequence[BaseMessage], add_messages] = field(default_factory=list)
    llm_calls: int = 0
    memory_context: str = ""
    system_prompt: str = ""
    
    # ========== 主图控制字段 ==========
    user_query: str = ""  # 用户当前查询
    current_action: CurrentAction = CurrentAction.NONE  # 当前操作
    intent_confidence: float = 0.0  # 意图识别置信度
    
    # ========== React 推理专用字段 ==========
    reasoning_step: int = 0  # 当前推理步数
    max_steps: int = 40  # 最大推理步数（≤40）
    last_action: str = ""  # 上一步动作
    reasoning_history: List[Dict[str, Any]] = field(default_factory=list)  # 推理历史
    
    # ========== RAG 相关字段 ==========
    rag_context: str = ""  # RAG 检索到的上下文
    rag_retrieved: bool = False  # 是否已经检索过
    rag_docs: List[Dict[str, Any]] = field(default_factory=list)  # 检索到的文档
    
    # ========== 错误处理字段 ==========
    errors: List[ErrorRecord] = field(default_factory=list)  # 错误列表
    current_error: Optional[ErrorRecord] = None  # 当前错误
    retry_action: Optional[str] = None  # 重试动作
    
    # ========== 子图结果占位 ==========
    news_result: Optional[Dict[str, Any]] = None  # 资讯子图结果
    dictionary_result: Optional[Dict[str, Any]] = None  # 词典子图结果
    contact_result: Optional[Dict[str, Any]] = None  # 通讯录子图结果
    
    # ========== 用户信息 ==========
    user_id: str = ""
    
    # ========== 执行状态 ==========
    current_phase: str = "init"
    error_message: str = ""
    final_result: str = ""
    success: bool = False
    
    # ========== 元数据 ==========
    start_time: Optional[str] = None
    end_time: Optional[str] = None
    debug_info: Dict[str, Any] = field(default_factory=dict)