lmnr 0.7.11__py3-none-any.whl → 0.7.12__py3-none-any.whl

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

@@ -0,0 +1,100 @@
+"""OpenTelemetry CUA instrumentation"""
+
+import logging
+from typing import Any, AsyncGenerator, Collection
+
+from lmnr.opentelemetry_lib.decorators import json_dumps
+from lmnr import Laminar
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import unwrap
+
+from opentelemetry.trace import Span
+from opentelemetry.trace.status import Status, StatusCode
+from wrapt import wrap_function_wrapper
+
+logger = logging.getLogger(__name__)
+
+_instruments = ("cua-agent >= 0.4.0",)
+
+
+def _wrap_run(
+    wrapped,
+    instance,
+    args,
+    kwargs,
+):
+    parent_span = Laminar.start_span("ComputerAgent.run")
+    instance._lmnr_parent_span = parent_span
+
+    try:
+        result: AsyncGenerator[dict[str, Any], None] = wrapped(*args, **kwargs)
+        return _abuild_from_streaming_response(parent_span, result)
+    except Exception as e:
+        if parent_span.is_recording():
+            parent_span.set_status(Status(StatusCode.ERROR))
+            parent_span.record_exception(e)
+            parent_span.end()
+        raise
+
+
+async def _abuild_from_streaming_response(
+    parent_span: Span, response: AsyncGenerator[dict[str, Any], None]
+) -> AsyncGenerator[dict[str, Any], None]:
+    with Laminar.use_span(parent_span, end_on_exit=True):
+        response_iter = aiter(response)
+        while True:
+            step = None
+            step_span = Laminar.start_span("ComputerAgent.step")
+            with Laminar.use_span(step_span):
+                try:
+                    step = await anext(response_iter)
+                    step_span.set_attribute("lmnr.span.output", json_dumps(step))
+                    try:
+                        # When processing tool calls, each output item is processed separately,
+                        # if the output is message, agent.step returns an empty array
+                        # https://github.com/trycua/cua/blob/17d670962970a1d1774daaec029ebf92f1f9235e/libs/python/agent/agent/agent.py#L459
+                        if len(step.get("output", [])) == 0:
+                            continue
+                    except Exception:
+                        pass
+                    if step_span.is_recording():
+                        step_span.end()
+                except StopAsyncIteration:
+                    # don't end on purpose, there is no iteration step here.
+                    break
+
+            if step is not None:
+                yield step
+
+
+class CuaAgentInstrumentor(BaseInstrumentor):
+    def __init__(self):
+        super().__init__()
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):
+        wrap_package = "agent.agent"
+        wrap_object = "ComputerAgent"
+        wrap_method = "run"
+        try:
+            wrap_function_wrapper(
+                wrap_package,
+                f"{wrap_object}.{wrap_method}",
+                _wrap_run,
+            )
+        except ModuleNotFoundError:
+            pass  # that's ok, we don't want to fail if some methods do not exist
+
+    def _uninstrument(self, **kwargs):
+        wrap_package = "agent.agent"
+        wrap_object = "ComputerAgent"
+        wrap_method = "run"
+        try:
+            unwrap(
+                f"{wrap_package}.{wrap_object}",
+                wrap_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/__init__.py

@@ -0,0 +1,477 @@
+"""OpenTelemetry CUA instrumentation"""
+
+import logging
+from typing import Collection
+
+from lmnr.opentelemetry_lib.decorators import json_dumps
+from lmnr.sdk.utils import get_input_from_func_args
+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>"

lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/shared/__init__.py

@@ -144,6 +144,11 @@ def _set_request_attributes(span, kwargs, instance=None):
     _set_span_attribute(
         span, SpanAttributes.LLM_IS_STREAMING, kwargs.get("stream") or False
     )
+    _set_span_attribute(
+        span,
+        SpanAttributes.LLM_REQUEST_REASONING_EFFORT,
+        kwargs.get("reasoning_effort"),
+    )
     if response_format := kwargs.get("response_format"):
         # backward-compatible check for
         # openai.types.shared_params.response_format_json_schema.ResponseFormatJSONSchema
@@ -263,6 +268,15 @@ def _set_response_attributes(span, response):
         SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS,
         prompt_tokens_details.get("cached_tokens", 0),
     )
+
+    if completion_token_details := dict(usage.get("completion_tokens_details", {})):
+        reasoning_tokens = completion_token_details.get("reasoning_tokens")
+        _set_span_attribute(
+            span,
+            SpanAttributes.LLM_USAGE_REASONING_TOKENS,
+            reasoning_tokens or 0,
+        )
+
     return
 
 

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

@@ -5,6 +5,7 @@ import threading
 import traceback
 from contextlib import asynccontextmanager
 from importlib.metadata import version
+from packaging.version import parse
 
 from opentelemetry import context as context_api
 from opentelemetry._events import EventLogger
@@ -19,7 +20,15 @@ LMNR_TRACE_CONTENT = "LMNR_TRACE_CONTENT"
 
 
 def is_openai_v1():
-    return _OPENAI_VERSION >= "1.0.0"
+    return parse(_OPENAI_VERSION) >= parse("1.0.0")
+
+
+def is_reasoning_supported():
+    # Reasoning has been introduced in OpenAI API on Dec 17, 2024
+    #     as per https://platform.openai.com/docs/changelog.
+    # The updated OpenAI library version is 1.58.0
+    #     as per https://pypi.org/project/openai/.
+    return parse(_OPENAI_VERSION) >= parse("1.58.0")
 
 
 def is_azure_openai(instance):