added a few type hints

This commit is contained in:
thiswillbeyourgithub 2024-09-12 15:50:18 +02:00
parent e872f5dc78
commit 65d5545cf0

View File

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