Create app.py

app.py

@@ -0,0 +1,402 @@
+import gradio as gr
+import spaces
+import os
+import logging
+import datetime
+from langchain_community.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import Chroma
+from huggingface_hub import InferenceClient, get_token
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Set HF_HOME for caching Hugging Face assets in persistent storage
+os.environ["HF_HOME"] = "/data/.huggingface"
+os.makedirs(os.environ["HF_HOME"], exist_ok=True)
+
+# Define persistent storage directories
+DATA_DIR = "/data"  # Root persistent storage directory
+DOCS_DIR = os.path.join(DATA_DIR, "documents")  # Subdirectory for uploaded PDFs
+CHROMA_DIR = os.path.join(DATA_DIR, "chroma_db")  # Subdirectory for Chroma vector store
+
+# Create directories if they don't exist
+os.makedirs(DOCS_DIR, exist_ok=True)
+os.makedirs(CHROMA_DIR, exist_ok=True)
+
+# Initialize Cerebras InferenceClient
+try:
+    token = get_token()
+    if not token:
+        logger.error("HF_TOKEN is not set in Space secrets")
+        client = None
+    else:
+        client = InferenceClient(
+            model="meta-llama/Llama-4-Scout-17B-16E-Instruct",
+            provider="cerebras",
+            token=token
+        )
+        logger.info("InferenceClient initialized successfully")
+except Exception as e:
+    logger.error(f"Failed to initialize InferenceClient: {str(e)}")
+    client = None
+
+# Global variables for vector store
+vectorstore = None
+retriever = None
+
+def list_uploaded_documents():
+    """List all uploaded documents in the persistent storage"""
+    try:
+        if not os.path.exists(DOCS_DIR):
+            return []
+        
+        files = os.listdir(DOCS_DIR)
+        pdf_files = [f for f in files if f.lower().endswith('.pdf')]
+        file_info = []
+        
+        for file in pdf_files:
+            file_path = os.path.join(DOCS_DIR, file)
+            file_size = os.path.getsize(file_path)
+            file_time = os.path.getmtime(file_path)
+            file_info.append({
+                "name": file,
+                "size": f"{file_size // 1024} KB",
+                "date": datetime.datetime.fromtimestamp(file_time).strftime('%Y-%m-%d %H:%M:%S')
+            })
+        
+        return file_info
+    except Exception as e:
+        logger.error(f"Error listing documents: {str(e)}")
+        return []
+
+def delete_document(filename):
+    """Delete a document from persistent storage and update the vector store"""
+    try:
+        if not filename:
+            return "No file selected for deletion"
+        
+        file_path = os.path.join(DOCS_DIR, filename)
+        if not os.path.exists(file_path):
+            return f"File {filename} does not exist"
+        
+        # Delete the file
+        os.remove(file_path)
+        logger.info(f"Deleted file {filename}")
+        
+        # Refresh the vector store
+        refresh_status = refresh_vector_store()
+        
+        return f"File {filename} deleted successfully! {refresh_status}"
+    except Exception as e:
+        logger.error(f"Error deleting document: {str(e)}")
+        return f"Error deleting document: {str(e)}"
+
+def preview_document(filename, max_pages=3):
+    """Generate a preview of the document's content"""
+    try:
+        if not filename:
+            return "No file selected for preview"
+        
+        file_path = os.path.join(DOCS_DIR, filename)
+        if not os.path.exists(file_path):
+            return f"File {filename} does not exist"
+        
+        loader = PyPDFLoader(file_path)
+        documents = loader.load()
+        
+        # Limit preview to first few pages
+        preview_docs = documents[:max_pages]
+        preview_text = f"Preview of {filename} (first {len(preview_docs)} pages):\n\n"
+        
+        for i, doc in enumerate(preview_docs):
+            preview_text += f"--- Page {i+1} ---\n{doc.page_content[:500]}...\n\n"
+        
+        return preview_text
+    except Exception as e:
+        logger.error(f"Error previewing document: {str(e)}")
+        return f"Error previewing document: {str(e)}"
+
+@spaces.GPU(duration=180)  # Use GPU for vector store recreation
+def refresh_vector_store():
+    """Rebuild the vector store from all available documents"""
+    global vectorstore, retriever
+    try:
+        if not os.path.exists(DOCS_DIR):
+            logger.warning("Documents directory does not exist")
+            return "No documents directory found"
+        
+        files = [f for f in os.listdir(DOCS_DIR) if f.lower().endswith('.pdf')]
+        if not files:
+            logger.warning("No PDF documents found")
+            
+            # Clear the vector store
+            if os.path.exists(CHROMA_DIR):
+                import shutil
+                shutil.rmtree(CHROMA_DIR)
+                os.makedirs(CHROMA_DIR, exist_ok=True)
+            
+            vectorstore = None
+            retriever = None
+            return "No PDF documents found. Vector store cleared."
+        
+        # Load and process all documents
+        all_texts = []
+        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        
+        for file in files:
+            file_path = os.path.join(DOCS_DIR, file)
+            try:
+                loader = PyPDFLoader(file_path)
+                documents = loader.load()
+                texts = text_splitter.split_documents(documents)
+                
+                # Add source file metadata to each chunk
+                for i, text in enumerate(texts):
+                    text.metadata["source"] = file
+                    
+                all_texts.extend(texts)
+                logger.info(f"Processed {file}, added {len(texts)} chunks")
+            except Exception as e:
+                logger.error(f"Error processing {file}: {str(e)}")
+        
+        # Create new vector store
+        if all_texts:
+            # Remove existing vector store
+            if os.path.exists(CHROMA_DIR):
+                import shutil
+                shutil.rmtree(CHROMA_DIR)
+                os.makedirs(CHROMA_DIR, exist_ok=True)
+            
+            vectorstore = Chroma.from_documents(
+                all_texts, embeddings, persist_directory=CHROMA_DIR
+            )
+            vectorstore.persist()
+            retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+            logger.info(f"Vector store recreated with {len(all_texts)} chunks from {len(files)} files")
+            return f"Vector store updated with {len(files)} documents!"
+        else:
+            logger.warning("No text chunks extracted from documents")
+            return "No content could be extracted from the PDF files"
+    except Exception as e:
+        logger.error(f"Error refreshing vector store: {str(e)}")
+        return f"Error refreshing vector store: {str(e)}"
+
+@spaces.GPU(duration=180)  # Use ZeroGPU (H200) for embedding generation, 180s timeout
+def initialize_rag(file):
+    global vectorstore, retriever
+    try:
+        # Debug file object properties
+        logger.info(f"File object: {type(file)}, Attributes: {dir(file)}")
+        logger.info(f"File name: {file.name}")
+
+        # Validate file
+        if not file or not file.name:
+            logger.error("No file provided or invalid file name")
+            return "Error: No file provided or invalid file name"
+        
+        # Verify temporary file exists and is accessible
+        if not os.path.exists(file.name):
+            logger.error(f"Temporary file {file.name} does not exist")
+            return f"Error: Temporary file {file.name} does not exist"
+        
+        # Check temporary file size
+        file_size = os.path.getsize(file.name)
+        logger.info(f"Temporary file size: {file_size} bytes")
+        if file_size == 0:
+            logger.error("Uploaded file is empty")
+            return "Error: Uploaded file is empty"
+
+        # Save uploaded file to persistent storage
+        file_name = os.path.basename(file.name)
+        file_path = os.path.join(DOCS_DIR, file_name)
+        
+        # Check if file exists and its size
+        should_save = True
+        if os.path.exists(file_path):
+            existing_size = os.path.getsize(file_path)
+            logger.info(f"Existing file {file_name} size: {existing_size} bytes")
+            if existing_size == 0:
+                logger.warning(f"Existing file {file_name} is empty, will overwrite")
+            else:
+                logger.info(f"File {file_name} already exists and is not empty, skipping save")
+                should_save = False
+
+        if should_save:
+            try:
+                with open(file.name, "rb") as src_file:
+                    file_content = src_file.read()
+                logger.info(f"Read {len(file_content)} bytes from temporary file")
+                if not file_content:
+                    logger.error("File content is empty after reading")
+                    return "Error: File content is empty after reading"
+                with open(file_path, "wb") as dst_file:
+                    dst_file.write(file_content)
+                    dst_file.flush()  # Ensure write completes
+                # Verify written file
+                written_size = os.path.getsize(file_path)
+                logger.info(f"Saved {file_name} to {file_path}, size: {written_size} bytes")
+                if written_size == 0:
+                    logger.error(f"Failed to write {file_name}, file is empty")
+                    return f"Error: Failed to write {file_name}, file is empty"
+            except PermissionError as e:
+                logger.error(f"Permission error writing to {file_path}: {str(e)}")
+                return f"Error: Permission denied writing to {file_path}"
+            except Exception as e:
+                logger.error(f"Error writing file to {file_path}: {str(e)}")
+                return f"Error writing file: {str(e)}"
+
+        # After saving the file, refresh the vector store
+        refresh_status = refresh_vector_store()
+        logger.info(f"Vector store refresh status: {refresh_status}")
+        
+        return f"Document '{file_name}' processed and saved! {refresh_status}"
+    except Exception as e:
+        logger.error(f"Error processing document: {str(e)}")
+        return f"Error processing document: {str(e)}"
+
+def query_documents(query, history, system_prompt, max_tokens, temperature):
+    global retriever, client
+    try:
+        if client is None:
+            logger.error("InferenceClient not initialized")
+            return history, "Error: InferenceClient not initialized. Check HF_TOKEN."
+        if retriever is None:
+            logger.error("No documents loaded")
+            return history, "Error: No documents loaded. Please upload a document first."
+
+        # Ensure history is a list of [user, assistant] lists
+        logger.info(f"History before processing: {history}")
+        if not isinstance(history, list):
+            logger.warning("History is not a list, resetting")
+            history = []
+        history = [[str(item[0]), str(item[1])] for item in history if isinstance(item, (list, tuple)) and len(item) == 2]
+
+        # Retrieve relevant documents
+        docs = retriever.get_relevant_documents(query)
+        
+        # Format context with source information
+        context_parts = []
+        for doc in docs:
+            source = doc.metadata.get('source', 'unknown')
+            page = doc.metadata.get('page', 'unknown')
+            context_parts.append(f"[Source: {source}, Page: {page}]\n{doc.page_content}")
+        
+        context = "\n\n".join(context_parts)
+
+        # Call Cerebras inference
+        logger.info("Calling Cerebras inference")
+        response = client.chat_completion(
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": f"Context: {context}\n\nQuery: {query}"}
+            ],
+            max_tokens=int(max_tokens),
+            temperature=float(temperature),
+            stream=False
+        )
+        answer = response.choices[0].message.content
+        logger.info("Inference successful")
+
+        # Update chat history with list format
+        history.append([query, answer])
+        logger.info(f"History after append: {history}")
+        return history, ""  # Clear the query input
+    except Exception as e:
+        logger.error(f"Error querying documents: {str(e)}")
+        return history, f"Error querying documents: {str(e)}"
+
+# Load existing vector store on startup
+try:
+    if os.path.exists(CHROMA_DIR):
+        logger.info("Loading existing vector store")
+        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+        vectorstore = Chroma(persist_directory=CHROMA_DIR, embedding_function=embeddings)
+        retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+        logger.info(f"Loaded vector store from {CHROMA_DIR}")
+except Exception as e:
+    logger.error(f"Error loading vector store: {str(e)}")
+
+# Create the Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# RAG Chatbot with Document Management")
+    
+    # File management tab
+    with gr.Tab("Document Management"):
+        # File upload
+        with gr.Row():
+            file_input = gr.File(label="Upload Document (PDF)", file_types=[".pdf"])
+            file_output = gr.Textbox(label="Upload Status")
+        
+        # Document listing and management
+        with gr.Row():
+            refresh_btn = gr.Button("Refresh Document List")
+            rebuild_vs_btn = gr.Button("Rebuild Vector Store")
+        
+        doc_list = gr.Dataframe(
+            headers=["name", "size", "date"],
+            label="Uploaded Documents"
+        )
+        
+        with gr.Row():
+            selected_doc = gr.Dropdown(label="Select Document")
+            preview_btn = gr.Button("Preview Document")
+            delete_btn = gr.Button("Delete Selected Document", variant="stop")
+        
+        doc_preview = gr.Textbox(label="Document Preview", lines=10)
+        delete_output = gr.Textbox(label="Operation Status")
+    
+    # Chat interface tab
+    with gr.Tab("Chat"):
+        chatbot = gr.Chatbot(label="Conversation")
+        
+        # Query and parameters
+        with gr.Row():
+            query_input = gr.Textbox(label="Query", placeholder="Ask about the document...")
+            system_prompt = gr.Textbox(
+                label="System Prompt",
+                value="You are a helpful assistant answering questions based on the provided document context. Only use the context provided to answer the question. If you don't know the answer, say so."
+            )
+        
+        with gr.Row():
+            max_tokens = gr.Slider(label="Max Tokens", minimum=50, maximum=2000, value=500, step=50)
+            temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, value=0.7, step=0.1)
+        
+        # Buttons
+        with gr.Row():
+            submit_btn = gr.Button("Send")
+            clear_btn = gr.Button("Clear Chat")
+    
+    # Event handlers for file management
+    def update_doc_list():
+        docs = list_uploaded_documents()
+        return docs, [d["name"] for d in docs]
+    
+    file_input.upload(initialize_rag, file_input, file_output).then(
+        update_doc_list, None, [doc_list, selected_doc]
+    )
+    
+    refresh_btn.click(update_doc_list, None, [doc_list, selected_doc])
+    rebuild_vs_btn.click(refresh_vector_store, None, delete_output)
+    preview_btn.click(preview_document, selected_doc, doc_preview)
+    delete_btn.click(delete_document, selected_doc, delete_output).then(
+        update_doc_list, None, [doc_list, selected_doc]
+    )
+    
+    # Event handlers for chat
+    submit_btn.click(
+        query_documents,
+        inputs=[query_input, chatbot, system_prompt, max_tokens, temperature],
+        outputs=[chatbot, query_input]
+    )
+    
+    clear_btn.click(lambda: [], None, chatbot)
+    
+    # Initialize document list on startup
+    demo.load(update_doc_list, None, [doc_list, selected_doc])
+
+if __name__ == "__main__":
+    demo.launch()