lmnr 0.4.53.dev0__py3-none-any.whl → 0.7.26__py3-none-any.whl

lmnr/opentelemetry_lib/opentelemetry/instrumentation/cua_computer/__init__.py

@@ -0,0 +1,476 @@
+"""OpenTelemetry CUA instrumentation"""
+
+import logging
+from typing import Collection
+
+from lmnr.sdk.utils import get_input_from_func_args, json_dumps
+from lmnr import Laminar
+from lmnr.opentelemetry_lib.tracing.context import get_current_context
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import unwrap
+
+from opentelemetry import trace
+from opentelemetry.trace import Span
+from opentelemetry.trace.status import Status, StatusCode
+from wrapt import wrap_function_wrapper
+
+from .utils import payload_to_placeholder
+
+logger = logging.getLogger(__name__)
+
+_instruments = ("cua-computer >= 0.4.0",)
+
+
+WRAPPED_METHODS = [
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "close",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "force_close",
+    },
+]
+WRAPPED_AMETHODS = [
+    {
+        "package": "computer.computer",
+        "object": "Computer",
+        "method": "__aenter__",
+        "action": "start_parent_span",
+    },
+    {
+        "package": "computer.computer",
+        "object": "Computer",
+        "method": "__aexit__",
+        "action": "end_parent_span",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "mouse_down",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "mouse_up",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "left_click",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "right_click",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "double_click",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "move_cursor",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "drag_to",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "drag",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "key_down",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "key_up",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "type_text",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "press",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "hotkey",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "scroll",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "scroll_down",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "scroll_up",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "screenshot",
+        "output_formatter": payload_to_placeholder,
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "get_screen_size",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "get_cursor_position",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "copy_to_clipboard",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "set_clipboard",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "file_exists",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "directory_exists",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "list_dir",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "read_text",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "write_text",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "read_bytes",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "write_bytes",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "delete_file",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "create_dir",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "delete_dir",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "get_file_size",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "run_command",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "get_accessibility_tree",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "to_screen_coordinates",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "get_active_window_bounds",
+    },
+    {
+        "package": "computer.interface.generic",
+        "object": "GenericComputerInterface",
+        "method": "to_screenshot_coordinates",
+    },
+]
+
+
+def _with_wrapper(func):
+    """Helper for providing tracer for wrapper functions. Includes metric collectors."""
+
+    def wrapper(
+        to_wrap,
+    ):
+        def wrapper(wrapped, instance, args, kwargs):
+            return func(
+                to_wrap,
+                wrapped,
+                instance,
+                args,
+                kwargs,
+            )
+
+        return wrapper
+
+    return wrapper
+
+
+def add_input_to_parent_span(span, instance):
+    # api_key is skipped on purpose
+    params = {}
+    if hasattr(instance, "display"):
+        params["display"] = instance.display
+    if hasattr(instance, "memory"):
+        params["memory"] = instance.memory
+    if hasattr(instance, "cpu"):
+        params["cpu"] = instance.cpu
+    if hasattr(instance, "os_type"):
+        params["os_type"] = instance.os_type
+    if hasattr(instance, "name"):
+        params["name"] = instance.name
+    if hasattr(instance, "image"):
+        params["image"] = instance.image
+    if hasattr(instance, "shared_directories"):
+        params["shared_directories"] = instance.shared_directories
+    if hasattr(instance, "use_host_computer_server"):
+        params["use_host_computer_server"] = instance.use_host_computer_server
+    if hasattr(instance, "verbosity"):
+        if (
+            isinstance(instance.verbosity, int)
+            and instance.verbosity in logging._levelToName
+        ):
+            params["verbosity"] = logging._levelToName[instance.verbosity]
+        else:
+            params["verbosity"] = instance.verbosity
+    if hasattr(instance, "telemetry_enabled"):
+        params["telemetry_enabled"] = instance.telemetry_enabled
+    if hasattr(instance, "provider_type"):
+        params["provider_type"] = instance.provider_type
+    if hasattr(instance, "port"):
+        params["port"] = instance.port
+    if hasattr(instance, "noVNC_port"):
+        params["noVNC_port"] = instance.noVNC_port
+    if hasattr(instance, "host"):
+        params["host"] = instance.host
+    if hasattr(instance, "storage"):
+        params["storage"] = instance.storage
+    if hasattr(instance, "ephemeral"):
+        params["ephemeral"] = instance.ephemeral
+    if hasattr(instance, "experiments"):
+        params["experiments"] = instance.experiments
+    span.set_attribute("lmnr.span.input", json_dumps(params))
+
+
+@_with_wrapper
+def _wrap(
+    to_wrap,
+    wrapped,
+    instance,
+    args,
+    kwargs,
+):
+    if to_wrap.get("action") == "start_parent_span":
+        parent_span = Laminar.start_span("computer.run")
+        add_input_to_parent_span(parent_span, instance)
+        result = wrapped(*args, **kwargs)
+        try:
+            instance._interface._lmnr_parent_span = parent_span
+        except Exception:
+            pass
+        return result
+    elif to_wrap.get("action") == "end_parent_span":
+        result = wrapped(*args, **kwargs)
+        try:
+            parent_span: Span = instance._interface._lmnr_parent_span
+            if parent_span and parent_span.is_recording():
+                parent_span.end()
+        except Exception:
+            pass
+        return result
+
+    # if there's no parent span, use
+    parent_span = trace.get_current_span(context=get_current_context())
+    try:
+        if instance._lmnr_parent_span:
+            parent_span: Span = instance._lmnr_parent_span
+    except Exception:
+        pass
+
+    with Laminar.use_span(parent_span):
+        instance_name = "interface"
+        with Laminar.start_as_current_span(
+            f"{instance_name}.{to_wrap.get('method')}", span_type="TOOL"
+        ) as span:
+            span.set_attribute(
+                "lmnr.span.input",
+                json_dumps(get_input_from_func_args(wrapped, True, args, kwargs)),
+            )
+            try:
+                result = wrapped(*args, **kwargs)
+            except Exception as e:  # pylint: disable=broad-except
+                span.set_status(Status(StatusCode.ERROR))
+                span.record_exception(e)
+                span.end()
+                raise
+            output_formatter = to_wrap.get("output_formatter") or (
+                lambda x: json_dumps(x)
+            )
+            span.set_attribute("lmnr.span.output", output_formatter(result))
+            return result
+
+
+@_with_wrapper
+async def _wrap_async(
+    to_wrap,
+    wrapped,
+    instance,
+    args,
+    kwargs,
+):
+    if to_wrap.get("action") == "start_parent_span":
+        parent_span = Laminar.start_span("computer.run")
+        add_input_to_parent_span(parent_span, instance)
+        result = await wrapped(*args, **kwargs)
+        try:
+            instance._interface._lmnr_parent_span = parent_span
+        except Exception:
+            pass
+        return result
+    elif to_wrap.get("action") == "end_parent_span":
+        result = await wrapped(*args, **kwargs)
+        try:
+            parent_span: Span = instance._interface._lmnr_parent_span
+            if parent_span and parent_span.is_recording():
+                parent_span.end()
+        except Exception:
+            pass
+        return result
+
+    # if there's no parent span, use
+    parent_span = trace.get_current_span(context=get_current_context())
+    try:
+        parent_span: Span = instance._lmnr_parent_span
+    except Exception:
+        pass
+
+    with Laminar.use_span(parent_span):
+        instance_name = "interface"
+        with Laminar.start_as_current_span(
+            f"{instance_name}.{to_wrap.get('method')}",
+            span_type="TOOL",
+        ) as span:
+            span.set_attribute(
+                "lmnr.span.input",
+                json_dumps(get_input_from_func_args(wrapped, True, args, kwargs)),
+            )
+            try:
+                result = await wrapped(*args, **kwargs)
+            except Exception as e:  # pylint: disable=broad-except
+                span.set_status(Status(StatusCode.ERROR))
+                span.record_exception(e)
+                span.end()
+                raise
+            output_formatter = to_wrap.get("output_formatter") or (
+                lambda x: json_dumps(x)
+            )
+            span.set_attribute("lmnr.span.output", output_formatter(result))
+            return result
+
+
+class CuaComputerInstrumentor(BaseInstrumentor):
+    def __init__(self):
+        super().__init__()
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):
+        for wrapped_method in WRAPPED_METHODS:
+            wrap_package = wrapped_method.get("package")
+            wrap_object = wrapped_method.get("object")
+            wrap_method = wrapped_method.get("method")
+
+            try:
+                wrap_function_wrapper(
+                    wrap_package,
+                    f"{wrap_object}.{wrap_method}",
+                    _wrap(wrapped_method),
+                )
+            except ModuleNotFoundError:
+                pass  # that's ok, we don't want to fail if some methods do not exist
+
+        for wrapped_method in WRAPPED_AMETHODS:
+            wrap_package = wrapped_method.get("package")
+            wrap_object = wrapped_method.get("object")
+            wrap_method = wrapped_method.get("method")
+            try:
+                wrap_function_wrapper(
+                    wrap_package,
+                    f"{wrap_object}.{wrap_method}",
+                    _wrap_async(wrapped_method),
+                )
+            except ModuleNotFoundError:
+                pass  # that's ok, we don't want to fail if some methods do not exist
+
+    def _uninstrument(self, **kwargs):
+        for wrapped_method in WRAPPED_METHODS:
+            wrap_package = wrapped_method.get("package")
+            wrap_object = wrapped_method.get("object")
+            try:
+                unwrap(
+                    f"{wrap_package}.{wrap_object}",
+                    wrapped_method.get("method"),
+                )
+            except ModuleNotFoundError:
+                pass  # that's ok, we don't want to fail if some methods do not exist
+
+        for wrapped_method in WRAPPED_AMETHODS:
+            wrap_package = wrapped_method.get("package")
+            wrap_object = wrapped_method.get("object")
+            try:
+                unwrap(
+                    f"{wrap_package}.{wrap_object}",
+                    wrapped_method.get("method"),
+                )
+            except ModuleNotFoundError:
+                pass  # that's ok, we don't want to fail if some methods do not exist

lmnr/opentelemetry_lib/opentelemetry/instrumentation/cua_computer/utils.py

@@ -0,0 +1,12 @@
+import base64
+import orjson
+
+
+def payload_to_base64url(payload_bytes: bytes) -> bytes:
+    data = base64.b64encode(payload_bytes).decode("utf-8")
+    url = f"data:image/png;base64,{data}"
+    return orjson.dumps({"base64url": url})
+
+
+def payload_to_placeholder(payload_bytes: bytes) -> str:
+    return "<BINARY_BLOB_SCREENSHOT>"