diff --git a/tests/test_agent/test_loop.py b/tests/test_agent/test_loop.py new file mode 100644 index 0000000..df218c5 --- /dev/null +++ b/tests/test_agent/test_loop.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Literal, overload + +from plyngent.agent import ( + AssistantMessageEvent, + ChatAgent, + MaxRoundsEvent, + TextDeltaEvent, + ToolCallEvent, + ToolRegistry, + ToolResultEvent, + run_chat_loop, + tool, +) +from plyngent.config.models import DatabaseConfig +from plyngent.lmproto.openai_compatible.model import ( + AnyChatMessage, + AssistantChatMessage, + AssistantFunctionTool, + AssistantFunctionToolCall, + ChatCompletionChoice, + ChatCompletionChunk, + ChatCompletionResponse, + ChatCompletionsParam, + UserChatMessage, +) +from plyngent.memory import MemoryStore + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + +class ScriptedClient: + """Returns scripted non-streaming chat completions in order.""" + + _responses: list[ChatCompletionResponse] + calls: list[ChatCompletionsParam] + + def __init__(self, responses: list[ChatCompletionResponse]) -> None: + self._responses = list(responses) + self.calls = [] + + @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]: + 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) + + async def _empty_stream(self) -> AsyncIterator[ChatCompletionChunk]: + if False: # pyright: ignore[reportUnreachable] + yield # type: ignore[misc] + return + + +def _response(message: AssistantChatMessage) -> ChatCompletionResponse: + return ChatCompletionResponse( + id="1", + object="chat.completion", + created=0, + model="test", + choices=[ + ChatCompletionChoice( + index=0, + message=message, + logprobs={}, + finish_reason="stop", + ) + ], + system_fingerprint="", + usage={}, + ) + + +async def test_run_chat_loop_text_only() -> None: + client = ScriptedClient( + [ + _response(AssistantChatMessage(content="hello")), + ] + ) + messages: list[AnyChatMessage] = [UserChatMessage(content="hi")] + events = [e async for e in run_chat_loop(client, messages, model="m")] + assert isinstance(events[0], AssistantMessageEvent) + assert isinstance(events[1], TextDeltaEvent) + assert events[1].content == "hello" + assert len(messages) == 2 # noqa: PLR2004 + assert len(client.calls) == 1 + + +async def test_run_chat_loop_with_tools() -> None: + @tool + def add(a: int, b: int) -> int: + return a + b + + registry = ToolRegistry([add]) + client = ScriptedClient( + [ + _response( + AssistantChatMessage( + content="", + tool_calls=[ + AssistantFunctionToolCall( + id="c1", + function=AssistantFunctionTool(name="add", arguments='{"a": 1, "b": 2}'), + ) + ], + ) + ), + _response(AssistantChatMessage(content="3")), + ] + ) + messages: list[AnyChatMessage] = [UserChatMessage(content="1+2")] + events = [e async for e in run_chat_loop(client, messages, model="m", tools=registry)] + types = [type(e) for e in events] + assert ToolCallEvent in types + assert ToolResultEvent in types + assert any(isinstance(e, TextDeltaEvent) and e.content == "3" for e in events) + assert len(client.calls) == 2 # noqa: PLR2004 + # second call includes tool result message + assert any(getattr(m, "tool_call_id", None) == "c1" for m in client.calls[1].messages) + + +async def test_max_rounds() -> None: + @tool + def ping() -> str: + return "pong" + + registry = ToolRegistry([ping]) + forever = _response( + AssistantChatMessage( + content="", + tool_calls=[ + AssistantFunctionToolCall( + id="c", + function=AssistantFunctionTool(name="ping", arguments="{}"), + ) + ], + ) + ) + client = ScriptedClient([forever, forever, forever]) + messages: list[AnyChatMessage] = [UserChatMessage(content="x")] + events = [e async for e in run_chat_loop(client, messages, model="m", tools=registry, max_rounds=2)] + assert any(isinstance(e, MaxRoundsEvent) and e.rounds == 2 for e in events) # noqa: PLR2004 + assert len(client.calls) == 2 # noqa: PLR2004 + + +async def test_chat_agent_memory_roundtrip() -> None: + store = await MemoryStore.open(DatabaseConfig()) + session = await store.create_session(name="t") + client = ScriptedClient([_response(AssistantChatMessage(content="yo"))]) + agent = ChatAgent(client, model="m", memory=store, session_id=session.sid) + events = [e async for e in agent.run("hi")] + assert any(isinstance(e, TextDeltaEvent) and e.content == "yo" for e in events) + + loaded = await store.list_messages(session.sid) + assert len(loaded) == 2 # noqa: PLR2004 + assert isinstance(loaded[0], UserChatMessage) + assert loaded[0].content == "hi" + + agent2 = ChatAgent(client, model="m", memory=store, session_id=session.sid) + await agent2.load_history() + assert len(agent2.messages) == 2 # noqa: PLR2004 + await store.close() diff --git a/tests/test_agent/test_tools.py b/tests/test_agent/test_tools.py new file mode 100644 index 0000000..0adf409 --- /dev/null +++ b/tests/test_agent/test_tools.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +from plyngent.agent import ToolRegistry, schema_from_callable, tool + + +def test_tool_decorator_infers_schema() -> None: + @tool + def add(a: int, b: int = 0) -> int: + """Add two numbers.""" + return a + b + + assert add.name == "add" + assert add.description == "Add two numbers." + assert add.parameters["type"] == "object" + assert add.parameters["properties"]["a"] == {"type": "integer"} + assert add.parameters["properties"]["b"] == {"type": "integer"} + assert add.parameters["required"] == ["a"] + + +def test_tool_decorator_overrides() -> None: + @tool(name="ping", description="health") + def health() -> str: + return "ok" + + assert health.name == "ping" + assert health.description == "health" + + +def test_schema_from_callable_optional() -> None: + def f(x: str | None = None) -> str: + return x or "" + + schema = schema_from_callable(f) + assert schema["properties"]["x"] == {"type": "string"} + assert "required" not in schema + + +async def test_registry_execute_sync_and_async() -> None: + @tool + def echo(text: str) -> str: + return text + + @tool + async def upper(text: str) -> str: + return text.upper() + + registry = ToolRegistry([echo, upper]) + assert await registry.execute("echo", '{"text": "hi"}') == "hi" + assert await registry.execute("upper", '{"text": "hi"}') == "HI" + + +async def test_registry_unknown_and_bad_json() -> None: + @tool + def noop() -> str: + return "ok" + + registry = ToolRegistry([noop]) + assert "unknown" in await registry.execute("missing", "{}") + assert "invalid" in await registry.execute("noop", "not-json") + + +async def test_registry_handler_error() -> None: + @tool + def boom() -> str: + msg = "nope" + raise RuntimeError(msg) + + registry = ToolRegistry([boom]) + result = await registry.execute("boom", "{}") + assert "failed" in result + + +def test_tool_items() -> None: + @tool + def f(x: int) -> int: + return x + + items = ToolRegistry([f]).tool_items() + assert len(items) == 1 + assert items[0].function.name == "f" + + +def test_schema_list_and_bool() -> None: + def g(flags: list[bool], *, ok: bool) -> None: + del flags, ok + + schema = schema_from_callable(g) + assert schema["properties"]["flags"]["type"] == "array" + assert schema["properties"]["flags"]["items"] == {"type": "boolean"} + assert schema["properties"]["ok"] == {"type": "boolean"}