pydantic-ai-slim 0.0.12__py3-none-any.whl → 0.0.14__py3-none-any.whl

pydantic_ai/models/openai.py

@@ -4,23 +4,28 @@ from collections.abc import AsyncIterator, Iterable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
+from itertools import chain
 from typing import Literal, Union, overload
 
 from httpx import AsyncClient as AsyncHTTPClient
 from typing_extensions import assert_never
 
 from .. import UnexpectedModelBehavior, _utils, result
+from .._utils import guard_tool_call_id as _guard_tool_call_id
 from ..messages import (
-    ArgsJson,
-    Message,
-    ModelAnyResponse,
-    ModelStructuredResponse,
-    ModelTextResponse,
-    RetryPrompt,
-    ToolCall,
-    ToolReturn,
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    ModelResponsePart,
+    RetryPromptPart,
+    SystemPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
 )
-from ..result import Cost
+from ..result import Usage
+from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
     AgentModel,
@@ -40,7 +45,7 @@ try:
 except ImportError as _import_error:
     raise ImportError(
         'Please install `openai` to use the OpenAI model, '
-        "you can use the `openai` optional group — `pip install 'pydantic-ai[openai]'`"
+        "you can use the `openai` optional group — `pip install 'pydantic-ai-slim[openai]'`"
     ) from _import_error
 
 OpenAIModelName = Union[ChatModel, str]
@@ -66,6 +71,7 @@ class OpenAIModel(Model):
         self,
         model_name: OpenAIModelName,
         *,
+        base_url: str | None = None,
         api_key: str | None = None,
         openai_client: AsyncOpenAI | None = None,
         http_client: AsyncHTTPClient | None = None,
@@ -76,22 +82,25 @@ class OpenAIModel(Model):
             model_name: The name of the OpenAI model to use. List of model names available
                 [here](https://github.com/openai/openai-python/blob/v1.54.3/src/openai/types/chat_model.py#L7)
                 (Unfortunately, despite being ask to do so, OpenAI do not provide `.inv` files for their API).
+            base_url: The base url for the OpenAI requests. If not provided, the `OPENAI_BASE_URL` environment variable
+                will be used if available. Otherwise, defaults to OpenAI's base url.
             api_key: The API key to use for authentication, if not provided, the `OPENAI_API_KEY` environment variable
                 will be used if available.
             openai_client: An existing
                 [`AsyncOpenAI`](https://github.com/openai/openai-python?tab=readme-ov-file#async-usage)
-                client to use, if provided, `api_key` and `http_client` must be `None`.
+                client to use. If provided, `base_url`, `api_key`, and `http_client` must be `None`.
             http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
         """
         self.model_name: OpenAIModelName = model_name
         if openai_client is not None:
             assert http_client is None, 'Cannot provide both `openai_client` and `http_client`'
+            assert base_url is None, 'Cannot provide both `openai_client` and `base_url`'
             assert api_key is None, 'Cannot provide both `openai_client` and `api_key`'
             self.client = openai_client
         elif http_client is not None:
-            self.client = AsyncOpenAI(api_key=api_key, http_client=http_client)
+            self.client = AsyncOpenAI(base_url=base_url, api_key=api_key, http_client=http_client)
         else:
-            self.client = AsyncOpenAI(api_key=api_key, http_client=cached_async_http_client())
+            self.client = AsyncOpenAI(base_url=base_url, api_key=api_key, http_client=cached_async_http_client())
 
     async def agent_model(
         self,
@@ -135,28 +144,34 @@ class OpenAIAgentModel(AgentModel):
     allow_text_result: bool
     tools: list[chat.ChatCompletionToolParam]
 
-    async def request(self, messages: list[Message]) -> tuple[ModelAnyResponse, result.Cost]:
-        response = await self._completions_create(messages, False)
-        return self._process_response(response), _map_cost(response)
+    async def request(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> tuple[ModelResponse, result.Usage]:
+        response = await self._completions_create(messages, False, model_settings)
+        return self._process_response(response), _map_usage(response)
 
     @asynccontextmanager
-    async def request_stream(self, messages: list[Message]) -> AsyncIterator[EitherStreamedResponse]:
-        response = await self._completions_create(messages, True)
+    async def request_stream(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> AsyncIterator[EitherStreamedResponse]:
+        response = await self._completions_create(messages, True, model_settings)
         async with response:
             yield await self._process_streamed_response(response)
 
     @overload
     async def _completions_create(
-        self, messages: list[Message], stream: Literal[True]
+        self, messages: list[ModelMessage], stream: Literal[True], model_settings: ModelSettings | None
     ) -> AsyncStream[ChatCompletionChunk]:
         pass
 
     @overload
-    async def _completions_create(self, messages: list[Message], stream: Literal[False]) -> chat.ChatCompletion:
+    async def _completions_create(
+        self, messages: list[ModelMessage], stream: Literal[False], model_settings: ModelSettings | None
+    ) -> chat.ChatCompletion:
         pass
 
     async def _completions_create(
-        self, messages: list[Message], stream: bool
+        self, messages: list[ModelMessage], stream: bool, model_settings: ModelSettings | None
     ) -> chat.ChatCompletion | AsyncStream[ChatCompletionChunk]:
         # standalone function to make it easier to override
         if not self.tools:
@@ -166,7 +181,10 @@ class OpenAIAgentModel(AgentModel):
         else:
             tool_choice = 'auto'
 
-        openai_messages = [self._map_message(m) for m in messages]
+        openai_messages = list(chain(*(self._map_message(m) for m in messages)))
+
+        model_settings = model_settings or {}
+
         return await self.client.chat.completions.create(
             model=self.model_name,
             messages=openai_messages,
@@ -176,27 +194,30 @@ class OpenAIAgentModel(AgentModel):
             tool_choice=tool_choice or NOT_GIVEN,
             stream=stream,
             stream_options={'include_usage': True} if stream else NOT_GIVEN,
+            max_tokens=model_settings.get('max_tokens', NOT_GIVEN),
+            temperature=model_settings.get('temperature', NOT_GIVEN),
+            top_p=model_settings.get('top_p', NOT_GIVEN),
+            timeout=model_settings.get('timeout', NOT_GIVEN),
         )
 
     @staticmethod
-    def _process_response(response: chat.ChatCompletion) -> ModelAnyResponse:
+    def _process_response(response: chat.ChatCompletion) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         timestamp = datetime.fromtimestamp(response.created, tz=timezone.utc)
         choice = response.choices[0]
+        items: list[ModelResponsePart] = []
+        if choice.message.content is not None:
+            items.append(TextPart(choice.message.content))
         if choice.message.tool_calls is not None:
-            return ModelStructuredResponse(
-                [ToolCall.from_json(c.function.name, c.function.arguments, c.id) for c in choice.message.tool_calls],
-                timestamp=timestamp,
-            )
-        else:
-            assert choice.message.content is not None, choice
-            return ModelTextResponse(choice.message.content, timestamp=timestamp)
+            for c in choice.message.tool_calls:
+                items.append(ToolCallPart.from_raw_args(c.function.name, c.function.arguments, c.id))
+        return ModelResponse(items, timestamp=timestamp)
 
     @staticmethod
     async def _process_streamed_response(response: AsyncStream[ChatCompletionChunk]) -> EitherStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         timestamp: datetime | None = None
-        start_cost = Cost()
+        start_usage = Usage()
         # the first chunk may contain enough information so we iterate until we get either `tool_calls` or `content`
         while True:
             try:
@@ -205,63 +226,73 @@ class OpenAIAgentModel(AgentModel):
                 raise UnexpectedModelBehavior('Streamed response ended without content or tool calls') from e
 
             timestamp = timestamp or datetime.fromtimestamp(chunk.created, tz=timezone.utc)
-            start_cost += _map_cost(chunk)
+            start_usage += _map_usage(chunk)
 
             if chunk.choices:
                 delta = chunk.choices[0].delta
 
                 if delta.content is not None:
-                    return OpenAIStreamTextResponse(delta.content, response, timestamp, start_cost)
+                    return OpenAIStreamTextResponse(delta.content, response, timestamp, start_usage)
                 elif delta.tool_calls is not None:
                     return OpenAIStreamStructuredResponse(
                         response,
                         {c.index: c for c in delta.tool_calls},
                         timestamp,
-                        start_cost,
+                        start_usage,
                     )
                 # else continue until we get either delta.content or delta.tool_calls
 
-    @staticmethod
-    def _map_message(message: Message) -> chat.ChatCompletionMessageParam:
+    @classmethod
+    def _map_message(cls, message: ModelMessage) -> Iterable[chat.ChatCompletionMessageParam]:
         """Just maps a `pydantic_ai.Message` to a `openai.types.ChatCompletionMessageParam`."""
-        if message.role == 'system':
-            # SystemPrompt ->
-            return chat.ChatCompletionSystemMessageParam(role='system', content=message.content)
-        elif message.role == 'user':
-            # UserPrompt ->
-            return chat.ChatCompletionUserMessageParam(role='user', content=message.content)
-        elif message.role == 'tool-return':
-            # ToolReturn ->
-            return chat.ChatCompletionToolMessageParam(
-                role='tool',
-                tool_call_id=_guard_tool_id(message),
-                content=message.model_response_str(),
-            )
-        elif message.role == 'retry-prompt':
-            # RetryPrompt ->
-            if message.tool_name is None:
-                return chat.ChatCompletionUserMessageParam(role='user', content=message.model_response())
-            else:
-                return chat.ChatCompletionToolMessageParam(
-                    role='tool',
-                    tool_call_id=_guard_tool_id(message),
-                    content=message.model_response(),
-                )
-        elif message.role == 'model-text-response':
-            # ModelTextResponse ->
-            return chat.ChatCompletionAssistantMessageParam(role='assistant', content=message.content)
-        elif message.role == 'model-structured-response':
-            assert (
-                message.role == 'model-structured-response'
-            ), f'Expected role to be "llm-tool-calls", got {message.role}'
-            # ModelStructuredResponse ->
-            return chat.ChatCompletionAssistantMessageParam(
-                role='assistant',
-                tool_calls=[_map_tool_call(t) for t in message.calls],
-            )
+        if isinstance(message, ModelRequest):
+            yield from cls._map_user_message(message)
+        elif isinstance(message, ModelResponse):
+            texts: list[str] = []
+            tool_calls: list[chat.ChatCompletionMessageToolCallParam] = []
+            for item in message.parts:
+                if isinstance(item, TextPart):
+                    texts.append(item.content)
+                elif isinstance(item, ToolCallPart):
+                    tool_calls.append(_map_tool_call(item))
+                else:
+                    assert_never(item)
+            message_param = chat.ChatCompletionAssistantMessageParam(role='assistant')
+            if texts:
+                # Note: model responses from this model should only have one text item, so the following
+                # shouldn't merge multiple texts into one unless you switch models between runs:
+                message_param['content'] = '\n\n'.join(texts)
+            if tool_calls:
+                message_param['tool_calls'] = tool_calls
+            yield message_param
         else:
             assert_never(message)
 
+    @classmethod
+    def _map_user_message(cls, message: ModelRequest) -> Iterable[chat.ChatCompletionMessageParam]:
+        for part in message.parts:
+            if isinstance(part, SystemPromptPart):
+                yield chat.ChatCompletionSystemMessageParam(role='system', content=part.content)
+            elif isinstance(part, UserPromptPart):
+                yield chat.ChatCompletionUserMessageParam(role='user', content=part.content)
+            elif isinstance(part, ToolReturnPart):
+                yield chat.ChatCompletionToolMessageParam(
+                    role='tool',
+                    tool_call_id=_guard_tool_call_id(t=part, model_source='OpenAI'),
+                    content=part.model_response_str(),
+                )
+            elif isinstance(part, RetryPromptPart):
+                if part.tool_name is None:
+                    yield chat.ChatCompletionUserMessageParam(role='user', content=part.model_response())
+                else:
+                    yield chat.ChatCompletionToolMessageParam(
+                        role='tool',
+                        tool_call_id=_guard_tool_call_id(t=part, model_source='OpenAI'),
+                        content=part.model_response(),
+                    )
+            else:
+                assert_never(part)
+
 
 @dataclass
 class OpenAIStreamTextResponse(StreamTextResponse):
@@ -270,7 +301,7 @@ class OpenAIStreamTextResponse(StreamTextResponse):
     _first: str | None
     _response: AsyncStream[ChatCompletionChunk]
     _timestamp: datetime
-    _cost: result.Cost
+    _usage: result.Usage
     _buffer: list[str] = field(default_factory=list, init=False)
 
     async def __anext__(self) -> None:
@@ -280,7 +311,7 @@ class OpenAIStreamTextResponse(StreamTextResponse):
             return None
 
         chunk = await self._response.__anext__()
-        self._cost += _map_cost(chunk)
+        self._usage += _map_usage(chunk)
         try:
             choice = chunk.choices[0]
         except IndexError:
@@ -296,8 +327,8 @@ class OpenAIStreamTextResponse(StreamTextResponse):
         yield from self._buffer
         self._buffer.clear()
 
-    def cost(self) -> Cost:
-        return self._cost
+    def usage(self) -> Usage:
+        return self._usage
 
     def timestamp(self) -> datetime:
         return self._timestamp
@@ -310,11 +341,11 @@ class OpenAIStreamStructuredResponse(StreamStructuredResponse):
     _response: AsyncStream[ChatCompletionChunk]
     _delta_tool_calls: dict[int, ChoiceDeltaToolCall]
     _timestamp: datetime
-    _cost: result.Cost
+    _usage: result.Usage
 
     async def __anext__(self) -> None:
         chunk = await self._response.__anext__()
-        self._cost += _map_cost(chunk)
+        self._usage += _map_usage(chunk)
         try:
             choice = chunk.choices[0]
         except IndexError:
@@ -335,48 +366,41 @@ class OpenAIStreamStructuredResponse(StreamStructuredResponse):
             else:
                 self._delta_tool_calls[new.index] = new
 
-    def get(self, *, final: bool = False) -> ModelStructuredResponse:
-        calls: list[ToolCall] = []
+    def get(self, *, final: bool = False) -> ModelResponse:
+        items: list[ModelResponsePart] = []
         for c in self._delta_tool_calls.values():
             if f := c.function:
                 if f.name is not None and f.arguments is not None:
-                    calls.append(ToolCall.from_json(f.name, f.arguments, c.id))
+                    items.append(ToolCallPart.from_raw_args(f.name, f.arguments, c.id))
 
-        return ModelStructuredResponse(calls, timestamp=self._timestamp)
+        return ModelResponse(items, timestamp=self._timestamp)
 
-    def cost(self) -> Cost:
-        return self._cost
+    def usage(self) -> Usage:
+        return self._usage
 
     def timestamp(self) -> datetime:
         return self._timestamp
 
 
-def _guard_tool_id(t: ToolCall | ToolReturn | RetryPrompt) -> str:
-    """Type guard that checks a `tool_id` is not None both for static typing and runtime."""
-    assert t.tool_id is not None, f'OpenAI requires `tool_id` to be set: {t}'
-    return t.tool_id
-
-
-def _map_tool_call(t: ToolCall) -> chat.ChatCompletionMessageToolCallParam:
-    assert isinstance(t.args, ArgsJson), f'Expected ArgsJson, got {t.args}'
+def _map_tool_call(t: ToolCallPart) -> chat.ChatCompletionMessageToolCallParam:
     return chat.ChatCompletionMessageToolCallParam(
-        id=_guard_tool_id(t),
+        id=_guard_tool_call_id(t=t, model_source='OpenAI'),
         type='function',
-        function={'name': t.tool_name, 'arguments': t.args.args_json},
+        function={'name': t.tool_name, 'arguments': t.args_as_json_str()},
     )
 
 
-def _map_cost(response: chat.ChatCompletion | ChatCompletionChunk) -> result.Cost:
+def _map_usage(response: chat.ChatCompletion | ChatCompletionChunk) -> result.Usage:
     usage = response.usage
     if usage is None:
-        return result.Cost()
+        return result.Usage()
     else:
         details: dict[str, int] = {}
         if usage.completion_tokens_details is not None:
             details.update(usage.completion_tokens_details.model_dump(exclude_none=True))
         if usage.prompt_tokens_details is not None:
             details.update(usage.prompt_tokens_details.model_dump(exclude_none=True))
-        return result.Cost(
+        return result.Usage(
             request_tokens=usage.prompt_tokens,
             response_tokens=usage.completion_tokens,
             total_tokens=usage.total_tokens,

pydantic_ai/models/test.py

@@ -9,18 +9,20 @@ from datetime import date, datetime, timedelta
 from typing import Any, Literal
 
 import pydantic_core
+from typing_extensions import assert_never
 
 from .. import _utils
 from ..messages import (
-    Message,
-    ModelAnyResponse,
-    ModelStructuredResponse,
-    ModelTextResponse,
-    RetryPrompt,
-    ToolCall,
-    ToolReturn,
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    RetryPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
 )
-from ..result import Cost
+from ..result import Usage
+from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
     AgentModel,
@@ -29,6 +31,7 @@ from . import (
     StreamStructuredResponse,
     StreamTextResponse,
 )
+from .function import _estimate_string_usage, _estimate_usage  # pyright: ignore[reportPrivateUsage]
 
 
 @dataclass
@@ -127,74 +130,85 @@ class TestAgentModel(AgentModel):
     result_tools: list[ToolDefinition]
     seed: int
 
-    async def request(self, messages: list[Message]) -> tuple[ModelAnyResponse, Cost]:
-        return self._request(messages), Cost()
+    async def request(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> tuple[ModelResponse, Usage]:
+        model_response = self._request(messages, model_settings)
+        usage = _estimate_usage([*messages, model_response])
+        return model_response, usage
 
     @asynccontextmanager
-    async def request_stream(self, messages: list[Message]) -> AsyncIterator[EitherStreamedResponse]:
-        msg = self._request(messages)
-        cost = Cost()
-        if isinstance(msg, ModelTextResponse):
-            yield TestStreamTextResponse(msg.content, cost)
+    async def request_stream(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> AsyncIterator[EitherStreamedResponse]:
+        msg = self._request(messages, model_settings)
+        usage = _estimate_usage(messages)
+
+        # TODO: Rework this once we make StreamTextResponse more general
+        texts: list[str] = []
+        tool_calls: list[ToolCallPart] = []
+        for item in msg.parts:
+            if isinstance(item, TextPart):
+                texts.append(item.content)
+            elif isinstance(item, ToolCallPart):
+                tool_calls.append(item)
+            else:
+                assert_never(item)
+
+        if texts:
+            yield TestStreamTextResponse('\n\n'.join(texts), usage)
         else:
-            yield TestStreamStructuredResponse(msg, cost)
+            yield TestStreamStructuredResponse(msg, usage)
 
     def gen_tool_args(self, tool_def: ToolDefinition) -> Any:
         return _JsonSchemaTestData(tool_def.parameters_json_schema, self.seed).generate()
 
-    def _request(self, messages: list[Message]) -> ModelAnyResponse:
+    def _request(self, messages: list[ModelMessage], model_settings: ModelSettings | None) -> ModelResponse:
         # if there are tools, the first thing we want to do is call all of them
-        if self.tool_calls and not any(m.role == 'model-structured-response' for m in messages):
-            calls = [ToolCall.from_dict(name, self.gen_tool_args(args)) for name, args in self.tool_calls]
-            return ModelStructuredResponse(calls=calls)
-
-        # get messages since the last model response
-        new_messages = _get_new_messages(messages)
-
-        # check if there are any retry prompts, if so retry them
-        new_retry_names = {m.tool_name for m in new_messages if isinstance(m, RetryPrompt)}
-        if new_retry_names:
-            calls = [
-                ToolCall.from_dict(name, self.gen_tool_args(args))
-                for name, args in self.tool_calls
-                if name in new_retry_names
-            ]
-            return ModelStructuredResponse(calls=calls)
+        if self.tool_calls and not any(isinstance(m, ModelResponse) for m in messages):
+            return ModelResponse(
+                parts=[ToolCallPart.from_raw_args(name, self.gen_tool_args(args)) for name, args in self.tool_calls]
+            )
+
+        if messages:
+            last_message = messages[-1]
+            assert isinstance(last_message, ModelRequest), 'Expected last message to be a `ModelRequest`.'
+
+            # check if there are any retry prompts, if so retry them
+            new_retry_names = {p.tool_name for p in last_message.parts if isinstance(p, RetryPromptPart)}
+            if new_retry_names:
+                return ModelResponse(
+                    parts=[
+                        ToolCallPart.from_raw_args(name, self.gen_tool_args(args))
+                        for name, args in self.tool_calls
+                        if name in new_retry_names
+                    ]
+                )
 
         if response_text := self.result.left:
             if response_text.value is None:
                 # build up details of tool responses
                 output: dict[str, Any] = {}
                 for message in messages:
-                    if isinstance(message, ToolReturn):
-                        output[message.tool_name] = message.content
+                    if isinstance(message, ModelRequest):
+                        for part in message.parts:
+                            if isinstance(part, ToolReturnPart):
+                                output[part.tool_name] = part.content
                 if output:
-                    return ModelTextResponse(content=pydantic_core.to_json(output).decode())
+                    return ModelResponse.from_text(pydantic_core.to_json(output).decode())
                 else:
-                    return ModelTextResponse(content='success (no tool calls)')
+                    return ModelResponse.from_text('success (no tool calls)')
             else:
-                return ModelTextResponse(content=response_text.value)
+                return ModelResponse.from_text(response_text.value)
         else:
             assert self.result_tools, 'No result tools provided'
             custom_result_args = self.result.right
             result_tool = self.result_tools[self.seed % len(self.result_tools)]
             if custom_result_args is not None:
-                return ModelStructuredResponse(calls=[ToolCall.from_dict(result_tool.name, custom_result_args)])
+                return ModelResponse(parts=[ToolCallPart.from_raw_args(result_tool.name, custom_result_args)])
             else:
                 response_args = self.gen_tool_args(result_tool)
-                return ModelStructuredResponse(calls=[ToolCall.from_dict(result_tool.name, response_args)])
-
-
-def _get_new_messages(messages: list[Message]) -> list[Message]:
-    last_model_index = None
-    for i, m in enumerate(messages):
-        if m.role in ('model-structured-response', 'model-text-response'):
-            last_model_index = i
-
-    if last_model_index is not None:
-        return messages[last_model_index + 1 :]
-    else:
-        return []
+                return ModelResponse(parts=[ToolCallPart.from_raw_args(result_tool.name, response_args)])
 
 
 @dataclass
@@ -202,7 +216,7 @@ class TestStreamTextResponse(StreamTextResponse):
     """A text response that streams test data."""
 
     _text: str
-    _cost: Cost
+    _usage: Usage
     _iter: Iterator[str] = field(init=False)
     _timestamp: datetime = field(default_factory=_utils.now_utc)
     _buffer: list[str] = field(default_factory=list, init=False)
@@ -217,14 +231,17 @@ class TestStreamTextResponse(StreamTextResponse):
         self._iter = iter(words)
 
     async def __anext__(self) -> None:
-        self._buffer.append(_utils.sync_anext(self._iter))
+        next_str = _utils.sync_anext(self._iter)
+        response_tokens = _estimate_string_usage(next_str)
+        self._usage += Usage(response_tokens=response_tokens, total_tokens=response_tokens)
+        self._buffer.append(next_str)
 
     def get(self, *, final: bool = False) -> Iterable[str]:
         yield from self._buffer
         self._buffer.clear()
 
-    def cost(self) -> Cost:
-        return self._cost
+    def usage(self) -> Usage:
+        return self._usage
 
     def timestamp(self) -> datetime:
         return self._timestamp
@@ -234,19 +251,19 @@ class TestStreamTextResponse(StreamTextResponse):
 class TestStreamStructuredResponse(StreamStructuredResponse):
     """A structured response that streams test data."""
 
-    _structured_response: ModelStructuredResponse
-    _cost: Cost
+    _structured_response: ModelResponse
+    _usage: Usage
     _iter: Iterator[None] = field(default_factory=lambda: iter([None]))
     _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
 
     async def __anext__(self) -> None:
         return _utils.sync_anext(self._iter)
 
-    def get(self, *, final: bool = False) -> ModelStructuredResponse:
+    def get(self, *, final: bool = False) -> ModelResponse:
         return self._structured_response
 
-    def cost(self) -> Cost:
-        return self._cost
+    def usage(self) -> Usage:
+        return self._usage
 
     def timestamp(self) -> datetime:
         return self._timestamp

pydantic_ai/models/vertexai.py

@@ -21,7 +21,7 @@ try:
 except ImportError as _import_error:
     raise ImportError(
         'Please install `google-auth` to use the VertexAI model, '
-        "you can use the `vertexai` optional group — `pip install 'pydantic-ai[vertexai]'`"
+        "you can use the `vertexai` optional group — `pip install 'pydantic-ai-slim[vertexai]'`"
     ) from _import_error
 
 VERTEX_AI_URL_TEMPLATE = (
@@ -114,7 +114,7 @@ class VertexAIModel(Model):
         allow_text_result: bool,
         result_tools: list[ToolDefinition],
     ) -> GeminiAgentModel:
-        url, auth = await self._ainit()
+        url, auth = await self.ainit()
         return GeminiAgentModel(
             http_client=self.http_client,
             model_name=self.model_name,
@@ -125,7 +125,11 @@ class VertexAIModel(Model):
             result_tools=result_tools,
         )
 
-    async def _ainit(self) -> tuple[str, BearerTokenAuth]:
+    async def ainit(self) -> tuple[str, BearerTokenAuth]:
+        """Initialize the model, setting the URL and auth.
+
+        This will raise an error if authentication fails.
+        """
         if self.url is not None and self.auth is not None:
             return self.url, self.auth