Enhanced app.py with RAG functionality: file upload, chunking with unstructured, embedding with sentence-transformers, vector storage with ChromaDB, and semantic retrieval QA

app.py

@@ -1,11 +1,109 @@
 import gradio as gr
 import random
+import os
+import tempfile
+from pathlib import Path
 from smolagents import GradioUI, CodeAgent, HfApiModel
+from sentence_transformers import SentenceTransformer
+import chromadb
+from unstructured.partition.auto import partition
+import numpy as np
+from typing import List, Optional
 
 # Import our custom tools from their modules
 from tools import DuckDuckGoSearchTool, WeatherInfoTool, HubStatsTool
 from retriever import load_guest_dataset
 
+# Initialize embedding model
+embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+
+# Initialize ChromaDB
+chroma_client = chromadb.Client()
+collection = chroma_client.get_or_create_collection(name="documents")
+
+class RAGDocumentProcessor:
+    def __init__(self, embedding_model, collection):
+        self.embedding_model = embedding_model
+        self.collection = collection
+    
+    def process_document(self, file_path: str) -> List[str]:
+        """Process document using unstructured for chunking"""
+        elements = partition(filename=file_path)
+        chunks = []
+        
+        # Group elements into meaningful chunks
+        current_chunk = ""
+        for element in elements:
+            text = str(element)
+            if len(current_chunk + text) > 1000:  # Max chunk size
+                if current_chunk:
+                    chunks.append(current_chunk.strip())
+                current_chunk = text
+            else:
+                current_chunk += " " + text
+        
+        if current_chunk:
+            chunks.append(current_chunk.strip())
+        
+        return chunks
+    
+    def add_document_to_vector_store(self, file_path: str, filename: str):
+        """Add document chunks to ChromaDB vector store"""
+        chunks = self.process_document(file_path)
+        
+        # Generate embeddings
+        embeddings = self.embedding_model.encode(chunks).tolist()
+        
+        # Create IDs and metadata
+        ids = [f"{filename}_{i}" for i in range(len(chunks))]
+        metadatas = [{"filename": filename, "chunk_id": i} for i in range(len(chunks))]
+        
+        # Add to collection
+        self.collection.add(
+            embeddings=embeddings,
+            documents=chunks,
+            metadatas=metadatas,
+            ids=ids
+        )
+        
+        return len(chunks)
+    
+    def semantic_search(self, query: str, n_results: int = 5) -> List[str]:
+        """Perform semantic search using ChromaDB"""
+        query_embedding = self.embedding_model.encode([query]).tolist()
+        
+        results = self.collection.query(
+            query_embeddings=query_embedding,
+            n_results=n_results
+        )
+        
+        return results['documents'][0] if results['documents'] else []
+
+# Initialize RAG processor
+rag_processor = RAGDocumentProcessor(embedding_model, collection)
+
+class SemanticRAGTool:
+    """Tool for semantic retrieval and QA using uploaded documents"""
+    
+    name = "semantic_rag_search"
+    description = "Search through uploaded documents using semantic similarity and provide context-aware responses"
+    
+    def __call__(self, query: str) -> str:
+        """Perform semantic search and return relevant context"""
+        relevant_docs = rag_processor.semantic_search(query, n_results=3)
+        
+        if not relevant_docs:
+            return "No relevant documents found. Please upload documents first."
+        
+        context = "\n\n".join(relevant_docs)
+        
+        response = f"Based on the uploaded documents, here's the relevant information:\n\n{context}\n\nThis information can help answer your query: {query}"
+        
+        return response
+
+# Initialize the semantic RAG tool
+semantic_rag_tool = SemanticRAGTool()
+
 # Initialize the Hugging Face model
 model = HfApiModel()
 
@@ -21,13 +119,102 @@ hub_stats_tool = HubStatsTool()
 # Load the guest dataset and initialize the guest info tool
 guest_info_tool = load_guest_dataset()
 
-# Create Alfred with all the tools
+def upload_and_process_file(file):
+    """Handle file upload and processing"""
+    if file is None:
+        return "No file uploaded."
+    
+    try:
+        # Get the file path
+        file_path = file.name
+        filename = Path(file_path).name
+        
+        # Process and add to vector store
+        num_chunks = rag_processor.add_document_to_vector_store(file_path, filename)
+        
+        return f"Successfully processed '{filename}' into {num_chunks} chunks and added to vector store."
+    
+    except Exception as e:
+        return f"Error processing file: {str(e)}"
+
+# Create Alfred with all the tools including the new RAG tool
 alfred = CodeAgent(
-    tools=[guest_info_tool, weather_info_tool, hub_stats_tool, search_tool], 
+    tools=[guest_info_tool, weather_info_tool, hub_stats_tool, search_tool, semantic_rag_tool], 
     model=model,
     add_base_tools=True,  # Add any additional base tools
     planning_interval=3   # Enable planning every 3 steps
 )
 
+# Create custom Gradio interface with file upload
+def create_rag_interface():
+    """Create enhanced Gradio interface with file upload and RAG capabilities"""
+    
+    with gr.Blocks(title="Production RAG Agent") as demo:
+        gr.Markdown("# Production RAG Agent with Document Upload")
+        gr.Markdown("Upload documents and ask questions using semantic search and AI reasoning.")
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                file_upload = gr.File(
+                    label="Upload Documents",
+                    file_types=[".pdf", ".docx", ".txt", ".md", ".html"],
+                    file_count="multiple"
+                )
+                upload_btn = gr.Button("Process Documents")
+                upload_status = gr.Textbox(
+                    label="Upload Status",
+                    interactive=False
+                )
+                
+            with gr.Column(scale=2):
+                # Embed the GradioUI from smolagents
+                chatbot = gr.Chatbot(label="AI Assistant")
+                msg_input = gr.Textbox(
+                    label="Message",
+                    placeholder="Ask questions about uploaded documents or anything else..."
+                )
+                send_btn = gr.Button("Send")
+                clear_btn = gr.Button("Clear")
+        
+        # File upload handler
+        upload_btn.click(
+            fn=lambda files: "\n".join([upload_and_process_file(file) for file in files]) if files else "No files selected.",
+            inputs=[file_upload],
+            outputs=[upload_status]
+        )
+        
+        # Chat functionality
+        def respond(message, history):
+            try:
+                # Use Alfred to generate response
+                response = alfred.run(message)
+                history.append((message, str(response)))
+                return history, ""
+            except Exception as e:
+                error_msg = f"Error: {str(e)}"
+                history.append((message, error_msg))
+                return history, ""
+        
+        send_btn.click(
+            respond,
+            inputs=[msg_input, chatbot],
+            outputs=[chatbot, msg_input]
+        )
+        
+        msg_input.submit(
+            respond,
+            inputs=[msg_input, chatbot],
+            outputs=[chatbot, msg_input]
+        )
+        
+        clear_btn.click(
+            lambda: ([], ""),
+            outputs=[chatbot, msg_input]
+        )
+    
+    return demo
+
 if __name__ == "__main__":
-    GradioUI(alfred).launch()
+    # Launch the enhanced RAG interface
+    demo = create_rag_interface()
+    demo.launch()