Merge branch 'main' into retry_docs

Author: DouweM

Makefile

@@ -53,16 +53,16 @@ typecheck-both: typecheck-pyright typecheck-mypy
 .PHONY: test
 test: ## Run tests and collect coverage data
 	@# To test using a specific version of python, run 'make install-all-python' then set environment variable PYTEST_PYTHON=3.10 or similar
-	$(if $(PYTEST_PYTHON),UV_PROJECT_ENVIRONMENT=.venv$(subst .,,$(PYTEST_PYTHON))) uv run $(if $(PYTEST_PYTHON),--python $(PYTEST_PYTHON)) coverage run -m pytest -n auto --dist=loadgroup --durations=20
+	COLUMNS=150 $(if $(PYTEST_PYTHON),UV_PROJECT_ENVIRONMENT=.venv$(subst .,,$(PYTEST_PYTHON))) uv run $(if $(PYTEST_PYTHON),--python $(PYTEST_PYTHON)) coverage run -m pytest -n auto --dist=loadgroup --durations=20
 	@uv run coverage combine
 	@uv run coverage report
 
 .PHONY: test-all-python
 test-all-python: ## Run tests on Python 3.10 to 3.13
-	UV_PROJECT_ENVIRONMENT=.venv310 uv run --python 3.10 --all-extras --all-packages coverage run -p -m pytest
-	UV_PROJECT_ENVIRONMENT=.venv311 uv run --python 3.11 --all-extras --all-packages coverage run -p -m pytest
-	UV_PROJECT_ENVIRONMENT=.venv312 uv run --python 3.12 --all-extras --all-packages coverage run -p -m pytest
-	UV_PROJECT_ENVIRONMENT=.venv313 uv run --python 3.13 --all-extras --all-packages coverage run -p -m pytest
+	COLUMNS=150 UV_PROJECT_ENVIRONMENT=.venv310 uv run --python 3.10 --all-extras --all-packages coverage run -p -m pytest
+	COLUMNS=150 UV_PROJECT_ENVIRONMENT=.venv311 uv run --python 3.11 --all-extras --all-packages coverage run -p -m pytest
+	COLUMNS=150 UV_PROJECT_ENVIRONMENT=.venv312 uv run --python 3.12 --all-extras --all-packages coverage run -p -m pytest
+	COLUMNS=150 UV_PROJECT_ENVIRONMENT=.venv313 uv run --python 3.13 --all-extras --all-packages coverage run -p -m pytest
 	@uv run coverage combine
 	@uv run coverage report
 

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -764,7 +764,9 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
 
     model_kind = provider_name
     if model_kind.startswith('gateway/'):
-        model_kind = provider_name.removeprefix('gateway/')
+        from ..providers.gateway import infer_gateway_model
+
+        return infer_gateway_model(model_kind.removeprefix('gateway/'), model_name=model_name)
     if model_kind in (
         'openai',
         'azure',

pydantic_ai_slim/pydantic_ai/models/bedrock.py

@@ -226,7 +226,7 @@ def __init__(
         self._model_name = model_name
 
         if isinstance(provider, str):
-            provider = infer_provider('gateway/bedrock' if provider == 'gateway' else provider)
+            provider = infer_provider('gateway/converse' if provider == 'gateway' else provider)
         self._provider = provider
         self.client = cast('BedrockRuntimeClient', provider.client)
 

pydantic_ai_slim/pydantic_ai/models/google.py

@@ -204,7 +204,7 @@ def __init__(
         self._model_name = model_name
 
         if isinstance(provider, str):
-            provider = infer_provider('gateway/google-vertex' if provider == 'gateway' else provider)
+            provider = infer_provider('gateway/gemini' if provider == 'gateway' else provider)
         self._provider = provider
         self.client = provider.client
 

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -375,7 +375,7 @@ def __init__(
         self._model_name = model_name
 
         if isinstance(provider, str):
-            provider = infer_provider('gateway/openai' if provider == 'gateway' else provider)
+            provider = infer_provider('gateway/chat' if provider == 'gateway' else provider)
         self._provider = provider
         self.client = provider.client
 
@@ -944,7 +944,7 @@ def __init__(
         self._model_name = model_name
 
         if isinstance(provider, str):
-            provider = infer_provider('gateway/openai' if provider == 'gateway' else provider)
+            provider = infer_provider('gateway/responses' if provider == 'gateway' else provider)
         self._provider = provider
         self.client = provider.client
 

pydantic_ai_slim/pydantic_ai/providers/__init__.py

@@ -158,8 +158,8 @@ def infer_provider(provider: str) -> Provider[Any]:
     if provider.startswith('gateway/'):
         from .gateway import gateway_provider
 
-        provider = provider.removeprefix('gateway/')
-        return gateway_provider(provider)
+        api_type = provider.removeprefix('gateway/')
+        return gateway_provider(api_type)
     elif provider in ('google-vertex', 'google-gla'):
         from .google import GoogleProvider
 

pydantic_ai_slim/pydantic_ai/providers/gateway.py

@@ -17,6 +17,7 @@
     from groq import AsyncGroq
     from openai import AsyncOpenAI
 
+    from pydantic_ai.models import Model
     from pydantic_ai.models.anthropic import AsyncAnthropicClient
     from pydantic_ai.providers import Provider
 
@@ -25,8 +26,11 @@
 
 @overload
 def gateway_provider(
-    upstream_provider: Literal['openai', 'openai-chat', 'openai-responses'],
+    api_type: Literal['chat', 'responses'],
+    /,
     *,
+    routing_group: str | None = None,
+    profile: str | None = None,
     api_key: str | None = None,
     base_url: str | None = None,
     http_client: httpx.AsyncClient | None = None,
@@ -35,8 +39,11 @@ def gateway_provider(
 
 @overload
 def gateway_provider(
-    upstream_provider: Literal['groq'],
+    api_type: Literal['groq'],
+    /,
     *,
+    routing_group: str | None = None,
+    profile: str | None = None,
     api_key: str | None = None,
     base_url: str | None = None,
     http_client: httpx.AsyncClient | None = None,
@@ -45,56 +52,77 @@ def gateway_provider(
 
 @overload
 def gateway_provider(
-    upstream_provider: Literal['google-vertex'],
+    api_type: Literal['anthropic'],
+    /,
     *,
+    routing_group: str | None = None,
+    profile: str | None = None,
     api_key: str | None = None,
     base_url: str | None = None,
-) -> Provider[GoogleClient]: ...
+    http_client: httpx.AsyncClient | None = None,
+) -> Provider[AsyncAnthropicClient]: ...
 
 
 @overload
 def gateway_provider(
-    upstream_provider: Literal['anthropic'],
+    api_type: Literal['converse'],
+    /,
     *,
+    routing_group: str | None = None,
+    profile: str | None = None,
     api_key: str | None = None,
     base_url: str | None = None,
-) -> Provider[AsyncAnthropicClient]: ...
+) -> Provider[BaseClient]: ...
 
 
 @overload
 def gateway_provider(
-    upstream_provider: Literal['bedrock'],
+    api_type: Literal['gemini'],
+    /,
     *,
+    routing_group: str | None = None,
+    profile: str | None = None,
     api_key: str | None = None,
     base_url: str | None = None,
-) -> Provider[BaseClient]: ...
+    http_client: httpx.AsyncClient | None = None,
+) -> Provider[GoogleClient]: ...
 
 
 @overload
 def gateway_provider(
-    upstream_provider: str,
+    api_type: str,
+    /,
     *,
+    routing_group: str | None = None,
+    profile: str | None = None,
     api_key: str | None = None,
     base_url: str | None = None,
 ) -> Provider[Any]: ...
 
 
-UpstreamProvider = Literal['openai', 'openai-chat', 'openai-responses', 'groq', 'google-vertex', 'anthropic', 'bedrock']
+APIType = Literal['chat', 'responses', 'gemini', 'converse', 'anthropic', 'groq']
 
 
 def gateway_provider(
-    upstream_provider: UpstreamProvider | str,
+    api_type: APIType | str,
+    /,
     *,
     # Every provider
+    routing_group: str | None = None,
+    profile: str | None = None,
     api_key: str | None = None,
     base_url: str | None = None,
-    # OpenAI, Groq & Anthropic
+    # OpenAI, Groq, Anthropic & Gemini - Only Bedrock doesn't have an HTTPX client.
     http_client: httpx.AsyncClient | None = None,
 ) -> Provider[Any]:
     """Create a new Gateway provider.
 
     Args:
-        upstream_provider: The upstream provider to use.
+        api_type: Determines the API type to use.
+        routing_group: The group of APIs that support the same models - the idea is that you can route the requests to
+            any provider in a routing group. The `pydantic-ai-gateway-routing-group` header will be added.
+        profile: A provider may have a profile, which is a unique identifier for the provider.
+            The `pydantic-ai-gateway-profile` header will be added.
         api_key: The API key to use for authentication. If not provided, the `PYDANTIC_AI_GATEWAY_API_KEY`
             environment variable will be used if available.
         base_url: The base URL to use for the Gateway. If not provided, the `PYDANTIC_AI_GATEWAY_BASE_URL`
@@ -109,18 +137,24 @@ def gateway_provider(
         )
 
     base_url = base_url or os.getenv('PYDANTIC_AI_GATEWAY_BASE_URL', GATEWAY_BASE_URL)
-    http_client = http_client or cached_async_http_client(provider=f'gateway/{upstream_provider}')
+    http_client = http_client or cached_async_http_client(provider=f'gateway/{api_type}')
     http_client.event_hooks = {'request': [_request_hook(api_key)]}
 
-    if upstream_provider in ('openai', 'openai-chat', 'openai-responses'):
+    if profile is not None:
+        http_client.headers.setdefault('pydantic-ai-gateway-profile', profile)
+
+    if routing_group is not None:
+        http_client.headers.setdefault('pydantic-ai-gateway-routing-group', routing_group)
+
+    if api_type in ('chat', 'responses'):
         from .openai import OpenAIProvider
 
-        return OpenAIProvider(api_key=api_key, base_url=_merge_url_path(base_url, 'openai'), http_client=http_client)
-    elif upstream_provider == 'groq':
+        return OpenAIProvider(api_key=api_key, base_url=_merge_url_path(base_url, api_type), http_client=http_client)
+    elif api_type == 'groq':
         from .groq import GroqProvider
 
         return GroqProvider(api_key=api_key, base_url=_merge_url_path(base_url, 'groq'), http_client=http_client)
-    elif upstream_provider == 'anthropic':
+    elif api_type == 'anthropic':
         from anthropic import AsyncAnthropic
 
         from .anthropic import AnthropicProvider
@@ -132,25 +166,25 @@ def gateway_provider(
                 http_client=http_client,
             )
         )
-    elif upstream_provider == 'bedrock':
+    elif api_type == 'converse':
         from .bedrock import BedrockProvider
 
         return BedrockProvider(
             api_key=api_key,
-            base_url=_merge_url_path(base_url, 'bedrock'),
+            base_url=_merge_url_path(base_url, api_type),
             region_name='pydantic-ai-gateway',  # Fake region name to avoid NoRegionError
         )
-    elif upstream_provider == 'google-vertex':
+    elif api_type == 'gemini':
         from .google import GoogleProvider
 
         return GoogleProvider(
             vertexai=True,
             api_key=api_key,
-            base_url=_merge_url_path(base_url, 'google-vertex'),
+            base_url=_merge_url_path(base_url, 'gemini'),
             http_client=http_client,
         )
     else:
-        raise UserError(f'Unknown upstream provider: {upstream_provider}')
+        raise UserError(f'Unknown API type: {api_type}')
 
 
 def _request_hook(api_key: str) -> Callable[[httpx.Request], Awaitable[httpx.Request]]:
@@ -182,3 +216,33 @@ def _merge_url_path(base_url: str, path: str) -> str:
         path: The path to merge.
     """
     return base_url.rstrip('/') + '/' + path.lstrip('/')
+
+
+def infer_gateway_model(api_type: APIType | str, *, model_name: str) -> Model:
+    """Infer the model class for a given API type."""
+    if api_type == 'chat':
+        from pydantic_ai.models.openai import OpenAIChatModel
+
+        return OpenAIChatModel(model_name=model_name, provider='gateway')
+    elif api_type == 'groq':
+        from pydantic_ai.models.groq import GroqModel
+
+        return GroqModel(model_name=model_name, provider='gateway')
+    elif api_type == 'responses':
+        from pydantic_ai.models.openai import OpenAIResponsesModel
+
+        return OpenAIResponsesModel(model_name=model_name, provider='gateway')
+    elif api_type == 'gemini':
+        from pydantic_ai.models.google import GoogleModel
+
+        return GoogleModel(model_name=model_name, provider='gateway')
+    elif api_type == 'converse':
+        from pydantic_ai.models.bedrock import BedrockConverseModel
+
+        return BedrockConverseModel(model_name=model_name, provider='gateway')
+    elif api_type == 'anthropic':
+        from pydantic_ai.models.anthropic import AnthropicModel
+
+        return AnthropicModel(model_name=model_name, provider='gateway')
+    else:
+        raise ValueError(f'Unknown API type: {api_type}')  # pragma: no cover

pydantic_ai_slim/pydantic_ai/ui/_adapter.py

@@ -19,6 +19,7 @@
 
 from pydantic_ai import DeferredToolRequests, DeferredToolResults
 from pydantic_ai.agent import AbstractAgent
+from pydantic_ai.agent.abstract import Instructions
 from pydantic_ai.builtin_tools import AbstractBuiltinTool
 from pydantic_ai.exceptions import UserError
 from pydantic_ai.messages import ModelMessage
@@ -200,6 +201,7 @@ def run_stream_native(
         message_history: Sequence[ModelMessage] | None = None,
         deferred_tool_results: DeferredToolResults | None = None,
         model: Model | KnownModelName | str | None = None,
+        instructions: Instructions[AgentDepsT] = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: UsageLimits | None = None,
@@ -216,6 +218,7 @@ def run_stream_native(
             message_history: History of the conversation so far.
             deferred_tool_results: Optional results for deferred tool calls in the message history.
             model: Optional model to use for this run, required if `model` was not set when creating the agent.
+            instructions: Optional additional instructions to use for this run.
             deps: Optional dependencies to use for this run.
             model_settings: Optional settings to use for this model's request.
             usage_limits: Optional limits on model request count or token usage.
@@ -251,6 +254,7 @@ def run_stream_native(
             model=model,
             deps=deps,
             model_settings=model_settings,
+            instructions=instructions,
             usage_limits=usage_limits,
             usage=usage,
             infer_name=infer_name,
@@ -265,6 +269,7 @@ def run_stream(
         message_history: Sequence[ModelMessage] | None = None,
         deferred_tool_results: DeferredToolResults | None = None,
         model: Model | KnownModelName | str | None = None,
+        instructions: Instructions[AgentDepsT] = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: UsageLimits | None = None,
@@ -282,6 +287,7 @@ def run_stream(
             message_history: History of the conversation so far.
             deferred_tool_results: Optional results for deferred tool calls in the message history.
             model: Optional model to use for this run, required if `model` was not set when creating the agent.
+            instructions: Optional additional instructions to use for this run.
             deps: Optional dependencies to use for this run.
             model_settings: Optional settings to use for this model's request.
             usage_limits: Optional limits on model request count or token usage.
@@ -298,6 +304,7 @@ def run_stream(
                 message_history=message_history,
                 deferred_tool_results=deferred_tool_results,
                 model=model,
+                instructions=instructions,
                 deps=deps,
                 model_settings=model_settings,
                 usage_limits=usage_limits,
@@ -318,6 +325,7 @@ async def dispatch_request(
         message_history: Sequence[ModelMessage] | None = None,
         deferred_tool_results: DeferredToolResults | None = None,
         model: Model | KnownModelName | str | None = None,
+        instructions: Instructions[AgentDepsT] = None,
         deps: AgentDepsT = None,
         output_type: OutputSpec[Any] | None = None,
         model_settings: ModelSettings | None = None,
@@ -338,6 +346,7 @@ async def dispatch_request(
             message_history: History of the conversation so far.
             deferred_tool_results: Optional results for deferred tool calls in the message history.
             model: Optional model to use for this run, required if `model` was not set when creating the agent.
+            instructions: Optional additional instructions to use for this run.
             deps: Optional dependencies to use for this run.
             model_settings: Optional settings to use for this model's request.
             usage_limits: Optional limits on model request count or token usage.
@@ -375,6 +384,7 @@ async def dispatch_request(
                 deps=deps,
                 output_type=output_type,
                 model=model,
+                instructions=instructions,
                 model_settings=model_settings,
                 usage_limits=usage_limits,
                 usage=usage,