diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index ea8bf230d05..ac3fcd33d2e 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json from collections.abc import Mapping from pathlib import Path from types import MappingProxyType @@ -8,6 +9,7 @@ from typing import Annotated, Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator, model_validator CASES_PATH: Final = Path(__file__).resolve().parent / "cost_tracking_cases.json" +PRIOR_RESPONSE_ID_MARKER: Final = "$PRIOR_RESPONSE_ID" class SearchContextCostPerQuery(BaseModel): @@ -312,6 +314,21 @@ class CostTrackingTestCase(BaseModel): usage: Final = self.response.body.get("usage") return isinstance(usage, dict) and isinstance(usage.get("cost"), (int, float)) + @property + def chains_prior_response(self) -> bool: + return self.request.get("previous_response_id") == PRIOR_RESPONSE_ID_MARKER + + @property + def can_chain_prior_response(self) -> bool: + return ( + self.chains_prior_response + and self.endpoint == "/v1/responses" + and isinstance(self.response, JsonResponse) + and isinstance(self.response.body.get("id"), str) + and not isinstance(self.expected, FailureExpected) + and not (isinstance(self.expected, ExactExpected) and self.expected.rollups) + ) + class BatchOutputLine(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") @@ -681,6 +698,11 @@ def data_errors() -> tuple[str, ...]: for marker in ('"id": "call_$REQUEST_ID"', '"id": "toolu_$REQUEST_ID"') ) ) + invalid_prior_response_chains: Final = sorted( + case.name + for case in CASES + if PRIOR_RESPONSE_ID_MARKER in json.dumps(case.request) and not case.can_chain_prior_response + ) return tuple( message for message in ( @@ -700,6 +722,10 @@ def data_errors() -> tuple[str, ...]: f"pinned tool IDs contain $REQUEST_ID: {invalid_pinned_tool_ids}" if invalid_pinned_tool_ids else None, + f"{PRIOR_RESPONSE_ID_MARKER} needs a non-rollup, non-failure /v1/responses JSON response with a string id" + f" as previous_response_id: {invalid_prior_response_chains}" + if invalid_prior_response_chains + else None, ) if message is not None ) diff --git a/tests/integration/cost_calculation/cost_tracking_cases.json b/tests/integration/cost_calculation/cost_tracking_cases.json index 17ebc793fae..09dfa66012f 100644 --- a/tests/integration/cost_calculation/cost_tracking_cases.json +++ b/tests/integration/cost_calculation/cost_tracking_cases.json @@ -27733,7 +27733,7 @@ "request": { "model": "$MODEL", "input": "continue this text", - "previous_response_id": "resp_scripted_prior" + "previous_response_id": "$PRIOR_RESPONSE_ID" }, "response": { "content_type": "application/json", diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index efab17acba4..82878634677 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -9,13 +9,14 @@ import time import uuid import wave import zlib +from collections.abc import Mapping from hashlib import sha256 from itertools import islice from typing import Final, cast import httpx import pytest -from integration._support.client import JSON_OBJECT, Gateway +from integration._support.client import JSON_OBJECT, Gateway, string_value from integration._support.upstream import delete_scenario, register_scenario from integration.cost_calculation.assertions import assert_exact, assert_recount from integration.cost_calculation.conftest import ( @@ -112,6 +113,19 @@ def _replace_model(value: JsonValue, model_name: str) -> JsonValue: return value +def _prime_prior_response( + gateway: Gateway, request_path: str, request_values: Mapping[str, JsonValue], key: str +) -> str: + primed: Final = gateway.request( + "POST", + request_path, + {field: value for field, value in request_values.items() if field != "previous_response_id"}, + key=key, + ) + assert primed.is_success, f"priming response failed: {primed.status_code}: {primed.text[:400]}" + return string_value(JSON_OBJECT.validate_json(primed.content)["id"]) + + @pytest.mark.parametrize("case", _CASES) def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) -> None: marker: Final = sha256(case.name.encode()).hexdigest()[:12] @@ -175,21 +189,6 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) if isinstance(expected, ExactExpected) and expected.rollups else None ) - request_body: Final = JSON_OBJECT.validate_python( - { - **base_request_values, - **( - {"model": fallback_deployment.model_name, "fallbacks": [model_name]} - if fallback_deployment is not None - else {} - ), - **( - {"user": end_user_id, "cache": {"no-cache": True}} - if end_user_id is not None - else {} - ), - } - ) request_headers: Final = ( { "x-pass-x-scripted-scenario": scenario_id, @@ -207,6 +206,27 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) if passthrough_provider is not None else case.endpoint ) + prior_response_id: Final = ( + _prime_prior_response(gateway, request_path, base_request_values, key) + if case.chains_prior_response + else None + ) + request_body: Final = JSON_OBJECT.validate_python( + { + **base_request_values, + **( + {"model": fallback_deployment.model_name, "fallbacks": [model_name]} + if fallback_deployment is not None + else {} + ), + **( + {"user": end_user_id, "cache": {"no-cache": True}} + if end_user_id is not None + else {} + ), + **({"previous_response_id": prior_response_id} if prior_response_id is not None else {}), + } + ) if case.disconnect_after_frames is not None: with gateway.client.stream( "POST", @@ -250,7 +270,7 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) assert response.is_success, f"{case.name}: proxy returned {response.status_code}: {response.text[:400]}" if case.response.content_type == "text/event-stream": _assert_stream_has_no_error(response.text) - rows: Final = poll_rows(key, len(responses)) + rows: Final = poll_rows(key, len(responses) + (prior_response_id is not None)) if isinstance(expected, RecountExpected): row: Final = rows[0] assert_recount(case.name, expected, row) diff --git a/tests/integration/proxy_config.yaml b/tests/integration/proxy_config.yaml index 34e48228dd7..a3b07f76d2f 100644 --- a/tests/integration/proxy_config.yaml +++ b/tests/integration/proxy_config.yaml @@ -5,7 +5,6 @@ general_settings: store_model_in_db: true disable_spend_logs: false proxy_batch_write_at: 1 - allow_unmanaged_response_ids: true litellm_settings: enable_redis_auth_cache: true cache: true