Update app.py

app.py

@@ -258,7 +258,7 @@ def setup_llm(model_name, temperature, api_key):
 
 
 
-def retrieve_from_vectorstore(vectorstore = vector_store, query = query, k = k, filter=None): 
+def retrieve_from_vectorstore(vectorstore = vectorstore, query = query, k = k, filter=None): 
     return vectorstore.similarity_search(query, k=k, filter=filter)
 
 def format_kadi_apy_library_context(docs):
@@ -306,7 +306,7 @@ def rag_workflow(query):
             Respond with only the exact corresponding option and do not include any additional comments, explanations, or text."
         """
     )
-    metadata_prediction = llm.predict(prompt)
+    library_usage_prediction = llm.predict(prompt)
 
     print("METADATA PREDICTION -------------------------:", metadata_prediction)
     print(metadata_prediction)
@@ -348,18 +348,8 @@ def rag_workflow(query):
     print("A", metadata_prediction)
     print(rewritten_query)
 
-
-    kadi_apy_docs = retrieve_from_vectorstore(
-        query, 
-        k=4,
-        filter={"usage": "doc"}
-    )
-    
-    kadi_apy_sourcecode = retrieve_from_vectorstore(
-        query, 
-        k=6,
-        filter={"usage": metadata_prediction}
-    )
+    kadi_apy_docs = vector_store.similarity_search(query, k=5, filter={"usage": "doc"})
+    kadi_apy_docs = vector_store.similarity_search(query, k=5, filter={"usage": library_usage_prediction})
 
     doc_context = format_kadi_api_doc_context(kadi_apy_docs)
     code_context = format_kadi_apy_library_context(kadi_apy_sourcecode)