Update app.py

app.py

@@ -1,384 +1,467 @@
-import base64
-import io
 import os
-import gradio as gr
-import numpy as np
-import matplotlib.pyplot as plt
-from typing import Dict, List, Tuple, Any
-import json
-from litellm import completion
+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)
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-CONVERSATION_PROMPT = """
-You are an engaging and insightful career advisor. Have natural conversations to learn about their career.
-Use an enthusiastic, supportive tone and show genuine interest in their journey.
-CONVERSATION STYLE:
-- Be warm and engaging
-- Show genuine interest in their experiences
-- Ask specific follow-up questions about details they mention
-- Keep the conversation flowing naturally
-- Use conversational language, not formal queries
-- Express enthusiasm about their achievements
-- Dig deeper into interesting points they make
-INFORMATION TO GATHER (through natural conversation):
-1. Current Role Details:
-   - Job title and responsibilities
-   - Company size and industry
-   - Team size and structure
-   - Project scope and impact
-   - Current compensation (base, bonus, equity)
-2. Experience Deep-Dive:
-   - Career progression story
-   - Leadership experience
-   - Major projects and achievements
-   - Technical skills and expertise
-   - Industry knowledge
-3. Educational Background:
-   - Degrees and certifications
-   - Specialized training
-   - Continuous learning
-4. Work Environment:
-   - Location and market
-   - Remote/hybrid setup
-   - Growth opportunities
-   - Company culture
-CONVERSATION FLOW:
-1. Start with: "Hi! I'd love to hear about your career journey. What kind of work are you doing currently?"
-2. After each response:
-   - Pick up on specific details they mentioned
-   - Ask engaging follow-up questions
-   - Show genuine interest in their experiences
-   - Build on previous information shared
-3. If they mention something interesting, probe deeper:
-   - "That project sounds fascinating! What were some unique challenges you faced?"
-   - "Leading a team must be exciting! How did you approach building and motivating your team?"
-   - "Interesting technology stack! What made you choose those specific tools?"
-4. When compensation is mentioned:
-   - Be tactful and professional
-   - Acknowledge their goals
-   - Ask about their desired growth
-5. Once you have enough information, say:
-   "I've got a good understanding of your career profile now! Would you like to see your personalized salary growth projection? Just click 'Generate Analysis' and I'll create a detailed forecast based on our discussion."
-IMPORTANT:
-- Keep conversation flowing naturally
-- Don't rush to collect information
-- Show genuine interest in their story
-- Ask insightful follow-up questions
-- Build rapport through discussion
-"""
+# 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
 
-EXTRACTION_PROMPT = """
-Analyze the conversation and extract numerical scores from 0 to 1 based on salary growth potential.
-SCORING GUIDELINES:
-1. Industry Score (0-1):
-   Industry Type & Growth:
-   - 1.0: Cutting-edge AI/ML companies
-   - 0.9: High-growth tech (cloud, cybersecurity)
-   - 0.8: Established tech companies
-   - 0.7: Finance/Healthcare tech
-   - 0.6: Traditional tech sectors
-   - 0.5: Non-tech industries
-   Company Position:
-   +0.1: Market leader
-   +0.1: High growth trajectory
-   -0.1: Declining market position
-2. Experience Score (0-1):
-   Years and Level:
-   - 1.0: 15+ years with executive experience
-   - 0.9: 10-15 years, senior leadership
-   - 0.8: 7-10 years, team leadership
-   - 0.7: 4-6 years, senior individual
-   - 0.6: 2-3 years, mid-level
-   - 0.5: 0-1 years, entry-level
-   Quality Indicators:
-   +0.1: Rapid promotions
-   +0.1: Significant achievements
-   +0.1: High-impact projects
-3. Education Score (0-1):
-   Formal Education:
-   - 1.0: PhD from top institution
-   - 0.9: Masters from top institution
-   - 0.8: Bachelors from top institution
-   - 0.7: Advanced degree
-   - 0.6: Bachelors degree
-   - 0.5: Other education
-   Additional Factors:
-   +0.1: Relevant certifications
-   +0.1: Continuous learning
-   +0.1: Field-specific expertise
-4. Skills Score (0-1):
-   Technical Depth:
-   - 1.0: Industry-leading expertise
-   - 0.9: Advanced technical leadership
-   - 0.8: Strong technical + leadership
-   - 0.7: Solid technical skills
-   - 0.6: Growing technical skills
-   - 0.5: Basic skill set
-   Breadth and Application:
-   +0.1: Multiple in-demand skills
-   +0.1: Proven implementation
-   +0.1: Cross-functional expertise
-5. Location Score (0-1):
-   Market Strength:
-   - 1.0: Major tech hubs (SF, NYC)
-   - 0.9: Growing tech hubs
-   - 0.8: Major cities
-   - 0.7: Regional tech centers
-   - 0.6: Smaller markets
-   - 0.5: Remote locations
-   Flexibility:
-   +0.1: Remote work option
-   +0.1: High-growth market
-   +0.1: Strategic location
-Return a JSON object with exactly these fields:
-{
-    "industry_score": float,
-    "experience_score": float,
-    "education_score": float,
-    "skills_score": float,
-    "location_score": float,
-    "current_salary": float
-}
-Base scores on available information. Make reasonable assumptions for missing data based on context clues.
-"""
+# 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)
 
-class CodeEnvironment:
-    """Environment for executing visualization code"""
-    
-    def __init__(self):
-        self.globals = {'np': np, 'plt': plt}
-        self.locals = {}
+# Embedding engine
+class EmbeddingEngine:
+    """Handle embeddings with a lightweight model"""
     
-    def execute(self, code: str, paths: np.ndarray = None) -> Dict[str, Any]:
-        """Execute visualization code and return results"""
-        if paths is not None:
-            self.globals['paths'] = paths
+    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
             
-        result = {'figures': [], 'error': None}
+        # 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:
-            exec(code, self.globals, self.locals)
-            buf = io.BytesIO()
-            plt.gcf().savefig(buf, format='png', dpi=300, bbox_inches='tight')
-            buf.seek(0)
-            result['figures'].append(buf.getvalue())
-            plt.close('all')
+            embeddings = self.model.encode(truncated_texts, batch_size=batch_size, 
+                                         show_progress_bar=False, convert_to_numpy=True)
+            return embeddings
         except Exception as e:
-            result['error'] = f"Visualization failed: {str(e)}"
-            plt.close('all')
-        return result
+            logger.error(f"Error generating embeddings: {e}")
+            # Return zero embeddings as fallback
+            return np.zeros((len(truncated_texts), self.vector_size))
 
-class SalarySimulator:
-    """Monte Carlo simulation for salary projections"""
+class VectorStoreManager:
+    """Manage Chroma vector store operations - upload, query, etc."""
     
-    def __init__(self, years: int = 5, num_paths: int = 1000):
-        self.years = years
-        self.num_paths = num_paths
-        
-    def run_simulation(self, profile: Dict[str, float]) -> np.ndarray:
-        """Generate salary growth paths using Monte Carlo simulation"""
-        paths = np.zeros((self.num_paths, self.years + 1))
-        paths[:, 0] = profile['current_salary']
+    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)
         
-        base_growth = 0.02 + (profile['industry_score'] * 0.04)
-        skill_premium = 0.01 + (profile['skills_score'] * 0.02)
-        exp_premium = 0.01 + (profile['experience_score'] * 0.02)
-        edu_premium = 0.005 + (profile['education_score'] * 0.015)
-        location_premium = 0.01 + (profile['location_score'] * 0.02)
+        # 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 = {}
         
-        volatility = 0.05 + (profile['industry_score'] * 0.05)
-        disruption_chance = 0.1
-        disruption_impact = 0.2
+        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)
         
-        for path in range(self.num_paths):
-            salary = paths[path, 0]
-            for year in range(1, self.years + 1):
-                growth = base_growth + skill_premium + exp_premium + edu_premium + location_premium
-                growth += np.random.normal(0, volatility)
-                if np.random.random() < disruption_chance:
-                    impact = disruption_impact * np.random.random()
-                    growth += impact if np.random.random() < 0.7 else -impact
-                growth = max(min(growth, 0.25), -0.1)
-                salary *= (1 + growth)
-                paths[path, year] = salary
-        return paths
+        return stats
 
-class CareerAdvisor:
-    """Main career advisor system"""
-    
-    def __init__(self, years: int = 5, num_paths: int = 1000):
-        self.chat_history = []
-        self.simulator = SalarySimulator(years, num_paths)
-        self.code_env = CodeEnvironment()
+class RAGSystem:
+    """Retrieval-Augmented Generation with multiple LLM providers"""
     
-    def reset(self):
-        """Reset conversation state"""
-        self.chat_history = []
+    def __init__(self, vector_store: VectorStoreManager):
+        self.vector_store = vector_store
+        self.openai_client = None
+        self.gemini_configured = False
     
-    def chat(self, message: str, api_key: str) -> str:
-        """Process user message and generate response"""
-        if not api_key.strip().startswith("sk-"):
-            return "Please enter a valid OpenAI API key starting with 'sk-'."
+    def setup_openai(self, api_key: str):
+        """Set up OpenAI client with API key"""
         try:
-            messages = [{"role": "system", "content": CONVERSATION_PROMPT}] + \
-                       self.chat_history + [{"role": "user", "content": message}]
-            response = completion(model="gpt-4o-mini", messages=messages, api_key=api_key)
-            self.chat_history.extend([
-                {"role": "user", "content": message},
-                {"role": "assistant", "content": response.choices[0].message.content}
-            ])
-            return response.choices[0].message.content
+            self.openai_client = OpenAI(api_key=api_key)
+            return True
         except Exception as e:
-            return f"Chat error: {str(e)}. Please check your API key or try again."
+            logger.error(f"Error initializing OpenAI client: {e}")
+            return False
     
-    def generate_analysis(self, api_key: str) -> Tuple[str, bytes]:
-        """Generate complete analysis with visualization"""
-        if not self.chat_history:
-            return "Please chat about your career first to generate an analysis.", None
+    def setup_gemini(self, api_key: str):
+        """Set up Gemini with API key"""
         try:
-            profile = self._extract_profile(api_key)
-            paths = self.simulator.run_simulation(profile)
-            
-            viz_code = """
-import matplotlib.pyplot as plt
-import numpy as np
-plt.style.use('dark_background')
-fig = plt.figure(figsize=(12, 16))
-ax1 = plt.subplot2grid((2, 1), (0, 0))
-for path in paths[::20]:
-    ax1.plot(range(paths.shape[1]), path, color='#4a90e2', alpha=0.1, linewidth=1)
-percentiles = [10, 25, 50, 75, 90]
-colors = ['#ff9999', '#ffcc99', '#ffffff', '#ffcc99', '#ff9999']
-labels = ['10th', '25th', 'Median', '75th', '90th']
-for p, color, label in zip(percentiles, colors, labels):
-    line = np.percentile(paths, p, axis=0)
-    ax1.plot(range(paths.shape[1]), line, color=color, linewidth=2, label=f'{label} percentile')
-ax1.set_title('Salary Growth Projections\n', fontsize=16, pad=20)
-ax1.set_xlabel('Years', fontsize=12)
-ax1.set_ylabel('Salary ($)', fontsize=12)
-ax1.grid(True, alpha=0.2)
-ax1.legend(fontsize=10)
-ax1.yaxis.set_major_formatter(plt.FuncFormatter(lambda x, p: f'${x:,.0f}'))
-ax1.set_xticks(range(paths.shape[1]))
-ax1.set_xticklabels(['Current'] + [f'Year {i+1}' for i in range(paths.shape[1]-1)])
-ax2 = plt.subplot2grid((2, 1), (1, 0))
-final_salaries = paths[:, -1]
-ax2.hist(final_salaries, bins=50, color='#4a90e2', alpha=0.7)
-ax2.set_title('Final Salary Distribution\n', fontsize=16, pad=20)
-ax2.set_xlabel('Salary ($)', fontsize=12)
-ax2.set_ylabel('Frequency', fontsize=12)
-ax2.grid(True, alpha=0.2)
-ax2.xaxis.set_major_formatter(plt.FuncFormatter(lambda x, p: f'${x:,.0f}'))
-for p, color in zip(percentiles, colors):
-    value = np.percentile(final_salaries, p)
-    ax2.axvline(x=value, color=color, linestyle='--', alpha=0.5)
-plt.tight_layout(pad=4)
-"""
-            viz_result = self.code_env.execute(viz_code, paths)
-            if viz_result['error']:
-                return f"Analysis generated, but {viz_result['error']}", None
-            summary = self._generate_summary(profile, paths)
-            return summary, viz_result['figures'][0]
+            genai.configure(api_key=api_key)
+            self.gemini_configured = True
+            return True
         except Exception as e:
-            return f"Analysis error: {str(e)}. Please ensure sufficient chat history.", None
+            logger.error(f"Error configuring Gemini: {e}")
+            return False
     
-    def _extract_profile(self, api_key: str) -> Dict[str, float]:
-        """Extract profile scores from conversation"""
-        conversation = "\n".join([f"{msg['role']}: {msg['content']}" for msg in self.chat_history])
-        messages = [
-            {"role": "system", "content": EXTRACTION_PROMPT},
-            {"role": "user", "content": f"Extract profile from:\n{conversation}"}
-        ]
-        response = completion(
-            model="gpt-4o-mini",
-            messages=messages,
-            api_key=api_key,
-            response_format={"type": "json_object"}
-        )
-        return json.loads(response.choices[0].message.content)
+    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_summary(self, profile: Dict[str, float], paths: np.ndarray) -> str:
-        """Generate analysis summary"""
-        final_salaries = paths[:, -1]
-        initial_salary = paths[0, 0]
-        cagr = (np.median(final_salaries) / initial_salary) ** (1/self.simulator.years) - 1
+    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}
         
-        return f"""
-        Career Profile Analysis
-        ======================
-        Current Situation:
-        • Salary: ${profile['current_salary']:,.2f}
-        • Industry Position: {profile['industry_score']:.2f}/1.0
-        • Experience Level: {profile['experience_score']:.2f}/1.0
-        • Education Rating: {profile['education_score']:.2f}/1.0
-        • Skills Assessment: {profile['skills_score']:.2f}/1.0
-        • Location Impact: {profile['location_score']:.2f}/1.0
-        {self.simulator.years}-Year Projection:
-        • Conservative (25th percentile): ${np.percentile(final_salaries, 25):,.2f}
-        • Most Likely (Median): ${np.percentile(final_salaries, 50):,.2f}
-        • Optimistic (75th percentile): ${np.percentile(final_salaries, 75):,.2f}
-        • Expected Annual Growth: {cagr*100:.1f}%
-        Key Insights:
-        • Your profile suggests {cagr*100:.1f}% annual growth potential
-        • {profile['industry_score']:.2f} industry score indicates {'strong' if profile['industry_score'] > 0.7 else 'moderate' if profile['industry_score'] > 0.5 else 'challenging'} growth environment
-        • Skills rating of {profile['skills_score']:.2f} suggests {'excellent' if profile['skills_score'] > 0.7 else 'good' if profile['skills_score'] > 0.5 else 'potential for'} career advancement
-        • Location score {profile['location_score']:.2f} {'enhances' if profile['location_score'] > 0.7 else 'supports' if profile['location_score'] > 0.5 else 'may limit'} opportunities
-        Based on {self.simulator.num_paths:,} simulated career paths
+        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)
 
-def create_interface():
-    """Create Gradio interface with configurable simulation parameters"""
-    advisor = None
+# 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)
     
-    def init_advisor(years: int, num_paths: int):
-        nonlocal advisor
-        advisor = CareerAdvisor(years=max(1, years), num_paths=max(100, num_paths))
-        advisor.reset()
+    # Define Gradio app
+    def rag_chat_wrapper(query, n_results, model_choice):
+        return rag_chat(query, n_results, model_choice, rag_system)
     
-    def user_message(message: str, history: List, api_key: str) -> Tuple[str, List]:
-        if not message.strip():
-            return "", history
-        if not advisor:
-            return "Please set simulation parameters first.", history
-        response = advisor.chat(message, api_key)
-        return "", history + [(message, response)]
+    def simple_query_wrapper(query, n_results):
+        return simple_query(query, n_results, vector_store)
     
-    def generate_analysis(api_key: str, history: List) -> Tuple[str, gr.Image]:
-        if not advisor or not history:
-            return "Please chat about your career and set parameters first.", None
-        summary, figure_data = advisor.generate_analysis(api_key)
-        return summary, figure_data if figure_data else None
+    def update_api_keys_wrapper(openai_key, gemini_key):
+        return update_api_keys(openai_key, gemini_key, rag_system)
     
-    with gr.Blocks(title="Monte Carlo Salary Prediction", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# 💰 Monte Carlo Simulation of Salary Prediction\nChat about your career to see your growth potential!")
+    # Create the Gradio interface
+    with gr.Blocks(title="Markdown RAG System") as app:
+        gr.Markdown("# RAG System with Multiple LLM Providers")
         
-        with gr.Row():
-            api_key = gr.Textbox(label="OpenAI API Key", placeholder="Enter your OpenAI API key", type="password")
-            years = gr.Number(label="Simulation Years", value=5, minimum=1, step=1)
-            num_paths = gr.Number(label="Number of Paths", value=1000, minimum=100, step=100)
+        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")
         
-        chatbot = gr.Chatbot(value=[], height=400, show_copy_button=True)
+        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("")
         
-        with gr.Row():
-            msg = gr.Textbox(label="Your message", placeholder="Tell me about your career...", lines=2)
-            send = gr.Button("Send", variant="primary", scale=0)
+        # Set up events
+        query_button.click(
+            fn=rag_chat_wrapper,
+            inputs=[query_input, num_results, model_choice],
+            outputs=response_output
+        )
         
-        analyze = gr.Button("Generate Analysis", variant="secondary", size="lg")
+        refresh_button.click(
+            fn=lambda: get_db_stats(vector_store),
+            inputs=None,
+            outputs=stats_display
+        )
         
-        with gr.Row():
-            analysis = gr.Textbox(label="Analysis", lines=10, show_copy_button=True)
-            plot = gr.Image(label="Projections", show_download_button=True, height=600)
+        search_button.click(
+            fn=simple_query_wrapper,
+            inputs=[search_input, search_num],
+            outputs=search_output
+        )
         
-        demo.load(lambda y, n: init_advisor(y, n), inputs=[years, num_paths], outputs=None)
-        msg.submit(user_message, inputs=[msg, chatbot, api_key], outputs=[msg, chatbot])
-        send.click(user_message, inputs=[msg, chatbot, api_key], outputs=[msg, chatbot])
-        analyze.click(generate_analysis, inputs=[api_key, chatbot], outputs=[analysis, plot])
+        save_keys_button.click(
+            fn=update_api_keys_wrapper,
+            inputs=[openai_key_input, gemini_key_input],
+            outputs=api_status
+        )
     
-    return demo
+    # Launch the interface
+    app.launch()
 
 if __name__ == "__main__":
-    demo = create_interface()
-    demo.launch()
+    main()