mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
fix(responses): retain final citations and converted tool indexes
This commit is contained in:
jinwukong
parent
39dd7a0428
commit
83b4f3ea15
2 changed files with 188 additions and 30 deletions
|
|
@ -1,9 +1,11 @@
|
|||
import time
|
||||
import uuid
|
||||
from collections.abc import Mapping, Sequence
|
||||
from itertools import accumulate
|
||||
from typing import Any, Final, cast
|
||||
|
||||
from openai.types.responses import ResponseFunctionToolCall, ResponseFunctionWebSearch, ResponseReasoningItem
|
||||
from pydantic import BaseModel
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params
|
||||
|
|
@ -162,6 +164,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
self._message_start_chunk_index: int = 0
|
||||
self._closed_message_text_length: int = 0
|
||||
self._completed_message_items: tuple[tuple[int, GenericResponseOutputItem], ...] = ()
|
||||
self._final_message_events_queued: bool = False
|
||||
self._pending_annotation_events: list[BaseLiteLLMOpenAIResponseObject] = []
|
||||
self._accumulated_reasoning_content_parts: list[str] = []
|
||||
self._accumulated_provider_specific_fields: dict[str, object] = {}
|
||||
|
|
@ -876,43 +879,39 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
)
|
||||
|
||||
def _message_annotations(
|
||||
self, annotations: list[GenericResponseOutputItemContentAnnotation]
|
||||
self,
|
||||
annotations: list[GenericResponseOutputItemContentAnnotation],
|
||||
*,
|
||||
start: int | None = None,
|
||||
end: int | None = None,
|
||||
) -> list[GenericResponseOutputItemContentAnnotation]:
|
||||
offset: Final = self._closed_message_text_length if start is None else start
|
||||
return [
|
||||
annotation.model_copy(
|
||||
update={
|
||||
key: value - self._closed_message_text_length
|
||||
key: value - offset
|
||||
for key, value in (("start_index", annotation.start_index), ("end_index", annotation.end_index))
|
||||
if value is not None
|
||||
}
|
||||
)
|
||||
for annotation in annotations
|
||||
if annotation.start_index is None or annotation.start_index >= self._closed_message_text_length
|
||||
if (annotation.start_index is None and end is None)
|
||||
or (
|
||||
annotation.start_index is not None
|
||||
and annotation.start_index >= offset
|
||||
and (end is None or annotation.start_index < end)
|
||||
)
|
||||
]
|
||||
|
||||
def _queue_message_done_events(self) -> None:
|
||||
def _store_message_segment(self) -> None:
|
||||
snapshot: Final = self.create_litellm_model_response(chunk_start=self._message_start_chunk_index)
|
||||
if snapshot is None:
|
||||
return
|
||||
self._sequence_number += 1
|
||||
text_done: Final = self.create_output_text_done_event(snapshot).model_copy(
|
||||
update={"sequence_number": self._sequence_number}
|
||||
)
|
||||
self._sequence_number += 1
|
||||
part_done: Final = self.create_output_content_part_done_event(snapshot).model_copy(
|
||||
update={"sequence_number": self._sequence_number}
|
||||
)
|
||||
self._sequence_number += 1
|
||||
item_done: Final = self.create_output_item_done_event(snapshot).model_copy(
|
||||
update={"sequence_number": self._sequence_number}
|
||||
)
|
||||
item_done: Final = self.create_output_item_done_event(snapshot)
|
||||
self._completed_message_items += (
|
||||
(self._message_output_index, GenericResponseOutputItem.model_validate(item_done.item.model_dump())),
|
||||
)
|
||||
self._pending_response_events.extend(self._pending_annotation_events)
|
||||
self._pending_annotation_events.clear()
|
||||
self._pending_response_events.extend((text_done, part_done, item_done))
|
||||
self._closed_message_text_length += len(text_done.text)
|
||||
self._closed_message_text_length += len(snapshot.choices[0].message.content or "")
|
||||
self._message_start_chunk_index = len(self.collected_chat_completion_chunks)
|
||||
self._message_active = False
|
||||
self.sent_annotation_events = False
|
||||
|
|
@ -920,6 +919,87 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
self.sent_output_content_part_done_event = True
|
||||
self.sent_output_item_done_event = True
|
||||
|
||||
def _messages_with_final_annotations(
|
||||
self, response: ModelResponse
|
||||
) -> tuple[tuple[int, GenericResponseOutputItem], ...]:
|
||||
annotations: Final = (
|
||||
LiteLLMCompletionResponsesConfig._transform_chat_completion_annotations_to_response_output_annotations(
|
||||
annotations=getattr(response.choices[0].message, "annotations", None)
|
||||
)
|
||||
)
|
||||
offsets: Final = tuple(
|
||||
accumulate((len(item.content[0].text or "") for _, item in self._completed_message_items), initial=0)
|
||||
)
|
||||
return tuple(
|
||||
(
|
||||
index,
|
||||
item.model_copy(
|
||||
update={
|
||||
"content": [
|
||||
item.content[0].model_copy(
|
||||
update={
|
||||
"annotations": self._message_annotations(
|
||||
annotations,
|
||||
start=offsets[position],
|
||||
end=offsets[position + 1]
|
||||
if position + 1 < len(self._completed_message_items)
|
||||
else None,
|
||||
)
|
||||
}
|
||||
)
|
||||
]
|
||||
}
|
||||
),
|
||||
)
|
||||
for position, (index, item) in enumerate(self._completed_message_items)
|
||||
)
|
||||
|
||||
def _queue_final_message_done_events(self, response: ModelResponse) -> None:
|
||||
if self._final_message_events_queued:
|
||||
return
|
||||
self._final_message_events_queued = True
|
||||
self._completed_message_items = self._messages_with_final_annotations(response)
|
||||
for index, item in self._completed_message_items:
|
||||
for annotation_index, annotation in enumerate(item.content[0].annotations or ()):
|
||||
self._sequence_number += 1
|
||||
self._pending_response_events.append(
|
||||
OutputTextAnnotationAddedEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED,
|
||||
item_id=item.id,
|
||||
output_index=index,
|
||||
content_index=0,
|
||||
annotation_index=annotation_index,
|
||||
annotation=annotation.model_dump(),
|
||||
).model_copy(update={"sequence_number": self._sequence_number})
|
||||
)
|
||||
self._sequence_number += 3
|
||||
self._pending_response_events.extend(
|
||||
(
|
||||
OutputTextDoneEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE,
|
||||
item_id=item.id,
|
||||
output_index=index,
|
||||
content_index=0,
|
||||
text=item.content[0].text or "",
|
||||
).model_copy(update={"sequence_number": self._sequence_number - 2}),
|
||||
ContentPartDoneEvent(
|
||||
type=ResponsesAPIStreamEvents.CONTENT_PART_DONE,
|
||||
item_id=item.id,
|
||||
output_index=index,
|
||||
content_index=0,
|
||||
part=ContentPartDonePartOutputText.model_validate(
|
||||
{**item.content[0].model_dump(), "logprobs": None}
|
||||
),
|
||||
).model_copy(update={"sequence_number": self._sequence_number - 1}),
|
||||
OutputItemDoneEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
|
||||
output_index=index,
|
||||
item=BaseLiteLLMOpenAIResponseObject(**item.model_dump()),
|
||||
sequence_number=self._sequence_number,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
def create_reasoning_output_item_done_event(
|
||||
self,
|
||||
reasoning_item_id: str,
|
||||
|
|
@ -991,11 +1071,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
self, litellm_complete_object: ModelResponse, *, store_message: bool = False
|
||||
) -> BaseLiteLLMOpenAIResponseObject | None:
|
||||
if store_message and self._message_active:
|
||||
self._queue_message_done_events()
|
||||
if self._pending_response_events:
|
||||
return self._pending_response_events.pop(0)
|
||||
self._store_message_segment()
|
||||
if store_message and self._completed_message_items:
|
||||
return None
|
||||
self._queue_final_message_done_events(litellm_complete_object)
|
||||
return self._pending_response_events.pop(0) if self._pending_response_events else None
|
||||
if self.sent_message_item_added_event is False:
|
||||
final_content: Final = litellm_complete_object.choices[0].message.content or ""
|
||||
if not final_content:
|
||||
|
|
@ -1187,7 +1266,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
or _delta_has_thinking_block(chunk.choices[0].delta)
|
||||
)
|
||||
):
|
||||
self._queue_message_done_events()
|
||||
self._store_message_segment()
|
||||
if (
|
||||
not self.sent_output_item_added_event
|
||||
and chunk.choices
|
||||
|
|
@ -1320,7 +1399,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
# This ensures we detect and queue annotation events from the annotation chunk
|
||||
if chunk.choices and hasattr(chunk.choices[0].delta, "annotations"):
|
||||
annotations: Final = chunk.choices[0].delta.annotations
|
||||
if annotations and self.sent_annotation_events is False:
|
||||
if annotations and self.sent_annotation_events is False and not allow_reasoning_resumption:
|
||||
# Store annotation events to emit them one by one
|
||||
response_annotations = self._message_annotations(
|
||||
LiteLLMCompletionResponsesConfig.transform_chat_completion_annotations_to_response_output_annotations(
|
||||
|
|
@ -1514,6 +1593,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
return self._tool_output_index_by_call_id.get(item.call_id or "", self._next_tool_output_index)
|
||||
if isinstance(item, ResponseFunctionWebSearch):
|
||||
return self._tool_output_index_by_call_id.get(item.id.removeprefix("ws_"), self._next_tool_output_index)
|
||||
if isinstance(item, BaseModel):
|
||||
item_id: Final[object] = item.model_dump(include={"id"}).get("id")
|
||||
if isinstance(item_id, str):
|
||||
return self._tool_output_index_by_call_id.get(item_id, self._next_tool_output_index)
|
||||
return self._next_tool_output_index
|
||||
|
||||
def _emit_terminal_response_event(
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ from litellm.types.llms.openai import (
|
|||
ChatCompletionThinkingBlock,
|
||||
ResponsesAPIStreamEvents,
|
||||
)
|
||||
from litellm.types.responses.main import build_web_search_call
|
||||
from litellm.types.responses.main import OutputCodeInterpreterCall, build_web_search_call
|
||||
from litellm.types.utils import (
|
||||
Delta,
|
||||
ModelResponse,
|
||||
|
|
@ -1611,8 +1611,9 @@ async def test_mixed_reasoning_and_text_chunk_preserves_both_segments() -> None:
|
|||
assert [item["content"][0]["text"] for item in output if item["type"] == "message"] == ["interim", "answer"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("late_annotations", [False, True])
|
||||
@pytest.mark.asyncio
|
||||
async def test_message_annotations_are_relative_to_each_segment() -> None:
|
||||
async def test_message_annotations_are_relative_to_each_segment(late_annotations: bool) -> None:
|
||||
intro: Final = ModelResponseStream(
|
||||
id=CHAT_COMPLETION_ID,
|
||||
model="test-model",
|
||||
|
|
@ -1660,11 +1661,22 @@ async def test_message_annotations_are_relative_to_each_segment() -> None:
|
|||
)
|
||||
],
|
||||
)
|
||||
delayed: Final = ModelResponseStream(
|
||||
id=CHAT_COMPLETION_ID,
|
||||
model="test-model",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
index=0,
|
||||
delta=Delta(annotations=[*intro.choices[0].delta.annotations, *answer.choices[0].delta.annotations]),
|
||||
finish_reason="stop",
|
||||
)
|
||||
],
|
||||
)
|
||||
iterator: Final = _build_iterator(
|
||||
[
|
||||
intro,
|
||||
_chunk("intro") if late_annotations else intro,
|
||||
_reasoning_block_chunk(ChatCompletionThinkingBlock(type="thinking", thinking="analysis", signature="sig")),
|
||||
answer,
|
||||
*([_chunk("answer"), delayed] if late_annotations else [answer]),
|
||||
]
|
||||
)
|
||||
events: Final = [json.loads(event.model_dump_json(exclude_none=True)) async for event in iterator]
|
||||
|
|
@ -1672,3 +1684,66 @@ async def test_message_annotations_are_relative_to_each_segment() -> None:
|
|||
assert [item["content"][0]["text"] for item in messages] == ["intro", "answer"]
|
||||
assert [item["content"][0]["annotations"][0]["start_index"] for item in messages] == [0, 0]
|
||||
assert [item["content"][0]["annotations"][0]["end_index"] for item in messages] == [5, 6]
|
||||
done_by_id: Final = {event["item"]["id"]: event for event in events if event["type"] == "response.output_item.done"}
|
||||
annotation_by_id: Final = {
|
||||
event["item_id"]: event for event in events if event["type"] == "response.output_text.annotation.added"
|
||||
}
|
||||
for message in messages:
|
||||
assert done_by_id[message["id"]]["item"]["content"] == message["content"]
|
||||
assert annotation_by_id[message["id"]]["annotation"] == message["content"][0]["annotations"][0]
|
||||
assert events.index(annotation_by_id[message["id"]]) < events.index(done_by_id[message["id"]])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_converted_code_tool_keeps_its_streamed_output_index() -> None:
|
||||
call_id: Final = "srvtoolu_code"
|
||||
code_result: Final = OutputCodeInterpreterCall(
|
||||
type="code_interpreter_call",
|
||||
id=call_id,
|
||||
code="print(1)",
|
||||
container_id="container_test",
|
||||
status="completed",
|
||||
outputs=[],
|
||||
)
|
||||
tool: Final = ModelResponseStream(
|
||||
id=CHAT_COMPLETION_ID,
|
||||
model="test-model",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
index=0,
|
||||
delta=Delta(
|
||||
tool_calls=[
|
||||
{
|
||||
"id": call_id,
|
||||
"type": "function",
|
||||
"index": 0,
|
||||
"function": {"name": "bash_code_execution", "arguments": '{"command":"print(1)"}'},
|
||||
}
|
||||
],
|
||||
provider_specific_fields={"code_interpreter_results": [code_result]},
|
||||
),
|
||||
)
|
||||
],
|
||||
)
|
||||
iterator: Final = _build_iterator(
|
||||
[
|
||||
_reasoning_block_chunk(
|
||||
ChatCompletionThinkingBlock(type="thinking", thinking="first", signature="first-sig")
|
||||
),
|
||||
tool,
|
||||
_reasoning_block_chunk(
|
||||
ChatCompletionThinkingBlock(type="thinking", thinking="second", signature="second-sig")
|
||||
),
|
||||
_chunk("answer", finish_reason="stop"),
|
||||
]
|
||||
)
|
||||
events: Final = [json.loads(event.model_dump_json(exclude_none=True)) async for event in iterator]
|
||||
tool_added: Final = next(
|
||||
event
|
||||
for event in events
|
||||
if event["type"] == "response.output_item.added" and event["item"].get("call_id") == call_id
|
||||
)
|
||||
output: Final = events[-1]["response"]["output"]
|
||||
assert [item["type"] for item in output] == ["reasoning", "code_interpreter_call", "reasoning", "message"]
|
||||
assert output[tool_added["output_index"]]["id"] == call_id
|
||||
assert output[tool_added["output_index"]]["code"] == "print(1)"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue