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:
mateo-berri 2026-09-03 01:25:32 -07:00
parent 1cfc402763
commit 2647c2890f
2 changed files with 63 additions and 3 deletions

View file

@ -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__ = [

View file

@ -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"]