pydantic-ai-slim 0.0.22__py3-none-any.whl → 0.0.24__py3-none-any.whl

pydantic_ai/models/vertexai.py

@@ -1,5 +1,7 @@
 from __future__ import annotations as _annotations
 
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta
 from pathlib import Path
@@ -7,11 +9,13 @@ from typing import Literal
 
 from httpx import AsyncClient as AsyncHTTPClient
 
+from .. import usage
 from .._utils import run_in_executor
 from ..exceptions import UserError
-from ..tools import ToolDefinition
-from . import Model, cached_async_http_client, check_allow_model_requests
-from .gemini import GeminiAgentModel, GeminiModelName
+from ..messages import ModelMessage, ModelResponse
+from ..settings import ModelSettings
+from . import ModelRequestParameters, StreamedResponse, cached_async_http_client
+from .gemini import GeminiModel, GeminiModelName
 
 try:
     import google.auth
@@ -52,19 +56,17 @@ The template is used thus:
 
 
 @dataclass(init=False)
-class VertexAIModel(Model):
+class VertexAIModel(GeminiModel):
     """A model that uses Gemini via the `*-aiplatform.googleapis.com` VertexAI API."""
 
-    model_name: GeminiModelName
     service_account_file: Path | str | None
     project_id: str | None
     region: VertexAiRegion
     model_publisher: Literal['google']
-    http_client: AsyncHTTPClient
     url_template: str
 
-    auth: BearerTokenAuth | None
-    url: str | None
+    _model_name: GeminiModelName = field(repr=False)
+    _system: str | None = field(default='google-vertex', repr=False)
 
     # TODO __init__ can be removed once we drop 3.9 and we can set kw_only correctly on the dataclass
     def __init__(
@@ -96,7 +98,7 @@ class VertexAIModel(Model):
                 [`VERTEX_AI_URL_TEMPLATE` docs][pydantic_ai.models.vertexai.VERTEX_AI_URL_TEMPLATE]
                 for more information.
         """
-        self.model_name = model_name
+        self._model_name = model_name
         self.service_account_file = service_account_file
         self.project_id = project_id
         self.region = region
@@ -104,35 +106,16 @@ class VertexAIModel(Model):
         self.http_client = http_client or cached_async_http_client()
         self.url_template = url_template
 
-        self.auth = None
-        self.url = None
+        self._auth = None
+        self._url = None
 
-    async def agent_model(
-        self,
-        *,
-        function_tools: list[ToolDefinition],
-        allow_text_result: bool,
-        result_tools: list[ToolDefinition],
-    ) -> GeminiAgentModel:
-        check_allow_model_requests()
-        url, auth = await self.ainit()
-        return GeminiAgentModel(
-            http_client=self.http_client,
-            model_name=self.model_name,
-            auth=auth,
-            url=url,
-            function_tools=function_tools,
-            allow_text_result=allow_text_result,
-            result_tools=result_tools,
-        )
-
-    async def ainit(self) -> tuple[str, BearerTokenAuth]:
+    async def ainit(self) -> None:
         """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
+        if self._url is not None and self._auth is not None:
+            return
 
         if self.service_account_file is not None:
             creds: BaseCredentials | ServiceAccountCredentials = _creds_from_file(self.service_account_file)
@@ -148,24 +131,45 @@ class VertexAIModel(Model):
                 raise UserError(f'No project_id provided and none found in {creds_source}')
             project_id = creds_project_id
         else:
-            if creds_project_id is not None and self.project_id != creds_project_id:
-                raise UserError(
-                    f'The project_id you provided does not match the one from {creds_source}: '
-                    f'{self.project_id!r} != {creds_project_id!r}'
-                )
             project_id = self.project_id
 
-        self.url = url = self.url_template.format(
+        self._url = self.url_template.format(
             region=self.region,
             project_id=project_id,
             model_publisher=self.model_publisher,
-            model=self.model_name,
+            model=self._model_name,
         )
-        self.auth = auth = BearerTokenAuth(creds)
-        return url, auth
+        self._auth = BearerTokenAuth(creds)
 
-    def name(self) -> str:
-        return f'google-vertex:{self.model_name}'
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> tuple[ModelResponse, usage.Usage]:
+        await self.ainit()
+        return await super().request(messages, model_settings, model_request_parameters)
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> AsyncIterator[StreamedResponse]:
+        await self.ainit()
+        async with super().request_stream(messages, model_settings, model_request_parameters) as value:
+            yield value
+
+    @property
+    def model_name(self) -> GeminiModelName:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str | None:
+        """The system / model provider."""
+        return self._system
 
 
 # pyright: reportUnknownMemberType=false

pydantic_ai/result.py

@@ -286,7 +286,7 @@ class StreamedRunResult(_BaseRunResult[ResultDataT], Generic[AgentDepsT, ResultD
                 await self._marked_completed(
                     _messages.ModelResponse(
                         parts=[_messages.TextPart(combined_validated_text)],
-                        model_name=self._stream_response.model_name(),
+                        model_name=self._stream_response.model_name,
                     )
                 )
 
@@ -347,7 +347,7 @@ class StreamedRunResult(_BaseRunResult[ResultDataT], Generic[AgentDepsT, ResultD
 
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""
-        return self._stream_response.timestamp()
+        return self._stream_response.timestamp
 
     async def validate_structured_result(
         self, message: _messages.ModelResponse, *, allow_partial: bool = False

{pydantic_ai_slim-0.0.22.dist-info → pydantic_ai_slim-0.0.24.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.0.22
+Version: 0.0.24
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>
 License-Expression: MIT
@@ -28,7 +28,7 @@ Requires-Dist: eval-type-backport>=0.2.0
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: logfire-api>=1.2.0
-Requires-Dist: pydantic-graph==0.0.22
+Requires-Dist: pydantic-graph==0.0.24
 Requires-Dist: pydantic>=2.10
 Provides-Extra: anthropic
 Requires-Dist: anthropic>=0.40.0; extra == 'anthropic'
@@ -41,7 +41,7 @@ Requires-Dist: logfire>=2.3; extra == 'logfire'
 Provides-Extra: mistral
 Requires-Dist: mistralai>=1.2.5; extra == 'mistral'
 Provides-Extra: openai
-Requires-Dist: openai>=1.59.0; extra == 'openai'
+Requires-Dist: openai>=1.61.0; extra == 'openai'
 Provides-Extra: vertexai
 Requires-Dist: google-auth>=2.36.0; extra == 'vertexai'
 Requires-Dist: requests>=2.32.3; extra == 'vertexai'

pydantic_ai_slim-0.0.24.dist-info/RECORD

@@ -0,0 +1,30 @@
+pydantic_ai/__init__.py,sha256=FbYetEgT6OO25u2KF5ZnFxKpz5DtnSpfckRXP4mjl8E,489
+pydantic_ai/_agent_graph.py,sha256=_uCNK12sEXmzB-Vk1OaCR-AMARYzEoInjKXZyx7ODjA,33136
+pydantic_ai/_griffe.py,sha256=RYRKiLbgG97QxnazbAwlnc74XxevGHLQet-FGfq9qls,3960
+pydantic_ai/_parts_manager.py,sha256=ARfDQY1_5AIY5rNl_M2fAYHEFCe03ZxdhgjHf9qeIKw,11872
+pydantic_ai/_pydantic.py,sha256=dROz3Hmfdi0C2exq88FhefDRVo_8S3rtkXnoUHzsz0c,8753
+pydantic_ai/_result.py,sha256=tN1pVulf_EM4bkBvpNUWPnUXezLY-sBrJEVCFdy2nLU,10264
+pydantic_ai/_system_prompt.py,sha256=602c2jyle2R_SesOrITBDETZqsLk4BZ8Cbo8yEhmx04,1120
+pydantic_ai/_utils.py,sha256=zfuY3NiPPsSM5J1q2JElfbfIa8S1ONGOlC7M-iyBVso,9430
+pydantic_ai/agent.py,sha256=q9nc_bTUuwxYVrzmUb5NkwOUWQm2XIsxpTuyyAuVdpU,44900
+pydantic_ai/exceptions.py,sha256=eGDKX6bGhgVxXBzu81Sk3iiAkXr0GUtgT7bD5Rxlqpg,2028
+pydantic_ai/format_as_xml.py,sha256=QE7eMlg5-YUMw1_2kcI3h0uKYPZZyGkgXFDtfZTMeeI,4480
+pydantic_ai/messages.py,sha256=kzXn4ZjlX9Sy2KXgFHWYbbwPk7TzTPdztzOJLWEONwU,17101
+pydantic_ai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pydantic_ai/result.py,sha256=HYTkmb-6fzbxXRlsPOyVic_EexEGN4JLfZfVBKpJtfM,18390
+pydantic_ai/settings.py,sha256=ntuWnke9UA18aByDxk9OIhN0tAgOaPdqCEkRf-wlp8Y,3059
+pydantic_ai/tools.py,sha256=lhupwm815lPlFFS79B0P61AyhUYtepA62LbZOCJrPEY,13205
+pydantic_ai/usage.py,sha256=60d9f6M7YEYuKMbqDGDogX4KsA73fhDtWyDXYXoIPaI,4948
+pydantic_ai/models/__init__.py,sha256=WCvo5bI2dWJbKPhdGWID-HFGrZGo3jvtHhgwaDuY75s,13219
+pydantic_ai/models/anthropic.py,sha256=lArkCv0oZQeG1SqyEHlqKUyrcEYzYrS5fkW_ICKlikU,17749
+pydantic_ai/models/cohere.py,sha256=yLUA3bTQg-4NNjpvD2DXkUIMluia8vETChsjYYN13lg,11016
+pydantic_ai/models/function.py,sha256=5AWAFo-AZ0UyX8wBAYc5nd6U6XpOFvnsYr8vVfatEXQ,10333
+pydantic_ai/models/gemini.py,sha256=CGaHve4wQxkGEa1hArgldvOBk79blGdOqdeLxMH_NQM,30846
+pydantic_ai/models/groq.py,sha256=M7P-9LrH1lm4wyIMiytsHf28bwkukOyVzw3S7F9sbic,14612
+pydantic_ai/models/mistral.py,sha256=1DA1AX-V46p7dI75-4aRwQ8TrQxdh5ZQaoMYZO7VRkI,25868
+pydantic_ai/models/openai.py,sha256=kUV13suTp1ueibH2uYWaRCX832GApImdxSiGqNyfWo4,17161
+pydantic_ai/models/test.py,sha256=6X0r78biqnvakiH1nFKYQFNOs4dWrS-yqe0RFBrKW0s,16873
+pydantic_ai/models/vertexai.py,sha256=9Kp_1KMBlbP8_HRJTuFnrkkFmlJ7yFhADQYjxOgIh9Y,9523
+pydantic_ai_slim-0.0.24.dist-info/METADATA,sha256=kSbEzpXrgMTSHX0ojDw-Mtf8QQFCFE5fIJTRZihLe1w,2839
+pydantic_ai_slim-0.0.24.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pydantic_ai_slim-0.0.24.dist-info/RECORD,,

pydantic_ai_slim-0.0.22.dist-info/RECORD

@@ -1,30 +0,0 @@
-pydantic_ai/__init__.py,sha256=FbYetEgT6OO25u2KF5ZnFxKpz5DtnSpfckRXP4mjl8E,489
-pydantic_ai/_agent_graph.py,sha256=IdznyPzgrLmfiw3733GA3gYWOA_6LYIEuXBKTBXOi0A,32942
-pydantic_ai/_griffe.py,sha256=RYRKiLbgG97QxnazbAwlnc74XxevGHLQet-FGfq9qls,3960
-pydantic_ai/_parts_manager.py,sha256=ARfDQY1_5AIY5rNl_M2fAYHEFCe03ZxdhgjHf9qeIKw,11872
-pydantic_ai/_pydantic.py,sha256=dROz3Hmfdi0C2exq88FhefDRVo_8S3rtkXnoUHzsz0c,8753
-pydantic_ai/_result.py,sha256=tN1pVulf_EM4bkBvpNUWPnUXezLY-sBrJEVCFdy2nLU,10264
-pydantic_ai/_system_prompt.py,sha256=602c2jyle2R_SesOrITBDETZqsLk4BZ8Cbo8yEhmx04,1120
-pydantic_ai/_utils.py,sha256=zfuY3NiPPsSM5J1q2JElfbfIa8S1ONGOlC7M-iyBVso,9430
-pydantic_ai/agent.py,sha256=RDKBnHU1yFBMwW9e7_BUYs4aRfbwJ69bpWp_ha2y9M0,44880
-pydantic_ai/exceptions.py,sha256=eGDKX6bGhgVxXBzu81Sk3iiAkXr0GUtgT7bD5Rxlqpg,2028
-pydantic_ai/format_as_xml.py,sha256=QE7eMlg5-YUMw1_2kcI3h0uKYPZZyGkgXFDtfZTMeeI,4480
-pydantic_ai/messages.py,sha256=kzXn4ZjlX9Sy2KXgFHWYbbwPk7TzTPdztzOJLWEONwU,17101
-pydantic_ai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pydantic_ai/result.py,sha256=0beXiUfiSb2kVc68iqLeg2qFQW6-mnTC8TFtiQayPDQ,18394
-pydantic_ai/settings.py,sha256=ntuWnke9UA18aByDxk9OIhN0tAgOaPdqCEkRf-wlp8Y,3059
-pydantic_ai/tools.py,sha256=lhupwm815lPlFFS79B0P61AyhUYtepA62LbZOCJrPEY,13205
-pydantic_ai/usage.py,sha256=60d9f6M7YEYuKMbqDGDogX4KsA73fhDtWyDXYXoIPaI,4948
-pydantic_ai/models/__init__.py,sha256=kVWd_WEK2dThHoG0Yha4t93w4Zaqkjcu9sKJ7nNqX_U,13245
-pydantic_ai/models/anthropic.py,sha256=v2m6zjaezLjtkHO2Rx67rmoN1iOqDeAp9fswJcqRMBA,16967
-pydantic_ai/models/cohere.py,sha256=wyjtD2uwkYYbsXLKwr-flmXSng_Atd-0jqtVDxDie14,10611
-pydantic_ai/models/function.py,sha256=jv2aw5K4KrMIRPFSTiGbBHcv16UZtGEe_1irjTudLg0,9826
-pydantic_ai/models/gemini.py,sha256=OsvtVqmKK4XmnLjPTRFwtCeXEuBOWfPrVZVZq0mSA30,28305
-pydantic_ai/models/groq.py,sha256=W-uosi9PtvgXp8yoF6w8QZxPtrVjX1fmRQDjQQ10meE,13603
-pydantic_ai/models/mistral.py,sha256=RuRlxTBPGXIpcwTdjX2596uN2tgHBvCDQCWieWa1A_Q,24793
-pydantic_ai/models/openai.py,sha256=ZwhYcpPIehDzmOBY_U1BKQhxYPtvKUeQjs9jDNr23bU,15547
-pydantic_ai/models/test.py,sha256=D_wBpRtrPcikopd5LBYjdAp3q-1gvuB9WnuopBQiRME,16659
-pydantic_ai/models/vertexai.py,sha256=bQZ5W8n4521AH3jTiyP1fOYxtNgCgkhjCbBeASv5Ap8,9388
-pydantic_ai_slim-0.0.22.dist-info/METADATA,sha256=lhhSBLTlHXj025p-n0u_X3EGOZnoSUTjAkXk59_WEV8,2839
-pydantic_ai_slim-0.0.22.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-pydantic_ai_slim-0.0.22.dist-info/RECORD,,