Update app.py

app.py

@@ -8,6 +8,7 @@ import json
 import asyncio
 from dataclasses import dataclass, asdict
 import logging
+import sys # Import sys for exiting if token is missing
 
 # Document processing imports
 import PyPDF2
@@ -27,7 +28,7 @@ from huggingface_hub import InferenceClient
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Get HF token from environment
+# --- Get HF token from environment and perform a crucial check ---
 HF_TOKEN = os.getenv('hf_token')
 
 if HF_TOKEN is None:
@@ -45,11 +46,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)
 
@@ -58,11 +59,11 @@ 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()
@@ -79,7 +80,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,
@@ -89,7 +90,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)
 
@@ -102,7 +103,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:
@@ -115,7 +116,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:
@@ -127,7 +128,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:
@@ -141,9 +142,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:
@@ -152,7 +153,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:
@@ -165,15 +166,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_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)
@@ -188,11 +189,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(
@@ -204,7 +205,7 @@ class IngestionAgent(BaseAgent):
                         }
                     )
                     all_documents.append(doc)
-        
+
         return all_documents
 
 # Retrieval Agent
@@ -215,7 +216,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:
@@ -226,16 +227,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:
@@ -245,7 +246,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}")
@@ -255,19 +256,25 @@ class RetrievalAgent(BaseAgent):
 class LLMResponseAgent(BaseAgent):
     def __init__(self):
         super().__init__("LLMResponseAgent")
+        # Use the global HF_TOKEN which is validated at script start
         self.client = InferenceClient(
             model="meta-llama/Llama-3.1-8B-Instruct",
-            token=HF_TOKEN
+            token=HF_TOKEN # Pass the token here
         )
-    
-    def format_prompt(self, query: str, context_chunks: List[Dict]) -> str:
-        """Format prompt with context and query"""
+
+    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.
+        """
         context_text = "\n\n".join([
             f"Source: {chunk['source']}\nContent: {chunk['content']}"
             for chunk in context_chunks
         ])
-        
-        prompt = f"""Based on the following context from uploaded documents, please answer the user's question.
+
+        # We are putting the RAG prompt into the 'user' input for the conversational model.
+        # This is a common way to use a conversational model for RAG if text_generation isn't available.
+        prompt_as_user_input = f"""Based on the following context from uploaded documents, please answer the user's question.
 
 Context:
 {context_text}
@@ -276,27 +283,40 @@ 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
-    
+Answer:""" # Keeping "Answer:" to guide the model to start generating the answer directly.
+        return prompt_as_user_input
+
     async def generate_response(self, query: str, context_chunks: List[Dict], trace_id: str) -> str:
-        """Generate response using LLM"""
+        """Generate response using LLM via the conversational task."""
         try:
-            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
+            # Format the RAG prompt as the user's input for the conversational model
+            formatted_input = self.format_prompt_for_conversational(query, context_chunks)
+
+            # Use the conversational task
+            response = self.client.conversational(
+                inputs=formatted_input, # This is the current user turn
+                # No past_user_inputs or generated_responses are provided initially
+                # to keep it stateless per query, akin to text_generation.
+                parameters={
+                    "temperature": 0.7,
+                    "max_new_tokens": 512,
+                    # Add other parameters if needed, e.g., do_sample, top_p, top_k
+                    # "do_sample": True,
+                    # "top_p": 0.95,
+                    # "top_k": 50,
+                }
             )
-            
-            return response
+
+            # The conversational response has a list of generated responses.
+            # We assume the first one is the primary answer.
+            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."
+
         except Exception as e:
-            logger.error(f"Error generating response: {e}")
+            logger.error(f"Error generating response with conversational LLM: {e}")
             return f"I apologize, but I encountered an error while generating the response: {str(e)}"
 
 # Coordinator Agent
@@ -305,72 +325,72 @@ class CoordinatorAgent(BaseAgent):
         super().__init__("CoordinatorAgent")
         self.ingestion_agent = IngestionAgent()
         self.retrieval_agent = RetrievalAgent()
-        self.llm_agent = LLMResponseAgent()
+        self.llm_agent = LLMResponseAgent() # LLMResponseAgent will use the global HF_TOKEN
         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.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,41 +402,49 @@ async def process_files(files):
     """Process uploaded files"""
     if not files:
         return "❌ Please upload at least one file."
-    
-    # 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)
-    
+        temp_dir = tempfile.gettempdir()
+        unique_filename = f"{uuid.uuid4()}_{os.path.basename(file.name)}"
+        temp_path = os.path.join(temp_dir, unique_filename)
+        try:
+            file_content = file.read()
+            with open(temp_path, 'wb') as f:
+                f.write(file_content)
+            file_paths.append(temp_path)
+        except Exception as e:
+            logger.error(f"Error saving uploaded file {file.name}: {e}")
+            return f"❌ Error saving uploaded file {file.name}: {e}"
+
     result = await coordinator.process_documents(file_paths)
-    
+
+    for path in file_paths:
+        try:
+            os.remove(path)
+        except Exception as e:
+            logger.warning(f"Could not remove temporary file {path}: {e}")
+
     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):  # Show top 3 sources
+        for i, chunk in enumerate(context_chunks[:3], 1):
             sources += f"{i}. {chunk['source']} (Chunk {chunk['chunk_id']})\n"
         response += sources
-    
-    # Add to chat history
+
     history.append((query, response))
-    
+
     return history, ""
 
-# Custom CSS
+# Custom CSS (unchanged)
 custom_css = """
 /* Main container styling */
 .gradio-container {
@@ -459,7 +487,7 @@ custom_css = """
 }
 
 /* Card styling */
-.setup-card, .upload-card, .chat-card {
+.upload-card, .chat-card {
     background: white !important;
     border-radius: 15px !important;
     padding: 2rem !important;
@@ -559,8 +587,8 @@ custom_css = """
     .header-title {
         font-size: 2rem !important;
     }
-    
-    .setup-card, .upload-card, .chat-card {
+
+    .upload-card, .chat-card {
         padding: 1.5rem !important;
     }
 }
@@ -586,9 +614,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
+            # Upload Tab (now the first tab)
             with gr.TabItem("📁 Upload Documents", elem_classes=["tab-nav"]):
                 gr.HTML("""
                 <div class="upload-card">
@@ -596,26 +624,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"],
                     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("""
@@ -624,13 +652,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",
@@ -638,11 +666,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?",
@@ -653,7 +681,7 @@ def create_interface():
                     inputs=query_input,
                     label="Example Questions"
                 )
-            
+
             # Architecture Tab
             with gr.TabItem("🏗️ Architecture", elem_classes=["tab-nav"]):
                 gr.HTML("""
@@ -662,38 +690,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": {
@@ -704,26 +732,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__":