diff --git a/src/agents/__init__.py b/src/agents/__init__.py index 6add01ba..adec7d99 100644 --- a/src/agents/__init__.py +++ b/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"] diff --git a/src/agents/chatbot/configuration.py b/src/agents/chatbot/configuration.py index 2c4fafd1..2296f354 100644 --- a/src/agents/chatbot/configuration.py +++ b/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." + }, + ) diff --git a/src/agents/chatbot/graph.py b/src/agents/chatbot/graph.py index 1c08bbc9..51c362d9 100644 --- a/src/agents/chatbot/graph.py +++ b/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 diff --git a/src/agents/registry.py b/src/agents/registry.py index 714725df..7c365ed2 100644 --- a/src/agents/registry.py +++ b/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} diff --git a/src/agents/utils.py b/src/agents/utils.py index 4f2bb20b..9da3ee33 100644 --- a/src/agents/utils.py +++ b/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: diff --git a/src/config/__init__.py b/src/config/__init__.py index 0682393d..6939af8b 100644 --- a/src/config/__init__.py +++ b/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 相同,则不修改 diff --git a/src/core/retriever.py b/src/core/retriever.py index c5c2118d..3181a870 100644 --- a/src/core/retriever.py +++ b/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"): diff --git a/src/models/__init__.py b/src/models/__init__.py index 99655a2c..9f4c0d8c 100644 --- a/src/models/__init__.py +++ b/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") diff --git a/src/routers/chat_router.py b/src/routers/chat_router.py index d6d2844a..977ecd27 100644 --- a/src/routers/chat_router.py +++ b/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}