diff --git a/google/genai/models.py b/google/genai/models.py index ddb378e50..cdd6cc158 100644 --- a/google/genai/models.py +++ b/google/genai/models.py @@ -1017,6 +1017,9 @@ def _EmbedContentResponse_from_mldev( if getv(from_object, ['metadata']) is not None: setv(to_object, ['metadata'], getv(from_object, ['metadata'])) + if getv(from_object, ['usageMetadata']) is not None: + setv(to_object, ['usage_metadata'], getv(from_object, ['usageMetadata'])) + return to_object @@ -1044,6 +1047,18 @@ def _EmbedContentResponse_from_vertex( if getv(from_object, ['metadata']) is not None: setv(to_object, ['metadata'], getv(from_object, ['metadata'])) + if getv(from_object, ['truncated']) is not None: + setv(to_object, ['truncated'], getv(from_object, ['truncated'])) + + if getv(from_object, ['usageMetadata']) is not None: + setv( + to_object, + ['usage_metadata'], + _UsageMetadata_from_vertex( + getv(from_object, ['usageMetadata']), to_object, root_object + ), + ) + if getv(root_object, ['embedding_api_type']) == 'EMBED_CONTENT': embedding = getv(from_object, ['embedding']) usage_metadata = getv(from_object, ['usageMetadata']) @@ -4180,6 +4195,86 @@ def _UpscaleImageResponse_from_vertex( return to_object +def _UsageMetadata_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, + root_object: Optional[Union[dict[str, Any], object]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ['promptTokenCount']) is not None: + setv( + to_object, + ['prompt_token_count'], + getv(from_object, ['promptTokenCount']), + ) + + if getv(from_object, ['cachedContentTokenCount']) is not None: + setv( + to_object, + ['cached_content_token_count'], + getv(from_object, ['cachedContentTokenCount']), + ) + + if getv(from_object, ['candidatesTokenCount']) is not None: + setv( + to_object, + ['response_token_count'], + getv(from_object, ['candidatesTokenCount']), + ) + + if getv(from_object, ['toolUsePromptTokenCount']) is not None: + setv( + to_object, + ['tool_use_prompt_token_count'], + getv(from_object, ['toolUsePromptTokenCount']), + ) + + if getv(from_object, ['thoughtsTokenCount']) is not None: + setv( + to_object, + ['thoughts_token_count'], + getv(from_object, ['thoughtsTokenCount']), + ) + + if getv(from_object, ['totalTokenCount']) is not None: + setv( + to_object, ['total_token_count'], getv(from_object, ['totalTokenCount']) + ) + + if getv(from_object, ['promptTokensDetails']) is not None: + setv( + to_object, + ['prompt_tokens_details'], + [item for item in getv(from_object, ['promptTokensDetails'])], + ) + + if getv(from_object, ['cacheTokensDetails']) is not None: + setv( + to_object, + ['cache_tokens_details'], + [item for item in getv(from_object, ['cacheTokensDetails'])], + ) + + if getv(from_object, ['candidatesTokensDetails']) is not None: + setv( + to_object, + ['response_tokens_details'], + [item for item in getv(from_object, ['candidatesTokensDetails'])], + ) + + if getv(from_object, ['toolUsePromptTokensDetails']) is not None: + setv( + to_object, + ['tool_use_prompt_tokens_details'], + [item for item in getv(from_object, ['toolUsePromptTokensDetails'])], + ) + + if getv(from_object, ['trafficType']) is not None: + setv(to_object, ['traffic_type'], getv(from_object, ['trafficType'])) + + return to_object + + def _VideoGenerationMask_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, diff --git a/google/genai/tests/models/test_embed_content.py b/google/genai/tests/models/test_embed_content.py index 8f565efef..dc64d87dd 100644 --- a/google/genai/tests/models/test_embed_content.py +++ b/google/genai/tests/models/test_embed_content.py @@ -78,7 +78,7 @@ def _get_bytes_from_file(relative_path: str) -> bytes: pytest_helper.TestTableItem( name='test_vertex_new_api_text_only_with_config', parameters=types.EmbedContentParameters( - model='gemini-embedding-2-preview', + model='gemini-embedding-2', contents=t.t_contents('What is your name?'), config={ 'output_dimensionality': 10, @@ -96,7 +96,7 @@ def _get_bytes_from_file(relative_path: str) -> bytes: pytest_helper.TestTableItem( name='test_new_api_text_only', parameters=types.EmbedContentParameters( - model='gemini-embedding-2-preview', + model='gemini-embedding-2', contents=t.t_contents('What is your name?'), config={ 'output_dimensionality': 100, @@ -120,7 +120,7 @@ def _get_bytes_from_file(relative_path: str) -> bytes: pytest_helper.TestTableItem( name='test_new_api_gcs_image_and_config', parameters=types.EmbedContentParameters( - model='gemini-embedding-2-preview', + model='gemini-embedding-2', contents=[ types.Content( parts=[ @@ -147,7 +147,7 @@ def _get_bytes_from_file(relative_path: str) -> bytes: pytest_helper.TestTableItem( name='test_new_api_inline_pdf', parameters=types.EmbedContentParameters( - model='gemini-embedding-2-preview', + model='gemini-embedding-2', contents=[ types.Content( parts=[ @@ -178,7 +178,7 @@ def _get_bytes_from_file(relative_path: str) -> bytes: pytest_helper.TestTableItem( name='test_vertex_inline_pdf_document_ocr', parameters=types.EmbedContentParameters( - model='gemini-embedding-2-preview', + model='gemini-embedding-2', contents=[ types.Content( parts=[ @@ -199,7 +199,7 @@ def _get_bytes_from_file(relative_path: str) -> bytes: pytest_helper.TestTableItem( name='test_vertex_inline_video_audio_track_extraction', parameters=types.EmbedContentParameters( - model='gemini-embedding-2-preview', + model='gemini-embedding-2', contents=[ types.Content( parts=[ @@ -229,7 +229,7 @@ def _get_bytes_from_file(relative_path: str) -> bytes: def test_gemini_embedding_2_content_combination(client): response = client.models.embed_content( - model='gemini-embedding-2-preview', + model='gemini-embedding-2', contents=[ 'The jetpack is cool', types.Part.from_bytes( @@ -253,6 +253,8 @@ def test_gemini_embedding_2_content_combination(client): assert statistics.token_count is not None assert statistics.tokens_details is not None assert len(statistics.tokens_details) > 0 + assert response.usage_metadata is not None + assert response.usage_metadata.traffic_type is not None @pytest.mark.asyncio @@ -268,7 +270,7 @@ async def test_async(client): @pytest.mark.asyncio async def test_async_new_api(client): response = await client.aio.models.embed_content( - model='gemini-embedding-2-preview', + model='gemini-embedding-2', contents=t.t_contents('What is your name?'), config={'output_dimensionality': 10}, ) diff --git a/google/genai/types.py b/google/genai/types.py index e37c171ba..e684d9b8a 100644 --- a/google/genai/types.py +++ b/google/genai/types.py @@ -693,6 +693,19 @@ class ModelStage(_common.CaseInSensitiveEnum): """Models in this stage are retired. These models cannot be used.""" +class ServiceTier(_common.CaseInSensitiveEnum): + """Pricing and performance service tier.""" + + UNSPECIFIED = 'unspecified' + """Default service tier, which is standard.""" + FLEX = 'flex' + """Flex service tier.""" + STANDARD = 'standard' + """Standard service tier.""" + PRIORITY = 'priority' + """Priority service tier.""" + + class MediaResolution(_common.CaseInSensitiveEnum): """The media resolution to use.""" @@ -1007,19 +1020,6 @@ class DocumentState(_common.CaseInSensitiveEnum): """Some `Chunks` of the `Document` failed processing.""" -class ServiceTier(_common.CaseInSensitiveEnum): - """Pricing and performance service tier.""" - - UNSPECIFIED = 'unspecified' - """Default service tier, which is standard.""" - FLEX = 'flex' - """Flex service tier.""" - STANDARD = 'standard' - """Standard service tier.""" - PRIORITY = 'priority' - """Priority service tier.""" - - class RubricContentType(_common.CaseInSensitiveEnum): """Represents the rubric content type.""" @@ -9396,6 +9396,102 @@ class EmbedContentMetadataDict(TypedDict, total=False): ] +class UsageMetadata(_common.BaseModel): + """Usage metadata about response(s).""" + + prompt_token_count: Optional[int] = Field( + default=None, + description="""The total number of tokens in the prompt. This includes any text, images, or other media provided in the request. When `cached_content` is set, this also includes the number of tokens in the cached content.""", + ) + cached_content_token_count: Optional[int] = Field( + default=None, + description="""Output only. The number of tokens in the cached content that was used for this request.""", + ) + response_token_count: Optional[int] = Field( + default=None, + description="""Total number of tokens across all the generated response candidates.""", + ) + tool_use_prompt_token_count: Optional[int] = Field( + default=None, + description="""Output only. The number of tokens in the results from tool executions, which are provided back to the model as input, if applicable.""", + ) + thoughts_token_count: Optional[int] = Field( + default=None, + description="""Output only. The number of tokens that were part of the model's generated "thoughts" output, if applicable.""", + ) + total_token_count: Optional[int] = Field( + default=None, + description="""The total number of tokens for the entire request. This is the sum of `prompt_token_count`, `candidates_token_count`, `tool_use_prompt_token_count`, and `thoughts_token_count`.""", + ) + prompt_tokens_details: Optional[list[ModalityTokenCount]] = Field( + default=None, + description="""Output only. A detailed breakdown of the token count for each modality in the prompt.""", + ) + cache_tokens_details: Optional[list[ModalityTokenCount]] = Field( + default=None, + description="""Output only. A detailed breakdown of the token count for each modality in the cached content.""", + ) + response_tokens_details: Optional[list[ModalityTokenCount]] = Field( + default=None, + description="""List of modalities that were returned in the response.""", + ) + tool_use_prompt_tokens_details: Optional[list[ModalityTokenCount]] = Field( + default=None, + description="""Output only. A detailed breakdown by modality of the token counts from the results of tool executions, which are provided back to the model as input.""", + ) + traffic_type: Optional[TrafficType] = Field( + default=None, + description="""Output only. The traffic type for this request. This field is not supported in Gemini API.""", + ) + service_tier: Optional[ServiceTier] = Field( + default=None, + description="""Output only. Service tier of the request. This field is not supported in Vertex AI.""", + ) + + +class UsageMetadataDict(TypedDict, total=False): + """Usage metadata about response(s).""" + + prompt_token_count: Optional[int] + """The total number of tokens in the prompt. This includes any text, images, or other media provided in the request. When `cached_content` is set, this also includes the number of tokens in the cached content.""" + + cached_content_token_count: Optional[int] + """Output only. The number of tokens in the cached content that was used for this request.""" + + response_token_count: Optional[int] + """Total number of tokens across all the generated response candidates.""" + + tool_use_prompt_token_count: Optional[int] + """Output only. The number of tokens in the results from tool executions, which are provided back to the model as input, if applicable.""" + + thoughts_token_count: Optional[int] + """Output only. The number of tokens that were part of the model's generated "thoughts" output, if applicable.""" + + total_token_count: Optional[int] + """The total number of tokens for the entire request. This is the sum of `prompt_token_count`, `candidates_token_count`, `tool_use_prompt_token_count`, and `thoughts_token_count`.""" + + prompt_tokens_details: Optional[list[ModalityTokenCountDict]] + """Output only. A detailed breakdown of the token count for each modality in the prompt.""" + + cache_tokens_details: Optional[list[ModalityTokenCountDict]] + """Output only. A detailed breakdown of the token count for each modality in the cached content.""" + + response_tokens_details: Optional[list[ModalityTokenCountDict]] + """List of modalities that were returned in the response.""" + + tool_use_prompt_tokens_details: Optional[list[ModalityTokenCountDict]] + """Output only. A detailed breakdown by modality of the token counts from the results of tool executions, which are provided back to the model as input.""" + + traffic_type: Optional[TrafficType] + """Output only. The traffic type for this request. This field is not supported in Gemini API.""" + + service_tier: Optional[ServiceTier] + """Output only. Service tier of the request. This field is not supported in Vertex AI.""" + + +UsageMetadataOrDict = Union[UsageMetadata, UsageMetadataDict] + + class EmbedContentResponse(_common.BaseModel): """Response for the embed_content method.""" @@ -9413,6 +9509,13 @@ class EmbedContentResponse(_common.BaseModel): description="""Gemini Enterprise Agent Platform only. Metadata about the request. """, ) + truncated: Optional[bool] = Field( + default=None, + description="""Whether the input content was truncated before generating the embedding. This field is not supported in Gemini API.""", + ) + usage_metadata: Optional[UsageMetadata] = Field( + default=None, description="""Usage metadata about the response(s).""" + ) class EmbedContentResponseDict(TypedDict, total=False): @@ -9430,6 +9533,12 @@ class EmbedContentResponseDict(TypedDict, total=False): """Gemini Enterprise Agent Platform only. Metadata about the request. """ + truncated: Optional[bool] + """Whether the input content was truncated before generating the embedding. This field is not supported in Gemini API.""" + + usage_metadata: Optional[UsageMetadataDict] + """Usage metadata about the response(s).""" + EmbedContentResponseOrDict = Union[ EmbedContentResponse, EmbedContentResponseDict @@ -20813,102 +20922,6 @@ class LiveServerToolCallCancellationDict(TypedDict, total=False): ] -class UsageMetadata(_common.BaseModel): - """Usage metadata about response(s).""" - - prompt_token_count: Optional[int] = Field( - default=None, - description="""The total number of tokens in the prompt. This includes any text, images, or other media provided in the request. When `cached_content` is set, this also includes the number of tokens in the cached content.""", - ) - cached_content_token_count: Optional[int] = Field( - default=None, - description="""Output only. The number of tokens in the cached content that was used for this request.""", - ) - response_token_count: Optional[int] = Field( - default=None, - description="""Total number of tokens across all the generated response candidates.""", - ) - tool_use_prompt_token_count: Optional[int] = Field( - default=None, - description="""Output only. The number of tokens in the results from tool executions, which are provided back to the model as input, if applicable.""", - ) - thoughts_token_count: Optional[int] = Field( - default=None, - description="""Output only. The number of tokens that were part of the model's generated "thoughts" output, if applicable.""", - ) - total_token_count: Optional[int] = Field( - default=None, - description="""The total number of tokens for the entire request. This is the sum of `prompt_token_count`, `candidates_token_count`, `tool_use_prompt_token_count`, and `thoughts_token_count`.""", - ) - prompt_tokens_details: Optional[list[ModalityTokenCount]] = Field( - default=None, - description="""Output only. A detailed breakdown of the token count for each modality in the prompt.""", - ) - cache_tokens_details: Optional[list[ModalityTokenCount]] = Field( - default=None, - description="""Output only. A detailed breakdown of the token count for each modality in the cached content.""", - ) - response_tokens_details: Optional[list[ModalityTokenCount]] = Field( - default=None, - description="""List of modalities that were returned in the response.""", - ) - tool_use_prompt_tokens_details: Optional[list[ModalityTokenCount]] = Field( - default=None, - description="""Output only. A detailed breakdown by modality of the token counts from the results of tool executions, which are provided back to the model as input.""", - ) - traffic_type: Optional[TrafficType] = Field( - default=None, - description="""Output only. The traffic type for this request. This field is not supported in Gemini API.""", - ) - service_tier: Optional[ServiceTier] = Field( - default=None, - description="""Output only. Service tier of the request. This field is not supported in Vertex AI.""", - ) - - -class UsageMetadataDict(TypedDict, total=False): - """Usage metadata about response(s).""" - - prompt_token_count: Optional[int] - """The total number of tokens in the prompt. This includes any text, images, or other media provided in the request. When `cached_content` is set, this also includes the number of tokens in the cached content.""" - - cached_content_token_count: Optional[int] - """Output only. The number of tokens in the cached content that was used for this request.""" - - response_token_count: Optional[int] - """Total number of tokens across all the generated response candidates.""" - - tool_use_prompt_token_count: Optional[int] - """Output only. The number of tokens in the results from tool executions, which are provided back to the model as input, if applicable.""" - - thoughts_token_count: Optional[int] - """Output only. The number of tokens that were part of the model's generated "thoughts" output, if applicable.""" - - total_token_count: Optional[int] - """The total number of tokens for the entire request. This is the sum of `prompt_token_count`, `candidates_token_count`, `tool_use_prompt_token_count`, and `thoughts_token_count`.""" - - prompt_tokens_details: Optional[list[ModalityTokenCountDict]] - """Output only. A detailed breakdown of the token count for each modality in the prompt.""" - - cache_tokens_details: Optional[list[ModalityTokenCountDict]] - """Output only. A detailed breakdown of the token count for each modality in the cached content.""" - - response_tokens_details: Optional[list[ModalityTokenCountDict]] - """List of modalities that were returned in the response.""" - - tool_use_prompt_tokens_details: Optional[list[ModalityTokenCountDict]] - """Output only. A detailed breakdown by modality of the token counts from the results of tool executions, which are provided back to the model as input.""" - - traffic_type: Optional[TrafficType] - """Output only. The traffic type for this request. This field is not supported in Gemini API.""" - - service_tier: Optional[ServiceTier] - """Output only. Service tier of the request. This field is not supported in Vertex AI.""" - - -UsageMetadataOrDict = Union[UsageMetadata, UsageMetadataDict] - - class LiveServerGoAway(_common.BaseModel): """Server will not be able to service client soon."""