Update ragchain.py

ragchain.py

@@ -1,13 +1,65 @@
-class RAGChain:
+class RAGBot:
     
     def __init__(self, llm, vector_store):
         """
-        Initialize the RAGChain with an LLM instance and a vector store.
+        Initialize the RAGChain with an LLM instance, a vector store, and a conversation history.
         """
         self.llm = llm
         self.vector_store = vector_store
+        self.conversation = []  
 
-    
+
+
+    def rag_chain(self, query):
+        """
+        Process a user query, handle history, retrieve contexts, and generate a response.
+        """
+        # Add the user query to the conversation history
+        self.add_to_conversation(user_query=query)
+        
+        # Rewrite query
+        rewritten_query = self.rewrite_query(query)
+
+        # Predict library usage
+        code_library_usage_prediction = self.predict_library_usage(query)
+        
+        # Retrieve contexts
+        doc_contexts = self.retrieve_contexts(query, k=5, filter={"usage": "doc"})
+        code_contexts = self.retrieve_contexts(rewritten_query, k=3, filter={"usage": code_library_usage_prediction})
+        
+        # Format contexts
+        formatted_doc_contexts = self.format_documents(doc_contexts)
+        formatted_code_contexts = self.format_documents(code_contexts)
+        
+        # Generate response
+        response = self.generate_response(query, formatted_doc_contexts, formatted_code_contexts)
+        
+        # Add the response to the existing query in the conversation history
+        self.add_to_conversation(llm_response=response)
+        
+        return response
+
+
+    def add_to_conversation(self, user_query=None, llm_response=None):
+        """
+        Add either the user's query, the LLM's response, or both to the conversation history.
+        """
+        if user_query and llm_response:
+            # Add a full query-response pair
+            self.conversation.append({"query": user_query, "response": llm_response})
+        elif user_query:
+            # Add a query with no response yet
+            self.conversation.append({"query": user_query, "response": None})
+        elif llm_response and self.conversation:
+            # Add a response to the most recent query
+            self.conversation[-1]["response"] = llm_response
+
+    def get_history(self):
+        """
+        Retrieve the entire conversation history.
+        """
+        return self.conversation
+        
     def rewrite_query(self, query):
         """
         Rewrite the user's query to align with the language and structure of the library's methods and documentation.