mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 22:44:59 +08:00
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:
@@ -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
@@ -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:
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
|
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
Reference in New Issue
Block a user