feat: add MediaModalities for ModalityTokenCount

MarkDaoust · copybara-github · commit fb2509c471a4 · 2025-03-25T17:48:32.000-07:00
PiperOrigin-RevId: 740549991
diff --git a/google/genai/tests/models/test_generate_content.py b/google/genai/tests/models/test_generate_content.py
@@ -1837,6 +1837,8 @@ def test_usage_metadata_part_types(client):
       [d.modality.name for d in usage_metadata.candidates_tokens_details]
   )
   assert modalities == ['TEXT']
+  assert isinstance(
+      usage_metadata.candidates_tokens_details[0].modality, types.MediaModality)
 
   assert usage_metadata.prompt_token_count
   assert usage_metadata.prompt_tokens_details
diff --git a/google/genai/types.py b/google/genai/types.py
@@ -343,6 +343,17 @@ class FileSource(_common.CaseInSensitiveEnum):
   GENERATED = 'GENERATED'
 
 
+class MediaModality(_common.CaseInSensitiveEnum):
+  """Server content modalities."""
+
+  MODALITY_UNSPECIFIED = 'MODALITY_UNSPECIFIED'
+  TEXT = 'TEXT'
+  IMAGE = 'IMAGE'
+  VIDEO = 'VIDEO'
+  AUDIO = 'AUDIO'
+  DOCUMENT = 'DOCUMENT'
+
+
 class VideoMetadata(_common.BaseModel):
   """Metadata describes the input video content."""
 
@@ -2861,7 +2872,7 @@ class GenerateContentResponsePromptFeedbackDict(TypedDict, total=False):
 class ModalityTokenCount(_common.BaseModel):
   """Represents token counting info for a single modality."""
 
-  modality: Optional[Modality] = Field(
+  modality: Optional[MediaModality] = Field(
       default=None,
       description="""The modality associated with this token count.""",
   )
@@ -2873,7 +2884,7 @@ class ModalityTokenCount(_common.BaseModel):
 class ModalityTokenCountDict(TypedDict, total=False):
   """Represents token counting info for a single modality."""
 
-  modality: Optional[Modality]
+  modality: Optional[MediaModality]
   """The modality associated with this token count."""
 
   token_count: Optional[int]

Original file line number	Diff line number	Diff line change
`@@ -1837,6 +1837,8 @@ def test_usage_metadata_part_types(client):`
`1837`	`1837`	`[d.modality.name for d in usage_metadata.candidates_tokens_details]`
`1838`	`1838`	`)`
`1839`	`1839`	`assert modalities == ['TEXT']`
	`1840`	`+ assert isinstance(`
	`1841`	`+ usage_metadata.candidates_tokens_details[0].modality, types.MediaModality)`
`1840`	`1842`
`1841`	`1843`	`assert usage_metadata.prompt_token_count`
`1842`	`1844`	`assert usage_metadata.prompt_tokens_details`