Merge pull request #102 from ant-xuexiao/feat_tool_knowledge

feat: add the knowledge tool

server/agent/stream.py

@@ -1,71 +1,82 @@
-import datetime
 import json
-import os
 import uuid
-from langchain.tools import tool
 from typing import AsyncIterator
 from langchain.agents import AgentExecutor
 from data_class import ChatData, Message
 from langchain.agents.format_scratchpad.openai_tools import (
     format_to_openai_tool_messages,
 )
+from langchain_core.messages import AIMessage, FunctionMessage, HumanMessage
 from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser
-from langchain.prompts import MessagesPlaceholder
+from langchain.prompts import PromptTemplate, MessagesPlaceholder
 from langchain_core.utils.function_calling import convert_to_openai_tool
 from langchain_core.prompts import ChatPromptTemplate
 from langchain.utilities.tavily_search import TavilySearchAPIWrapper
 from langchain.tools.tavily_search import TavilySearchResults
 from langchain_openai import ChatOpenAI
 from uilts.env import get_env_variable
-from tools import issue
-from tools import sourcecode
-from langchain_core.messages import AIMessage, FunctionMessage, HumanMessage
-
+from tools import issue, sourcecode, knowledge
 TAVILY_API_KEY =  get_env_variable("TAVILY_API_KEY")
 
+prompt_template = """
+# Character
+You are a skilled assistant dedicated to Ant Design, capable of delivering comprehensive insights and solutions pertaining to Ant Design. You excel in fixing code issues correlated with Ant Design.
+
+## Skills
+### Skill 1: Engaging Interaction
+Your primary role involves engaging with users, offering them in-depth responses to their Ant Design inquiries in a conversational fashion.
+
+### Skill 2: Insightful Information Search
+For queries that touch upon unfamiliar zones, you are equipped with two powerful knowledge lookup tools, used to gather necessary details:
+   - search_knowledge: This is your initial resource for queries concerning ambiguous topics about Ant Design. While using this, ensure to retain the user's original query language for the highest accuracy possible. Therefore, a specific question like 'Ant Design的新特性是什么?' should be searched as 'Ant Design的新特性是什么?'.
+   - tavily_search_results_json: Should search_knowledge fail to accommodate the required facts, this tool would be the next step.
+
+### Skill 3: Expert Issue Solver
+In case of specific issues reported by users, you are to aid them using a selection of bespoke tools, curated as per the issue nature and prescribed steps. The common instances cater to:
+   - Routine engagement with the user.
+   - Employment of certain tools such as create_issue, get_issues, search_issues, search_code etc. when the user is facing a specific hurdle.
+
+## Constraints:
+- Maintain a strict focus on Ant Design in your responses; if confronted with unrelated queries, politely notify the user of your confines and steer them towards asking questions relevant to Ant Design.
+- Your tool utilization choices should be driven by the nature of the inquiry and recommended actions.
+- While operating tools for searching information, keep the user's original language to attain utmost precision.
+- With your multilingual capability, always respond in the user's language. If the inquiry popped is in English, your response should mirror that; same goes for Chinese or any other language.
+"""
+
 prompt = ChatPromptTemplate.from_messages(
     [
         (
             "system",
-            "Character: Ant Design Bot"
-            "You are a question-and-answer robot specifically dedicated to providing solutions for Ant Design. You are well-versed in Ant Design-related knowledge and capable of assisting in fixing code issues related to it. However, you are unable to answer questions not related to Ant Design."
-            "Skill 1: Normal Conversation"
-            "In most cases, you will have ordinary conversations with users, answering questions about Ant Design."
-            "Skill 2: Issue Handling"
-            "When users ask about issues, you will choose to use the appropriate tools for handling based on the description of actions. There could be two scenarios:"
-            "1. Talk with the user as normal. "
-            "2. If they ask you about issues, use a tool"
-            "Constraints:"
-            "Only answer questions related to Ant Design; if users pose unrelated questions, you need to inform the user that you cannot answer and guide them to ask questions related to Ant Design."
-            "Decide whether to use relevant tools for handling based on the nature of the question and the description of actions."
-            "You must answer questions in the language the user inputs. You are programmed with the ability to recognize multiple languages, allowing you to understand and respond in the same language used by the user. If a user poses a question in English, you should answer in English; if a user communicates in Chinese, you should respond in Chinese."
+            prompt_template,
         ),
         MessagesPlaceholder(variable_name="chat_history"),
         ("user", "{input}"),
         MessagesPlaceholder(variable_name="agent_scratchpad"),
     ]
 )
 
-@tool
-def get_datetime() -> datetime:
-    """Get current date time."""
-    return datetime.now().strftime("%H:%M:%S")
-
 TOOL_MAPPING = {
-    "get_datetime": get_datetime,
+    "search_knowledge": knowledge.search_knowledge,
     "create_issue": issue.create_issue,
     "get_issues": issue.get_issues,
     "search_issues": issue.search_issues,
     "search_code": sourcecode.search_code,
 }
-TOOLS = ["get_datetime", "create_issue", "get_issues", "search_issues",  "search_code"]
+TOOLS = ["search_knowledge", "create_issue", "get_issues", "search_issues",  "search_code"]
 
+def init_default_tools():
+    # init Tavily 
+    search = TavilySearchAPIWrapper()
+    tavily_tool = TavilySearchResults(api_wrapper=search)
+
+    return [tavily_tool]
+
 
 def _create_agent_with_tools(open_api_key: str) -> AgentExecutor:
     llm = ChatOpenAI(model="gpt-4-1106-preview", temperature=0.2, streaming=True, max_tokens=1500, openai_api_key=open_api_key)
-    search = TavilySearchAPIWrapper()
-    tavily_tool = TavilySearchResults(api_wrapper=search)
-    tools = [tavily_tool]
+
+
+    tools = init_default_tools();
 
     for requested_tool in TOOLS:
         if requested_tool not in TOOL_MAPPING:
@@ -91,9 +102,9 @@ def _create_agent_with_tools(open_api_key: str) -> AgentExecutor:
         | llm_with_tools
         | OpenAIToolsAgentOutputParser()
     )
-    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True).with_config(
-        {"run_name": "agent"}
-    )
+
+    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, handle_parsing_errors=True,
+max_iterations = 5)
     return agent_executor
 
 
@@ -115,6 +126,7 @@ async def agent_chat(input_data: ChatData, open_api_key: str) -> AsyncIterator[s
         messages = input_data.messages
         agent_executor = _create_agent_with_tools(open_api_key)
         print(chat_history_transform(messages))
+
         async for event in agent_executor.astream_events(
             {
                 "input": messages[len(messages) - 1].content,
@@ -123,22 +135,6 @@ async def agent_chat(input_data: ChatData, open_api_key: str) -> AsyncIterator[s
             version="v1",
         ):
             kind = event["event"]
-            if kind == "on_chain_start":
-                if (
-                    event["name"] == "agent"
-                ): 
-                    print(
-                        f"Starting agent: {event['name']} "
-                        f"with input: {event['data'].get('input')}"
-                    )
-            elif kind == "on_chain_end":
-                if (
-                    event["name"] == "agent"
-                ): 
-                    print (
-                        f"Done agent: {event['name']} "
-                        f"with output: {event['data'].get('output')['output']}"
-                    )
             if kind == "on_chat_model_stream":
                 uid =  str(uuid.uuid4())
                 content = event["data"]["chunk"].content

server/main.py

@@ -4,6 +4,7 @@
 from fastapi import FastAPI
 from fastapi.responses import StreamingResponse
 from starlette.middleware.sessions import SessionMiddleware
+from fastapi.middleware.cors import CORSMiddleware
 
 from agent import stream
 
@@ -27,6 +28,14 @@
     secret_key = session_secret_key,
 )
 
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # 明确指定允许的源
+    allow_credentials=True,
+    allow_methods=["*"],  # 允许所有方法
+    allow_headers=["*"],  # 允许所有头部
+)
+
 app.include_router(health_checker.router)
 app.include_router(github.router)
 app.include_router(rag.router)

server/tools/knowledge.py

@@ -0,0 +1,19 @@
+from langchain.tools import tool
+from rag import retrieval
+
+
+@tool
+def search_knowledge(
+      query: str,
+):
+    """
+    Search for information based on the query.  When use this tool, do not translate the search query. Use the original query language to search. eg: When user's question is 'Ant Design 有哪些新特性?', the query should be 'Ant Design 有哪些新特性?'.
+
+    :param query: The user's question.
+    """
+    try:
+        return retrieval.search_knowledge(query)
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        return None
+