114 lines
4.2 KiB
Python
114 lines
4.2 KiB
Python
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:
|
|
"""结合对话历史生成回复"""
|
|
all_messages = self.conversation_history + state.messages
|
|
|
|
if all_messages:
|
|
response = self.model.invoke(all_messages)
|
|
state.messages.append(response)
|
|
|
|
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:]
|