diff --git a/requirements.txt b/requirements.txt index 8e77b9a7..790c4f8a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -26,4 +26,9 @@ uvicorn[standard] fastapi python-multipart tavily-python -rapidocr_onnxruntime \ No newline at end of file +rapidocr_onnxruntime +langchain +langsmith +langgraph +langchain-openai +langchain-community \ No newline at end of file diff --git a/src/agents/__init__.py b/src/agents/__init__.py new file mode 100644 index 00000000..3759c5b9 --- /dev/null +++ b/src/agents/__init__.py @@ -0,0 +1,8 @@ + + +class AgentManager: + def __init__(self): + self.agents = {} + + def add_agent(self, agent_id, agent): + self.agents[agent_id] = agent diff --git a/src/agents/chatbot/__init__.py b/src/agents/chatbot/__init__.py new file mode 100644 index 00000000..92e9d0b0 --- /dev/null +++ b/src/agents/chatbot/__init__.py @@ -0,0 +1,3 @@ +from .graph import graph + +__all__ = ["graph"] \ No newline at end of file diff --git a/src/agents/chatbot/configuration.py b/src/agents/chatbot/configuration.py new file mode 100644 index 00000000..bdd8b788 --- /dev/null +++ b/src/agents/chatbot/configuration.py @@ -0,0 +1,16 @@ +from dataclasses import dataclass, field +from typing import List + +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 + + +@dataclass(kw_only=True) +class ChatbotConfiguration(Configuration): + name: str = "chatbot" + llm: ChatOpenAI = field(default_factory=select_model().chat_open_ai) + tools: List[Tool] = field(default_factory=list) \ No newline at end of file diff --git a/src/agents/chatbot/graph.py b/src/agents/chatbot/graph.py new file mode 100644 index 00000000..99b00695 --- /dev/null +++ b/src/agents/chatbot/graph.py @@ -0,0 +1,35 @@ +"""Define a simple chatbot agent. + +This agent returns a predefined response without using an actual LLM. +""" + +from typing import Any, Dict + +from langchain_core.runnables import RunnableConfig +from langgraph.graph import StateGraph + +from src.agents.registry import State, BaseAgent +from src.agents.chatbot.configuration import ChatbotConfiguration +from src.agents.tools_factory import search_tool + +class ChatbotAgent(BaseAgent): + def __init__(self, configuration: ChatbotConfiguration): + super().__init__(configuration) + self.llm = configuration.llm + self.tools = [search_tool] + + async def _get_tools(self): + return self.tools + + async def llm_call(self, state: State) -> Dict[str, Any]: + return await self.llm.invoke(state.messages) + + async def get_graph(self): + workflow = StateGraph(State, config_schema=ChatbotConfiguration) + workflow.add_node("chat", self.chat) + workflow.add_edge("__start__", "chat") + return workflow.compile() + + async def invoke(self, messages: List[str], config: RunnableConfig) -> Dict[str, Any]: + return await self.get_graph().invoke({"messages": messages}, config) + diff --git a/src/agents/registry.py b/src/agents/registry.py new file mode 100644 index 00000000..04915770 --- /dev/null +++ b/src/agents/registry.py @@ -0,0 +1,39 @@ +from typing import Type, Annotated, Optional, TypedDict +from abc import abstractmethod + +from langchain_core.tools import BaseTool +from langchain_openai import ChatOpenAI +from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.graph.message import add_messages +from langchain_core.messages import BaseMessage, ToolMessage, AIMessage, filter_messages +from langgraph.graph.state import CompiledStateGraph + + +from dataclasses import dataclass +class State(TypedDict): + """ + 定义一个基础 State 供 各类 graph 继承, 其中: + 1. messages 为所有 graph 的核心信息队列, 所有聊天工作流均应该将关键信息补充到此队列中; + 2. history 为所有工作流单次启动时获取 history_len 的 messages 所用(节约成本, 及防止单轮对话 tokens 占用长度达到 llm 支持上限), + history 中的信息理应是可以被丢弃的. + """ + messages: Annotated[list[BaseMessage], add_messages] + history: Optional[list[BaseMessage]] + + +@dataclass(kw_only=True) +class Configuration: + """ + 定义一个基础 Configuration 供 各类 graph 继承 + """ + llm: ChatOpenAI + + +class BaseAgent(): + + def __init__(self, configuration: Configuration): + self.configuration = configuration + + @abstractmethod + def get_graph(self) -> CompiledStateGraph: + pass \ No newline at end of file diff --git a/src/agents/tools_factory.py b/src/agents/tools_factory.py new file mode 100644 index 00000000..b9f2780e --- /dev/null +++ b/src/agents/tools_factory.py @@ -0,0 +1,16 @@ +from pydantic import BaseModel, Field +from langchain_core.tools import Tool, StructuredTool + +from src.utils.web_search import WebSearcher + +class SearchArgsSchema(BaseModel): + query: str = Field(..., description="The query to search the web for") + +search_tool = Tool( + name="search", + description="Search the web for information", + func=WebSearcher().search, + args_schema=SearchArgsSchema, + return_direct=True, +) + diff --git a/src/core/knowledgebase.py b/src/core/knowledgebase.py index 43cb00c3..dd2ac706 100644 --- a/src/core/knowledgebase.py +++ b/src/core/knowledgebase.py @@ -275,13 +275,6 @@ class KnowledgeBase: #* Below is the code for retriever # ################################### - def get_retriever(self, db_id): - db = self.get_kb_by_id(db_id) - if db is None: - raise Exception(f"database not found, {db_id}") - - return db.retriever - def query(self, query, db_id, **kwargs): db = self.get_kb_by_id(db_id) diff --git a/src/models/chat_model.py b/src/models/chat_model.py index fe5c9c82..3a671e21 100644 --- a/src/models/chat_model.py +++ b/src/models/chat_model.py @@ -1,6 +1,7 @@ import os from openai import OpenAI from src.utils import logger, get_docker_safe_url +from langchain_openai import ChatOpenAI class OpenAIBase(): def __init__(self, api_key, base_url, model_name): @@ -8,7 +9,9 @@ class OpenAIBase(): self.base_url = base_url self.client = OpenAI(api_key=api_key, base_url=base_url) self.model_name = model_name - # logger.debug(f"{self.get_models()=}") + self.chat_open_ai = ChatOpenAI(model=model_name, + api_key=api_key, + base_url=base_url) def predict(self, message, stream=False): if isinstance(message, str): @@ -37,7 +40,7 @@ class OpenAIBase(): stream=False, ) return response.choices[0].message - + def get_models(self): try: return self.client.models.list()