diff --git a/src/plyngent/lmproto/openai_compatible/client.py b/src/plyngent/lmproto/openai_compatible/client.py index 0004774..fce7bd1 100644 --- a/src/plyngent/lmproto/openai_compatible/client.py +++ b/src/plyngent/lmproto/openai_compatible/client.py @@ -25,6 +25,10 @@ if TYPE_CHECKING: from niquests.models import AsyncResponse +_HTTP_ERROR = 400 +_ERROR_BODY_PREVIEW = 500 + + async def _close_async_response(resp: AsyncResponse) -> None: """Close a streaming response; support sync or async close().""" for name in ("aclose", "close"): @@ -37,6 +41,43 @@ async def _close_async_response(resp: AsyncResponse) -> None: return +def sse_data_payload(line: bytes) -> bytes | None | Literal[False]: + """Parse one SSE line. + + Returns: + ``bytes`` payload after ``data: ``, + ``None`` to skip (comment/empty/non-data), + ``False`` when the stream is finished (``[DONE]``). + """ + stripped = line.strip() + if not stripped: + return None + if stripped in {b"data: [DONE]", b"[DONE]"}: + return False + if not stripped.startswith(b"data: "): + return None + payload = stripped[6:].strip() + if payload == b"[DONE]": + return False + if not payload: + return None + return payload + + +def http_error_message(status_code: int, body: bytes | str | None = None) -> str | None: + """Return an error string for HTTP ``status_code`` >= 400, else ``None``.""" + if status_code < _HTTP_ERROR: + return None + msg = f"chat completions HTTP {status_code}" + if body is None: + return msg + if isinstance(body, (bytes, bytearray)): + text = bytes(body[:_ERROR_BODY_PREVIEW]).decode(errors="replace") + else: + text = str(body)[:_ERROR_BODY_PREVIEW] + return f"{msg}: {text}" if text else msg + + class BaseOpenAIClient: session: AsyncSession encoder: msgspec.json.Encoder @@ -54,22 +95,20 @@ class BaseOpenAIClient: async def _ensure_ok(self, resp: object) -> None: """Raise if the HTTP response is an error (stream or non-stream).""" - http_error = 400 status = getattr(resp, "status_code", None) - if not isinstance(status, int) or status < http_error: + if not isinstance(status, int): return - body = "" content = getattr(resp, "content", None) - if content is not None: - if isinstance(content, (bytes, bytearray)): - body = bytes(content[:500]).decode(errors="replace") - else: - body = str(content)[:500] + body: bytes | str | None = None + if isinstance(content, bytes | bytearray): + body = bytes(content) + elif isinstance(content, str): + body = content + msg = http_error_message(status, body) + if msg is None: + return if hasattr(resp, "close") or hasattr(resp, "aclose"): await _close_async_response(cast("AsyncResponse", resp)) - msg = f"chat completions HTTP {status}" - if body: - msg = f"{msg}: {body}" raise RuntimeError(msg) async def _parse_sse(self, resp: AsyncResponse) -> AsyncIterator[ChatCompletionChunk]: @@ -79,17 +118,12 @@ class BaseOpenAIClient: async for raw in resp.iter_lines(): if not raw: continue - line = bytes(raw).strip() - if line in {b"data: [DONE]", b"[DONE]"}: - break - if not line.startswith(b"data: "): + parsed = sse_data_payload(bytes(raw)) + if parsed is None: continue - payload = line[6:].strip() - if payload == b"[DONE]": + if parsed is False: break - if not payload: - continue - yield self.chunk_decoder.decode(payload) + yield self.chunk_decoder.decode(parsed) finally: await _close_async_response(resp) diff --git a/tests/test_lmproto/test_sse_parse.py b/tests/test_lmproto/test_sse_parse.py index c6ecb17..878d014 100644 --- a/tests/test_lmproto/test_sse_parse.py +++ b/tests/test_lmproto/test_sse_parse.py @@ -1,57 +1,71 @@ from __future__ import annotations -from typing import TYPE_CHECKING, cast - -from plyngent.lmproto.openai_compatible.client import BaseOpenAIClient -from plyngent.lmproto.openai_compatible.config import OpenAIConfig - -if TYPE_CHECKING: - from collections.abc import AsyncIterator - - from niquests.models import AsyncResponse +from plyngent.lmproto.openai_compatible.client import http_error_message, sse_data_payload +from plyngent.lmproto.openai_compatible.model import ChatCompletionChunk -class _FakeResp: - status_code: int - closed: bool - content: bytes - _lines: list[bytes] - - def __init__(self, lines: list[bytes], *, status_code: int = 200) -> None: - self.status_code = status_code - self.closed = False - self._lines = lines - self.content = b"" - - async def iter_lines(self) -> AsyncIterator[bytes]: - for line in self._lines: - yield line - - def close(self) -> None: - self.closed = True +def test_sse_data_payload_done_variants() -> None: + assert sse_data_payload(b"data: [DONE]") is False + assert sse_data_payload(b"[DONE]") is False + assert sse_data_payload(b"data: [DONE]\n") is False + assert sse_data_payload(b"data: [DONE] ") is False -async def test_parse_sse_stops_at_done() -> None: - client = BaseOpenAIClient(OpenAIConfig(access_key_or_token="t", base_url="http://x")) - # Minimal OpenAI-style chunk payload - chunk = ( +def test_sse_data_payload_skip() -> None: + assert sse_data_payload(b"") is None + assert sse_data_payload(b": comment") is None + assert sse_data_payload(b"event: message") is None + assert sse_data_payload(b"data: ") is None + + +def test_sse_data_payload_json() -> None: + line = ( b'data: {"id":"1","object":"chat.completion.chunk","created":0,' b'"model":"m","choices":[{"index":0,"delta":{"content":"hi"},' b'"finish_reason":null}]}' ) - resp = _FakeResp([chunk, b"", b"data: [DONE]", b"data: should-not-parse"]) - out = [c async for c in client._parse_sse(cast("AsyncResponse", cast("object", resp)))] - assert len(out) == 1 - assert out[0].choices[0].delta.content == "hi" - assert resp.closed is True + payload = sse_data_payload(line) + assert isinstance(payload, bytes) + chunk = msgspec_decode_chunk(payload) + assert chunk.choices[0].delta.content == "hi" -async def test_parse_sse_http_error() -> None: - import pytest +def msgspec_decode_chunk(payload: bytes) -> ChatCompletionChunk: + import msgspec - client = BaseOpenAIClient(OpenAIConfig(access_key_or_token="t", base_url="http://x")) - resp = _FakeResp([], status_code=500) - resp.content = b'{"error":"boom"}' - with pytest.raises(RuntimeError, match="500"): - _ = [c async for c in client._parse_sse(cast("AsyncResponse", cast("object", resp)))] - assert resp.closed is True + return msgspec.json.decode(payload, type=ChatCompletionChunk) + + +def test_sse_stream_stops_at_done() -> None: + """Simulate a line iterator: only pre-DONE data lines are decoded.""" + lines = [ + b'data: {"id":"1","object":"chat.completion.chunk","created":0,' + b'"model":"m","choices":[{"index":0,"delta":{"content":"a"},' + b'"finish_reason":null}]}', + b"", + b"data: [DONE]", + b'data: {"id":"1","choices":[{"delta":{"content":"never"}}]}', + ] + payloads: list[bytes] = [] + for line in lines: + parsed = sse_data_payload(line) + if parsed is None: + continue + if parsed is False: + break + payloads.append(parsed) + assert len(payloads) == 1 + assert msgspec_decode_chunk(payloads[0]).choices[0].delta.content == "a" + + +def test_http_error_message() -> None: + assert http_error_message(200) is None + assert http_error_message(399) is None + err = http_error_message(500, b'{"error":"boom"}') + assert err is not None + assert "500" in err + assert "boom" in err + err503 = http_error_message(503, "unavailable") + assert err503 is not None + assert "503" in err503 + assert "unavailable" in err503