Delete app.py

app.py

@@ -1,410 +0,0 @@
-import gradio as gr
-import spaces
-import os
-import logging
-import datetime
-from langchain.document_loaders import PyPDFLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.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
-
-@spaces.CPU  # Use CPU for file operations
-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 []
-
-@spaces.CPU  # Use CPU for file operations
-def delete_document(filename):
-    """Delete a document from persistent storage and update the vector store"""
-    global vectorstore, retriever
-    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}")
-        
-        # We should also update the vector store to remove documents from this file
-        # This requires re-creating the vector store from all remaining documents
-        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)}"
-
-@spaces.CPU  # Use CPU for file operations
-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
-                    text.metadata["chunk_id"] = i
-                    text.metadata["page"] = text.metadata.get("page", 0)
-                    
-                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 and {len(all_texts)} chunks!"
-        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 GPU for embedding generation
-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 successfully, 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",
-            interactive=False
-        )
-        
-        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()