mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix(oci): keep the stream's own finish reason instead of a synthetic stop chunk
The OCI wrapper overrides chunk_creator wholesale, so it never recorded the finish reason or marked the terminal chunk as sent. The shared end-of-stream finalizer then appended a synthetic chunk whose finish_reason was always stop, which downgraded a tool_calls completion for any client that reads the finish reason off the last chunk.
This commit is contained in:
parent
1cfc402763
commit
2647c2890f
2 changed files with 63 additions and 3 deletions
|
|
@ -745,7 +745,13 @@ class OCIStreamWrapper(CustomStreamWrapper):
|
|||
# single-event case (terminal chunk carries the only copy of the text).
|
||||
self._cohere_text_emitted = False
|
||||
|
||||
def _with_stream_identity(self, parsed: ModelResponseStream) -> ModelResponseStream:
|
||||
def _emit_chunk(self, parsed: ModelResponseStream) -> ModelResponseStream:
|
||||
for choice in parsed.choices:
|
||||
if getattr(choice.delta, "tool_calls", None):
|
||||
self.tool_call = True
|
||||
if choice.finish_reason is not None:
|
||||
self.received_finish_reason = choice.finish_reason
|
||||
self.sent_last_chunk = True
|
||||
return self.model_response_creator(chunk={"choices": parsed.choices})
|
||||
|
||||
def chunk_creator(self, chunk: Any) -> ModelResponseStream | None:
|
||||
|
|
@ -780,8 +786,8 @@ class OCIStreamWrapper(CustomStreamWrapper):
|
|||
if getattr(choice.delta, "content", None):
|
||||
self._cohere_text_emitted = True
|
||||
break
|
||||
return self._with_stream_identity(result)
|
||||
return self._with_stream_identity(handle_generic_stream_chunk(dict_chunk))
|
||||
return self._emit_chunk(result)
|
||||
return self._emit_chunk(handle_generic_stream_chunk(dict_chunk))
|
||||
|
||||
|
||||
__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")
|
||||
|
||||
|
||||
_GENERIC_TOOL_CALL_EVENT = (
|
||||
'data: {"index":0,"message":{"role":"ASSISTANT","content":[],'
|
||||
'"toolCalls":[{"type":"FUNCTION","id":"call_1","name":"get_weather","arguments":"{}"}]}}'
|
||||
)
|
||||
_GENERIC_TOOL_TERMINAL_EVENT = (
|
||||
'data: {"index":0,"message":{"role":"ASSISTANT","content":[]},"finishReason":"TOOL_CALLS"}'
|
||||
)
|
||||
|
||||
|
||||
def _drain_stream(model: str, events: list[str]) -> list:
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {"custom_llm_provider": "oci", "litellm_params": {}}
|
||||
wrapper = OCIStreamWrapper(
|
||||
completion_stream=iter(events),
|
||||
model=model,
|
||||
custom_llm_provider="oci",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
return list(wrapper)
|
||||
|
||||
|
||||
class TestOCIStreamWrapperTerminalChunk:
|
||||
"""OCI's ``chunk_creator`` override bypasses the shared handler's
|
||||
finish-reason bookkeeping, so the shared end-of-stream finalizer used to
|
||||
append a synthetic ``stop`` chunk after OCI's own terminal chunk, silently
|
||||
downgrading a ``tool_calls`` completion for any client that reads the
|
||||
finish reason off the last chunk."""
|
||||
|
||||
def test_generic_tool_call_stream_ends_on_tool_calls(self):
|
||||
chunks = _drain_stream(
|
||||
_STREAM_GENERIC_MODEL,
|
||||
[_GENERIC_TOOL_CALL_EVENT, _GENERIC_TOOL_TERMINAL_EVENT, "data: [DONE]"],
|
||||
)
|
||||
|
||||
assert [chunk.choices[0].finish_reason for chunk in chunks] == [None, "tool_calls"]
|
||||
assert len({chunk.id for chunk in chunks}) == 1
|
||||
|
||||
def test_generic_text_stream_emits_exactly_one_finish_reason(self):
|
||||
chunks = _drain_stream(
|
||||
_STREAM_GENERIC_MODEL,
|
||||
[_GENERIC_TEXT_EVENT.format(text="1"), _GENERIC_TERMINAL_EVENT, "data: [DONE]"],
|
||||
)
|
||||
|
||||
assert [chunk.choices[0].finish_reason for chunk in chunks] == [None, "stop"]
|
||||
|
||||
def test_cohere_stream_emits_exactly_one_finish_reason(self):
|
||||
chunks = _drain_stream(
|
||||
_STREAM_COHERE_MODEL,
|
||||
[_COHERE_TEXT_EVENT.format(text="123"), _COHERE_TERMINAL_EVENT],
|
||||
)
|
||||
|
||||
assert [chunk.choices[0].finish_reason for chunk in chunks] == [None, "stop"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue