Update app.py

app.py

@@ -4,193 +4,48 @@ import logging
 from pathlib import Path
 import json
 from datetime import datetime
-from typing import List, Dict, Any, Optional, Tuple, Union
-import traceback
+from typing import List, Dict, Any, Optional
 
-# Configure detailed logging with file output
-LOG_DIR = "logs"
-os.makedirs(LOG_DIR, exist_ok=True)
-log_file = os.path.join(LOG_DIR, f"rag_system_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log")
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
-# Set up root logger with both file and console handlers
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-    handlers=[
-        logging.FileHandler(log_file),
-        logging.StreamHandler(sys.stdout)
-    ]
-)
-logger = logging.getLogger("rag_system")
-logger.info(f"Starting RAG system. Log file: {log_file}")
-
-# Importing necessary libraries with error handling
-try:
-    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
-    logger.info("All required libraries successfully imported")
-except ImportError as e:
-    logger.critical(f"Failed to import required libraries: {e}")
-    print(f"ERROR: Missing required libraries. Please install with: pip install -r requirements.txt")
-    print(f"Specific error: {e}")
-    sys.exit(1)
-
-# Version info for tracking
-VERSION = "1.1.0"
-logger.info(f"RAG System Version: {VERSION}")
-
-# Custom CSS for better UI
-custom_css = """
-    .gradio-container {
-        max-width: 1200px;
-        margin: auto;
-    }
-    .gr-prose h1 {
-        font-size: 2.5rem;
-        margin-bottom: 1rem;
-        color: #1a5276;
-    }
-    .gr-prose h3 {
-        font-size: 1.25rem;
-        font-weight: 600;
-        margin-top: 1rem;
-        margin-bottom: 0.5rem;
-        color: #2874a6;
-    }
-    .container {
-        margin: 0 auto;
-        padding: 2rem;
-    }
-    .gr-box {
-        border-radius: 8px;
-        box-shadow: 0 1px 3px rgba(0,0,0,0.12), 0 1px 2px rgba(0,0,0,0.24);
-        padding: 1rem;
-        margin-bottom: 1rem;
-        background-color: #f9f9f9;
-    }
-    .footer {
-        text-align: center;
-        font-size: 0.8rem;
-        color: #666;
-        margin-top: 2rem;
-    }
-"""
+# 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 system.
-    
-    This class centralizes all configuration parameters for the application,
-    making it easier to modify settings and ensure consistency.
-    
-    Attributes:
-        local_dir (str): Directory for ChromaDB persistence
-        embedding_model (str): Name of the embedding model to use
-        collection_name (str): Name of the ChromaDB collection
-        default_top_k (int): Default number of results to return
-        openai_model (str): Default OpenAI model to use
-        gemini_model (str): Default Gemini model to use
-        temperature (float): Temperature setting for LLM generation
-        max_tokens (int): Maximum tokens for LLM response
-        system_name (str): Name of the system for UI
-        context_limit (int): Maximum characters to include in context
-    """
-    
+    """Configuration for vector store and RAG"""
     def __init__(self, 
-                 local_dir: str = "./chroma_db",
+                 local_dir: str = ".",
                  embedding_model: str = "all-MiniLM-L6-v2",
-                 collection_name: str = "markdown_docs",
-                 default_top_k: int = 8,  # Increased from 5 to 8 for more context
-                 openai_model: str = "gpt-4o-mini",
-                 gemini_model: str = "gemini-1.5-flash",
-                 temperature: float = 0.3,
-                 max_tokens: int = 2000,  # Increased from 1000 to 2000 for more comprehensive responses
-                 system_name: str = "Document Knowledge Assistant",
-                 context_limit: int = 16000):  # Increased context limit for more comprehensive context
+                 collection_name: str = "markdown_docs"):
         self.local_dir = local_dir
         self.embedding_model = embedding_model
         self.collection_name = collection_name
-        self.default_top_k = default_top_k
-        self.openai_model = openai_model
-        self.gemini_model = gemini_model
-        self.temperature = temperature
-        self.max_tokens = max_tokens
-        self.system_name = system_name
-        self.context_limit = context_limit
-        
-        # Create local directory if it doesn't exist
-        os.makedirs(local_dir, exist_ok=True)
-        
-        logger.info(f"Initialized configuration: {self.__dict__}")
-    
-    def to_dict(self) -> Dict[str, Any]:
-        """Convert configuration to dictionary for serialization"""
-        return self.__dict__
-    
-    @classmethod
-    def from_file(cls, config_path: str) -> 'Config':
-        """Load configuration from JSON file"""
-        try:
-            with open(config_path, 'r') as f:
-                config_dict = json.load(f)
-            logger.info(f"Loaded configuration from {config_path}")
-            return cls(**config_dict)
-        except Exception as e:
-            logger.error(f"Failed to load configuration from {config_path}: {e}")
-            logger.info("Using default configuration")
-            return cls()
-    
-    def save_to_file(self, config_path: str) -> bool:
-        """Save configuration to JSON file"""
-        try:
-            with open(config_path, 'w') as f:
-                json.dump(self.to_dict(), f, indent=2)
-            logger.info(f"Saved configuration to {config_path}")
-            return True
-        except Exception as e:
-            logger.error(f"Failed to save configuration to {config_path}: {e}")
-            return False
 
+# Embedding engine
 class EmbeddingEngine:
-    """
-    Handle embeddings with a lightweight model.
-    
-    This class manages the embedding model used to convert text to vector 
-    representations for semantic search.
-    
-    Attributes:
-        model (SentenceTransformer): The loaded embedding model
-        model_name (str): Name of the successfully loaded model
-        vector_size (int): Dimension of the embedding vectors
-        device (str): Device used for inference ('cuda' or 'cpu')
-    """
+    """Handle embeddings with a lightweight model"""
     
     def __init__(self, model_name="all-MiniLM-L6-v2"):
-        """
-        Initialize the embedding engine with the specified model.
-        
-        Args:
-            model_name (str): Name of the embedding model to load
-            
-        Raises:
-            SystemExit: If no embedding model could be loaded
-        """
         # Use GPU if available
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        logger.info(f"Using device for embeddings: {self.device}")
+        logger.info(f"Using device: {self.device}")
         
         # Try multiple model options in order of preference
         model_options = [
             model_name,
-            "all-MiniLM-L6-v2",      # Good balance of speed and quality
-            "paraphrase-MiniLM-L3-v2", # Faster but less accurate
-            "all-mpnet-base-v2"      # Higher quality but larger model
+            "all-MiniLM-L6-v2",
+            "paraphrase-MiniLM-L3-v2",
+            "all-mpnet-base-v2"  # Higher quality but larger model
         ]
         
         self.model = None
@@ -198,99 +53,47 @@ class EmbeddingEngine:
         # Try each model in order until one works
         for model_option in model_options:
             try:
-                logger.info(f"Attempting to load embedding model: {model_option}")
+                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 embedding model: {model_option}")
+                logger.info(f"Successfully loaded model: {model_option}")
                 self.model_name = model_option
                 self.vector_size = self.model.get_sentence_embedding_dimension()
-                logger.info(f"Embedding vector size: {self.vector_size}")
                 break
                 
             except Exception as e:
-                logger.warning(f"Failed to load embedding model {model_option}: {str(e)}")
+                logger.warning(f"Failed to load model {model_option}: {str(e)}")
         
         if self.model is None:
-            error_msg = "Failed to load any embedding model. Please check your internet connection or install models locally."
-            logger.critical(error_msg)
-            raise SystemExit(error_msg)
-    
-    def embed(self, texts: List[str]) -> np.ndarray:
-        """
-        Generate embeddings for a list of texts.
-        
-        Args:
-            texts (List[str]): List of texts to embed
-            
-        Returns:
-            np.ndarray: Array of embeddings
-            
-        Raises:
-            ValueError: If the input is invalid
-            RuntimeError: If embedding fails
-        """
-        if not texts:
-            raise ValueError("Cannot embed empty list of texts")
-        
-        try:
-            embeddings = self.model.encode(texts, convert_to_numpy=True)
-            return embeddings
-        except Exception as e:
-            logger.error(f"Error generating embeddings: {e}")
-            raise RuntimeError(f"Failed to generate embeddings: {e}")
+            logger.error("Failed to load any embedding model. Exiting.")
+            sys.exit(1)
 
 class VectorStoreManager:
-    """
-    Manage Chroma vector store operations - upload, query, etc.
-    
-    This class provides an interface to the ChromaDB vector database,
-    handling document storage, retrieval, and management.
-    
-    Attributes:
-        config (Config): Configuration parameters
-        client (chromadb.PersistentClient): ChromaDB client
-        collection (chromadb.Collection): The active ChromaDB collection
-        embedding_engine (EmbeddingEngine): Engine for generating embeddings
-    """
+    """Manage Chroma vector store operations - upload, query, etc."""
     
     def __init__(self, config: Config):
-        """
-        Initialize the vector store manager.
-        
-        Args:
-            config (Config): Configuration parameters
-            
-        Raises:
-            SystemExit: If the vector store cannot be initialized
-        """
         self.config = config
             
         # Initialize Chroma client (local persistence)
         logger.info(f"Initializing Chroma at {config.local_dir}")
-        try:
-            self.client = chromadb.PersistentClient(path=config.local_dir)
-            logger.info("ChromaDB client initialized successfully")
-        except Exception as e:
-            error_msg = f"Failed to initialize ChromaDB client: {e}"
-            logger.critical(error_msg)
-            raise SystemExit(error_msg)
+        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 embedding model: {self.embedding_engine.model_name}")
+            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 or create a new one
+            # Try to get existing collection
             try:
                 self.collection = self.client.get_collection(
                     name=config.collection_name,
@@ -298,7 +101,7 @@ class VectorStoreManager:
                 )
                 logger.info(f"Using existing collection: {config.collection_name}")
             except Exception as e:
-                logger.warning(f"Error getting collection: {e}")
+                logger.error(f"Error getting collection: {e}")
                 # Attempt to get a list of available collections
                 collections = self.client.list_collections()
                 if collections:
@@ -319,28 +122,14 @@ class VectorStoreManager:
                     logger.info(f"Created new collection: {config.collection_name}")
                 
         except Exception as e:
-            error_msg = f"Error initializing Chroma collection: {e}"
-            logger.critical(error_msg)
-            raise SystemExit(error_msg)
+            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.
-        
-        Args:
-            query_text (str): The query text
-            n_results (int): Number of results to return
-            
-        Returns:
-            List[Dict]: List of results with document text, metadata, and similarity score
+        Query the vector store with a text query
         """
-        if not query_text.strip():
-            logger.warning("Empty query received")
-            return []
-        
         try:
-            logger.info(f"Querying vector store with: '{query_text[:50]}...' (top {n_results})")
-            
             # Query the collection
             search_results = self.collection.query(
                 query_texts=[query_text],
@@ -354,106 +143,26 @@ class VectorStoreManager:
                 for i in range(len(search_results["documents"][0])):
                     results.append({
                         'document': search_results["documents"][0][i],
-                        'metadata': search_results["metadatas"][0][i] if search_results["metadatas"] else {},
-                        'score': 1.0 - search_results["distances"][0][i],  # Convert distance to similarity
-                        'distance': search_results["distances"][0][i]
+                        'metadata': search_results["metadatas"][0][i],
+                        'score': 1.0 - search_results["distances"][0][i]  # Convert distance to similarity
                     })
-                
-                logger.info(f"Found {len(results)} results for query")
-            else:
-                logger.info("No results found for query")
             
             return results
         except Exception as e:
             logger.error(f"Error querying collection: {e}")
-            logger.debug(traceback.format_exc())
             return []
 
-    def add_document(self, 
-                    document: str, 
-                    doc_id: str, 
-                    metadata: Dict[str, Any]) -> bool:
-        """
-        Add a document to the vector store.
-        
-        Args:
-            document (str): The document text
-            doc_id (str): Unique identifier for the document
-            metadata (Dict[str, Any]): Metadata about the document
-            
-        Returns:
-            bool: True if successful, False otherwise
-        """
-        try:
-            logger.info(f"Adding document '{doc_id}' to vector store")
-            
-            # Add the document to the collection
-            self.collection.add(
-                documents=[document],
-                ids=[doc_id],
-                metadatas=[metadata]
-            )
-            
-            logger.info(f"Successfully added document '{doc_id}'")
-            return True
-        except Exception as e:
-            logger.error(f"Error adding document to collection: {e}")
-            return False
-
-    def delete_document(self, doc_id: str) -> bool:
-        """
-        Delete a document from the vector store.
-        
-        Args:
-            doc_id (str): ID of the document to delete
-            
-        Returns:
-            bool: True if successful, False otherwise
-        """
-        try:
-            logger.info(f"Deleting document '{doc_id}' from vector store")
-            self.collection.delete(ids=[doc_id])
-            logger.info(f"Successfully deleted document '{doc_id}'")
-            return True
-        except Exception as e:
-            logger.error(f"Error deleting document from collection: {e}")
-            return False
-
     def get_statistics(self) -> Dict[str, Any]:
-        """
-        Get statistics about the vector store.
-        
-        Returns:
-            Dict[str, Any]: Statistics about the vector store
-        """
-        stats = {
-            'collection_name': self.config.collection_name,
-            'embedding_model': self.embedding_engine.model_name,
-            'embedding_dimensions': self.embedding_engine.vector_size,
-            'device': self.embedding_engine.device
-        }
+        """Get statistics about the vector store"""
+        stats = {}
         
         try:
             # Get collection count
-            collection_count = self.collection.count()
-            stats['total_documents'] = collection_count
+            collection_info = self.collection.count()
+            stats['total_documents'] = collection_info
             
-            # Get unique metadata values
-            if collection_count > 0:
-                try:
-                    # Get a sample of document metadata
-                    sample_results = self.collection.get(limit=min(collection_count, 100))
-                    if sample_results and 'metadatas' in sample_results and sample_results['metadatas']:
-                        # Count unique files if filename exists in metadata
-                        filenames = set()
-                        for metadata in sample_results['metadatas']:
-                            if 'filename' in metadata:
-                                filenames.add(metadata['filename'])
-                        stats['unique_files'] = len(filenames)
-                except Exception as e:
-                    logger.warning(f"Error getting metadata statistics: {e}")
-            
-            logger.info(f"Vector store statistics: {stats}")
+            # 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)
@@ -461,635 +170,274 @@ class VectorStoreManager:
         return stats
 
 class RAGSystem:
-    """
-    Retrieval-Augmented Generation with multiple LLM providers.
-    
-    This class handles the RAG workflow: retrieval of relevant documents,
-    formatting context, and generating responses with different LLM providers.
-    
-    Attributes:
-        vector_store (VectorStoreManager): Manager for vector store operations
-        openai_client (Optional[OpenAI]): OpenAI client
-        gemini_configured (bool): Whether Gemini API is configured
-        config (Config): Configuration parameters
-    """
+    """Retrieval-Augmented Generation with multiple LLM providers"""
     
-    def __init__(self, vector_store: VectorStoreManager, config: Config):
-        """
-        Initialize the RAG system.
-        
-        Args:
-            vector_store (VectorStoreManager): Vector store manager
-            config (Config): Configuration parameters
-        """
+    def __init__(self, vector_store: VectorStoreManager):
         self.vector_store = vector_store
-        self.config = config
         self.openai_client = None
         self.gemini_configured = False
-        
-        logger.info("Initialized RAG system")
     
-    def setup_openai(self, api_key: str) -> bool:
-        """
-        Set up OpenAI client with API key.
-        
-        Args:
-            api_key (str): OpenAI API key
-            
-        Returns:
-            bool: True if successful, False otherwise
-        """
-        if not api_key.strip():
-            logger.warning("Empty OpenAI API key provided")
-            return False
-            
+    def setup_openai(self, api_key: str):
+        """Set up OpenAI client with API key"""
         try:
-            logger.info("Setting up OpenAI client")
             self.openai_client = OpenAI(api_key=api_key)
-            # Test the API key with a simple request
-            response = self.openai_client.chat.completions.create(
-                model=self.config.openai_model,
-                messages=[
-                    {"role": "system", "content": "You are a helpful assistant."},
-                    {"role": "user", "content": "Test connection"}
-                ],
-                max_tokens=10
-            )
-            logger.info("OpenAI client configured successfully")
             return True
         except Exception as e:
             logger.error(f"Error initializing OpenAI client: {e}")
-            self.openai_client = None
             return False
     
-    def setup_gemini(self, api_key: str) -> bool:
-        """
-        Set up Gemini with API key.
-        
-        Args:
-            api_key (str): Google AI API key
-            
-        Returns:
-            bool: True if successful, False otherwise
-        """
-        if not api_key.strip():
-            logger.warning("Empty Gemini API key provided")
-            return False
-            
+    def setup_gemini(self, api_key: str):
+        """Set up Gemini with API key"""
         try:
-            logger.info("Setting up Gemini client")
             genai.configure(api_key=api_key)
-            
-            # Test the API key with a simple request
-            model = genai.GenerativeModel(self.config.gemini_model)
-            response = model.generate_content("Test connection")
-            
             self.gemini_configured = True
-            logger.info("Gemini client configured successfully")
             return True
         except Exception as e:
             logger.error(f"Error configuring Gemini: {e}")
-            self.gemini_configured = False
             return False
     
     def format_context(self, documents: List[Dict]) -> str:
-        """
-        Format retrieved documents into context for the LLM.
-        
-        Args:
-            documents (List[Dict]): List of retrieved documents
-            
-        Returns:
-            str: Formatted context for the LLM
-        """
+        """Format retrieved documents into context for the LLM"""
         if not documents:
-            logger.warning("No documents provided for context formatting")
             return "No relevant documents found."
         
-        logger.info(f"Formatting {len(documents)} documents for context")
         context_parts = []
-        
         for i, doc in enumerate(documents):
             metadata = doc['metadata']
-            # Extract document metadata in a robust way
             title = metadata.get('title', metadata.get('filename', 'Unknown document'))
             
-            # Format header with just essential metadata for cleaner context
-            header = f"Document {i+1} - {title}"
-            
-            # For readability, limit length of context document (using config value)
+            # For readability, limit length of context document
             doc_text = doc['document']
-            if len(doc_text) > (self.config.context_limit // len(documents)):  
-                # Divide context limit among the documents
-                max_length = self.config.context_limit // len(documents)
-                doc_text = doc_text[:max_length] + "... [Document truncated for brevity]"
+            if len(doc_text) > 10000:  # Limit long documents in context
+                doc_text = doc_text[:10000] + "... [Document truncated for context]"
                 
-            context_parts.append(f"{header}:\n{doc_text}\n")
+            context_parts.append(f"Document {i+1} - {title}:\n{doc_text}\n")
         
-        full_context = "\n".join(context_parts)
-        logger.info(f"Created context with {len(full_context)} characters")
-        
-        return full_context
+        return "\n".join(context_parts)
     
     def generate_response_openai(self, query: str, context: str) -> str:
-        """
-        Generate a response using OpenAI model with context.
-        
-        Args:
-            query (str): User query
-            context (str): Formatted document context
-            
-        Returns:
-            str: Generated response
-        """
+        """Generate a response using OpenAI model with context"""
         if not self.openai_client:
-            logger.warning("OpenAI API key not configured for response generation")
-            return "Please configure an OpenAI API key to use this feature. Enter your API key in the field and click 'Save API Key'."
+            return "Error: OpenAI API key not configured. Please enter an API key in the API key field."
         
-        # Improved system prompt for better, more comprehensive responses
         system_prompt = """
-        You are an exceptionally helpful, clear, and friendly AI research assistant. Your goal is to provide comprehensive, well-structured, and insightful answers based on the provided document context.
-
-        Guidelines for your response:
-        
-        1. USE ONLY the information contained in the provided context documents to form your answer. If the context doesn't contain enough information to provide a complete answer, acknowledge this limitation clearly.
-        
-        2. Always provide well-structured, detailed responses between 300-500 words that thoroughly address the user's question.
-        
-        3. Format your response with clear headings, bullet points, or numbered lists when appropriate to enhance readability.
-        
-        4. Cite your sources by referring to the document numbers (e.g., "According to Document 1...") to support your claims.
-        
-        5. Use a friendly, conversational, and supportive tone that makes complex information accessible.
-        
-        6. If different documents offer conflicting information, acknowledge these differences and present both perspectives without bias.
-        
-        7. When appropriate, organize information into logical categories or chronological order to improve clarity.
-        
-        8. Use examples from the documents to illustrate key points when available.
-        
-        9. Conclude with a brief summary of the main points if the answer is complex.
-        
-        10. Remember to stay focused on the user's specific question while providing sufficient context for complete understanding.
+        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:
-            logger.info(f"Generating response with OpenAI ({self.config.openai_model})")
-            
-            start_time = datetime.now()
             response = self.openai_client.chat.completions.create(
-                model=self.config.openai_model,
+                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=self.config.temperature,
-                max_tokens=self.config.max_tokens,
+                temperature=0.3,  # Lower temperature for more factual responses
+                max_tokens=5000,
             )
-            
-            generation_time = (datetime.now() - start_time).total_seconds()
-            response_text = response.choices[0].message.content
-            
-            logger.info(f"Generated response with OpenAI in {generation_time:.2f} seconds")
-            return response_text
+            return response.choices[0].message.content
         except Exception as e:
-            error_msg = f"Error generating response with OpenAI: {str(e)}"
-            logger.error(error_msg)
-            return f"I encountered an error while generating your response. Please try again or check your API key. Error details: {str(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.
-        
-        Args:
-            query (str): User query
-            context (str): Formatted document context
-            
-        Returns:
-            str: Generated response
-        """
+        """Generate a response using Gemini with context"""
         if not self.gemini_configured:
-            logger.warning("Gemini API key not configured for response generation")
-            return "Please configure a Google AI API key to use this feature. Enter your API key in the field and click 'Save API Key'."
+            return "Error: Google AI API key not configured. Please enter an API key in the API key field."
         
-        # Improved Gemini prompt for more comprehensive and user-friendly responses
         prompt = f"""
-        You are a knowledgeable and friendly research assistant who excels at providing clear, comprehensive, and well-structured responses. Your goal is to help users understand complex information from documents in an accessible way.
-
-        **Guidelines for Your Response:**
         
-        - Create a detailed, well-organized response of approximately 300-500 words that thoroughly addresses the user's question.
-        - Use ONLY information from the provided context documents.
-        - Structure your answer with clear paragraphs, and use headings, bullet points, or numbered lists when appropriate.
-        - Maintain a friendly, conversational tone that makes information accessible and engaging.
-        - When citing information, reference specific documents by number (e.g., "As mentioned in Document 2...").
-        - If the context doesn't contain enough information for a complete answer, acknowledge this limitation while providing what you can from the available context.
-        - If documents contain conflicting information, present both perspectives fairly.
-        - Conclude with a brief summary if the topic is complex.
-        
-        **Context Documents:**
+<prompt>
+    <system>
+        <name>Loss Dog</name>
+        <role>You are a highly intelligent AI specializing in labor market analysis, job trends, and skillset forecasting. You utilize a combination of structured data from sources like the Bureau of Labor Statistics (BLS) and the World Economic Forum (WEF), alongside advanced retrieval-augmented generation (RAG) techniques.</role>
+        <goal>Your mission is to provide insightful, data-driven, and comprehensive answers to users seeking career and job market intelligence. You must ensure clarity, depth, and practical relevance in all responses.</goal>
+        <personality>
+            <tone>Friendly, professional, and engaging</tone>
+            <depth>Detailed, nuanced, and well-explained</depth>
+            <clarity>Well-structured with headings, citations, and easy-to-follow breakdowns</clarity>
+        </personality>
+        <methodology>
+            <data_sources>
+                <source>Bureau of Labor Statistics (BLS)</source>
+                <source>World Economic Forum (WEF) reports</source>
+                <source>Market research studies</source>
+                <source>Industry whitepapers</source>
+                <source>Company hiring trends</source>
+            </data_sources>
+            <reasoning_strategy>
+                <if_data_available>
+                    <response>
+                        Use precise statistics, industry insights, and expert analyses from retrieved sources to craft an evidence-based answer.
+                    </response>
+                </if_data_available>
+                <if_data_unavailable>
+                    <response>
+                        Clearly state that the exact data is unavailable. However, provide a **comprehensive explanation** using logical deduction, adjacent industry trends, historical patterns, and economic principles.
+                    </response>
+                </if_data_unavailable>
+            </reasoning_strategy>
+            <output_expectations>
+                <length>100-500 words, depending on complexity and sources available</length>
+                <structure>
+                    <section>Introduction (sets context and purpose)</section>
+                    <section>Data-backed analysis (citing retrieved sources)</section>
+                    <section>Logical deduction and reasoning (when necessary)</section>
+                    <section>Conclusion (summarizes insights and provides actionable takeaways)</section>
+                </structure>
+                <citation_style>Clearly cite data sources within the response (e.g., "According to BLS 2024 report...").</citation_style>
+                <engagement>Encourage follow-up questions and deeper exploration where relevant.</engagement>
+            </output_expectations>
+        </methodology>
+    </system>
+        Context:
         {context}
         
-        **User's Question:** 
-        {query}
-
-        **Your Response:** 
+        Question: {query}
         """
-
+        
         try:
-            logger.info(f"Generating response with Gemini ({self.config.gemini_model})")
-            
-            start_time = datetime.now()
-            model = genai.GenerativeModel(self.config.gemini_model)
-            
-            generation_config = {
-                "temperature": self.config.temperature,
-                "max_output_tokens": self.config.max_tokens,
-                "top_p": 0.9,
-                "top_k": 40
-            }
-            
-            response = model.generate_content(
-                prompt,
-                generation_config=generation_config
-            )
-            
-            generation_time = (datetime.now() - start_time).total_seconds()
-            response_text = response.text
-            
-            logger.info(f"Generated response with Gemini in {generation_time:.2f} seconds")
-            return response_text
+            model = genai.GenerativeModel('gemini-1.5-flash')
+            response = model.generate_content(prompt)
+            return response.text
         except Exception as e:
-            error_msg = f"Error generating response with Gemini: {str(e)}"
-            logger.error(error_msg)
-            return f"I encountered an error while generating your response. Please try again or check your API key. Error details: {str(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") -> Tuple[str, str]:
-        """
-        Retrieve relevant documents and generate a response using the specified model.
-        
-        Args:
-            query (str): User query
-            n_results (int): Number of documents to retrieve
-            model (str): Model provider to use ('openai' or 'gemini')
-            
-        Returns:
-            Tuple[str, str]: (Generated response, Search results)
-        """
-        if not query.strip():
-            logger.warning("Empty query received")
-            return "Please enter a question to get a response.", "No search performed."
-        
-        logger.info(f"Processing query: '{query[:50]}...' with {model} model")
-        
+    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)
         
-        # Format search results (for logs and hidden UI component)
-        # We'll format this in a way that's more useful for reference but not shown in UI
-        formatted_results = []
-        for i, res in enumerate(documents):
-            metadata = res['metadata']
-            title = metadata.get('title', metadata.get('filename', 'Unknown'))
-            score = res['score']
-            
-            # Only include a very brief preview for reference
-            preview = res['document'][:100] + '...' if len(res['document']) > 100 else res['document']
-            formatted_results.append(f"Document {i+1}: {title} (Relevance: {score:.2f})")
-        
-        search_output_text = "\n".join(formatted_results) if formatted_results else "No relevant documents found."
-        
         if not documents:
-            logger.warning("No relevant documents found")
-            return "I couldn't find relevant information in the knowledge base to answer your question. Could you try rephrasing your question or ask about a different topic?", search_output_text
+            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":
-            response = self.generate_response_openai(query, context)
+            return self.generate_response_openai(query, context)
         elif model == "gemini":
-            response = self.generate_response_gemini(query, context)
+            return self.generate_response_gemini(query, context)
         else:
-            error_msg = f"Unknown model: {model}"
-            logger.error(error_msg)
-            return error_msg, search_output_text
-        
-        return response, search_output_text
-
-def get_db_stats(vector_store: VectorStoreManager) -> str:
-    """
-    Function to get vector store statistics.
-    
-    Args:
-        vector_store (VectorStoreManager): Vector store manager
-        
-    Returns:
-        str: Formatted statistics string
-    """
-    try:
-        stats = vector_store.get_statistics()
-        total_docs = stats.get('total_documents', 0)
-        
-        stats_text = f"Documents in knowledge base: {total_docs}"
-        return stats_text
-    except Exception as e:
-        logger.error(f"Error getting statistics: {e}")
-        return "Error getting database statistics"
-
-# Helper function for loading documents (can be expanded in future versions)
-def load_document(file_path: str, chunk_size: int = 2000, chunk_overlap: int = 200) -> bool:
-    """
-    Load a document into the vector store.
-    
-    Args:
-        file_path (str): Path to the document
-        chunk_size (int): Size of chunks to split the document into
-        chunk_overlap (int): Overlap between chunks
-        
-    Returns:
-        bool: True if successful, False otherwise
-    """
-    try:
-try:
-        logger.info(f"Loading document: {file_path}")
-        
-        # Initialize components
-        config = Config()
-        vector_store = VectorStoreManager(config)
-        
-        # Read the file with different encodings if needed
-        content = None
-        encodings = ['utf-8', 'latin-1', 'cp1252']
-        
-        for encoding in encodings:
-            try:
-                with open(file_path, 'r', encoding=encoding) as f:
-                    content = f.read()
-                logger.info(f"Successfully read file with {encoding} encoding")
-                break
-            except UnicodeDecodeError:
-                logger.warning(f"Failed to read with {encoding} encoding, trying next...")
-        
-        if content is None:
-            logger.error(f"Failed to read file with any encoding: {file_path}")
-            return False
-        
-        # Extract metadata
-        file_name = os.path.basename(file_path)
-        file_ext = os.path.splitext(file_name)[1].lower()
-        file_size = os.path.getsize(file_path)
-        file_mtime = os.path.getmtime(file_path)
-        
-        # Try to extract title from content for better reference
-        title = file_name
-        try:
-            # Simple heuristic to find a title (first non-empty line)
-            lines = content.split('\n')
-            for line in lines:
-                line = line.strip()
-                if line and len(line) < 100:  # Reasonable title length
-                    title = line
-                    break
-        except:
-            pass
-        
-        # Create metadata
-        metadata = {
-            'filename': file_name,
-            'title': title,
-            'path': file_path,
-            'extension': file_ext,
-            'size': file_size,
-            'modified': datetime.fromtimestamp(file_mtime).isoformat(),
-            'created_at': datetime.now().isoformat()
-        }
-        
-        # Generate a unique ID for the document
-        doc_id = f"{file_name}_{hash(content)}"
-        
-        # Add to vector store
-        success = vector_store.add_document(content, doc_id, metadata)
-        
-        logger.info(f"Document loaded successfully: {file_path}" if success else f"Failed to load document: {file_path}")
-        return success
-    
-    except Exception as e:
-        logger.error(f"Error loading document {file_path}: {e}")
-        logger.error(traceback.format_exc())
-        return False
+            return f"Unknown model: {model}"
 
+# Main function to run the application
 def main():
-    """Main function to run the RAG application"""
-    # Path for configuration file
-    CONFIG_FILE_PATH = "rag_config.json"
+    # Initialize the system with current directory as the Chroma location
+    config = Config(
+        local_dir=".",  # Look for Chroma files in current directory
+        collection_name="markdown_docs"
+    )
     
     try:
-        # Try to load configuration from file, or use defaults
-        if os.path.exists(CONFIG_FILE_PATH):
-            config = Config.from_file(CONFIG_FILE_PATH)
-        else:
-            config = Config(
-                local_dir="./chroma_db",  # Store Chroma files in dedicated directory
-                collection_name="markdown_docs"
-            )
-            # Save default configuration
-            config.save_to_file(CONFIG_FILE_PATH)
-        
-        print(f"Starting Document Knowledge Assistant v{VERSION}")
-        print(f"Log file: {log_file}")
-        
         # Initialize vector store manager with existing collection
         vector_store = VectorStoreManager(config)
         
         # Initialize RAG system without API keys initially
-        rag_system = RAGSystem(vector_store, config)
+        rag_system = RAGSystem(vector_store)
         
-        # Create the Gradio interface with custom CSS
-        with gr.Blocks(title="Document Knowledge Assistant", css=custom_css) as app:
-            gr.Markdown(f"# Document Knowledge Assistant v{VERSION}")
-            gr.Markdown("Ask questions about your documents and get comprehensive AI-powered answers")
+        # Create the Gradio interface
+        with gr.Blocks(title="Document RAG System") as app:
+            gr.Markdown("# Document RAG System")
             
-            # Main layout
             with gr.Row():
-                # Left column for asking questions
-                with gr.Column(scale=3):
-                    with gr.Box():
-                        gr.Markdown("### Ask Your Question")
-                        query_input = gr.Textbox(
-                            label="", 
-                            placeholder="What would you like to know about your documents?",
-                            lines=3
-                        )
-                        
-                        with gr.Row():
-                            query_button = gr.Button("Ask Question", variant="primary", scale=3)
-                            clear_button = gr.Button("Clear", variant="secondary", scale=1)
-                    
-                    with gr.Box():
-                        gr.Markdown("### Answer")
-                        response_output = gr.Markdown()
-                
-                # Right column for settings
                 with gr.Column(scale=1):
                     # API Keys and model selection
-                    with gr.Accordion("AI Model Settings", open=True):
-                        gr.Markdown("### AI Configuration")
-                        model_choice = gr.Radio(
-                            choices=["openai", "gemini"], 
-                            value="openai", 
-                            label="AI Provider",
-                            info=f"Select your preferred AI model"
-                        )
-                        
-                        api_key_input = gr.Textbox(
-                            label="API Key",
-                            placeholder="Enter your API key here...",
-                            type="password",
-                            info="Your key is not stored between sessions"
-                        )
-                        
-                        save_key_button = gr.Button("Save API Key", variant="primary")
-                        api_status = gr.Markdown("")
+                    model_choice = gr.Radio(
+                        choices=["openai", "gemini"], 
+                        value="openai", 
+                        label="Choose LLM Provider",
+                        info="Select which model to use (GPT-4o mini or Gemini 1.5 Flash)"
+                    )
                     
-                    # Advanced search controls
-                    with gr.Accordion("Advanced Settings", open=False):
-                        gr.Markdown("### Search & Response Settings")
-                        num_results = gr.Slider(
-                            minimum=3, 
-                            maximum=15, 
-                            value=config.default_top_k, 
-                            step=1, 
-                            label="Documents to search",
-                            info="Higher values provide more context"
-                        )
-                        
-                        temperature_slider = gr.Slider(
-                            minimum=0.0,
-                            maximum=1.0,
-                            value=config.temperature,
-                            step=0.05,
-                            label="Creativity",
-                            info="Lower = more factual, Higher = more creative"
-                        )
-                        
-                        max_tokens_slider = gr.Slider(
-                            minimum=500,
-                            maximum=4000,
-                            value=config.max_tokens,
-                            step=100,
-                            label="Response Length",
-                            info="Maximum words in response"
-                        )
+                    api_key_input = gr.Textbox(
+                        label="API Key",
+                        placeholder="Enter your API key here...",
+                        type="password"
+                    )
                     
-                    # Database stats - simplified
-                    with gr.Accordion("System Info", open=False):
-                        stats_display = gr.Markdown(get_db_stats(vector_store))
-                        
-                        gr.Markdown(f"""
-                        **System Details:**
-                        - Version: {VERSION}
-                        - Embedding: {vector_store.embedding_engine.model_name}
-                        - Device: {vector_store.embedding_engine.device}
-                        """)
-                        refresh_button = gr.Button("Refresh", variant="secondary", size="sm")
-            
-            # Hidden element for search results (not visible to user)
-            with gr.Accordion("Debug Information", open=False, visible=False):
-                search_output = gr.Markdown()
-            
-            # Query history at the bottom (optional section)
-            with gr.Accordion("Recent Questions", open=False):
-                history_list = gr.Dataframe(
-                    headers=["Time", "Question", "Model"],
-                    datatype=["str", "str", "str"],
-                    row_count=5,
-                    col_count=(3, "fixed"),
-                    interactive=False
-                )
+                    save_key_button = gr.Button("Save API Key", variant="primary")
+                    api_status = gr.Markdown("")
+                    
+                    # Search controls
+                    num_results = gr.Slider(
+                        minimum=1, 
+                        maximum=10, 
+                        value=10, 
+                        step=1, 
+                        label="Number of documents to retrieve"
+                    )
+                    
+                    # Database stats
+                    gr.Markdown("### Database Statistics")
+                    stats_display = gr.Textbox(
+                        label="", 
+                        value=get_db_stats(vector_store),
+                        lines=2
+                    )
+                    refresh_button = gr.Button("Refresh Stats")
                 
-            # Footer
-            gr.Markdown(
-                """<div class="footer">Document Knowledge Assistant helps you get insights from your documents using AI. 
-                Powered by Retrieval Augmented Generation.</div>"""
-            )
-            
-            # Query history storage
-            query_history = []
+                with gr.Column(scale=2):
+                    # Query and response
+                    query_input = gr.Textbox(
+                        label="Your Question", 
+                        placeholder="Ask a question about your documents...",
+                        lines=2
+                    )
+                    
+                    query_button = gr.Button("Ask Question", variant="primary")
+                    
+                    gr.Markdown("### Response")
+                    response_output = gr.Markdown()
+                    
+                    gr.Markdown("### Document Search Results")
+                    search_output = gr.Markdown()
             
             # Function to update API key based on selected model
             def update_api_key(api_key, model):
-                if not api_key.strip():
-                    return "❌ API key cannot be empty"
-                    
                 if model == "openai":
                     success = rag_system.setup_openai(api_key)
-                    model_name = f"OpenAI {config.openai_model}"
+                    model_name = "OpenAI GPT-4o mini"
                 else:
                     success = rag_system.setup_gemini(api_key)
-                    model_name = f"Google {config.gemini_model}"
+                    model_name = "Google Gemini 1.5 Flash"
                 
                 if success:
-                    return f"✅ {model_name} connected successfully"
+                    return f"✅ {model_name} API key configured successfully"
                 else:
-                    return f"❌ Connection failed. Please check your API key and try again."
+                    return f"❌ Failed to configure {model_name} API key"
             
             # Query function that returns both response and search results
-            def query_and_search(query, n_results, model, temperature, max_tokens):
-                # Update configuration with current UI values
-                config.temperature = float(temperature)
-                config.max_tokens = int(max_tokens)
+            def query_and_search(query, n_results, model):
+                # Get search results first
+                results = vector_store.query(query, n_results=int(n_results))
                 
-                start_time = datetime.now()
+                # Format search results
+                formatted_results = []
+                for i, res in enumerate(results):
+                    metadata = res['metadata']
+                    title = metadata.get('title', metadata.get('filename', 'Unknown'))
+                    preview = res['document'][:500] + '...' if len(res['document']) > 500 else res['document']
+                    formatted_results.append(f"**Result {i+1}** (Similarity: {res['score']:.2f})\n"
+                                          f"**Source:** {title}\n"
+                                          f"**Preview:**\n{preview}\n\n---\n")
                 
-                if not query.strip():
-                    return "Please enter a question to get an answer.", "", query_history[-5:] if query_history else []
+                search_output_text = "\n".join(formatted_results) if formatted_results else "No results found."
                 
-                try:
-                    # Verify that API keys are configured
-                    if (model == "openai" and rag_system.openai_client is None) or \
-                       (model == "gemini" and not rag_system.gemini_configured):
-                        return "Please configure your API key first. Enter your API key in the settings panel and click 'Save API Key'.", "", query_history[-5:] if query_history else []
-                    
-                    # Call the RAG system's query and generate function
-                    response, search_output_text = rag_system.query_and_generate(
-                        query=query, 
-                        n_results=int(n_results),
-                        model=model
-                    )
-                    
-                    # Add to history
-                    timestamp = datetime.now().strftime("%H:%M")
-                    query_history.append([timestamp, query, model])
-                    
-                    # Keep only the last 100 queries
-                    if len(query_history) > 100:
-                        query_history.pop(0)
-                    
-                    # Update the history display with the most recent entries (reverse chronological)
-                    recent_history = list(reversed(query_history[-5:])) if len(query_history) >= 5 else list(reversed(query_history))
-                    
-                    # Calculate elapsed time
-                    elapsed_time = (datetime.now() - start_time).total_seconds()
-                    
-                    # Add subtle timing information to the response
-                    response_with_timing = f"{response}\n\n<small>Answered in {elapsed_time:.1f}s</small>"
-                    
-                    return response_with_timing, search_output_text, recent_history
+                # Generate response if we have results
+                response = "No documents found to answer your question."
+                if results:
+                    context = rag_system.format_context(results)
+                    if model == "openai":
+                        response = rag_system.generate_response_openai(query, context)
+                    else:
+                        response = rag_system.generate_response_gemini(query, context)
                 
-                except Exception as e:
-                    error_msg = f"Error processing query: {str(e)}"
-                    logger.error(error_msg)
-                    logger.error(traceback.format_exc())
-                    return "I encountered an error while processing your question. Please try again or check your API key settings.", "", query_history[-5:] if query_history else []
-            
-            # Function to clear the input and results
-            def clear_inputs():
-                return "", "", "", query_history[-5:] if query_history else []
+                return response, search_output_text
             
             # Set up events
             save_key_button.click(
@@ -1100,8 +448,8 @@ def main():
             
             query_button.click(
                 fn=query_and_search,
-                inputs=[query_input, num_results, model_choice, temperature_slider, max_tokens_slider],
-                outputs=[response_output, search_output, history_list]
+                inputs=[query_input, num_results, model_choice],
+                outputs=[response_output, search_output]
             )
             
             refresh_button.click(
@@ -1109,84 +457,24 @@ def main():
                 inputs=None,
                 outputs=stats_display
             )
-            
-            clear_button.click(
-                fn=clear_inputs,
-                inputs=None,
-                outputs=[query_input, response_output, search_output, history_list]
-            )
-            
-            # Handle Enter key in query input
-            query_input.submit(
-                fn=query_and_search,
-                inputs=[query_input, num_results, model_choice, temperature_slider, max_tokens_slider],
-                outputs=[response_output, search_output, history_list]
-            )
-            
-            # Auto-fill examples
-            examples = [
-                ["What are the main features of this application?"],
-                ["How does the retrieval augmented generation work?"],
-                ["Can you explain the embedding models used in this system?"],
-            ]
-            
-            gr.Examples(
-                examples=examples,
-                inputs=query_input,
-                outputs=[response_output, search_output, history_list],
-                fn=lambda q: query_and_search(q, num_results.value, model_choice.value, temperature_slider.value, max_tokens_slider.value),
-                cache_examples=False,
-            )
         
-        # Launch the interface with a nice theme
-        app.launch(
-            share=False,  # Set to True to create a public link
-            server_name="0.0.0.0",  # Listen on all interfaces
-            server_port=7860,  # Default Gradio port
-            debug=False,  # Set to True during development
-            auth=None,  # Add (username, password) tuple for basic auth
-            favicon_path="favicon.ico" if os.path.exists("favicon.ico") else None,
-            show_error=True
-        )
-    
+        # Launch the interface
+        app.launch()
+        
     except Exception as e:
-        logger.critical(f"Error starting application: {e}")
-        print(f"Error starting application: {e}")
+        logger.error(f"Error initializing application: {e}")
+        print(f"Error: {e}")
         sys.exit(1)
 
+# Helper function to get database stats
+def get_db_stats(vector_store):
+    """Function to get vector store statistics"""
+    try:
+        stats = vector_store.get_statistics()
+        return f"Total documents: {stats.get('total_documents', 0)}"
+    except Exception as e:
+        logger.error(f"Error getting statistics: {e}")
+        return "Error getting database statistics"
+
 if __name__ == "__main__":
-    # Parse command line arguments
-    if len(sys.argv) > 1:
-        if sys.argv[1] == "--load" and len(sys.argv) > 2:
-            # Load documents mode
-            print(f"Document Knowledge Assistant v{VERSION}")
-            print(f"Loading documents into knowledge base...")
-            
-            success_count = 0
-            failed_count = 0
-            
-            for file_path in sys.argv[2:]:
-                if os.path.exists(file_path):
-                    success = load_document(file_path)
-                    if success:
-                        success_count += 1
-                        print(f"✅ Successfully loaded: {file_path}")
-                    else:
-                        failed_count += 1
-                        print(f"❌ Failed to load: {file_path}")
-                else:
-                    failed_count += 1
-                    print(f"❌ File not found: {file_path}")
-            
-            print(f"\nLoading complete: {success_count} documents loaded, {failed_count} failed")
-            sys.exit(0)
-        elif sys.argv[1] == "--help":
-            print(f"Document Knowledge Assistant v{VERSION}")
-            print("Usage:")
-            print("  python rag_system.py                    # Start the web UI")
-            print("  python rag_system.py --load file1 file2 # Load documents into the knowledge base")
-            print("  python rag_system.py --help             # Show this help message")
-            sys.exit(0)
-    
-    # Start the web UI
     main()