Use provider on models inference (#1234)

Author: Kludex

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -12,7 +12,7 @@
 from dataclasses import dataclass, field
 from datetime import datetime
 from functools import cache
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import httpx
 from typing_extensions import Literal
@@ -383,6 +383,7 @@ def infer_model(model: Model | KnownModelName) -> Model:
 
     try:
         provider, model_name = model.split(':', maxsplit=1)
+        provider = cast(str, provider)
     except ValueError:
         model_name = model
         # TODO(Marcelo): We should deprecate this way.
@@ -414,22 +415,19 @@ def infer_model(model: Model | KnownModelName) -> Model:
     elif provider == 'groq':
         from .groq import GroqModel
 
-        # TODO(Marcelo): Missing provider API.
-        return GroqModel(model_name)
+        return GroqModel(model_name, provider=provider)
     elif provider == 'mistral':
         from .mistral import MistralModel
 
-        # TODO(Marcelo): Missing provider API.
-        return MistralModel(model_name)
+        return MistralModel(model_name, provider=provider)
     elif provider == 'anthropic':
         from .anthropic import AnthropicModel
 
-        # TODO(Marcelo): Missing provider API.
-        return AnthropicModel(model_name)
+        return AnthropicModel(model_name, provider=provider)
     elif provider == 'bedrock':
         from .bedrock import BedrockConverseModel
 
-        return BedrockConverseModel(model_name)
+        return BedrockConverseModel(model_name, provider=provider)
     else:
         raise UserError(f'Unknown model: {model}')
 

tests/test_examples.py

@@ -86,6 +86,8 @@ def test_docs_examples(  # noqa: C901
     env.set('GEMINI_API_KEY', 'testing')
     env.set('GROQ_API_KEY', 'testing')
     env.set('CO_API_KEY', 'testing')
+    env.set('MISTRAL_API_KEY', 'testing')
+    env.set('ANTHROPIC_API_KEY', 'testing')
 
     sys.path.append('tests/example_modules')