Update app.py

app.py

@@ -8,7 +8,7 @@ import json
 import asyncio
 from dataclasses import dataclass, asdict
 import logging
-import sys # Import sys for exiting if token is missing
+import sys
 
 # Document processing imports
 import PyPDF2
@@ -29,7 +29,7 @@ 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')
+HF_TOKEN = os.getenv('HF_TOKEN')
 
 if HF_TOKEN is None:
     logger.error("FATAL ERROR: HuggingFace token (HF_TOKEN) environment variable is not set.")
@@ -70,7 +70,7 @@ class MCPCommunicator:
             if message.receiver == agent_name:
                 return message
             # Re-queue if not for this agent
-            await self.message_queue.put(message)
+            await self.mcp.put(message) # Corrected: Use mcp.put instead of message_queue.put directly
 
 # Global MCP instance
 mcp = MCPCommunicator()
@@ -171,7 +171,7 @@ class IngestionAgent(BaseAgent):
         """Process uploaded documents and return chunked documents"""
         all_documents = []
 
-        for file_path in files:
+        for file_path in files: # file_path is already the path to the temporary file
             file_ext = os.path.splitext(file_path)[1].lower()
             filename = os.path.basename(file_path)
 
@@ -256,10 +256,9 @@ 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 # Pass the token here
+            token=HF_TOKEN
         )
 
     def format_prompt_for_conversational(self, query: str, context_chunks: List[Dict]) -> str:
@@ -272,8 +271,6 @@ class LLMResponseAgent(BaseAgent):
             for chunk in context_chunks
         ])
 
-        # 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:
@@ -283,32 +280,22 @@ 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:""" # Keeping "Answer:" to guide the model to start generating the answer directly.
+Answer:"""
         return prompt_as_user_input
 
     async def generate_response(self, query: str, context_chunks: List[Dict], trace_id: str) -> str:
         """Generate response using LLM via the conversational task."""
         try:
-            # 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.
+                inputs=formatted_input,
                 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,
                 }
             )
 
-            # 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:
@@ -325,7 +312,7 @@ class CoordinatorAgent(BaseAgent):
         super().__init__("CoordinatorAgent")
         self.ingestion_agent = IngestionAgent()
         self.retrieval_agent = RetrievalAgent()
-        self.llm_agent = LLMResponseAgent() # LLMResponseAgent will use the global HF_TOKEN
+        self.llm_agent = LLMResponseAgent()
         self.documents_processed = False
 
     async def process_documents(self, files: List[str]) -> str:
@@ -333,7 +320,6 @@ class CoordinatorAgent(BaseAgent):
         trace_id = str(uuid.uuid4())
 
         try:
-            # Step 1: Ingestion
             await self.send_mcp_message(
                 "IngestionAgent",
                 "DOCUMENT_INGESTION_REQUEST",
@@ -350,7 +336,6 @@ class CoordinatorAgent(BaseAgent):
                 trace_id
             )
 
-            # Step 2: Create vector store
             await self.retrieval_agent.create_vector_store(documents, trace_id)
 
             self.documents_processed = True
@@ -369,7 +354,6 @@ class CoordinatorAgent(BaseAgent):
         trace_id = str(uuid.uuid4())
 
         try:
-            # Step 1: Retrieval
             await self.send_mcp_message(
                 "RetrievalAgent",
                 "RETRIEVAL_REQUEST",
@@ -377,9 +361,8 @@ class CoordinatorAgent(BaseAgent):
                 trace_id
             )
 
-            context_chunks = await self.retrieval_agent.retrieve_relevant_chunks(query, k=5, trace_id=trace_id)
+            context_chunks = await self.retrierieval_agent.retrieve_relevant_chunks(query, k=5, trace_id=trace_id)
 
-            # Step 2: LLM Response
             await self.send_mcp_message(
                 "LLMResponseAgent",
                 "LLM_GENERATION_REQUEST",
@@ -399,31 +382,19 @@ class CoordinatorAgent(BaseAgent):
 coordinator = CoordinatorAgent()
 
 async def process_files(files):
-    """Process uploaded files"""
+    """Process uploaded files (already temporary file paths from Gradio)"""
     if not files:
         return "❌ Please upload at least one file."
 
-    file_paths = []
-    for file in files:
-        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}"
+    # 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)
+    result = await coordinator.process_documents(file_paths_for_ingestion)
 
-    for path in file_paths:
-        try:
-            os.remove(path)
-        except Exception as e:
-            logger.warning(f"Could not remove temporary file {path}: {e}")
+    # 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.
 
     return result
 
@@ -629,6 +600,8 @@ def create_interface():
                     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"]
                 )