mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(cost): detect file_citation annotations moved to provider_specific_fields
File search cost detection only looked at message.annotations, but partition_annotations now moves every non-url_citation annotation into message.provider_specific_fields["annotations"]. Chat completions that cited a file were therefore no longer billed for file search. Fixes #49126
This commit is contained in:
parent
4ce724e3df
commit
157eb7adbb
2 changed files with 117 additions and 0 deletions
|
|
@ -618,6 +618,18 @@ class StandardBuiltInToolCostTracking:
|
|||
for annotation in annotations:
|
||||
if annotation.get("type", None) == annotation_type:
|
||||
return True
|
||||
# Non-url_citation annotations (e.g. file_citation) are moved to
|
||||
# provider_specific_fields["annotations"] by partition_annotations.
|
||||
provider_specific_fields: object = getattr(message, "provider_specific_fields", None)
|
||||
if isinstance(provider_specific_fields, dict):
|
||||
provider_annotations: object = provider_specific_fields.get("annotations")
|
||||
if isinstance(provider_annotations, list):
|
||||
for provider_annotation in provider_annotations:
|
||||
if (
|
||||
isinstance(provider_annotation, dict)
|
||||
and provider_annotation.get("type") == annotation_type
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -0,0 +1,105 @@
|
|||
"""
|
||||
File search detection must find file_citation annotations after they are
|
||||
moved to message.provider_specific_fields["annotations"].
|
||||
|
||||
See https://github.com/BerriAI/litellm/issues/49126.
|
||||
"""
|
||||
|
||||
from typing import Final
|
||||
|
||||
import litellm
|
||||
from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS
|
||||
from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
|
||||
StandardBuiltInToolCostTracking,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
|
||||
convert_to_model_response_object,
|
||||
)
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
_FILE_CITATION: Final = {
|
||||
"type": "file_citation",
|
||||
"file_id": "file-abc123",
|
||||
"filename": "doc.pdf",
|
||||
"index": 0,
|
||||
}
|
||||
|
||||
|
||||
def _streamed_response(annotations: tuple[object, ...]) -> ModelResponse:
|
||||
chunks: Final = [
|
||||
litellm.ModelResponseStream(
|
||||
id="chatcmpl-1",
|
||||
model="gpt-4o",
|
||||
choices=[{"index": 0, "delta": {"role": "assistant", "content": "See the doc."}}],
|
||||
),
|
||||
litellm.ModelResponseStream(
|
||||
id="chatcmpl-1",
|
||||
model="gpt-4o",
|
||||
choices=[
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"annotations": list(annotations)},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
),
|
||||
]
|
||||
response: Final = litellm.stream_chunk_builder(chunks=chunks)
|
||||
assert isinstance(response, ModelResponse)
|
||||
return response
|
||||
|
||||
|
||||
def _non_streamed_response(annotations: tuple[object, ...]) -> ModelResponse:
|
||||
response: Final = convert_to_model_response_object(
|
||||
response_object={
|
||||
"id": "chatcmpl-1",
|
||||
"object": "chat.completion",
|
||||
"created": 0,
|
||||
"model": "gpt-4o",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "See the doc.",
|
||||
"annotations": list(annotations),
|
||||
},
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
},
|
||||
model_response_object=ModelResponse(),
|
||||
)
|
||||
assert isinstance(response, ModelResponse)
|
||||
return response
|
||||
|
||||
|
||||
def test_streamed_file_citation_is_detected_as_file_search() -> None:
|
||||
response: Final = _streamed_response((_FILE_CITATION,))
|
||||
assert StandardBuiltInToolCostTracking.response_object_includes_file_search_call(response_object=response)
|
||||
|
||||
|
||||
def test_non_streamed_file_citation_is_detected_as_file_search() -> None:
|
||||
response: Final = _non_streamed_response((_FILE_CITATION,))
|
||||
assert StandardBuiltInToolCostTracking.response_object_includes_file_search_call(response_object=response)
|
||||
|
||||
|
||||
def test_url_citation_only_is_not_file_search() -> None:
|
||||
url_citation: Final = {
|
||||
"type": "url_citation",
|
||||
"url_citation": {"start_index": 0, "end_index": 3, "title": "t", "url": "https://example.com"},
|
||||
}
|
||||
response: Final = _streamed_response((url_citation,))
|
||||
assert not StandardBuiltInToolCostTracking.response_object_includes_file_search_call(response_object=response)
|
||||
|
||||
|
||||
def test_streamed_file_citation_bills_file_search_cost() -> None:
|
||||
response: Final = _streamed_response((_FILE_CITATION,))
|
||||
cost: Final = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools(
|
||||
model="gpt-4o",
|
||||
response_object=response,
|
||||
custom_llm_provider="openai",
|
||||
standard_built_in_tools_params={"file_search": {"type": "file_search", "vector_store_ids": ["vs_1"]}},
|
||||
)
|
||||
assert cost == OPENAI_FILE_SEARCH_COST_PER_1K_CALLS
|
||||
Loading…
Add table
Reference in a new issue