Revert changes in lib/utils/api_key_manager/components/ai_providers_setup.py to the state before commit 036fde9.
This commit is contained in:
ajaysi
@@ -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}
|
||||
|
||||
Reference in New Issue
Block a user