core/lmproto+agent: stream via chat_completions(stream=True)

Keep library async-def→async-iterator shape; remove raw-lines path;
typed StreamToolCallDelta merge; scripted clients emit real chunks.
This commit is contained in:
2026-07-14 22:40:24 +08:00
parent 4548828354
commit bd6d98ceda
6 changed files with 173 additions and 184 deletions
+3 -1
View File
@@ -15,7 +15,9 @@ if TYPE_CHECKING:
class ChatClient(Protocol):
"""Structural protocol for OpenAI-compatible chat completion clients.
``chat_completions_raw_lines`` is optional (checked dynamically in the loop).
When ``stream=True``, implementations may return an async iterator from an
async method (``await client.chat_completions(..., stream=True)`` then
``async for chunk in stream``). That library shape is intentional.
"""
@overload
+22 -34
View File
@@ -4,11 +4,13 @@ from typing import TYPE_CHECKING, cast
from msgspec import UNSET
from plyngent.lmproto.openai_compatible.client import merge_stream_tool_calls
from plyngent.lmproto.openai_compatible.model import (
AnyAssistantToolCall,
AssistantChatMessage,
AssistantFunctionToolCall,
ChatCompletionsParam,
StreamToolCallDelta,
ToolChatMessage,
)
from plyngent.typedef import Unset # noqa: TC001
@@ -79,46 +81,34 @@ async def _stream_and_build_assistant(
client: ChatClient,
param: ChatCompletionsParam,
) -> tuple[AssistantChatMessage, list[TextDeltaEvent]]:
"""Stream a round, return the assistant message and any text deltas.
"""Stream one completion via the normal client API.
Requires ``client.chat_completions_raw_lines`` (async generator of raw SSE
payload bytes). Falls back to non-streaming if unavailable.
Pattern: ``stream = await client.chat_completions(..., stream=True)`` then
``async for chunk in stream``. That return type (async function → async
iterator) is accepted as the library interface.
"""
raw_lines_attr = getattr(client, "chat_completions_raw_lines", None)
if raw_lines_attr is None:
return await _non_stream_assistant(client, param)
# Async generators are not awaitable — call to get the iterator.
stream_iter = raw_lines_attr(param)
raw_lines: list[bytes] = []
stream = await client.chat_completions(param, stream=True)
content_parts: list[str] = []
text_events: list[TextDeltaEvent] = []
stream_decoder = getattr(client, "stream_decoder", None)
tool_deltas: list[StreamToolCallDelta] = []
async for raw in stream_iter:
raw_lines.append(raw)
if stream_decoder is None:
continue
try:
chunk = stream_decoder.decode(raw)
except Exception: # noqa: BLE001 — skip malformed SSE payloads
continue
async for chunk in stream:
if not chunk.choices:
continue
delta_text = chunk.choices[0].delta.content
if isinstance(delta_text, str) and delta_text:
content_parts.append(delta_text)
text_events.append(TextDeltaEvent(content=delta_text))
choice = chunk.choices[0]
delta = choice.delta
if isinstance(delta.content, str) and delta.content:
content_parts.append(delta.content)
text_events.append(TextDeltaEvent(content=delta.content))
if delta.tool_calls is not UNSET and delta.tool_calls:
tool_deltas.extend(delta.tool_calls)
full_content = "".join(content_parts)
tool_calls: list[AnyAssistantToolCall] | Unset = UNSET
# Reconstruct tool calls from raw lines whenever present (do not rely solely
# on finish_reason — some providers omit it on intermediate chunks).
from plyngent.lmproto.openai_compatible.client import merge_stream_tool_calls
calls = merge_stream_tool_calls(raw_lines)
if calls:
tool_calls = cast("list[AnyAssistantToolCall]", calls)
if tool_deltas:
calls = merge_stream_tool_calls(tool_deltas)
if calls:
tool_calls = cast("list[AnyAssistantToolCall]", calls)
assistant = AssistantChatMessage(
content=full_content or None,
@@ -140,10 +130,8 @@ async def run_chat_loop(
) -> AsyncIterator[AgentEvent]:
"""Multi-round chat/tool loop; mutates ``messages`` in place and yields events.
When ``stream=True``, text tokens yield as they arrive; tool calls are
reconstructed from raw SSE payloads when the client supports raw streaming.
Continues until the model returns no tool calls, or ``max_rounds`` is hit.
When ``stream=True``, uses ``chat_completions(..., stream=True)`` and yields
text deltas as chunks arrive; tool calls are merged from stream deltas.
"""
tool_items: Sequence[AnyToolItem] | None = None
if tools is not None and len(tools) > 0:
+23 -106
View File
@@ -1,9 +1,10 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Any, Literal, cast, overload
from typing import TYPE_CHECKING, Literal, overload
import msgspec
import niquests
from msgspec import UNSET
from niquests.auth import BearerTokenAuth
from .config import OpenAIConfig # noqa: TC001
@@ -13,11 +14,11 @@ from .model import (
ChatCompletionChunk,
ChatCompletionResponse,
ChatCompletionsParam,
StreamChatCompletionChunk,
StreamToolCallDelta,
)
if TYPE_CHECKING:
from collections.abc import AsyncIterator, Mapping
from collections.abc import AsyncIterator
from niquests.async_session import AsyncSession
from niquests.models import AsyncResponse
@@ -28,7 +29,6 @@ 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(
@@ -38,38 +38,15 @@ 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()
async for line in lines:
if not line or line == b"data: [DONE]":
continue
if line.startswith(b"data: "):
yield self.chunk_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).
The HTTP response is closed when the iterator finishes or is closed early
(e.g. task cancellation during streaming).
"""
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,
)
try:
async for line in resp.iter_lines():
if not line or line == b"data: [DONE]":
continue
if line.startswith(b"data: "):
yield line[6:]
yield self.chunk_decoder.decode(line[6:])
finally:
# Best-effort abort of the in-flight stream (level-2 cancel toward HTTP).
aclose = getattr(resp, "aclose", None)
if callable(aclose):
await aclose() # pyright: ignore[reportGeneralTypeIssues]
@@ -96,6 +73,7 @@ class OpenAIClient(BaseOpenAIClient):
async def chat_completions(
self, param: ChatCompletionsParam, *, stream: bool = False
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
# Library pattern: async def returns AsyncIterator when stream=True.
param = msgspec.structs.replace(param, stream=stream)
data = self.encoder.encode(param)
if stream:
@@ -116,92 +94,31 @@ class OpenAIClient(BaseOpenAIClient):
return self.decoder.decode(resp.content)
def _as_mapping(value: object) -> Mapping[str, Any] | None:
if isinstance(value, dict):
return cast("Mapping[str, Any]", value)
return None
def _as_list(value: object) -> list[object] | None:
if isinstance(value, list):
return cast("list[object]", value)
return None
def _dstr(d: Mapping[str, Any], key: str, default: str = "") -> str:
v = d.get(key)
return v if isinstance(v, str) else default
def _merge_tool_entry(merge: dict[int, dict[str, str]], tc: Mapping[str, Any]) -> None:
idx_raw = tc.get("index", 0)
idx = idx_raw if isinstance(idx_raw, int) else 0
if idx not in merge:
merge[idx] = {"id": "", "name": "", "arguments": ""}
entry = merge[idx]
raw_id = tc.get("id")
if isinstance(raw_id, str) and raw_id:
entry["id"] = raw_id
fn = _as_mapping(tc.get("function", {}))
if fn is None:
return
name = _dstr(fn, "name")
if name:
entry["name"] = name
args = _dstr(fn, "arguments")
if args:
entry["arguments"] = entry["arguments"] + args
def _stream_choices(raw_line: bytes) -> list[Mapping[str, Any]]:
"""Extract ``choices`` list from a raw SSE payload byte string, or empty list."""
import json as _json
try:
raw_data: object = _json.loads(raw_line)
except _json.JSONDecodeError:
return []
data = _as_mapping(raw_data)
if data is None:
return []
choices = _as_list(data.get("choices", []))
if choices is None:
return []
result: list[Mapping[str, Any]] = []
for item in choices:
mapping = _as_mapping(item)
if mapping is not None:
result.append(mapping)
return result
def merge_stream_tool_calls(raw_lines: list[bytes]) -> list[AssistantFunctionToolCall]:
"""Accumulate streaming tool-call deltas by index across raw SSE payload bytes."""
def merge_stream_tool_calls(deltas: list[StreamToolCallDelta]) -> list[AssistantFunctionToolCall]:
"""Accumulate streaming tool-call deltas by index into complete tool calls."""
merge: dict[int, dict[str, str]] = {}
for raw in raw_lines:
for choice in _stream_choices(raw):
delta = _as_mapping(choice.get("delta", {}))
if delta is None:
continue
calls = _as_list(delta.get("tool_calls", []))
if calls is None:
continue
for tc in calls:
mapping = _as_mapping(tc)
if mapping is not None:
_merge_tool_entry(merge, mapping)
for delta in deltas:
if delta.index not in merge:
merge[delta.index] = {"id": "", "name": "", "arguments": ""}
entry = merge[delta.index]
if isinstance(delta.id, str) and delta.id:
entry["id"] = delta.id
if delta.function is not UNSET:
fn = delta.function
if isinstance(fn.name, str) and fn.name:
entry["name"] = fn.name
if isinstance(fn.arguments, str) and fn.arguments:
entry["arguments"] = entry["arguments"] + fn.arguments
result: list[AssistantFunctionToolCall] = []
for idx in sorted(merge):
entry = merge[idx]
entry_id = entry["id"]
fn_name = entry["name"]
if not entry_id or not fn_name:
if not entry["id"] or not entry["name"]:
continue
result.append(
AssistantFunctionToolCall(
id=entry_id,
function=AssistantFunctionTool(name=fn_name, arguments=entry["arguments"]),
id=entry["id"],
function=AssistantFunctionTool(name=entry["name"], arguments=entry["arguments"]),
)
)
return result
+14 -23
View File
@@ -219,10 +219,23 @@ class ChatCompletionResponse(Struct):
service_tier: ServiceTier | Unset = UNSET
# Streaming deltas may be partial (id/name on first chunk, arguments later).
class StreamFunctionDelta(Struct):
name: str | Unset = UNSET
arguments: str | Unset = UNSET
class StreamToolCallDelta(Struct):
index: int = 0
id: str | Unset = UNSET
type: Literal["function"] | Unset = UNSET
function: StreamFunctionDelta | Unset = UNSET
class DeltaMessage(Struct):
role: RoleAssistant | Unset = UNSET
content: str | Unset = UNSET
tool_calls: list[AnyAssistantToolCall] | Unset = UNSET
tool_calls: list[StreamToolCallDelta] | Unset = UNSET
class ChunkChoice(Struct):
@@ -239,25 +252,3 @@ 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