Bug/ipc

lf_toolkit/evaluation/image_upload.py

@@ -101,7 +101,7 @@ def get_aws_signed_request(full_url, buffer, mime_type):
     aws_request = AWSRequest(
         method='PUT',
         url=full_url,
-        data=buffer,
+        data=data,
         headers=headers
     )
 

lf_toolkit/io/rpc_handler.py

@@ -4,6 +4,7 @@
 from jsonrpcserver import Success
 from jsonrpcserver import async_dispatch
 
+from ..shared import Command
 from .handler import Handler
 
 
@@ -23,10 +24,10 @@ async def dispatch(self, req: str) -> str:
         )
 
 
-def jsonrpc_handler(handler: Handler, name: str):
+def jsonrpc_handler(handler: Handler, name: Command):
     async def wrapped(req: dict):
         try:
-            result = await handler.handle(name, req)
+            result = await handler.handle(name, {"params": req})
             return Success(result)
         except Exception as e:
             return Error(0, str(e), e)

lf_toolkit/io/stdio_server.py

@@ -2,6 +2,7 @@
 
 from typing import Optional
 
+import anyio
 from anyio.streams.file import FileReadStream
 from anyio.streams.file import FileWriteStream
 from anyio.streams.stapled import StapledByteStream
@@ -15,17 +16,18 @@
 class StdioClient(StreamIO):
 
     def __init__(self):
+        self._stdout_buffer = sys.stdout.buffer
         self.stream = StapledByteStream(
-            FileWriteStream(sys.stdout),
-            FileReadStream(sys.stdin),
+            FileWriteStream(self._stdout_buffer),
+            FileReadStream(sys.stdin.buffer),
         )
 
     async def read(self, size: int) -> bytes:
         return await self.stream.receive(size)
 
     async def write(self, data: bytes):
         await self.stream.send(data)
-        await self.stream.flush()
+        await anyio.to_thread.run_sync(self._stdout_buffer.flush)
 
     async def close(self):
         await self.stream.aclose()
@@ -37,10 +39,11 @@ class StdioServer(StreamServer):
 
     def __init__(self, handler: Optional[Handler] = None):
         super().__init__(handler)
-        self._client = StdioClient()
 
     def wrap_io(self, client: StreamIO) -> StreamIO:
         return PrefixStreamIO(client)
 
     async def run(self):
+        print("StdioServer started", file=sys.stderr, flush=True)
+        self._client = StdioClient()
         await self._handle_client(self._client)

lf_toolkit/io/stream_io.py

@@ -63,10 +63,11 @@ async def read(self, size: int) -> bytes:
         if content_length == 0:
             raise ValueError("Content-Length header not found or is zero")
 
-        if content_length > size:
-            raise ValueError("Content-Length is larger than the read size")
-
-        return await self.base.read(content_length)
+        data = b""
+        while len(data) < content_length:
+            chunk = await self.base.read(content_length - len(data))
+            data += chunk
+        return data
 
     async def write(self, data: bytes):
         response_headers_str = f"Content-Length: {len(data)}\r\n\r\n"
@@ -84,22 +85,27 @@ async def _handle_client(self, client: StreamIO):
 
         while True:
             try:
+                import sys
+                print("waiting for data...", file=sys.stderr, flush=True)
                 data = await io.read(4096)
+                print(f"got data: {data[:80]}", file=sys.stderr, flush=True)
 
                 if not data:
-                    # print("Received empty data")
                     break
 
+                print("dispatching...", file=sys.stderr, flush=True)
                 response = await self.dispatch(data.decode("utf-8"))
+                print(f"got response: {str(response)[:80]}", file=sys.stderr, flush=True)
 
                 await io.write(response.encode("utf-8"))
+                print("wrote response", file=sys.stderr, flush=True)
             except anyio.EndOfStream:
-                # print("Client disconnected")
                 break
             except anyio.ClosedResourceError:
-                # print("Client disconnected")
                 break
             except Exception as e:
-                print(f"Exception: {e}")
-            finally:
-                await client.close()
+                import traceback
+                traceback.print_exc(file=sys.stderr)
+                break
+
+        await client.close()

pyproject.toml

@@ -59,6 +59,7 @@ pytest-cov = "5.0.0"
 ##########################
 # extras
 ##########################
+pytest-anyio = "^0.0.0"
 
 [tool.poetry.extras]
 parsing = ["antlr4-python3-runtime", "lark", "latex2sympy"]