fix(oci): pin one response id per streamed completion, skip the [DONE] sentinel

OCIStreamWrapper.chunk_creator built every chunk straight from the apiFormat
handlers, so it never reached model_response_creator and OCI streams came back
with a fresh chatcmpl id, a drifting created value and no model on every chunk.
Both exits now go through the shared creator.

The GENERIC apiFormat also closes its stream with a literal `data: [DONE]` line,
which chunk_creator json-parsed and turned into a 500 on every OCI streaming
completion. It is skipped now.
This commit is contained in:
mateo-berri 2026-09-03 00:01:44 -07:00
parent ff17e8b987
commit 0135e1d634
2 changed files with 157 additions and 4 deletions

View file

@ -745,13 +745,21 @@ class OCIStreamWrapper(CustomStreamWrapper):
# single-event case (terminal chunk carries the only copy of the text).
self._cohere_text_emitted = False
def chunk_creator(self, chunk: Any) -> ModelResponseStream:
def _with_stream_identity(self, parsed: ModelResponseStream) -> ModelResponseStream:
model_response: Final = self.model_response_creator()
model_response.choices = parsed.choices
return model_response
def chunk_creator(self, chunk: Any) -> ModelResponseStream | None:
if not isinstance(chunk, str):
raise ValueError(f"Chunk is not a string: {chunk}")
if not chunk.startswith("data:"):
raise ValueError(f"Chunk does not start with 'data:': {chunk}")
payload: Final = chunk[5:].strip()
if payload == "[DONE]":
return None
try:
dict_chunk: Final = json.loads(chunk[5:])
dict_chunk: Final = json.loads(payload)
except json.JSONDecodeError as e:
raise OCIError(
status_code=500,
@ -774,8 +782,8 @@ class OCIStreamWrapper(CustomStreamWrapper):
if getattr(choice.delta, "content", None):
self._cohere_text_emitted = True
break
return result
return handle_generic_stream_chunk(dict_chunk)
return self._with_stream_identity(result)
return self._with_stream_identity(handle_generic_stream_chunk(dict_chunk))
__all__ = [

View file

@ -1900,3 +1900,148 @@ class TestOCIImageUrlTransformation:
adapt_messages_to_generic_oci_standard(messages)
assert "image_url" in str(exc_info.value)
# ---------------------------------------------------------------------------
# OCIStreamWrapper: per-stream identity and the GENERIC `[DONE]` sentinel
# ---------------------------------------------------------------------------
import itertools
from unittest.mock import patch
from litellm.llms.oci.chat.transformation import OCIStreamWrapper, _iter_sse_events
_STREAM_GENERIC_MODEL = "xai.grok-4"
_STREAM_COHERE_MODEL = "cohere.command-latest"
_GENERIC_TEXT_EVENT = (
'data: {{"index":0,"message":{{"role":"ASSISTANT","content":[{{"type":"TEXT","text":"{text}"}}]}},"pad":"aaa"}}'
)
_GENERIC_TERMINAL_EVENT = (
'data: {"message":{"role":"ASSISTANT","content":[{"type":"TEXT","text":""}]},"finishReason":"stop","pad":"a"}'
)
_COHERE_TEXT_EVENT = 'data: {{"apiFormat":"COHERE","text":"{text}","pad":"aaaaaa"}}'
_COHERE_TERMINAL_EVENT = (
'data: {"apiFormat":"COHERE","text":"123","finishReason":"COMPLETE",'
'"chatHistory":[{"role":"USER","message":"count"},{"role":"CHATBOT","message":"123"}]}'
)
def _make_stream_wrapper(model: str) -> OCIStreamWrapper:
logging_obj = MagicMock()
logging_obj.model_call_details = {"custom_llm_provider": "oci", "litellm_params": {}}
return OCIStreamWrapper(
completion_stream=iter([]),
model=model,
custom_llm_provider="oci",
logging_obj=logging_obj,
)
def _ticking_clock():
"""A ``time.time`` stand-in that advances a full second on every call.
Without it the whole test runs inside one wall-clock second, so a per-chunk
``created`` would coincidentally match and the drift would go unnoticed.
"""
return itertools.count(1_700_000_000.0)
class TestOCIStreamWrapperIdentityPinning:
"""One OCI streaming completion must present one id, one created and the
wrapper's model on every chunk, the way every other provider does."""
def test_generic_stream_shares_one_id_created_and_model(self):
wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL)
events = [
_GENERIC_TEXT_EVENT.format(text="1"),
_GENERIC_TEXT_EVENT.format(text="2"),
_GENERIC_TEXT_EVENT.format(text="3"),
_GENERIC_TERMINAL_EVENT,
]
with patch("time.time", side_effect=_ticking_clock()):
chunks = [wrapper.chunk_creator(event) for event in events]
assert len(chunks) == 4
assert len({chunk.id for chunk in chunks}) == 1
assert chunks[0].id.startswith("chatcmpl-")
assert len({chunk.created for chunk in chunks}) == 1
assert {chunk.model for chunk in chunks} == {_STREAM_GENERIC_MODEL}
assert [chunk.choices[0].delta.content for chunk in chunks[:3]] == ["1", "2", "3"]
assert chunks[-1].choices[0].finish_reason == "stop"
assert all(chunk._hidden_params["custom_llm_provider"] == "oci" for chunk in chunks)
def test_cohere_stream_shares_one_id_created_and_model(self):
"""Rebuilding each chunk through the shared creator must not disturb the
Cohere bookkeeping that suppresses the terminal event's repeated text."""
wrapper = _make_stream_wrapper(_STREAM_COHERE_MODEL)
events = [
_COHERE_TEXT_EVENT.format(text="1"),
_COHERE_TEXT_EVENT.format(text="2"),
_COHERE_TEXT_EVENT.format(text="3"),
_COHERE_TERMINAL_EVENT,
]
with patch("time.time", side_effect=_ticking_clock()):
chunks = [wrapper.chunk_creator(event) for event in events]
assert len(chunks) == 4
assert len({chunk.id for chunk in chunks}) == 1
assert len({chunk.created for chunk in chunks}) == 1
assert {chunk.model for chunk in chunks} == {_STREAM_COHERE_MODEL}
assert [chunk.choices[0].delta.content for chunk in chunks[:3]] == ["1", "2", "3"]
assert chunks[-1].choices[0].finish_reason == "stop"
assert chunks[-1].choices[0].delta.content is None
assert wrapper._cohere_text_emitted is True
def test_id_is_pinned_to_the_wrapper_response_id(self):
wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL)
first = wrapper.chunk_creator(_GENERIC_TEXT_EVENT.format(text="1"))
assert wrapper.response_id == first.id
assert wrapper.created == first.created
class TestOCIStreamWrapperDoneSentinel:
"""OCI's GENERIC apiFormat closes the stream with a literal `[DONE]` line;
parsing it as JSON turned every streaming completion into a 500."""
@pytest.mark.parametrize("done_event", ["data: [DONE]", "data:[DONE]", "data: [DONE] "])
def test_done_sentinel_returns_none(self, done_event):
wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL)
assert wrapper.chunk_creator(done_event) is None
def test_done_sentinel_off_the_sse_splitter_is_skipped(self):
wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL)
wire = (
f"{_GENERIC_TEXT_EVENT.format(text='1')}\n\n"
f"{_GENERIC_TEXT_EVENT.format(text='2')}\n\n"
f"{_GENERIC_TERMINAL_EVENT}\n\n"
"data: [DONE]\n\n"
)
events = list(_iter_sse_events(iter([wire])))
assert events[-1] == "data: [DONE]"
chunks = [wrapper.chunk_creator(event) for event in events]
assert chunks[-1] is None
emitted = [chunk for chunk in chunks if chunk is not None]
assert len(emitted) == 3
assert len({chunk.id for chunk in emitted}) == 1
def test_unparseable_payload_still_raises_oci_error(self):
from litellm.llms.oci.common_utils import OCIError
wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL)
with pytest.raises(OCIError, match="Chunk cannot be parsed as JSON"):
wrapper.chunk_creator("data: not-json-at-all")
def test_done_lookalike_payload_still_raises_oci_error(self):
from litellm.llms.oci.common_utils import OCIError
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")