Update ragchain.py

ragchain.py

@@ -63,14 +63,12 @@ class RAGChain:
         )
         return self.llm.predict(prompt)
 
-    def retrieve_contexts(self, query, library_usage_prediction):
+    def retrieve_contexts(self, query, k, filter = None):
         """
         Retrieve relevant documents and source code based on the query and library usage prediction.
         """
-        doc_contexts = self.vector_store.similarity_search(query, k=5, filter={"usage": "doc"})
-        code_contexts = self.vector_store.similarity_search(query, k=5, filter={"usage": library_usage_prediction})
-        
-        return doc_contexts, code_contexts
+        context = self.vector_store.similarity_search(query = query, k=k, filter=filter)       
+        return context
 
     def format_documents(self, documents):
         formatted_docs = []