diff --git a/src/plyngent/agent/client.py b/src/plyngent/agent/client.py index 8c942f1..d217024 100644 --- a/src/plyngent/agent/client.py +++ b/src/plyngent/agent/client.py @@ -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 diff --git a/src/plyngent/agent/loop.py b/src/plyngent/agent/loop.py index 6dda594..617f439 100644 --- a/src/plyngent/agent/loop.py +++ b/src/plyngent/agent/loop.py @@ -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: diff --git a/src/plyngent/lmproto/openai_compatible/client.py b/src/plyngent/lmproto/openai_compatible/client.py index 5a167da..8586874 100644 --- a/src/plyngent/lmproto/openai_compatible/client.py +++ b/src/plyngent/lmproto/openai_compatible/client.py @@ -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 diff --git a/src/plyngent/lmproto/openai_compatible/model.py b/src/plyngent/lmproto/openai_compatible/model.py index a767968..2e25608 100644 --- a/src/plyngent/lmproto/openai_compatible/model.py +++ b/src/plyngent/lmproto/openai_compatible/model.py @@ -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 diff --git a/tests/test_agent/test_loop.py b/tests/test_agent/test_loop.py index b89f8ff..c6d4d85 100644 --- a/tests/test_agent/test_loop.py +++ b/tests/test_agent/test_loop.py @@ -3,6 +3,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, Literal, overload import pytest +from msgspec import UNSET from plyngent.agent import ( AssistantMessageEvent, @@ -25,6 +26,10 @@ from plyngent.lmproto.openai_compatible.model import ( ChatCompletionChunk, ChatCompletionResponse, ChatCompletionsParam, + ChunkChoice, + DeltaMessage, + StreamFunctionDelta, + StreamToolCallDelta, UserChatMessage, ) from plyngent.memory import MemoryStore @@ -33,8 +38,79 @@ if TYPE_CHECKING: 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: - """Returns scripted non-streaming chat completions in order.""" + """Scripted chat completions; supports stream=True via chunked responses.""" _responses: list[ChatCompletionResponse] calls: list[ChatCompletionsParam] @@ -57,16 +133,16 @@ class ScriptedClient: self, param: ChatCompletionsParam, *, stream: bool = False ) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]: self.calls.append(param) - if stream: - return self._empty_stream() if not self._responses: msg = "no more scripted responses" 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]: - empty: list[ChatCompletionChunk] = [] - for chunk in empty: + async def _as_stream(self, response: ChatCompletionResponse) -> AsyncIterator[ChatCompletionChunk]: + for chunk in _chunks_from_response(response): yield chunk @@ -290,15 +366,15 @@ async def test_chat_agent_retry_after_failure() -> None: if self.calls == 1: msg = "temporary" raise RuntimeError(msg) + response = _response(AssistantChatMessage(content="recovered")) if stream: - async def empty() -> AsyncIterator[ChatCompletionChunk]: - empty_chunks: list[ChatCompletionChunk] = [] - for chunk in empty_chunks: + async def as_stream() -> AsyncIterator[ChatCompletionChunk]: + for chunk in _chunks_from_response(response): yield chunk - return empty() - return _response(AssistantChatMessage(content="recovered")) + return as_stream() + return response client = FlakyClient() agent = ChatAgent(client, model="m", memory=store, session_id=session.sid) diff --git a/tests/test_cli/test_retry.py b/tests/test_cli/test_retry.py index 0e1c61b..3e0e168 100644 --- a/tests/test_cli/test_retry.py +++ b/tests/test_cli/test_retry.py @@ -13,6 +13,8 @@ from plyngent.lmproto.openai_compatible.model import ( ChatCompletionChunk, ChatCompletionResponse, ChatCompletionsParam, + ChunkChoice, + DeltaMessage, UserChatMessage, ) 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: calls: int fail_times: int @@ -66,15 +86,10 @@ class FlakyClient: if self.calls <= self.fail_times: msg = f"fail-{self.calls}" raise RuntimeError(msg) + response = _response(AssistantChatMessage(content="ok")) if stream: - - async def empty() -> AsyncIterator[ChatCompletionChunk]: - empty_chunks: list[ChatCompletionChunk] = [] - for chunk in empty_chunks: - yield chunk - - return empty() - return _response(AssistantChatMessage(content="ok")) + return _as_stream(response) + return response async def test_sleep_cancellable_completes() -> None: