pydantic-ai-slim 0.0.31__tar.gz → 0.0.33__tar.gz

{pydantic_ai_slim-0.0.31 → pydantic_ai_slim-0.0.33}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.0.31
+Version: 0.0.33
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>
 License-Expression: MIT
@@ -28,12 +28,12 @@ Requires-Dist: eval-type-backport>=0.2.0
 Requires-Dist: exceptiongroup; python_version < '3.11'
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
-Requires-Dist: logfire-api>=1.2.0
 Requires-Dist: opentelemetry-api>=1.28.0
-Requires-Dist: pydantic-graph==0.0.31
+Requires-Dist: pydantic-graph==0.0.33
 Requires-Dist: pydantic>=2.10
+Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: anthropic
-Requires-Dist: anthropic>=0.40.0; extra == 'anthropic'
+Requires-Dist: anthropic>=0.49.0; extra == 'anthropic'
 Provides-Extra: cohere
 Requires-Dist: cohere>=5.13.11; extra == 'cohere'
 Provides-Extra: duckduckgo

{pydantic_ai_slim-0.0.31 → pydantic_ai_slim-0.0.33}/pydantic_ai/_agent_graph.py

@@ -2,13 +2,14 @@ from __future__ import annotations as _annotations
 
 import asyncio
 import dataclasses
+import json
 from collections.abc import AsyncIterator, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar
 from dataclasses import field
 from typing import Any, Generic, Literal, Union, cast
 
-import logfire_api
+from opentelemetry.trace import Span, Tracer
 from typing_extensions import TypeGuard, TypeVar, assert_never
 
 from pydantic_graph import BaseNode, Graph, GraphRunContext
@@ -42,17 +43,6 @@ __all__ = (
     'capture_run_messages',
 )
 
-_logfire = logfire_api.Logfire(otel_scope='pydantic-ai')
-
-# while waiting for https://github.com/pydantic/logfire/issues/745
-try:
-    import logfire._internal.stack_info
-except ImportError:
-    pass
-else:
-    from pathlib import Path
-
-    logfire._internal.stack_info.NON_USER_CODE_PREFIXES += (str(Path(__file__).parent.absolute()),)
 
 T = TypeVar('T')
 S = TypeVar('S')
@@ -105,7 +95,8 @@ class GraphAgentDeps(Generic[DepsT, ResultDataT]):
 
     function_tools: dict[str, Tool[DepsT]] = dataclasses.field(repr=False)
 
-    run_span: logfire_api.LogfireSpan
+    run_span: Span
+    tracer: Tracer
 
 
 class AgentNode(BaseNode[GraphAgentState, GraphAgentDeps[DepsT, Any], result.FinalResult[NodeRunEndT]]):
@@ -330,7 +321,9 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
         ctx.state.run_step += 1
 
         model_settings = merge_model_settings(ctx.deps.model_settings, None)
-        with _logfire.span('preparing model request params {run_step=}', run_step=ctx.state.run_step):
+        with ctx.deps.tracer.start_as_current_span(
+            'preparing model request params', attributes=dict(run_step=ctx.state.run_step)
+        ):
             model_request_parameters = await _prepare_request_parameters(ctx)
         return model_settings, model_request_parameters
 
@@ -380,26 +373,12 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
         self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
     ) -> AsyncIterator[AsyncIterator[_messages.HandleResponseEvent]]:
         """Process the model response and yield events for the start and end of each function tool call."""
-        with _logfire.span('handle model response', run_step=ctx.state.run_step) as handle_span:
-            stream = self._run_stream(ctx)
-            yield stream
+        stream = self._run_stream(ctx)
+        yield stream
 
-            # Run the stream to completion if it was not finished:
-            async for _event in stream:
-                pass
-
-            # Set the next node based on the final state of the stream
-            next_node = self._next_node
-            if isinstance(next_node, End):
-                handle_span.set_attribute('result', next_node.data)
-                handle_span.message = 'handle model response -> final result'
-            elif tool_responses := self._tool_responses:
-                # TODO: We could drop `self._tool_responses` if we drop this set_attribute
-                #   I'm thinking it might be better to just create a span for the handling of each tool
-                #   than to set an attribute here.
-                handle_span.set_attribute('tool_responses', tool_responses)
-                tool_responses_str = ' '.join(r.part_kind for r in tool_responses)
-                handle_span.message = f'handle model response -> {tool_responses_str}'
+        # Run the stream to completion if it was not finished:
+        async for _event in stream:
+            pass
 
     async def _run_stream(
         self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
@@ -494,10 +473,29 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
         if tool_responses:
             messages.append(_messages.ModelRequest(parts=tool_responses))
 
-        run_span.set_attribute('usage', usage)
-        run_span.set_attribute(
-            'all_messages_events',
-            [InstrumentedModel.event_to_dict(e) for e in InstrumentedModel.messages_to_otel_events(messages)],
+        run_span.set_attributes(
+            {
+                **usage.opentelemetry_attributes(),
+                'all_messages_events': json.dumps(
+                    [InstrumentedModel.event_to_dict(e) for e in InstrumentedModel.messages_to_otel_events(messages)]
+                ),
+                'final_result': final_result.data
+                if isinstance(final_result.data, str)
+                else json.dumps(InstrumentedModel.serialize_any(final_result.data)),
+            }
+        )
+        run_span.set_attributes(
+            {
+                'logfire.json_schema': json.dumps(
+                    {
+                        'type': 'object',
+                        'properties': {
+                            'all_messages_events': {'type': 'array'},
+                            'final_result': {'type': 'object'},
+                        },
+                    }
+                ),
+            }
         )
 
         # End the run with self.data
@@ -619,7 +617,10 @@ async def process_function_tools(
 
     # Run all tool tasks in parallel
     results_by_index: dict[int, _messages.ModelRequestPart] = {}
-    with _logfire.span('running {tools=}', tools=[call.tool_name for _, call in calls_to_run]):
+    tool_names = [call.tool_name for _, call in calls_to_run]
+    with ctx.deps.tracer.start_as_current_span(
+        'running tools', attributes={'tools': tool_names, 'logfire.msg': f'running tools: {", ".join(tool_names)}'}
+    ):
         # TODO: Should we wrap each individual tool call in a dedicated span?
         tasks = [asyncio.create_task(tool.run(call, run_context), name=call.tool_name) for tool, call in calls_to_run]
         pending = tasks

{pydantic_ai_slim-0.0.31 → pydantic_ai_slim-0.0.33}/pydantic_ai/_pydantic.py

@@ -6,7 +6,7 @@ This module has to use numerous internal Pydantic APIs and is therefore brittle
 from __future__ import annotations as _annotations
 
 from inspect import Parameter, signature
-from typing import TYPE_CHECKING, Any, Callable, TypedDict, cast, get_origin
+from typing import TYPE_CHECKING, Any, Callable, TypedDict, cast
 
 from pydantic import ConfigDict
 from pydantic._internal import _decorators, _generate_schema, _typing_extra
@@ -15,6 +15,7 @@ from pydantic.fields import FieldInfo
 from pydantic.json_schema import GenerateJsonSchema
 from pydantic.plugin._schema_validator import create_schema_validator
 from pydantic_core import SchemaValidator, core_schema
+from typing_extensions import get_origin
 
 from ._griffe import doc_descriptions
 from ._utils import check_object_json_schema, is_model_like
@@ -223,8 +224,7 @@ def _build_schema(
 
 
 def _is_call_ctx(annotation: Any) -> bool:
+    """Return whether the annotation is the `RunContext` class, parameterized or not."""
     from .tools import RunContext
 
-    return annotation is RunContext or (
-        _typing_extra.is_generic_alias(annotation) and get_origin(annotation) is RunContext
-    )
+    return annotation is RunContext or get_origin(annotation) is RunContext

{pydantic_ai_slim-0.0.31 → pydantic_ai_slim-0.0.33}/pydantic_ai/_result.py

@@ -1,14 +1,14 @@
 from __future__ import annotations as _annotations
 
 import inspect
-import sys
-import types
 from collections.abc import Awaitable, Iterable, Iterator
 from dataclasses import dataclass, field
-from typing import Any, Callable, Generic, Literal, Union, cast, get_args, get_origin
+from typing import Any, Callable, Generic, Literal, Union, cast
 
 from pydantic import TypeAdapter, ValidationError
-from typing_extensions import TypeAliasType, TypedDict, TypeVar
+from typing_extensions import TypedDict, TypeVar, get_args, get_origin
+from typing_inspection import typing_objects
+from typing_inspection.introspection import is_union_origin
 
 from . import _utils, messages as _messages
 from .exceptions import ModelRetry
@@ -248,23 +248,12 @@ def extract_str_from_union(response_type: Any) -> _utils.Option[Any]:
 
 
 def get_union_args(tp: Any) -> tuple[Any, ...]:
-    """Extract the arguments of a Union type if `response_type` is a union, otherwise return an empty union."""
-    if isinstance(tp, TypeAliasType):
+    """Extract the arguments of a Union type if `response_type` is a union, otherwise return an empty tuple."""
+    if typing_objects.is_typealiastype(tp):
         tp = tp.__value__
 
     origin = get_origin(tp)
-    if origin_is_union(origin):
+    if is_union_origin(origin):
         return get_args(tp)
     else:
         return ()
-
-
-if sys.version_info < (3, 10):
-
-    def origin_is_union(tp: type[Any] | None) -> bool:
-        return tp is Union
-
-else:
-
-    def origin_is_union(tp: type[Any] | None) -> bool:
-        return tp is Union or tp is types.UnionType

{pydantic_ai_slim-0.0.31 → pydantic_ai_slim-0.0.33}/pydantic_ai/agent.py

@@ -8,7 +8,7 @@ from copy import deepcopy
 from types import FrameType
 from typing import Any, Callable, Generic, cast, final, overload
 
-import logfire_api
+from opentelemetry.trace import NoOpTracer, use_span
 from typing_extensions import TypeGuard, TypeVar, deprecated
 
 from pydantic_graph import End, Graph, GraphRun, GraphRunContext
@@ -58,17 +58,6 @@ __all__ = (
     'UserPromptNode',
 )
 
-_logfire = logfire_api.Logfire(otel_scope='pydantic-ai')
-
-# while waiting for https://github.com/pydantic/logfire/issues/745
-try:
-    import logfire._internal.stack_info
-except ImportError:
-    pass
-else:
-    from pathlib import Path
-
-    logfire._internal.stack_info.NON_USER_CODE_PREFIXES += (str(Path(__file__).parent.absolute()),)
 
 T = TypeVar('T')
 S = TypeVar('S')
@@ -123,6 +112,9 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
     The type of the result data, used to validate the result data, defaults to `str`.
     """
 
+    instrument: bool
+    """Automatically instrument with OpenTelemetry. Will use Logfire if it's configured."""
+
     _deps_type: type[AgentDepsT] = dataclasses.field(repr=False)
     _result_tool_name: str = dataclasses.field(repr=False)
     _result_tool_description: str | None = dataclasses.field(repr=False)
@@ -155,6 +147,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
+        instrument: bool = False,
     ):
         """Create an agent.
 
@@ -184,6 +177,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
                 [override the model][pydantic_ai.Agent.override] for testing.
             end_strategy: Strategy for handling tool calls that are requested alongside a final result.
                 See [`EndStrategy`][pydantic_ai.agent.EndStrategy] for more information.
+            instrument: Automatically instrument with OpenTelemetry. Will use Logfire if it's configured.
         """
         if model is None or defer_model_check:
             self.model = model
@@ -194,6 +188,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         self.name = name
         self.model_settings = model_settings
         self.result_type = result_type
+        self.instrument = instrument
 
         self._deps_type = deps_type
 
@@ -396,6 +391,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         if infer_name and self.name is None:
             self._infer_name(inspect.currentframe())
         model_used = self._get_model(model)
+        del model
 
         deps = self._get_deps(deps)
         new_message_index = len(message_history) if message_history else 0
@@ -425,14 +421,20 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         model_settings = merge_model_settings(self.model_settings, model_settings)
         usage_limits = usage_limits or _usage.UsageLimits()
 
-        # Build the deps object for the graph
-        run_span = _logfire.span(
-            '{agent_name} run {prompt=}',
-            prompt=user_prompt,
-            agent=self,
-            model_name=model_used.model_name if model_used else 'no-model',
-            agent_name=self.name or 'agent',
+        if isinstance(model_used, InstrumentedModel):
+            tracer = model_used.tracer
+        else:
+            tracer = NoOpTracer()
+        agent_name = self.name or 'agent'
+        run_span = tracer.start_span(
+            'agent run',
+            attributes={
+                'model_name': model_used.model_name if model_used else 'no-model',
+                'agent_name': agent_name,
+                'logfire.msg': f'{agent_name} run',
+            },
         )
+
         graph_deps = _agent_graph.GraphAgentDeps[AgentDepsT, RunResultDataT](
             user_deps=deps,
             prompt=user_prompt,
@@ -447,6 +449,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
             result_validators=result_validators,
             function_tools=self._function_tools,
             run_span=run_span,
+            tracer=tracer,
         )
         start_node = _agent_graph.UserPromptNode[AgentDepsT](
             user_prompt=user_prompt,
@@ -460,7 +463,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
             state=state,
             deps=graph_deps,
             infer_name=False,
-            span=run_span,
+            span=use_span(run_span, end_on_exit=True),
         ) as graph_run:
             yield AgentRun(graph_run)
 
@@ -1116,7 +1119,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         else:
             raise exceptions.UserError('`model` must be set either when creating the agent or when calling it.')
 
-        if not isinstance(model_, InstrumentedModel):
+        if self.instrument and not isinstance(model_, InstrumentedModel):
             model_ = InstrumentedModel(model_)
 
         return model_

{pydantic_ai_slim-0.0.31 → pydantic_ai_slim-0.0.33}/pydantic_ai/models/__init__.py

@@ -28,9 +28,11 @@ if TYPE_CHECKING:
 
 
 KnownModelName = Literal[
+    'anthropic:claude-3-7-sonnet-latest',
     'anthropic:claude-3-5-haiku-latest',
     'anthropic:claude-3-5-sonnet-latest',
     'anthropic:claude-3-opus-latest',
+    'claude-3-7-sonnet-latest',
     'claude-3-5-haiku-latest',
     'claude-3-5-sonnet-latest',
     'claude-3-opus-latest',
@@ -47,6 +49,8 @@ KnownModelName = Literal[
     'cohere:command-r-plus-04-2024',
     'cohere:command-r-plus-08-2024',
     'cohere:command-r7b-12-2024',
+    'deepseek:deepseek-chat',
+    'deepseek:deepseek-reasoner',
     'google-gla:gemini-1.0-pro',
     'google-gla:gemini-1.5-flash',
     'google-gla:gemini-1.5-flash-8b',
@@ -56,6 +60,7 @@ KnownModelName = Literal[
     'google-gla:gemini-exp-1206',
     'google-gla:gemini-2.0-flash',
     'google-gla:gemini-2.0-flash-lite-preview-02-05',
+    'google-gla:gemini-2.0-pro-exp-02-05',
     'google-vertex:gemini-1.0-pro',
     'google-vertex:gemini-1.5-flash',
     'google-vertex:gemini-1.5-flash-8b',
@@ -65,6 +70,7 @@ KnownModelName = Literal[
     'google-vertex:gemini-exp-1206',
     'google-vertex:gemini-2.0-flash',
     'google-vertex:gemini-2.0-flash-lite-preview-02-05',
+    'google-vertex:gemini-2.0-pro-exp-02-05',
     'gpt-3.5-turbo',
     'gpt-3.5-turbo-0125',
     'gpt-3.5-turbo-0301',
@@ -316,54 +322,52 @@ def infer_model(model: Model | KnownModelName) -> Model:
         from .test import TestModel
 
         return TestModel()
-    elif model.startswith('cohere:'):
-        from .cohere import CohereModel
 
-        return CohereModel(model[7:])
-    elif model.startswith('openai:'):
-        from .openai import OpenAIModel
+    try:
+        provider, model_name = model.split(':')
+    except ValueError:
+        model_name = model
+        # TODO(Marcelo): We should deprecate this way.
+        if model_name.startswith(('gpt', 'o1', 'o3')):
+            provider = 'openai'
+        elif model_name.startswith('claude'):
+            provider = 'anthropic'
+        elif model_name.startswith('gemini'):
+            provider = 'google-gla'
+        else:
+            raise UserError(f'Unknown model: {model}')
+
+    if provider == 'vertexai':
+        provider = 'google-vertex'
+
+    if provider == 'cohere':
+        from .cohere import CohereModel
 
-        return OpenAIModel(model[7:])
-    elif model.startswith(('gpt', 'o1', 'o3')):
+        # TODO(Marcelo): Missing provider API.
+        return CohereModel(model_name)
+    elif provider in ('deepseek', 'openai'):
         from .openai import OpenAIModel
 
-        return OpenAIModel(model)
-    elif model.startswith('google-gla'):
-        from .gemini import GeminiModel
-
-        return GeminiModel(model[11:])
-    # backwards compatibility with old model names (ex, gemini-1.5-flash -> google-gla:gemini-1.5-flash)
-    elif model.startswith('gemini'):
+        return OpenAIModel(model_name, provider=provider)
+    elif provider in ('google-gla', 'google-vertex'):
         from .gemini import GeminiModel
 
-        # noinspection PyTypeChecker
-        return GeminiModel(model)
-    elif model.startswith('groq:'):
+        return GeminiModel(model_name, provider=provider)
+    elif provider == 'groq':
         from .groq import GroqModel
 
-        return GroqModel(model[5:])
-    elif model.startswith('google-vertex'):
-        from .vertexai import VertexAIModel
-
-        return VertexAIModel(model[14:])
-    # backwards compatibility with old model names (ex, vertexai:gemini-1.5-flash -> google-vertex:gemini-1.5-flash)
-    elif model.startswith('vertexai:'):
-        from .vertexai import VertexAIModel
-
-        return VertexAIModel(model[9:])
-    elif model.startswith('mistral:'):
+        # TODO(Marcelo): Missing provider API.
+        return GroqModel(model_name)
+    elif provider == 'mistral':
         from .mistral import MistralModel
 
-        return MistralModel(model[8:])
-    elif model.startswith('anthropic'):
-        from .anthropic import AnthropicModel
-
-        return AnthropicModel(model[10:])
-    # backwards compatibility with old model names (ex, claude-3-5-sonnet-latest -> anthropic:claude-3-5-sonnet-latest)
-    elif model.startswith('claude'):
+        # TODO(Marcelo): Missing provider API.
+        return MistralModel(model_name)
+    elif provider == 'anthropic':
         from .anthropic import AnthropicModel
 
-        return AnthropicModel(model)
+        # TODO(Marcelo): Missing provider API.
+        return AnthropicModel(model_name)
     else:
         raise UserError(f'Unknown model: {model}')
 

{pydantic_ai_slim-0.0.31 → pydantic_ai_slim-0.0.33}/pydantic_ai/models/anthropic.py

@@ -42,6 +42,7 @@ from . import (
 try:
     from anthropic import NOT_GIVEN, APIStatusError, AsyncAnthropic, AsyncStream
     from anthropic.types import (
+        ContentBlock,
         ImageBlockParam,
         Message as AnthropicMessage,
         MessageParam,
@@ -69,6 +70,7 @@ except ImportError as _import_error:
     ) from _import_error
 
 LatestAnthropicModelNames = Literal[
+    'claude-3-7-sonnet-latest',
     'claude-3-5-haiku-latest',
     'claude-3-5-sonnet-latest',
     'claude-3-opus-latest',
@@ -423,7 +425,7 @@ class AnthropicStreamedResponse(StreamedResponse):
     _timestamp: datetime
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
-        current_block: TextBlock | ToolUseBlock | None = None
+        current_block: ContentBlock | None = None
         current_json: str = ''
 
         async for event in self._response:

{pydantic_ai_slim-0.0.31 → pydantic_ai_slim-0.0.33}/pydantic_ai/models/gemini.py

@@ -8,12 +8,14 @@ from contextlib import asynccontextmanager
 from copy import deepcopy
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Annotated, Any, Literal, Protocol, Union, cast
+from typing import Annotated, Any, Literal, Protocol, Union, cast, overload
 from uuid import uuid4
 
 import pydantic
 from httpx import USE_CLIENT_DEFAULT, AsyncClient as AsyncHTTPClient, Response as HTTPResponse
-from typing_extensions import NotRequired, TypedDict, assert_never
+from typing_extensions import NotRequired, TypedDict, assert_never, deprecated
+
+from pydantic_ai.providers import Provider, infer_provider
 
 from .. import ModelHTTPError, UnexpectedModelBehavior, UserError, _utils, usage
 from ..messages import (
@@ -53,6 +55,7 @@ LatestGeminiModelNames = Literal[
     'gemini-exp-1206',
     'gemini-2.0-flash',
     'gemini-2.0-flash-lite-preview-02-05',
+    'gemini-2.0-pro-exp-02-05',
 ]
 """Latest Gemini models."""
 
@@ -81,17 +84,39 @@ class GeminiModel(Model):
     Apart from `__init__`, all methods are private or match those of the base class.
     """
 
-    http_client: AsyncHTTPClient = field(repr=False)
+    client: AsyncHTTPClient = field(repr=False)
 
     _model_name: GeminiModelName = field(repr=False)
+    _provider: Literal['google-gla', 'google-vertex'] | Provider[AsyncHTTPClient] | None = field(repr=False)
     _auth: AuthProtocol | None = field(repr=False)
     _url: str | None = field(repr=False)
     _system: str | None = field(default='google-gla', repr=False)
 
+    @overload
+    def __init__(
+        self,
+        model_name: GeminiModelName,
+        *,
+        provider: Literal['google-gla', 'google-vertex'] | Provider[AsyncHTTPClient] = 'google-gla',
+    ) -> None: ...
+
+    @deprecated('Use the `provider` argument instead of the `api_key`, `http_client`, and `url_template` arguments.')
+    @overload
     def __init__(
         self,
         model_name: GeminiModelName,
         *,
+        provider: None = None,
+        api_key: str | None = None,
+        http_client: AsyncHTTPClient | None = None,
+        url_template: str = 'https://generativelanguage.googleapis.com/v1beta/models/{model}:',
+    ) -> None: ...
+
+    def __init__(
+        self,
+        model_name: GeminiModelName,
+        *,
+        provider: Literal['google-gla', 'google-vertex'] | Provider[AsyncHTTPClient] | None = None,
         api_key: str | None = None,
         http_client: AsyncHTTPClient | None = None,
         url_template: str = 'https://generativelanguage.googleapis.com/v1beta/models/{model}:',
@@ -100,6 +125,7 @@ class GeminiModel(Model):
 
         Args:
             model_name: The name of the model to use.
+            provider: The provider to use for the model.
             api_key: The API key to use for authentication, if not provided, the `GEMINI_API_KEY` environment variable
                 will be used if available.
             http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
@@ -108,14 +134,24 @@ class GeminiModel(Model):
                 `model` is substituted with the model name, and `function` is added to the end of the URL.
         """
         self._model_name = model_name
-        if api_key is None:
-            if env_api_key := os.getenv('GEMINI_API_KEY'):
-                api_key = env_api_key
+        self._provider = provider
+
+        if provider is not None:
+            if isinstance(provider, str):
+                self._system = provider
+                self.client = infer_provider(provider).client
             else:
-                raise UserError('API key must be provided or set in the GEMINI_API_KEY environment variable')
-        self.http_client = http_client or cached_async_http_client()
-        self._auth = ApiKeyAuth(api_key)
-        self._url = url_template.format(model=model_name)
+                self._system = provider.name
+                self.client = provider.client
+        else:
+            if api_key is None:
+                if env_api_key := os.getenv('GEMINI_API_KEY'):
+                    api_key = env_api_key
+                else:
+                    raise UserError('API key must be provided or set in the GEMINI_API_KEY environment variable')
+            self.client = http_client or cached_async_http_client()
+            self._auth = ApiKeyAuth(api_key)
+            self._url = url_template.format(model=model_name)
 
     @property
     def auth(self) -> AuthProtocol:
@@ -216,17 +252,19 @@ class GeminiModel(Model):
         if generation_config:
             request_data['generation_config'] = generation_config
 
-        url = self.url + ('streamGenerateContent' if streamed else 'generateContent')
-
         headers = {
             'Content-Type': 'application/json',
             'User-Agent': get_user_agent(),
-            **await self.auth.headers(),
         }
+        if self._provider is None:  # pragma: no cover
+            url = self.url + ('streamGenerateContent' if streamed else 'generateContent')
+            headers.update(await self.auth.headers())
+        else:
+            url = f'/{self._model_name}:{"streamGenerateContent" if streamed else "generateContent"}'
 
         request_json = _gemini_request_ta.dump_json(request_data, by_alias=True)
 
-        async with self.http_client.stream(
+        async with self.client.stream(
             'POST',
             url,
             content=request_json,