Update app.py

app.py

@@ -1,3 +1,5 @@
+Updated Multi-agent RAG-based LLM Model
+
 import os
 import multiprocessing
 import concurrent.futures
@@ -15,12 +17,69 @@ import gradio as gr
 import re
 from threading import Thread
 
+class Agent:
+    def __init__(self, name, role, doc_retrieval_gen, tokenizer):
+        self.name = name
+        self.role = role
+        self.doc_retrieval_gen = doc_retrieval_gen
+        self.tokenizer = tokenizer
+
+    def generate_response(self, query, context):
+        if self.role == "Information Retrieval":
+            return self.retriever_logic(query, context)
+        elif self.role == "Content Analysis":
+            return self.analyzer_logic(query, context)
+        elif self.role == "Response Generation":
+            return self.generator_logic(query, context)
+        elif self.role == "Task Coordination":
+            return self.coordinator_logic(query, context)
+
+    def retriever_logic(self, query, all_splits):
+        query_embedding = self.doc_retrieval_gen.embeddings.encode(query, convert_to_tensor=True).cpu().numpy()
+        distances, indices = self.doc_retrieval_gen.gpu_index.search(np.array([query_embedding]), k=3)
+        relevant_docs = [all_splits[i] for i in indices[0] if distances[0][i] <= 1]
+        return relevant_docs
+
+    def analyzer_logic(self, query, relevant_docs):
+        analysis_prompt = f"Analyze the following documents in relation to the query: '{query}'\n\nDocuments:\n"
+        for doc in relevant_docs:
+            analysis_prompt += f"- {doc.page_content}\n"
+        analysis_prompt += "\nProvide a concise analysis of the key points relevant to the query."
+        
+        input_ids = self.tokenizer.encode(analysis_prompt, return_tensors="pt").to(self.doc_retrieval_gen.model.device)
+        analysis = self.doc_retrieval_gen.model.generate(input_ids, max_length=200, num_return_sequences=1)
+        return self.tokenizer.decode(analysis[0], skip_special_tokens=True)
+
+    def generator_logic(self, query, analyzed_content):
+        generation_prompt = f"Based on the following analysis, generate a comprehensive answer to the query: '{query}'\n\nAnalysis:\n{analyzed_content}\n\nGenerate a detailed response:"
+        
+        input_ids = self.tokenizer.encode(generation_prompt, return_tensors="pt").to(self.doc_retrieval_gen.model.device)
+        response = self.doc_retrieval_gen.model.generate(input_ids, max_length=300, num_return_sequences=1)
+        return self.tokenizer.decode(response[0], skip_special_tokens=True)
+
+    def coordinator_logic(self, query, final_response):
+        coordination_prompt = f"As a coordinator, review and refine the following response to the query: '{query}'\n\nResponse:\n{final_response}\n\nProvide a final, polished answer:"
+        
+        input_ids = self.tokenizer.encode(coordination_prompt, return_tensors="pt").to(self.doc_retrieval_gen.model.device)
+        coordinated_response = self.doc_retrieval_gen.model.generate(input_ids, max_length=350, num_return_sequences=1)
+        return self.tokenizer.decode(coordinated_response[0], skip_special_tokens=True)
+
 class DocumentRetrievalAndGeneration:
     def __init__(self, embedding_model_name, lm_model_id, data_folder):
         self.all_splits = self.load_documents(data_folder)
         self.embeddings = SentenceTransformer(embedding_model_name)
         self.gpu_index = self.create_faiss_index()
         self.tokenizer, self.model = self.initialize_llm(lm_model_id)
+        self.agents = self.initialize_agents()
+
+    def initialize_agents(self):
+        agents = [
+            Agent("Retriever", "Information Retrieval", self, self.tokenizer),
+            Agent("Analyzer", "Content Analysis", self, self.tokenizer),
+            Agent("Generator", "Response Generation", self, self.tokenizer),
+            Agent("Coordinator", "Task Coordination", self, self.tokenizer)
+        ]
+        return agents
 
     def load_documents(self, folder_path):
         loader = DirectoryLoader(folder_path, loader_cls=TextLoader)
@@ -59,6 +118,30 @@ class DocumentRetrievalAndGeneration:
         )
         return tokenizer, model
 
+    def coordinate_agents(self, query):
+        coordinator = next(agent for agent in self.agents if agent.name == "Coordinator")
+        
+        # Step 1: Information Retrieval
+        retriever = next(agent for agent in self.agents if agent.name == "Retriever")
+        relevant_docs = retriever.generate_response(query, self.all_splits)
+        
+        # Step 2: Content Analysis
+        analyzer = next(agent for agent in self.agents if agent.name == "Analyzer")
+        analyzed_content = analyzer.generate_response(query, relevant_docs)
+        
+        # Step 3: Response Generation
+        generator = next(agent for agent in self.agents if agent.name == "Generator")
+        final_response = generator.generate_response(query, analyzed_content)
+        
+        # Step 4: Coordination and Refinement
+        coordinated_response = coordinator.generate_response(query, final_response)
+        
+        return coordinated_response, "\n".join([doc.page_content for doc in relevant_docs])
+
+    def query_and_generate_response(self, query):
+        return self.coordinate_agents(query)
+
+        
     def generate_response_with_timeout(self, input_ids, max_new_tokens=1000):
         try:
             streamer = TextIteratorStreamer(self.tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
@@ -86,6 +169,7 @@ class DocumentRetrievalAndGeneration:
             print(f"Error in generate_response_with_timeout: {str(e)}")
             return "Text generation process encountered an error"
 
+
     def query_and_generate_response(self, query):
         similarityThreshold = 1
         query_embedding = self.embeddings.encode(query, convert_to_tensor=True).cpu().numpy()
@@ -142,8 +226,8 @@ class DocumentRetrievalAndGeneration:
         return solution_text, content
 
     def qa_infer_gradio(self, query):
-        response = self.query_and_generate_response(query)
-        return response
+        response, related_queries = self.query_and_generate_response(query)
+        return response, related_queries
 
 if __name__ == "__main__":
     embedding_model_name = 'flax-sentence-embeddings/all_datasets_v3_MiniLM-L12'