From 60409395b4d598d562c8f1efbcf72100d295ec3f Mon Sep 17 00:00:00 2001 From: ajaysi Date: Wed, 16 Apr 2025 19:24:05 +0530 Subject: [PATCH] Revert changes in lib/utils/api_key_manager/components/ai_providers_setup.py to the state before commit 036fde9. --- .../components/ai_providers_setup.py | 276 ++++++------------ 1 file changed, 86 insertions(+), 190 deletions(-) diff --git a/lib/utils/api_key_manager/components/ai_providers_setup.py b/lib/utils/api_key_manager/components/ai_providers_setup.py index 198d7011..1ab045f6 100644 --- a/lib/utils/api_key_manager/components/ai_providers_setup.py +++ b/lib/utils/api_key_manager/components/ai_providers_setup.py @@ -5,214 +5,110 @@ from loguru import logger import streamlit as st import os import sys -# Corrected import: Assuming validation functions are in validation.py in the parent directory -from ..validation import ( - test_openai_api_key, - test_gemini_api_key, - # test_anthropic_api_key, # Keep commented if not used or add if needed - # test_deepseek_api_key, # Keep commented if not used or add if needed - test_mistral_api_key -) - -# Helper function to validate a specific provider's key -def _validate_provider_key(provider_name: str, key_value: str) -> bool: - """Validate the API key for a given provider.""" - if not key_value: - logger.debug(f"Validation: Key for {provider_name} is empty.") - return False - try: - logger.debug(f"Validating key for {provider_name}...") - if provider_name == "openai": - # Ensure the function exists in validation.py - if callable(getattr(sys.modules[__name__], 'test_openai_api_key', None)): - is_valid = test_openai_api_key(key_value) - else: - logger.error("test_openai_api_key not found in validation module") - is_valid = False # Assume invalid if test func missing - elif provider_name == "gemini": - if callable(getattr(sys.modules[__name__], 'test_gemini_api_key', None)): - is_valid = test_gemini_api_key(key_value) - else: - logger.error("test_gemini_api_key not found in validation module") - is_valid = False - elif provider_name == "mistral": - if callable(getattr(sys.modules[__name__], 'test_mistral_api_key', None)): - is_valid = test_mistral_api_key(key_value) - else: - logger.error("test_mistral_api_key not found in validation module") - is_valid = False - else: - logger.warning(f"Validation not implemented for provider: {provider_name}") - return False # Or True if unknown providers are allowed without validation - - logger.info(f"Validation result for {provider_name}: {'Valid' if is_valid else 'Invalid'}") - return is_valid - except Exception as e: - logger.error(f"Error validating key for {provider_name}: {e}", exc_info=True) - return False - -# Callback function for handling API key input changes -def _handle_api_key_change(provider_name: str, api_key_manager): - """Save and validate API key when input changes.""" - key_input_widget_key = f"{provider_name}_key_input" - status_widget_key = f"{provider_name}_status" - - # Check if the input widget key exists in session state - if key_input_widget_key not in st.session_state: - logger.warning(f"Input widget key '{key_input_widget_key}' not found in session state.") - return - - key_value = st.session_state[key_input_widget_key] - current_status = st.session_state.get(status_widget_key) - - logger.debug(f"Handling change for {provider_name}. Key: {'***' if key_value else 'Empty'}. Current status: {current_status}") - - # If key is empty, reset status - if not key_value: - api_key_manager.save_api_key(provider_name, "") # Ensure empty key is saved - st.session_state[status_widget_key] = "unsaved" - logger.info(f"Cleared API key for {provider_name}.") - return - - # Set status to saving/validating - st.session_state[status_widget_key] = "saving" - st.rerun() # Rerun to show the spinner immediately - - try: - # Save the key using the manager - logger.debug(f"Saving key for {provider_name}...") - api_key_manager.save_api_key(provider_name, key_value) - logger.info(f"Saved API key for {provider_name}.") - - # Validate the key - is_valid = _validate_provider_key(provider_name, key_value) - - # Update status based on validation result - if is_valid: - st.session_state[status_widget_key] = "valid" - else: - st.session_state[status_widget_key] = "invalid" - - except Exception as e: - logger.error(f"Error during saving/validation for {provider_name}: {e}", exc_info=True) - st.session_state[status_widget_key] = "error" def render_ai_providers_setup(api_key_manager) -> Dict[str, Any]: """ - Render the AI providers setup component with immediate feedback. + Render the AI providers setup component. Args: api_key_manager: API key manager instance Returns: - Dict[str, Any]: Component state (not directly used here, handled by state manager) + Dict[str, Any]: Component state """ try: logger.info("[render_ai_providers_setup] Rendering AI providers setup") - # Initialize status in session state if not present - for provider in ["openai", "gemini", "mistral"]: - status_key = f"{provider}_status" - if status_key not in st.session_state: - # Check if a key exists and try to validate it on first load - existing_key = api_key_manager.get_api_key(provider) - if existing_key: - if _validate_provider_key(provider, existing_key): - st.session_state[status_key] = "valid" - else: - # Keep it unsaved/invalid on load if pre-existing key is bad - # Or maybe set to invalid? Let's choose unsaved for now. - st.session_state[status_key] = "invalid" - else: - st.session_state[status_key] = "unsaved" - # Display section header - st.header("Step 1: Configure AI Providers") + st.header("Step 1: Select AI Providers") st.markdown(""" Configure your AI providers to enable advanced content generation capabilities. - Enter your API keys below. They will be validated automatically. + Choose and set up the AI services you want to use. """) - # --- OpenAI --- - st.subheader("OpenAI (Required)") - st.markdown("Get your API key from: [OpenAI Dashboard](https://platform.openai.com/account/api-keys)") - openai_key = api_key_manager.get_api_key("openai") - st.text_input( - "OpenAI API Key", - value=openai_key if openai_key else "", - type="password", - key="openai_key_input", - on_change=_handle_api_key_change, - args=("openai", api_key_manager) - ) - # Feedback Area for OpenAI - openai_status = st.session_state.get("openai_status", "unsaved") - feedback_placeholder_openai = st.empty() - if openai_status == "saving": - feedback_placeholder_openai.info("Validating OpenAI key...", icon="⏳") - elif openai_status == "valid": - feedback_placeholder_openai.success("OpenAI key saved and valid!", icon="✅") - elif openai_status == "invalid": - feedback_placeholder_openai.error("Invalid OpenAI key. Please check and try again.", icon="❌") - elif openai_status == "error": - feedback_placeholder_openai.error("Error saving/validating OpenAI key.", icon="⚠️") + # Create columns for different providers + col1, col2 = st.columns(2) - # --- Google Gemini --- - st.subheader("Google Gemini (Required)") - st.markdown("Get your API key from: [Google AI Studio](https://makersuite.google.com/app/apikey)") - gemini_key = api_key_manager.get_api_key("gemini") - st.text_input( - "Gemini API Key", - value=gemini_key if gemini_key else "", - type="password", - key="gemini_key_input", - on_change=_handle_api_key_change, - args=("gemini", api_key_manager) - ) - # Feedback Area for Gemini - gemini_status = st.session_state.get("gemini_status", "unsaved") - feedback_placeholder_gemini = st.empty() - if gemini_status == "saving": - feedback_placeholder_gemini.info("Validating Gemini key...", icon="⏳") - elif gemini_status == "valid": - feedback_placeholder_gemini.success("Gemini key saved and valid!", icon="✅") - elif gemini_status == "invalid": - feedback_placeholder_gemini.error("Invalid Gemini key. Please check and try again.", icon="❌") - elif gemini_status == "error": - feedback_placeholder_gemini.error("Error saving/validating Gemini key.", icon="⚠️") - - # --- Mistral AI (Optional) --- - st.subheader("Mistral AI (Optional)") - st.markdown("Get your API key from: [Mistral Platform](https://console.mistral.ai/api-keys/)") - mistral_key = api_key_manager.get_api_key("mistral") - st.text_input( - "Mistral API Key", - value=mistral_key if mistral_key else "", - type="password", - key="mistral_key_input", - on_change=_handle_api_key_change, - args=("mistral", api_key_manager) - ) - # Feedback Area for Mistral - mistral_status = st.session_state.get("mistral_status", "unsaved") - feedback_placeholder_mistral = st.empty() - if mistral_status == "saving": - feedback_placeholder_mistral.info("Validating Mistral key...", icon="⏳") - elif mistral_status == "valid": - feedback_placeholder_mistral.success("Mistral key saved and valid!", icon="✅") - elif mistral_status == "invalid": - feedback_placeholder_mistral.error("Invalid Mistral key. Please check and try again.", icon="❌") - elif mistral_status == "error": - feedback_placeholder_mistral.error("Error saving/validating Mistral key.", icon="⚠️") - - # --- Final Notes --- - st.info("Note: At least one AI provider (OpenAI or Google Gemini) must have a valid API key to proceed.") + with col1: + st.subheader("OpenAI") + st.markdown(""" + OpenAI's GPT models provide powerful natural language processing capabilities. + + Get your API key from: [OpenAI Dashboard](https://platform.openai.com/account/api-keys) + """) + + openai_key = api_key_manager.get_api_key("openai") + openai_input = st.text_input( + "OpenAI API Key", + value=openai_key if openai_key else "", + type="password", + key="openai_key_input" + ) + + with col2: + st.subheader("Google Gemini") + st.markdown(""" + Google's Gemini models offer advanced AI capabilities. + + Get your API key from: [Google AI Studio](https://makersuite.google.com/app/apikey) + """) + + gemini_key = api_key_manager.get_api_key("gemini") + gemini_input = st.text_input( + "Gemini API Key", + value=gemini_key if gemini_key else "", + type="password", + key="gemini_key_input" + ) + + # Optional AI Provider + st.subheader("Additional AI Provider (Optional)") + col1, col2 = st.columns(2) + + with col1: + st.markdown(""" + Mistral AI provides an alternative model for content generation. + + Get your API key from: [Mistral Platform](https://console.mistral.ai/api-keys/) + """) + + mistral_key = api_key_manager.get_api_key("mistral") + mistral_input = st.text_input( + "Mistral API Key (Optional)", + value=mistral_key if mistral_key else "", + type="password", + key="mistral_key_input" + ) + + # Add a note about saving + st.info(""" + Note: At least one AI provider (OpenAI or Google Gemini) is required. + Click Continue to save your keys and proceed. + """) + + # Save keys if they've changed when proceeding to next step + if st.session_state.get('wizard_current_step', 1) > 1: + if openai_input != openai_key: + api_key_manager.save_api_key("openai", openai_input) + logger.info("[render_ai_providers_setup] OpenAI API key saved") + + if gemini_input != gemini_key: + api_key_manager.save_api_key("gemini", gemini_input) + logger.info("[render_ai_providers_setup] Gemini API key saved") + + if mistral_input != mistral_key: + api_key_manager.save_api_key("mistral", mistral_input) + logger.info("[render_ai_providers_setup] Mistral API key saved") + + # Validate that at least one required provider is configured + if not (openai_input or gemini_input): + st.error("Please configure at least one AI provider (OpenAI or Google Gemini) to proceed.") + return {"current_step": 1, "can_proceed": False} + + return {"current_step": 1, "can_proceed": bool(openai_input or gemini_input)} - # Return value is not strictly needed if navigation relies on session state status - return {} - except Exception as e: - error_msg = f"Error rendering AI providers setup: {str(e)}" - logger.error(f"[render_ai_providers_setup] {error_msg}", exc_info=True) + error_msg = f"Error in AI providers setup: {str(e)}" + logger.error(f"[render_ai_providers_setup] {error_msg}") st.error(error_msg) - return {"error": error_msg} + return {"current_step": 1, "error": error_msg}