Skip to content

Commit 2958fb1

Browse files
jaycee-licopybara-github
jaycee-li
authored and
copybara-github
committed
feat: add usage_metadata (including traffic_type) and truncated to EmbedContentResponse for Gemini Enterprise API
PiperOrigin-RevId: 1017527289
1 parent 0f29f4d commit 2958fb1

3 files changed

Lines changed: 227 additions & 117 deletions

File tree

‎google/genai/models.py‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1017,6 +1017,9 @@ def _EmbedContentResponse_from_mldev(
10171017
if getv(from_object, ['metadata']) is not None:
10181018
setv(to_object, ['metadata'], getv(from_object, ['metadata']))
10191019

1020+
if getv(from_object, ['usageMetadata']) is not None:
1021+
setv(to_object, ['usage_metadata'], getv(from_object, ['usageMetadata']))
1022+
10201023
return to_object
10211024

10221025

@@ -1044,6 +1047,18 @@ def _EmbedContentResponse_from_vertex(
10441047
if getv(from_object, ['metadata']) is not None:
10451048
setv(to_object, ['metadata'], getv(from_object, ['metadata']))
10461049

1050+
if getv(from_object, ['truncated']) is not None:
1051+
setv(to_object, ['truncated'], getv(from_object, ['truncated']))
1052+
1053+
if getv(from_object, ['usageMetadata']) is not None:
1054+
setv(
1055+
to_object,
1056+
['usage_metadata'],
1057+
_UsageMetadata_from_vertex(
1058+
getv(from_object, ['usageMetadata']), to_object, root_object
1059+
),
1060+
)
1061+
10471062
if getv(root_object, ['embedding_api_type']) == 'EMBED_CONTENT':
10481063
embedding = getv(from_object, ['embedding'])
10491064
usage_metadata = getv(from_object, ['usageMetadata'])
@@ -4180,6 +4195,86 @@ def _UpscaleImageResponse_from_vertex(
41804195
return to_object
41814196

41824197

4198+
def _UsageMetadata_from_vertex(
4199+
from_object: Union[dict[str, Any], object],
4200+
parent_object: Optional[dict[str, Any]] = None,
4201+
root_object: Optional[Union[dict[str, Any], object]] = None,
4202+
) -> dict[str, Any]:
4203+
to_object: dict[str, Any] = {}
4204+
if getv(from_object, ['promptTokenCount']) is not None:
4205+
setv(
4206+
to_object,
4207+
['prompt_token_count'],
4208+
getv(from_object, ['promptTokenCount']),
4209+
)
4210+
4211+
if getv(from_object, ['cachedContentTokenCount']) is not None:
4212+
setv(
4213+
to_object,
4214+
['cached_content_token_count'],
4215+
getv(from_object, ['cachedContentTokenCount']),
4216+
)
4217+
4218+
if getv(from_object, ['candidatesTokenCount']) is not None:
4219+
setv(
4220+
to_object,
4221+
['response_token_count'],
4222+
getv(from_object, ['candidatesTokenCount']),
4223+
)
4224+
4225+
if getv(from_object, ['toolUsePromptTokenCount']) is not None:
4226+
setv(
4227+
to_object,
4228+
['tool_use_prompt_token_count'],
4229+
getv(from_object, ['toolUsePromptTokenCount']),
4230+
)
4231+
4232+
if getv(from_object, ['thoughtsTokenCount']) is not None:
4233+
setv(
4234+
to_object,
4235+
['thoughts_token_count'],
4236+
getv(from_object, ['thoughtsTokenCount']),
4237+
)
4238+
4239+
if getv(from_object, ['totalTokenCount']) is not None:
4240+
setv(
4241+
to_object, ['total_token_count'], getv(from_object, ['totalTokenCount'])
4242+
)
4243+
4244+
if getv(from_object, ['promptTokensDetails']) is not None:
4245+
setv(
4246+
to_object,
4247+
['prompt_tokens_details'],
4248+
[item for item in getv(from_object, ['promptTokensDetails'])],
4249+
)
4250+
4251+
if getv(from_object, ['cacheTokensDetails']) is not None:
4252+
setv(
4253+
to_object,
4254+
['cache_tokens_details'],
4255+
[item for item in getv(from_object, ['cacheTokensDetails'])],
4256+
)
4257+
4258+
if getv(from_object, ['candidatesTokensDetails']) is not None:
4259+
setv(
4260+
to_object,
4261+
['response_tokens_details'],
4262+
[item for item in getv(from_object, ['candidatesTokensDetails'])],
4263+
)
4264+
4265+
if getv(from_object, ['toolUsePromptTokensDetails']) is not None:
4266+
setv(
4267+
to_object,
4268+
['tool_use_prompt_tokens_details'],
4269+
[item for item in getv(from_object, ['toolUsePromptTokensDetails'])],
4270+
)
4271+
4272+
if getv(from_object, ['trafficType']) is not None:
4273+
setv(to_object, ['traffic_type'], getv(from_object, ['trafficType']))
4274+
4275+
return to_object
4276+
4277+
41834278
def _VideoGenerationMask_to_vertex(
41844279
from_object: Union[dict[str, Any], object],
41854280
parent_object: Optional[dict[str, Any]] = None,

‎google/genai/tests/models/test_embed_content.py‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ def _get_bytes_from_file(relative_path: str) -> bytes:
7878
pytest_helper.TestTableItem(
7979
name='test_vertex_new_api_text_only_with_config',
8080
parameters=types.EmbedContentParameters(
81-
model='gemini-embedding-2-preview',
81+
model='gemini-embedding-2',
8282
contents=t.t_contents('What is your name?'),
8383
config={
8484
'output_dimensionality': 10,
@@ -96,7 +96,7 @@ def _get_bytes_from_file(relative_path: str) -> bytes:
9696
pytest_helper.TestTableItem(
9797
name='test_new_api_text_only',
9898
parameters=types.EmbedContentParameters(
99-
model='gemini-embedding-2-preview',
99+
model='gemini-embedding-2',
100100
contents=t.t_contents('What is your name?'),
101101
config={
102102
'output_dimensionality': 100,
@@ -120,7 +120,7 @@ def _get_bytes_from_file(relative_path: str) -> bytes:
120120
pytest_helper.TestTableItem(
121121
name='test_new_api_gcs_image_and_config',
122122
parameters=types.EmbedContentParameters(
123-
model='gemini-embedding-2-preview',
123+
model='gemini-embedding-2',
124124
contents=[
125125
types.Content(
126126
parts=[
@@ -147,7 +147,7 @@ def _get_bytes_from_file(relative_path: str) -> bytes:
147147
pytest_helper.TestTableItem(
148148
name='test_new_api_inline_pdf',
149149
parameters=types.EmbedContentParameters(
150-
model='gemini-embedding-2-preview',
150+
model='gemini-embedding-2',
151151
contents=[
152152
types.Content(
153153
parts=[
@@ -178,7 +178,7 @@ def _get_bytes_from_file(relative_path: str) -> bytes:
178178
pytest_helper.TestTableItem(
179179
name='test_vertex_inline_pdf_document_ocr',
180180
parameters=types.EmbedContentParameters(
181-
model='gemini-embedding-2-preview',
181+
model='gemini-embedding-2',
182182
contents=[
183183
types.Content(
184184
parts=[
@@ -199,7 +199,7 @@ def _get_bytes_from_file(relative_path: str) -> bytes:
199199
pytest_helper.TestTableItem(
200200
name='test_vertex_inline_video_audio_track_extraction',
201201
parameters=types.EmbedContentParameters(
202-
model='gemini-embedding-2-preview',
202+
model='gemini-embedding-2',
203203
contents=[
204204
types.Content(
205205
parts=[
@@ -229,7 +229,7 @@ def _get_bytes_from_file(relative_path: str) -> bytes:
229229

230230
def test_gemini_embedding_2_content_combination(client):
231231
response = client.models.embed_content(
232-
model='gemini-embedding-2-preview',
232+
model='gemini-embedding-2',
233233
contents=[
234234
'The jetpack is cool',
235235
types.Part.from_bytes(
@@ -253,6 +253,8 @@ def test_gemini_embedding_2_content_combination(client):
253253
assert statistics.token_count is not None
254254
assert statistics.tokens_details is not None
255255
assert len(statistics.tokens_details) > 0
256+
assert response.usage_metadata is not None
257+
assert response.usage_metadata.traffic_type is not None
256258

257259

258260
@pytest.mark.asyncio
@@ -268,7 +270,7 @@ async def test_async(client):
268270
@pytest.mark.asyncio
269271
async def test_async_new_api(client):
270272
response = await client.aio.models.embed_content(
271-
model='gemini-embedding-2-preview',
273+
model='gemini-embedding-2',
272274
contents=t.t_contents('What is your name?'),
273275
config={'output_dimensionality': 10},
274276
)

0 commit comments

Comments
 (0)