From 4a6f31d3317bbdfb4f29c1dad9fa08e98ca34ed6 Mon Sep 17 00:00:00 2001 From: Marcus Messer Date: Tue, 4 Aug 2026 17:15:26 +0100 Subject: [PATCH] Added chat and chat/health handler registration, dispatch, and testing --- lf_toolkit/chat/__init__.py | 12 +++- lf_toolkit/io/base_server.py | 17 +++++ lf_toolkit/io/handler.py | 30 +++++++- lf_toolkit/io/rpc_handler.py | 4 +- lf_toolkit/shared/command.py | 2 +- tests/io/base_server_test.py | 30 ++++++++ tests/io/handler_test.py | 130 +++++++++++++++++++++++++++++++++++ tests/io/rpc_handler_test.py | 71 +++++++++++++++++++ 8 files changed, 290 insertions(+), 6 deletions(-) create mode 100644 tests/io/base_server_test.py create mode 100644 tests/io/handler_test.py create mode 100644 tests/io/rpc_handler_test.py diff --git a/lf_toolkit/chat/__init__.py b/lf_toolkit/chat/__init__.py index 39026da..e2cae8a 100644 --- a/lf_toolkit/chat/__init__.py +++ b/lf_toolkit/chat/__init__.py @@ -1,7 +1,17 @@ +from ..shared.mued_api_v0_1_0 import ChatCapabilities +from ..shared.mued_api_v0_1_0 import ChatHealthResponse from ..shared.mued_api_v0_1_0 import ChatRequest from ..shared.mued_api_v0_1_0 import ChatResponse from ..shared.mued_api_v0_1_0 import Message from .params import ChatParams from .result import ChatResult -__all__ = ["ChatRequest", "ChatResponse", "Message", "ChatParams", "ChatResult"] +__all__ = [ + "ChatRequest", + "ChatResponse", + "ChatCapabilities", + "ChatHealthResponse", + "Message", + "ChatParams", + "ChatResult", +] diff --git a/lf_toolkit/io/base_server.py b/lf_toolkit/io/base_server.py index f84764e..bdb8245 100644 --- a/lf_toolkit/io/base_server.py +++ b/lf_toolkit/io/base_server.py @@ -7,6 +7,9 @@ from typing import Optional from typing import Union +from ..chat import ChatHealthResponse +from ..chat import ChatRequest +from ..chat import ChatResponse from ..evaluation import Result as EvaluationResult from ..preview import Result as PreviewResult from ..shared import Params @@ -22,6 +25,14 @@ [Any, Params], Union[PreviewResult, Awaitable[PreviewResult]] ] +ChatFunction = Callable[ + [ChatRequest], Union[ChatResponse, Awaitable[ChatResponse]] +] + +ChatHealthFunction = Callable[ + [], Union[ChatHealthResponse, Awaitable[ChatHealthResponse]] +] + class BaseServer(ABC): @@ -43,6 +54,12 @@ def eval(self, fn: EvaluationFunction): def preview(self, fn: PreviewFunction): return handler_decorator(self._handler, "preview", fn) + def chat(self, fn: ChatFunction): + return handler_decorator(self._handler, "chat", fn) + + def chat_health(self, fn: ChatHealthFunction): + return handler_decorator(self._handler, "chat/health", fn) + def handler_decorator(registry: Handler, name: str, fn): @wraps(fn) diff --git a/lf_toolkit/io/handler.py b/lf_toolkit/io/handler.py index 059ed04..e438e94 100644 --- a/lf_toolkit/io/handler.py +++ b/lf_toolkit/io/handler.py @@ -7,6 +7,9 @@ import anyio +from ..chat import ChatHealthResponse +from ..chat import ChatRequest +from ..chat import ChatResponse from ..evaluation import Result as EvaluationResult from ..shared import Command from ..shared import Params @@ -14,7 +17,10 @@ class Handler(ABC): - _handlers: Dict[str, Callable] = {} + _handlers: Dict[str, Callable] + + def __init__(self): + self._handlers = {} @abstractmethod async def dispatch(self, req: str) -> str: @@ -61,11 +67,29 @@ async def handle_healthcheck(self, req: dict): from .healthcheck import run_healthcheck return await anyio.to_thread.run_sync(run_healthcheck) + async def handle_chat(self, req: dict): + params = req["params"] + chat_request = ChatRequest.model_validate(params) + + result = await self._call_user_handler("chat", chat_request) + + if isinstance(result, ChatResponse): + return result.model_dump(mode="json", exclude_none=True) + + return result + + async def handle_chat_health(self, req: dict): + result = await self._call_user_handler("chat/health") + + if isinstance(result, ChatHealthResponse): + return result.model_dump(mode="json", exclude_none=True) + + return result + async def handle(self, name: Command, req: dict) -> dict: - handler = getattr(self, f"handle_{name}", None) + handler = getattr(self, f"handle_{name.replace('/', '_')}", None) if handler is None: raise ValueError(f"No handler for '{name}'") return await handler(req) - diff --git a/lf_toolkit/io/rpc_handler.py b/lf_toolkit/io/rpc_handler.py index 1eaff8a..bc47a71 100644 --- a/lf_toolkit/io/rpc_handler.py +++ b/lf_toolkit/io/rpc_handler.py @@ -11,8 +11,10 @@ class JsonRpcHandler(Handler): def __init__(self): + super().__init__() self._methods = { - name: jsonrpc_handler(self, name) for name in ["eval", "preview", "healthcheck"] + name: jsonrpc_handler(self, name) + for name in ["eval", "preview", "healthcheck", "chat", "chat/health"] } async def dispatch(self, req: str) -> str: diff --git a/lf_toolkit/shared/command.py b/lf_toolkit/shared/command.py index 6453a9d..bed6105 100644 --- a/lf_toolkit/shared/command.py +++ b/lf_toolkit/shared/command.py @@ -1,4 +1,4 @@ from typing import Literal -Command = Literal["eval", "preview"] +Command = Literal["eval", "preview", "chat", "chat/health"] diff --git a/tests/io/base_server_test.py b/tests/io/base_server_test.py new file mode 100644 index 0000000..86f3ee6 --- /dev/null +++ b/tests/io/base_server_test.py @@ -0,0 +1,30 @@ +from lf_toolkit.chat import ChatHealthResponse +from lf_toolkit.chat import ChatRequest +from lf_toolkit.chat import ChatResponse +from lf_toolkit.io.base_server import BaseServer + + +class ConcreteServer(BaseServer): + async def run(self): + pass + + +class TestBaseServerChatRegistration: + + def test_chat_registers_under_chat_name(self): + server = ConcreteServer() + + @server.chat + def chat_fn(request: ChatRequest) -> ChatResponse: + raise NotImplementedError + + assert server._handler._handlers["chat"] is chat_fn + + def test_chat_health_registers_under_chat_slash_health_name(self): + server = ConcreteServer() + + @server.chat_health + def chat_health_fn() -> ChatHealthResponse: + raise NotImplementedError + + assert server._handler._handlers["chat/health"] is chat_health_fn diff --git a/tests/io/handler_test.py b/tests/io/handler_test.py new file mode 100644 index 0000000..6dc48c8 --- /dev/null +++ b/tests/io/handler_test.py @@ -0,0 +1,130 @@ +import pytest + +from lf_toolkit.chat import ChatCapabilities +from lf_toolkit.chat import ChatHealthResponse +from lf_toolkit.chat import ChatRequest +from lf_toolkit.chat import ChatResponse +from lf_toolkit.chat import Message +from lf_toolkit.io.file_server import FileHandler +from lf_toolkit.shared.mued_api_v0_1_0 import DataPolicySupport +from lf_toolkit.shared.mued_api_v0_1_0 import HealthStatus +from lf_toolkit.shared.mued_api_v0_1_0 import Role + +pytest_plugins = ('pytest_asyncio',) + + +class TestHandleChat: + + @pytest.fixture + def handler(self): + return FileHandler() + + @pytest.mark.asyncio + async def test_calls_registered_handler_with_chat_request(self, handler): + received = {} + + def chat_fn(request: ChatRequest) -> ChatResponse: + received["request"] = request + return ChatResponse(output=Message(role=Role.ASSISTANT, content="hi there")) + + handler.register("chat", chat_fn) + + result = await handler.handle_chat({ + "params": {"messages": [{"role": "USER", "content": "hello"}]} + }) + + assert isinstance(received["request"], ChatRequest) + assert received["request"].messages[0].content == "hello" + assert result == {"output": {"role": "ASSISTANT", "content": "hi there"}} + + @pytest.mark.asyncio + async def test_passes_through_non_chat_response_result(self, handler): + handler.register("chat", lambda request: {"output": {"role": "ASSISTANT", "content": "raw"}}) + + result = await handler.handle_chat({ + "params": {"messages": [{"role": "USER", "content": "hello"}]} + }) + + assert result == {"output": {"role": "ASSISTANT", "content": "raw"}} + + @pytest.mark.asyncio + async def test_raises_when_no_handler_registered(self, handler): + with pytest.raises(ValueError, match="No user handler for 'chat'"): + await handler.handle_chat({ + "params": {"messages": [{"role": "USER", "content": "hello"}]} + }) + + +class TestHandleChatHealth: + + @pytest.fixture + def handler(self): + return FileHandler() + + @pytest.mark.asyncio + async def test_calls_registered_handler_with_no_arguments(self, handler): + def chat_health_fn() -> ChatHealthResponse: + return ChatHealthResponse( + status=HealthStatus.OK, + capabilities=ChatCapabilities( + supportsChat=True, + supportsDataPolicy=DataPolicySupport.NOT_SUPPORTED, + ), + ) + + handler.register("chat/health", chat_health_fn) + + result = await handler.handle_chat_health({"params": {}}) + + assert result == { + "status": "OK", + "capabilities": {"supportsChat": True, "supportsDataPolicy": "NOT_SUPPORTED"}, + } + + @pytest.mark.asyncio + async def test_passes_through_non_chat_health_response_result(self, handler): + handler.register("chat/health", lambda: {"status": "OK"}) + + result = await handler.handle_chat_health({"params": {}}) + + assert result == {"status": "OK"} + + +class TestHandleDispatch: + """Covers the name -> method lookup in Handler.handle, including the + 'chat/health' slash normalisation.""" + + @pytest.fixture + def handler(self): + return FileHandler() + + @pytest.mark.asyncio + async def test_dispatches_chat_to_handle_chat(self, handler): + handler.register("chat", lambda request: ChatResponse( + output=Message(role=Role.ASSISTANT, content="ok") + )) + + result = await handler.handle("chat", { + "params": {"messages": [{"role": "USER", "content": "hi"}]} + }) + + assert result == {"output": {"role": "ASSISTANT", "content": "ok"}} + + @pytest.mark.asyncio + async def test_dispatches_chat_slash_health_to_handle_chat_health(self, handler): + handler.register("chat/health", lambda: ChatHealthResponse( + status=HealthStatus.OK, + capabilities=ChatCapabilities( + supportsChat=True, + supportsDataPolicy=DataPolicySupport.NOT_SUPPORTED, + ), + )) + + result = await handler.handle("chat/health", {"params": {}}) + + assert result["status"] == "OK" + + @pytest.mark.asyncio + async def test_unknown_command_raises(self, handler): + with pytest.raises(ValueError, match="No handler for 'unknown'"): + await handler.handle("unknown", {"params": {}}) diff --git a/tests/io/rpc_handler_test.py b/tests/io/rpc_handler_test.py new file mode 100644 index 0000000..2f924a0 --- /dev/null +++ b/tests/io/rpc_handler_test.py @@ -0,0 +1,71 @@ +import ujson +import pytest + +from lf_toolkit.chat import ChatCapabilities +from lf_toolkit.chat import ChatHealthResponse +from lf_toolkit.chat import ChatRequest +from lf_toolkit.chat import ChatResponse +from lf_toolkit.chat import Message +from lf_toolkit.io.rpc_handler import JsonRpcHandler +from lf_toolkit.shared.mued_api_v0_1_0 import DataPolicySupport +from lf_toolkit.shared.mued_api_v0_1_0 import HealthStatus +from lf_toolkit.shared.mued_api_v0_1_0 import Role + +pytest_plugins = ('pytest_asyncio',) + + +class TestJsonRpcHandlerChat: + + @pytest.fixture + def handler(self): + return JsonRpcHandler() + + def test_registers_chat_methods(self, handler): + assert "chat" in handler._methods + assert "chat/health" in handler._methods + + @pytest.mark.asyncio + async def test_dispatch_chat_round_trip(self, handler): + def chat_fn(request: ChatRequest) -> ChatResponse: + last = request.messages[-1] + return ChatResponse(output=Message(role=Role.ASSISTANT, content=f"echo: {last.content}")) + + handler.register("chat", chat_fn) + + # go-ethereum's rpc.Client sends a single positional params array, + # not a params object -- this is the actual wire shape shimmy produces. + req = ujson.dumps({ + "jsonrpc": "2.0", + "method": "chat", + "params": [{"messages": [{"role": "USER", "content": "hi"}]}], + "id": 1, + }) + + response = ujson.loads(await handler.dispatch(req)) + + assert response["result"] == {"output": {"role": "ASSISTANT", "content": "echo: hi"}} + + @pytest.mark.asyncio + async def test_dispatch_chat_health_round_trip(self, handler): + def chat_health_fn() -> ChatHealthResponse: + return ChatHealthResponse( + status=HealthStatus.OK, + capabilities=ChatCapabilities( + supportsChat=True, + supportsDataPolicy=DataPolicySupport.NOT_SUPPORTED, + ), + ) + + handler.register("chat/health", chat_health_fn) + + req = ujson.dumps({ + "jsonrpc": "2.0", + "method": "chat/health", + "params": [{}], + "id": 2, + }) + + response = ujson.loads(await handler.dispatch(req)) + + assert response["result"]["status"] == "OK" + assert response["result"]["capabilities"]["supportsChat"] is True