Update app.py

app.py

@@ -768,103 +768,20 @@ def get_db_stats(vector_store: VectorStoreManager) -> str:
 # The error occurs because there's likely a `try` statement without a matching `except` or `finally`
 # Here are the possible locations and fixes:
 
-# Option 1: If it's in the custom CSS definition around that line number
-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;
-    }
-"""  # Make sure this closing triple quote is present
 
-# Option 2: If it's in the main function, make sure all try blocks have matching except clauses
-def main():
-    """Main function to run the RAG application"""
-    # Path for configuration file
-    CONFIG_FILE_PATH = "rag_config.json"
-    
-    # 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}")
-    
-    try:
-        # Initialize vector store manager with existing collection
-        vector_store = VectorStoreManager(config)
-        
-        # Initialize RAG system without API keys initially
-        rag_system = RAGSystem(vector_store, config)
-        
-        # Create the Gradio interface with custom CSS
-        with gr.Blocks(title="Document Knowledge Assistant", css=custom_css) as app:
-            # Interface code here...
-            pass  # Replace with actual UI code
-            
-        # Launch the interface
-        app.launch(
-            share=False,
-            server_name="0.0.0.0",
-            server_port=7860,
-            debug=False
-        )
-            
-    except Exception as e:  # Make sure there's an except clause for every try
-        logger.critical(f"Error starting application: {e}")
-        print(f"Error starting application: {e}")
-        sys.exit(1)
 
-# Option 3: If it's in an Example or Markdown section, make sure triple quotes are properly closed
-
-# Option 4: Complete fix for the most likely scenario - an incomplete try block in the main function
-def main_fixed():
+def main():
     """Main function to run the RAG application"""
     # Path for configuration file
     CONFIG_FILE_PATH = "rag_config.json"
     
-    # Try to load configuration from file, or use defaults
     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",
+                local_dir="./chroma_db",  # Store Chroma files in dedicated directory
                 collection_name="markdown_docs"
             )
             # Save default configuration
@@ -879,17 +796,244 @@ def main_fixed():
         # Initialize RAG system without API keys initially
         rag_system = RAGSystem(vector_store, config)
         
-        # Create the Gradio interface
+        # Create the Gradio interface with custom CSS
         with gr.Blocks(title="Document Knowledge Assistant", css=custom_css) as app:
-            # UI code here
-            pass  # Replace with actual UI code
+            gr.Markdown(f"# Document Knowledge Assistant v{VERSION}")
+            gr.Markdown("Ask questions about your documents and get comprehensive AI-powered answers")
+            
+            # 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("")
+                    
+                    # 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"
+                        )
+                    
+                    # 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
+                )
+                
+            # 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 = []
+            
+            # 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}"
+                else:
+                    success = rag_system.setup_gemini(api_key)
+                    model_name = f"Google {config.gemini_model}"
+                
+                if success:
+                    return f"✅ {model_name} connected successfully"
+                else:
+                    return f"❌ Connection failed. Please check your API key and try again."
+            
+            # 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)
+                
+                start_time = datetime.now()
+                
+                if not query.strip():
+                    return "Please enter a question to get an answer.", "", query_history[-5:] if query_history else []
+                
+                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
+                
+                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 []
+            
+            # Set up events
+            save_key_button.click(
+                fn=update_api_key,
+                inputs=[api_key_input, model_choice],
+                outputs=api_status
+            )
+            
+            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]
+            )
+            
+            refresh_button.click(
+                fn=lambda: get_db_stats(vector_store),
+                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
+        # Launch the interface with a nice theme
         app.launch(
-            share=False,
-            server_name="0.0.0.0",
-            server_port=7860,
-            debug=False
+            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
         )
     except Exception as e:
         logger.critical(f"Error starting application: {e}")