Add RAG capability with document upload and management

app.py

@@ -15,8 +15,13 @@ from modules.citation import generate_citations, format_citations
 from modules.server_cache import get_cached_result, cache_result
 from modules.status_logger import log_request
 from modules.server_monitor import ServerMonitor
+from modules.rag.rag_chain import RAGChain
+from modules.rag.vector_store import VectorStore
+from langchain.docstore.document import Document
 
 server_monitor = ServerMonitor()
+rag_chain = RAGChain()
+vector_store = VectorStore()
 
 # Cat-themed greeting function
 def get_cat_greeting():
@@ -123,8 +128,8 @@ def run_startup_check():
     return wrapper
 
 # Enhanced streaming with markdown support
-async def research_assistant(query, history):
-    log_request("Research started", query=query)
+async def research_assistant(query, history, use_rag=False):
+    log_request("Research started", query=query, use_rag=use_rag)
 
     # Add typing indicator
     history.append((query, "🔄 Searching for information..."))
@@ -188,9 +193,19 @@ async def research_assistant(query, history):
     if any(keyword in lower_query for keyword in space_keywords):
         context_section += f"\nSpace Weather Context: {space_weather_data}"
 
-    # Build the enriched input with context only when needed
+    # Build the enriched input
     enriched_input = f"{validated_query}\n\n{answer_content}Search Results:\n{search_content}{context_section}"
 
+    # If RAG is enabled, use it
+    if use_rag:
+        history[-1] = (query, "📚 Searching document database...")
+        yield history
+        
+        rag_result = rag_chain.query(validated_query)
+        if rag_result["status"] == "success":
+            enriched_input = rag_result["prompt"]
+            context_section += f"\n\nDocument Context:\n" + "\n\n".join([doc.page_content for doc in rag_result["context_docs"][:2]])
+
     server_status = server_monitor.check_server_status()
     if not server_status["available"]:
         wait_time = server_status["estimated_wait"]
@@ -298,11 +313,39 @@ class AsyncGeneratorWrapper:
             raise StopIteration
         return item
 
-def research_assistant_wrapper(query, history):
-    async_gen = research_assistant(query, history)
+def research_assistant_wrapper(query, history, use_rag):
+    async_gen = research_assistant(query, history, use_rag)
     wrapper = AsyncGeneratorWrapper(async_gen)
     return wrapper
 
+# Document upload function
+def upload_documents(files):
+    """Upload and process documents for RAG"""
+    try:
+        documents = []
+        for file in files:
+            # For PDF files
+            if file.name.endswith('.pdf'):
+                from PyPDF2 import PdfReader
+                reader = PdfReader(file.name)
+                text = ""
+                for page in reader.pages:
+                    text += page.extract_text()
+                documents.append(Document(page_content=text, metadata={"source": file.name}))
+            # For text files
+            else:
+                with open(file.name, 'r') as f:
+                    text = f.read()
+                documents.append(Document(page_content=text, metadata={"source": file.name}))
+        
+        result = vector_store.add_documents(documents)
+        if result["status"] == "success":
+            return f"✅ Successfully added {result['count']} document chunks to the knowledge base!"
+        else:
+            return f"❌ Error adding documents: {result['message']}"
+    except Exception as e:
+        return f"❌ Error processing documents: {str(e)}"
+
 # Performance dashboard data
 def get_performance_stats():
     """Get performance statistics from Redis"""
@@ -344,14 +387,16 @@ with gr.Blocks(
                     gr.Markdown("## How to Use")
                     gr.Markdown("""
                     1. Enter a research question in the input box
-                    2. Click Submit or press Enter
-                    3. Watch as the response streams in real-time
-                    4. Review sources at the end of each response
+                    2. Toggle 'Use Document Knowledge' to enable RAG
+                    3. Click Submit or press Enter
+                    4. Watch as the response streams in real-time
+                    5. Review sources at the end of each response
                     
                     ## Features
                     - 🔍 Web search integration
                     - 🌤️ Context-aware weather data (only when relevant)
                     - 🌌 Context-aware space weather data (only when relevant)
+                    - 📚 RAG (Retrieval-Augmented Generation) with document database
                     - 📚 Real-time citations
                     - ⚡ Streaming output
                     """)
@@ -368,6 +413,11 @@ with gr.Blocks(
                         placeholder="Ask a complex research question...",
                         lines=3
                     )
+                    use_rag = gr.Checkbox(
+                        label="📚 Use Document Knowledge (RAG)", 
+                        value=False,
+                        info="Enable to search uploaded documents for context"
+                    )
                     with gr.Row():
                         submit_btn = gr.Button("Submit Research Query", variant="primary")
                         clear_btn = gr.Button("Clear Conversation")
@@ -384,6 +434,25 @@ with gr.Blocks(
                         label="Example Questions"
                     )
         
+        with gr.TabItem("📚 Document Management"):
+            gr.Markdown("## Upload Documents for RAG")
+            gr.Markdown("Upload PDF or text files to add them to the knowledge base for document-based queries.")
+            file_upload = gr.File(
+                file_types=[".pdf", ".txt"], 
+                file_count="multiple",
+                label="Upload Documents"
+            )
+            upload_btn = gr.Button("📤 Upload Documents")
+            upload_output = gr.Textbox(label="Upload Status", interactive=False)
+            clear_docs_btn = gr.Button("🗑️ Clear All Documents")
+            
+            gr.Markdown("## Current Documents")
+            doc_list = gr.Textbox(
+                label="Document List", 
+                value="No documents uploaded yet",
+                interactive=False
+            )
+        
         with gr.TabItem("📊 Performance"):
             perf_refresh_btn = gr.Button("🔄 Refresh Stats")
             perf_display = gr.JSON(label="System Statistics")
@@ -432,9 +501,9 @@ While you wait, why not prepare some treats? I'll be ready to hunt for knowledge
         startup_check_result = run_startup_check()
         return update_status()
     
-    def respond(message, history):
+    def respond(message, history, use_rag_flag):
         # Get streaming response
-        for updated_history in research_assistant_wrapper(message, history):
+        for updated_history in research_assistant_wrapper(message, history, use_rag_flag):
             yield updated_history, update_status()
     
     def clear_conversation():
@@ -452,17 +521,21 @@ While you wait, why not prepare some treats? I'll be ready to hunt for knowledge
     check_btn.click(refresh_status, outputs=status_display)
     submit_btn.click(
         respond, 
-        [msg, chat_history], 
+        [msg, chat_history, use_rag], 
         [chatbot, status_display]
     )
     msg.submit(
         respond, 
-        [msg, chat_history], 
+        [msg, chat_history, use_rag], 
         [chatbot, status_display]
     )
     
     clear_btn.click(clear_conversation, outputs=[chat_history, chatbot])
     
+    # Document management
+    upload_btn.click(upload_documents, file_upload, upload_output)
+    clear_docs_btn.click(lambda: vector_store.delete_collection(), None, upload_output)
+    
     # Performance dashboard
     perf_refresh_btn.click(update_performance_stats, outputs=perf_display)
 

modules/rag/rag_chain.py

@@ -0,0 +1,55 @@
+from langchain.chains import RetrievalQA
+from langchain.llms import OpenAI
+from langchain.prompts import PromptTemplate
+from modules.rag.vector_store import VectorStore
+from modules.analyzer import client
+import os
+
+class RAGChain:
+    def __init__(self):
+        self.vector_store = VectorStore()
+        self.retriever = self.vector_store.vector_store.as_retriever(
+            search_type="similarity",
+            search_kwargs={"k": 5}
+        )
+        
+        # Custom prompt template
+        self.prompt_template = """
+        You are an AI research assistant with access to a document database. 
+        Use the following pieces of context to answer the question at the end. 
+        If you don't know the answer, just say that you don't know, don't try to make up an answer.
+        
+        Context: {context}
+        
+        Question: {question}
+        
+        Answer:
+        """
+        
+        self.prompt = PromptTemplate(
+            template=self.prompt_template,
+            input_variables=["context", "question"]
+        )
+    
+    def query(self, question):
+        """Query the RAG system"""
+        try:
+            # Search for relevant documents
+            search_result = self.vector_store.search(question)
+            if search_result["status"] != "success":
+                return {"status": "error", "message": search_result["message"]}
+            
+            # Format context
+            context = "\n\n".join([doc.page_content for doc in search_result["documents"]])
+            
+            # Create enhanced prompt
+            enhanced_prompt = self.prompt.format(context=context, question=question)
+            
+            # For streaming, we'll return the prompt for the analyzer to handle
+            return {
+                "status": "success", 
+                "prompt": enhanced_prompt,
+                "context_docs": search_result["documents"]
+            }
+        except Exception as e:
+            return {"status": "error", "message": str(e)}

modules/rag/vector_store.py

@@ -0,0 +1,82 @@
+import os
+import chromadb
+from chromadb.utils import embedding_functions
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores import Chroma
+from langchain.embeddings import SentenceTransformerEmbeddings
+from langchain.docstore.document import Document
+import uuid
+
+class VectorStore:
+    def __init__(self):
+        # Initialize embedding function
+        self.embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
+        
+        # Initialize ChromaDB client
+        self.client = chromadb.PersistentClient(path="./chroma_db")
+        
+        # Create or get collection
+        self.collection = self.client.get_or_create_collection(
+            name="research_documents",
+            embedding_function=embedding_functions.SentenceTransformerEmbeddingFunction(
+                model_name="all-MiniLM-L6-v2"
+            )
+        )
+        
+        # Initialize LangChain vector store
+        self.vector_store = Chroma(
+            collection_name="research_documents",
+            embedding_function=self.embedding_function,
+            persist_directory="./chroma_db"
+        )
+        
+        # Initialize text splitter
+        self.text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=1000,
+            chunk_overlap=200,
+            length_function=len,
+        )
+    
+    def add_documents(self, documents):
+        """Add documents to the vector store"""
+        try:
+            # Split documents into chunks
+            split_docs = []
+            for doc in documents:
+                splits = self.text_splitter.split_text(doc.page_content)
+                for i, split in enumerate(splits):
+                    split_docs.append(Document(
+                        page_content=split,
+                        metadata={**doc.metadata, "chunk": i}
+                    ))
+            
+            # Add to vector store
+            ids = [str(uuid.uuid4()) for _ in split_docs]
+            self.vector_store.add_documents(split_docs, ids=ids)
+            
+            return {"status": "success", "count": len(split_docs)}
+        except Exception as e:
+            return {"status": "error", "message": str(e)}
+    
+    def search(self, query, k=5):
+        """Search for relevant documents"""
+        try:
+            # Perform similarity search
+            docs = self.vector_store.similarity_search(query, k=k)
+            return {"status": "success", "documents": docs}
+        except Exception as e:
+            return {"status": "error", "message": str(e)}
+    
+    def delete_collection(self):
+        """Delete the entire collection"""
+        try:
+            self.client.delete_collection("research_documents")
+            self.collection = self.client.get_or_create_collection(
+                name="research_documents",
+                embedding_function=embedding_functions.SentenceTransformerEmbeddingFunction(
+                    model_name="all-MiniLM-L6-v2"
+                )
+            )
+            return {"status": "success"}
+        except Exception as e:
+            return {"status": "error", "message": str(e)}

requirements.txt

@@ -5,3 +5,10 @@ redis
 aiohttp
 requests
 python-dotenv
+langchain
+langchain-community
+langchain-openai
+chromadb
+sentence-transformers
+pypdf
+python-multipart