agent demo

requirements.txt

@@ -26,4 +26,9 @@ uvicorn[standard]
 fastapi
 python-multipart
 tavily-python
-rapidocr_onnxruntime
+rapidocr_onnxruntime
+langchain
+langsmith
+langgraph
+langchain-openai
+langchain-community

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

src/agents/chatbot/__init__.py

@@ -0,0 +1,3 @@
+from .graph import graph
+
+__all__ = ["graph"]

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)

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)
+

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

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,
+)
+

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)
 

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()