From 65d5545cf0cd85c1b4a5e55266289b42879c4673 Mon Sep 17 00:00:00 2001
From: thiswillbeyourgithub
 <26625900+thiswillbeyourgithub@users.noreply.github.com>
Date: Thu, 12 Sep 2024 15:50:18 +0200
Subject: [PATCH] added a few type hints

---
 backend/open_webui/apps/rag/utils.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/backend/open_webui/apps/rag/utils.py b/backend/open_webui/apps/rag/utils.py
index 29b12d0b0..9da08e28a 100644
--- a/backend/open_webui/apps/rag/utils.py
+++ b/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: