mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix(oci): record finish state so [DONE] fallback keeps tool_calls reason
This commit is contained in:
parent
1cfc402763
commit
06db628d2d
2 changed files with 76 additions and 1 deletions
|
|
@ -748,6 +748,24 @@ class OCIStreamWrapper(CustomStreamWrapper):
|
|||
def _with_stream_identity(self, parsed: ModelResponseStream) -> ModelResponseStream:
|
||||
return self.model_response_creator(chunk={"choices": parsed.choices})
|
||||
|
||||
def _record_terminal_state(self, parsed: ModelResponseStream) -> None:
|
||||
"""Mirror the finish state the parent wrapper's ``chunk_creator`` would
|
||||
normally record, so its ``StopIteration`` fallback emits a trailing
|
||||
chunk with the real finish reason instead of defaulting to ``"stop"``.
|
||||
|
||||
Without this, skipping ``[DONE]`` lets the source iterator exhaust
|
||||
cleanly and the parent's ``finish_reason_handler()`` runs with
|
||||
``received_finish_reason`` and ``tool_call`` both unset, producing a
|
||||
trailing ``"stop"`` chunk that overwrites an earlier ``"tool_calls"``
|
||||
finish reason for consumers (and ``stream_chunk_builder``) that keep
|
||||
the last non-null value.
|
||||
"""
|
||||
for choice in parsed.choices:
|
||||
if getattr(choice, "finish_reason", None):
|
||||
self.received_finish_reason = choice.finish_reason
|
||||
if getattr(choice.delta, "tool_calls", None):
|
||||
self.tool_call = True
|
||||
|
||||
def chunk_creator(self, chunk: Any) -> ModelResponseStream | None:
|
||||
if not isinstance(chunk, str):
|
||||
raise ValueError(f"Chunk is not a string: {chunk}")
|
||||
|
|
@ -780,8 +798,11 @@ class OCIStreamWrapper(CustomStreamWrapper):
|
|||
if getattr(choice.delta, "content", None):
|
||||
self._cohere_text_emitted = True
|
||||
break
|
||||
self._record_terminal_state(result)
|
||||
return self._with_stream_identity(result)
|
||||
return self._with_stream_identity(handle_generic_stream_chunk(dict_chunk))
|
||||
generic_result: Final = handle_generic_stream_chunk(dict_chunk)
|
||||
self._record_terminal_state(generic_result)
|
||||
return self._with_stream_identity(generic_result)
|
||||
|
||||
|
||||
__all__ = [
|
||||
|
|
|
|||
|
|
@ -2041,3 +2041,57 @@ class TestOCIStreamWrapperDoneSentinel:
|
|||
wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL)
|
||||
with pytest.raises(OCIError, match="Chunk cannot be parsed as JSON"):
|
||||
wrapper.chunk_creator("data: [DONE] trailing garbage")
|
||||
|
||||
|
||||
class TestOCIStreamWrapperRecordsFinishState:
|
||||
"""After skipping ``[DONE]`` the source iterator ends normally, so
|
||||
``CustomStreamWrapper.__next__`` falls back to ``finish_reason_handler()``
|
||||
to emit a trailing chunk. That helper reads ``received_finish_reason`` and
|
||||
``tool_call``, so ``chunk_creator`` must mirror both onto ``self`` or the
|
||||
trailing chunk defaults to ``"stop"`` and overwrites an earlier
|
||||
``"tool_calls"`` finish reason for consumers that keep the last one."""
|
||||
|
||||
def test_generic_tool_calls_terminal_chunk_records_finish_state(self):
|
||||
wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL)
|
||||
tool_call_event = (
|
||||
'data: {"index":0,"message":{"role":"ASSISTANT","content":null,'
|
||||
'"toolCalls":[{"id":"call_1","type":"FUNCTION","name":"get_weather","arguments":"{}"}]},'
|
||||
'"finishReason":"TOOL_CALLS"}'
|
||||
)
|
||||
|
||||
wrapper.chunk_creator(tool_call_event)
|
||||
|
||||
assert wrapper.received_finish_reason == "tool_calls"
|
||||
assert wrapper.tool_call is True
|
||||
|
||||
trailing = wrapper.finish_reason_handler()
|
||||
assert trailing.choices[0].finish_reason == "tool_calls"
|
||||
|
||||
def test_generic_stop_terminal_chunk_records_finish_state(self):
|
||||
wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL)
|
||||
|
||||
wrapper.chunk_creator(_GENERIC_TERMINAL_EVENT)
|
||||
|
||||
assert wrapper.received_finish_reason == "stop"
|
||||
assert wrapper.tool_call is False
|
||||
|
||||
trailing = wrapper.finish_reason_handler()
|
||||
assert trailing.choices[0].finish_reason == "stop"
|
||||
|
||||
def test_cohere_terminal_chunk_records_finish_state(self):
|
||||
wrapper = _make_stream_wrapper(_STREAM_COHERE_MODEL)
|
||||
|
||||
wrapper.chunk_creator(_COHERE_TERMINAL_EVENT)
|
||||
|
||||
assert wrapper.received_finish_reason == "stop"
|
||||
|
||||
trailing = wrapper.finish_reason_handler()
|
||||
assert trailing.choices[0].finish_reason == "stop"
|
||||
|
||||
def test_non_terminal_generic_chunk_leaves_finish_state_untouched(self):
|
||||
wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL)
|
||||
|
||||
wrapper.chunk_creator(_GENERIC_TEXT_EVENT.format(text="hi"))
|
||||
|
||||
assert wrapper.received_finish_reason is None
|
||||
assert wrapper.tool_call is False
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue