jandres/GenAIExamples
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

fix bugs in DocIndexRetriever (#1770)

Browse Source 
Signed-off-by: minmin-intel <minmin.hou@intel.com>
This commit is contained in:
minmin-intel
2025-04-09 18:45:46 -07:00
committed by GitHub
parent 00d7a65dd8
commit 411bb28f41
12 changed files with 152 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
DocIndexRetriever/tests/test.py
View File

@@ -2,16 +2,17 @@
# SPDX-License-Identifier: Apache-2.0
import os
from typing import Any
import requests
def search_knowledge_base(query: str) -> str:
def search_knowledge_base(query: str, args: Any) -> str:
"""Search the knowledge base for a specific query."""
url = os.environ.get("RETRIEVAL_TOOL_URL")
url = os.environ.get("RETRIEVAL_TOOL_URL", "http://localhost:8889/v1/retrievaltool")
print(url)
proxies = {"http": ""}
payload = {"messages": query, "k": 5, "top_n": 2}
payload = {"messages": query, "k": args.k, "top_n": args.top_n}
response = requests.post(url, json=payload, proxies=proxies)
print(response)
if "documents" in response.json():
@@ -33,6 +34,16 @@ def search_knowledge_base(query: str) -> str:
if __name__ == "__main__":
resp = search_knowledge_base("What is OPEA?")
# resp = search_knowledge_base("Thriller")
import argparse
parser = argparse.ArgumentParser(description="Test the knowledge base search.")
parser.add_argument("--k", type=int, default=5, help="retriever top k")
parser.add_argument("--top_n", type=int, default=2, help="reranker top n")
args = parser.parse_args()
resp = search_knowledge_base("What is OPEA?", args)
print(resp)
if not resp.startswith("Error"):
print("Test successful!")