Explorar o código

added a few type hints

thiswillbeyourgithub hai 7 meses
pai
achega
65d5545cf0
Modificáronse 1 ficheiros con 4 adicións e 4 borrados
  1. 4 4
      backend/open_webui/apps/rag/utils.py

+ 4 - 4
backend/open_webui/apps/rag/utils.py

@@ -48,7 +48,7 @@ def query_doc_with_hybrid_search(
     k: int,
     reranking_function,
     r: float,
-):
+) -> dict:
     try:
         collection = CHROMA_CLIENT.get_collection(name=collection_name)
         documents = collection.get()  # get all documents
@@ -93,7 +93,7 @@ def query_doc_with_hybrid_search(
         raise e
 
 
-def merge_and_sort_query_results(query_results, k, reverse=False):
+def merge_and_sort_query_results(query_results: list[dict], k: int, reverse: bool = False) -> list[dict]:
     # Initialize lists to store combined data
     combined_distances = []
     combined_documents = []
@@ -139,7 +139,7 @@ def query_collection(
     query: str,
     embedding_function,
     k: int,
-):
+) -> dict:
     results = []
     for collection_name in collection_names:
         if collection_name:
@@ -166,7 +166,7 @@ def query_collection_with_hybrid_search(
     k: int,
     reranking_function,
     r: float,
-):
+) -> dict:
     results = []
     failed = 0
     for collection_name in collection_names: