Delete app.py

app.py

@@ -1,467 +0,0 @@
-import os
-import sys
-import logging
-from pathlib import Path
-import json
-import hashlib
-from datetime import datetime
-import threading
-import queue
-from typing import List, Dict, Any, Tuple, Optional
-
-# Configure logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
-# Importing necessary libraries
-import torch
-import numpy as np
-from sentence_transformers import SentenceTransformer
-import chromadb
-from chromadb.utils import embedding_functions
-import gradio as gr
-from openai import OpenAI
-import google.generativeai as genai
-
-# Configuration class
-class Config:
-    """Configuration for vector store and RAG"""
-    def __init__(self, 
-                 local_dir: str = "./chroma_data",
-                 batch_size: int = 20,
-                 max_workers: int = 4,
-                 embedding_model: str = "all-MiniLM-L6-v2",
-                 collection_name: str = "markdown_docs"):
-        self.local_dir = local_dir
-        self.batch_size = batch_size
-        self.max_workers = max_workers
-        self.checkpoint_file = Path(local_dir) / "checkpoint.json"
-        self.embedding_model = embedding_model
-        self.collection_name = collection_name
-        
-        # Create local directory for checkpoints and Chroma
-        Path(local_dir).mkdir(parents=True, exist_ok=True)
-
-# Embedding engine
-class EmbeddingEngine:
-    """Handle embeddings with a lightweight model"""
-    
-    def __init__(self, model_name="all-MiniLM-L6-v2"):
-        # Use GPU if available
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        logger.info(f"Using device: {self.device}")
-        
-        # Try multiple model options in order of preference
-        model_options = [
-            model_name,
-            "all-MiniLM-L6-v2",
-            "paraphrase-MiniLM-L3-v2",
-            "all-mpnet-base-v2"  # Higher quality but larger model
-        ]
-        
-        self.model = None
-        
-        # Try each model in order until one works
-        for model_option in model_options:
-            try:
-                logger.info(f"Attempting to load model: {model_option}")
-                self.model = SentenceTransformer(model_option)
-                
-                # Move model to device
-                self.model.to(self.device)
-                
-                logger.info(f"Successfully loaded model: {model_option}")
-                self.model_name = model_option
-                self.vector_size = self.model.get_sentence_embedding_dimension()
-                break
-                
-            except Exception as e:
-                logger.warning(f"Failed to load model {model_option}: {str(e)}")
-        
-        if self.model is None:
-            logger.error("Failed to load any embedding model. Exiting.")
-            sys.exit(1)
-
-    def encode(self, text, batch_size=32):
-        """Get embedding for a text or list of texts"""
-        # Handle single text
-        if isinstance(text, str):
-            texts = [text]
-        else:
-            texts = text
-            
-        # Truncate texts if necessary to avoid tokenization issues
-        truncated_texts = [t[:50000] if len(t) > 50000 else t for t in texts]
-        
-        # Generate embeddings
-        try:
-            embeddings = self.model.encode(truncated_texts, batch_size=batch_size, 
-                                         show_progress_bar=False, convert_to_numpy=True)
-            return embeddings
-        except Exception as e:
-            logger.error(f"Error generating embeddings: {e}")
-            # Return zero embeddings as fallback
-            return np.zeros((len(truncated_texts), self.vector_size))
-
-class VectorStoreManager:
-    """Manage Chroma vector store operations - upload, query, etc."""
-    
-    def __init__(self, config: Config):
-        self.config = config
-            
-        # Initialize Chroma client (local persistence)
-        logger.info(f"Initializing Chroma at {config.local_dir}")
-        self.client = chromadb.PersistentClient(path=config.local_dir)
-        
-        # Get or create collection
-        try:
-            # Initialize embedding model
-            logger.info("Loading embedding model...")
-            self.embedding_engine = EmbeddingEngine(config.embedding_model)
-            logger.info(f"Using model: {self.embedding_engine.model_name}")
-            
-            # Create embedding function
-            sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
-                model_name=self.embedding_engine.model_name
-            )
-            
-            # Try to get existing collection
-            try:
-                self.collection = self.client.get_collection(
-                    name=config.collection_name,
-                    embedding_function=sentence_transformer_ef
-                )
-                logger.info(f"Using existing collection: {config.collection_name}")
-            except:
-                # Create new collection if it doesn't exist
-                self.collection = self.client.create_collection(
-                    name=config.collection_name,
-                    embedding_function=sentence_transformer_ef,
-                    metadata={"hnsw:space": "cosine"}
-                )
-                logger.info(f"Created new collection: {config.collection_name}")
-                
-        except Exception as e:
-            logger.error(f"Error initializing Chroma collection: {e}")
-            sys.exit(1)
-    
-    def query(self, query_text: str, n_results: int = 5) -> List[Dict]:
-        """
-        Query the vector store with a text query
-        """
-        try:
-            # Query the collection
-            search_results = self.collection.query(
-                query_texts=[query_text],
-                n_results=n_results,
-                include=["documents", "metadatas", "distances"]
-            )
-            
-            # Format results
-            results = []
-            if search_results["documents"] and len(search_results["documents"][0]) > 0:
-                for i in range(len(search_results["documents"][0])):
-                    results.append({
-                        'document': search_results["documents"][0][i],
-                        'metadata': search_results["metadatas"][0][i],
-                        'score': 1.0 - search_results["distances"][0][i]  # Convert distance to similarity
-                    })
-            
-            return results
-        except Exception as e:
-            logger.error(f"Error querying collection: {e}")
-            return []
-
-    def get_statistics(self) -> Dict[str, Any]:
-        """Get statistics about the vector store"""
-        stats = {}
-        
-        try:
-            # Get collection count
-            collection_info = self.collection.count()
-            stats['total_documents'] = collection_info
-            
-            # Estimate unique files - with no chunking, each document is a file
-            stats['unique_files'] = collection_info
-        except Exception as e:
-            logger.error(f"Error getting statistics: {e}")
-            stats['error'] = str(e)
-        
-        return stats
-
-class RAGSystem:
-    """Retrieval-Augmented Generation with multiple LLM providers"""
-    
-    def __init__(self, vector_store: VectorStoreManager):
-        self.vector_store = vector_store
-        self.openai_client = None
-        self.gemini_configured = False
-    
-    def setup_openai(self, api_key: str):
-        """Set up OpenAI client with API key"""
-        try:
-            self.openai_client = OpenAI(api_key=api_key)
-            return True
-        except Exception as e:
-            logger.error(f"Error initializing OpenAI client: {e}")
-            return False
-    
-    def setup_gemini(self, api_key: str):
-        """Set up Gemini with API key"""
-        try:
-            genai.configure(api_key=api_key)
-            self.gemini_configured = True
-            return True
-        except Exception as e:
-            logger.error(f"Error configuring Gemini: {e}")
-            return False
-    
-    def format_context(self, documents: List[Dict]) -> str:
-        """Format retrieved documents into context for the LLM"""
-        if not documents:
-            return "No relevant documents found."
-        
-        context_parts = []
-        for i, doc in enumerate(documents):
-            metadata = doc['metadata']
-            title = metadata.get('title', metadata.get('filename', 'Unknown document'))
-            
-            # For readability, limit length of context document
-            doc_text = doc['document']
-            if len(doc_text) > 10000:  # Limit long documents in context
-                doc_text = doc_text[:10000] + "... [Document truncated for context]"
-                
-            context_parts.append(f"Document {i+1} - {title}:\n{doc_text}\n")
-        
-        return "\n".join(context_parts)
-    
-    def generate_response_openai(self, query: str, context: str) -> str:
-        """Generate a response using OpenAI model with context"""
-        if not self.openai_client:
-            return "Error: OpenAI API key not configured. Please enter an API key in the settings tab."
-        
-        system_prompt = """
-        You are a helpful assistant that answers questions based on the context provided.
-        Use the information from the context to answer the user's question.
-        If the context doesn't contain the information needed, say so clearly.
-        Always cite the specific sections from the context that you used in your answer.
-        """
-        
-        try:
-            response = self.openai_client.chat.completions.create(
-                model="gpt-4o-mini",  # Use GPT-4o mini
-                messages=[
-                    {"role": "system", "content": system_prompt},
-                    {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}"}
-                ],
-                temperature=0.3,  # Lower temperature for more factual responses
-                max_tokens=1000,
-            )
-            return response.choices[0].message.content
-        except Exception as e:
-            logger.error(f"Error generating response with OpenAI: {e}")
-            return f"Error generating response with OpenAI: {str(e)}"
-    
-    def generate_response_gemini(self, query: str, context: str) -> str:
-        """Generate a response using Gemini with context"""
-        if not self.gemini_configured:
-            return "Error: Google AI API key not configured. Please enter an API key in the settings tab."
-        
-        prompt = f"""
-        You are a helpful assistant that answers questions based on the context provided.
-        Use the information from the context to answer the user's question.
-        If the context doesn't contain the information needed, say so clearly.
-        Always cite the specific sections from the context that you used in your answer.
-        
-        Context:
-        {context}
-        
-        Question: {query}
-        """
-        
-        try:
-            model = genai.GenerativeModel('gemini-1.5-flash')
-            response = model.generate_content(prompt)
-            return response.text
-        except Exception as e:
-            logger.error(f"Error generating response with Gemini: {e}")
-            return f"Error generating response with Gemini: {str(e)}"
-    
-    def query_and_generate(self, query: str, n_results: int = 5, model: str = "openai") -> str:
-        """Retrieve relevant documents and generate a response using the specified model"""
-        # Query vector store
-        documents = self.vector_store.query(query, n_results=n_results)
-        
-        if not documents:
-            return "No relevant documents found to answer your question."
-        
-        # Format context
-        context = self.format_context(documents)
-        
-        # Generate response with the appropriate model
-        if model == "openai":
-            return self.generate_response_openai(query, context)
-        elif model == "gemini":
-            return self.generate_response_gemini(query, context)
-        else:
-            return f"Unknown model: {model}"
-
-def rag_chat(query, n_results, model_choice, rag_system):
-    """Function to handle RAG chat queries"""
-    return rag_system.query_and_generate(query, n_results=int(n_results), model=model_choice)
-
-def simple_query(query, n_results, vector_store):
-    """Function to handle simple vector store queries"""
-    results = vector_store.query(query, n_results=int(n_results))
-    
-    # Format results for display
-    formatted = []
-    for i, res in enumerate(results):
-        metadata = res['metadata']
-        title = metadata.get('title', metadata.get('filename', 'Unknown'))
-        # Limit preview text for display
-        preview = res['document'][:800] + '...' if len(res['document']) > 800 else res['document']
-        formatted.append(f"**Result {i+1}** (Similarity: {res['score']:.2f})\n\n"
-                       f"**Source:** {title}\n\n"
-                       f"**Content:**\n{preview}\n\n"
-                       f"---\n")
-    
-    return "\n".join(formatted) if formatted else "No results found."
-
-def get_db_stats(vector_store):
-    """Function to get vector store statistics"""
-    stats = vector_store.get_statistics()
-    return (f"Total documents: {stats.get('total_documents', 0)}\n"
-           f"Unique files: {stats.get('unique_files', 0)}")
-
-def update_api_keys(openai_key, gemini_key, rag_system):
-    """Update API keys for the RAG system"""
-    success_msg = []
-    
-    if openai_key:
-        if rag_system.setup_openai(openai_key):
-            success_msg.append("✅ OpenAI API key configured successfully")
-        else:
-            success_msg.append("❌ Failed to configure OpenAI API key")
-    
-    if gemini_key:
-        if rag_system.setup_gemini(gemini_key):
-            success_msg.append("✅ Google AI API key configured successfully")
-        else:
-            success_msg.append("❌ Failed to configure Google AI API key")
-    
-    if not success_msg:
-        return "Please enter at least one API key"
-    
-    return "\n".join(success_msg)
-
-# Main function to run the application
-def main():
-    # Set up paths for existing Chroma database
-    chroma_dir = Path("./chroma_data")
-    
-    # Initialize the system
-    config = Config(
-        local_dir=str(chroma_dir),
-        collection_name="markdown_docs"
-    )
-    
-    # Initialize vector store manager with existing collection
-    vector_store = VectorStoreManager(config)
-    
-    # Initialize RAG system without API keys initially
-    rag_system = RAGSystem(vector_store)
-    
-    # Define Gradio app
-    def rag_chat_wrapper(query, n_results, model_choice):
-        return rag_chat(query, n_results, model_choice, rag_system)
-    
-    def simple_query_wrapper(query, n_results):
-        return simple_query(query, n_results, vector_store)
-    
-    def update_api_keys_wrapper(openai_key, gemini_key):
-        return update_api_keys(openai_key, gemini_key, rag_system)
-    
-    # Create the Gradio interface
-    with gr.Blocks(title="Markdown RAG System") as app:
-        gr.Markdown("# RAG System with Multiple LLM Providers")
-        
-        with gr.Tab("Chat with Documents"):
-            with gr.Row():
-                with gr.Column(scale=3):
-                    query_input = gr.Textbox(label="Question", placeholder="Ask a question about your documents...")
-                    num_results = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="Number of documents to retrieve")
-                    model_choice = gr.Radio(
-                        choices=["openai", "gemini"], 
-                        value="openai", 
-                        label="Choose LLM Provider",
-                        info="Select which model to use for generating answers"
-                    )
-                    query_button = gr.Button("Ask", variant="primary")
-                
-                with gr.Column(scale=7):
-                    response_output = gr.Markdown(label="Response")
-            
-            # Database stats
-            stats_display = gr.Textbox(label="Database Statistics", value=get_db_stats(vector_store))
-            refresh_button = gr.Button("Refresh Statistics")
-        
-        with gr.Tab("Document Search"):
-            search_input = gr.Textbox(label="Search Query", placeholder="Search your documents...")
-            search_num = gr.Slider(minimum=1, maximum=20, value=5, step=1, label="Number of results")
-            search_button = gr.Button("Search", variant="primary")
-            search_output = gr.Markdown(label="Search Results")
-        
-        with gr.Tab("Settings"):
-            gr.Markdown("""
-            ## API Keys Configuration
-            
-            This application can use either OpenAI's GPT-4o-mini or Google's Gemini 1.5 Flash for generating responses.
-            You need to provide at least one API key to use the chat functionality.
-            """)
-            
-            openai_key_input = gr.Textbox(
-                label="OpenAI API Key",
-                placeholder="Enter your OpenAI API key here...",
-                type="password"
-            )
-            
-            gemini_key_input = gr.Textbox(
-                label="Google AI API Key",
-                placeholder="Enter your Google AI API key here...",
-                type="password"
-            )
-            
-            save_keys_button = gr.Button("Save API Keys", variant="primary")
-            api_status = gr.Markdown("")
-        
-        # Set up events
-        query_button.click(
-            fn=rag_chat_wrapper,
-            inputs=[query_input, num_results, model_choice],
-            outputs=response_output
-        )
-        
-        refresh_button.click(
-            fn=lambda: get_db_stats(vector_store),
-            inputs=None,
-            outputs=stats_display
-        )
-        
-        search_button.click(
-            fn=simple_query_wrapper,
-            inputs=[search_input, search_num],
-            outputs=search_output
-        )
-        
-        save_keys_button.click(
-            fn=update_api_keys_wrapper,
-            inputs=[openai_key_input, gemini_key_input],
-            outputs=api_status
-        )
-    
-    # Launch the interface
-    app.launch()
-
-if __name__ == "__main__":
-    main()