[全量] 初始化项目代码、配置、文档及Agent协同harness

langchain-chat/server/agent/tools/policy_knowledgebase_search.py

@@ -0,0 +1,42 @@
+from server.chat.knowledge_base_chat import knowledge_base_chat
+from configs import VECTOR_SEARCH_TOP_K, SCORE_THRESHOLD, MAX_TOKENS, LLM_MODELS
+import json
+import asyncio
+from typing import List, Tuple, Any
+from pydantic import BaseModel, Field
+
+
+class PolicyKnowledgeInput(BaseModel):
+    location: str = Field(description="The policy related query to be searched")
+
+
+async def policy_knowledgebase_search_iter(query: str) -> tuple[str | Any, list[Any] | Any]:
+    response = await knowledge_base_chat(query=query,
+                                         fileName=None,
+                                         knowledge_base_name_list = ["t_policy_total_bge_new_v1"],
+                                         model_name=LLM_MODELS[0],
+                                         temperature=0.01,
+                                         history=[],
+                                         top_k=VECTOR_SEARCH_TOP_K,
+                                         max_tokens=MAX_TOKENS,
+                                         prompt_name="default",
+                                         score_threshold=SCORE_THRESHOLD,
+                                         stream=False)
+
+    contents = ""
+    docs = []
+    async for data in response.body_iterator:  # 这里的data是一个json字符串
+        data = json.loads(data)
+        print("data>>>>>", data)
+        contents = data["answer"] if "answer" in data.keys() else contents
+        docs = data["docs"] if "docs" in data.keys() else docs
+    return contents, docs
+
+
+def policy_knowledgebase_search(query: str) -> tuple[str | Any, list[Any] | Any]:
+    return asyncio.run(policy_knowledgebase_search_iter(query))
+
+
+if __name__ == "__main__":
+    result = policy_knowledgebase_search("大数据男女比例")
+    print("答案:", result)