mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
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:
parent
264d05b1fa
commit
a215b78d99
2 changed files with 30 additions and 25 deletions
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue