mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 07:17:54 +08:00
NCBM
6d7c239718
Registry cases (none / once / round / grant / YOLO / denial / error / handler record / closest answer wins), the wire cap keeping the line, the three policy paths that record their answer, and a pretty detail that ignores the header. The existing confirm tests now assert the line they gain.
258 lines
8.3 KiB
Python
258 lines
8.3 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import TYPE_CHECKING, Literal, overload
|
|
|
|
from msgspec import UNSET
|
|
|
|
from plyngent.agent import ChatAgent, ToolRegistry, tool
|
|
from plyngent.lmproto.openai_compatible.model import (
|
|
AssistantChatMessage,
|
|
AssistantFunctionTool,
|
|
AssistantFunctionToolCall,
|
|
ChatCompletionChoice,
|
|
ChatCompletionChunk,
|
|
ChatCompletionResponse,
|
|
ChatCompletionsParam,
|
|
ChunkChoice,
|
|
DeltaMessage,
|
|
StreamFunctionDelta,
|
|
StreamToolCallDelta,
|
|
ToolChatMessage,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncIterator, Mapping
|
|
|
|
|
|
def _response(message: AssistantChatMessage) -> ChatCompletionResponse:
|
|
return ChatCompletionResponse(
|
|
id="1",
|
|
object="chat.completion",
|
|
created=0,
|
|
model="t",
|
|
choices=[ChatCompletionChoice(index=0, message=message, logprobs={}, finish_reason="stop")],
|
|
system_fingerprint="",
|
|
usage={},
|
|
)
|
|
|
|
|
|
class ParallelScriptedClient:
|
|
step: int
|
|
|
|
def __init__(self) -> None:
|
|
self.step = 0
|
|
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[False] = False
|
|
) -> ChatCompletionResponse: ...
|
|
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[True]
|
|
) -> AsyncIterator[ChatCompletionChunk]: ...
|
|
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: bool = False
|
|
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
|
|
del param
|
|
self.step += 1
|
|
if self.step == 1:
|
|
msg = AssistantChatMessage(
|
|
content="",
|
|
tool_calls=[
|
|
AssistantFunctionToolCall(
|
|
id="1",
|
|
function=AssistantFunctionTool(name="slow_a", arguments="{}"),
|
|
),
|
|
AssistantFunctionToolCall(
|
|
id="2",
|
|
function=AssistantFunctionTool(name="slow_b", arguments="{}"),
|
|
),
|
|
],
|
|
)
|
|
else:
|
|
msg = AssistantChatMessage(content="done")
|
|
response = _response(msg)
|
|
if stream:
|
|
return self._stream(response)
|
|
return response
|
|
|
|
async def _stream(self, response: ChatCompletionResponse) -> AsyncIterator[ChatCompletionChunk]:
|
|
message = response.choices[0].message
|
|
if isinstance(message.content, str) and message.content:
|
|
yield ChatCompletionChunk(
|
|
id="1",
|
|
object="chat.completion.chunk",
|
|
created=0,
|
|
model="t",
|
|
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,
|
|
),
|
|
)
|
|
)
|
|
yield ChatCompletionChunk(
|
|
id="1",
|
|
object="chat.completion.chunk",
|
|
created=0,
|
|
model="t",
|
|
choices=[
|
|
ChunkChoice(
|
|
index=0,
|
|
delta=DeltaMessage(tool_calls=deltas),
|
|
finish_reason="tool_calls",
|
|
)
|
|
],
|
|
)
|
|
|
|
|
|
async def test_parallel_tool_calls_run_concurrently() -> None:
|
|
started = 0
|
|
expected_parallel = 2
|
|
gate = asyncio.Event()
|
|
wait_timeout = 2.0
|
|
|
|
@tool(register=False)
|
|
async def slow_a() -> str:
|
|
nonlocal started
|
|
started += 1
|
|
if started >= expected_parallel:
|
|
_ = gate.set()
|
|
_ = await asyncio.wait_for(gate.wait(), timeout=wait_timeout)
|
|
return "a"
|
|
|
|
@tool(register=False)
|
|
async def slow_b() -> str:
|
|
nonlocal started
|
|
started += 1
|
|
if started >= expected_parallel:
|
|
_ = gate.set()
|
|
_ = await asyncio.wait_for(gate.wait(), timeout=wait_timeout)
|
|
return "b"
|
|
|
|
registry = ToolRegistry([slow_a, slow_b])
|
|
agent = ChatAgent(
|
|
ParallelScriptedClient(),
|
|
model="m",
|
|
tools=registry,
|
|
stream=True,
|
|
parallel_tools=True,
|
|
)
|
|
events = [e async for e in agent.run("go")]
|
|
tool_msgs = [e for e in agent.messages if isinstance(e, ToolChatMessage)]
|
|
assert {m.content for m in tool_msgs} == {"a", "b"}
|
|
assert any(getattr(e, "content", None) == "done" for e in events)
|
|
|
|
|
|
async def test_same_file_writes_apply_in_call_order() -> None:
|
|
"""A batch's writes to one file run one at a time, in call order."""
|
|
order: list[str] = []
|
|
|
|
@tool(register=False, writes=("path",))
|
|
async def write_a(path: str) -> str:
|
|
order.append(f"a start {path}")
|
|
await asyncio.sleep(0) # an unserialized sibling would interleave here
|
|
order.append(f"a end {path}")
|
|
return "a"
|
|
|
|
@tool(register=False, writes=("path",))
|
|
async def write_b(path: str) -> str:
|
|
order.append(f"b start {path}")
|
|
await asyncio.sleep(0)
|
|
order.append(f"b end {path}")
|
|
return "b"
|
|
|
|
registry = ToolRegistry([write_a, write_b])
|
|
results = await asyncio.gather(
|
|
registry.execute("write_a", '{"path": "a.txt"}'),
|
|
registry.execute("write_b", '{"path": "./a.txt"}'),
|
|
)
|
|
assert results == ["a", "b"]
|
|
assert order == ["a start a.txt", "a end a.txt", "b start ./a.txt", "b end ./a.txt"]
|
|
|
|
|
|
async def test_different_files_still_write_concurrently() -> None:
|
|
"""Serialization is per file: two files of one batch still overlap."""
|
|
overlaps = 0
|
|
gate = asyncio.Event()
|
|
|
|
async def _rendezvous() -> None:
|
|
nonlocal overlaps
|
|
overlaps += 1
|
|
if overlaps == 2:
|
|
_ = gate.set()
|
|
_ = await asyncio.wait_for(gate.wait(), timeout=2.0)
|
|
|
|
@tool(register=False, writes=("path",))
|
|
async def write_a(path: str) -> str:
|
|
del path
|
|
await _rendezvous()
|
|
return "a"
|
|
|
|
@tool(register=False, writes=("path",))
|
|
async def write_b(path: str) -> str:
|
|
del path
|
|
await _rendezvous()
|
|
return "b"
|
|
|
|
registry = ToolRegistry([write_a, write_b])
|
|
results = await asyncio.gather(
|
|
registry.execute("write_a", '{"path": "a.txt"}'),
|
|
registry.execute("write_b", '{"path": "b.txt"}'),
|
|
)
|
|
assert results == ["a", "b"]
|
|
|
|
|
|
async def test_a_write_holds_its_file_across_the_confirm() -> None:
|
|
"""The file is locked before the prompt, so a sibling write cannot overtake it."""
|
|
order: list[str] = []
|
|
prompt_shown = asyncio.Event()
|
|
answered = asyncio.Event()
|
|
|
|
@tool(register=False, writes=("path",))
|
|
async def write_a(path: str) -> str:
|
|
del path
|
|
order.append("a")
|
|
return "a"
|
|
|
|
@tool(register=False, writes=("path",))
|
|
async def write_b(path: str) -> str:
|
|
del path
|
|
order.append("b")
|
|
return "b"
|
|
|
|
async def on_confirm(name: str, args: Mapping[str, object], reason: str) -> bool:
|
|
del args, reason
|
|
order.append(f"confirm {name}")
|
|
prompt_shown.set()
|
|
await answered.wait()
|
|
return True
|
|
|
|
def danger(name: str, args: Mapping[str, object]) -> str | None:
|
|
del args
|
|
return "overwrite" if name == "write_a" else None
|
|
|
|
registry = ToolRegistry([write_a, write_b], danger=danger, on_confirm=on_confirm)
|
|
first = asyncio.create_task(registry.execute("write_a", '{"path": "a.txt"}'))
|
|
await prompt_shown.wait()
|
|
second = asyncio.create_task(registry.execute("write_b", '{"path": "a.txt"}'))
|
|
await asyncio.sleep(0) # the second write waits for the file, not for the human
|
|
assert order == ["confirm write_a"]
|
|
answered.set()
|
|
assert await asyncio.gather(first, second) == ["[approval] user (once)\na", "b"]
|
|
assert order == ["confirm write_a", "a", "b"]
|