Added chat and chat/health handler registration, dispatch, and testing (#9)

Author: m-messer

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",
+]

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)

lf_toolkit/io/handler.py

@@ -7,14 +7,20 @@
 
 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
 
 
 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)
-

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:

lf_toolkit/shared/command.py

@@ -1,4 +1,4 @@
 from typing import Literal
 
 
-Command = Literal["eval", "preview"]
+Command = Literal["eval", "preview", "chat", "chat/health"]

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

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": {}})

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