lmnr 0.6.16__py3-none-any.whl → 0.7.26__py3-none-any.whl

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

@@ -0,0 +1 @@
+__version__ = "0.40.14"

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

@@ -0,0 +1,388 @@
+"""OpenTelemetry OpenHands AI instrumentation"""
+
+import sys
+from typing import Collection
+
+from lmnr.opentelemetry_lib.tracing.attributes import (
+    ASSOCIATION_PROPERTIES,
+    SESSION_ID,
+    USER_ID,
+)
+from lmnr.opentelemetry_lib.utils.wrappers import _with_tracer_wrapper
+from lmnr.sdk.log import get_default_logger
+from lmnr.sdk.utils import get_input_from_func_args, json_dumps
+from lmnr import Laminar
+from lmnr.version import __version__
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import unwrap
+from opentelemetry.trace import get_tracer, Tracer
+from wrapt import wrap_function_wrapper
+
+logger = get_default_logger(__name__)
+
+_instruments = ("openhands-ai >= 0.9.0", "openhands-aci >= 0.1.0")
+parent_spans = {}
+
+
+def is_message_action(event) -> bool:
+    """Check if event has action attribute equal to 'message'."""
+    return event and hasattr(event, "action") and event.action == "message"
+
+
+def is_user_message(event) -> bool:
+    """Check if event is a message action from user source."""
+    return (
+        is_message_action(event) and hasattr(event, "source") and event.source == "user"
+    )
+
+
+def is_agent_message(event) -> bool:
+    """Check if event is a message action from agent source."""
+    return (
+        is_message_action(event)
+        and hasattr(event, "source")
+        and event.source == "agent"
+    )
+
+
+def is_agent_state_changed_to(event, state: str) -> bool:
+    """Check if event is an agent_state_changed observation with specific state."""
+    return (
+        event
+        and hasattr(event, "observation")
+        and event.observation == "agent_state_changed"
+        and hasattr(event, "agent_state")
+        and event.agent_state == state
+    )
+
+
+def get_handle_action_action(event) -> str:
+    """Get the action of the handle_action event."""
+    if event and hasattr(event, "action"):
+        try:
+            return event.action.value
+        except Exception:
+            return event.action
+    return None
+
+
+WRAPPED_METHODS = [
+    {
+        "package": "openhands.agenthub.codeact_agent.codeact_agent",
+        "object": "CodeActAgent",
+        "methods": [
+            {"method": "step"},
+            {"method": "response_to_actions"},
+        ],
+    },
+    {
+        "package": "openhands.controller.agent_controller",
+        "object": "AgentController",
+        "methods": [
+            {"method": "_step", "async": True},
+            {
+                "method": "_handle_action",
+                "async": True,
+                "span_type": "TOOL",
+            },
+            {"method": "_handle_observation", "async": True},
+            {"method": "_handle_message_action", "async": True},
+            {"method": "on_event"},
+            {"method": "save_state"},
+            {"method": "get_trajectory"},
+            {"method": "start_delegate"},
+            {"method": "end_delegate"},
+            {"method": "_is_stuck"},
+        ],
+    },
+]
+
+
+@_with_tracer_wrapper
+def _wrap_on_event(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
+    """Wrapper for on_event."""
+    controller_id = instance.id
+    user_id = instance.user_id
+    event = kwargs.get("event", args[0] if len(args) > 0 else None)
+    start_event = False
+    finish_event = False
+    user_message = ""
+    agent_message = ""
+    span_name = to_wrap.get("span_name")
+    span_type = to_wrap.get("span_type", "DEFAULT")
+    if event and hasattr(event, "action") and event.action == "system":
+        return wrapped(*args, **kwargs)
+
+    event_type = None
+    subtype = None
+    if event and hasattr(event, "action"):
+        event_type = "action"
+        try:
+            subtype = event.action.value
+        except Exception:
+            subtype = event.action
+    elif event and hasattr(event, "observation"):
+        event_type = "observation"
+        try:
+            subtype = event.observation.value
+        except Exception:
+            subtype = event.observation
+    if event_type and subtype:
+        span_name = f"event.{event_type}.{subtype}"
+        span_type = "EVENT"
+
+    # start trace on user message
+    if is_user_message(event):
+        user_message = event.content if hasattr(event, "content") else ""
+        start_event = True
+    # end trace on agent state change to finished or error
+    if is_agent_state_changed_to(event, "stopped") or is_agent_state_changed_to(
+        event, "awaiting_user_input"
+    ):
+        finish_event = True
+
+    if is_agent_state_changed_to(event, "user_rejected"):
+        agent_message = "<user_rejected>"
+
+    if is_agent_message(event):
+        agent_message = event.content if hasattr(event, "content") else ""
+
+    if start_event:
+        if controller_id in parent_spans:
+            logger.debug(
+                "Received a message, but already have a span for this trace. Resetting span."
+            )
+            parent_spans[controller_id].end()
+            del parent_spans[controller_id]
+        parent_span = Laminar.start_span("conversation.turn", span_type="DEFAULT")
+        if user_id:
+            parent_span.set_attribute(f"{ASSOCIATION_PROPERTIES}.{USER_ID}", user_id)
+        if user_message:
+            parent_span.set_attribute("lmnr.span.input", user_message)
+        parent_span.set_attribute(
+            f"{ASSOCIATION_PROPERTIES}.{SESSION_ID}", controller_id
+        )
+        parent_spans[controller_id] = parent_span
+
+    if controller_id in parent_spans:
+        with Laminar.use_span(parent_spans[controller_id]):
+            result = _wrap_sync_method_inner(
+                tracer,
+                {**to_wrap, "span_name": span_name, "span_type": span_type},
+                wrapped,
+                instance,
+                args,
+                kwargs,
+            )
+            if agent_message:
+                parent_spans[controller_id].set_attribute(
+                    "lmnr.span.output", agent_message
+                )
+            if finish_event:
+                parent_spans[controller_id].end()
+                del parent_spans[controller_id]
+            return result
+
+    return wrapped(*args, **kwargs)
+
+
+@_with_tracer_wrapper
+async def _wrap_handle_action(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
+    """Wrapper for on_event."""
+    event = kwargs.get("event", args[0] if len(args) > 0 else None)
+    if event and hasattr(event, "action"):
+        if event.action == "system":
+            return await wrapped(*args, **kwargs)
+    action_name = get_handle_action_action(event)
+    if action_name and action_name != "message":
+        to_wrap["span_name"] = f"action.{action_name}"
+    controller_id = instance.id
+    if controller_id not in parent_spans:
+        return await wrapped(*args, **kwargs)
+    return await _wrap_async_method_inner(
+        tracer, to_wrap, wrapped, instance, args, kwargs
+    )
+
+
+def _wrap_sync_method_inner(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
+    """Wrapper for synchronous methods."""
+    span_name = to_wrap.get("span_name")
+
+    with Laminar.start_as_current_span(
+        span_name,
+        span_type=to_wrap.get("span_type", "DEFAULT"),
+        input=json_dumps(
+            get_input_from_func_args(
+                wrapped, to_wrap.get("object") is not None, args, kwargs
+            )
+        ),
+    ) as span:
+        try:
+            result = wrapped(*args, **kwargs)
+
+            # Capture output
+            if not to_wrap.get("ignore_output"):
+                span.set_attribute("lmnr.span.output", json_dumps(result))
+            return result
+
+        except Exception as e:
+            span.record_exception(e)
+            raise
+
+
+@_with_tracer_wrapper
+def _wrap_sync_method(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
+    instance_id = None
+    if to_wrap.get("object") == "AgentController":
+        instance_id = instance.id
+    if to_wrap.get("object") == "ActionExecutionClient" and hasattr(instance, "sid"):
+        instance_id = instance.sid
+    if instance_id is not None and instance_id not in parent_spans:
+        return wrapped(*args, **kwargs)
+    return _wrap_sync_method_inner(tracer, to_wrap, wrapped, instance, args, kwargs)
+
+
+async def _wrap_async_method_inner(
+    tracer: Tracer, to_wrap, wrapped, instance, args, kwargs
+):
+    """Wrapper for asynchronous methods."""
+    span_name = to_wrap.get("span_name")
+    instance_id = None
+    if to_wrap.get("object") == "AgentController":
+        instance_id = instance.id
+    if to_wrap.get("object") == "ActionExecutionClient" and hasattr(instance, "sid"):
+        instance_id = instance.sid
+    if instance_id is not None and instance_id not in parent_spans:
+        return await wrapped(*args, **kwargs)
+
+    with Laminar.start_as_current_span(
+        span_name,
+        span_type=to_wrap.get("span_type", "DEFAULT"),
+        input=json_dumps(
+            get_input_from_func_args(
+                wrapped, to_wrap.get("object") is not None, args, kwargs
+            )
+        ),
+    ) as span:
+        try:
+            result = await wrapped(*args, **kwargs)
+
+            # Capture output
+            if not to_wrap.get("ignore_output"):
+                span.set_attribute("lmnr.span.output", json_dumps(result))
+            return result
+
+        except Exception as e:
+            span.record_exception(e)
+            raise
+
+
+@_with_tracer_wrapper
+async def _wrap_async_method(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
+    """Wrapper for asynchronous methods."""
+    return await _wrap_async_method_inner(
+        tracer, to_wrap, wrapped, instance, args, kwargs
+    )
+
+
+class OpenHandsInstrumentor(BaseInstrumentor):
+    """An instrumentor for OpenHands AI."""
+
+    def __init__(self):
+        super().__init__()
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):
+        """Instrument OpenHands AI methods."""
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, __version__, tracer_provider)
+
+        for wrapped_config in WRAPPED_METHODS:
+            wrap_package = wrapped_config.get("package")
+
+            wrap_object = wrapped_config.get("object")
+            methods = wrapped_config.get("methods", [])
+
+            for method_config in methods:
+
+                wrap_method = method_config.get("method")
+                async_wrap = method_config.get("async", False)
+                windows_only = method_config.get("windows_only", False)
+                if windows_only and sys.platform != "win32":
+                    continue
+
+                # Create the method configuration for the wrapper
+                method_wrapper_config = {
+                    "package": wrap_package,
+                    "object": wrap_object,
+                    "method": wrap_method,
+                    "span_name": method_config.get(
+                        "span_name",
+                        f"{wrap_object}.{wrap_method}" if wrap_object else wrap_method,
+                    ),
+                    "span_type": method_config.get("span_type", "DEFAULT"),
+                    "async": async_wrap,
+                }
+
+                # Determine the target for wrapping
+                if wrap_object:
+                    target = f"{wrap_object}.{wrap_method}"
+                else:
+                    target = wrap_method
+
+                if wrap_object == "AgentController" and wrap_method == "on_event":
+                    wrap_function_wrapper(
+                        wrap_package,
+                        target,
+                        _wrap_on_event(tracer, method_wrapper_config),
+                    )
+                    continue
+                if wrap_object == "AgentController" and wrap_method == "_handle_action":
+                    wrap_function_wrapper(
+                        wrap_package,
+                        target,
+                        _wrap_handle_action(tracer, method_wrapper_config),
+                    )
+                    continue
+
+                try:
+                    if async_wrap:
+                        wrap_function_wrapper(
+                            wrap_package,
+                            target,
+                            _wrap_async_method(tracer, method_wrapper_config),
+                        )
+                    else:
+                        wrap_function_wrapper(
+                            wrap_package,
+                            target,
+                            _wrap_sync_method(tracer, method_wrapper_config),
+                        )
+                except (ModuleNotFoundError, AttributeError) as e:
+                    logger.debug(f"Could not instrument {wrap_package}.{target}: {e}")
+
+    def _uninstrument(self, **kwargs):
+        """Remove OpenHands AI instrumentation."""
+        for wrapped_config in WRAPPED_METHODS:
+            wrap_package = wrapped_config.get("package")
+            wrap_object = wrapped_config.get("object")
+            methods = wrapped_config.get("methods", [])
+
+            for method_config in methods:
+                wrap_method = method_config.get("method")
+
+                # Determine the module path for unwrapping
+                if wrap_object:
+                    module_path = f"{wrap_package}.{wrap_object}"
+                else:
+                    module_path = wrap_package
+
+                try:
+                    unwrap(module_path, wrap_method)
+                except (AttributeError, ValueError) as e:
+                    logger.debug(
+                        f"Could not uninstrument {module_path}.{wrap_method}: {e}"
+                    )

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

@@ -0,0 +1,69 @@
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import unwrap
+from opentelemetry.trace import TraceFlags, SpanContext
+from typing import Collection
+from wrapt import wrap_function_wrapper
+import logging
+
+
+def _wrap_span_context(fn, instance, args, kwargs):
+    """
+    DataDog does something to the OpenTelemetry Contexts, so that when any code
+    tries to access the current active span, it returns a non-recording span.
+
+    There is nothing wrong about that per se, but they create their
+    NonRecordingSpan from an invalid SpanContext, because they don't
+    wrap the trace flags int/bitmap into a TraceFlags object.
+
+    It is an easy to miss bug, because `TraceFlags.SAMPLED` looks like an
+    instance of `TraceFlags`, but is actually just an integer 1, and  the
+    proper way to create it is actually
+    `TraceFlags(TraceFlags.SAMPLED)` or `TraceFlags(0x1)`.
+
+    This is a problem because the trace flags are used to determine if a span
+    is sampled or not. If the trace flags are not wrapped, then the check
+    for sampling will fail, causing any span creation to fail, and sometimes
+    breaking the entire application.
+
+    Issue: https://github.com/DataDog/dd-trace-py/issues/12585
+    PR: https://github.com/DataDog/dd-trace-py/pull/12596
+    The PR only fixed the issue in one place, but it is still there in other places.
+    https://github.com/DataDog/dd-trace-py/pull/12596#issuecomment-2718239507
+
+    https://github.com/DataDog/dd-trace-py/blob/a8419a40fe9e73e0a84c4cab53094c384480a5a6/ddtrace/internal/opentelemetry/context.py#L83
+
+    We patch the `get_span_context` method to return a valid SpanContext.
+    """
+    res = fn(*args, **kwargs)
+
+    new_span_context = SpanContext(
+        trace_id=res.trace_id,
+        span_id=res.span_id,
+        is_remote=res.is_remote,
+        trace_state=res.trace_state,
+        trace_flags=TraceFlags(res.trace_flags),
+    )
+
+    return new_span_context
+
+
+class OpentelemetryInstrumentor(BaseInstrumentor):
+    def __init__(self):
+        super().__init__()
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ("opentelemetry-api>=1.0.0",)
+
+    def _instrument(self, **kwargs):
+        try:
+            wrap_function_wrapper(
+                "opentelemetry.trace.span",
+                "NonRecordingSpan.get_span_context",
+                _wrap_span_context,
+            )
+
+        except Exception as e:
+            logging.debug(f"Error wrapping SpanContext: {e}")
+
+    def _uninstrument(self, **kwargs):
+        unwrap("opentelemetry.trace.span", "NonRecordingSpan.get_span_context")

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

@@ -1,6 +1,5 @@
-from lmnr.opentelemetry_lib.decorators import json_dumps
 from lmnr.sdk.browser.utils import with_tracer_wrapper
-from lmnr.sdk.utils import get_input_from_func_args
+from lmnr.sdk.utils import get_input_from_func_args, json_dumps
 from lmnr.version import __version__
 
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
@@ -25,14 +24,14 @@ WRAPPED_METHODS = [
     {
         "package": "skyvern.library.skyvern",
         "object": "Skyvern",  # Class name
-        "method": "run_task", # Method name
+        "method": "run_task",  # Method name
         "span_name": "Skyvern.run_task",
         "span_type": "DEFAULT",
     },
     {
         "package": "skyvern.webeye.scraper.scraper",
         # No "object" field for module-level functions
-        "method": "get_interactable_element_tree", # Function name
+        "method": "get_interactable_element_tree",  # Function name
         "span_name": "get_interactable_element_tree",
         "span_type": "DEFAULT",
     },
@@ -43,31 +42,31 @@ WRAPPED_METHODS = [
         "span_name": "ForgeAgent.execute_step",
         "span_type": "DEFAULT",
     },
-    {  
-        "package": "skyvern.services.task_v2_service",  
-        "method": "initialize_task_v2",  
-        "span_name": "initialize_task_v2",  
-        "span_type": "DEFAULT",  
+    {
+        "package": "skyvern.services.task_v2_service",
+        "method": "initialize_task_v2",
+        "span_name": "initialize_task_v2",
+        "span_type": "DEFAULT",
     },
-    {  
-        "package": "skyvern.services.task_v2_service",  
-        "method": "run_task_v2_helper",  
-        "span_name": "run_task_v2_helper",  
-        "span_type": "DEFAULT",  
+    {
+        "package": "skyvern.services.task_v2_service",
+        "method": "run_task_v2_helper",
+        "span_name": "run_task_v2_helper",
+        "span_type": "DEFAULT",
     },
-    {  
+    {
         "package": "skyvern.forge.sdk.workflow.models.block",
         "object": "Block",
-        "method": "_generate_workflow_run_block_description",  
-        "span_name": "Block._generate_workflow_run_block_description",  
-        "span_type": "DEFAULT",  
-    },       
-    {  
-        "package": "skyvern.webeye.actions.handler",  
-        "method": "extract_information_for_navigation_goal",  
-        "span_name": "extract_information_for_navigation_goal",  
-        "span_type": "DEFAULT",  
-    },  
+        "method": "_generate_workflow_run_block_description",
+        "span_name": "Block._generate_workflow_run_block_description",
+        "span_type": "DEFAULT",
+    },
+    {
+        "package": "skyvern.webeye.actions.handler",
+        "method": "extract_information_for_navigation_goal",
+        "span_name": "extract_information_for_navigation_goal",
+        "span_type": "DEFAULT",
+    },
 ]
 
 
@@ -77,36 +76,36 @@ async def _wrap(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
     attributes = {
         "lmnr.span.type": to_wrap.get("span_type"),
     }
-    
+
     attributes["lmnr.span.input"] = json_dumps(
         get_input_from_func_args(wrapped, True, args, kwargs)
     )
-    
+
     with tracer.start_as_current_span(span_name, attributes=attributes) as span:
-        try:  
-            result = await wrapped(*args, **kwargs)  
-                
-            to_serialize = result  
-            serialized = (  
-                to_serialize.model_dump_json()  
-                if isinstance(to_serialize, pydantic.BaseModel)  
-                else json_dumps(to_serialize)  
-            )  
-            span.set_attribute("lmnr.span.output", serialized)  
-            return result  
-                
-        except Exception as e:  
-            span.record_exception(e)  
+        try:
+            result = await wrapped(*args, **kwargs)
+
+            to_serialize = result
+            serialized = (
+                to_serialize.model_dump_json()
+                if isinstance(to_serialize, pydantic.BaseModel)
+                else json_dumps(to_serialize)
+            )
+            span.set_attribute("lmnr.span.output", serialized)
+            return result
+
+        except Exception as e:
+            span.record_exception(e)
             raise
 
 
-def instrument_llm_handler(tracer: Tracer):  
-    from skyvern.forge import app  
-      
-    # Store the original handler  
-    original_handler = app.LLM_API_HANDLER  
-      
-    async def wrapped_llm_handler(*args, **kwargs):  
+def instrument_llm_handler(tracer: Tracer):
+    from skyvern.forge import app
+
+    # Store the original handler
+    original_handler = app.LLM_API_HANDLER
+
+    async def wrapped_llm_handler(*args, **kwargs):
 
         prompt_name = kwargs.get("prompt_name", "")
 
@@ -115,13 +114,13 @@ def instrument_llm_handler(tracer: Tracer):
         else:
             span_name = "app.LLM_API_HANDLER"
 
-        attributes = {  
-            "lmnr.span.type": "DEFAULT",  
-        }  
+        attributes = {
+            "lmnr.span.type": "DEFAULT",
+        }
 
-        with tracer.start_as_current_span(span_name, attributes=attributes) as span:  
-            try:  
-                result = await original_handler(*args, **kwargs)  
+        with tracer.start_as_current_span(span_name, attributes=attributes) as span:
+            try:
+                result = await original_handler(*args, **kwargs)
 
                 to_serialize = result
                 serialized = (
@@ -129,13 +128,13 @@ def instrument_llm_handler(tracer: Tracer):
                     if isinstance(to_serialize, pydantic.BaseModel)
                     else json_dumps(to_serialize)
                 )
-                span.set_attribute("lmnr.span.output", serialized)  
-                return result  
-            except Exception as e:  
-                span.record_exception(e)  
+                span.set_attribute("lmnr.span.output", serialized)
+                return result
+            except Exception as e:
+                span.record_exception(e)
                 raise
-      
-    # Replace the global handler  
+
+    # Replace the global handler
     app.LLM_API_HANDLER = wrapped_llm_handler
 
 
@@ -190,4 +189,3 @@ class SkyvernInstrumentor(BaseInstrumentor):
                 module_path = wrap_package
 
             unwrap(module_path, wrap_method)
-