diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 732c72c54ae..d80a3f4cbd6 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -744,11 +744,12 @@ def _steps_that_ran(result: PipelineExecutionResult) -> frozenset[str]: def _whole_response_log_waits_for_its_pipelines(data: Mapping[str, object], call_type: str) -> bool: """ The success log is parked until the response hooks finish, and the headers are built - after them, except on the LLM pass-through route. + after them. The LLM pass-through route builds its headers first, and the generateContent + adapter logs its nested completion before the response hooks run. """ logging_obj: Final = data.get("litellm_logging_obj") return ( - call_type != CallTypes.allm_passthrough_route.value + call_type not in (CallTypes.allm_passthrough_route.value, CallTypes.agenerate_content.value) and isinstance(logging_obj, Logging) and logging_obj.enqueue_deferred_logging is not None ) diff --git a/tests/integration/observability/test_guardrail_pipeline_applied_header.py b/tests/integration/observability/test_guardrail_pipeline_applied_header.py index b62e38a1a37..bebef167f92 100644 --- a/tests/integration/observability/test_guardrail_pipeline_applied_header.py +++ b/tests/integration/observability/test_guardrail_pipeline_applied_header.py @@ -90,7 +90,6 @@ _FINE_TUNING_JOB: Final[dict[str, JsonValue]] = { "validation_file": None, } _VIDEO: Final = { - "id": "video_pipeline", "object": "video", "model": "sora-2", "status": "queued", @@ -240,7 +239,7 @@ def _completion(request: Request) -> Reply: if request.target in ("/v1/images/generations", "/v1/images/edits"): return Reply(body=json.dumps({"created": 1, "data": [{"b64_json": "aW1n", "revised_prompt": None}]}).encode()) if request.target == "/v1/videos": - return Reply(body=json.dumps(_VIDEO).encode()) + return Reply(body=json.dumps({**_VIDEO, "id": f"video_{uuid.uuid4().hex}"}).encode()) if request.target == "/v1/fine_tuning/jobs": return Reply(body=json.dumps(_FINE_TUNING_JOB).encode()) if request.target.split("?", 1)[0].endswith(":generateContent"): @@ -823,7 +822,7 @@ def test_a_cache_hit_names_every_response_step_that_ran_in_the_spend_log_and_the expected: Final = [_RESPONSE_FIRST, _RESPONSE_SECOND] assert [first.status_code, cached.status_code] == [200, 200], (first.text, cached.text) assert _steps_that_ran(rig) == (_RESPONSE_FIRST, _RESPONSE_SECOND) * 2 - assert [row["cache_hit"] for row in rows] == ["False", "True"], rows + assert [row["cache_hit"] == "True" for row in rows] == [False, True], rows assert [sorted(response.headers.get(_HEADER, "").split(",")) for response in (first, cached)] == [expected] * 2 assert [sorted(object_value(row["metadata"]).get("applied_guardrails") or []) for row in rows] == [expected] * 2 assert [sorted(ast.literal_eval(str(span["attributes"]["metadata.applied_guardrails"]))) for span in spans] == [ diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index a0ba2723443..99a06f4e141 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1366,6 +1366,9 @@ async def test_pre_call_hook_names_a_step_that_ran_in_the_litellm_metadata_a_lat pytest.param( True, "allm_passthrough_route", "metadata", ["first", "second"], id="pass-through-headers-already-out" ), + pytest.param( + True, "agenerate_content", "metadata", ["first", "second"], id="generate-content-adapter-logged-its-call" + ), ], ) async def test_post_call_pipeline_drops_the_steps_that_did_not_run_only_while_nothing_has_read_them(