mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(vertex): preserve Gemini Embedding 2 usageMetadata for cost tracking (#31354)
* fix(vertex): preserve Gemini Embedding 2 usageMetadata for cost tracking * style(vertex): apply ruff format to batch_embed_content_transformation * fix(vertex): bill files/ image refs in Gemini embedContent at per-image rate Resolved files/... references whose mime type is an image were not detected by _is_image_element, so image_count stayed 0 and generic_cost_per_token fell back to the text token rate instead of input_cost_per_image. Thread the resolved_files mapping into the usage builder so resolved image references are counted and billed per image. Also modernize the _flatten_input return annotation to satisfy the ruff UP006 strict gate. * fix(vertex): bill Gemini embedding audio per-second and stop video+audio double-billing Audio-only embedContent responses set audio_tokens, but generic_cost_per_token only charges audio via input_cost_per_audio_token. gemini-embedding-2 prices audio via input_cost_per_audio_per_second, so spend stayed at $0. Plumb a new audio_length_seconds field through PromptTokensDetailsWrapper, parse it in _parse_prompt_tokens_details, and bill it from _calculate_input_cost. The vertex embedding transformation derives audio_length_seconds from audio_tokens using the documented 32 tokens/sec Gemini rate. The 1-token text floor that protects video billing only fired when no other modality was billable, but audio presence flipped that flag, leaving text_tokens at zero for video+audio responses. generic_cost_per_token then rewrote text_tokens to prompt_tokens minus audio_tokens (the video token count), charging video tokens as text on top of the per-second video cost. The rewrite trigger is text_tokens == 0 and image_count == 0; align the floor with that trigger and ignore audio_tokens. --------- Co-authored-by: Cursor Agent <cursoragent@cursor.com>
This commit is contained in:
Mateo Wang
GitHub
parent
997c7a2676
commit
6e3540856c
6 changed files with 383 additions and 13 deletions
|
|
@ -485,6 +485,7 @@ class PromptTokensDetailsResult(TypedDict):
|
|||
character_count: int
|
||||
image_count: int
|
||||
video_length_seconds: float
|
||||
audio_length_seconds: float
|
||||
|
||||
|
||||
def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
|
||||
|
|
@ -535,6 +536,13 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
|
|||
)
|
||||
or 0.0
|
||||
)
|
||||
audio_length_seconds = (
|
||||
cast(
|
||||
Optional[float],
|
||||
getattr(usage.prompt_tokens_details, "audio_length_seconds", 0),
|
||||
)
|
||||
or 0.0
|
||||
)
|
||||
|
||||
return PromptTokensDetailsResult(
|
||||
cache_hit_tokens=cache_hit_tokens,
|
||||
|
|
@ -546,6 +554,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
|
|||
character_count=character_count,
|
||||
image_count=image_count,
|
||||
video_length_seconds=float(video_length_seconds),
|
||||
audio_length_seconds=float(audio_length_seconds),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -667,6 +676,14 @@ def _calculate_input_cost(
|
|||
prompt_tokens_details["video_length_seconds"],
|
||||
)
|
||||
|
||||
### AUDIO LENGTH COST
|
||||
if prompt_tokens_details["audio_length_seconds"]:
|
||||
prompt_cost += calculate_cost_component(
|
||||
model_info,
|
||||
"input_cost_per_audio_per_second",
|
||||
prompt_tokens_details["audio_length_seconds"],
|
||||
)
|
||||
|
||||
return prompt_cost
|
||||
|
||||
|
||||
|
|
@ -743,6 +760,7 @@ def generic_cost_per_token(
|
|||
character_count=0,
|
||||
image_count=0,
|
||||
video_length_seconds=0.0,
|
||||
audio_length_seconds=0.0,
|
||||
)
|
||||
if usage.prompt_tokens_details:
|
||||
prompt_tokens_details = _parse_prompt_tokens_details(usage)
|
||||
|
|
|
|||
|
|
@ -274,6 +274,7 @@ class GoogleBatchEmbeddings(VertexLLM):
|
|||
model_response=model_response,
|
||||
model=model,
|
||||
response_json=_json_response,
|
||||
resolved_files=resolved_files,
|
||||
)
|
||||
else:
|
||||
_predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore
|
||||
|
|
@ -377,6 +378,7 @@ class GoogleBatchEmbeddings(VertexLLM):
|
|||
model_response=model_response,
|
||||
model=model,
|
||||
response_json=_json_response,
|
||||
resolved_files=resolved_files,
|
||||
)
|
||||
else:
|
||||
_predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore
|
||||
|
|
|
|||
|
|
@ -4,7 +4,10 @@ Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /batc
|
|||
Why separate file? Make it easy to see how transformation works
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
from collections.abc import Mapping
|
||||
from typing import Dict, List, Optional, Sequence, Tuple
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.types.llms.vertex_ai import (
|
||||
BlobType,
|
||||
|
|
@ -13,10 +16,17 @@ from litellm.types.llms.vertex_ai import (
|
|||
FileDataType,
|
||||
GeminiEmbeddingInput,
|
||||
PartType,
|
||||
PromptTokensDetails,
|
||||
UsageMetadata,
|
||||
VertexAIBatchEmbeddingsRequestBody,
|
||||
VertexAIBatchEmbeddingsResponseObject,
|
||||
)
|
||||
from litellm.types.utils import Embedding, EmbeddingResponse, Usage
|
||||
from litellm.types.utils import (
|
||||
Embedding,
|
||||
EmbeddingResponse,
|
||||
PromptTokensDetailsWrapper,
|
||||
Usage,
|
||||
)
|
||||
from litellm.utils import get_formatted_prompt, token_counter
|
||||
|
||||
SUPPORTED_EMBEDDING_MIME_TYPES = {
|
||||
|
|
@ -294,11 +304,133 @@ def transform_openai_input_gemini_embed_content(
|
|||
return request_body
|
||||
|
||||
|
||||
_IMAGE_MIME_TYPES = frozenset({"image/png", "image/jpeg"})
|
||||
_VIDEO_TOKENS_PER_SECOND = 258.0
|
||||
_AUDIO_TOKENS_PER_SECOND = 32.0
|
||||
_usage_metadata_adapter = TypeAdapter(UsageMetadata)
|
||||
|
||||
|
||||
def _parse_usage_metadata(raw_usage_metadata: object) -> Optional[UsageMetadata]:
|
||||
if not isinstance(raw_usage_metadata, dict):
|
||||
return None
|
||||
try:
|
||||
return _usage_metadata_adapter.validate_python(raw_usage_metadata)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _flatten_input(input: GeminiEmbeddingInput) -> tuple[str, ...]:
|
||||
if isinstance(input, str):
|
||||
return (input,)
|
||||
return tuple(
|
||||
sub
|
||||
for element in input
|
||||
for sub in (element if isinstance(element, list) else [element])
|
||||
)
|
||||
|
||||
|
||||
def _is_image_element(
|
||||
element: str,
|
||||
resolved_files: Mapping[str, Mapping[str, str]],
|
||||
) -> bool:
|
||||
if element.startswith("data:") and ";base64," in element:
|
||||
try:
|
||||
mime_type, _ = _parse_data_url(element)
|
||||
except ValueError:
|
||||
return False
|
||||
return mime_type in _IMAGE_MIME_TYPES
|
||||
if _is_gcs_url(element):
|
||||
try:
|
||||
return _infer_mime_type_from_gcs_url(element) in _IMAGE_MIME_TYPES
|
||||
except ValueError:
|
||||
return False
|
||||
if _is_file_reference(element):
|
||||
file_info = resolved_files.get(element)
|
||||
return file_info is not None and file_info.get("mime_type") in _IMAGE_MIME_TYPES
|
||||
return False
|
||||
|
||||
|
||||
def _count_input_images(
|
||||
input: GeminiEmbeddingInput,
|
||||
resolved_files: Mapping[str, Mapping[str, str]],
|
||||
) -> int:
|
||||
return sum(
|
||||
1
|
||||
for element in _flatten_input(input)
|
||||
if _is_image_element(element, resolved_files)
|
||||
)
|
||||
|
||||
|
||||
def _tokens_for_modality(details: Sequence[PromptTokensDetails], modality: str) -> int:
|
||||
return sum(
|
||||
detail["tokenCount"] for detail in details if detail["modality"] == modality
|
||||
)
|
||||
|
||||
|
||||
def _fallback_usage(input: GeminiEmbeddingInput, model: str) -> Usage:
|
||||
if _is_multimodal_input(input):
|
||||
return Usage(prompt_tokens=0, total_tokens=0)
|
||||
input_text = get_formatted_prompt(data={"input": input}, call_type="embedding")
|
||||
prompt_tokens = token_counter(model=model, text=input_text)
|
||||
return Usage(prompt_tokens=prompt_tokens, total_tokens=prompt_tokens)
|
||||
|
||||
|
||||
def _usage_from_embed_content_response(
|
||||
input: GeminiEmbeddingInput,
|
||||
model: str,
|
||||
raw_usage_metadata: object,
|
||||
resolved_files: Mapping[str, Mapping[str, str]],
|
||||
) -> Usage:
|
||||
usage_metadata = _parse_usage_metadata(raw_usage_metadata)
|
||||
if usage_metadata is None:
|
||||
return _fallback_usage(input, model)
|
||||
|
||||
prompt_tokens = usage_metadata.get("promptTokenCount", 0)
|
||||
total_tokens = usage_metadata.get("totalTokenCount") or prompt_tokens
|
||||
|
||||
details: Sequence[PromptTokensDetails] = (
|
||||
usage_metadata.get("promptTokensDetails") or ()
|
||||
)
|
||||
text_tokens = _tokens_for_modality(details, "TEXT")
|
||||
audio_tokens = _tokens_for_modality(details, "AUDIO")
|
||||
video_tokens = _tokens_for_modality(details, "VIDEO")
|
||||
image_count = _count_input_images(input, resolved_files)
|
||||
|
||||
video_length_seconds = (
|
||||
video_tokens / _VIDEO_TOKENS_PER_SECOND if video_tokens > 0 else 0.0
|
||||
)
|
||||
audio_length_seconds = (
|
||||
audio_tokens / _AUDIO_TOKENS_PER_SECOND if audio_tokens > 0 else 0.0
|
||||
)
|
||||
|
||||
# generic_cost_per_token rewrites text_tokens to the full prompt minus
|
||||
# other modalities when both text_tokens and image_count are zero. For
|
||||
# video, that misallocates video tokens to text; a 1-token floor sidesteps
|
||||
# the rewrite and keeps billing on input_cost_per_video_per_second.
|
||||
needs_video_text_floor = (
|
||||
video_length_seconds > 0 and text_tokens == 0 and image_count == 0
|
||||
)
|
||||
resolved_text_tokens = 1 if needs_video_text_floor else text_tokens
|
||||
|
||||
return Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
total_tokens=total_tokens,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
text_tokens=resolved_text_tokens,
|
||||
audio_tokens=audio_tokens,
|
||||
image_count=image_count,
|
||||
video_length_seconds=video_length_seconds,
|
||||
audio_length_seconds=audio_length_seconds,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def process_embed_content_response(
|
||||
input: GeminiEmbeddingInput,
|
||||
model_response: EmbeddingResponse,
|
||||
model: str,
|
||||
response_json: dict,
|
||||
resolved_files: Mapping[str, Mapping[str, str]] | None = None,
|
||||
) -> EmbeddingResponse:
|
||||
"""
|
||||
Process Gemini embedContent response (single embedding for multimodal input).
|
||||
|
|
@ -308,6 +440,8 @@ def process_embed_content_response(
|
|||
model_response: EmbeddingResponse to populate
|
||||
model: Model name
|
||||
response_json: Raw JSON response from embedContent endpoint
|
||||
resolved_files: Mapping of file references (files/abc) to {mime_type, uri},
|
||||
used to bill resolved image references at the per-image rate
|
||||
|
||||
Returns:
|
||||
EmbeddingResponse with single embedding
|
||||
|
|
@ -327,14 +461,11 @@ def process_embed_content_response(
|
|||
|
||||
model_response.data = [openai_embedding]
|
||||
model_response.model = model
|
||||
|
||||
if _is_multimodal_input(input):
|
||||
prompt_tokens = 0
|
||||
else:
|
||||
input_text = get_formatted_prompt(data={"input": input}, call_type="embedding")
|
||||
prompt_tokens = token_counter(model=model, text=input_text)
|
||||
model_response.usage = Usage(
|
||||
prompt_tokens=prompt_tokens, total_tokens=prompt_tokens
|
||||
model_response.usage = _usage_from_embed_content_response(
|
||||
input=input,
|
||||
model=model,
|
||||
raw_usage_metadata=response_json.get("usageMetadata"),
|
||||
resolved_files=resolved_files or {},
|
||||
)
|
||||
|
||||
return model_response
|
||||
|
|
|
|||
|
|
@ -1539,6 +1539,9 @@ class PromptTokensDetailsWrapper(
|
|||
video_length_seconds: Optional[float] = None
|
||||
"""Length of videos sent to the model. Used for Vertex AI multimodal embeddings."""
|
||||
|
||||
audio_length_seconds: Optional[float] = None
|
||||
"""Length of audio sent to the model. Used for multimodal embeddings priced per audio-second."""
|
||||
|
||||
cache_creation_tokens: Optional[int] = None
|
||||
"""Number of cache creation tokens sent to the model. Used for Anthropic prompt caching."""
|
||||
|
||||
|
|
@ -1553,6 +1556,8 @@ class PromptTokensDetailsWrapper(
|
|||
del self.image_count
|
||||
if self.video_length_seconds is None:
|
||||
del self.video_length_seconds
|
||||
if self.audio_length_seconds is None:
|
||||
del self.audio_length_seconds
|
||||
if self.web_search_requests is None:
|
||||
del self.web_search_requests
|
||||
if self.cache_creation_tokens is None:
|
||||
|
|
|
|||
|
|
@ -951,6 +951,7 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details():
|
|||
"character_count": 0,
|
||||
"image_count": 0,
|
||||
"video_length_seconds": 0.0,
|
||||
"audio_length_seconds": 0.0,
|
||||
}
|
||||
|
||||
model_info: ModelInfo = {}
|
||||
|
|
|
|||
|
|
@ -10,9 +10,11 @@ Covers:
|
|||
|
||||
import pytest
|
||||
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
|
||||
from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import (
|
||||
_build_part_for_input,
|
||||
_is_multimodal_input,
|
||||
process_embed_content_response,
|
||||
process_response,
|
||||
transform_openai_input_gemini_content,
|
||||
transform_openai_input_gemini_embed_content,
|
||||
|
|
@ -72,7 +74,9 @@ class TestBuildPartForInput:
|
|||
assert part["file_data"]["file_uri"] == GCS_URL
|
||||
|
||||
def test_file_reference_resolved(self):
|
||||
resolved = {"files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"}}
|
||||
resolved = {
|
||||
"files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"}
|
||||
}
|
||||
part = _build_part_for_input("files/abc", resolved_files=resolved)
|
||||
assert part["file_data"] is not None
|
||||
assert part["file_data"]["mime_type"] == "image/jpeg"
|
||||
|
|
@ -94,7 +98,9 @@ class TestTransformOpenaiInputGeminiContent:
|
|||
|
||||
def test_multiple_texts(self):
|
||||
result = transform_openai_input_gemini_content(
|
||||
input=["hello", "world"], model="gemini-embedding-2-preview", optional_params={}
|
||||
input=["hello", "world"],
|
||||
model="gemini-embedding-2-preview",
|
||||
optional_params={},
|
||||
)
|
||||
assert len(result["requests"]) == 2
|
||||
assert result["requests"][0]["content"]["parts"][0]["text"] == "hello"
|
||||
|
|
@ -109,7 +115,10 @@ class TestTransformOpenaiInputGeminiContent:
|
|||
)
|
||||
assert len(result["requests"]) == 2
|
||||
# First request is text
|
||||
assert result["requests"][0]["content"]["parts"][0]["text"] == "The food was delicious"
|
||||
assert (
|
||||
result["requests"][0]["content"]["parts"][0]["text"]
|
||||
== "The food was delicious"
|
||||
)
|
||||
# Second request is image
|
||||
assert result["requests"][1]["content"]["parts"][0]["inline_data"] is not None
|
||||
|
||||
|
|
@ -288,3 +297,207 @@ class TestProcessResponse:
|
|||
model="gemini-embedding-2-preview",
|
||||
optional_params={},
|
||||
)
|
||||
|
||||
|
||||
class TestProcessEmbedContentResponseUsage:
|
||||
"""Gemini Embedding 2 embedContent usageMetadata must drive spend.
|
||||
|
||||
Regression for multimodal calls recording prompt_tokens=0 / spend=$0.
|
||||
"""
|
||||
|
||||
MODEL = "gemini-embedding-2"
|
||||
|
||||
def test_multimodal_image_preserves_usage_metadata(self):
|
||||
response_json = {
|
||||
"embedding": {"values": [0.1, 0.2, 0.3]},
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 258,
|
||||
"totalTokenCount": 258,
|
||||
"promptTokensDetails": [{"modality": "IMAGE", "tokenCount": 258}],
|
||||
},
|
||||
}
|
||||
result = process_embed_content_response(
|
||||
input=[IMAGE_DATA_URI],
|
||||
model_response=EmbeddingResponse(),
|
||||
model=self.MODEL,
|
||||
response_json=response_json,
|
||||
)
|
||||
assert result.usage.prompt_tokens == 258
|
||||
assert result.usage.total_tokens == 258
|
||||
assert result.usage.prompt_tokens_details.image_count == 1
|
||||
|
||||
prompt_cost, _ = generic_cost_per_token(
|
||||
model=self.MODEL,
|
||||
usage=result.usage,
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
assert prompt_cost > 0
|
||||
|
||||
def test_text_modality_detail_populated(self):
|
||||
response_json = {
|
||||
"embedding": {"values": [0.1, 0.2]},
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 12,
|
||||
"totalTokenCount": 12,
|
||||
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}],
|
||||
},
|
||||
}
|
||||
result = process_embed_content_response(
|
||||
input="a short caption",
|
||||
model_response=EmbeddingResponse(),
|
||||
model=self.MODEL,
|
||||
response_json=response_json,
|
||||
)
|
||||
assert result.usage.prompt_tokens == 12
|
||||
assert result.usage.prompt_tokens_details.text_tokens == 12
|
||||
|
||||
prompt_cost, _ = generic_cost_per_token(
|
||||
model=self.MODEL,
|
||||
usage=result.usage,
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
assert prompt_cost > 0
|
||||
|
||||
def test_video_modality_derives_seconds_and_text_floor(self):
|
||||
response_json = {
|
||||
"embedding": {"values": [0.1]},
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 516,
|
||||
"totalTokenCount": 516,
|
||||
"promptTokensDetails": [{"modality": "VIDEO", "tokenCount": 516}],
|
||||
},
|
||||
}
|
||||
result = process_embed_content_response(
|
||||
input=["gs://bucket/clip.mp4"],
|
||||
model_response=EmbeddingResponse(),
|
||||
model=self.MODEL,
|
||||
response_json=response_json,
|
||||
)
|
||||
assert result.usage.prompt_tokens == 516
|
||||
assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx(
|
||||
2.0
|
||||
)
|
||||
assert result.usage.prompt_tokens_details.text_tokens == 1
|
||||
|
||||
def test_missing_usage_metadata_does_not_estimate_from_base64(self):
|
||||
response_json = {"embedding": {"values": [0.1, 0.2]}}
|
||||
result = process_embed_content_response(
|
||||
input=[IMAGE_DATA_URI],
|
||||
model_response=EmbeddingResponse(),
|
||||
model=self.MODEL,
|
||||
response_json=response_json,
|
||||
)
|
||||
assert result.usage.prompt_tokens == 0
|
||||
assert result.usage.total_tokens == 0
|
||||
|
||||
def test_missing_usage_metadata_text_falls_back_to_token_counter(self):
|
||||
response_json = {"embedding": {"values": [0.1, 0.2]}}
|
||||
result = process_embed_content_response(
|
||||
input="hello world this is plain text",
|
||||
model_response=EmbeddingResponse(),
|
||||
model=self.MODEL,
|
||||
response_json=response_json,
|
||||
)
|
||||
assert result.usage.prompt_tokens > 0
|
||||
|
||||
def test_file_reference_image_billed_per_image_not_text(self):
|
||||
"""files/... image refs must bill per-image, not at the text token rate."""
|
||||
response_json = {
|
||||
"embedding": {"values": [0.1, 0.2, 0.3]},
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 258,
|
||||
"totalTokenCount": 258,
|
||||
"promptTokensDetails": [{"modality": "IMAGE", "tokenCount": 258}],
|
||||
},
|
||||
}
|
||||
result = process_embed_content_response(
|
||||
input=["files/img123"],
|
||||
model_response=EmbeddingResponse(),
|
||||
model=self.MODEL,
|
||||
response_json=response_json,
|
||||
resolved_files={
|
||||
"files/img123": {
|
||||
"mime_type": "image/png",
|
||||
"uri": "https://example.com/img123",
|
||||
}
|
||||
},
|
||||
)
|
||||
assert result.usage.prompt_tokens_details.image_count == 1
|
||||
assert result.usage.prompt_tokens_details.text_tokens == 0
|
||||
|
||||
prompt_cost, _ = generic_cost_per_token(
|
||||
model=self.MODEL,
|
||||
usage=result.usage,
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
assert prompt_cost == pytest.approx(0.00012)
|
||||
|
||||
def test_file_reference_non_image_not_counted_as_image(self):
|
||||
"""A files/... ref resolving to a non-image mime must not be image-counted."""
|
||||
response_json = {
|
||||
"embedding": {"values": [0.1, 0.2]},
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 64,
|
||||
"totalTokenCount": 64,
|
||||
"promptTokensDetails": [{"modality": "AUDIO", "tokenCount": 64}],
|
||||
},
|
||||
}
|
||||
result = process_embed_content_response(
|
||||
input=["files/clip1"],
|
||||
model_response=EmbeddingResponse(),
|
||||
model=self.MODEL,
|
||||
response_json=response_json,
|
||||
resolved_files={
|
||||
"files/clip1": {
|
||||
"mime_type": "audio/mpeg",
|
||||
"uri": "https://example.com/clip1",
|
||||
}
|
||||
},
|
||||
)
|
||||
assert result.usage.prompt_tokens_details.image_count == 0
|
||||
assert result.usage.prompt_tokens_details.audio_tokens == 64
|
||||
assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx(
|
||||
2.0
|
||||
)
|
||||
|
||||
prompt_cost, _ = generic_cost_per_token(
|
||||
model=self.MODEL,
|
||||
usage=result.usage,
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
assert prompt_cost == pytest.approx(2.0 * 0.00016)
|
||||
|
||||
def test_video_plus_audio_does_not_double_bill_text(self):
|
||||
"""Video+audio responses must not get video tokens reassigned to text."""
|
||||
response_json = {
|
||||
"embedding": {"values": [0.1]},
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 580,
|
||||
"totalTokenCount": 580,
|
||||
"promptTokensDetails": [
|
||||
{"modality": "VIDEO", "tokenCount": 516},
|
||||
{"modality": "AUDIO", "tokenCount": 64},
|
||||
],
|
||||
},
|
||||
}
|
||||
result = process_embed_content_response(
|
||||
input=["gs://bucket/clip.mp4"],
|
||||
model_response=EmbeddingResponse(),
|
||||
model=self.MODEL,
|
||||
response_json=response_json,
|
||||
)
|
||||
assert result.usage.prompt_tokens_details.text_tokens == 1
|
||||
assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx(
|
||||
2.0
|
||||
)
|
||||
assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx(
|
||||
2.0
|
||||
)
|
||||
|
||||
prompt_cost, _ = generic_cost_per_token(
|
||||
model=self.MODEL,
|
||||
usage=result.usage,
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
# 1 floor text token at 2e-7 + 2s of video at 7.9e-4 + 2s of audio at 1.6e-4
|
||||
assert prompt_cost == pytest.approx(1 * 2e-7 + 2 * 0.00079 + 2 * 0.00016)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue