pydantic-ai-slim 0.0.32__tar.gz → 0.0.34__tar.gz

{pydantic_ai_slim-0.0.32 → pydantic_ai_slim-0.0.34}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.0.32
+Version: 0.0.34
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>
 License-Expression: MIT
@@ -29,10 +29,15 @@ Requires-Dist: exceptiongroup; python_version < '3.11'
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: opentelemetry-api>=1.28.0
-Requires-Dist: pydantic-graph==0.0.32
+Requires-Dist: pydantic-graph==0.0.34
 Requires-Dist: pydantic>=2.10
+Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: anthropic
 Requires-Dist: anthropic>=0.49.0; extra == 'anthropic'
+Provides-Extra: cli
+Requires-Dist: argcomplete>=3.5.0; extra == 'cli'
+Requires-Dist: prompt-toolkit>=3; extra == 'cli'
+Requires-Dist: rich>=13; extra == 'cli'
 Provides-Extra: cohere
 Requires-Dist: cohere>=5.13.11; extra == 'cohere'
 Provides-Extra: duckduckgo

pydantic_ai_slim-0.0.34/pydantic_ai/_cli.py

@@ -0,0 +1,225 @@
+from __future__ import annotations as _annotations
+
+import argparse
+import asyncio
+import sys
+from collections.abc import Sequence
+from datetime import datetime, timezone
+from importlib.metadata import version
+from pathlib import Path
+from typing import cast
+
+from typing_inspection.introspection import get_literal_values
+
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.models import KnownModelName
+from pydantic_graph.nodes import End
+
+try:
+    import argcomplete
+    from prompt_toolkit import PromptSession
+    from prompt_toolkit.auto_suggest import AutoSuggestFromHistory, Suggestion
+    from prompt_toolkit.buffer import Buffer
+    from prompt_toolkit.document import Document
+    from prompt_toolkit.history import FileHistory
+    from rich.console import Console, ConsoleOptions, RenderResult
+    from rich.live import Live
+    from rich.markdown import CodeBlock, Markdown
+    from rich.status import Status
+    from rich.syntax import Syntax
+    from rich.text import Text
+except ImportError as _import_error:
+    raise ImportError(
+        'Please install `rich`, `prompt-toolkit` and `argcomplete` to use the PydanticAI CLI, '
+        "you can use the `cli` optional group — `pip install 'pydantic-ai-slim[cli]'`"
+    ) from _import_error
+
+from pydantic_ai.agent import Agent
+from pydantic_ai.messages import ModelMessage, PartDeltaEvent, TextPartDelta
+
+__version__ = version('pydantic-ai')
+
+
+class SimpleCodeBlock(CodeBlock):
+    def __rich_console__(self, console: Console, options: ConsoleOptions) -> RenderResult:  # pragma: no cover
+        code = str(self.text).rstrip()
+        yield Text(self.lexer_name, style='dim')
+        yield Syntax(code, self.lexer_name, theme=self.theme, background_color='default', word_wrap=True)
+        yield Text(f'/{self.lexer_name}', style='dim')
+
+
+Markdown.elements['fence'] = SimpleCodeBlock
+
+
+def cli(args_list: Sequence[str] | None = None) -> int:  # noqa: C901  # pragma: no cover
+    parser = argparse.ArgumentParser(
+        prog='pai',
+        description=f"""\
+PydanticAI CLI v{__version__}\n\n
+
+Special prompt:
+* `/exit` - exit the interactive mode
+* `/markdown` - show the last markdown output of the last question
+* `/multiline` - toggle multiline mode
+""",
+        formatter_class=argparse.RawTextHelpFormatter,
+    )
+    parser.add_argument('prompt', nargs='?', help='AI Prompt, if omitted fall into interactive mode')
+    parser.add_argument(
+        '--model',
+        nargs='?',
+        help='Model to use, it should be "<provider>:<model>" e.g. "openai:gpt-4o". If omitted it will default to "openai:gpt-4o"',
+        default='openai:gpt-4o',
+    ).completer = argcomplete.ChoicesCompleter(list(get_literal_values(KnownModelName)))  # type: ignore[reportPrivateUsage]
+    parser.add_argument('--no-stream', action='store_true', help='Whether to stream responses from OpenAI')
+    parser.add_argument('--version', action='store_true', help='Show version and exit')
+
+    argcomplete.autocomplete(parser)
+    args = parser.parse_args(args_list)
+
+    console = Console()
+    console.print(f'pai - PydanticAI CLI v{__version__}', style='green bold', highlight=False)
+    if args.version:
+        return 0
+
+    now_utc = datetime.now(timezone.utc)
+    tzname = now_utc.astimezone().tzinfo.tzname(now_utc)  # type: ignore
+    try:
+        agent = Agent(
+            model=args.model or 'openai:gpt-4o',
+            system_prompt=f"""\
+    Help the user by responding to their request, the output should be concise and always written in markdown.
+    The current date and time is {datetime.now()} {tzname}.
+    The user is running {sys.platform}.""",
+        )
+    except UserError:
+        console.print(f'[red]Invalid model "{args.model}"[/red]')
+        return 1
+
+    stream = not args.no_stream
+
+    if prompt := cast(str, args.prompt):
+        try:
+            asyncio.run(ask_agent(agent, prompt, stream, console))
+        except KeyboardInterrupt:
+            pass
+        return 0
+
+    history = Path.home() / '.pai-prompt-history.txt'
+    session = PromptSession(history=FileHistory(str(history)))  # type: ignore
+    multiline = False
+    messages: list[ModelMessage] = []
+
+    while True:
+        try:
+            auto_suggest = CustomAutoSuggest(['/markdown', '/multiline', '/exit'])
+            text = cast(str, session.prompt('pai ➤ ', auto_suggest=auto_suggest, multiline=multiline))
+        except (KeyboardInterrupt, EOFError):
+            return 0
+
+        if not text.strip():
+            continue
+
+        ident_prompt = text.lower().strip(' ').replace(' ', '-').lstrip(' ')
+        if ident_prompt == '/markdown':
+            try:
+                parts = messages[-1].parts
+            except IndexError:
+                console.print('[dim]No markdown output available.[/dim]')
+                continue
+            for part in parts:
+                if part.part_kind == 'text':
+                    last_content = part.content
+                    console.print('[dim]Last markdown output of last question:[/dim]\n')
+                    console.print(Syntax(last_content, lexer='markdown', background_color='default'))
+
+            continue
+        if ident_prompt == '/multiline':
+            multiline = not multiline
+            if multiline:
+                console.print(
+                    'Enabling multiline mode. '
+                    '[dim]Press [Meta+Enter] or [Esc] followed by [Enter] to accept input.[/dim]'
+                )
+            else:
+                console.print('Disabling multiline mode.')
+            continue
+        if ident_prompt == '/exit':
+            console.print('[dim]Exiting…[/dim]')
+            return 0
+
+        try:
+            messages = asyncio.run(ask_agent(agent, text, stream, console, messages))
+        except KeyboardInterrupt:
+            return 0
+
+
+async def ask_agent(
+    agent: Agent,
+    prompt: str,
+    stream: bool,
+    console: Console,
+    messages: list[ModelMessage] | None = None,
+) -> list[ModelMessage]:  # pragma: no cover
+    status: None | Status = Status('[dim]Working on it…[/dim]', console=console)
+    live = Live('', refresh_per_second=15, console=console)
+    status.start()
+
+    async with agent.iter(prompt, message_history=messages) as agent_run:
+        console.print('\nResponse:', style='green')
+
+        content: str = ''
+        interrupted = False
+        try:
+            node = agent_run.next_node
+            while not isinstance(node, End):
+                node = await agent_run.next(node)
+                if Agent.is_model_request_node(node):
+                    async with node.stream(agent_run.ctx) as handle_stream:
+                        # NOTE(Marcelo): It took me a lot of time to figure out how to stop `status` and start `live`
+                        # in a context manager, so I had to do it manually with `stop` and `start` methods.
+                        # PR welcome to simplify this code.
+                        if status is not None:
+                            status.stop()
+                            status = None
+                        if not live.is_started:
+                            live.start()
+                        async for event in handle_stream:
+                            if isinstance(event, PartDeltaEvent) and isinstance(event.delta, TextPartDelta):
+                                if stream:
+                                    content += event.delta.content_delta
+                                    live.update(Markdown(content))
+        except KeyboardInterrupt:
+            interrupted = True
+        finally:
+            live.stop()
+
+        if interrupted:
+            console.print('[dim]Interrupted[/dim]')
+
+        assert agent_run.result
+        if not stream:
+            content = agent_run.result.data
+            console.print(Markdown(content))
+        return agent_run.result.all_messages()
+
+
+class CustomAutoSuggest(AutoSuggestFromHistory):
+    def __init__(self, special_suggestions: list[str] | None = None):  # pragma: no cover
+        super().__init__()
+        self.special_suggestions = special_suggestions or []
+
+    def get_suggestion(self, buffer: Buffer, document: Document) -> Suggestion | None:  # pragma: no cover
+        # Get the suggestion from history
+        suggestion = super().get_suggestion(buffer, document)
+
+        # Check for custom suggestions
+        text = document.text_before_cursor.strip()
+        for special in self.special_suggestions:
+            if special.startswith(text):
+                return Suggestion(special[len(text) :])
+        return suggestion
+
+
+def app():  # pragma: no cover
+    sys.exit(cli())

{pydantic_ai_slim-0.0.32 → pydantic_ai_slim-0.0.34}/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.32 → pydantic_ai_slim-0.0.34}/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.32 → pydantic_ai_slim-0.0.34}/pydantic_ai/agent.py

@@ -6,7 +6,7 @@ from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
 from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
 from copy import deepcopy
 from types import FrameType
-from typing import Any, Callable, Generic, cast, final, overload
+from typing import Any, Callable, ClassVar, Generic, cast, final, overload
 
 from opentelemetry.trace import NoOpTracer, use_span
 from typing_extensions import TypeGuard, TypeVar, deprecated
@@ -25,7 +25,7 @@ from . import (
     result,
     usage as _usage,
 )
-from .models.instrumented import InstrumentedModel
+from .models.instrumented import InstrumentationSettings, InstrumentedModel
 from .result import FinalResult, ResultDataT, StreamedRunResult
 from .settings import ModelSettings, merge_model_settings
 from .tools import (
@@ -56,6 +56,7 @@ __all__ = (
     'CallToolsNode',
     'ModelRequestNode',
     'UserPromptNode',
+    'InstrumentationSettings',
 )
 
 
@@ -112,8 +113,10 @@ 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."""
+    instrument: InstrumentationSettings | bool | None
+    """Options to automatically instrument with OpenTelemetry."""
+
+    _instrument_default: ClassVar[InstrumentationSettings | bool] = False
 
     _deps_type: type[AgentDepsT] = dataclasses.field(repr=False)
     _result_tool_name: str = dataclasses.field(repr=False)
@@ -147,7 +150,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
-        instrument: bool = False,
+        instrument: InstrumentationSettings | bool | None = None,
     ):
         """Create an agent.
 
@@ -177,7 +180,12 @@ 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.
+            instrument: Set to True to automatically instrument with OpenTelemetry,
+                which will use Logfire if it's configured.
+                Set to an instance of [`InstrumentationSettings`][pydantic_ai.agent.InstrumentationSettings] to customize.
+                If this isn't set, then the last value set by
+                [`Agent.instrument_all()`][pydantic_ai.Agent.instrument_all]
+                will be used, which defaults to False.
         """
         if model is None or defer_model_check:
             self.model = model
@@ -213,6 +221,11 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
             else:
                 self._register_tool(Tool(tool))
 
+    @staticmethod
+    def instrument_all(instrument: InstrumentationSettings | bool = True) -> None:
+        """Set the instrumentation options for all agents where `instrument` is not set."""
+        Agent._instrument_default = instrument
+
     @overload
     async def run(
         self,
@@ -422,7 +435,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         usage_limits = usage_limits or _usage.UsageLimits()
 
         if isinstance(model_used, InstrumentedModel):
-            tracer = model_used.tracer
+            tracer = model_used.options.tracer
         else:
             tracer = NoOpTracer()
         agent_name = self.name or 'agent'
@@ -1119,8 +1132,15 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         else:
             raise exceptions.UserError('`model` must be set either when creating the agent or when calling it.')
 
-        if self.instrument and not isinstance(model_, InstrumentedModel):
-            model_ = InstrumentedModel(model_)
+        instrument = self.instrument
+        if instrument is None:
+            instrument = self._instrument_default
+
+        if instrument and not isinstance(model_, InstrumentedModel):
+            if instrument is True:
+                instrument = InstrumentationSettings()
+
+            model_ = InstrumentedModel(model_, instrument)
 
         return model_
 

{pydantic_ai_slim-0.0.32 → pydantic_ai_slim-0.0.34}/pydantic_ai/messages.py

@@ -189,7 +189,10 @@ class ToolReturnPart:
             return {'return_value': tool_return_ta.dump_python(self.content, mode='json')}
 
     def otel_event(self) -> Event:
-        return Event('gen_ai.tool.message', body={'content': self.content, 'role': 'tool', 'id': self.tool_call_id})
+        return Event(
+            'gen_ai.tool.message',
+            body={'content': self.content, 'role': 'tool', 'id': self.tool_call_id, 'name': self.tool_name},
+        )
 
 
 error_details_ta = pydantic.TypeAdapter(list[pydantic_core.ErrorDetails], config=pydantic.ConfigDict(defer_build=True))
@@ -244,7 +247,13 @@ class RetryPromptPart:
             return Event('gen_ai.user.message', body={'content': self.model_response(), 'role': 'user'})
         else:
             return Event(
-                'gen_ai.tool.message', body={'content': self.model_response(), 'role': 'tool', 'id': self.tool_call_id}
+                'gen_ai.tool.message',
+                body={
+                    'content': self.model_response(),
+                    'role': 'tool',
+                    'id': self.tool_call_id,
+                    'name': self.tool_name,
+                },
             )
 
 

{pydantic_ai_slim-0.0.32 → pydantic_ai_slim-0.0.34}/pydantic_ai/models/__init__.py

@@ -49,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',
@@ -320,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.32 → pydantic_ai_slim-0.0.34}/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 (
@@ -82,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}:',
@@ -101,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.
@@ -109,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:
@@ -217,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,