diff --git a/app/routes/chat.py b/app/routes/chat.py index 1779d44..aba1f5d 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -54,7 +54,16 @@ def _chunk_to_dict(chunk) -> dict: - """Normalize a litellm chunk (Pydantic model or dict) into a plain dict.""" + """Normalize a litellm chunk (Pydantic model or dict) into a plain dict. + + Returns an empty dict for None chunks — some LiteLLM stream wrappers + yield None as a heartbeat/keepalive signal, and crashing on those + would kill the stream for a non-event. Other unexpected types fall + through to dict() which will raise TypeError if the object is not + iterable, surfacing the bug rather than silently swallowing it. + """ + if chunk is None: + return {} if isinstance(chunk, dict): return chunk if hasattr(chunk, "model_dump"): diff --git a/tests/unit/test_chunk_to_dict_guard.py b/tests/unit/test_chunk_to_dict_guard.py new file mode 100644 index 0000000..a7885c2 --- /dev/null +++ b/tests/unit/test_chunk_to_dict_guard.py @@ -0,0 +1,32 @@ +"""Regression: _chunk_to_dict must handle None chunks gracefully. + +Some LiteLLM stream wrappers yield None as a heartbeat/keepalive signal. +Without a guard, dict(None) raises TypeError and kills the stream for a +non-event. The function must return an empty dict for None so the stream +continues uninterrupted.""" + +from app.routes.chat import _chunk_to_dict + + +def test_chunk_to_dict_none_returns_empty(): + assert _chunk_to_dict(None) == {} + + +def test_chunk_to_dict_dict_passthrough(): + d = {"id": "1", "model": "gpt-4o"} + assert _chunk_to_dict(d) is d + + +def test_chunk_to_dict_pydantic_model(): + class FakeModel: + def model_dump(self, exclude_none=False): + return {"id": "1", "model": "gpt-4o"} + + result = _chunk_to_dict(FakeModel()) + assert result == {"id": "1", "model": "gpt-4o"} + + +def test_chunk_to_dict_iterable_fallback(): + # A tuple of pairs is iterable and dict()-able + result = _chunk_to_dict((("key", "value"),)) + assert result == {"key": "value"}