diff --git a/litellm/integrations/datadog/datadog_cost_management.py b/litellm/integrations/datadog/datadog_cost_management.py index d25f9355d7b..d4b75fc0bce 100644 --- a/litellm/integrations/datadog/datadog_cost_management.py +++ b/litellm/integrations/datadog/datadog_cost_management.py @@ -205,16 +205,16 @@ class DatadogCostManagementLogger(CustomBatchLogger): # see _RESERVED_TAG_KEYS for the rationale. if self.cost_tag_keys: allow: Final = set(self.cost_tag_keys) + client_allow: Final = allow - {"is_silent_experiment"} for rt in log.get("request_tags") or []: if not isinstance(rt, str) or ":" not in rt: continue k, _, v = rt.partition(":") - if k in allow and v: + if k in client_allow and v: self._set_custom_tag(tags, k, v) for k, v in metadata.items(): if k in allow and v is not None and not isinstance(v, (dict, list)): self._set_custom_tag(tags, k, str(v)) - client_allow: Final = allow - {"is_silent_experiment"} for nested_key in ("spend_logs_metadata", "requester_metadata"): nested = metadata.get(nested_key) if isinstance(nested, dict): diff --git a/tests/unit/integrations/datadog/test_datadog_cost_management.py b/tests/unit/integrations/datadog/test_datadog_cost_management.py index 5dea4987087..3601db12909 100644 --- a/tests/unit/integrations/datadog/test_datadog_cost_management.py +++ b/tests/unit/integrations/datadog/test_datadog_cost_management.py @@ -316,22 +316,28 @@ async def test_extract_tags_nested_metadata_allowlisted(clean_env): @pytest.mark.parametrize( - ("metadata", "expected_tag"), + ("metadata", "request_tags", "expected_tag"), [ - ({"is_silent_experiment": True, "requester_metadata": {"is_silent_experiment": False}}, "true"), + ( + {"is_silent_experiment": True, "requester_metadata": {"is_silent_experiment": False}}, + ["is_silent_experiment:false"], + "true", + ), ( { "requester_metadata": {"is_silent_experiment": True}, "spend_logs_metadata": {"is_silent_experiment": "true"}, }, + [], None, ), + ({}, ["is_silent_experiment:true"], None), ], - ids=["mirrored call", "direct call with a client flag"], + ids=["mirrored call", "direct call with a client flag", "direct call with a client request tag"], ) @pytest.mark.asyncio async def test_extract_tags_takes_the_silent_experiment_flag_only_from_the_router( - clean_env: None, metadata: dict[str, object], expected_tag: str | None + clean_env: None, metadata: dict[str, object], request_tags: list[str], expected_tag: str | None ): logger: Final = DatadogCostManagementLogger(cost_tag_keys=["is_silent_experiment"]) log: Final = StandardLoggingPayload( @@ -340,6 +346,7 @@ async def test_extract_tags_takes_the_silent_experiment_flag_only_from_the_route response_cost=0.01, startTime=time.time(), metadata=metadata, + request_tags=request_tags, ) assert logger._extract_tags(log).get("is_silent_experiment") == expected_tag diff --git a/tests/unit/test_router_silent_experiment.py b/tests/unit/test_router_silent_experiment.py index bf175012090..b696a61aa3c 100644 --- a/tests/unit/test_router_silent_experiment.py +++ b/tests/unit/test_router_silent_experiment.py @@ -58,6 +58,15 @@ class _RecordingLogger(CustomLogger): with self._logged: self._logged.wait_for(logged_enough, timeout=timeout) + def judge_flags(self) -> dict[str, object]: + return { + call["litellm_params"]["metadata"]["judge_of"]: call["standard_logging_object"]["metadata"].get( + "is_silent_experiment" + ) + for call in self.success_kwargs + if "judge_of" in call["litellm_params"]["metadata"] + } + def shadow_successes(self) -> list[dict[str, object]]: return [ call @@ -116,15 +125,6 @@ class _JudgingLogger(_RecordingLogger): ) -> None: await self._judge(kwargs) - def judge_flags(self) -> dict[str, object]: - return { - call["litellm_params"]["metadata"]["judge_of"]: call["standard_logging_object"]["metadata"].get( - "is_silent_experiment" - ) - for call in self.success_kwargs - if "judge_of" in call["litellm_params"]["metadata"] - } - def _installed(logger: _RecordingLogger) -> Generator[_RecordingLogger]: settle_loop: Final = asyncio.new_event_loop() @@ -163,7 +163,7 @@ async def _wait_for_all_successes(logger: _RecordingLogger, expected: int) -> No await asyncio.to_thread(logger.wait_until, lambda: len(logger.success_kwargs) >= expected) -async def _wait_for_judges(logger: _JudgingLogger, judged: set[str]) -> None: +async def _wait_for_judges(logger: _RecordingLogger, judged: set[str]) -> None: await asyncio.to_thread(logger.wait_until, lambda: judged <= logger.judge_flags().keys()) @@ -986,12 +986,8 @@ def test_silent_experiment_does_not_launch_from_a_shadow_request(run_silent_expe acompletion.assert_not_awaited() -@pytest.mark.asyncio -@pytest.mark.parametrize("shadow_outcome", ["pong", "litellm.InternalServerError"]) -async def test_calls_a_callback_starts_while_logging_a_mirrored_call_are_not_marked_as_mirrored( - judging_logger: _JudgingLogger, shadow_outcome: str -) -> None: - router: Final = Router( +def _judged_router(shadow_outcome: str) -> Router: + return Router( model_list=[ { "model_name": "primary-model", @@ -1012,7 +1008,39 @@ async def test_calls_a_callback_starts_while_logging_a_mirrored_call_are_not_mar }, ] ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("shadow_outcome", ["pong", "litellm.InternalServerError"]) +async def test_calls_a_callback_starts_while_logging_a_mirrored_call_are_not_marked_as_mirrored( + judging_logger: _JudgingLogger, shadow_outcome: str +) -> None: + router: Final = _judged_router(shadow_outcome) await router.acompletion(model="primary-model", messages=[{"role": "user", "content": "hi"}]) await _wait_for_judges(judging_logger, {"primary-model", "shadow-a"}) assert judging_logger.judge_flags() == {"primary-model": None, "shadow-a": None} + + +def _start_judge_call( + kwargs: Mapping[str, object], completion_response: object, start_time: object, end_time: object +) -> None: + litellm.completion( + model="openai/gpt-5.4-nano", + api_key="fake-key", + messages=[{"role": "user", "content": "score it"}], + mock_response="ok", + metadata={"judge_of": kwargs["litellm_params"]["metadata"]["model_group"]}, + ) + + +@pytest.mark.asyncio +async def test_calls_a_failure_function_callback_starts_for_a_mirrored_call_are_not_marked_as_mirrored( + sync_and_async_recording_logger: _SyncAndAsyncRecordingLogger, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "failure_callback", [_start_judge_call]) + router: Final = _judged_router("litellm.InternalServerError") + await router.acompletion(model="primary-model", messages=[{"role": "user", "content": "hi"}]) + await _wait_for_judges(sync_and_async_recording_logger, {"shadow-a"}) + + assert sync_and_async_recording_logger.judge_flags() == {"shadow-a": None}