55 lines
1.8 KiB
Python
55 lines
1.8 KiB
Python
|
from typing import Any, Dict, Optional
|
||
|
|
from langchain_core.messages import HumanMessage
|
||
|
|
from langgraph.graph import StateGraph, END
|
||
|
|
|
||
|
|
from services.llm_factory import create_chat_model
|
||
|
|
from .state import AgentState
|
||
|
|
from . import nodes
|
||
|
|
|
||
|
|
|
||
|
|
class BaseAgent:
|
||
|
|
"""包含通用功能的基础代理类"""
|
||
|
|
|
||
|
|
def __init__(self, model_section: Optional[str] = None):
|
||
|
|
self.model = create_chat_model(model_section)
|
||
|
|
self.graph = self._build_graph()
|
||
|
|
|
||
|
|
def _build_graph(self) -> StateGraph:
|
||
|
|
"""构建代理状态图"""
|
||
|
|
workflow = StateGraph(AgentState)
|
||
|
|
|
||
|
|
workflow.add_node("process_input", nodes.process_input)
|
||
|
|
workflow.add_node("normalize_input", self._normalize_input)
|
||
|
|
workflow.add_node("generate_response", self._generate_response)
|
||
|
|
|
||
|
|
workflow.add_edge("process_input", "normalize_input")
|
||
|
|
workflow.add_edge("normalize_input", "generate_response")
|
||
|
|
workflow.add_edge("generate_response", END)
|
||
|
|
|
||
|
|
workflow.set_entry_point("process_input")
|
||
|
|
|
||
|
|
return workflow.compile()
|
||
|
|
|
||
|
|
def _generate_response(self, state: AgentState) -> AgentState:
|
||
|
|
"""使用 LLM 生成回复"""
|
||
|
|
return nodes.generate_response(state, self.model)
|
||
|
|
|
||
|
|
def _normalize_input(self, state: AgentState) -> AgentState:
|
||
|
|
"""规范化用户输入"""
|
||
|
|
return nodes.normalize_input(state, self.model)
|
||
|
|
|
||
|
|
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", {}),
|
||
|
|
"final_step": result.get("current_step", "unknown")
|
||
|
|
}
|