Update code runner with mime type in stdout

Author: ajhai

llmstack/common/runner/proto/runner.proto

@@ -1,6 +1,22 @@
 syntax = "proto3";
 import "google/protobuf/struct.proto";
 
+enum ContentMimeType {
+  TEXT = 0; // text
+  JSON = 1; // json
+  HTML = 2; // html
+  PNG = 3; // png
+  JPEG = 4; // jpeg
+  SVG = 5; // svg
+  PDF = 6; // pdf
+  LATEX = 7; // latex
+}
+
+message Content {
+  ContentMimeType mime_type = 1; // mime type of content
+  bytes data = 2; // data of content
+}
+
 enum BrowserCommandType {
   GOTO = 0; // go to url
   TERMINATE = 1; // terminate session
@@ -123,7 +139,7 @@ message RestrictedPythonCodeRunnerRequest {
 message RestrictedPythonCodeRunnerResponse {
   RemoteBrowserState state = 1; // state of runner
   google.protobuf.Struct local_variables = 2; // local variable after running code
-  repeated string stdout = 3; // stdout of code
+  repeated Content stdout = 3; // stdout of code
   string stderr = 4; // stderr of code
   int32 exit_code = 5; // exit code of code
 }

llmstack/common/runner/proto/runner_pb2.py

@@ -10,11 +10,19 @@ COPY: BrowserCommandType
 DESCRIPTOR: _descriptor.FileDescriptor
 ENTER: BrowserCommandType
 GOTO: BrowserCommandType
+HTML: ContentMimeType
+JPEG: ContentMimeType
+JSON: ContentMimeType
+LATEX: ContentMimeType
+PDF: ContentMimeType
+PNG: ContentMimeType
 RUNNING: RemoteBrowserState
 SCROLL_X: BrowserCommandType
 SCROLL_Y: BrowserCommandType
+SVG: ContentMimeType
 TERMINATE: BrowserCommandType
 TERMINATED: RemoteBrowserState
+TEXT: ContentMimeType
 TIMEOUT: RemoteBrowserState
 TYPE: BrowserCommandType
 WAIT: BrowserCommandType
@@ -119,6 +127,14 @@ class BrowserTextAreaField(_message.Message):
     text: str
     def __init__(self, selector: _Optional[str] = ..., text: _Optional[str] = ...) -> None: ...
 
+class Content(_message.Message):
+    __slots__ = ["data", "mime_type"]
+    DATA_FIELD_NUMBER: _ClassVar[int]
+    MIME_TYPE_FIELD_NUMBER: _ClassVar[int]
+    data: bytes
+    mime_type: ContentMimeType
+    def __init__(self, mime_type: _Optional[_Union[ContentMimeType, str]] = ..., data: _Optional[bytes] = ...) -> None: ...
+
 class PlaywrightBrowserRequest(_message.Message):
     __slots__ = ["session_data", "steps", "stream_video", "timeout", "url"]
     SESSION_DATA_FIELD_NUMBER: _ClassVar[int]
@@ -200,8 +216,11 @@ class RestrictedPythonCodeRunnerResponse(_message.Message):
     local_variables: _struct_pb2.Struct
     state: RemoteBrowserState
     stderr: str
-    stdout: _containers.RepeatedScalarFieldContainer[str]
-    def __init__(self, state: _Optional[_Union[RemoteBrowserState, str]] = ..., local_variables: _Optional[_Union[_struct_pb2.Struct, _Mapping]] = ..., stdout: _Optional[_Iterable[str]] = ..., stderr: _Optional[str] = ..., exit_code: _Optional[int] = ...) -> None: ...
+    stdout: _containers.RepeatedCompositeFieldContainer[Content]
+    def __init__(self, state: _Optional[_Union[RemoteBrowserState, str]] = ..., local_variables: _Optional[_Union[_struct_pb2.Struct, _Mapping]] = ..., stdout: _Optional[_Iterable[_Union[Content, _Mapping]]] = ..., stderr: _Optional[str] = ..., exit_code: _Optional[int] = ...) -> None: ...
+
+class ContentMimeType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+    __slots__ = []
 
 class BrowserCommandType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     __slots__ = []

llmstack/common/runner/proto/runner_pb2.pyi

@@ -5,6 +5,7 @@
 import os
 import re
 import subprocess
+import time
 from concurrent import futures
 from typing import Iterator
 
@@ -20,6 +21,8 @@
 from llmstack.common.runner.playwright.browser import Playwright
 from llmstack.common.runner.proto.runner_pb2 import (
     TERMINATE,
+    Content,
+    ContentMimeType,
     PlaywrightBrowserRequest,
     PlaywrightBrowserResponse,
     RemoteBrowserRequest,
@@ -194,13 +197,13 @@ def GetRestrictedPythonCodeRunner(
     ) -> Iterator[RestrictedPythonCodeRunnerResponse]:
         from google.protobuf.json_format import MessageToDict, ParseDict
         from google.protobuf.struct_pb2 import Struct, Value
-        from RestrictedPython import compile_restricted, safe_builtins
+        from RestrictedPython import compile_restricted
         from RestrictedPython.Guards import (
             guarded_iter_unpack_sequence,
             guarded_unpack_sequence,
             safe_builtins,
         )
-        from RestrictedPython.PrintCollector import PrintCollector
+        from RestrictedPython.transformer import IOPERATOR_TO_STR
 
         class CustomPrint(object):
             def __init__(self):
@@ -211,7 +214,12 @@ def write(self, text):
                 if self.enabled:
                     if text and text.strip():
                         log_line = "{0}".format(text)
-                        self.lines.append(log_line)
+                        self.lines.append(
+                            (
+                                time.time(),
+                                Content(data=bytes(log_line.encode("utf-8")), mime_type=ContentMimeType.TEXT),
+                            )
+                        )
 
             def enable(self):
                 self.enabled = True
@@ -256,7 +264,6 @@ async def execute_restricted_code(source_code, input_data={}):
             mathplot_lib_display = []
 
             def custom_pyplot_show():
-                import base64
                 import io
 
                 from matplotlib.backends.backend_agg import (
@@ -284,17 +291,10 @@ def custom_pyplot_show():
                 # Rewind the buffer to start
                 buf.seek(0)
 
-                # Encode the buffer contents to base64
-                encoded = base64.b64encode(buf.read())
-
-                # Convert to string for a cleaner return format
-                encoded_string = encoded.decode("utf-8")
+                mathplot_lib_display.append((time.time(), Content(data=buf.read(), mime_type=ContentMimeType.PNG)))
 
                 # Cleanup by closing the buffer
                 buf.close()
-                mathplot_lib_display.append(f"data:image/png;base64,{encoded_string}")
-                # Return the base64 string
-                return encoded_string
 
             def custom_import(name, globals=None, locals=None, fromlist=(), level=0):
                 module = __import__(name, globals, locals, fromlist, level)
@@ -341,7 +341,11 @@ def custom_import(name, globals=None, locals=None, fromlist=(), level=0):
                 if isinstance(v, (int, float, str, bool, list, dict, tuple, type(None), Value, Struct))
             }
             # Return the result and any printed output
-            return {**local_variables, **{"display": mathplot_lib_display}}, custom_print.lines, errors
+            return (
+                local_variables,
+                [x[1] for x in sorted(custom_print.lines + mathplot_lib_display, key=lambda x: x[0])],
+                errors,
+            )
 
         yield RestrictedPythonCodeRunnerResponse(state=RemoteBrowserState.RUNNING)
 

llmstack/common/runner/server.py

@@ -0,0 +1,22 @@
+from enum import Enum
+
+from pydantic import BaseModel, Field
+
+
+class ContentMimeType(str, Enum):
+    TEXT = "text/plain"
+    JSON = "application/json"
+    HTML = "text/html"
+    PNG = "image/png"
+    JPEG = "image/jpeg"
+    SVG = "image/svg+xml"
+    PDF = "application/pdf"
+    LATEX = "application/x-latex"
+
+    def __str__(self):
+        return self.value
+
+
+class Content(BaseModel):
+    data: str = Field(description="The content data")
+    mime_type: ContentMimeType = Field(description="The content mime type", default=ContentMimeType.TEXT)