Update app.py

app.py

@@ -13,6 +13,10 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
 from langchain.docstore.document import Document
+from langchain.chains import LLMChain, RetrievalQA, ConversationalRetrievalChain
+from langchain.memory import ConversationBufferMemory
+from langchain.prompts import PromptTemplate
+from langchain_community.llms import HuggingFaceHub
 
 # Import document parsers
 import PyPDF2
@@ -30,8 +34,12 @@ HF_TOKEN = os.getenv("hf_token")
 if not HF_TOKEN:
     raise ValueError("HuggingFace token not found in environment variables")
 
-# Initialize HuggingFace Inference Client
-client = InferenceClient(model="meta-llama/Llama-3.1-8B-Instruct", token=HF_TOKEN)
+# Initialize HuggingFace LLM
+llm = HuggingFaceHub(
+    repo_id="meta-llama/Llama-3.1-8B-Instruct",
+    huggingfacehub_api_token=HF_TOKEN,
+    model_kwargs={"temperature": 0.7, "max_length": 512}
+)
 
 # Initialize embeddings
 embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
@@ -197,29 +205,60 @@ class IngestionAgent:
         self.message_bus.publish(response)
 
 class RetrievalAgent:
-    """Agent responsible for embedding and semantic retrieval"""
+    """Agent responsible for embedding and semantic retrieval using LangChain"""
     
     def __init__(self, message_bus: MessageBus):
         self.name = "RetrievalAgent"
         self.message_bus = message_bus
         self.message_bus.subscribe(self.name, self.handle_message)
         self.vector_store = None
+        self.retriever = None
+        self.qa_chain = None
+        self.conversation_chain = None
+        self.memory = ConversationBufferMemory(
+            memory_key="chat_history",
+            return_messages=True,
+            output_key="answer"
+        )
     
     def handle_message(self, message: MCPMessage):
         """Handle incoming MCP messages"""
         if message.type == "INGESTION_COMPLETE":
             self.create_vector_store(message)
         elif message.type == "RETRIEVAL_REQUEST":
-            self.retrieve_context(message)
+            self.process_query(message)
     
     def create_vector_store(self, message: MCPMessage):
-        """Create vector store from processed documents"""
+        """Create vector store and chains from processed documents"""
         documents = message.payload.get("documents", [])
         
         if documents:
             try:
                 self.vector_store = FAISS.from_documents(documents, embeddings)
-                logger.info(f"Vector store created with {len(documents)} documents")
+                self.retriever = self.vector_store.as_retriever(
+                    search_type="similarity",
+                    search_kwargs={"k": 3}
+                )
+                
+                # Create QA chain
+                self.qa_chain = RetrievalQA.from_chain_type(
+                    llm=llm,
+                    chain_type="stuff",
+                    retriever=self.retriever,
+                    return_source_documents=True,
+                    verbose=True
+                )
+                
+                # Create conversational chain
+                self.conversation_chain = ConversationalRetrievalChain.from_llm(
+                    llm=llm,
+                    retriever=self.retriever,
+                    memory=self.memory,
+                    return_source_documents=True,
+                    verbose=True
+                )
+                
+                logger.info(f"Vector store and chains created with {len(documents)} documents")
                 
                 # Notify completion
                 response = MCPMessage(
@@ -233,102 +272,60 @@ class RetrievalAgent:
             except Exception as e:
                 logger.error(f"Error creating vector store: {e}")
     
-    def retrieve_context(self, message: MCPMessage):
-        """Retrieve relevant context for a query"""
+    def process_query(self, message: MCPMessage):
+        """Process query using conversational retrieval chain"""
         query = message.payload.get("query", "")
-        k = message.payload.get("k", 3)
+        use_conversation = message.payload.get("use_conversation", True)
         
-        if self.vector_store and query:
-            try:
-                docs = self.vector_store.similarity_search(query, k=k)
-                context = [{"content": doc.page_content, "source": doc.metadata.get("source", "")} 
-                           for doc in docs]
-                
-                response = MCPMessage(
-                    sender=self.name,
-                    receiver="LLMResponseAgent",
-                    msg_type="CONTEXT_RESPONSE",
-                    trace_id=message.trace_id,
-                    payload={
-                        "query": query,
-                        "retrieved_context": context,
-                        "top_chunks": [doc.page_content for doc in docs]
-                    }
-                )
-                self.message_bus.publish(response)
-            except Exception as e:
-                logger.error(f"Error retrieving context: {e}")
-
-class LLMResponseAgent:
-    """Agent responsible for generating LLM responses"""
-    
-    def __init__(self, message_bus: MessageBus):
-        self.name = "LLMResponseAgent"
-        self.message_bus = message_bus
-        self.message_bus.subscribe(self.name, self.handle_message)
-    
-    def handle_message(self, message: MCPMessage):
-        """Handle incoming MCP messages"""
-        if message.type == "CONTEXT_RESPONSE":
-            self.generate_response(message)
-    
-    def generate_response(self, message: MCPMessage):
-        """Generate response using retrieved context"""
-        query = message.payload.get("query", "")
-        context = message.payload.get("retrieved_context", [])
-        
-        # Build context string
-        context_text = "\n\n".join([f"Source: {ctx['source']}\nContent: {ctx['content']}" 
-                                      for ctx in context])
+        if not self.qa_chain or not query:
+            return
         
-        # Create messages for conversational format
-        messages = [
-            {
-                "role": "system",
-                "content": "You are a helpful assistant. Based on the provided context below, answer the user's question accurately and comprehensively. Cite the sources if possible.",
-            },
-            {
-                "role": "user", 
-                "content": f"Context:\n\n{context_text}\n\nQuestion: {query}"
-            }
-        ]
-
         try:
-            # Use client.chat_completion for conversational models
-            response_stream = client.chat_completion(
-                messages=messages,
-                max_tokens=512,
-                temperature=0.7,
-                stream=True
-            )
+            if use_conversation and self.conversation_chain:
+                # Use conversational chain for context-aware responses
+                result = self.conversation_chain({"question": query})
+                answer = result["answer"]
+                source_docs = result.get("source_documents", [])
+            else:
+                # Use simple QA chain
+                result = self.qa_chain({"query": query})
+                answer = result["result"]
+                source_docs = result.get("source_documents", [])
+            
+            # Format sources
+            sources = []
+            for doc in source_docs:
+                sources.append({
+                    "content": doc.page_content[:200] + "...",
+                    "source": doc.metadata.get("source", "Unknown")
+                })
             
-            # Send streaming response
             response = MCPMessage(
                 sender=self.name,
                 receiver="CoordinatorAgent",
-                msg_type="LLM_RESPONSE_STREAM",
+                msg_type="CHAIN_RESPONSE",
                 trace_id=message.trace_id,
                 payload={
                     "query": query,
-                    "response_stream": response_stream,
-                    "context": context
+                    "answer": answer,
+                    "sources": sources
                 }
             )
             self.message_bus.publish(response)
             
         except Exception as e:
-            logger.error(f"Error generating response: {e}")
-            # Send an error stream back
-            error_msg = f"Error from LLM: {e}"
-            def error_generator():
-                yield error_msg
-            
+            logger.error(f"Error processing query: {e}")
+            # Send error response
             response = MCPMessage(
                 sender=self.name,
                 receiver="CoordinatorAgent",
-                msg_type="LLM_RESPONSE_STREAM",
+                msg_type="CHAIN_RESPONSE",
                 trace_id=message.trace_id,
-                payload={"response_stream": error_generator()}
+                payload={
+                    "query": query,
+                    "answer": f"Error processing query: {str(e)}",
+                    "sources": []
+                }
             )
             self.message_bus.publish(response)
 
@@ -339,15 +336,15 @@ class CoordinatorAgent:
         self.name = "CoordinatorAgent"
         self.message_bus = message_bus
         self.message_bus.subscribe(self.name, self.handle_message)
-        self.current_response_stream = None
         self.vector_store_ready = False
+        self.current_response = None
     
     def handle_message(self, message: MCPMessage):
         """Handle incoming MCP messages"""
         if message.type == "VECTORSTORE_READY":
             self.vector_store_ready = True
-        elif message.type == "LLM_RESPONSE_STREAM":
-            self.current_response_stream = message.payload.get("response_stream")
+        elif message.type == "CHAIN_RESPONSE":
+            self.current_response = message.payload
     
     def process_files(self, files):
         """Process uploaded files"""
@@ -367,59 +364,56 @@ class CoordinatorAgent:
         
         return f"Processing {len(files)} files: {', '.join([os.path.basename(fp) for fp in file_paths])}"
     
-    def handle_query(self, query: str, history: List) -> Generator[str, None, None]:
-        """Handle user query and return streaming response"""
+    def handle_query(self, query: str):
+        """Handle user query using LangChain chains"""
         if not self.vector_store_ready:
-            yield "Please upload and process documents first."
-            return
+            return "Please upload and process documents first."
         
         # Send retrieval request
         message = MCPMessage(
             sender=self.name,
             receiver="RetrievalAgent",
             msg_type="RETRIEVAL_REQUEST",
-            payload={"query": query}
+            payload={"query": query, "use_conversation": True}
         )
         self.message_bus.publish(message)
         
-        # Wait for response and stream
+        # Wait for response
         import time
-        timeout = 20  # seconds
+        timeout = 30  # seconds
         start_time = time.time()
         
-        while not self.current_response_stream and (time.time() - start_time) < timeout:
+        while not self.current_response and (time.time() - start_time) < timeout:
             time.sleep(0.1)
         
-        if self.current_response_stream:
-            try:
-                # Stream tokens directly
-                for chunk in self.current_response_stream:
-                    # The token is in chunk.choices[0].delta.content for chat_completion
-                    token = chunk.choices[0].delta.content
-                    if token:
-                        yield token
-            except Exception as e:
-                yield f"Error streaming response: {e}"
-            finally:
-                self.current_response_stream = None # Reset for next query
+        if self.current_response:
+            response = self.current_response
+            self.current_response = None  # Reset for next query
+            
+            # Format response with sources
+            answer = response.get("answer", "No answer generated.")
+            sources = response.get("sources", [])
+            
+            if sources:
+                source_text = "\n\n**Sources:**\n"
+                for i, source in enumerate(sources, 1):
+                    source_text += f"{i}. {source['source']}: {source['content']}\n"
+                answer += source_text
+            
+            return answer
         else:
-            yield "Timeout: No response received from LLM agent."
+            return "Timeout: No response received from the system."
 
 # Initialize agents
 ingestion_agent = IngestionAgent(message_bus)
 retrieval_agent = RetrievalAgent(message_bus)
-llm_response_agent = LLMResponseAgent(message_bus)
 coordinator_agent = CoordinatorAgent(message_bus)
 
 def create_interface():
     """Create ChatGPT-style Gradio interface"""
     
     with gr.Blocks(
-        theme=gr.themes.Base(
-            primary_hue="yellow",
-            secondary_hue="gray",
-            neutral_hue="slate"
-        ),
+        theme=gr.themes.Base(),
         css="""
         /* Dark theme styling */
         .gradio-container {
@@ -478,24 +472,6 @@ def create_interface():
             margin-bottom: 1rem;
         }
         
-        .gradio-chatbot .message {
-            background: rgba(45, 45, 45, 0.8) !important;
-            border: 1px solid rgba(255, 193, 7, 0.1) !important;
-            border-radius: 12px !important;
-            margin: 0.5rem 0 !important;
-        }
-        
-        .gradio-chatbot .message.user {
-            background: rgba(255, 193, 7, 0.1) !important;
-            border-color: rgba(255, 193, 7, 0.3) !important;
-            margin-left: 20% !important;
-        }
-        
-        .gradio-chatbot .message.bot {
-            background: rgba(45, 45, 45, 0.9) !important;
-            margin-right: 20% !important;
-        }
-        
         /* Input area */
         .input-area {
             background: rgba(45, 45, 45, 0.6);
@@ -515,11 +491,6 @@ def create_interface():
             transition: all 0.3s ease;
         }
         
-        .upload-area:hover {
-            background: rgba(255, 193, 7, 0.1);
-            border-color: rgba(255, 193, 7, 0.5);
-        }
-        
         /* Buttons */
         .primary-btn {
             background: linear-gradient(135deg, #ffc107 0%, #ff8f00 100%) !important;
@@ -527,13 +498,6 @@ def create_interface():
             border: none !important;
             border-radius: 8px !important;
             font-weight: 600 !important;
-            transition: all 0.3s ease !important;
-        }
-        
-        .primary-btn:hover {
-            background: linear-gradient(135deg, #ffcd38 0%, #ffa726 100%) !important;
-            transform: translateY(-1px);
-            box-shadow: 0 4px 12px rgba(255, 193, 7, 0.3);
         }
         
         /* Text inputs */
@@ -544,17 +508,6 @@ def create_interface():
             border-radius: 8px !important;
         }
         
-        .gradio-textbox input:focus, .gradio-textbox textarea:focus {
-            border-color: #ffc107 !important;
-            box-shadow: 0 0 0 2px rgba(255, 193, 7, 0.2) !important;
-        }
-        
-        /* File component */
-        .gradio-file {
-            background: transparent !important;
-            border: none !important;
-        }
-        
         /* Processing indicator */
         .processing-indicator {
             background: rgba(255, 193, 7, 0.1);
@@ -564,76 +517,40 @@ def create_interface():
             margin: 0.5rem 0;
             color: #ffc107;
             text-align: center;
-            animation: pulse 2s infinite;
-        }
-        
-        @keyframes pulse {
-            0%, 100% { opacity: 1; }
-            50% { opacity: 0.7; }
-        }
-        
-        /* Hide labels for cleaner look */
-        .gradio-textbox label,
-        .gradio-file label {
-            display: none !important;
-        }
-        
-        /* Scrollbar styling */
-        ::-webkit-scrollbar {
-            width: 8px;
-        }
-        
-        ::-webkit-scrollbar-track {
-            background: rgba(45, 45, 45, 0.3);
-            border-radius: 4px;
-        }
-        
-        ::-webkit-scrollbar-thumb {
-            background: rgba(255, 193, 7, 0.5);
-            border-radius: 4px;
-        }
-        
-        ::-webkit-scrollbar-thumb:hover {
-            background: rgba(255, 193, 7, 0.7);
         }
         """,
         title="Agentic RAG Assistant"
     ) as iface:
 
         # Header
-        with gr.Row(elem_classes=["header"]):
+        with gr.Row():
             with gr.Column():
                 gr.HTML("""
-                <div style="display: flex; align-items: center; gap: 1rem;">
-                    <div>
-                        <h1>🤖 Agentic RAG Assistant</h1>
-                        <p>Upload documents and ask questions - powered by Multi-Agent Architecture</p>
-                    </div>
+                <div class="header">
+                    <h1>🤖 Agentic RAG Assistant</h1>
+                    <p>Upload documents and ask questions - powered by LangChain Multi-Agent Architecture</p>
                 </div>
                 """)
 
         # Main chat container
-        with gr.Row(elem_classes=["chat-container"]):
+        with gr.Row():
             with gr.Column():
                 
                 # Chatbot
                 chatbot = gr.Chatbot(
                     value=[],
                     height=500,
-                    show_copy_button=True,
-                    bubble_full_width=False,
-                    elem_classes=["gradio-chatbot"]
+                    show_copy_button=True
                 )
                 
                 # Input area
-                with gr.Column(elem_classes=["input-area"]):
+                with gr.Column():
                     
-                    # File upload (initially hidden, shows when needed)
+                    # File upload
                     file_upload = gr.File(
                         file_count="multiple",
                         file_types=[".pdf", ".pptx", ".csv", ".docx", ".txt", ".md"],
-                        elem_classes=["upload-area"],
-                        visible=True
+                        label="Upload Documents"
                     )
                     
                     # Processing status
@@ -643,27 +560,10 @@ def create_interface():
                     with gr.Row():
                         msg_input = gr.Textbox(
                             placeholder="Upload documents above, then ask your questions here...",
-                            scale=5,
-                            max_lines=3,
-                            autofocus=True
-                        )
-                        send_btn = gr.Button(
-                            "Send", 
-                            scale=1, 
-                            elem_classes=["primary-btn"]
+                            label="Message",
+                            scale=4
                         )
-                    
-                    # Quick examples
-                    gr.Examples(
-                        examples=[
-                            "What are the main topics in the documents?",
-                            "Summarize the key findings",
-                            "What metrics are mentioned?",
-                            "What are the recommendations?"
-                        ],
-                        inputs=msg_input,
-                        elem_classes=["examples"]
-                    )
+                        send_btn = gr.Button("Send", scale=1, variant="primary")
 
         # State to track document processing
         doc_processed = gr.State(False)
@@ -686,9 +586,9 @@ def create_interface():
                 
                 # Wait a moment for processing to complete
                 import time
-                time.sleep(2)
+                time.sleep(3)
                 
-                success_html = f"""
+                success_html = """
                 <div style="background: rgba(76, 175, 80, 0.1); border: 1px solid rgba(76, 175, 80, 0.3); 
                            border-radius: 8px; padding: 0.75rem; color: #4caf50; text-align: center;">
                     ✅ Documents processed successfully! You can now ask questions.
@@ -712,15 +612,11 @@ def create_interface():
             if not message.strip():
                 return history, message
             
-            # Add user message
-            history.append([message, None])
+            # Get response from coordinator
+            response = coordinator_agent.handle_query(message)
             
-            # Generate response
-            response = ""
-            for token in coordinator_agent.handle_query(message, history):
-                response += token
-                history[-1][1] = response
-                yield history, ""
+            # Add to chat history
+            history.append([message, response])
             
             return history, ""
 
@@ -735,15 +631,13 @@ def create_interface():
         send_btn.click(
             respond,
             inputs=[msg_input, chatbot, doc_processed],
-            outputs=[chatbot, msg_input],
-            show_progress=True
+            outputs=[chatbot, msg_input]
         )
 
         msg_input.submit(
             respond,
             inputs=[msg_input, chatbot, doc_processed],
-            outputs=[chatbot, msg_input],
-            show_progress=True
+            outputs=[chatbot, msg_input]
         )
 
     return iface