diff --git a/pydantic_ai/__init__.py b/pydantic_ai/__init__.py index 5d0832598..1dea86b60 100644 --- a/pydantic_ai/__init__.py +++ b/pydantic_ai/__init__.py @@ -1,7 +1,8 @@ from importlib.metadata import version from .agent import Agent -from .shared import AgentError, CallContext, ModelRetry, UnexpectedModelBehaviour, UserError +from .call_typing import CallContext +from .exceptions import AgentError, ModelRetry, UnexpectedModelBehaviour, UserError __all__ = 'Agent', 'AgentError', 'CallContext', 'ModelRetry', 'UnexpectedModelBehaviour', 'UserError', '__version__' __version__ = version('pydantic_ai') diff --git a/pydantic_ai/_pydantic.py b/pydantic_ai/_pydantic.py index 770434069..7ba1bddc3 100644 --- a/pydantic_ai/_pydantic.py +++ b/pydantic_ai/_pydantic.py @@ -23,7 +23,7 @@ if TYPE_CHECKING: from . import _retriever - from .shared import AgentDeps + from .call_typing import AgentDeps __all__ = 'function_schema', 'LazyTypeAdapter' @@ -118,7 +118,7 @@ def function_schema(either_function: _retriever.RetrieverEitherFunc[AgentDeps, _ var_positional_field = field_name if errors: - from .shared import UserError + from .exceptions import UserError error_details = '\n '.join(errors) raise UserError(f'Error generating schema for {function.__qualname__}:\n {error_details}') @@ -307,7 +307,7 @@ def _infer_docstring_style(doc: str) -> DocstringStyle: def _is_call_ctx(annotation: Any) -> bool: - from .shared import CallContext + from .call_typing import CallContext return annotation is CallContext or ( _typing_extra.is_generic_alias(annotation) and get_origin(annotation) is CallContext diff --git a/pydantic_ai/_result.py b/pydantic_ai/_result.py index bbeae0b15..216a25a2a 100644 --- a/pydantic_ai/_result.py +++ b/pydantic_ai/_result.py @@ -11,8 +11,10 @@ from typing_extensions import Self, TypeAliasType, TypedDict from . import _utils, messages +from .call_typing import AgentDeps, CallContext +from .exceptions import ModelRetry from .messages import LLMToolCalls, ToolCall -from .shared import AgentDeps, CallContext, ModelRetry, ResultData +from .result import ResultData # A function that always takes `ResultData` and returns `ResultData`, # but may or maybe not take `CallInfo` as a first argument, and may or may not be async. diff --git a/pydantic_ai/_retriever.py b/pydantic_ai/_retriever.py index 6054ff288..e96d6c1e7 100644 --- a/pydantic_ai/_retriever.py +++ b/pydantic_ai/_retriever.py @@ -11,7 +11,8 @@ from typing_extensions import Concatenate, ParamSpec from . import _pydantic, _utils, messages -from .shared import AgentDeps, CallContext, ModelRetry +from .call_typing import AgentDeps, CallContext +from .exceptions import ModelRetry # retrieval function parameters P = ParamSpec('P') diff --git a/pydantic_ai/_system_prompt.py b/pydantic_ai/_system_prompt.py index 4ffdf005c..0206bb7ae 100644 --- a/pydantic_ai/_system_prompt.py +++ b/pydantic_ai/_system_prompt.py @@ -6,7 +6,7 @@ from typing import Any, Callable, Generic, Union, cast from . import _utils -from .shared import AgentDeps, CallContext +from .call_typing import AgentDeps, CallContext # A function that may or maybe not take `CallContext` as an argument, and may or may not be async. # Usage `SystemPromptFunc[AgentDeps]` diff --git a/pydantic_ai/agent.py b/pydantic_ai/agent.py index eaffc780f..c68d0a410 100644 --- a/pydantic_ai/agent.py +++ b/pydantic_ai/agent.py @@ -9,8 +9,9 @@ from pydantic import ValidationError from typing_extensions import assert_never -from . import _result, _retriever as _r, _system_prompt, _utils, messages as _messages, models, shared -from .shared import AgentDeps, ResultData +from . import _result, _retriever as _r, _system_prompt, _utils, exceptions, messages as _messages, models, result +from .call_typing import AgentDeps +from .result import ResultData __all__ = 'Agent', 'KnownModelName' KnownModelName = Literal[ @@ -74,7 +75,7 @@ async def run( message_history: list[_messages.Message] | None = None, model: models.Model | KnownModelName | None = None, deps: AgentDeps | None = None, - ) -> shared.RunResult[ResultData]: + ) -> result.RunResult[ResultData]: """Run the agent with a user prompt in async mode. Args: @@ -92,7 +93,7 @@ async def run( model_ = self.model custom_model = None else: - raise shared.UserError('`model` must be set either when creating the agent or when calling it.') + raise exceptions.UserError('`model` must be set either when creating the agent or when calling it.') if deps is None: deps = self._default_deps @@ -115,7 +116,7 @@ async def run( for retriever in self._retrievers.values(): retriever.reset() - cost = shared.Cost() + cost = result.Cost() with _logfire.span( 'agent run {prompt=}', prompt=user_prompt, agent=self, custom_model=custom_model, model_name=model_.name() @@ -139,17 +140,17 @@ async def run( run_span.set_attribute('cost', cost) handle_span.set_attribute('result', left.value) handle_span.message = 'handle model response -> final result' - return shared.RunResult(left.value, cost, messages, new_message_index) + return result.RunResult(left.value, cost, messages, new_message_index) else: tool_responses = either.right handle_span.set_attribute('tool_responses', tool_responses) response_msgs = ' '.join(m.role for m in tool_responses) handle_span.message = f'handle model response -> {response_msgs}' messages.extend(tool_responses) - except (ValidationError, shared.UnexpectedModelBehaviour) as e: + except (ValidationError, exceptions.UnexpectedModelBehaviour) as e: run_span.set_attribute('messages', messages) # noinspection PyTypeChecker - raise shared.AgentError(messages, model_) from e + raise exceptions.AgentError(messages, model_) from e def run_sync( self, @@ -158,7 +159,7 @@ def run_sync( message_history: list[_messages.Message] | None = None, model: models.Model | KnownModelName | None = None, deps: AgentDeps | None = None, - ) -> shared.RunResult[ResultData]: + ) -> result.RunResult[ResultData]: """Run the agent with a user prompt synchronously. This is a convenience method that wraps `self.run` with `asyncio.run()`. @@ -295,7 +296,7 @@ async def _handle_model_response( retriever = self._retrievers.get(call.tool_name) if retriever is None: # should this be a retry error? - raise shared.UnexpectedModelBehaviour(f'Unknown function name: {call.tool_name!r}') + raise exceptions.UnexpectedModelBehaviour(f'Unknown function name: {call.tool_name!r}') coros.append(retriever.run(deps, call)) new_messages = await asyncio.gather(*coros) return _utils.Either(right=new_messages) @@ -312,7 +313,7 @@ async def _validate_result( def _incr_result_retry(self) -> None: self._current_result_retry += 1 if self._current_result_retry > self._max_result_retries: - raise shared.UnexpectedModelBehaviour( + raise exceptions.UnexpectedModelBehaviour( f'Exceeded maximum retries ({self._max_result_retries}) for result validation' ) diff --git a/pydantic_ai/call_typing.py b/pydantic_ai/call_typing.py new file mode 100644 index 000000000..5f8dcb0f9 --- /dev/null +++ b/pydantic_ai/call_typing.py @@ -0,0 +1,17 @@ +from __future__ import annotations as _annotations + +from dataclasses import dataclass +from typing import Generic, TypeVar + +__all__ = 'AgentDeps', 'CallContext' + +AgentDeps = TypeVar('AgentDeps') + + +@dataclass +class CallContext(Generic[AgentDeps]): + """Information about the current call.""" + + deps: AgentDeps + retry: int + tool_name: str | None diff --git a/pydantic_ai/shared.py b/pydantic_ai/exceptions.py similarity index 54% rename from pydantic_ai/shared.py rename to pydantic_ai/exceptions.py index 73aa9c7ee..361643e0a 100644 --- a/pydantic_ai/shared.py +++ b/pydantic_ai/exceptions.py @@ -1,8 +1,7 @@ from __future__ import annotations as _annotations import json -from dataclasses import dataclass -from typing import TYPE_CHECKING, Generic, TypeVar +from typing import TYPE_CHECKING from pydantic import ValidationError @@ -11,85 +10,7 @@ if TYPE_CHECKING: from .models import Model -__all__ = ( - 'AgentDeps', - 'ResultData', - 'Cost', - 'RunResult', - 'ModelRetry', - 'CallContext', - 'AgentError', - 'UserError', - 'UnexpectedModelBehaviour', -) - -AgentDeps = TypeVar('AgentDeps') -ResultData = TypeVar('ResultData') - - -@dataclass -class Cost: - """Cost of a request or run.""" - - request_tokens: int | None = None - response_tokens: int | None = None - total_tokens: int | None = None - details: dict[str, int] | None = None - - def __add__(self, other: Cost) -> Cost: - counts: dict[str, int] = {} - for field in 'request_tokens', 'response_tokens', 'total_tokens': - self_value = getattr(self, field) - other_value = getattr(other, field) - if self_value is not None or other_value is not None: - counts[field] = (self_value or 0) + (other_value or 0) - - details = self.details.copy() if self.details is not None else None - if other.details is not None: - details = details or {} - for key, value in other.details.items(): - details[key] = details.get(key, 0) + value - - return Cost(**counts, details=details or None) - - -@dataclass -class RunResult(Generic[ResultData]): - """Result of a run.""" - - response: ResultData - cost: Cost - _all_messages: list[messages.Message] - _new_message_index: int - - def all_messages(self) -> list[messages.Message]: - """Return the history of messages.""" - # this is a method to be consistent with the other methods - return self._all_messages - - def all_messages_json(self) -> bytes: - """Return the history of messages as JSON bytes.""" - return messages.MessagesTypeAdapter.dump_json(self.all_messages()) - - def new_messages(self) -> list[messages.Message]: - """Return new messages associated with this run. - - System prompts and any messages from older runs are excluded. - """ - return self.all_messages()[self._new_message_index :] - - def new_messages_json(self) -> bytes: - """Return new messages from [new_messages][] as JSON bytes.""" - return messages.MessagesTypeAdapter.dump_json(self.new_messages()) - - -@dataclass -class CallContext(Generic[AgentDeps]): - """Information about the current call.""" - - deps: AgentDeps - retry: int - tool_name: str | None +__all__ = 'ModelRetry', 'AgentError', 'UserError', 'UnexpectedModelBehaviour' class ModelRetry(Exception): diff --git a/pydantic_ai/messages.py b/pydantic_ai/messages.py index 2c9fe55ac..865533138 100644 --- a/pydantic_ai/messages.py +++ b/pydantic_ai/messages.py @@ -28,7 +28,7 @@ class UserPrompt: role: Literal['user'] = 'user' -return_value_object = _pydantic.LazyTypeAdapter(dict[str, Any]) +tool_return_value_object = _pydantic.LazyTypeAdapter(dict[str, Any]) @dataclass @@ -43,14 +43,14 @@ def model_response_str(self) -> str: if isinstance(self.content, str): return self.content else: - content = return_value_object.validate_python(self.content) - return return_value_object.dump_json(content).decode() + content = tool_return_value_object.validate_python(self.content) + return tool_return_value_object.dump_json(content).decode() def model_response_object(self) -> dict[str, Any]: if isinstance(self.content, str): return {'return_value': self.content} else: - return return_value_object.validate_python(self.content) + return tool_return_value_object.validate_python(self.content) @dataclass @@ -113,4 +113,4 @@ class LLMToolCalls: LLMMessage = Union[LLMResponse, LLMToolCalls] Message = Union[SystemPrompt, UserPrompt, ToolReturn, RetryPrompt, LLMMessage] -MessagesTypeAdapter = pydantic.TypeAdapter(list[Annotated[Message, pydantic.Field(discriminator='role')]]) +MessagesTypeAdapter = _pydantic.LazyTypeAdapter(list[Annotated[Message, pydantic.Field(discriminator='role')]]) diff --git a/pydantic_ai/models/__init__.py b/pydantic_ai/models/__init__.py index 212c3f464..92402a42e 100644 --- a/pydantic_ai/models/__init__.py +++ b/pydantic_ai/models/__init__.py @@ -18,7 +18,7 @@ if TYPE_CHECKING: from .._utils import ObjectJsonSchema from ..agent import KnownModelName - from ..shared import Cost + from ..result import Cost class Model(ABC): @@ -73,7 +73,7 @@ def infer_model(model: Model | KnownModelName) -> Model: # noinspection PyTypeChecker return GeminiModel(model) # pyright: ignore[reportArgumentType] else: - from ..shared import UserError + from ..exceptions import UserError raise UserError(f'Unknown model: {model}') diff --git a/pydantic_ai/models/function.py b/pydantic_ai/models/function.py index f2adfa48b..2ba3a6aea 100644 --- a/pydantic_ai/models/function.py +++ b/pydantic_ai/models/function.py @@ -6,7 +6,7 @@ from typing_extensions import TypeAlias -from .. import shared +from .. import result from ..messages import LLMMessage, Message from . import AbstractToolDefinition, AgentModel, Model @@ -58,5 +58,5 @@ class FunctionAgentModel(AgentModel): function: FunctionDef agent_info: AgentInfo - async def request(self, messages: list[Message]) -> tuple[LLMMessage, shared.Cost]: - return self.function(messages, self.agent_info), shared.Cost() + async def request(self, messages: list[Message]) -> tuple[LLMMessage, result.Cost]: + return self.function(messages, self.agent_info), result.Cost() diff --git a/pydantic_ai/models/gemini.py b/pydantic_ai/models/gemini.py index 5ae51499a..8db2568e0 100644 --- a/pydantic_ai/models/gemini.py +++ b/pydantic_ai/models/gemini.py @@ -27,7 +27,7 @@ from pydantic import Field from typing_extensions import assert_never -from .. import _pydantic, _utils, shared +from .. import _pydantic, _utils, exceptions, result from ..messages import ( ArgsObject, LLMMessage, @@ -67,7 +67,7 @@ def __init__( if env_api_key := os.getenv('GEMINI_API_KEY'): api_key = env_api_key else: - raise shared.UserError('API key must be provided or set in the GEMINI_API_KEY environment variable') + raise exceptions.UserError('API key must be provided or set in the GEMINI_API_KEY environment variable') self.api_key = api_key self.http_client = http_client or cached_async_http_client() self.url_template = url_template @@ -109,7 +109,7 @@ class GeminiAgentModel(AgentModel): tool_config: _GeminiToolConfig | None url_template: str - async def request(self, messages: list[Message]) -> tuple[LLMMessage, shared.Cost]: + async def request(self, messages: list[Message]) -> tuple[LLMMessage, result.Cost]: response = await self.make_request(messages) return self.process_response(response), response.usage_metadata.as_cost() @@ -138,7 +138,7 @@ async def make_request(self, messages: list[Message]) -> _GeminiResponse: url = self.url_template.format(model=self.model_name) r = await self.http_client.post(url, content=request_json, headers=headers) if r.status_code != 200: - raise shared.UnexpectedModelBehaviour(f'Unexpected response from gemini {r.status_code}', r.text) + raise exceptions.UnexpectedModelBehaviour(f'Unexpected response from gemini {r.status_code}', r.text) return _gemini_response_ta.validate_json(r.content) @staticmethod @@ -153,7 +153,7 @@ def process_response(response: _GeminiResponse) -> LLMMessage: parts = cast(list[_GeminiTextPart], parts) return LLMResponse(content=''.join(part.text for part in parts)) else: - raise shared.UnexpectedModelBehaviour( + raise exceptions.UnexpectedModelBehaviour( f'Unexpected response from Gemini, expected all parts to be function calls or text, got: {parts!r}' ) @@ -366,11 +366,11 @@ class _GeminiUsageMetaData: total_token_count: Annotated[int, Field(alias='totalTokenCount')] cached_content_token_count: Annotated[int | None, Field(alias='cachedContentTokenCount')] = None - def as_cost(self) -> shared.Cost: + def as_cost(self) -> result.Cost: details: dict[str, int] = {} if self.cached_content_token_count is not None: details['cached_content_token_count'] = self.cached_content_token_count - return shared.Cost( + return result.Cost( request_tokens=self.prompt_token_count, response_tokens=self.candidates_token_count, total_tokens=self.total_token_count, @@ -430,7 +430,7 @@ def _simplify(self, schema: dict[str, Any], refs_stack: tuple[str, ...]) -> None # noinspection PyTypeChecker key = re.sub(r'^#/\$defs/', '', ref) if key in refs_stack: - raise shared.UserError('Recursive `$ref`s in JSON Schema are not supported by Gemini') + raise exceptions.UserError('Recursive `$ref`s in JSON Schema are not supported by Gemini') refs_stack += (key,) schema_def = self.defs[key] self._simplify(schema_def, refs_stack) @@ -451,7 +451,7 @@ def _simplify(self, schema: dict[str, Any], refs_stack: tuple[str, ...]) -> None def _object(self, schema: dict[str, Any], refs_stack: tuple[str, ...]) -> None: ad_props = schema.pop('additionalProperties', None) if ad_props: - raise shared.UserError('Additional properties in JSON Schema are not supported by Gemini') + raise exceptions.UserError('Additional properties in JSON Schema are not supported by Gemini') if properties := schema.get('properties'): # pragma: no branch for value in properties.values(): diff --git a/pydantic_ai/models/openai.py b/pydantic_ai/models/openai.py index 7fb41a7b5..38f72cfef 100644 --- a/pydantic_ai/models/openai.py +++ b/pydantic_ai/models/openai.py @@ -10,7 +10,7 @@ from openai.types import ChatModel, chat from typing_extensions import assert_never -from .. import shared +from .. import result from ..messages import ( ArgsJson, LLMMessage, @@ -84,7 +84,7 @@ class OpenAIAgentModel(AgentModel): allow_text_result: bool tools: list[chat.ChatCompletionToolParam] - async def request(self, messages: list[Message]) -> tuple[LLMMessage, shared.Cost]: + async def request(self, messages: list[Message]) -> tuple[LLMMessage, result.Cost]: response = await self.completions_create(messages) return self.process_response(response), _map_cost(response) @@ -174,17 +174,17 @@ def _map_tool_call(t: ToolCall) -> chat.ChatCompletionMessageToolCallParam: ) -def _map_cost(response: chat.ChatCompletion) -> shared.Cost: +def _map_cost(response: chat.ChatCompletion) -> result.Cost: usage = response.usage if usage is None: - return shared.Cost() + return result.Cost() 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 shared.Cost( + return result.Cost( request_tokens=usage.prompt_tokens, response_tokens=usage.completion_tokens, total_tokens=usage.total_tokens, diff --git a/pydantic_ai/models/test.py b/pydantic_ai/models/test.py index e428e84d9..1d0ee00e3 100644 --- a/pydantic_ai/models/test.py +++ b/pydantic_ai/models/test.py @@ -14,7 +14,7 @@ import pydantic_core -from .. import _utils, shared +from .. import _utils, result from ..messages import LLMMessage, LLMResponse, LLMToolCalls, Message, RetryPrompt, ToolCall, ToolReturn from . import AbstractToolDefinition, AgentModel, Model @@ -102,8 +102,8 @@ class TestAgentModel(AgentModel): step: int = 0 last_message_count: int = 0 - async def request(self, messages: list[Message]) -> tuple[LLMMessage, shared.Cost]: - cost = shared.Cost() + async def request(self, messages: list[Message]) -> tuple[LLMMessage, result.Cost]: + cost = result.Cost() if self.step == 0: calls = [ToolCall.from_object(name, self.gen_retriever_args(args)) for name, args in self.retriever_calls] self.step += 1 diff --git a/pydantic_ai/result.py b/pydantic_ai/result.py new file mode 100644 index 000000000..e236be9e0 --- /dev/null +++ b/pydantic_ai/result.py @@ -0,0 +1,66 @@ +from __future__ import annotations as _annotations + +from dataclasses import dataclass +from typing import Generic, TypeVar + +from . import messages + +__all__ = 'ResultData', 'Cost', 'RunResult' + +ResultData = TypeVar('ResultData') + + +@dataclass +class Cost: + """Cost of a request or run.""" + + request_tokens: int | None = None + response_tokens: int | None = None + total_tokens: int | None = None + details: dict[str, int] | None = None + + def __add__(self, other: Cost) -> Cost: + counts: dict[str, int] = {} + for field in 'request_tokens', 'response_tokens', 'total_tokens': + self_value = getattr(self, field) + other_value = getattr(other, field) + if self_value is not None or other_value is not None: + counts[field] = (self_value or 0) + (other_value or 0) + + details = self.details.copy() if self.details is not None else None + if other.details is not None: + details = details or {} + for key, value in other.details.items(): + details[key] = details.get(key, 0) + value + + return Cost(**counts, details=details or None) + + +@dataclass +class RunResult(Generic[ResultData]): + """Result of a run.""" + + response: ResultData + cost: Cost + _all_messages: list[messages.Message] + _new_message_index: int + + def all_messages(self) -> list[messages.Message]: + """Return the history of messages.""" + # this is a method to be consistent with the other methods + return self._all_messages + + def all_messages_json(self) -> bytes: + """Return the history of messages as JSON bytes.""" + return messages.MessagesTypeAdapter.dump_json(self.all_messages()) + + def new_messages(self) -> list[messages.Message]: + """Return new messages associated with this run. + + System prompts and any messages from older runs are excluded. + """ + return self.all_messages()[self._new_message_index :] + + def new_messages_json(self) -> bytes: + """Return new messages from [new_messages][] as JSON bytes.""" + return messages.MessagesTypeAdapter.dump_json(self.new_messages()) diff --git a/tests/conftest.py b/tests/conftest.py index 5ced3e62c..ff86e5af1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -15,7 +15,7 @@ from _pytest.assertion.rewrite import AssertionRewritingHook from typing_extensions import TypeAlias -__all__ = 'IsNow', 'TestEnv' +__all__ = 'IsNow', 'TestEnv', 'ClientWithHandler' if TYPE_CHECKING: diff --git a/tests/models/test_gemini.py b/tests/models/test_gemini.py index 53de14d92..5b800cafc 100644 --- a/tests/models/test_gemini.py +++ b/tests/models/test_gemini.py @@ -1,3 +1,4 @@ +# pyright: reportPrivateUsage=false from __future__ import annotations as _annotations import json @@ -25,19 +26,19 @@ ) from pydantic_ai.models.gemini import ( GeminiModel, - _gemini_response_ta, # pyright: ignore[reportPrivateUsage] - _GeminiCandidates, # pyright: ignore[reportPrivateUsage] - _GeminiContent, # pyright: ignore[reportPrivateUsage] - _GeminiFunction, # pyright: ignore[reportPrivateUsage] - _GeminiFunctionCallingConfig, # pyright: ignore[reportPrivateUsage] - _GeminiFunctionCallPart, # pyright: ignore[reportPrivateUsage] - _GeminiResponse, # pyright: ignore[reportPrivateUsage] - _GeminiTextPart, # pyright: ignore[reportPrivateUsage] - _GeminiToolConfig, # pyright: ignore[reportPrivateUsage] - _GeminiTools, # pyright: ignore[reportPrivateUsage] - _GeminiUsageMetaData, # pyright: ignore[reportPrivateUsage] + _gemini_response_ta, + _GeminiCandidates, + _GeminiContent, + _GeminiFunction, + _GeminiFunctionCallingConfig, + _GeminiFunctionCallPart, + _GeminiResponse, + _GeminiTextPart, + _GeminiToolConfig, + _GeminiTools, + _GeminiUsageMetaData, ) -from pydantic_ai.shared import Cost +from pydantic_ai.result import Cost from tests.conftest import ClientWithHandler, IsNow, TestEnv pytestmark = pytest.mark.anyio diff --git a/tests/models/test_openai.py b/tests/models/test_openai.py index 60e9b9a92..38162026e 100644 --- a/tests/models/test_openai.py +++ b/tests/models/test_openai.py @@ -25,7 +25,7 @@ UserPrompt, ) from pydantic_ai.models.openai import OpenAIModel -from pydantic_ai.shared import Cost +from pydantic_ai.result import Cost from tests.conftest import IsNow pytestmark = pytest.mark.anyio diff --git a/tests/test_agent.py b/tests/test_agent.py index b225f2493..3a408e353 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -22,7 +22,7 @@ ) from pydantic_ai.models.function import AgentInfo, FunctionModel from pydantic_ai.models.test import TestModel -from pydantic_ai.shared import Cost, RunResult +from pydantic_ai.result import Cost, RunResult from tests.conftest import IsNow