Update app.py

app.py

@@ -8,7 +8,6 @@ import json
 import asyncio
 from dataclasses import dataclass, asdict
 import logging
-import sys
 
 # Document processing imports
 import PyPDF2
@@ -28,14 +27,8 @@ from huggingface_hub import InferenceClient
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# --- Get HF token from environment and perform a crucial check ---
-HF_TOKEN = os.getenv('HF_TOKEN')
-
-if HF_TOKEN is None:
-    logger.error("FATAL ERROR: HuggingFace token (HF_TOKEN) environment variable is not set.")
-    logger.error("Please set the HF_TOKEN environment variable before running the application.")
-    sys.exit("HuggingFace token (HF_TOKEN) is not set. Exiting.")
-
+# Get HF token from environment
+HF_TOKEN = os.getenv('hf_token')
 
 # MCP Message Structure
 @dataclass
@@ -46,11 +39,11 @@ class MCPMessage:
     trace_id: str
     payload: Dict[str, Any]
     timestamp: str = None
-
+    
     def __post_init__(self):
         if self.timestamp is None:
             self.timestamp = datetime.now().isoformat()
-
+    
     def to_dict(self):
         return asdict(self)
 
@@ -59,18 +52,18 @@ class MCPCommunicator:
     def __init__(self):
         self.message_queue = asyncio.Queue()
         self.subscribers = {}
-
+    
     async def send_message(self, message: MCPMessage):
         logger.info(f"MCP: {message.sender} -> {message.receiver}: {message.type}")
         await self.message_queue.put(message)
-
+    
     async def receive_message(self, agent_name: str) -> MCPMessage:
         while True:
             message = await self.message_queue.get()
             if message.receiver == agent_name:
                 return message
             # Re-queue if not for this agent
-            await self.mcp.put(message) # Corrected: Use mcp.put instead of message_queue.put directly
+            await self.message_queue.put(message)
 
 # Global MCP instance
 mcp = MCPCommunicator()
@@ -80,7 +73,7 @@ class BaseAgent:
     def __init__(self, name: str):
         self.name = name
         self.mcp = mcp
-
+    
     async def send_mcp_message(self, receiver: str, msg_type: str, payload: Dict[str, Any], trace_id: str):
         message = MCPMessage(
             sender=self.name,
@@ -90,7 +83,7 @@ class BaseAgent:
             payload=payload
         )
         await self.mcp.send_message(message)
-
+    
     async def receive_mcp_message(self) -> MCPMessage:
         return await self.mcp.receive_message(self.name)
 
@@ -103,7 +96,7 @@ class IngestionAgent(BaseAgent):
             chunk_overlap=200,
             length_function=len,
         )
-
+    
     def parse_pdf(self, file_path: str) -> str:
         """Parse PDF file and extract text"""
         try:
@@ -116,7 +109,7 @@ class IngestionAgent(BaseAgent):
         except Exception as e:
             logger.error(f"Error parsing PDF: {e}")
             return ""
-
+    
     def parse_docx(self, file_path: str) -> str:
         """Parse DOCX file and extract text"""
         try:
@@ -128,7 +121,7 @@ class IngestionAgent(BaseAgent):
         except Exception as e:
             logger.error(f"Error parsing DOCX: {e}")
             return ""
-
+    
     def parse_pptx(self, file_path: str) -> str:
         """Parse PPTX file and extract text"""
         try:
@@ -142,9 +135,9 @@ class IngestionAgent(BaseAgent):
                 text += "\n"
             return text
         except Exception as e:
-                logger.error(f"Error parsing PPTX: {e}")
-                return ""
-
+            logger.error(f"Error parsing PPTX: {e}")
+            return ""
+    
     def parse_csv(self, file_path: str) -> str:
         """Parse CSV file and convert to text"""
         try:
@@ -153,7 +146,7 @@ class IngestionAgent(BaseAgent):
         except Exception as e:
             logger.error(f"Error parsing CSV: {e}")
             return ""
-
+    
     def parse_txt_md(self, file_path: str) -> str:
         """Parse TXT or MD file"""
         try:
@@ -166,15 +159,15 @@ class IngestionAgent(BaseAgent):
         except Exception as e:
             logger.error(f"Error parsing TXT/MD: {e}")
             return ""
-
+    
     async def process_documents(self, files: List[str], trace_id: str) -> List[LCDocument]:
         """Process uploaded documents and return chunked documents"""
         all_documents = []
-
-        for file_path in files: # file_path is already the path to the temporary file
+        
+        for file_path in files:
             file_ext = os.path.splitext(file_path)[1].lower()
             filename = os.path.basename(file_path)
-
+            
             # Parse based on file extension
             if file_ext == '.pdf':
                 content = self.parse_pdf(file_path)
@@ -189,11 +182,11 @@ class IngestionAgent(BaseAgent):
             else:
                 logger.warning(f"Unsupported file type: {file_ext}")
                 continue
-
+            
             if content.strip():
                 # Split content into chunks
                 chunks = self.text_splitter.split_text(content)
-
+                
                 # Create LangChain documents
                 for i, chunk in enumerate(chunks):
                     doc = LCDocument(
@@ -205,7 +198,7 @@ class IngestionAgent(BaseAgent):
                         }
                     )
                     all_documents.append(doc)
-
+        
         return all_documents
 
 # Retrieval Agent
@@ -216,7 +209,7 @@ class RetrievalAgent(BaseAgent):
             model_name="sentence-transformers/all-MiniLM-L6-v2"
         )
         self.vector_store = None
-
+    
     async def create_vector_store(self, documents: List[LCDocument], trace_id: str):
         """Create vector store from documents"""
         try:
@@ -227,16 +220,16 @@ class RetrievalAgent(BaseAgent):
                 logger.warning("No documents to create vector store")
         except Exception as e:
             logger.error(f"Error creating vector store: {e}")
-
+    
     async def retrieve_relevant_chunks(self, query: str, k: int = 5, trace_id: str = None) -> List[Dict]:
         """Retrieve relevant chunks for a query"""
         if not self.vector_store:
             return []
-
+        
         try:
             # Similarity search
             docs = self.vector_store.similarity_search(query, k=k)
-
+            
             # Format results
             results = []
             for doc in docs:
@@ -246,7 +239,7 @@ class RetrievalAgent(BaseAgent):
                     "chunk_id": doc.metadata.get("chunk_id", 0),
                     "file_type": doc.metadata.get("file_type", "Unknown")
                 })
-
+            
             return results
         except Exception as e:
             logger.error(f"Error retrieving chunks: {e}")
@@ -260,18 +253,15 @@ class LLMResponseAgent(BaseAgent):
             model="meta-llama/Llama-3.1-8B-Instruct",
             token=HF_TOKEN
         )
-
-    def format_prompt_for_conversational(self, query: str, context_chunks: List[Dict]) -> str:
-        """
-        Format prompt with context and query as a single 'user' input
-        suitable for a conversational model.
-        """
+    
+    def format_prompt(self, query: str, context_chunks: List[Dict]) -> str:
+        """Format prompt with context and query"""
         context_text = "\n\n".join([
             f"Source: {chunk['source']}\nContent: {chunk['content']}"
             for chunk in context_chunks
         ])
-
-        prompt_as_user_input = f"""Based on the following context from uploaded documents, please answer the user's question.
+        
+        prompt = f"""Based on the following context from uploaded documents, please answer the user's question.
 
 Context:
 {context_text}
@@ -281,29 +271,26 @@ Question: {query}
 Please provide a comprehensive answer based on the context above. If the context doesn't contain enough information to fully answer the question, please mention what information is available and what might be missing.
 
 Answer:"""
-        return prompt_as_user_input
-
+        
+        return prompt
+    
     async def generate_response(self, query: str, context_chunks: List[Dict], trace_id: str) -> str:
-        """Generate response using LLM via the conversational task."""
+        """Generate response using LLM"""
         try:
-            formatted_input = self.format_prompt_for_conversational(query, context_chunks)
-
-            response = self.client.conversational(
-                inputs=formatted_input,
-                parameters={
-                    "temperature": 0.7,
-                    "max_new_tokens": 512,
-                }
+            prompt = self.format_prompt(query, context_chunks)
+            
+            # Generate response using HuggingFace Inference
+            response = self.client.text_generation(
+                prompt,
+                max_new_tokens=512,
+                temperature=0.7,
+                do_sample=True,
+                return_full_text=False
             )
-
-            if response.generated_responses:
-                return response.generated_responses[0]
-            else:
-                logger.warning("LLM generated an empty response via conversational API.")
-                return "I apologize, the model did not generate a response."
-
+            
+            return response
         except Exception as e:
-            logger.error(f"Error generating response with conversational LLM: {e}")
+            logger.error(f"Error generating response: {e}")
             return f"I apologize, but I encountered an error while generating the response: {str(e)}"
 
 # Coordinator Agent
@@ -314,66 +301,70 @@ class CoordinatorAgent(BaseAgent):
         self.retrieval_agent = RetrievalAgent()
         self.llm_agent = LLMResponseAgent()
         self.documents_processed = False
-
+    
     async def process_documents(self, files: List[str]) -> str:
         """Orchestrate document processing"""
         trace_id = str(uuid.uuid4())
-
+        
         try:
+            # Step 1: Ingestion
             await self.send_mcp_message(
-                "IngestionAgent",
-                "DOCUMENT_INGESTION_REQUEST",
-                {"files": files},
+                "IngestionAgent", 
+                "DOCUMENT_INGESTION_REQUEST", 
+                {"files": files}, 
                 trace_id
             )
-
+            
             documents = await self.ingestion_agent.process_documents(files, trace_id)
-
+            
             await self.send_mcp_message(
-                "RetrievalAgent",
-                "VECTOR_STORE_CREATE_REQUEST",
-                {"documents": len(documents)},
+                "RetrievalAgent", 
+                "VECTOR_STORE_CREATE_REQUEST", 
+                {"documents": len(documents)}, 
                 trace_id
             )
-
+            
+            # Step 2: Create vector store
             await self.retrieval_agent.create_vector_store(documents, trace_id)
-
+            
             self.documents_processed = True
-
+            
             return f"Successfully processed {len(documents)} document chunks from {len(files)} files."
-
+            
         except Exception as e:
             logger.error(f"Error in document processing: {e}")
             return f"Error processing documents: {str(e)}"
-
+    
     async def answer_query(self, query: str) -> tuple[str, List[Dict]]:
         """Orchestrate query answering"""
         if not self.documents_processed:
             return "Please upload and process documents first.", []
-
+        
         trace_id = str(uuid.uuid4())
-
+        
         try:
+            # Step 1: Retrieval
             await self.send_mcp_message(
-                "RetrievalAgent",
-                "RETRIEVAL_REQUEST",
-                {"query": query},
+                "RetrievalAgent", 
+                "RETRIEVAL_REQUEST", 
+                {"query": query}, 
                 trace_id
             )
-
-            context_chunks = await self.retrierieval_agent.retrieve_relevant_chunks(query, k=5, trace_id=trace_id)
-
+            
+            context_chunks = await self.retrieval_agent.retrieve_relevant_chunks(query, k=5, trace_id=trace_id)
+            
+            # Step 2: LLM Response
             await self.send_mcp_message(
-                "LLMResponseAgent",
-                "LLM_GENERATION_REQUEST",
-                {"query": query, "context_chunks": len(context_chunks)},
+                "LLMResponseAgent", 
+                "LLM_GENERATION_REQUEST", 
+                {"query": query, "context_chunks": len(context_chunks)}, 
                 trace_id
             )
-
+            
             response = await self.llm_agent.generate_response(query, context_chunks, trace_id)
-
+            
             return response, context_chunks
-
+            
         except Exception as e:
             logger.error(f"Error in query processing: {e}")
             return f"Error processing query: {str(e)}", []
@@ -382,40 +373,44 @@ class CoordinatorAgent(BaseAgent):
 coordinator = CoordinatorAgent()
 
 async def process_files(files):
-    """Process uploaded files (already temporary file paths from Gradio)"""
+    """Process uploaded files"""
     if not files:
         return "❌ Please upload at least one file."
-
-    # Gradio's gr.File component with default type="filepath" already provides
-    # the temporary file paths. We just need to pass them to the ingestion agent.
-    # The 'files' variable here is already a list of strings (filepaths).
-    file_paths_for_ingestion = [file.name for file in files] # Extract the actual path string
-
-    result = await coordinator.process_documents(file_paths_for_ingestion)
-
-    # Cleanup is handled by Gradio's tmp directory cleanup, but explicit removal is also fine.
-    # No need to manually create temp files or read content.
-
+    
+    # Save uploaded files to temporary directory
+    file_paths = []
+    for file in files:
+        # Handle file path - Gradio returns file path as string
+        if hasattr(file, 'name'):
+            file_path = file.name
+        else:
+            file_path = str(file)
+        file_paths.append(file_path)
+    
+    result = await coordinator.process_documents(file_paths)
+    
     return result
 
 async def answer_question(query, history):
     """Answer user question"""
     if not query.strip():
         return history, ""
-
+    
     response, context_chunks = await coordinator.answer_query(query)
-
+    
+    # Format response with sources
     if context_chunks:
         sources = "\n\n**Sources:**\n"
-        for i, chunk in enumerate(context_chunks[:3], 1):
+        for i, chunk in enumerate(context_chunks[:3], 1):  # Show top 3 sources
             sources += f"{i}. {chunk['source']} (Chunk {chunk['chunk_id']})\n"
         response += sources
-
+    
+    # Add to chat history
     history.append((query, response))
-
+    
     return history, ""
 
-# Custom CSS (unchanged)
+# Custom CSS
 custom_css = """
 /* Main container styling */
 .gradio-container {
@@ -458,7 +453,7 @@ custom_css = """
 }
 
 /* Card styling */
-.upload-card, .chat-card {
+.setup-card, .upload-card, .chat-card {
     background: white !important;
     border-radius: 15px !important;
     padding: 2rem !important;
@@ -558,8 +553,8 @@ custom_css = """
     .header-title {
         font-size: 2rem !important;
     }
-
-    .upload-card, .chat-card {
+    
+    .setup-card, .upload-card, .chat-card {
         padding: 1.5rem !important;
     }
 }
@@ -585,9 +580,9 @@ def create_interface():
             <p class="header-subtitle">Multi-Format Document QA using Model Context Protocol (MCP)</p>
         </div>
         """)
-
+        
         with gr.Tabs() as tabs:
-            # Upload Tab (now the first tab)
+            # Upload Tab
             with gr.TabItem("📁 Upload Documents", elem_classes=["tab-nav"]):
                 gr.HTML("""
                 <div class="upload-card">
@@ -595,28 +590,26 @@ def create_interface():
                     <p>Upload your documents in any supported format: PDF, DOCX, PPTX, CSV, TXT, or Markdown.</p>
                 </div>
                 """)
-
+                
                 file_upload = gr.File(
                     label="Choose Files",
                     file_count="multiple",
                     file_types=[".pdf", ".docx", ".pptx", ".csv", ".txt", ".md"],
-                    # type="filepath" is the default, but explicitly setting it helps clarify
-                    type="filepath",
                     elem_classes=["file-upload"]
                 )
-
+                
                 upload_button = gr.Button(
-                    "Process Documents",
+                    "Process Documents", 
                     variant="primary",
                     elem_classes=["primary-button"]
                 )
-
+                
                 upload_status = gr.Textbox(
                     label="Processing Status",
                     interactive=False,
                     elem_classes=["input-container"]
                 )
-
+            
             # Chat Tab
             with gr.TabItem("💬 Chat", elem_classes=["tab-nav"]):
                 gr.HTML("""
@@ -625,13 +618,13 @@ def create_interface():
                     <p>Ask questions about your uploaded documents. The AI will provide answers based on the document content.</p>
                 </div>
                 """)
-
+                
                 chatbot = gr.Chatbot(
                     label="Conversation",
                     height=400,
                     elem_classes=["chat-container"]
                 )
-
+                
                 with gr.Row():
                     query_input = gr.Textbox(
                         label="Your Question",
@@ -639,11 +632,11 @@ def create_interface():
                         elem_classes=["input-container"]
                     )
                     ask_button = gr.Button(
-                        "Ask",
+                        "Ask", 
                         variant="primary",
                         elem_classes=["primary-button"]
                     )
-
+                
                 gr.Examples(
                     examples=[
                         "What are the main topics covered in the documents?",
@@ -654,7 +647,7 @@ def create_interface():
                     inputs=query_input,
                     label="Example Questions"
                 )
-
+            
             # Architecture Tab
             with gr.TabItem("🏗️ Architecture", elem_classes=["tab-nav"]):
                 gr.HTML("""
@@ -663,38 +656,38 @@ def create_interface():
                     <p>This system uses an agentic architecture with Model Context Protocol (MCP) for inter-agent communication.</p>
                 </div>
                 """)
-
+                
                 gr.Markdown("""
                 ## 🔄 Agent Flow Diagram
-
+                
                 ```
                 User Upload → CoordinatorAgent → IngestionAgent → RetrievalAgent → LLMResponseAgent
                      ↓              ↓                ↓               ↓              ↓
                  Documents    MCP Messages    Text Chunks    Vector Store    Final Response
                 ```
-
+                
                 ## 🤖 Agent Descriptions
-
+                
                 - **CoordinatorAgent**: Orchestrates the entire workflow and manages MCP communication
                 - **IngestionAgent**: Parses and preprocesses documents (PDF, DOCX, PPTX, CSV, TXT, MD)
                 - **RetrievalAgent**: Handles embeddings and semantic retrieval using FAISS
                 - **LLMResponseAgent**: Generates final responses using Llama-3.1-8B-Instruct
-
+                
                 ## 🔗 Tech Stack
-
+                
                 - **Frontend**: Gradio with custom CSS
                 - **LLM**: Meta Llama-3.1-8B-Instruct (via HuggingFace Inference)
                 - **Embeddings**: sentence-transformers/all-MiniLM-L6-v2
                 - **Vector Store**: FAISS
                 - **Document Processing**: PyPDF2, python-docx, python-pptx, pandas
                 - **Framework**: LangChain for document handling
-
+                
                 ## 📨 MCP Message Example
-
+                
                 ```json
                 {
                   "sender": "RetrievalAgent",
-                  "receiver": "LLMResponseAgent",
+                  "receiver": "LLMResponseAgent", 
                   "type": "RETRIEVAL_RESULT",
                   "trace_id": "rag-457",
                   "payload": {
@@ -705,26 +698,26 @@ def create_interface():
                 }
                 ```
                 """)
-
+        
         # Event handlers
         upload_button.click(
             fn=process_files,
             inputs=[file_upload],
             outputs=[upload_status]
         )
-
+        
         ask_button.click(
             fn=answer_question,
             inputs=[query_input, chatbot],
             outputs=[chatbot, query_input]
         )
-
+        
         query_input.submit(
             fn=answer_question,
             inputs=[query_input, chatbot],
             outputs=[chatbot, query_input]
         )
-
+    
     return demo
 
 if __name__ == "__main__":