diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 54213f40ac0..c1d16b05184 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -616,13 +616,27 @@ class StandardBuiltInToolCostTracking: message: Message | None = getattr(choice, "message", None) if message is None: continue - annotations: list[ChatCompletionAnnotation] | None = getattr(message, "annotations", None) - if annotations: - for annotation in annotations: - if annotation.get("type", None) == annotation_type: - return True + for annotation in StandardBuiltInToolCostTracking._message_annotations(message): + if annotation.get("type", None) == annotation_type: + return True return False + @staticmethod + def _message_annotations(message: Message) -> list[ChatCompletionAnnotation]: + """Every annotation on the message, wherever the split put it. + + `Message.annotations` is typed for `url_citation` only, so file citations travel + under `provider_specific_fields["annotations"]` and were invisible to the cost + detector, which billed streamed file search at 0.0. + """ + annotations: list[ChatCompletionAnnotation] = list(getattr(message, "annotations", None) or []) + provider_specific_fields = getattr(message, "provider_specific_fields", None) + if isinstance(provider_specific_fields, Mapping): + nested = cast("Mapping[str, object]", provider_specific_fields).get("annotations") + if isinstance(nested, list): + annotations.extend(cast("list[ChatCompletionAnnotation]", nested)) + return annotations + @staticmethod def response_includes_output_type( response_object: ResponsesAPIResponse, diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 9da2632943c..646a9fdeda6 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -1230,3 +1230,45 @@ def test_vertex_gemini_deployment_priced_web_search_stays_per_prompt( def test_built_in_tool_lookups_ignore_malformed_request_values(kwargs: dict[str, object]) -> None: assert StandardBuiltInToolCostTracking.get_web_search_options(kwargs) is None assert StandardBuiltInToolCostTracking.get_file_search_tool_call(kwargs) is None + + +def _response_with_annotations(**message_kwargs) -> ModelResponse: + from litellm.types.utils import Choices, Message + + message = Message(content="", role="assistant", **message_kwargs) + return ModelResponse(choices=[Choices(message=message)], model="gpt-4o") + + +def test_file_citation_under_provider_specific_fields_is_a_file_search_call(): + """A streamed file citation must still be billed. + + `Message.annotations` is typed for `url_citation` only, so file citations sit in + `provider_specific_fields["annotations"]` after the provider-specific split. The + detector only read the former, so streamed file search was priced at 0. + """ + response = _response_with_annotations( + provider_specific_fields={ + "annotations": [{"type": "file_citation", "file_id": "file-abc123", "filename": "doc.pdf", "index": 0}] + } + ) + + assert StandardBuiltInToolCostTracking.response_object_includes_file_search_call(response) is True + + +def test_url_citation_is_not_a_file_search_call(): + """A url citation is a web search result, not a file search one.""" + response = _response_with_annotations( + provider_specific_fields={ + "annotations": [ + { + "type": "url_citation", + "url": "https://example.com", + "title": "example", + "start_index": 0, + "end_index": 1, + } + ] + } + ) + + assert StandardBuiltInToolCostTracking.response_object_includes_file_search_call(response) is False