ailine/backend/app/graph/state.py

"""
主图状态定义 - 扩展版
Main Graph State Definition - Extended
"""

from enum import Enum, auto
from typing import Optional, Dict, Any, Annotated, Sequence, TypedDict
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()


@dataclass
class MainGraphState:
    """
    主图状态 - 兼容旧代码 + 新增子图功能
    
    包含：
    1. 旧代码的MessagesState兼容性字段
    2. 主图控制字段
    3. 子图结果占位
    4. 用户信息
    """
    # ========== 兼容性字段（保留旧的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  # 意图识别置信度
    
    # ========== 子图结果占位 ==========
    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)