diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index b1af7ed2ec3..96dc1e58338 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -8,6 +8,7 @@ from typing import ( Any, AsyncIterator, Dict, + Iterator, List, Optional, Protocol, @@ -190,6 +191,13 @@ def get_vendor_from_model(model: str) -> OCIVendors: STREAMING_TIMEOUT = 60 * 5 +def _split_sse_text(text: str) -> Iterator[str]: + """Split a possibly-batched SSE text block into individual event strings.""" + for chunk in text.split("\n\n"): + if chunk: + yield chunk.strip() + + class OCIChatConfig(BaseConfig): """ Configuration class for OCI's API interface. @@ -1110,8 +1118,12 @@ class OCIChatConfig(BaseConfig): completion_stream = response.iter_text() + def split_chunks(stream: Iterator[str]) -> Iterator[str]: + for item in stream: + yield from _split_sse_text(item) + streaming_response = OCIStreamWrapper( - completion_stream=completion_stream, + completion_stream=split_chunks(completion_stream), model=model, custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, @@ -1155,12 +1167,10 @@ class OCIChatConfig(BaseConfig): completion_stream = response.aiter_text() - async def split_chunks(completion_stream: AsyncIterator[str]): - async for item in completion_stream: - for chunk in item.split("\n\n"): - if not chunk: - continue - yield chunk.strip() + async def split_chunks(stream: AsyncIterator[str]) -> AsyncIterator[str]: + async for item in stream: + for chunk in _split_sse_text(item): + yield chunk streaming_response = OCIStreamWrapper( completion_stream=split_chunks(completion_stream), diff --git a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py index 3b53f9de714..4248d65316f 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py +++ b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py @@ -11,7 +11,7 @@ import litellm sys.path.insert(0, os.path.abspath("../../../../..")) from litellm import ModelResponse -from litellm.llms.oci.chat.transformation import OCIChatConfig, OCIRequestWrapper, version +from litellm.llms.oci.chat.transformation import OCIChatConfig, OCIRequestWrapper, version, _split_sse_text TEST_MODEL_NAME = "xai.grok-4" TEST_MODEL = f"oci/{TEST_MODEL_NAME}" @@ -741,3 +741,73 @@ class TestOCISignerSupport: ) assert wrapper.path_url == "/api/v1/chat" + + +class TestSplitSseText: + """Tests for the _split_sse_text helper used by both sync and async streaming.""" + + def test_single_event(self): + result = list(_split_sse_text("data: {\"id\": \"1\"}")) + assert result == ["data: {\"id\": \"1\"}"] + + def test_multiple_events(self): + text = "data: {\"id\": \"1\"}\n\ndata: {\"id\": \"2\"}" + result = list(_split_sse_text(text)) + assert result == ["data: {\"id\": \"1\"}", "data: {\"id\": \"2\"}"] + + def test_empty_string(self): + assert list(_split_sse_text("")) == [] + + def test_only_separators(self): + assert list(_split_sse_text("\n\n\n\n")) == [] + + def test_strips_whitespace(self): + text = " data: {\"id\": \"1\"} \n\n data: {\"id\": \"2\"} " + result = list(_split_sse_text(text)) + assert result == ["data: {\"id\": \"1\"}", "data: {\"id\": \"2\"}"] + + def test_three_events(self): + text = "data: {\"a\":1}\n\ndata: {\"b\":2}\n\ndata: {\"c\":3}" + result = list(_split_sse_text(text)) + assert len(result) == 3 + + +class TestSyncStreamSplitChunks: + """Tests that the sync streaming path correctly splits batched SSE chunks.""" + + def test_sync_stream_wrapper_splits_batched_chunks(self): + """Verify get_sync_custom_stream_wrapper splits multi-event HTTP chunks.""" + from unittest.mock import MagicMock, patch + + config = OCIChatConfig() + + # Simulate an HTTP response whose iter_text() yields a batched chunk + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.iter_text.return_value = iter([ + "data: {\"id\":\"1\"}\n\ndata: {\"id\":\"2\"}", + "data: {\"id\":\"3\"}", + ]) + + mock_client = MagicMock() + mock_client.post.return_value = mock_response + + with patch("litellm.llms.oci.chat.transformation.track_llm_api_timing", lambda: lambda f: f): + result = config.get_sync_custom_stream_wrapper( + model=TEST_MODEL_NAME, + custom_llm_provider="oci", + logging_obj=MagicMock(), + api_base="https://example.com/chat", + headers={}, + data={"test": "data"}, + messages=[], + client=mock_client, + ) + + # Consume the stream wrapper's completion_stream to verify splitting + chunks = list(result.completion_stream) + assert chunks == [ + "data: {\"id\":\"1\"}", + "data: {\"id\":\"2\"}", + "data: {\"id\":\"3\"}", + ]