pydantic-ai-slim 0.0.22__py3-none-any.whl → 0.0.23__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,35 @@ 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
 
 
 # pyright: reportUnknownMemberType=false

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.0.22
+Version: 0.0.23
 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.23
 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.22.dist-info → pydantic_ai_slim-0.0.23.dist-info}/RECORD

@@ -1,12 +1,12 @@
 pydantic_ai/__init__.py,sha256=FbYetEgT6OO25u2KF5ZnFxKpz5DtnSpfckRXP4mjl8E,489
-pydantic_ai/_agent_graph.py,sha256=IdznyPzgrLmfiw3733GA3gYWOA_6LYIEuXBKTBXOi0A,32942
+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=RDKBnHU1yFBMwW9e7_BUYs4aRfbwJ69bpWp_ha2y9M0,44880
+pydantic_ai/agent.py,sha256=1jTnra4feBI0cMq-uR9NPVMkoQv9KRd7AVZsy8R1OAs,44918
 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
@@ -15,16 +15,16 @@ 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,,
+pydantic_ai/models/__init__.py,sha256=m5iGkMQQOmzV1j4wzuskIxXP3kbSM002A29kPEDusKI,13001
+pydantic_ai/models/anthropic.py,sha256=nt7vuoQw_O5rBmEkncTPAd_RQ23j40I2pIoXWChIrJ4,17260
+pydantic_ai/models/cohere.py,sha256=Nn7KbhSzWGR1a5ot2PtCEd0fZajhv5eL9qvTHJwUnps,10772
+pydantic_ai/models/function.py,sha256=JuqULRz0TmhEqmPrbLeN6iqFrpP_k85B6rTgeeE7Tcg,9887
+pydantic_ai/models/gemini.py,sha256=wXMtgjK9yepoBNb5LpFfD4eVyDejYlP9slrRFPXz98g,30181
+pydantic_ai/models/groq.py,sha256=Or2CDLWDn-DCxeOdKZAXSbQLe7l_uLdPQZE7Iwo_QL4,14138
+pydantic_ai/models/mistral.py,sha256=2lsTIoC3FhrYSyKS88yXQ4mZNu8LAqRQzKITp0EvzmY,25385
+pydantic_ai/models/openai.py,sha256=1px0ykE-4xCHSaQJkfiudQkFTw0PHtjmlpidevDTTok,16682
+pydantic_ai/models/test.py,sha256=eoenx_2RfcVCsao6RMDCGk0AtNBq7GIoSXi02Wep-k8,16428
+pydantic_ai/models/vertexai.py,sha256=4SuVviOFgxcnOe6CjA_z8X0Hzozs6fcMUo5PvafYc68,9279
+pydantic_ai_slim-0.0.23.dist-info/METADATA,sha256=XWCt3p_Ib1xE-2UABp0RVREc_e_WXXC0tdfjWiCrKbg,2839
+pydantic_ai_slim-0.0.23.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pydantic_ai_slim-0.0.23.dist-info/RECORD,,