mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
Merge 62b14cd8dd into 0aa7079f04
This commit is contained in:
commit
a30afc3ea9
2 changed files with 61 additions and 5 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue