feat: Add country info tool to agent

src/agent/tool_divisions.py

@@ -113,16 +113,16 @@ def complete_administrative_division(input_text, data):
         for city in province['cities']:
             if input_text in city['name']:
                 return {
-                    'province': province['province'],
-                    'city': city['name'],
-                    'counties': city['counties']
+                    '省': province['province'],
+                    '市': city['name'],
+                    '县(区)': city['counties']
                 }
             for county in city['counties']:
                 if input_text in county:
                     return {
-                        'province': province['province'],
-                        'city': city['name'],
-                        'county': county
+                        '省': province['province'],
+                        '市': city['name'],
+                        '县(区)': county
                     }
     return None
 
@@ -132,9 +132,22 @@ class AdministrativeDivisionArgs(BaseModel):
 
 class AdministrativeDivision(BaseTool):
     name = "administrative_division"
-    description = "根据输入补全行政区划信息，明确具体的省、市、县信息。"
+    description = "根据输入补全行政区划信息，明确具体的省、市、县信息。比如输入县，补全所属省市，输入市则补全省级以及下辖所有县区"
     args_schema: Type[BaseModel] = AdministrativeDivisionArgs
 
     def _run(self, input_text: str) -> str:
         result = complete_administrative_division(input_text, divisions)
-        return result
+        return result
\ No newline at end of file
+    
+
+class CountryInfoArgs(BaseModel):
+    country: str = Field(...,description="县(区)这个级别的行政区划名称")
+
+# fake tool for testing
+class CountryInfo(BaseTool):
+    name = "country_info"
+    description = "查询特定县(区)相关的地质情况"
+    args_schema: Type[BaseModel] = CountryInfoArgs
+
+    def _run(self, country: str) -> str:
+        return f"{country} 的地质情况是：xxxxxx"
\ No newline at end of file

src/config/prompts.py

@@ -45,7 +45,7 @@ PROMPT_AGENT_HUMAN = """{input}\n\n{agent_scratchpad}\n (请注意，无论如
 ##################################################################################################################
 
 # 结合历史对话信息，对用户提问进行扩展，生成不同角度的多个提问。用于 RAG 场景中
-PROMPT_QUERY_EXTEND = """作为一个向量检索助手，你的任务是结合历史记录，从不同角度，为“原问题”生成三个不同版本的“检索词”，从而提高向量检索的语义丰富度，提高向量检索的精度。生成的问题要求指向对象清晰明确，并与“原问题语言相同”。例如：
+PROMPT_QUERY_EXTEND = """作为一个向量检索助手，你的任务是结合历史记录，从不同角度，为“原问题”生成三个不同版本的“检索词”，从而提高向量检索的语义丰富度，提高向量检索的精度。生成的问题要求指向对象清晰明确，并与“原问题语言相同”。如果问题中包含多个问询对象，可以将其拆分为多个问题。例如：
 历史记录: 
 '''
 '''
@@ -67,7 +67,7 @@ A: 当前对话是关于 Nginx 的介绍和使用等。
 Q: 报错 "no connection"
 A: 报错"no connection"可能是因为……
 '''
-原问题: 怎么解决
+原问题: 这是什么原因要怎么解决
 检索词: ["Nginx报错"no connection"如何解决？","造成'no connection'报错的原因。","Nginx提示'no connection'，要怎么办？"]
 ----------------
 历史记录: 
@@ -96,11 +96,9 @@ A: 关于 FatGPT 的介绍和使用等问题。
 ----------------
 历史记录:
 '''
-Q: FastGPT 如何收费？
-A: FastGPT 收费可以参考……
 '''
-原问题: 你知道 laf 么？
-检索词: ["laf 的官网地址是多少？","laf 的使用教程。","laf 有什么特点和优势。"]
+原问题: 北京和上海那个天气好？
+检索词: ["北京的天气情况如何","上海的天气情况如何","北京和上海的天气那个更舒适"]
 ----------------
 历史记录:
 '''

src/server/agent.py

@@ -24,3 +24,6 @@ class Agent:
         #         yield step  
         return self.agent_executor.invoke(input=prompt_args)
         
+    def stream(self, prompt_args: dict = {}):
+        for step in self.agent_executor.stream(prompt_args):
+            yield step
\ No newline at end of file

test/agent_test.py

@@ -18,7 +18,7 @@ from pydantic import BaseModel, Field
 from src.server.agent import Agent
 from src.config.prompts import PROMPT_AGENT_SYS, PROMPT_AGENT_HUMAN
 
-from src.agent.tool_divisions import AdministrativeDivision
+from src.agent.tool_divisions import AdministrativeDivision, CountryInfo
 
 class CalcInput(BaseModel):
     a: int = Field(...,description="第一个数")
@@ -38,6 +38,7 @@ class Calc(BaseTool):
 
 
 
+
 tools = [AdministrativeDivision()]
 
 llm = ChatOpenAI(
@@ -82,9 +83,9 @@ def test_add():
     #     print(step)
      
 def test_agent_division():
-    tools = [AdministrativeDivision()]
+    tools = [AdministrativeDivision(),CountryInfo()]
     agent = Agent(llm=llm, tools=tools, prompt=prompt, verbose=True)
-    res = agent.exec(prompt_args={"input": "介绍下陇南市武都区的基本情况"})
+    res = agent.exec(prompt_args={"input": "我想知道陇南市西和县和文县的降雨量谁的多"})
     print(res)
 
 if __name__ == "__main__":