test(integration): audio, image and per-unit cost cases

tests/integration/_support/client.py

@@ -67,6 +67,21 @@ def request(
             headers={"Authorization": f"Bearer {self.key if key is None else key}"},
         )
 
+    def request_multipart(
+        self,
+        path: str,
+        fields: Mapping[str, str],
+        files: Mapping[str, tuple[str, bytes, str]],
+        *,
+        key: str | None = None,
+    ) -> httpx.Response:
+        return self.client.post(
+            path,
+            data=fields,
+            files=files,
+            headers={"Authorization": f"Bearer {self.key if key is None else key}"},
+        )
+
     def post(self, path: str, body: Mapping[str, JsonValue], *, key: str | None = None) -> dict[str, JsonValue]:
         response: Final = self.request("POST", path, body, key=key)
         assert response.status_code == 200, f"POST {path}: {response.status_code} {response.text}"

tests/integration/_support/upstream.py

@@ -22,6 +22,7 @@
 
 from _fake_openai_endpoint_server import chat_completions, completions, embeddings, health, moderations
 from integration.cost_calculation.cost_tracking_case import (
+    BinaryResponse,
     EventStreamResponse,
     JsonResponse,
     SseResponse,
@@ -212,6 +213,11 @@ def _response(response: StoredResponse, scenario_id: str) -> Response:
                     media_type=response.content_type,
                     status_code=response.status,
                 )
+            case BinaryResponse():
+                return Response(
+                    content=b"\x00" * response.length,
+                    media_type=response.content_type,
+                )
             case SseResponse():
                 stream_body: Final = ("\n\n".join(response.frames) + "\n\n").replace(
                     "$REQUEST_ID", scenario_id

tests/integration/contracts.json

@@ -1336,6 +1336,54 @@
     "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[us.anthropic.claude-opus-5-v1:0-stream_full_usage]": [
       "quota_management.spend_tracking.scripted_wire.logs_cost"
     ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[whisper-next-transcriptions-per-second]": [
+    "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[whisper-verbose-next-transcriptions-duration]": [
+    "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-4o-transcribe-next-transcriptions-tokens]": [
+    "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[nova-next-transcriptions-per-second]": [
+    "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-whisper-next-transcriptions-deployment]": [
+    "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[tts-next-speech-per-character]": [
+    "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[tts-next-hd-speech-per-character]": [
+    "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-tts-next-speech-deployment]": [
+    "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[dall-e-3-next-images-standard]": [
+    "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[dall-e-3-next-images-hd]": [
+    "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[dall-e-3-next-images-wide]": [
+    "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[dall-e-3-next-images-two]": [
+    "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-image-next-images-low]": [
+    "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[imagen-next-images-one]": [
+    "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[amazon-nova-canvas-next-images-one]": [
+    "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
+    "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-image-next-images-edit]": [
+    "quota_management.spend_tracking.cost_matrix.logs_cost"
+    ],
     "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[azure-text-embeddings-4-large-deployment]": [
       "quota_management.spend_tracking.cost_matrix.logs_cost"
     ],

tests/integration/cost_calculation/conftest.py

@@ -40,14 +40,14 @@ class CostRow(BaseModel):
     model_config = ConfigDict(extra="ignore")
 
     spend: float | None = None
+    status: str | None = None
     prompt_tokens: int | None = None
     completion_tokens: int | None = None
     metadata: CostMetadata | None = None
 
     @property
-    def breakdown(self) -> CostBreakdown:
-        assert self.metadata is not None and self.metadata.cost_breakdown is not None
-        return self.metadata.cost_breakdown
+    def breakdown(self) -> CostBreakdown | None:
+        return self.metadata.cost_breakdown if self.metadata is not None else None
 
 
 class FailureRow(BaseModel):
@@ -63,8 +63,7 @@ def approx_equal(actual: float, expected: float) -> bool:
     return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2)
 
 
-def assert_total_is_sum_of_components(row: CostRow, context: str) -> None:
-    breakdown: Final = row.breakdown
+def assert_total_is_sum_of_components(row: CostRow, breakdown: CostBreakdown, context: str) -> None:
     total: Final = sum(
         cost or 0.0
         for cost in (breakdown.input_cost, breakdown.output_cost, breakdown.tool_usage_cost)
@@ -83,15 +82,16 @@ def _row(value: Mapping[str, object]) -> CostRow | None:
     metadata_value: Final = value.get("metadata")
     metadata: Final = json.loads(metadata_value) if isinstance(metadata_value, str) else metadata_value
     parsed: Final = CostRow.model_validate({**value, "metadata": metadata})
-    return parsed if parsed.metadata and parsed.metadata.cost_breakdown else None
+    return parsed if parsed.metadata is not None or (parsed.spend is not None and parsed.status is not None) else None
 
 
 def poll_cost_row(key: str) -> CostRow:
     digest: Final = sha256(key.encode()).hexdigest()
 
     def read() -> CostRow | None:
         rows: Final = read_rows(
-            'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
+            'SELECT spend, status, metadata, prompt_tokens, completion_tokens '
+            'FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
             (digest,),
         )
         return next((parsed for row in rows if (parsed := _row(row)) is not None), None)

tests/integration/cost_calculation/cost_tracking_case.py

@@ -46,8 +46,15 @@ class CostMapEntry(BaseModel):
     cache_creation_input_token_cost_above_200k_tokens: float | None = None
     output_cost_per_reasoning_token: float | None = None
     input_cost_per_audio_token: float | None = None
+    input_cost_per_second: float | None = None
+    output_cost_per_second: float | None = None
+    input_cost_per_character: float | None = None
+    output_cost_per_character: float | None = None
+    input_cost_per_image: float | None = None
+    output_cost_per_image: float | None = None
     output_cost_per_audio_token: float | None = None
     input_cost_per_image_token: float | None = None
+    output_cost_per_image_token: float | None = None
     input_cost_per_video_token: float | None = None
     input_cost_per_token_above_200k_tokens: float | None = None
     output_cost_per_token_above_200k_tokens: float | None = None
@@ -69,6 +76,22 @@ class Deployment(BaseModel):
     base_model: str | None = None
 
 
+class WavUpload(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    kind: Literal["wav"]
+    seconds: float
+
+
+class PngUpload(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    kind: Literal["png"]
+
+
+Upload: TypeAlias = Annotated[WavUpload | PngUpload, Field(discriminator="kind")]
+
+
 class JsonResponse(BaseModel):
     model_config = ConfigDict(frozen=True, extra="forbid")
 
@@ -98,8 +121,15 @@ class EventStreamResponse(BaseModel):
     events: tuple[EventStreamEvent, ...]
 
 
+class BinaryResponse(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    content_type: Literal["audio/mpeg"]
+    length: int
+
+
 StoredResponse: TypeAlias = Annotated[
-    JsonResponse | SseResponse | EventStreamResponse,
+    JsonResponse | SseResponse | EventStreamResponse | BinaryResponse,
     Field(discriminator="content_type"),
 ]
 
@@ -116,6 +146,7 @@ class ExactExpected(BaseModel):
     cache_creation_cost: float | None = None
     reasoning_cost: float | None = None
     tool_usage_cost: float | None = None
+    breakdown_persisted: bool = True
 
 
 class RecountRates(BaseModel):
@@ -160,8 +191,13 @@ class CostTrackingTestCase(BaseModel):
         "/v1/rerank",
         "/v1/completions",
         "/v1/moderations",
+        "/v1/audio/transcriptions",
+        "/v1/audio/speech",
+        "/v1/images/generations",
+        "/v1/images/edits",
     ] = "/v1/chat/completions"
     deployment: Deployment | None = None
+    upload: Upload | None = None
     request: dict[str, JsonValue]
     response: StoredResponse
     expected: Expected
@@ -176,7 +212,11 @@ def litellm_model(self) -> str:
         prefix: Final = (
             "openai"
             if provider == "openai"
-            and (self.endpoint == "/v1/responses" or self.rates.mode in {"chat", "embedding", "moderation"})
+            and (
+                self.endpoint == "/v1/responses"
+                or self.rates.mode
+                in {"chat", "embedding", "moderation", "audio_transcription", "audio_speech", "image_generation"}
+            )
             else "openai/responses"
             if provider == "openai"
             else _PROVIDER_PREFIXES.get(provider)
@@ -210,11 +250,13 @@ class _CasesFile(BaseModel):
 _PROVIDER_PREFIXES: Final[Mapping[str, str]] = MappingProxyType(
     {
         "anthropic": "anthropic",
+        "bedrock": "bedrock",
         "bedrock_converse": "bedrock/converse",
+        "deepgram": "deepgram",
         "text-completion-openai": "text-completion-openai",
         "cohere": "cohere",
-        "bedrock": "bedrock",
         "vertex_ai-language-models": "vertex_ai",
+        "vertex_ai-image-models": "vertex_ai",
         "vertex_ai-embedding-models": "vertex_ai",
         "gemini": "",
         "together_ai": "",
@@ -225,25 +267,29 @@ class _CasesFile(BaseModel):
 _LITELLM_PARAMS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType(
     {
         "anthropic": MappingProxyType({}),
-        "bedrock_converse": MappingProxyType(
+        "bedrock": MappingProxyType(
             {
                 "aws_access_key_id": "AKIASCRIPTEDPROVIDER",
                 "aws_secret_access_key": "scripted-secret",
                 "aws_region_name": "us-east-1",
             }
         ),
-        "text-completion-openai": MappingProxyType({}),
-        "cohere": MappingProxyType({}),
-        "bedrock": MappingProxyType(
+        "bedrock_converse": MappingProxyType(
             {
                 "aws_access_key_id": "AKIASCRIPTEDPROVIDER",
                 "aws_secret_access_key": "scripted-secret",
                 "aws_region_name": "us-east-1",
             }
         ),
+        "deepgram": MappingProxyType({}),
+        "text-completion-openai": MappingProxyType({}),
+        "cohere": MappingProxyType({}),
         "vertex_ai-language-models": MappingProxyType(
             {"vertex_project": "cc-scripted-project", "vertex_location": "us-central1"}
         ),
+        "vertex_ai-image-models": MappingProxyType(
+            {"vertex_project": "cc-scripted-project", "vertex_location": "us-central1"}
+        ),
         "vertex_ai-embedding-models": MappingProxyType(
             {"vertex_project": "cc-scripted-project", "vertex_location": "us-central1"}
         ),