mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
core/lmproto: stop SSE stream at data: [DONE]
Previously DONE was skipped with continue, so keep-alive connections could hang forever after the model finished. Also fail fast on HTTP error statuses for stream and non-stream completions.
This commit is contained in:
@@ -1,7 +1,7 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import inspect
|
import inspect
|
||||||
from typing import TYPE_CHECKING, Literal, overload
|
from typing import TYPE_CHECKING, Literal, cast, overload
|
||||||
|
|
||||||
import msgspec
|
import msgspec
|
||||||
import niquests
|
import niquests
|
||||||
@@ -52,13 +52,44 @@ class BaseOpenAIClient:
|
|||||||
self.decoder = msgspec.json.Decoder(ChatCompletionResponse)
|
self.decoder = msgspec.json.Decoder(ChatCompletionResponse)
|
||||||
self.chunk_decoder = msgspec.json.Decoder(ChatCompletionChunk)
|
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]:
|
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:
|
try:
|
||||||
async for line in resp.iter_lines():
|
await self._ensure_ok(resp)
|
||||||
if not line or line == b"data: [DONE]":
|
async for raw in resp.iter_lines():
|
||||||
|
if not raw:
|
||||||
continue
|
continue
|
||||||
if line.startswith(b"data: "):
|
line = bytes(raw).strip()
|
||||||
yield self.chunk_decoder.decode(line[6:])
|
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:
|
finally:
|
||||||
await _close_async_response(resp)
|
await _close_async_response(resp)
|
||||||
|
|
||||||
@@ -97,6 +128,7 @@ class OpenAIClient(BaseOpenAIClient):
|
|||||||
headers={"Content-Type": "application/json"},
|
headers={"Content-Type": "application/json"},
|
||||||
stream=False,
|
stream=False,
|
||||||
)
|
)
|
||||||
|
await self._ensure_ok(resp)
|
||||||
assert resp.content is not None
|
assert resp.content is not None
|
||||||
return self.decoder.decode(resp.content)
|
return self.decoder.decode(resp.content)
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
Reference in New Issue
Block a user