Skip to content

Commit 93bd5e6

Browse files
ccurmeccurme
authored
feat(vertexai): support model profiles (#1374)
1 parent 8d6b248 commit 93bd5e6

File tree

7 files changed

+348
-9
lines changed

7 files changed

+348
-9
lines changed

‎libs/vertexai/langchain_google_vertexai/chat_models.py‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,11 @@
2828
AsyncCallbackManagerForLLMRun,
2929
CallbackManagerForLLMRun,
3030
)
31-
from langchain_core.language_models import LanguageModelInput
31+
from langchain_core.language_models import (
32+
LanguageModelInput,
33+
ModelProfile,
34+
ModelProfileRegistry,
35+
)
3236
from langchain_core.language_models.chat_models import (
3337
BaseChatModel,
3438
LangSmithParams,
@@ -115,6 +119,8 @@
115119
ToolConfig as GapicToolConfig,
116120
VideoMetadata,
117121
)
122+
123+
from langchain_google_vertexai.data._profiles import _PROFILES
118124
from langchain_google_vertexai._base import _VertexAICommon
119125
from langchain_google_vertexai._compat import _convert_from_v1_to_vertex
120126
from langchain_google_vertexai._image_utils import (
@@ -178,6 +184,14 @@
178184
]
179185

180186

187+
_MODEL_PROFILES = cast("ModelProfileRegistry", _PROFILES)
188+
189+
190+
def _get_default_model_profile(model_name: str) -> ModelProfile:
191+
default = _MODEL_PROFILES.get(model_name) or {}
192+
return default.copy()
193+
194+
181195
_FUNCTION_CALL_THOUGHT_SIGNATURES_MAP_KEY = (
182196
"__gemini_function_call_thought_signatures__"
183197
)
@@ -1867,6 +1881,14 @@ def validate_environment(self) -> Self:
18671881

18681882
return self
18691883

1884+
@model_validator(mode="after")
1885+
def _set_model_profile(self) -> Self:
1886+
"""Set model profile if not overridden."""
1887+
if self.profile is None:
1888+
model_id = re.sub(r"-\d{3}$", "", self.model_name.replace("models/", ""))
1889+
self.profile = _get_default_model_profile(model_id)
1890+
return self
1891+
18701892
def _prepare_params(
18711893
self,
18721894
stop: list[str] | None = None,

‎libs/vertexai/langchain_google_vertexai/data/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
"""Model profile data. All edits should be made in profile_augmentations.toml."""

‎libs/vertexai/langchain_google_vertexai/data/_profiles.py‎

Lines changed: 291 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,291 @@
1+
"""Auto-generated model profiles.
2+
3+
DO NOT EDIT THIS FILE MANUALLY.
4+
This file is generated by the langchain-profiles CLI tool.
5+
6+
It contains data derived from the models.dev project.
7+
8+
Source: https://github.com/sst/models.dev
9+
License: MIT License
10+
11+
To update these data, refer to the instructions here:
12+
13+
https://docs.langchain.com/oss/python/langchain/models#updating-or-overwriting-profile-data
14+
"""
15+
16+
from typing import Any
17+
18+
_PROFILES: dict[str, dict[str, Any]] = {
19+
"gemini-embedding-001": {
20+
"max_input_tokens": 2048,
21+
"max_output_tokens": 3072,
22+
"image_inputs": False,
23+
"audio_inputs": False,
24+
"video_inputs": False,
25+
"image_outputs": False,
26+
"audio_outputs": False,
27+
"video_outputs": False,
28+
"reasoning_output": False,
29+
"tool_calling": False,
30+
"image_url_inputs": True,
31+
"pdf_inputs": True,
32+
"image_tool_message": True,
33+
"tool_choice": True,
34+
"structured_output": True,
35+
},
36+
"gemini-2.5-flash-preview-05-20": {
37+
"max_input_tokens": 1048576,
38+
"max_output_tokens": 65536,
39+
"image_inputs": True,
40+
"audio_inputs": True,
41+
"pdf_inputs": True,
42+
"video_inputs": True,
43+
"image_outputs": False,
44+
"audio_outputs": False,
45+
"video_outputs": False,
46+
"reasoning_output": True,
47+
"tool_calling": True,
48+
"image_url_inputs": True,
49+
"image_tool_message": True,
50+
"tool_choice": True,
51+
"structured_output": True,
52+
},
53+
"gemini-flash-lite-latest": {
54+
"max_input_tokens": 1048576,
55+
"max_output_tokens": 65536,
56+
"image_inputs": True,
57+
"audio_inputs": True,
58+
"pdf_inputs": True,
59+
"video_inputs": True,
60+
"image_outputs": False,
61+
"audio_outputs": False,
62+
"video_outputs": False,
63+
"reasoning_output": True,
64+
"tool_calling": True,
65+
"image_url_inputs": True,
66+
"image_tool_message": True,
67+
"tool_choice": True,
68+
"structured_output": True,
69+
},
70+
"gemini-3-pro-preview": {
71+
"max_input_tokens": 1048576,
72+
"max_output_tokens": 65536,
73+
"image_inputs": True,
74+
"audio_inputs": True,
75+
"pdf_inputs": True,
76+
"video_inputs": True,
77+
"image_outputs": False,
78+
"audio_outputs": False,
79+
"video_outputs": False,
80+
"reasoning_output": True,
81+
"tool_calling": True,
82+
"structured_output": True,
83+
"image_url_inputs": True,
84+
"image_tool_message": True,
85+
"tool_choice": True,
86+
},
87+
"gemini-2.5-flash": {
88+
"max_input_tokens": 1048576,
89+
"max_output_tokens": 65536,
90+
"image_inputs": True,
91+
"audio_inputs": True,
92+
"pdf_inputs": True,
93+
"video_inputs": True,
94+
"image_outputs": False,
95+
"audio_outputs": False,
96+
"video_outputs": False,
97+
"reasoning_output": True,
98+
"tool_calling": True,
99+
"image_url_inputs": True,
100+
"image_tool_message": True,
101+
"tool_choice": True,
102+
"structured_output": True,
103+
},
104+
"gemini-flash-latest": {
105+
"max_input_tokens": 1048576,
106+
"max_output_tokens": 65536,
107+
"image_inputs": True,
108+
"audio_inputs": True,
109+
"pdf_inputs": True,
110+
"video_inputs": True,
111+
"image_outputs": False,
112+
"audio_outputs": False,
113+
"video_outputs": False,
114+
"reasoning_output": True,
115+
"tool_calling": True,
116+
"image_url_inputs": True,
117+
"image_tool_message": True,
118+
"tool_choice": True,
119+
"structured_output": True,
120+
},
121+
"gemini-2.5-pro-preview-05-06": {
122+
"max_input_tokens": 1048576,
123+
"max_output_tokens": 65536,
124+
"image_inputs": True,
125+
"audio_inputs": True,
126+
"pdf_inputs": True,
127+
"video_inputs": True,
128+
"image_outputs": False,
129+
"audio_outputs": False,
130+
"video_outputs": False,
131+
"reasoning_output": True,
132+
"tool_calling": True,
133+
"image_url_inputs": True,
134+
"image_tool_message": True,
135+
"tool_choice": True,
136+
"structured_output": True,
137+
},
138+
"gemini-2.0-flash-lite": {
139+
"max_input_tokens": 1048576,
140+
"max_output_tokens": 8192,
141+
"image_inputs": True,
142+
"audio_inputs": True,
143+
"pdf_inputs": True,
144+
"video_inputs": True,
145+
"image_outputs": False,
146+
"audio_outputs": False,
147+
"video_outputs": False,
148+
"reasoning_output": False,
149+
"tool_calling": True,
150+
"image_url_inputs": True,
151+
"image_tool_message": True,
152+
"tool_choice": True,
153+
"structured_output": True,
154+
},
155+
"gemini-2.0-flash": {
156+
"max_input_tokens": 1048576,
157+
"max_output_tokens": 8192,
158+
"image_inputs": True,
159+
"audio_inputs": True,
160+
"pdf_inputs": True,
161+
"video_inputs": True,
162+
"image_outputs": False,
163+
"audio_outputs": False,
164+
"video_outputs": False,
165+
"reasoning_output": False,
166+
"tool_calling": True,
167+
"image_url_inputs": True,
168+
"image_tool_message": True,
169+
"tool_choice": True,
170+
"structured_output": True,
171+
},
172+
"gemini-2.5-flash-lite": {
173+
"max_input_tokens": 1048576,
174+
"max_output_tokens": 65536,
175+
"image_inputs": True,
176+
"audio_inputs": True,
177+
"pdf_inputs": True,
178+
"video_inputs": True,
179+
"image_outputs": False,
180+
"audio_outputs": False,
181+
"video_outputs": False,
182+
"reasoning_output": True,
183+
"tool_calling": True,
184+
"image_url_inputs": True,
185+
"image_tool_message": True,
186+
"tool_choice": True,
187+
"structured_output": True,
188+
},
189+
"gemini-2.5-pro-preview-06-05": {
190+
"max_input_tokens": 1048576,
191+
"max_output_tokens": 65536,
192+
"image_inputs": True,
193+
"audio_inputs": True,
194+
"pdf_inputs": True,
195+
"video_inputs": True,
196+
"image_outputs": False,
197+
"audio_outputs": False,
198+
"video_outputs": False,
199+
"reasoning_output": True,
200+
"tool_calling": True,
201+
"image_url_inputs": True,
202+
"image_tool_message": True,
203+
"tool_choice": True,
204+
"structured_output": True,
205+
},
206+
"gemini-2.5-flash-lite-preview-06-17": {
207+
"max_input_tokens": 65536,
208+
"max_output_tokens": 65536,
209+
"image_inputs": True,
210+
"audio_inputs": True,
211+
"pdf_inputs": True,
212+
"video_inputs": True,
213+
"image_outputs": False,
214+
"audio_outputs": False,
215+
"video_outputs": False,
216+
"reasoning_output": True,
217+
"tool_calling": True,
218+
"image_url_inputs": True,
219+
"image_tool_message": True,
220+
"tool_choice": True,
221+
"structured_output": True,
222+
},
223+
"gemini-2.5-flash-preview-09-2025": {
224+
"max_input_tokens": 1048576,
225+
"max_output_tokens": 65536,
226+
"image_inputs": True,
227+
"audio_inputs": True,
228+
"pdf_inputs": True,
229+
"video_inputs": True,
230+
"image_outputs": False,
231+
"audio_outputs": False,
232+
"video_outputs": False,
233+
"reasoning_output": True,
234+
"tool_calling": True,
235+
"image_url_inputs": True,
236+
"image_tool_message": True,
237+
"tool_choice": True,
238+
"structured_output": True,
239+
},
240+
"gemini-2.5-flash-preview-04-17": {
241+
"max_input_tokens": 1048576,
242+
"max_output_tokens": 65536,
243+
"image_inputs": True,
244+
"audio_inputs": True,
245+
"pdf_inputs": True,
246+
"video_inputs": True,
247+
"image_outputs": False,
248+
"audio_outputs": False,
249+
"video_outputs": False,
250+
"reasoning_output": True,
251+
"tool_calling": True,
252+
"image_url_inputs": True,
253+
"image_tool_message": True,
254+
"tool_choice": True,
255+
"structured_output": True,
256+
},
257+
"gemini-2.5-pro": {
258+
"max_input_tokens": 1048576,
259+
"max_output_tokens": 65536,
260+
"image_inputs": True,
261+
"audio_inputs": True,
262+
"pdf_inputs": True,
263+
"video_inputs": True,
264+
"image_outputs": False,
265+
"audio_outputs": False,
266+
"video_outputs": False,
267+
"reasoning_output": True,
268+
"tool_calling": True,
269+
"image_url_inputs": True,
270+
"image_tool_message": True,
271+
"tool_choice": True,
272+
"structured_output": True,
273+
},
274+
"gemini-2.5-flash-lite-preview-09-2025": {
275+
"max_input_tokens": 1048576,
276+
"max_output_tokens": 65536,
277+
"image_inputs": True,
278+
"audio_inputs": True,
279+
"pdf_inputs": True,
280+
"video_inputs": True,
281+
"image_outputs": False,
282+
"audio_outputs": False,
283+
"video_outputs": False,
284+
"reasoning_output": True,
285+
"tool_calling": True,
286+
"image_url_inputs": True,
287+
"image_tool_message": True,
288+
"tool_choice": True,
289+
"structured_output": True,
290+
},
291+
}

‎libs/vertexai/langchain_google_vertexai/data/profile_augmentations.toml‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
provider = "google-vertex"
2+
3+
[overrides]
4+
image_url_inputs = true
5+
pdf_inputs = true
6+
image_tool_message = true
7+
tool_choice = true
8+
structured_output = true

‎libs/vertexai/pyproject.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,10 @@ license = {text = "MIT"}
99
readme = "README.md"
1010
authors = []
1111

12-
version = "3.0.3"
12+
version = "3.1.0"
1313
requires-python = ">=3.10.0,<4.0.0"
1414
dependencies = [
15-
"langchain-core>=1.0.0,<2.0.0",
15+
"langchain-core>=1.1.0,<2.0.0",
1616
"google-cloud-aiplatform>=1.97.0,<2.0.0",
1717
"google-cloud-storage>=2.18.0,<4.0.0",
1818
"httpx>=0.28.0,<1.0.0",

‎libs/vertexai/tests/unit_tests/test_chat_models.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -354,6 +354,23 @@ def test_init_client_with_custom_model_kwargs() -> None:
354354
assert default_params["thinking"] == {"type": "enabled", "budget_tokens": 1024}
355355

356356

357+
def test_profile() -> None:
358+
model = ChatVertexAI(
359+
model="gemini-2.0-flash", project="test-project", location="moon-dark1"
360+
)
361+
assert model.profile
362+
assert not model.profile["reasoning_output"]
363+
364+
model = ChatVertexAI(
365+
model="gemini-2.5-flash", project="test-project", location="moon-dark1"
366+
)
367+
assert model.profile
368+
assert model.profile["reasoning_output"]
369+
370+
model = ChatVertexAI(model="foo", project="test-project", location="moon-dark1")
371+
assert model.profile == {}
372+
373+
357374
@pytest.mark.parametrize(
358375
("model", "location"),
359376
[

0 commit comments

Comments
 (0)