fix: infer azure provider and environment for key (#1422)

Author: Wh1isper

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -427,7 +427,7 @@ def infer_model(model: Model | KnownModelName | str) -> Model:
         from .cohere import CohereModel
 
         return CohereModel(model_name, provider=provider)
-    elif provider in ('deepseek', 'openai'):
+    elif provider in ('deepseek', 'openai', 'azure'):
         from .openai import OpenAIModel
 
         return OpenAIModel(model_name, provider=provider)

pydantic_ai_slim/pydantic_ai/providers/__init__.py

@@ -52,6 +52,10 @@ def infer_provider(provider: str) -> Provider[Any]:
         from .deepseek import DeepSeekProvider
 
         return DeepSeekProvider()
+    elif provider == 'azure':
+        from .azure import AzureProvider
+
+        return AzureProvider()
     elif provider == 'google-vertex':
         from .google_vertex import GoogleVertexProvider
 

pydantic_ai_slim/pydantic_ai/providers/azure.py

@@ -87,9 +87,9 @@ def __init__(
                     'Must provide one of the `azure_endpoint` argument or the `AZURE_OPENAI_ENDPOINT` environment variable'
                 )
 
-            if not api_key and 'OPENAI_API_KEY' not in os.environ:  # pragma: no cover
+            if not api_key and 'AZURE_OPENAI_API_KEY' not in os.environ:  # pragma: no cover
                 raise UserError(
-                    'Must provide one of the `api_key` argument or the `OPENAI_API_KEY` environment variable'
+                    'Must provide one of the `api_key` argument or the `AZURE_OPENAI_API_KEY` environment variable'
                 )
 
             if not api_version and 'OPENAI_API_VERSION' not in os.environ:  # pragma: no cover

tests/models/test_model.py

@@ -11,6 +11,7 @@
     ('OPENAI_API_KEY', 'openai:gpt-3.5-turbo', 'gpt-3.5-turbo', 'openai', 'openai', 'OpenAIModel'),
     ('OPENAI_API_KEY', 'gpt-3.5-turbo', 'gpt-3.5-turbo', 'openai', 'openai', 'OpenAIModel'),
     ('OPENAI_API_KEY', 'o1', 'o1', 'openai', 'openai', 'OpenAIModel'),
+    ('AZURE_OPENAI_API_KEY', 'azure:gpt-3.5-turbo', 'gpt-3.5-turbo', 'azure', 'azure', 'OpenAIModel'),
     ('GEMINI_API_KEY', 'google-gla:gemini-1.5-flash', 'gemini-1.5-flash', 'google-gla', 'gemini', 'GeminiModel'),
     ('GEMINI_API_KEY', 'gemini-1.5-flash', 'gemini-1.5-flash', 'google-gla', 'gemini', 'GeminiModel'),
     (

tests/providers/test_provider_names.py

@@ -15,6 +15,7 @@
     from openai import OpenAIError
 
     from pydantic_ai.providers.anthropic import AnthropicProvider
+    from pydantic_ai.providers.azure import AzureProvider
     from pydantic_ai.providers.cohere import CohereProvider
     from pydantic_ai.providers.deepseek import DeepSeekProvider
     from pydantic_ai.providers.google_gla import GoogleGLAProvider
@@ -28,6 +29,7 @@
         ('cohere', CohereProvider, 'CO_API_KEY'),
         ('deepseek', DeepSeekProvider, 'DEEPSEEK_API_KEY'),
         ('openai', OpenAIProvider, 'OPENAI_API_KEY'),
+        ('azure', AzureProvider, 'AZURE_OPENAI'),
         ('google-vertex', GoogleVertexProvider, None),
         ('google-gla', GoogleGLAProvider, 'GEMINI_API_KEY'),
         ('groq', GroqProvider, 'GROQ_API_KEY'),