fix(sap): address review feedback on stream chunk validation

Type the chunk payload as dict[str, object] and narrow choices with an
isinstance guard before iterating, fix the in-place mutation suppression
that claimed the loop never mutated, build the litellm Usage with
model_validate instead of spreading a dict[str, Any], and make the usage
regression test go through model_dump_json, the path that actually raised
MockValSer. model_dump() self-heals via __getattr__, so the old test passed
on unfixed code.

Carry the same typing into to_openai_chunk and the payload literals it
validates, so no basedpyright rule counts higher than before the change.
This commit is contained in:
ZOU Yi (BD/SWD-WDE1) 2026-09-14 11:05:04 +08:00
parent 264d05b1fa
commit a215b78d99
2 changed files with 30 additions and 25 deletions

View file

@ -49,14 +49,16 @@ class _StreamParser:
@staticmethod
def _validate_chunk(
payload: dict, # mutable-ok: normalized in place (pops empty logprobs) before validation
payload: dict[str, object], # mutable-ok: normalized in place (pops empty logprobs) before validation
) -> OpenAIChatCompletionChunk:
for choice in payload.get("choices") or []: # mutable-ok: only iterated, never mutated
if isinstance(choice, dict) and not choice.get("logprobs"):
choice.pop("logprobs", None)
choices: Final = payload.get("choices")
if isinstance(choices, list):
for choice in choices:
if isinstance(choice, dict) and not choice.get("logprobs"):
choice.pop("logprobs", None) # mutable-ok: pops the logprobs key in-place before model_validate
chunk = OpenAIChatCompletionChunk.model_validate(payload)
if chunk.usage is not None:
chunk.usage = Usage(**chunk.usage.model_dump())
chunk.usage = Usage.model_validate(chunk.usage.model_dump())
return chunk
@staticmethod
@ -68,25 +70,24 @@ class _StreamParser:
if not orc:
return None
return _StreamParser._validate_chunk(
{
"id": orc.get("id") or evt.get("request_id") or "stream-chunk",
"object": orc.get("object") or "chat.completion.chunk",
"created": orc.get("created") or evt.get("created") or _now_ts(),
"model": orc.get("model") or "unknown",
"choices": [
{
"index": c.get("index", 0),
"delta": c.get("delta") or {},
"finish_reason": c.get("finish_reason"),
}
for c in (orc.get("choices") or [])
],
}
)
payload: Final[dict[str, object]] = {
"id": orc.get("id") or evt.get("request_id") or "stream-chunk",
"object": orc.get("object") or "chat.completion.chunk",
"created": orc.get("created") or evt.get("created") or _now_ts(),
"model": orc.get("model") or "unknown",
"choices": [
{
"index": c.get("index", 0),
"delta": c.get("delta") or {},
"finish_reason": c.get("finish_reason"),
}
for c in (orc.get("choices") or [])
],
}
return _StreamParser._validate_chunk(payload)
@staticmethod
def to_openai_chunk(event_obj: dict) -> OpenAIChatCompletionChunk | None:
def to_openai_chunk(event_obj: dict[str, object]) -> OpenAIChatCompletionChunk | None:
"""
Accepts:
- {"final_result": <openai-style CHUNK>} (IMPORTANT: this is just another chunk, NOT terminal)
@ -102,7 +103,10 @@ class _StreamParser:
# FINAL RESULT IS *NOT* TERMINAL: treat it as the next chunk
if "final_result" in event_obj:
fr: Final = event_obj["final_result"] or {}
final_result: Final = event_obj["final_result"]
if not isinstance(final_result, dict):
return None
fr: Final[dict[str, object]] = final_result
# ensure it looks like an OpenAI chunk
if "object" not in fr:
fr["object"] = "chat.completion.chunk"

View file

@ -1,3 +1,4 @@
import json
import httpx
from unittest.mock import patch, PropertyMock
@ -267,7 +268,7 @@ def test_validate_chunk_without_usage_keeps_none():
assert chunk.usage is None
def test_validated_usage_survives_nested_model_dump():
def test_validated_usage_survives_nested_model_dump_json():
from litellm.llms.sap.chat.handler import _StreamParser
from litellm.types.utils import ModelResponseStream
@ -276,7 +277,7 @@ def test_validated_usage_survives_nested_model_dump():
model_response = ModelResponseStream()
setattr(model_response, "usage", chunk.usage)
dumped = model_response.model_dump()
dumped = json.loads(model_response.model_dump_json())
assert dumped["usage"]["total_tokens"] == 62528