From 665c3f435d50dd2da6eab365b7ad6106f9990d32 Mon Sep 17 00:00:00 2001 From: worldmozara Date: Tue, 14 Jul 2026 22:11:25 +0800 Subject: [PATCH] core/lmproto: add stream decoder and raw SSE tool-call merger StreamChatCompletionChunk tolerant of partial tool calls; merge_stream_tool_calls accumulates by index from raw bytes. --- .../lmproto/openai_compatible/client.py | 110 +++++++++++++++++- .../lmproto/openai_compatible/model.py | 22 ++++ 2 files changed, 131 insertions(+), 1 deletion(-) diff --git a/src/plyngent/lmproto/openai_compatible/client.py b/src/plyngent/lmproto/openai_compatible/client.py index 62ee035..054152f 100644 --- a/src/plyngent/lmproto/openai_compatible/client.py +++ b/src/plyngent/lmproto/openai_compatible/client.py @@ -5,7 +5,14 @@ import niquests from niquests.auth import BearerTokenAuth from .config import OpenAIConfig # noqa: TC001 -from .model import ChatCompletionChunk, ChatCompletionResponse, ChatCompletionsParam +from .model import ( + AssistantFunctionTool, + AssistantFunctionToolCall, + ChatCompletionChunk, + ChatCompletionResponse, + ChatCompletionsParam, + StreamChatCompletionChunk, +) if TYPE_CHECKING: from collections.abc import AsyncIterator @@ -19,6 +26,7 @@ class BaseOpenAIClient: encoder: msgspec.json.Encoder decoder: msgspec.json.Decoder[ChatCompletionResponse] chunk_decoder: msgspec.json.Decoder[ChatCompletionChunk] + stream_decoder: msgspec.json.Decoder[StreamChatCompletionChunk] def __init__(self, config: OpenAIConfig) -> None: self.session = niquests.AsyncSession( @@ -28,6 +36,7 @@ class BaseOpenAIClient: self.encoder = msgspec.json.Encoder() self.decoder = msgspec.json.Decoder(ChatCompletionResponse) self.chunk_decoder = msgspec.json.Decoder(ChatCompletionChunk) + self.stream_decoder = msgspec.json.Decoder(StreamChatCompletionChunk) async def _parse_sse(self, resp: AsyncResponse) -> AsyncIterator[ChatCompletionChunk]: lines = resp.iter_lines() @@ -37,6 +46,33 @@ class BaseOpenAIClient: if line.startswith(b"data: "): yield self.chunk_decoder.decode(line[6:]) + async def _parse_sse_stream(self, resp: AsyncResponse) -> AsyncIterator[StreamChatCompletionChunk]: + """Parse SSE using the tolerant ::class:`StreamChatCompletionChunk` decoder.""" + lines = resp.iter_lines() + async for line in lines: + if not line or line == b"data: [DONE]": + continue + if line.startswith(b"data: "): + yield self.stream_decoder.decode(line[6:]) + + async def chat_completions_raw_lines(self, param: ChatCompletionsParam) -> AsyncIterator[bytes]: + """Yield raw SSE ``data: `` payload bytes for manual accumulation. + + Each yielded bytes object is a complete JSON object (no ``data: `` prefix). + """ + data = self.encoder.encode(msgspec.structs.replace(param, stream=True)) + resp = await self.session.post( + "/chat/completions", + data=data, + headers={"Content-Type": "application/json"}, + stream=True, + ) + async for line in resp.iter_lines(): + if not line or line == b"data: [DONE]": + continue + if line.startswith(b"data: "): + yield line[6:] + class OpenAIClient(BaseOpenAIClient): def __init__(self, config: OpenAIConfig) -> None: @@ -73,3 +109,75 @@ class OpenAIClient(BaseOpenAIClient): ) assert resp.content is not None return self.decoder.decode(resp.content) + + +def _merge_tool_entry( + merge: dict[int, dict[str, object]], + tc: dict[str, object], +) -> None: + idx = tc.get("index", 0) + if not isinstance(idx, int): + idx = 0 + if idx not in merge: + merge[idx] = {"id": "", "function": {"name": "", "arguments": ""}} + entry = merge[idx] + if isinstance(tc.get("id"), str) and tc["id"]: + entry["id"] = tc["id"] + fn_raw = tc.get("function", {}) + if isinstance(fn_raw, dict): + entry_fn = entry["function"] + if isinstance(entry_fn, dict): + if isinstance(fn_raw.get("name"), str) and fn_raw["name"]: + entry_fn["name"] = fn_raw["name"] + if isinstance(fn_raw.get("arguments"), str) and fn_raw["arguments"]: + entry_fn["arguments"] = entry_fn.get("arguments", "") + fn_raw["arguments"] + + +def _stream_choices(raw_line: bytes) -> list[dict[str, object]]: + """Extract ``choices`` list from a raw SSE payload byte string, or empty list.""" + import json as _json + + try: + data: object = _json.loads(raw_line) + except _json.JSONDecodeError: + return [] + if not isinstance(data, dict): + return [] + choices = data.get("choices", []) + if not isinstance(choices, list): + return [] + return [c for c in choices if isinstance(c, dict)] + + +def merge_stream_tool_calls(raw_lines: list[bytes]) -> list[AssistantFunctionToolCall]: + """Accumulate streaming tool-call deltas by index across raw SSE payload bytes.""" + merge: dict[int, dict[str, object]] = {} + for raw in raw_lines: + for choice in _stream_choices(raw): + delta = choice.get("delta", {}) + if not isinstance(delta, dict): + continue + for tc in delta.get("tool_calls", []): + if isinstance(tc, dict): + _merge_tool_entry(merge, tc) + + result: list[AssistantFunctionToolCall] = [] + for idx in sorted(merge): + entry = merge[idx] + if not isinstance(entry, dict): + continue + entry_id: object = entry.get("id", "") + entry_fn: object = entry.get("function", {}) + if not isinstance(entry_fn, dict): + continue + fn_name: object = entry_fn.get("name", "") + fn_args: object = entry_fn.get("arguments", "") + if not entry_id or not fn_name: + continue + result.append( + AssistantFunctionToolCall( + id=str(entry_id), + function=AssistantFunctionTool(name=str(fn_name), arguments=str(fn_args or "")), + ) + ) + return result diff --git a/src/plyngent/lmproto/openai_compatible/model.py b/src/plyngent/lmproto/openai_compatible/model.py index e14a7c0..a767968 100644 --- a/src/plyngent/lmproto/openai_compatible/model.py +++ b/src/plyngent/lmproto/openai_compatible/model.py @@ -239,3 +239,25 @@ class ChatCompletionChunk(Struct): model: str choices: list[ChunkChoice] usage: dict[str, Any] | Unset = UNSET + + +# -- Streaming-tolerant models (ignore partial tool calls in stream) -- + +class StreamDelta(Struct): + """Delta that accepts any extra fields (partial tool calls), ignoring them.""" + content: str | Unset = UNSET + + +class StreamChunkChoice(Struct): + index: int + delta: StreamDelta + finish_reason: FinishReason | None | Unset = UNSET + + +class StreamChatCompletionChunk(Struct): + id: str + object: Literal["chat.completion.chunk"] + created: int + model: str + choices: list[StreamChunkChoice] + usage: dict[str, Any] | Unset = UNSET