Update app.py

app.py

@@ -18,54 +18,8 @@ from threading import Thread
 from llama_index.core import VectorStoreIndex, Document
 from llama_index.core.tools import QueryEngineTool, ToolMetadata
 from llama_index.agent.openai import OpenAIAgent
-
-
-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)
+from llama_index.llms.openai import OpenAI
+from llama_index.embeddings.openai import OpenAIEmbedding
 
 class MultiDocumentAgentSystem:
     def __init__(self, documents_dict, llm, embed_model):
@@ -77,8 +31,8 @@ class MultiDocumentAgentSystem:
 
     def create_document_agents(self, documents_dict):
         for doc_name, doc_content in documents_dict.items():
-            vector_index = VectorStoreIndex.from_documents([Document(doc_content)])
-            summary_index = VectorStoreIndex.from_documents([Document(doc_content)])
+            vector_index = VectorStoreIndex.from_documents([Document(text=doc_content)])
+            summary_index = VectorStoreIndex.from_documents([Document(text=doc_content)])
             
             vector_query_engine = vector_index.as_query_engine(similarity_top_k=2)
             summary_query_engine = summary_index.as_query_engine()
@@ -131,27 +85,15 @@ class MultiDocumentAgentSystem:
 
     def query(self, user_input):
         return self.top_agent.chat(user_input)
-        
+
 class DocumentRetrievalAndGeneration:
     def __init__(self, embedding_model_name, lm_model_id, data_folder):
-        self.all_splits = self.load_documents(data_folder)
+        self.documents_dict = 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()
-        documents_dict = self.load_documents(data_folder)
-        self.multi_doc_system = MultiDocumentAgentSystem(documents_dict, self.model, self.embeddings)
-
-
-
-    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
+        self.llm = OpenAI(temperature=0, model="gpt-3.5-turbo")
+        self.embed_model = OpenAIEmbedding()
+        self.multi_doc_system = MultiDocumentAgentSystem(self.documents_dict, self.llm, self.embed_model)
 
     def load_documents(self, folder_path):
         documents_dict = {}
@@ -163,15 +105,6 @@ class DocumentRetrievalAndGeneration:
                     documents_dict[file_name[:-4]] = content  # Use filename without .txt as key
         return documents_dict
 
-    def create_faiss_index(self):
-        all_texts = [split.page_content for split in self.all_splits]
-        embeddings = self.embeddings.encode(all_texts, convert_to_tensor=True).cpu().numpy()
-        index = faiss.IndexFlatL2(embeddings.shape[1])
-        index.add(embeddings)
-        gpu_resource = faiss.StandardGpuResources()
-        gpu_index = faiss.index_cpu_to_gpu(gpu_resource, 0, index)
-        return gpu_index
-
     def initialize_llm(self, model_id):
         quantization_config = BitsAndBytesConfig(
             load_in_4bit=True,
@@ -184,36 +117,10 @@ class DocumentRetrievalAndGeneration:
             model_id,
             torch_dtype=torch.bfloat16,
             device_map="auto",
-           
             quantization_config=quantization_config
         )
         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):
-        response = self.multi_doc_system.query(query)
-        return str(response), ""
-
-        
     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)
@@ -241,61 +148,9 @@ 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()
-        distances, indices = self.gpu_index.search(np.array([query_embedding]), k=3)
-        print("Distance", distances, "indices", indices)
-        content = ""
-        filtered_results = []
-        for idx, distance in zip(indices[0], distances[0]):
-            if distance <= similarityThreshold:
-                filtered_results.append(idx)
-            for i in filtered_results:
-                print(self.all_splits[i].page_content)
-            content += "-" * 50 + "\n"
-            content += self.all_splits[idx].page_content + "\n"
-            print("CHUNK", idx)
-            print("Distance:", distance)
-            print("indices:", indices)
-            print(self.all_splits[idx].page_content)
-            print("############################")
-
-        conversation = [
-            {"role": "system", "content": "You are a knowledgeable assistant with access to a comprehensive database."},
-            {"role": "user", "content": f"""
-            I need you to answer my question and provide related information in a specific format.
-            I have provided five relatable json files {content}, choose the most suitable chunks for answering the query.
-            RETURN ONLY SOLUTION without additional comments, sign-offs, retrived chunks, refrence to any Ticket or extra phrases. Be direct and to the point.
-            IF THERE IS NO ANSWER RELATABLE IN RETRIEVED CHUNKS, RETURN "NO SOLUTION AVAILABLE".
-            DO NOT GIVE REFRENCE TO ANY CHUNKS OR TICKETS,BE ON POINT.
-            
-            Here's my question:
-            Query: {query}
-            Solution==>
-            """}
-        ]
-        #Include a final answer without additional comments, sign-offs, or extra phrases. Be direct and to the point.
-        input_ids = self.tokenizer.apply_chat_template(conversation, return_tensors="pt").to(self.model.device)
-        
-        start_time = datetime.now()
-        generated_response = self.generate_response_with_timeout(input_ids)
-        elapsed_time = datetime.now() - start_time
-
-        print("Generated response:", generated_response)
-        print("Time elapsed:", elapsed_time)
-        print("Device in use:", self.model.device)
-
-        solution_text = generated_response.strip()
-        if "Solution:" in solution_text:
-            solution_text = solution_text.split("Solution:", 1)[1].strip()
-
-        # Post-processing to remove "assistant" prefix
-        solution_text = re.sub(r'^assistant\s*', '', solution_text, flags=re.IGNORECASE)
-        solution_text = solution_text.strip()
-
-        return solution_text, content
+        response = self.multi_doc_system.query(query)
+        return str(response), ""
 
     def qa_infer_gradio(self, query):
         response, related_queries = self.query_and_generate_response(query)