diff --git a/src/inference_endpoint/endpoint_client/adapter_protocol.py b/src/inference_endpoint/endpoint_client/adapter_protocol.py index a96faaa0a..1f58c59a3 100644 --- a/src/inference_endpoint/endpoint_client/adapter_protocol.py +++ b/src/inference_endpoint/endpoint_client/adapter_protocol.py @@ -127,12 +127,11 @@ def parse_sse_chunk(cls, buffer: bytes, end_pos: int) -> list[Any]: """ json_docs = cls.SSE_DATA_PATTERN.findall(buffer[:end_pos]) parsed: list[Any] = [] - # Note: if one frame is malformed, remaining frames are skipped - try: - for json_doc in json_docs: + for json_doc in json_docs: + try: content = cls.decode_sse_message(json_doc) if content is not None: parsed.append(content) - except (msgspec.DecodeError, msgspec.ValidationError) as exc: - logger.warning("skipping malformed SSE batch (%s)", type(exc).__name__) + except (msgspec.DecodeError, msgspec.ValidationError) as exc: + logger.warning("skipping malformed SSE frame (%s)", type(exc).__name__) return parsed diff --git a/src/inference_endpoint/sglang/adapter.py b/src/inference_endpoint/sglang/adapter.py index a6614fa35..c19e3d2a4 100644 --- a/src/inference_endpoint/sglang/adapter.py +++ b/src/inference_endpoint/sglang/adapter.py @@ -122,7 +122,7 @@ def decode_sse_message(cls, json_bytes: bytes) -> SGLangSSEDelta: # the output_tokens is the delta from the previous chunk token_delta = resp.output_ids - has_retractions = resp.meta_info.total_retractions > 0 + has_retractions = resp.meta_info.num_retractions > 0 return SGLangSSEDelta( text=total_text, token_delta=token_delta, diff --git a/src/inference_endpoint/sglang/types.py b/src/inference_endpoint/sglang/types.py index 9e75b3a4d..89f71eed8 100644 --- a/src/inference_endpoint/sglang/types.py +++ b/src/inference_endpoint/sglang/types.py @@ -52,13 +52,13 @@ class SGLangGenerateRequest( class MetaInfo(msgspec.Struct, frozen=True, kw_only=True, omit_defaults=True, gc=False): # type: ignore[call-arg] id: str - finish_reason: dict[str, Any] + finish_reason: dict[str, Any] | None prompt_tokens: int weight_version: str - total_retractions: int + num_retractions: int completion_tokens: int cached_tokens: int - e2e_latency: float + e2e_latency: float | None = None class SGLangGenerateResponse( diff --git a/tests/performance/sglang/test_sglang_adapter.py b/tests/performance/sglang/test_sglang_adapter.py index f614cc0d7..3e0fa90e8 100644 --- a/tests/performance/sglang/test_sglang_adapter.py +++ b/tests/performance/sglang/test_sglang_adapter.py @@ -64,7 +64,7 @@ def make_response_bytes(text: str, n_tokens: int) -> bytes: "finish_reason": {"type": "stop"}, "prompt_tokens": 10, "weight_version": "v1", - "total_retractions": 0, + "num_retractions": 0, "completion_tokens": n_tokens, "cached_tokens": 0, "e2e_latency": 0.1, diff --git a/tests/unit/endpoint_client/test_adapter_protocol.py b/tests/unit/endpoint_client/test_adapter_protocol.py index 591dc4236..b348b93ff 100644 --- a/tests/unit/endpoint_client/test_adapter_protocol.py +++ b/tests/unit/endpoint_client/test_adapter_protocol.py @@ -22,10 +22,11 @@ from inference_endpoint.config.schema import ModelParams from inference_endpoint.dataset_manager.transforms import Transform from inference_endpoint.endpoint_client.adapter_protocol import HttpRequestAdapter +from inference_endpoint.sglang.adapter import SGLangGenerateAdapter class _SimpleAdapter(HttpRequestAdapter): - """Minimal adapter: raises DecodeError for b'BAD', returns {"x": 1} otherwise.""" + """Decode JSON objects with integer values.""" @classmethod def dataset_transforms(cls, model_params: ModelParams) -> list[Transform]: @@ -41,15 +42,47 @@ def decode_response(cls, response_bytes: bytes, query_id: str) -> Any: @classmethod def decode_sse_message(cls, json_bytes: bytes) -> Any: - if json_bytes.strip() == b"BAD": - raise msgspec.DecodeError("bad frame") - return {"x": 1} + return msgspec.json.decode(json_bytes, type=dict[str, int]) @pytest.mark.unit -def test_parse_sse_chunk_skips_bad_frame_and_keeps_valid(): +@pytest.mark.parametrize("bad_frame", [b'{"ok":}', b'{"ok":"bad"}']) +def test_parse_sse_chunk_skips_bad_frame_and_keeps_valid(bad_frame): """A malformed SSE frame is skipped; surrounding valid frames are preserved.""" - buffer = b'data: {"ok":1}\n\ndata: BAD\n\ndata: {"ok":1}\n\n' + buffer = b'data: {"ok":1}\n\ndata: ' + bad_frame + b'\n\ndata: {"ok":2}\n\n' end_pos = len(buffer) result = _SimpleAdapter.parse_sse_chunk(buffer, end_pos) - assert result == [{"x": 1}, {"x": 1}] + assert result == [{"ok": 1}, {"ok": 2}] + + +@pytest.mark.unit +@pytest.mark.parametrize("finish_reason", [None, {"type": "stop"}]) +@pytest.mark.parametrize("latency", [{}, {"e2e_latency": None}, {"e2e_latency": 0.1}]) +@pytest.mark.parametrize("retractions", [0, 2]) +def test_sglang_response_metadata(finish_reason, latency, retractions): + response = msgspec.json.encode( + { + "text": "Hello", + "output_ids": [10], + "meta_info": { + "id": "sample", + "finish_reason": finish_reason, + "prompt_tokens": 3, + "weight_version": "v1", + "num_retractions": retractions, + "completion_tokens": 1, + "cached_tokens": 0, + **latency, + }, + } + ) + result = SGLangGenerateAdapter.decode_response(response, "sample") + assert result.get_response_output_string() == "Hello" + assert result.metadata == {"token_ids": [10], "n_tokens": 1} + frame = b"data: " + response + b"\n\n" + chunks = SGLangGenerateAdapter.parse_sse_chunk(frame, len(frame)) + assert len(chunks) == 1 + assert chunks[0].text == "Hello" + assert chunks[0].token_delta == [10] + assert chunks[0].total_completion_tokens == 1 + assert chunks[0].has_retractions is (retractions > 0)