diff --git a/pydantic_ai_slim/pydantic_ai/messages.py b/pydantic_ai_slim/pydantic_ai/messages.py index 379d70efd..2bea7dd87 100644 --- a/pydantic_ai_slim/pydantic_ai/messages.py +++ b/pydantic_ai_slim/pydantic_ai/messages.py @@ -5,6 +5,7 @@ from collections.abc import Sequence from dataclasses import dataclass, field, replace from datetime import datetime +from functools import lru_cache from mimetypes import guess_type from typing import TYPE_CHECKING, Annotated, Any, Literal, Union, cast, overload @@ -312,7 +313,7 @@ def __init__( def _infer_media_type(self) -> str: """Return the media type of the document, based on the url.""" - type_, _ = guess_type(self.url) + type_, _ = self._guess_type_cached(self.url) if type_ is None: raise ValueError(f'Unknown document file extension: {self.url}') return type_ @@ -329,6 +330,11 @@ def format(self) -> DocumentFormat: except KeyError as e: raise ValueError(f'Unknown document media type: {media_type}') from e + @staticmethod + @lru_cache(maxsize=1024) + def _guess_type_cached(url: str): + return guess_type(url) + @dataclass(repr=False) class BinaryContent: diff --git a/pydantic_ai_slim/pydantic_ai/profiles/google.py b/pydantic_ai_slim/pydantic_ai/profiles/google.py index 9178d7dd4..0c8b1a3cc 100644 --- a/pydantic_ai_slim/pydantic_ai/profiles/google.py +++ b/pydantic_ai_slim/pydantic_ai/profiles/google.py @@ -6,7 +6,7 @@ from . import ModelProfile from ._json_schema import JsonSchema, JsonSchemaTransformer - +import time def google_model_profile(model_name: str) -> ModelProfile | None: """Get the model profile for a Google model.""" @@ -32,6 +32,7 @@ def __init__(self, schema: JsonSchema, *, strict: bool | None = None): super().__init__(schema, strict=strict, prefer_inlined_defs=True, simplify_nullable_unions=True) def transform(self, schema: JsonSchema) -> JsonSchema: + time.sleep(0.001) # Note: we need to remove `additionalProperties: False` since it is currently mishandled by Gemini additional_properties = schema.pop( 'additionalProperties', None