ailine/backend/app/main_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 app.main_graph.graph import add_messages
from langchain_core.messages import BaseMessage


# ========== 新的类型 ==========
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:
    """
    - 旧代码的 MessagesState 兼容性字段
    - React 推理控制字段
    - 循环和错误处理
    - 子图结果占位
    - 用户信息
    """
    # ========== 旧 MessagesState 兼容性字段 ==========
    messages: Annotated[Sequence[BaseMessage], add_messages] = field(default_factory=list)
    llm_calls: int = 0
    memory_context: str = ""
    system_prompt: str = ""
    turns_since_last_summary: int = 0  # 新增：来自旧状态
    last_token_usage: Dict[str, Any] = field(default_factory=dict)  # 新增：来自旧状态
    last_elapsed_time: float = 0.0  # 新增：来自旧状态
    
    # ========== 主图控制字段 ==========
    user_query: str = ""
    current_action: CurrentAction = CurrentAction.NONE
    intent_confidence: float = 0.0
    
    # ========== React 推理专用字段 ==========
    reasoning_step: int = 0
    max_steps: int = 10  # 从 40 改到 10，避免过长循环
    last_action: str = ""
    reasoning_history: List[Dict[str, Any]] = field(default_factory=list)
    
    # ========== RAG 相关字段 ==========
    rag_context: str = ""
    rag_retrieved: bool = False
    rag_docs: List[Dict[str, Any]] = field(default_factory=list)
    
    # ========== 联网搜索相关字段 ==========
    web_search_results: List[str] = 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)