Langchain-Chatchat/server/agent/tools_factory/search_local_knowledgebase.py
2024-03-22 20:39:25 +08:00

40 lines
1.4 KiB
Python

from urllib.parse import urlencode
from server.pydantic_v1 import BaseModel, Field
from server.knowledge_base.kb_doc_api import search_docs
from configs import TOOL_CONFIG
def search_knowledgebase(query: str, database: str, config: dict):
docs = search_docs(
query=query,
knowledge_base_name=database,
top_k=config["top_k"],
score_threshold=config["score_threshold"])
context = ""
source_documents = []
for inum, doc in enumerate(docs):
filename = doc.metadata.get("source")
parameters = urlencode({"knowledge_base_name": database, "file_name": filename})
url = f"download_doc?" + parameters
text = f"""出处 [{inum + 1}] [{filename}]({url}) \n\n{doc.page_content}\n\n"""
source_documents.append(text)
if len(source_documents) == 0:
context = "没有找到相关文档,请更换关键词重试"
else:
for doc in source_documents:
context += doc + "\n"
return context
class SearchKnowledgeInput(BaseModel):
database: str = Field(description="Database for Knowledge Search")
query: str = Field(description="Query for Knowledge Search")
def search_local_knowledgebase(database: str, query: str):
tool_config = TOOL_CONFIG["search_local_knowledgebase"]
return search_knowledgebase(query=query, database=database, config=tool_config)