优化基础智能体体验

src/agents/__init__.py

@@ -4,27 +4,19 @@ class AgentManager:
     def __init__(self):
         self.agents = {}
 
-    def add_agent(self, agent_id, agent_class, configuration_class):
-        self.agents[agent_id] = {
-            "agent_class": agent_class,
-            "configuration_class": configuration_class
-        }
+    def add_agent(self, agent_id, agent_class):
+        self.agents[agent_id] = agent_class
 
     def get_runnable_agent(self, agent_id, **kwargs):
         agent_class = self.get_agent(agent_id)
-        configuration_class = self.get_configuration(agent_id)
-        configuration = configuration_class(**kwargs)
-        return agent_class(configuration)
+        return agent_class()
 
     def get_agent(self, agent_id):
-        return self.agents[agent_id]["agent_class"]
-
-    def get_configuration(self, agent_id):
-        return self.agents[agent_id]["configuration_class"]
+        return self.agents[agent_id]
 
 
 agent_manager = AgentManager()
-agent_manager.add_agent("chatbot", ChatbotAgent, ChatbotConfiguration)
+agent_manager.add_agent("chatbot", ChatbotAgent)
 
 __all__ = ["agent_manager"]
 

src/agents/chatbot/configuration.py

@@ -1,10 +1,9 @@
 from dataclasses import dataclass, field
+from datetime import datetime, timezone
 
 from langchain_core.tools import tool
 from langchain_openai import ChatOpenAI
 
-from src import config
-from src.models import select_model
 from src.agents.registry import Configuration
 
 def get_default_requirements():
@@ -17,8 +16,21 @@ def multiply(first_int: int, second_int: int) -> int:
 
 @dataclass(kw_only=True)
 class ChatbotConfiguration(Configuration):
-    requirements: list[str] = field(default_factory=get_default_requirements)
-    llm: ChatOpenAI | None = None
-    model_provider: str = "zhipu"
-    model_name: str = "glm-4-plus"
+    """Chatbot 的配置"""
+
+    system_prompt: str = field(
+        default=f"You are a helpful assistant. Now is {datetime.now(tz=timezone.utc).isoformat()}",
+        metadata={
+            "description": "The system prompt to use for the agent's interactions. "
+            "This prompt sets the context and behavior for the agent."
+        },
+    )
+
+    model: str = field(
+        default="zhipu/glm-4-plus",
+        metadata={
+            "description": "The name of the language model to use for the agent's main interactions. "
+            "Should be in the form: provider/model-name."
+        },
+    )
 

src/agents/chatbot/graph.py

@@ -7,41 +7,42 @@ from langchain_core.runnables import RunnableConfig
 from langgraph.graph import StateGraph, START, END
 from langgraph.prebuilt import ToolNode, tools_condition
 from langgraph.checkpoint.memory import MemorySaver
+from langchain_community.tools.tavily_search import TavilySearchResults
+
 
 from src.agents.registry import State, BaseAgent
+from src.agents.utils import load_chat_model
 from src.agents.chatbot.configuration import ChatbotConfiguration, multiply
 
 class ChatbotAgent(BaseAgent):
     name = "chatbot"
     description = "A chatbot that can answer questions and help with tasks."
     _graph_cache = None
-    config_schema = ChatbotConfiguration
+    config_schema = ChatbotConfiguration.to_dict()
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
     def _get_tools(self, config_schema: RunnableConfig):
         """根据配置获取工具"""
-        conf = ChatbotConfiguration.from_runnable_config(config_schema)
-        tools = [multiply]
-        if not conf:
-            return tools
-
-        if conf.get("use_web", None):
-            from langchain_community.tools.tavily_search import TavilySearchResults
-            tools.append(TavilySearchResults(max_results=10))
-
+        tools = [multiply, TavilySearchResults(max_results=10)]
         return tools
 
-    def llm_call(self, state: State, config: RunnableConfig) -> dict[str, Any]:
-        model = self.llm.bind_tools(self._get_tools(config))
+    def llm_call(self, state: State, config: RunnableConfig = None) -> dict[str, Any]:
+        """调用 llm 模型"""
+        config_schema = config or {}
+        conf = ChatbotConfiguration.from_runnable_config(config_schema)
+        model = load_chat_model(conf.model)
+        model_with_tools = model.bind_tools(self._get_tools(config_schema))
 
-        res = model.invoke(state["messages"])
+        res = model_with_tools.invoke(
+            [{"role": "system", "content": conf.system_prompt}, *state["messages"]]
+        )
         return {"messages": [res]}
 
     def get_graph(self, config_schema: RunnableConfig = None):
         """构建图"""
-        workflow = StateGraph(State)
+        workflow = StateGraph(State, config_schema=ChatbotConfiguration)
         workflow.add_node("chatbot", self.llm_call)
         workflow.add_node("tools", ToolNode(tools=self._get_tools(config_schema)))
         workflow.add_edge(START, "chatbot")
@@ -60,12 +61,13 @@ class ChatbotAgent(BaseAgent):
         for event in graph.stream({"messages": messages}, stream_mode="values", config=config_schema):
             yield event["messages"]
 
-    def stream_messages(self, messages: list[str], config: RunnableConfig = None):
-        graph = self.get_graph(config)
-        for msg, metadata in graph.stream({"messages": messages}, stream_mode="messages", config=config):
+    def stream_messages(self, messages: list[str], config_schema: RunnableConfig = None):
+        graph = self.get_graph(config_schema)
+        conf = ChatbotConfiguration.from_runnable_config(config_schema)
+        for msg, metadata in graph.stream({"messages": messages}, stream_mode="messages", config=config_schema):
             msg_type = msg.type
 
-            return_keys = config.get("configurable", {}).get("return_keys", [])
+            return_keys =conf.return_keys
             if not return_keys or msg_type in return_keys:
                 yield msg, metadata
 

src/agents/registry.py

@@ -9,6 +9,8 @@ from langchain_core.messages import BaseMessage
 from langgraph.graph.state import CompiledStateGraph
 from langgraph.graph.message import add_messages
 
+from src.config import SimpleConfig
+
 
 class State(TypedDict):
     """
@@ -22,13 +24,11 @@ class State(TypedDict):
 
 
 @dataclass(kw_only=True)
-class Configuration:
+class Configuration(SimpleConfig):
     """
     定义一个基础 Configuration 供 各类 graph 继承
     """
 
-    user_id: str = field(metadata={"description": "Unique identifier for the user."})
-
     @classmethod
     def from_runnable_config(
         cls, config: Optional[RunnableConfig] = None
@@ -40,7 +40,7 @@ class Configuration:
 
     @classmethod
     def to_dict(cls):
-        pass
+        return {f.name: getattr(cls, f.name) for f in fields(cls) if f.init}
 
 
 

src/agents/utils.py

@@ -1,8 +1,22 @@
+from src.models import select_model
 from src.agents.registry import BaseAgent
+from langchain_core.language_models import BaseChatModel
 from langchain_core.runnables import RunnableConfig
 from langchain_core.messages import AIMessageChunk, ToolMessage
 
 
+
+
+def load_chat_model(fully_specified_name: str) -> BaseChatModel:
+    """Load a chat model from a fully specified name.
+
+    Args:
+        fully_specified_name (str): String in the format 'provider/model'.
+    """
+    provider, model = fully_specified_name.split("/", maxsplit=1)
+    return select_model(model_name=model, model_provider=provider).chat_open_ai
+
+
 def agent_cli(agent: BaseAgent, config: RunnableConfig = None):
     config = config or {}
     if "configurable" not in config:

src/config/__init__.py

@@ -16,7 +16,7 @@ DEFAULT_MOCK_API = 'this_is_mock_api_key_in_frontend'
 class SimpleConfig(dict):
 
     def __key(self, key):
-        return "" if key is None else key.lower()  # 目前忘记了这里为什么要 lower 了，只能说配置项最好不要有大写的
+        return "" if key is None else key  # 目前忘记了这里为什么要 lower 了，只能说配置项最好不要有大写的
 
     def __str__(self):
         return json.dumps(self)
@@ -180,15 +180,15 @@ class Config(SimpleConfig):
         """
         获取安全的配置，即过滤掉 api_key
         """
-        
+
         config = json.loads(str(self))
-        
+
         # 过滤掉 api_key
         for model in config.get("custom_models", []):
             model["api_key"] = DEFAULT_MOCK_API if model.get("api_key") else ""
 
         return config
-    
+
     def compare_custom_models(self, value):
         """
         比较 custom_models 中的 api_key，如果输入的 api_key 与当前的 api_key 相同，则不修改

src/core/retriever.py

@@ -133,7 +133,7 @@ class Retriever:
         """重写查询"""
         model_provider = config.model_provider_lite
         model_name = config.model_name_lite
-        model = select_model(config, model_provider=model_provider, model_name=model_name)
+        model = select_model(model_provider=model_provider, model_name=model_name)
         if refs["meta"].get("mode") == "search":  # 如果是搜索模式，就使用 meta 的配置，否则就使用全局的配置
             rewrite_query_span = refs["meta"].get("use_rewrite_query", "off")
         else:
@@ -159,7 +159,7 @@ class Retriever:
         query = refs.get("rewritten_query", query)
         model_provider = config.model_provider_lite
         model_name = config.model_name_lite
-        model = select_model(config, model_provider=model_provider, model_name=model_name)
+        model = select_model(model_provider=model_provider, model_name=model_name)
 
         entities = []
         if refs["meta"].get("use_graph"):

src/models/__init__.py

@@ -1,15 +1,17 @@
 import os
 
+from src import config
 from src.utils.logging_config import logger
 from src.models.chat_model import OpenAIBase
 
 
-def select_model(config, model_provider=None, model_name=None):
-
+def select_model(model_provider=None, model_name=None):
+    """根据模型提供者选择模型"""
     model_provider = model_provider or config.model_provider
     model_info = config.model_names.get(model_provider, {})
     model_name = model_name or config.model_name or model_info.get("default", "")
 
+
     logger.info(f"Selecting model from `{model_provider}` with `{model_name}`")
 
     if model_provider in [
@@ -39,7 +41,7 @@ def select_model(config, model_provider=None, model_name=None):
         return OpenModel(model_name)
 
     elif model_provider == "custom":
-        model_info = next((x for x in config.custom_models if x["custom_id"] == model_name), None)
+        model_info = next((x for x in conf.custom_models if x["custom_id"] == model_name), None)
         if model_info is None:
             raise ValueError(f"Model {model_name} not found in custom models")
 

src/routers/chat_router.py

@@ -27,7 +27,7 @@ def chat_post(
         history: list | None = Body(None),
         thread_id: str | None = Body(None)):
 
-    model = select_model(config)
+    model = select_model()
     meta["server_model_name"] = model.model_name
     history_manager = HistoryManager(history)
     logger.debug(f"Received query: {query} with meta: {meta}")
@@ -96,7 +96,7 @@ def chat_post(
 
 @chat.post("/call")
 async def call(query: str = Body(...), meta: dict = Body(None)):
-    model = select_model(config, model_provider=meta.get("model_provider"), model_name=meta.get("model_name"))
+    model = select_model(model_provider=meta.get("model_provider"), model_name=meta.get("model_name"))
     async def predict_async(query):
         loop = asyncio.get_event_loop()
         return await loop.run_in_executor(executor, model.predict, query)
@@ -113,7 +113,7 @@ async def call(query: str = Body(...), meta: dict = Body(None)):
         loop = asyncio.get_event_loop()
         model_provider = meta.get("model_provider", config.model_provider_lite)
         model_name = meta.get("model_name", config.model_name_lite)
-        model = select_model(config, model_provider=model_provider, model_name=model_name)
+        model = select_model(model_provider=model_provider, model_name=model_name)
         return await loop.run_in_executor(executor, model.predict, query)
 
     response = await predict_async(query)
@@ -124,8 +124,9 @@ async def call(query: str = Body(...), meta: dict = Body(None)):
 @chat.get("/agent")
 async def get_agent():
     agents = [{
-        "name": agent["agent_class"].name,
-        "description": agent["agent_class"].description
+        "name": agent.name,
+        "description": agent.description,
+        "config_schema": agent.config_schema
     } for agent in agent_manager.agents.values()]
     return {"agents": agents}