diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5d21b2762cd..97f3ca68e42 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -748,6 +748,7 @@ class Logging(LiteLLMLoggingBaseClass): stream_options = None litellm_request_debug: bool = False streamed_anthropic_message_id: str | None = None + served_agentic_follow_up: bool = False bridged_model_response: ModelResponse | None = None classifier_input: Mapping[str, JsonValue] | None = None baseline_cache_context: "BaselineCacheContext | None" = None @@ -5400,6 +5401,10 @@ class Logging(LiteLLMLoggingBaseClass): def record_streamed_anthropic_message_id(self, message_id: str) -> None: self.streamed_anthropic_message_id = message_id + def record_served_agentic_follow_up(self) -> None: + """This call served an agentic follow-up's reply, so its own spend row is keyed by its upstream reply.""" + self.served_agentic_follow_up = True + def record_bridged_model_response(self, response: ModelResponse) -> None: """The chat completion a /v1/messages response was translated from. @@ -5515,6 +5520,8 @@ class Logging(LiteLLMLoggingBaseClass): json_mode=False, litellm_params={}, ) + if self.served_agentic_follow_up: + return self._first_round_row(result, httpx_response) elif result is None: verbose_logger.warning("LiteLLM: the anthropic_messages stream assembled no response, logging an empty one") return litellm.ModelResponse(model=self.model) @@ -5539,6 +5546,15 @@ class Logging(LiteLLMLoggingBaseClass): set_response_cost_in_hidden_params(result, reported_cost) return result + def _first_round_row(self, reparsed: ModelResponse, upstream: Response) -> ModelResponse: + body: Final = _LITELLM_PARAMS.validate_python(upstream.json()) + upstream_id: Final = _provider_response_id(body) + row: Final = reparsed if upstream_id is None else reparsed.model_copy(update={"id": upstream_id}) + set_response_cost_in_hidden_params( + row, native_messages_reported_cost(body.get("usage"), self.model_call_details.get("custom_llm_provider")) + ) + return row + def _translate_responses_api_response_to_model_response( self, result: ResponsesAPIResponse, request: BridgedChatRequest | None = None ) -> ModelResponse: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index c34b696bb5d..8213f9dc779 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -6188,6 +6188,7 @@ class BaseLLMHTTPHandler: **kwargs_for_followup, } ) + logging_obj.record_served_agentic_follow_up() if callback is not None: try: diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 51fbcaa54bf..80782fa89d4 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -7202,6 +7202,93 @@ def test_handle_anthropic_messages_response_logging_reparses_the_served_body_wit assert logged.usage.completion_tokens == 7 # type: ignore[attr-defined] +def _upstream_messages_reply(usage: Mapping[str, object], message_id: str | None = None) -> dict[str, object]: + return { + **({} if message_id is None else {"id": message_id}), + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "answer"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": dict(usage), + } + + +def _logging_obj_answered_by( + upstream_reply: Mapping[str, object], served_agentic_follow_up: bool, custom_llm_provider: str = "bedrock" +) -> LitellmLogging: + logging_obj: Final = _anthropic_messages_logging_obj() + logging_obj.optional_params = {} + logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider + logging_obj.model_call_details["httpx_response"] = httpx.Response( + 200, json=dict(upstream_reply), request=httpx.Request("POST", "https://upstream.test/v1/messages") + ) + if served_agentic_follow_up: + logging_obj.record_served_agentic_follow_up() + return logging_obj + + +def test_websearch_first_round_spend_row_keeps_its_own_id_when_the_follow_up_reply_is_served() -> None: + from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id + + first_round: Final = _upstream_messages_reply({"input_tokens": 100, "output_tokens": 10}, "msg_first_round") + follow_up: Final = _upstream_messages_reply({"input_tokens": 300, "output_tokens": 30}, "msg_follow_up") + + first_row: Final = _logging_obj_answered_by(first_round, True)._anthropic_messages_logged_response( + result=dict(follow_up) + ) + follow_up_row: Final = _logging_obj_answered_by(follow_up, False)._anthropic_messages_logged_response( + result=dict(follow_up) + ) + + first_usage: Final = first_row.model_dump()["usage"] + assert (first_usage["prompt_tokens"], first_usage["completion_tokens"]) == (100, 10) + assert get_spend_logs_id("anthropic_messages", first_row.model_dump(), {"litellm_call_id": "first"}) == ( + "msg_first_round" + ) + assert get_spend_logs_id("anthropic_messages", follow_up_row.model_dump(), {"litellm_call_id": "follow"}) == ( + "msg_follow_up" + ) + + +@pytest.mark.parametrize( + ("upstream_id", "served_agentic_follow_up"), + [ + pytest.param("msg_native", False, id="post_call_hook_replaced_the_served_reply"), + pytest.param(None, True, id="follow_up_served_and_upstream_reply_has_no_id"), + ], +) +def test_anthropic_messages_spend_row_keeps_the_served_id( + upstream_id: str | None, served_agentic_follow_up: bool +) -> None: + upstream: Final = _upstream_messages_reply({"input_tokens": 12, "output_tokens": 7}, upstream_id) + served: Final = _upstream_messages_reply({"input_tokens": 12, "output_tokens": 7}, "msg_served") + + logged: Final = _logging_obj_answered_by(upstream, served_agentic_follow_up)._anthropic_messages_logged_response( + result=served + ) + + assert logged.id == "msg_served" + + +def test_websearch_first_round_spend_row_bills_the_first_upstream_reported_cost() -> None: + from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER + + first_round: Final = _upstream_messages_reply( + {"input_tokens": 100, "output_tokens": 10, "opencode": {"list_cost": 0.25}}, "msg_first_round" + ) + follow_up: Final = _upstream_messages_reply( + {"input_tokens": 300, "output_tokens": 30, "opencode": {"list_cost": 1.5}}, "msg_follow_up" + ) + + first_row: Final = _logging_obj_answered_by(first_round, True, "opencode_go")._anthropic_messages_logged_response( + result=follow_up + ) + + assert first_row.hidden_params["additional_headers"][RESPONSE_COST_HEADER] == 0.25 + + @pytest.mark.parametrize( "compaction_block", [ diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 08f38352735..b51abfa1ad0 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -5579,6 +5579,66 @@ async def test_responses_agentic_followup_sends_the_plans_request_param_over_a_s assert followup_calls[0]["prompt_cache_key"] == "from-plan-params" +def _messages_reply(message_id: str, input_tokens: int, output_tokens: int) -> dict[str, object]: + return { + "id": message_id, + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "answer"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens}, + } + + +@pytest.mark.asyncio +async def test_anthropic_agentic_follow_up_keeps_the_first_rounds_spend_row_on_its_own_id() -> None: + from datetime import datetime, timezone + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id + + request_messages: Final = [{"role": "user", "content": "search"}] + first_round: Final = Logging( + model="claude-sonnet-5-5", + messages=request_messages, + stream=False, + call_type="anthropic_messages", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), + litellm_call_id="first-round", + function_id="first-round", + ) + first_round.optional_params = {} + first_round.model_call_details["httpx_response"] = httpx.Response( + 200, + json=_messages_reply("msg_first_round", 100, 10), + request=httpx.Request("POST", "https://api.anthropic.com/v1/messages"), + ) + follow_up_upstream: Final = httpx.MockTransport( + lambda request: httpx.Response(200, json=_messages_reply("msg_follow_up", 300, 30)) + ) + + served: Final = await BaseLLMHTTPHandler()._execute_anthropic_agentic_plan( + plan=AgenticLoopPlan(run_agentic_loop=True, request_patch=AgenticLoopRequestPatch(messages=request_messages)), + model="anthropic/claude-sonnet-5-5", + messages=request_messages, + anthropic_messages_optional_request_params={}, + logging_obj=first_round, + kwargs={"api_key": "test-key", "client": AsyncHTTPHandler(transport=follow_up_upstream)}, + depth=0, + max_loops=3, + fingerprints=[], + fingerprint="fp", + ) + + assert isinstance(served, dict) + assert served["id"] == "msg_follow_up" + first_round_row: Final = first_round._anthropic_messages_logged_response(result=served).model_dump() + first_round_id: Final = get_spend_logs_id("anthropic_messages", first_round_row, {"litellm_call_id": "first-round"}) + assert first_round_id == "msg_first_round" + + class _PlannedChatLoop(CustomLogger): def __init__(self, plan: AgenticLoopPlan, claimed_call_name: str = "litellm_web_search") -> None: super().__init__()