more_dots/agent/conversation.py

from typing import Dict, Any, List, Optional
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langgraph.graph import StateGraph, END

from .graph import BaseAgent
from .state import AgentState


class ConversationAgent(BaseAgent):
    """处理多轮对话的代理"""

    def __init__(self, model_section: Optional[str] = None):
        super().__init__(model_section)
        self.conversation_history: List[BaseMessage] = []

    def _build_graph(self) -> StateGraph:
        """构建对话专用图"""
        workflow = StateGraph(AgentState)

        workflow.add_node("analyze_intent", self._analyze_intent)
        workflow.add_node("normalize_input", self._normalize_input)
        workflow.add_node("generate_sql", self._generate_sql)
        workflow.add_node("generate_response", self._generate_response)
        workflow.add_node("update_context", self._update_context)

        workflow.add_edge("analyze_intent", "normalize_input")
        workflow.add_edge("normalize_input", "generate_sql")
        workflow.add_edge("generate_sql", "generate_response")
        workflow.add_edge("generate_response", "update_context")
        workflow.add_edge("update_context", END)

        workflow.set_entry_point("analyze_intent")

        return workflow.compile()

    def _analyze_intent(self, state: AgentState) -> AgentState:
        """分析用户意图与对话上下文"""
        user_message = state.messages[-1] if state.messages else None

        if user_message and isinstance(user_message, HumanMessage):
            content = user_message.content.lower()

            if any(word in content for word in ["hello", "hi", "hey", "greetings"]):
                state.context["intent"] = "greeting"
            elif any(word in content for word in ["help", "assist", "support"]):
                state.context["intent"] = "help"
            elif "?" in content:
                state.context["intent"] = "question"
            else:
                state.context["intent"] = "general"

        state.current_step = "intent_analyzed"
        return state

    def _generate_response(self, state: AgentState) -> AgentState:
        """优先返回 SQL 执行结果，其次返回生成 SQL，再回退到模型回复"""
        from . import nodes
        state = nodes.generate_response(state, self.model)
        state.current_step = "response_generated"
        return state

    def _generate_sql(self, state: AgentState) -> AgentState:
        """生成 SQL"""
        from . import nodes
        return nodes.generate_sql(state, self.model)

    def _update_context(self, state: AgentState) -> AgentState:
        """更新对话上下文与历史"""
        for message in state.messages:
            if isinstance(message, (HumanMessage, AIMessage)):
                self.conversation_history.append(message)

        if len(self.conversation_history) > 10:
            self.conversation_history = self.conversation_history[-10:]

        state.current_step = "context_updated"
        return state

    def run(self, user_input: str, **kwargs) -> Dict[str, Any]:
        """运行对话并维护历史"""
        initial_state = AgentState(
            messages=[HumanMessage(content=user_input)],
            context=kwargs
        )

        result = self.graph.invoke(initial_state)

        return {
            "messages": result.get("messages", []),
            "context": result.get("context", {}),
            "conversation_history": self.conversation_history,
            "final_step": result.get("current_step", "unknown")
        }

    def stream_run(self, user_input: str):
        """流式运行对话并维护历史"""
        all_messages = self.conversation_history + [HumanMessage(content=user_input)]
        full_text = ""

        for chunk in self.model.stream(all_messages):
            if hasattr(chunk, "content") and chunk.content:
                full_text += chunk.content
                yield chunk.content

        self.conversation_history.append(HumanMessage(content=user_input))
        self.conversation_history.append(AIMessage(content=full_text))

        if len(self.conversation_history) > 10:
            self.conversation_history = self.conversation_history[-10:]