62 lines
1.8 KiB
Python
Executable File
62 lines
1.8 KiB
Python
Executable File
import requests
|
|
|
|
def retrieval_from_knowledge_base(
|
|
query: str,
|
|
topk: int
|
|
) -> str:
|
|
"""
|
|
Retrieval for knowledge from the knowledge base based on the specified query and returns the topk results.
|
|
|
|
Parameters:
|
|
query (str): The query for knowledge retrieval.
|
|
topk (int): The number of top results to return, default is 3.
|
|
|
|
Returns:
|
|
str: The result of the knowledge retrieval in JSON format.
|
|
"""
|
|
url = 'http://127.0.0.1:7080/v1/chat-messages'
|
|
headers = {
|
|
'Authorization': f'Bearer app-uJgo3TQKcS1O9PMCDHko71Fp',
|
|
'Content-Type': 'application/json'
|
|
}
|
|
data = {
|
|
"inputs": {"topK": topk},
|
|
"query": query,
|
|
"response_mode": "blocking",
|
|
"user": "tangger",
|
|
"files": []
|
|
}
|
|
|
|
response = requests.post(url, headers=headers, json=data)
|
|
|
|
if response.status_code == 524:
|
|
print("Server is not responding. Please try again later. Maybe GPU was down in the container.")
|
|
return None
|
|
|
|
try:
|
|
result = response.json()
|
|
except ValueError:
|
|
return [{"error": "Response is not in JSON format"}]
|
|
|
|
useful_results = []
|
|
try:
|
|
answer = eval(result.get("answer", "[]"))
|
|
for item in answer:
|
|
metadata = item.get("metadata", {})
|
|
useful_info = {
|
|
"id": metadata.get("document_id"),
|
|
"title": item.get("title"),
|
|
"content": item.get("content"),
|
|
"metadata": None,
|
|
"embedding": None,
|
|
"score": metadata.get("score")
|
|
}
|
|
useful_results.append(useful_info)
|
|
except Exception as e:
|
|
return [{"error": f"Error processing result: {e}", "status": "TERMINATE"}]
|
|
if useful_results == []:
|
|
useful_results = "NULL"
|
|
return str(useful_results)
|
|
|
|
|