Update app.py

app.py

@@ -255,45 +255,11 @@ def setup_llm(model_name, temperature, api_key):
     llm = ChatGroq(model=model_name, temperature=temperature, api_key=api_key)
     return llm
 
-def retrieve_from_vectorstore(vectorstore, query, k):
-    retrieved_docs = vectorstore.similarity_search(query, k=k)    
-    return retrieved_docs
-
-
-
-def retrieve_within_kadiApy_docs(vectorstore, query, k):
-    filter_criteria = {"usage": "doc"}
-    retrieved_docs = vectorstore.similarity_search(query=query, k=k, filter=filter_criteria)
-    return retrieved_docs
-
-def retrieve_within_kadiApy_library(vectorstore, query, k):
-    filter_criteria = {"usage": "library"}
-    retrieved_docs = vectorstore.similarity_search(query=query, k=k, filter=filter_criteria)
-    return retrieved_docs
-
-def retrieve_within_kadiApy_cli_library(vectorstore, query, k):
-    filter_criteria = {"usage": "cli_library"}
-    retrieved_docs = vectorstore.similarity_search(query=query, k=k, filter=filter_criteria)
-    return retrieved_docs
-
-def retrieve_within_kadiApy_cli_library_excluding_cli_commands(vectorstore, query, k):
-    filter_criteria = {"usage": "cli_library"}
-    filtered_docs = [doc for doc in documents if "command" not in doc.metadata]
-    return filtered_docs
-    
-    
-def retrieve_kadiApy_cli_commands(vectorstore, query, k):
-    filter_criteria = {"type": "command"}
-    results = vectorstore.similarity_search(query=query, k=k, filter=filter_criteria)
-    return results
-
-
-
-def retrieve_docs_from_vectorstore(vectorstore, query, k): 
-    return vectorstore.similarity_search(query, k=k)
 
 
 
+def retrieve_from_vectorstore(vectorstore = vector_store, query, k, filter=None): 
+    return vectorstore.similarity_search(query, k=k, filter=filter)
 
 def format_kadi_apy_library_context(docs):
     doc_context = []
@@ -378,23 +344,30 @@ def rag_workflow(query):
 
     rewritten_query_response = llm.invoke(rewrite_prompt)
     rewritten_query = rewritten_query_response.content.strip()
-    print("AAAAArewritten_query -------------------------:", metadata_prediction)
+    
+    print("A", metadata_prediction)
     print(rewritten_query)
 
 
+    kadi_apy_docs = retrieve_from_vectorstore(
+        query, 
+        k=4,
+        filter={"usage": "doc"}
+    )
     
-    kadi_apy_docs = retrieve_within_kadiApy_docs (vector_store, query, k = 5)
-    kadi_apy_library_docs = vector_store.similarity_search(query, filter={"usage": metadata_prediction})
-        
+    kadi_apy_sourcecode = retrieve_from_vectorstore(
+        query, 
+        k=6,
+        filter={"usage": metadata_prediction}
+    )
+
     doc_context = format_kadi_api_doc_context(kadi_apy_docs)
-    code_context = format_kadi_apy_library_context(kadi_apy_library_docs)
+    code_context = format_kadi_apy_library_context(kadi_apy_sourcecode)
 
     
-    # Assuming kadi_apy_library_docs is a list of Document objects
-
     print("HERE WE GHOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOO")
     print("::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::")
-    for doc in kadi_apy_library_docs:
+    for doc in kadi_apy_sourcecode:
         print(doc.metadata.get("source", "Unknown Type"))
         print("\n")
     print("::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::")