app.py

import gradio as gr
import os
import json
import uuid
import asyncio
from datetime import datetime
from typing import List, Dict, Any, Optional, Generator
import logging

# Import required libraries
from huggingface_hub import InferenceClient
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.docstore.document import Document

# Import document parsers
import PyPDF2
from pptx import Presentation
import pandas as pd
from docx import Document as DocxDocument
import io

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Get HuggingFace token from environment
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 embeddings
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

class MCPMessage:
    """Model Context Protocol Message Structure"""
    def __init__(self, sender: str, receiver: str, msg_type: str,
                 trace_id: str = None, payload: Dict = None):
        self.sender = sender
        self.receiver = receiver
        self.type = msg_type
        self.trace_id = trace_id or str(uuid.uuid4())
        self.payload = payload or {}
        self.timestamp = datetime.now().isoformat()
    
    def to_dict(self):
        return {
            "sender": self.sender,
            "receiver": self.receiver,
            "type": self.type,
            "trace_id": self.trace_id,
            "payload": self.payload,
            "timestamp": self.timestamp
        }

class MessageBus:
    """In-memory message bus for MCP communication"""
    def __init__(self):
        self.messages = []
        self.subscribers = {}
    
    def publish(self, message: MCPMessage):
        """Publish message to the bus"""
        self.messages.append(message)
        logger.info(f"Message published: {message.sender} -> {message.receiver} [{message.type}]")
        
        # Notify subscribers
        if message.receiver in self.subscribers:
            for callback in self.subscribers[message.receiver]:
                callback(message)
    
    def subscribe(self, agent_name: str, callback):
        """Subscribe agent to receive messages"""
        if agent_name not in self.subscribers:
            self.subscribers[agent_name] = []
        self.subscribers[agent_name].append(callback)

# Global message bus
message_bus = MessageBus()

class IngestionAgent:
    """Agent responsible for document parsing and preprocessing"""
    
    def __init__(self, message_bus: MessageBus):
        self.name = "IngestionAgent"
        self.message_bus = message_bus
        self.message_bus.subscribe(self.name, self.handle_message)
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200
        )
    
    def handle_message(self, message: MCPMessage):
        """Handle incoming MCP messages"""
        if message.type == "INGESTION_REQUEST":
            self.process_documents(message)
    
    def parse_pdf(self, file_path: str) -> str:
        """Parse PDF document"""
        try:
            with open(file_path, 'rb') as file:
                pdf_reader = PyPDF2.PdfReader(file)
                text = ""
                for page in pdf_reader.pages:
                    text += page.extract_text()
                return text
        except Exception as e:
            logger.error(f"Error parsing PDF: {e}")
            return ""
    
    def parse_pptx(self, file_path: str) -> str:
        """Parse PPTX document"""
        try:
            prs = Presentation(file_path)
            text = ""
            for slide in prs.slides:
                for shape in slide.shapes:
                    if hasattr(shape, "text"):
                        text += shape.text + "\n"
            return text
        except Exception as e:
            logger.error(f"Error parsing PPTX: {e}")
            return ""
    
    def parse_csv(self, file_path: str) -> str:
        """Parse CSV document"""
        try:
            df = pd.read_csv(file_path)
            return df.to_string()
        except Exception as e:
            logger.error(f"Error parsing CSV: {e}")
            return ""
    
    def parse_docx(self, file_path: str) -> str:
        """Parse DOCX document"""
        try:
            doc = DocxDocument(file_path)
            text = ""
            for paragraph in doc.paragraphs:
                text += paragraph.text + "\n"
            return text
        except Exception as e:
            logger.error(f"Error parsing DOCX: {e}")
            return ""
    
    def parse_txt(self, file_path: str) -> str:
        """Parse TXT/Markdown document"""
        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                return file.read()
        except Exception as e:
            logger.error(f"Error parsing TXT: {e}")
            return ""
    
    def process_documents(self, message: MCPMessage):
        """Process uploaded documents"""
        files = message.payload.get("files", [])
        processed_docs = []
        
        for file_path in files:
            file_ext = os.path.splitext(file_path)[1].lower()
            
            # Parse document based on file type
            if file_ext == '.pdf':
                text = self.parse_pdf(file_path)
            elif file_ext == '.pptx':
                text = self.parse_pptx(file_path)
            elif file_ext == '.csv':
                text = self.parse_csv(file_path)
            elif file_ext == '.docx':
                text = self.parse_docx(file_path)
            elif file_ext in ['.txt', '.md']:
                text = self.parse_txt(file_path)
            else:
                logger.warning(f"Unsupported file type: {file_ext}")
                continue
            
            if text:
                # Split text into chunks
                chunks = self.text_splitter.split_text(text)
                docs = [Document(page_content=chunk, metadata={"source": file_path}) 
                        for chunk in chunks]
                processed_docs.extend(docs)
        
        # Send processed documents to RetrievalAgent
        response = MCPMessage(
            sender=self.name,
            receiver="RetrievalAgent",
            msg_type="INGESTION_COMPLETE",
            trace_id=message.trace_id,
            payload={"documents": processed_docs}
        )
        self.message_bus.publish(response)

class RetrievalAgent:
    """Agent responsible for embedding and semantic retrieval"""
    
    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
    
    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)
    
    def create_vector_store(self, message: MCPMessage):
        """Create vector store 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")
                
                # Notify completion
                response = MCPMessage(
                    sender=self.name,
                    receiver="CoordinatorAgent",
                    msg_type="VECTORSTORE_READY",
                    trace_id=message.trace_id,
                    payload={"status": "ready"}
                )
                self.message_bus.publish(response)
            except Exception as e:
                logger.error(f"Error creating vector store: {e}")
    
    def retrieve_context(self, message: MCPMessage):
        """Retrieve relevant context for a query"""
        query = message.payload.get("query", "")
        k = message.payload.get("k", 3)
        
        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])
        
        # 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
            )
            
            # Send streaming response
            response = MCPMessage(
                sender=self.name,
                receiver="CoordinatorAgent",
                msg_type="LLM_RESPONSE_STREAM",
                trace_id=message.trace_id,
                payload={
                    "query": query,
                    "response_stream": response_stream,
                    "context": context
                }
            )
            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
            
            response = MCPMessage(
                sender=self.name,
                receiver="CoordinatorAgent",
                msg_type="LLM_RESPONSE_STREAM",
                trace_id=message.trace_id,
                payload={"response_stream": error_generator()}
            )
            self.message_bus.publish(response)

class CoordinatorAgent:
    """Coordinator agent that orchestrates the entire workflow"""
    
    def __init__(self, message_bus: MessageBus):
        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
    
    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")
    
    def process_files(self, files):
        """Process uploaded files"""
        if not files:
            return "No files uploaded."
        
        file_paths = [file.name for file in files]
        
        # Send ingestion request
        message = MCPMessage(
            sender=self.name,
            receiver="IngestionAgent",
            msg_type="INGESTION_REQUEST",
            payload={"files": file_paths}
        )
        self.message_bus.publish(message)
        
        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"""
        if not self.vector_store_ready:
            yield "Please upload and process documents first."
            return
        
        # Send retrieval request
        message = MCPMessage(
            sender=self.name,
            receiver="RetrievalAgent",
            msg_type="RETRIEVAL_REQUEST",
            payload={"query": query}
        )
        self.message_bus.publish(message)
        
        # Wait for response and stream
        import time
        timeout = 20  # seconds
        start_time = time.time()
        
        while not self.current_response_stream 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
                    if hasattr(chunk, 'choices') and chunk.choices:
                        token = chunk.choices[0].delta.content
                        if token:
                            yield token
                    else:
                        # Fallback for different response format
                        if hasattr(chunk, 'token'):
                            yield chunk.token
                        elif isinstance(chunk, str):
                            yield chunk
            except Exception as e:
                yield f"Error streaming response: {e}"
            finally:
                self.current_response_stream = None # Reset for next query
        else:
            yield "Timeout: No response received from LLM agent."

# 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 Claude-style Gradio interface with integrated file upload"""
    
    with gr.Blocks(
        theme=gr.themes.Base(),
        css="""
        /* Claude-inspired dark theme */
        .gradio-container {
            background-color: #0f0f0f !important;
            color: #ffffff !important;
            height: 100vh !important;
            max-width: none !important;
            padding: 0 !important;
            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', 'Roboto', sans-serif !important;
        }
        
        /* Main layout */
        .main-container {
            display: flex;
            flex-direction: column;
            height: 100vh;
            background: #0f0f0f;
            max-width: 800px;
            margin: 0 auto;
        }
        
        /* Header - minimal like Claude */
        .header {
            background: #0f0f0f;
            border-bottom: 1px solid #2a2a2a;
            padding: 1rem 2rem;
            text-align: center;
        }
        
        .header h1 {
            color: #ffffff;
            margin: 0;
            font-size: 1.25rem;
            font-weight: 500;
        }
        
        /* Chat container - full height like Claude */
        .chat-container {
            flex: 1;
            display: flex;
            flex-direction: column;
            padding: 0;
            background: #0f0f0f;
        }
        
        /* Chatbot area - Claude style */
        .gradio-chatbot {
            flex: 1 !important;
            height: calc(100vh - 200px) !important;
            max-height: none !important;
            background: #0f0f0f !important;
            border: none !important;
            padding: 1rem !important;
            overflow-y: auto !important;
        }
        
        /* Messages styling - Claude-like bubbles */
        .message.user {
            background: #2a2a2a !important;
            border-radius: 16px !important;
            padding: 12px 16px !important;
            margin: 8px 0 !important;
            color: #ffffff !important;
        }
        
        .message.bot {
            background: transparent !important;
            border-radius: 16px !important;
            padding: 12px 16px !important;
            margin: 8px 0 !important;
            color: #ffffff !important;
            border-left: 3px solid #ff6b35 !important;
            padding-left: 16px !important;
        }
        
        /* Input area - fixed at bottom like Claude */
        .input-container {
            position: sticky;
            bottom: 0;
            background: #0f0f0f;
            border-top: 1px solid #2a2a2a;
            padding: 1rem 2rem 2rem 2rem;
        }
        
        .input-wrapper {
            background: #2a2a2a;
            border-radius: 24px;
            border: 1px solid #404040;
            padding: 4px;
            display: flex;
            align-items: end;
            gap: 8px;
            transition: border-color 0.2s;
        }
        
        .input-wrapper:focus-within {
            border-color: #ff6b35 !important;
        }
        
        /* File upload button - integrated like Claude */
        .file-btn {
            background: transparent !important;
            border: none !important;
            color: #999999 !important;
            padding: 8px !important;
            border-radius: 20px !important;
            cursor: pointer !important;
            display: flex !important;
            align-items: center !important;
            justify-content: center !important;
            min-width: 36px !important;
            height: 36px !important;
            transition: all 0.2s !important;
        }
        
        .file-btn:hover {
            background: #404040 !important;
            color: #ffffff !important;
        }
        
        /* Text input - Claude style */
        .gradio-textbox {
            flex: 1 !important;
        }
        
        .gradio-textbox textarea, .gradio-textbox input {
            background: transparent !important;
            border: none !important;
            color: #ffffff !important;
            resize: none !important;
            padding: 12px 16px !important;
            font-size: 16px !important;
            line-height: 1.5 !important;
            max-height: 200px !important;
            min-height: 24px !important;
        }
        
        .gradio-textbox textarea:focus, .gradio-textbox input:focus {
            outline: none !important;
            box-shadow: none !important;
        }
        
        /* Send button - Claude style */
        .send-btn {
            background: #ff6b35 !important;
            border: none !important;
            border-radius: 20px !important;
            color: #ffffff !important;
            padding: 8px !important;
            min-width: 36px !important;
            height: 36px !important;
            display: flex !important;
            align-items: center !important;
            justify-content: center !important;
            cursor: pointer !important;
            transition: all 0.2s !important;
        }
        
        .send-btn:hover {
            background: #e55a2b !important;
        }
        
        .send-btn:disabled {
            background: #404040 !important;
            cursor: not-allowed !important;
        }
        
        /* File upload area - hidden by default */
        .file-upload {
            display: none !important;
        }
        
        /* Status messages */
        .status-message {
            background: #1a1a1a;
            border: 1px solid #404040;
            border-radius: 12px;
            padding: 12px 16px;
            margin: 8px 0;
            color: #cccccc;
            font-size: 14px;
            text-align: center;
        }
        
        .status-success {
            border-color: #22c55e !important;
            color: #22c55e !important;
            background: rgba(34, 197, 94, 0.1) !important;
        }
        
        .status-error {
            border-color: #ef4444 !important;
            color: #ef4444 !important;
            background: rgba(239, 68, 68, 0.1) !important;
        }
        
        .status-processing {
            border-color: #ff6b35 !important;
            color: #ff6b35 !important;
            background: rgba(255, 107, 53, 0.1) !important;
        }
        
        /* File list styling */
        .file-list {
            background: #1a1a1a;
            border-radius: 8px;
            padding: 8px 12px;
            margin: 4px 0;
            font-size: 14px;
            color: #cccccc;
            border-left: 3px solid #ff6b35;
        }
        
        /* Hide Gradio elements we don't want */
        .gradio-file {
            display: none !important;
        }
        
        /* Scrollbar styling */
        ::-webkit-scrollbar {
            width: 6px;
        }
        
        ::-webkit-scrollbar-track {
            background: transparent;
        }
        
        ::-webkit-scrollbar-thumb {
            background: #404040;
            border-radius: 3px;
        }
        
        ::-webkit-scrollbar-thumb:hover {
            background: #555555;
        }
        """,
        title="Agentic RAG Assistant"
    ) as iface:

        # Header - minimal like Claude
        gr.HTML("""
        <div class="header">
            <h1>Agentic RAG Assistant</h1>
        </div>
        """)

        # Main chat container
        with gr.Row():
            with gr.Column(scale=1, elem_classes=["main-container"]):
                
                # Chat area
                chatbot = gr.Chatbot(
                    elem_classes=["gradio-chatbot"],
                    show_copy_button=True,
                    type="messages",
                    placeholder="Hello! Upload documents using the 📎 button below, then ask me anything about them.",
                    avatar_images=(None, None),
                    bubble_full_width=False
                )
                
                # Input area - Claude style with integrated file upload
                with gr.Row(elem_classes=["input-container"]):
                    with gr.Column():
                        with gr.Row(elem_classes=["input-wrapper"]):
                            
                            # Hidden file upload
                            file_upload = gr.File(
                                file_count="multiple",
                                file_types=[".pdf", ".pptx", ".csv", ".docx", ".txt", ".md"],
                                elem_classes=["file-upload"],
                                visible=False
                            )
                            
                            # File upload button (📎 icon)
                            file_btn = gr.Button(
                                "📎",
                                elem_classes=["file-btn"],
                                size="sm",
                                scale=0
                            )
                            
                            # Message input
                            msg_input = gr.Textbox(
                                placeholder="Message Agentic RAG Assistant...",
                                show_label=False,
                                elem_classes=["message-input"],
                                scale=1,
                                container=False,
                                autofocus=True,
                                lines=1,
                                max_lines=10
                            )
                            
                            # Send button
                            send_btn = gr.Button(
                                "↑",
                                elem_classes=["send-btn"],
                                size="sm",
                                scale=0
                            )

        # State management
        doc_processed = gr.State(False)
        uploaded_files = gr.State([])
        
        # JavaScript for file upload integration
        gr.HTML("""
        <script>
        document.addEventListener('DOMContentLoaded', function() {
            // Make file button trigger file upload
            const fileBtn = document.querySelector('.file-btn');
            const fileInput = document.querySelector('.file-upload input[type="file"]');
            
            if (fileBtn && fileInput) {
                fileBtn.addEventListener('click', function(e) {
                    e.preventDefault();
                    e.stopPropagation();
                    fileInput.click();
                });
            }
            
            // Auto-resize textarea
            const textarea = document.querySelector('.message-input textarea');
            if (textarea) {
                textarea.addEventListener('input', function() {
                    this.style.height = 'auto';
                    this.style.height = Math.min(this.scrollHeight, 200) + 'px';
                });
            }
        });
        </script>
        """)
        
        # Event handlers
        def handle_file_upload(files, history):
            """Handle file uploads and show in chat"""
            if not files:
                return history, [], False, gr.update()
            
            # Add file upload message to chat
            file_names = [f.name.split('/')[-1] if hasattr(f, 'name') else str(f) for f in files]
            file_list = "\n".join([f"📄 {name}" for name in file_names])
            
            history.append({
                "role": "user", 
                "content": f"📎 Uploaded {len(files)} file(s):\n{file_list}"
            })
            
            # Show processing message
            history.append({
                "role": "assistant", 
                "content": "📄 Processing documents... This may take a moment."
            })
            
            try:
                # Process files
                result = coordinator_agent.process_files(files)
                
                # Wait for processing
                import time
                time.sleep(3)
                
                # Update last message with success
                history[-1]["content"] = f"✅ Successfully processed {len(files)} document(s)! You can now ask questions about your documents."
                
                return history, files, True, gr.update(value=None)
                
            except Exception as e:
                history[-1]["content"] = f"❌ Error processing documents: {str(e)}"
                return history, [], False, gr.update(value=None)

        def respond(message, history, doc_ready, files):
            """Handle user messages"""
            if not message.strip():
                return history, ""
            
            # Add user message
            history.append({"role": "user", "content": message})
            
            if not doc_ready and not any(cmd in message.lower() for cmd in ['hello', 'hi', 'help']):
                history.append({
                    "role": "assistant", 
                    "content": "👋 Hello! I'm your RAG assistant. Please upload some documents first using the 📎 button, then I can help you analyze and answer questions about them."
                })
                return history, ""
            
            # Handle general queries without documents
            if not doc_ready:
                if any(cmd in message.lower() for cmd in ['hello', 'hi', 'help']):
                    history.append({
                        "role": "assistant", 
                        "content": "👋 Hello! I'm an AI assistant specialized in document analysis. Upload documents using the 📎 button above, and I'll help you:\n\n• Summarize content\n• Answer questions\n• Extract key insights\n• Find specific information\n\nSupported formats: PDF, DOCX, PPTX, CSV, TXT, MD"
                    })
                else:
                    history.append({
                        "role": "assistant", 
                        "content": "Please upload documents first using the 📎 button above, then I can help answer questions about them."
                    })
                return history, ""
            
            # Add assistant message placeholder
            history.append({"role": "assistant", "content": ""})
            
            # Stream response
            try:
                for token in coordinator_agent.handle_query(message, history):
                    if token:
                        history[-1]["content"] += token
                        yield history, ""
            except Exception as e:
                history[-1]["content"] = f"❌ Sorry, I encountered an error: {str(e)}"
                yield history, ""
        
        # Event bindings
        
        # File upload handling
        file_upload.upload(
            handle_file_upload,
            inputs=[file_upload, chatbot],
            outputs=[chatbot, uploaded_files, doc_processed, file_upload]
        )
        
        # File button click (handled by JavaScript)
        file_btn.click(None, None, None, js="""
        function() {
            document.querySelector('.file-upload input[type="file"]').click();
        }
        """)
        
        # Send button click
        send_btn.click(
            respond,
            inputs=[msg_input, chatbot, doc_processed, uploaded_files],
            outputs=[chatbot, msg_input]
        )
        
        # Enter key submit
        msg_input.submit(
            respond,
            inputs=[msg_input, chatbot, doc_processed, uploaded_files],
            outputs=[chatbot, msg_input]
        )

    return iface
# Launch the application
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860
    )