revert: original rag pipeline

backend/apps/rag/utils.py

@@ -18,6 +18,9 @@ from langchain.retrievers import (
     EnsembleRetriever,
 )
 
+from sentence_transformers import CrossEncoder
+
+from typing import Optional
 from config import SRC_LOG_LEVELS, CHROMA_CLIENT
 
 
@@ -28,12 +31,14 @@ log.setLevel(SRC_LOG_LEVELS["RAG"])
 def query_embeddings_doc(
     collection_name: str,
     query: str,
-    k: int,
-    r: float,
     embeddings_function,
-    reranking_function,
+    k: int,
+    reranking_function: Optional[CrossEncoder] = None,
+    r: Optional[float] = None,
 ):
     try:
+
+        if reranking_function:
             # if you use docker use the model from the environment variable
             collection = CHROMA_CLIENT.get_collection(name=collection_name)
 
@@ -71,7 +76,19 @@ def query_embeddings_doc(
                 "documents": [[d.page_content for d in result]],
                 "metadatas": [[d.metadata for d in result]],
             }
+        else:
+            # if you use docker use the model from the environment variable
+            query_embeddings = embeddings_function(query)
+
+            log.info(f"query_embeddings_doc {query_embeddings}")
+            collection = CHROMA_CLIENT.get_collection(name=collection_name)
+
+            result = collection.query(
+                query_embeddings=[query_embeddings],
+                n_results=k,
+            )
 
+            log.info(f"query_embeddings_doc:result {result}")
         return result
     except Exception as e:
         raise e