Update app.py

app.py

@@ -15,6 +15,11 @@ import gradio as gr
 import re
 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
@@ -62,6 +67,71 @@ class Agent:
         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 MultiDocumentAgentSystem:
+    def __init__(self, documents_dict, llm, embed_model):
+        self.llm = llm
+        self.embed_model = embed_model
+        self.document_agents = {}
+        self.create_document_agents(documents_dict)
+        self.top_agent = self.create_top_agent()
+
+    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_query_engine = vector_index.as_query_engine(similarity_top_k=2)
+            summary_query_engine = summary_index.as_query_engine()
+
+            query_engine_tools = [
+                QueryEngineTool(
+                    query_engine=vector_query_engine,
+                    metadata=ToolMetadata(
+                        name=f"vector_tool_{doc_name}",
+                        description=f"Useful for specific questions about {doc_name}",
+                    ),
+                ),
+                QueryEngineTool(
+                    query_engine=summary_query_engine,
+                    metadata=ToolMetadata(
+                        name=f"summary_tool_{doc_name}",
+                        description=f"Useful for summarizing content about {doc_name}",
+                    ),
+                ),
+            ]
+
+            self.document_agents[doc_name] = OpenAIAgent.from_tools(
+                query_engine_tools,
+                llm=self.llm,
+                verbose=True,
+                system_prompt=f"You are an agent designed to answer queries about {doc_name}.",
+            )
+
+    def create_top_agent(self):
+        all_tools = []
+        for doc_name, agent in self.document_agents.items():
+            doc_tool = QueryEngineTool(
+                query_engine=agent,
+                metadata=ToolMetadata(
+                    name=f"tool_{doc_name}",
+                    description=f"Use this tool for questions about {doc_name}",
+                ),
+            )
+            all_tools.append(doc_tool)
+
+        obj_index = VectorStoreIndex.from_objects(all_tools, embed_model=self.embed_model)
+
+        return OpenAIAgent.from_tools(
+            all_tools,
+            llm=self.llm,
+            verbose=True,
+            system_prompt="You are an agent designed to answer queries about multiple documents.",
+            tool_retriever=obj_index.as_retriever(similarity_top_k=3),
+        )
+
+    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)
@@ -69,6 +139,10 @@ class DocumentRetrievalAndGeneration:
         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 = [
@@ -80,15 +154,14 @@ class DocumentRetrievalAndGeneration:
         return agents
 
     def load_documents(self, folder_path):
-        loader = DirectoryLoader(folder_path, loader_cls=TextLoader)
-        documents = loader.load()
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=5000, chunk_overlap=250)
-        all_splits = text_splitter.split_documents(documents)
-        print('Length of documents:', len(documents))
-        print("LEN of all_splits", len(all_splits))
-        for i in range(3):
-            print(all_splits[i].page_content)
-        return all_splits
+        documents_dict = {}
+        for file_name in os.listdir(folder_path):
+            if file_name.endswith('.txt'):
+                file_path = os.path.join(folder_path, file_name)
+                with open(file_path, 'r', encoding='utf-8') as file:
+                    content = file.read()
+                    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]
@@ -137,7 +210,8 @@ class DocumentRetrievalAndGeneration:
         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)
+        response = self.multi_doc_system.query(query)
+        return str(response), ""
 
         
     def generate_response_with_timeout(self, input_ids, max_new_tokens=1000):