custom exception for missing api key

Author: njbrake

src/any_llm/__init__.py

@@ -1,7 +1,11 @@
 from any_llm.main import (
     completion,
 )
+from any_llm.exceptions import (
+    MissingApiKeyError,
+)
 
 __all__ = [
     "completion",
+    "MissingApiKeyError",
 ]

src/any_llm/exceptions.py

@@ -0,0 +1,23 @@
+"""Custom exceptions for any-llm package."""
+
+
+class MissingApiKeyError(Exception):
+    """Exception raised when an API key is missing or not provided."""
+    
+    def __init__(self, provider_name: str, env_var_name: str) -> None:
+        """Initialize the exception.
+        
+        Args:
+            provider_name: Name of the provider (e.g., "OpenAI", "Google", "Mistral")
+            env_var_name: Name of the environment variable that should contain the API key
+            message: Optional custom message. If not provided, a default message will be used.
+        """
+        self.provider_name = provider_name
+        self.env_var_name = env_var_name
+        
+        message = (
+            f"No {provider_name} API key provided. "
+            f"Please provide it in the config or set the {env_var_name} environment variable."
+        )
+        
+        super().__init__(message) 

src/any_llm/providers/google/google.py

@@ -14,6 +14,7 @@
 from openai.types.chat.chat_completion_message import ChatCompletionMessage
 from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall, Function
 from any_llm.provider import Provider, ApiConfig
+from any_llm.exceptions import MissingApiKeyError
 
 DEFAULT_TEMPERATURE = 0.7
 
@@ -196,8 +197,10 @@ def __init__(self, config: ApiConfig) -> None:
             self.location = os.getenv("GOOGLE_REGION", "us-central1")
 
             if not self.project_id:
-                msg = "GOOGLE_PROJECT_ID environment variable is required for Vertex AI"
-                raise ValueError(msg)
+                raise MissingApiKeyError(
+                    "Google Vertex AI",
+                    "GOOGLE_PROJECT_ID",
+                )
 
             # Initialize client for Vertex AI
             self.client = genai.Client(vertexai=True, project=self.project_id, location=self.location)
@@ -207,8 +210,10 @@ def __init__(self, config: ApiConfig) -> None:
             api_key = getattr(config, "api_key", None) or os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
 
             if not api_key:
-                msg = "API key is required for Gemini Developer API. Provide it via ApiConfig or set GEMINI_API_KEY/GOOGLE_API_KEY environment variable"
-                raise ValueError(msg)
+                raise MissingApiKeyError(
+                    "Google Gemini Developer API",
+                    "GEMINI_API_KEY/GOOGLE_API_KEY",
+                )
 
             # Initialize client for Gemini Developer API
             self.client = genai.Client(api_key=api_key)

src/any_llm/providers/mistral/mistral.py

@@ -15,6 +15,7 @@
 from openai.types.chat.chat_completion_message import ChatCompletionMessage
 from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall, Function
 from any_llm.provider import Provider, ApiConfig
+from any_llm.exceptions import MissingApiKeyError
 
 
 def _convert_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
@@ -97,8 +98,7 @@ def __init__(self, config: ApiConfig) -> None:
         if not config.api_key:
             config.api_key = os.getenv("MISTRAL_API_KEY")
         if not config.api_key:
-            msg = "No Mistral API key provided. Please provide it in the config or set the MISTRAL_API_KEY environment variable."
-            raise ValueError(msg)
+            raise MissingApiKeyError("Mistral", "MISTRAL_API_KEY")
         self.client = Mistral(api_key=config.api_key, server_url=config.api_base)
 
     def completion(

src/any_llm/providers/openai/base.py

@@ -6,6 +6,7 @@
 from openai.types.chat.chat_completion import ChatCompletion
 
 from any_llm.provider import Provider, ApiConfig
+from any_llm.exceptions import MissingApiKeyError
 
 
 class BaseOpenAIProvider(Provider, ABC):
@@ -34,8 +35,7 @@ def __init__(self, config: ApiConfig) -> None:
             client_kwargs["base_url"] = config.api_base
 
         if not config.api_key and not os.getenv(self.ENV_API_KEY_NAME):
-            msg = f"No {self.PROVIDER_NAME} API key provided. Please provide it in the config or set the {self.ENV_API_KEY_NAME} environment variable."
-            raise ValueError(msg)
+            raise MissingApiKeyError(self.PROVIDER_NAME, self.ENV_API_KEY_NAME)
 
         # Get API key from environment if not provided in config
         api_key = config.api_key or os.getenv(self.ENV_API_KEY_NAME)

test_exception.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+"""Simple test to verify the custom exception works correctly."""
+
+from any_llm.exceptions import MissingApiKeyError
+
+# Test the custom exception
+try:
+    raise MissingApiKeyError("Test Provider", "TEST_API_KEY")
+except MissingApiKeyError as e:
+    print(f"✅ Custom exception works: {e}")
+    print(f"   Provider: {e.provider_name}")
+    print(f"   Env var: {e.env_var_name}")
+
+# Test with custom message
+try:
+    raise MissingApiKeyError("Test Provider", "TEST_API_KEY", "Custom error message")
+except MissingApiKeyError as e:
+    print(f"✅ Custom message works: {e}")
+
+print("All tests passed! 🎉") 

tests/integration/test_completion.py

@@ -1,6 +1,7 @@
 import httpx
 import pytest
 from any_llm import completion
+from any_llm.exceptions import MissingApiKeyError
 
 
 # Use small models for testing to make sure they work
@@ -25,10 +26,8 @@ def test_providers(provider: str) -> None:
     model_id = provider_model_map[provider]
     try:
         result = completion(f"{provider}/{model_id}", messages=[{"role": "user", "content": "Hello"}])
-    except ValueError as e:
-        if "API key provided" in str(e):
-            pytest.skip(f"{provider} API key not provided, skipping")
-        raise e
+    except MissingApiKeyError:
+        pytest.skip(f"{provider} API key not provided, skipping")
     except (httpx.HTTPStatusError, httpx.ConnectError):
         if provider == "ollama":
             pytest.skip("Ollama is not set up, skipping")