feat: better API key error handling

- replaced API key error handling with custom exception
This commit is contained in:
marginal23326
2025-03-18 08:04:07 +06:00
parent e8ef2342c7
commit 10fdfce47f
2 changed files with 12 additions and 15 deletions

View File

@@ -10,7 +10,6 @@ from langchain_mistralai import ChatMistralAI
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_ollama import ChatOllama
from langchain_openai import AzureChatOpenAI, ChatOpenAI
import gradio as gr
from .llm import DeepSeekR1ChatOpenAI, DeepSeekR1ChatOllama
@@ -36,7 +35,7 @@ def get_llm_model(provider: str, **kwargs):
env_var = f"{provider.upper()}_API_KEY"
api_key = kwargs.get("api_key", "") or os.getenv(env_var, "")
if not api_key:
handle_api_key_error(provider, env_var)
raise MissingAPIKeyError(provider, env_var)
kwargs["api_key"] = api_key
if provider == "anthropic":
@@ -184,6 +183,7 @@ def update_model_dropdown(llm_provider, api_key=None, base_url=None):
"""
Update the model name dropdown with predefined models for the selected provider.
"""
import gradio as gr
# Use API keys from .env if not provided
if not api_key:
api_key = os.getenv(f"{llm_provider.upper()}_API_KEY", "")
@@ -196,16 +196,12 @@ def update_model_dropdown(llm_provider, api_key=None, base_url=None):
else:
return gr.Dropdown(choices=[], value="", interactive=True, allow_custom_value=True)
def handle_api_key_error(provider: str, env_var: str):
"""
Handles the missing API key error by raising a gr.Error with a clear message.
"""
provider_display = PROVIDER_DISPLAY_NAMES.get(provider, provider.upper())
raise gr.Error(
f"💥 {provider_display} API key not found! 🔑 Please set the "
f"`{env_var}` environment variable or provide it in the UI."
)
class MissingAPIKeyError(Exception):
"""Custom exception for missing API key."""
def __init__(self, provider: str, env_var: str):
provider_display = PROVIDER_DISPLAY_NAMES.get(provider, provider.upper())
super().__init__(f"💥 {provider_display} API key not found! 🔑 Please set the "
f"`{env_var}` environment variable or provide it in the UI.")
def encode_image(img_path):

View File

@@ -34,7 +34,7 @@ from src.controller.custom_controller import CustomController
from gradio.themes import Citrus, Default, Glass, Monochrome, Ocean, Origin, Soft, Base
from src.utils.default_config_settings import default_config, load_config_from_file, save_config_to_file, \
save_current_config, update_ui_from_config
from src.utils.utils import update_model_dropdown, get_latest_files, capture_screenshot
from src.utils.utils import update_model_dropdown, get_latest_files, capture_screenshot, MissingAPIKeyError
# Global variables for persistence
_global_browser = None
@@ -245,8 +245,9 @@ async def run_browser_agent(
gr.update(interactive=True) # Re-enable run button
)
except gr.Error:
raise
except MissingAPIKeyError as e:
logger.error(str(e))
raise gr.Error(str(e), print_exception=False)
except Exception as e:
import traceback