Merge 62b14cd8dd into bd8eeedfc2

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,

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