From 1371a8fdf9968577507d4fecb020c9edbc86e523 Mon Sep 17 00:00:00 2001 From: Michael Slezak Date: Wed, 11 Mar 2026 23:30:10 -0600 Subject: [PATCH] Add tool call annotation support. --- .gitignore | 3 + README.md | 2 +- examples/.env | 1 + examples/simple_agent.py | 20 +++-- src/llpsdk/__init__.py | 4 + src/llpsdk/client.py | 72 +++------------ src/llpsdk/handler.py | 21 ++++- src/llpsdk/message.py | 28 ++++++ src/llpsdk/tool_call.py | 35 ++++++++ tests/test_handler.py | 31 +++++-- tests/test_tool_call.py | 189 +++++++++++++++++++++++++++++++++++++++ 11 files changed, 324 insertions(+), 82 deletions(-) create mode 100644 src/llpsdk/tool_call.py create mode 100644 tests/test_tool_call.py diff --git a/.gitignore b/.gitignore index 9ff2485..c05f861 100644 --- a/.gitignore +++ b/.gitignore @@ -54,3 +54,6 @@ dmypy.json # OS .DS_Store Thumbs.db + +# Secrets +.env diff --git a/README.md b/README.md index 69d2251..9f68a12 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ import asyncio, os import llpsdk as llp # Define a callback handler for processing messages -async def on_message(msg): +async def on_message(annotater, msg): # Process the prompt with your agent. # Replace this with your own processing logic. response = msg.prompt diff --git a/examples/.env b/examples/.env index 7e7970b..71cfb59 100644 --- a/examples/.env +++ b/examples/.env @@ -1 +1,2 @@ LLP_URL="ws://localhost:4000/agent/websocket" +LLP_API_KEY= diff --git a/examples/simple_agent.py b/examples/simple_agent.py index 3eb2121..9a1a1df 100644 --- a/examples/simple_agent.py +++ b/examples/simple_agent.py @@ -1,5 +1,6 @@ """Agent example demonstrating basic LLP SDK usage.""" import asyncio +from datetime import timedelta import llpsdk as llp import os from dotenv import load_dotenv @@ -9,16 +10,22 @@ async def main() -> None: """Run a simple agent that connects, sends presence, and sends a message.""" load_dotenv() platform_url = os.getenv("LLP_URL") + api_key = os.getenv("LLP_API_KEY") + if platform_url is None: raise Exception("LLP_URL env var is not defined") - cfg = llp.Config() + if api_key is None: + raise Exception("LLP_API_KEY env var is not defined") + + cfg = llp.Config(platform_url=platform_url) cfg.platform_url = platform_url - client = llp.Client("simple-agent", "testkey", cfg) + client = llp.Client("simple-agent", api_key, cfg) # Set up handlers - async def on_message(msg: llp.TextMessage) -> llp.TextMessage: - print("Feed msg.prompt into your agent and return the response.") + async def on_message(annotater: llp.Annotater, msg: llp.TextMessage) -> llp.TextMessage: + tc = msg.tool_call("get_weather", '{"city":"Seattle"}', "rainy", timedelta(seconds=1)) + await annotater.annotate_tool_call(tc) return msg.reply("this is my response") # Register handlers @@ -30,11 +37,6 @@ async def on_message(msg: llp.TextMessage) -> llp.TextMessage: await client.connect() print(f"Connected! Session ID: {client.session_id}") - # Send a message - msg = llp.TextMessage(recipient="echo-agent", prompt="Hello from Python!") - print(f"Sending message to {msg.recipient}...") - await client.send_message(msg) - # Keep running print("Agent running. Press Ctrl+C to exit...") await asyncio.Event().wait() diff --git a/src/llpsdk/__init__.py b/src/llpsdk/__init__.py index 4add306..6f2c864 100644 --- a/src/llpsdk/__init__.py +++ b/src/llpsdk/__init__.py @@ -11,12 +11,14 @@ PlatformError, TimeoutError, ) +from .handler import Annotater from .message import ( AuthenticatedResponse, PresenceMessage, TextMessage, ) from .presence import ConnectionStatus, PresenceStatus +from .tool_call import ToolCall __all__ = [ "Client", @@ -33,6 +35,8 @@ "AlreadyClosedError", "TimeoutError", "InvalidStatusError", + "Annotater", + "ToolCall", ] __version__ = "0.1.0" diff --git a/src/llpsdk/client.py b/src/llpsdk/client.py index bee19c1..81e90bd 100644 --- a/src/llpsdk/client.py +++ b/src/llpsdk/client.py @@ -23,6 +23,7 @@ PresenceMessage, TextMessage, ) +from .tool_call import ToolCall from .presence import ConnectionStatus, PresenceStatus @@ -64,10 +65,6 @@ def __init__(self, name: str, api_key: str, config: Optional[Config] = None) -> self._write_task: Optional[asyncio.Task[None]] = None self._stop_event = asyncio.Event() - # Pending messages (for request/response) - self._pending_lock = asyncio.Lock() - self._pending: Dict[str, asyncio.Future[TextMessage]] = {} - # Auth future (for waiting on authentication) self._auth_future: Optional[asyncio.Future[AuthenticatedResponse]] = None @@ -170,7 +167,7 @@ async def close(self) -> None: async with self._presence_lock: self._presence = PresenceStatus.unavailable - async def send_async_message(self, message: TextMessage) -> None: + async def _send_async_message(self, message: TextMessage) -> None: """ Send a message asynchronously (fire-and-forget). @@ -187,50 +184,22 @@ async def send_async_message(self, message: TextMessage) -> None: await self._send(message.encode()) - async def send_message(self, message: TextMessage, timeout: float = 10.0) -> TextMessage: + async def annotate_tool_call(self, tool_call: ToolCall) -> None: """ - Send a message and wait for response. + Send a tool call annotation to the platform for telemetry. Args: - message: Message to send - timeout: Response timeout in seconds - - Returns: - Response message + tool_call: The tool call to annotate (created via TextMessage.tool_call() or + TextMessage.tool_call_exception()) Raises: - ValueError: If message ID is empty NotAuthenticatedError: If not authenticated - TimeoutError: If no response within timeout """ - # ID is REQUIRED for synchronous send - if not message._id: - raise ValueError("Message ID is required for send_message()") - async with self._status_lock: if self._status != ConnectionStatus.AUTHENTICATED: - raise NotAuthenticatedError("Must connect before sending messages") + raise NotAuthenticatedError("Must connect before annotating tool calls") - # Create future for this message - response_future: asyncio.Future[TextMessage] = asyncio.get_event_loop().create_future() - - async with self._pending_lock: - self._pending[message._id] = response_future - - try: - # Send message asynchronously - await self.send_async_message(message) - - # Wait for response with timeout - response = await asyncio.wait_for(response_future, timeout=timeout) - return response - - except asyncio.TimeoutError: - raise TimeoutError(f"No response within {timeout}s") - finally: - # Clean up - async with self._pending_lock: - self._pending.pop(message._id, None) + await self._send(tool_call.encode()) # Properties @@ -415,15 +384,9 @@ async def _handle_message(self, msg_dict: Dict[str, Any]) -> None: self._auth_future.set_exception(error) return - # Check if this error is for a pending message - if error.id: - async with self._pending_lock: - if error.id in self._pending: - future = self._pending.pop(error.id) - if not future.done(): - future.set_exception(error) - return + return + if msg_type == "ack": return if msg_type == "authenticated": @@ -438,21 +401,10 @@ async def _handle_message(self, msg_dict: Dict[str, Any]) -> None: return if msg_type == "message": - msg_id = msg_dict.get("id", "") - - async with self._pending_lock: - if msg_id in self._pending: - future = self._pending[msg_id] - if not future.done(): - tm = TextMessage.decode(msg_dict) - future.set_result(tm) - return - - # Not a response, call message handler tm = TextMessage.decode(msg_dict) - reply = await self._handlers.call_message(tm) + reply = await self._handlers.call_message(self, tm) if reply is not None: - await self.send_async_message(reply) + await self._send_async_message(reply) return async def _handle_disconnect(self) -> None: diff --git a/src/llpsdk/handler.py b/src/llpsdk/handler.py index aae84eb..bfb447f 100644 --- a/src/llpsdk/handler.py +++ b/src/llpsdk/handler.py @@ -3,13 +3,24 @@ import asyncio from typing import Awaitable, Callable, Optional, Union +from typing import Protocol, runtime_checkable + from .message import PresenceMessage, TextMessage +from .tool_call import ToolCall + + +@runtime_checkable +class Annotater(Protocol): + """Protocol for annotating tool calls for telemetry.""" + + async def annotate_tool_call(self, tool_call: ToolCall) -> None: ... + # Handler type signatures (supports both sync and async) PresenceHandler = Union[ Callable[[PresenceMessage], None], Callable[[PresenceMessage], Awaitable[None]] ] -MessageHandler = Callable[[TextMessage], Awaitable[TextMessage]] +MessageHandler = Callable[["Annotater", TextMessage], Awaitable[TextMessage]] class HandlerRegistry: @@ -36,9 +47,11 @@ async def call_presence(self, update: PresenceMessage) -> None: else: self._on_presence(update) - async def call_message(self, message: TextMessage) -> Optional[TextMessage]: - """Call the message handler if set.""" + async def call_message( + self, annotater: Annotater, message: TextMessage + ) -> Optional[TextMessage]: + """Call the message handler if set, passing annotater for tool call telemetry.""" if self._on_message is not None: - result = await self._on_message(message) + result = await self._on_message(annotater, message) return result return None diff --git a/src/llpsdk/message.py b/src/llpsdk/message.py index b72dbb7..a2f4278 100644 --- a/src/llpsdk/message.py +++ b/src/llpsdk/message.py @@ -4,10 +4,12 @@ import uuid import json from dataclasses import dataclass +from datetime import timedelta from typing import Any, Dict, Optional from llpsdk.errors import TextMessageEmptyError from llpsdk.presence import PresenceStatus +from llpsdk.tool_call import ToolCall class TextMessage: @@ -57,6 +59,32 @@ def encode(self) -> str: def has_attachment(self) -> bool: return self.attachment != "" + def tool_call(self, name: str, parameters: str, result: str, duration: timedelta) -> ToolCall: + """Create a successful ToolCall annotation from this message.""" + return ToolCall( + id=self._id, + recipient=self.sender, + name=name, + parameters=parameters, + result=result, + threw_exception=False, + duration=duration, + ) + + def tool_call_exception( + self, name: str, parameters: str, error: Exception, duration: timedelta + ) -> ToolCall: + """Create a failed ToolCall annotation from this message.""" + return ToolCall( + id=self._id, + recipient=self.sender, + name=name, + parameters=parameters, + result=str(error), + threw_exception=True, + duration=duration, + ) + @staticmethod def decode(msg: Dict[str, Any]) -> "TextMessage": """Decodes a JSON dict into a TextMessage object""" diff --git a/src/llpsdk/tool_call.py b/src/llpsdk/tool_call.py new file mode 100644 index 0000000..87839b3 --- /dev/null +++ b/src/llpsdk/tool_call.py @@ -0,0 +1,35 @@ +"""ToolCall message type for telemetry annotation.""" + +import json +from dataclasses import dataclass +from datetime import timedelta +from typing import Optional + + +@dataclass +class ToolCall: + """Tool call annotation sent to the platform for telemetry.""" + + id: Optional[str] + recipient: str + name: str + parameters: str + result: str + threw_exception: bool + duration: timedelta + + def encode(self) -> str: + """Encode a ToolCall into serialized JSON.""" + data = { + "type": "tool_call", + "id": self.id, + "data": { + "to": self.recipient, + "name": self.name, + "parameters": self.parameters, + "result": self.result, + "threw_exception": self.threw_exception, + "duration_ms": int(self.duration.total_seconds() * 1000), + }, + } + return json.dumps(data) diff --git a/tests/test_handler.py b/tests/test_handler.py index c0e5e05..988408d 100644 --- a/tests/test_handler.py +++ b/tests/test_handler.py @@ -2,9 +2,19 @@ import pytest -from llpsdk.handler import HandlerRegistry +from llpsdk.handler import Annotater, HandlerRegistry from llpsdk.message import PresenceMessage, TextMessage from llpsdk.presence import PresenceStatus +from llpsdk.tool_call import ToolCall + + +class _FakeAnnotater: + async def annotate_tool_call(self, tool_call: ToolCall) -> None: + pass + + +_annotater = _FakeAnnotater() + @pytest.mark.asyncio async def test_async_message_handler(): @@ -12,13 +22,14 @@ async def test_async_message_handler(): registry = HandlerRegistry() calls = [] - async def async_handler(msg: TextMessage) -> TextMessage: + async def async_handler(ann: Annotater, msg: TextMessage) -> TextMessage: calls.append(("async", msg.prompt)) return msg.reply("test") registry.set_message(async_handler) msg = TextMessage("alice", "World") - await registry.call_message(msg) + msg.sender = "bob" + await registry.call_message(_annotater, msg) assert len(calls) == 1 assert calls[0] == ("async", "World") @@ -64,7 +75,7 @@ async def test_no_handler_set(): registry = HandlerRegistry() # Should not raise - await registry.call_message(TextMessage("alice", "test")) + await registry.call_message(_annotater, TextMessage("alice", "test")) await registry.call_presence(PresenceMessage(sender="alice", status=PresenceStatus.available)) @@ -74,18 +85,22 @@ async def test_handler_replacement(): registry = HandlerRegistry() calls = [] - async def handler1(msg: TextMessage) -> TextMessage: + async def handler1(ann: Annotater, msg: TextMessage) -> TextMessage: calls.append("handler1") return msg.reply("test") - async def handler2(msg: TextMessage) -> TextMessage: + async def handler2(ann: Annotater, msg: TextMessage) -> TextMessage: calls.append("handler2") return msg.reply("test") registry.set_message(handler1) - await registry.call_message(TextMessage("alice", "test1")) + msg1 = TextMessage("alice", "test1") + msg1.sender = "bob" + await registry.call_message(_annotater, msg1) registry.set_message(handler2) - await registry.call_message(TextMessage("alice", "test2")) + msg2 = TextMessage("alice", "test2") + msg2.sender = "bob" + await registry.call_message(_annotater, msg2) assert calls == ["handler1", "handler2"] diff --git a/tests/test_tool_call.py b/tests/test_tool_call.py new file mode 100644 index 0000000..e5bbfc4 --- /dev/null +++ b/tests/test_tool_call.py @@ -0,0 +1,189 @@ +"""Tests for ToolCall annotation feature.""" + +import json +from datetime import timedelta +from unittest.mock import AsyncMock + +import pytest + +from llpsdk.handler import Annotater, HandlerRegistry +from llpsdk.message import TextMessage +from llpsdk.tool_call import ToolCall + + +# --- ToolCall encode --- + + +def test_tool_call_encode(): + tc = ToolCall( + id="msg-1", + recipient="alice", + name="search", + parameters='{"query": "weather"}', + result="Sunny, 72°F", + threw_exception=False, + duration=timedelta(milliseconds=150), + ) + data = json.loads(tc.encode()) + + assert data["type"] == "tool_call" + assert data["id"] == "msg-1" + assert data["data"]["to"] == "alice" + assert data["data"]["name"] == "search" + assert data["data"]["parameters"] == '{"query": "weather"}' + assert data["data"]["result"] == "Sunny, 72°F" + assert data["data"]["threw_exception"] is False + assert data["data"]["duration_ms"] == 150 + + +def test_tool_call_encode_exception(): + tc = ToolCall( + id="msg-2", + recipient="bob", + name="fetch_data", + parameters="{}", + result="Connection refused", + threw_exception=True, + duration=timedelta(seconds=1), + ) + data = json.loads(tc.encode()) + + assert data["data"]["threw_exception"] is True + assert data["data"]["duration_ms"] == 1000 + + +def test_tool_call_encode_duration_fractional_seconds(): + tc = ToolCall( + id="msg-3", + recipient="bob", + name="slow_tool", + parameters="{}", + result="done", + threw_exception=False, + duration=timedelta(milliseconds=1500), + ) + data = json.loads(tc.encode()) + assert data["data"]["duration_ms"] == 1500 + + +# --- TextMessage factory methods --- + + +def test_text_message_tool_call_factory(): + msg = TextMessage("bob", "hello") + msg._id = "msg-42" + msg.sender = "alice" + + tc = msg.tool_call("weather", '{"city": "NYC"}', "Sunny", timedelta(milliseconds=200)) + + assert tc.id == "msg-42" + assert tc.recipient == "alice" + assert tc.name == "weather" + assert tc.parameters == '{"city": "NYC"}' + assert tc.result == "Sunny" + assert tc.threw_exception is False + assert tc.duration == timedelta(milliseconds=200) + + +def test_text_message_tool_call_exception_factory(): + msg = TextMessage("bob", "hello") + msg._id = "msg-43" + msg.sender = "carol" + + err = ValueError("service unavailable") + tc = msg.tool_call_exception("fetch", '{"url": "..."}', err, timedelta(seconds=2)) + + assert tc.id == "msg-43" + assert tc.recipient == "carol" + assert tc.name == "fetch" + assert tc.result == "service unavailable" + assert tc.threw_exception is True + assert tc.duration == timedelta(seconds=2) + + +# --- Annotater Protocol --- + + +def test_annotater_protocol_satisfied(): + """Client-like objects that implement annotate_tool_call satisfy Annotater.""" + + class FakeClient: + async def annotate_tool_call(self, tool_call: ToolCall) -> None: + pass + + assert isinstance(FakeClient(), Annotater) + + +def test_annotater_protocol_not_satisfied(): + """Objects without annotate_tool_call do not satisfy Annotater.""" + + class NotAnAnnotater: + pass + + assert not isinstance(NotAnAnnotater(), Annotater) + + +# --- HandlerRegistry with annotater --- + + +@pytest.mark.asyncio +async def test_message_handler_receives_annotater(): + """Message handler receives the annotater as its first argument.""" + registry = HandlerRegistry() + received_annotater = [] + + class FakeAnnotater: + async def annotate_tool_call(self, tool_call: ToolCall) -> None: + pass + + annotater = FakeAnnotater() + + async def handler(ann: Annotater, msg: TextMessage) -> TextMessage: + received_annotater.append(ann) + return msg.reply("ok") + + registry.set_message(handler) + msg = TextMessage("alice", "hello") + msg.sender = "bob" + await registry.call_message(annotater, msg) + + assert len(received_annotater) == 1 + assert received_annotater[0] is annotater + + +@pytest.mark.asyncio +async def test_handler_can_annotate_tool_call(): + """Handler can call annotate_tool_call on the annotater it receives.""" + registry = HandlerRegistry() + annotated: list[ToolCall] = [] + + class FakeAnnotater: + async def annotate_tool_call(self, tool_call: ToolCall) -> None: + annotated.append(tool_call) + + annotater = FakeAnnotater() + + async def handler(ann: Annotater, msg: TextMessage) -> TextMessage: + tc = msg.tool_call("lookup", "{}", "result", timedelta(milliseconds=50)) + await ann.annotate_tool_call(tc) + return msg.reply("done") + + registry.set_message(handler) + msg = TextMessage("alice", "hello") + msg._id = "msg-99" + msg.sender = "bob" + await registry.call_message(annotater, msg) + + assert len(annotated) == 1 + assert annotated[0].name == "lookup" + assert annotated[0].id == "msg-99" + + +@pytest.mark.asyncio +async def test_no_message_handler_with_annotater(): + """call_message returns None without error when no handler is set.""" + registry = HandlerRegistry() + annotater = AsyncMock(spec=Annotater) + + result = await registry.call_message(annotater, TextMessage("alice", "test")) + assert result is None