diff --git a/src/plyngent/lmproto/openai_compatible/client.py b/src/plyngent/lmproto/openai_compatible/client.py index 29778e9..0004774 100644 --- a/src/plyngent/lmproto/openai_compatible/client.py +++ b/src/plyngent/lmproto/openai_compatible/client.py @@ -1,7 +1,7 @@ from __future__ import annotations import inspect -from typing import TYPE_CHECKING, Literal, overload +from typing import TYPE_CHECKING, Literal, cast, overload import msgspec import niquests @@ -52,13 +52,44 @@ class BaseOpenAIClient: self.decoder = msgspec.json.Decoder(ChatCompletionResponse) self.chunk_decoder = msgspec.json.Decoder(ChatCompletionChunk) + 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: + 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] + 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]: + """Yield SSE chunks; stop at ``data: [DONE]`` so we do not hang on keep-alive.""" try: - async for line in resp.iter_lines(): - if not line or line == b"data: [DONE]": + await self._ensure_ok(resp) + async for raw in resp.iter_lines(): + if not raw: continue - if line.startswith(b"data: "): - yield self.chunk_decoder.decode(line[6:]) + line = bytes(raw).strip() + if line in {b"data: [DONE]", b"[DONE]"}: + break + if not line.startswith(b"data: "): + continue + payload = line[6:].strip() + if payload == b"[DONE]": + break + if not payload: + continue + yield self.chunk_decoder.decode(payload) finally: await _close_async_response(resp) @@ -97,6 +128,7 @@ class OpenAIClient(BaseOpenAIClient): headers={"Content-Type": "application/json"}, stream=False, ) + await self._ensure_ok(resp) assert resp.content is not None return self.decoder.decode(resp.content) diff --git a/tests/test_lmproto/test_sse_parse.py b/tests/test_lmproto/test_sse_parse.py new file mode 100644 index 0000000..c6ecb17 --- /dev/null +++ b/tests/test_lmproto/test_sse_parse.py @@ -0,0 +1,57 @@ +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 + + +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 + + +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 = ( + 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 + + +async def test_parse_sse_http_error() -> None: + import pytest + + 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