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): class ChatClient(Protocol):
"""Structural protocol for OpenAI-compatible chat completion clients. """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 @overload
+22 -34
View File
@@ -4,11 +4,13 @@ from typing import TYPE_CHECKING, cast
from msgspec import UNSET from msgspec import UNSET
from plyngent.lmproto.openai_compatible.client import merge_stream_tool_calls
from plyngent.lmproto.openai_compatible.model import ( from plyngent.lmproto.openai_compatible.model import (
AnyAssistantToolCall, AnyAssistantToolCall,
AssistantChatMessage, AssistantChatMessage,
AssistantFunctionToolCall, AssistantFunctionToolCall,
ChatCompletionsParam, ChatCompletionsParam,
StreamToolCallDelta,
ToolChatMessage, ToolChatMessage,
) )
from plyngent.typedef import Unset # noqa: TC001 from plyngent.typedef import Unset # noqa: TC001
@@ -79,46 +81,34 @@ async def _stream_and_build_assistant(
client: ChatClient, client: ChatClient,
param: ChatCompletionsParam, param: ChatCompletionsParam,
) -> tuple[AssistantChatMessage, list[TextDeltaEvent]]: ) -> 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 Pattern: ``stream = await client.chat_completions(..., stream=True)`` then
payload bytes). Falls back to non-streaming if unavailable. ``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) stream = await client.chat_completions(param, stream=True)
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] = []
content_parts: list[str] = [] content_parts: list[str] = []
text_events: list[TextDeltaEvent] = [] text_events: list[TextDeltaEvent] = []
stream_decoder = getattr(client, "stream_decoder", None) tool_deltas: list[StreamToolCallDelta] = []
async for raw in stream_iter: async for chunk in stream:
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
if not chunk.choices: if not chunk.choices:
continue continue
delta_text = chunk.choices[0].delta.content choice = chunk.choices[0]
if isinstance(delta_text, str) and delta_text: delta = choice.delta
content_parts.append(delta_text) if isinstance(delta.content, str) and delta.content:
text_events.append(TextDeltaEvent(content=delta_text)) 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) full_content = "".join(content_parts)
tool_calls: list[AnyAssistantToolCall] | Unset = UNSET tool_calls: list[AnyAssistantToolCall] | Unset = UNSET
# Reconstruct tool calls from raw lines whenever present (do not rely solely if tool_deltas:
# on finish_reason — some providers omit it on intermediate chunks). calls = merge_stream_tool_calls(tool_deltas)
from plyngent.lmproto.openai_compatible.client import merge_stream_tool_calls if calls:
tool_calls = cast("list[AnyAssistantToolCall]", calls)
calls = merge_stream_tool_calls(raw_lines)
if calls:
tool_calls = cast("list[AnyAssistantToolCall]", calls)
assistant = AssistantChatMessage( assistant = AssistantChatMessage(
content=full_content or None, content=full_content or None,
@@ -140,10 +130,8 @@ async def run_chat_loop(
) -> AsyncIterator[AgentEvent]: ) -> AsyncIterator[AgentEvent]:
"""Multi-round chat/tool loop; mutates ``messages`` in place and yields events. """Multi-round chat/tool loop; mutates ``messages`` in place and yields events.
When ``stream=True``, text tokens yield as they arrive; tool calls are When ``stream=True``, uses ``chat_completions(..., stream=True)`` and yields
reconstructed from raw SSE payloads when the client supports raw streaming. text deltas as chunks arrive; tool calls are merged from stream deltas.
Continues until the model returns no tool calls, or ``max_rounds`` is hit.
""" """
tool_items: Sequence[AnyToolItem] | None = None tool_items: Sequence[AnyToolItem] | None = None
if tools is not None and len(tools) > 0: if tools is not None and len(tools) > 0:
+23 -106
View File
@@ -1,9 +1,10 @@
from __future__ import annotations from __future__ import annotations
from typing import TYPE_CHECKING, Any, Literal, cast, overload from typing import TYPE_CHECKING, Literal, overload
import msgspec import msgspec
import niquests import niquests
from msgspec import UNSET
from niquests.auth import BearerTokenAuth from niquests.auth import BearerTokenAuth
from .config import OpenAIConfig # noqa: TC001 from .config import OpenAIConfig # noqa: TC001
@@ -13,11 +14,11 @@ from .model import (
ChatCompletionChunk, ChatCompletionChunk,
ChatCompletionResponse, ChatCompletionResponse,
ChatCompletionsParam, ChatCompletionsParam,
StreamChatCompletionChunk, StreamToolCallDelta,
) )
if TYPE_CHECKING: if TYPE_CHECKING:
from collections.abc import AsyncIterator, Mapping from collections.abc import AsyncIterator
from niquests.async_session import AsyncSession from niquests.async_session import AsyncSession
from niquests.models import AsyncResponse from niquests.models import AsyncResponse
@@ -28,7 +29,6 @@ class BaseOpenAIClient:
encoder: msgspec.json.Encoder encoder: msgspec.json.Encoder
decoder: msgspec.json.Decoder[ChatCompletionResponse] decoder: msgspec.json.Decoder[ChatCompletionResponse]
chunk_decoder: msgspec.json.Decoder[ChatCompletionChunk] chunk_decoder: msgspec.json.Decoder[ChatCompletionChunk]
stream_decoder: msgspec.json.Decoder[StreamChatCompletionChunk]
def __init__(self, config: OpenAIConfig) -> None: def __init__(self, config: OpenAIConfig) -> None:
self.session = niquests.AsyncSession( self.session = niquests.AsyncSession(
@@ -38,38 +38,15 @@ class BaseOpenAIClient:
self.encoder = msgspec.json.Encoder() self.encoder = msgspec.json.Encoder()
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)
self.stream_decoder = msgspec.json.Decoder(StreamChatCompletionChunk)
async def _parse_sse(self, resp: AsyncResponse) -> AsyncIterator[ChatCompletionChunk]: 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: try:
async for line in resp.iter_lines(): async for line in resp.iter_lines():
if not line or line == b"data: [DONE]": if not line or line == b"data: [DONE]":
continue continue
if line.startswith(b"data: "): if line.startswith(b"data: "):
yield line[6:] yield self.chunk_decoder.decode(line[6:])
finally: finally:
# Best-effort abort of the in-flight stream (level-2 cancel toward HTTP).
aclose = getattr(resp, "aclose", None) aclose = getattr(resp, "aclose", None)
if callable(aclose): if callable(aclose):
await aclose() # pyright: ignore[reportGeneralTypeIssues] await aclose() # pyright: ignore[reportGeneralTypeIssues]
@@ -96,6 +73,7 @@ class OpenAIClient(BaseOpenAIClient):
async def chat_completions( async def chat_completions(
self, param: ChatCompletionsParam, *, stream: bool = False self, param: ChatCompletionsParam, *, stream: bool = False
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]: ) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
# Library pattern: async def returns AsyncIterator when stream=True.
param = msgspec.structs.replace(param, stream=stream) param = msgspec.structs.replace(param, stream=stream)
data = self.encoder.encode(param) data = self.encoder.encode(param)
if stream: if stream:
@@ -116,92 +94,31 @@ class OpenAIClient(BaseOpenAIClient):
return self.decoder.decode(resp.content) return self.decoder.decode(resp.content)
def _as_mapping(value: object) -> Mapping[str, Any] | None: def merge_stream_tool_calls(deltas: list[StreamToolCallDelta]) -> list[AssistantFunctionToolCall]:
if isinstance(value, dict): """Accumulate streaming tool-call deltas by index into complete tool calls."""
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."""
merge: dict[int, dict[str, str]] = {} merge: dict[int, dict[str, str]] = {}
for raw in raw_lines: for delta in deltas:
for choice in _stream_choices(raw): if delta.index not in merge:
delta = _as_mapping(choice.get("delta", {})) merge[delta.index] = {"id": "", "name": "", "arguments": ""}
if delta is None: entry = merge[delta.index]
continue if isinstance(delta.id, str) and delta.id:
calls = _as_list(delta.get("tool_calls", [])) entry["id"] = delta.id
if calls is None: if delta.function is not UNSET:
continue fn = delta.function
for tc in calls: if isinstance(fn.name, str) and fn.name:
mapping = _as_mapping(tc) entry["name"] = fn.name
if mapping is not None: if isinstance(fn.arguments, str) and fn.arguments:
_merge_tool_entry(merge, mapping) entry["arguments"] = entry["arguments"] + fn.arguments
result: list[AssistantFunctionToolCall] = [] result: list[AssistantFunctionToolCall] = []
for idx in sorted(merge): for idx in sorted(merge):
entry = merge[idx] entry = merge[idx]
entry_id = entry["id"] if not entry["id"] or not entry["name"]:
fn_name = entry["name"]
if not entry_id or not fn_name:
continue continue
result.append( result.append(
AssistantFunctionToolCall( AssistantFunctionToolCall(
id=entry_id, id=entry["id"],
function=AssistantFunctionTool(name=fn_name, arguments=entry["arguments"]), function=AssistantFunctionTool(name=entry["name"], arguments=entry["arguments"]),
) )
) )
return result return result
+14 -23
View File
@@ -219,10 +219,23 @@ class ChatCompletionResponse(Struct):
service_tier: ServiceTier | Unset = UNSET 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): class DeltaMessage(Struct):
role: RoleAssistant | Unset = UNSET role: RoleAssistant | Unset = UNSET
content: str | Unset = UNSET content: str | Unset = UNSET
tool_calls: list[AnyAssistantToolCall] | Unset = UNSET tool_calls: list[StreamToolCallDelta] | Unset = UNSET
class ChunkChoice(Struct): class ChunkChoice(Struct):
@@ -239,25 +252,3 @@ class ChatCompletionChunk(Struct):
model: str model: str
choices: list[ChunkChoice] choices: list[ChunkChoice]
usage: dict[str, Any] | Unset = UNSET 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
+88 -12
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
from typing import TYPE_CHECKING, Literal, overload from typing import TYPE_CHECKING, Literal, overload
import pytest import pytest
from msgspec import UNSET
from plyngent.agent import ( from plyngent.agent import (
AssistantMessageEvent, AssistantMessageEvent,
@@ -25,6 +26,10 @@ from plyngent.lmproto.openai_compatible.model import (
ChatCompletionChunk, ChatCompletionChunk,
ChatCompletionResponse, ChatCompletionResponse,
ChatCompletionsParam, ChatCompletionsParam,
ChunkChoice,
DeltaMessage,
StreamFunctionDelta,
StreamToolCallDelta,
UserChatMessage, UserChatMessage,
) )
from plyngent.memory import MemoryStore from plyngent.memory import MemoryStore
@@ -33,8 +38,79 @@ if TYPE_CHECKING:
from collections.abc import AsyncIterator from collections.abc import AsyncIterator
def _chunks_from_response(response: ChatCompletionResponse) -> list[ChatCompletionChunk]:
"""Turn a full response into stream chunks (library-style stream=True path)."""
message = response.choices[0].message
chunks: list[ChatCompletionChunk] = []
if isinstance(message.content, str) and message.content:
chunks.append(
ChatCompletionChunk(
id=response.id,
object="chat.completion.chunk",
created=response.created,
model=response.model,
choices=[
ChunkChoice(
index=0,
delta=DeltaMessage(content=message.content),
finish_reason=None,
)
],
)
)
tool_calls = message.tool_calls
if tool_calls is not UNSET and tool_calls:
deltas: list[StreamToolCallDelta] = []
for i, call in enumerate(tool_calls):
if isinstance(call, AssistantFunctionToolCall):
deltas.append(
StreamToolCallDelta(
index=i,
id=call.id,
type="function",
function=StreamFunctionDelta(
name=call.function.name,
arguments=call.function.arguments,
),
)
)
if deltas:
chunks.append(
ChatCompletionChunk(
id=response.id,
object="chat.completion.chunk",
created=response.created,
model=response.model,
choices=[
ChunkChoice(
index=0,
delta=DeltaMessage(tool_calls=deltas),
finish_reason="tool_calls",
)
],
)
)
if not chunks:
chunks.append(
ChatCompletionChunk(
id=response.id,
object="chat.completion.chunk",
created=response.created,
model=response.model,
choices=[
ChunkChoice(
index=0,
delta=DeltaMessage(),
finish_reason=response.choices[0].finish_reason or "stop",
)
],
)
)
return chunks
class ScriptedClient: class ScriptedClient:
"""Returns scripted non-streaming chat completions in order.""" """Scripted chat completions; supports stream=True via chunked responses."""
_responses: list[ChatCompletionResponse] _responses: list[ChatCompletionResponse]
calls: list[ChatCompletionsParam] calls: list[ChatCompletionsParam]
@@ -57,16 +133,16 @@ class ScriptedClient:
self, param: ChatCompletionsParam, *, stream: bool = False self, param: ChatCompletionsParam, *, stream: bool = False
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]: ) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
self.calls.append(param) self.calls.append(param)
if stream:
return self._empty_stream()
if not self._responses: if not self._responses:
msg = "no more scripted responses" msg = "no more scripted responses"
raise RuntimeError(msg) raise RuntimeError(msg)
return self._responses.pop(0) response = self._responses.pop(0)
if stream:
return self._as_stream(response)
return response
async def _empty_stream(self) -> AsyncIterator[ChatCompletionChunk]: async def _as_stream(self, response: ChatCompletionResponse) -> AsyncIterator[ChatCompletionChunk]:
empty: list[ChatCompletionChunk] = [] for chunk in _chunks_from_response(response):
for chunk in empty:
yield chunk yield chunk
@@ -290,15 +366,15 @@ async def test_chat_agent_retry_after_failure() -> None:
if self.calls == 1: if self.calls == 1:
msg = "temporary" msg = "temporary"
raise RuntimeError(msg) raise RuntimeError(msg)
response = _response(AssistantChatMessage(content="recovered"))
if stream: if stream:
async def empty() -> AsyncIterator[ChatCompletionChunk]: async def as_stream() -> AsyncIterator[ChatCompletionChunk]:
empty_chunks: list[ChatCompletionChunk] = [] for chunk in _chunks_from_response(response):
for chunk in empty_chunks:
yield chunk yield chunk
return empty() return as_stream()
return _response(AssistantChatMessage(content="recovered")) return response
client = FlakyClient() client = FlakyClient()
agent = ChatAgent(client, model="m", memory=store, session_id=session.sid) agent = ChatAgent(client, model="m", memory=store, session_id=session.sid)
+23 -8
View File
@@ -13,6 +13,8 @@ from plyngent.lmproto.openai_compatible.model import (
ChatCompletionChunk, ChatCompletionChunk,
ChatCompletionResponse, ChatCompletionResponse,
ChatCompletionsParam, ChatCompletionsParam,
ChunkChoice,
DeltaMessage,
UserChatMessage, UserChatMessage,
) )
from plyngent.memory import MemoryStore from plyngent.memory import MemoryStore
@@ -40,6 +42,24 @@ def _response(message: AssistantChatMessage) -> ChatCompletionResponse:
) )
async def _as_stream(response: ChatCompletionResponse) -> AsyncIterator[ChatCompletionChunk]:
message = response.choices[0].message
content = message.content if isinstance(message.content, str) else ""
yield ChatCompletionChunk(
id=response.id,
object="chat.completion.chunk",
created=response.created,
model=response.model,
choices=[
ChunkChoice(
index=0,
delta=DeltaMessage(content=content),
finish_reason="stop",
)
],
)
class FlakyClient: class FlakyClient:
calls: int calls: int
fail_times: int fail_times: int
@@ -66,15 +86,10 @@ class FlakyClient:
if self.calls <= self.fail_times: if self.calls <= self.fail_times:
msg = f"fail-{self.calls}" msg = f"fail-{self.calls}"
raise RuntimeError(msg) raise RuntimeError(msg)
response = _response(AssistantChatMessage(content="ok"))
if stream: if stream:
return _as_stream(response)
async def empty() -> AsyncIterator[ChatCompletionChunk]: return response
empty_chunks: list[ChatCompletionChunk] = []
for chunk in empty_chunks:
yield chunk
return empty()
return _response(AssistantChatMessage(content="ok"))
async def test_sleep_cancellable_completes() -> None: async def test_sleep_cancellable_completes() -> None: