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

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

@@ -0,0 +1,451 @@
+import importlib
+import os
+import sys
+
+from lmnr import Laminar
+from lmnr.opentelemetry_lib.tracing import get_current_context
+from lmnr.opentelemetry_lib.tracing.attributes import SPAN_IDS_PATH, SPAN_PATH
+from lmnr.sdk.log import get_default_logger
+from lmnr.sdk.utils import get_input_from_func_args, is_method, json_dumps
+
+from opentelemetry import trace
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import unwrap
+from opentelemetry.trace import Status, StatusCode
+from typing import Any, Collection
+from typing_extensions import TypedDict
+from wrapt import FunctionWrapper, wrap_function_wrapper
+
+from .proxy import start_proxy, release_proxy, get_cc_proxy_base_url, set_trace_to_proxy
+
+logger = get_default_logger(__name__)
+
+
+class SpanContextPayload(TypedDict):
+    trace_id: str
+    span_id: str
+    project_api_key: str
+    span_ids_path: list[str]
+    span_path: list[str]
+    laminar_url: str
+
+
+_instruments = ("claude-agent-sdk >= 0.1.0",)
+
+WRAPPED_METHODS = [
+    {
+        "package": "claude_agent_sdk.client",
+        "object": "ClaudeSDKClient",
+        "method": "connect",
+        "class_name": "ClaudeSDKClient",
+        "is_async": True,
+        "is_streaming": False,
+        # start proxy on connection
+        "is_start_proxy": True,
+        "is_publish_span_context": True,  # TODO: is there a need to publish span context here?
+        "is_release_proxy": False,
+    },
+    {
+        "package": "claude_agent_sdk.client",
+        "object": "ClaudeSDKClient",
+        "method": "query",
+        "class_name": "ClaudeSDKClient",
+        "is_async": True,
+        "is_streaming": False,
+        # only publish span context here as start/close managed by connect/disconnect
+        "is_start_proxy": False,
+        "is_publish_span_context": True,
+        "is_release_proxy": False,
+    },
+    {
+        "package": "claude_agent_sdk.client",
+        "object": "ClaudeSDKClient",
+        "method": "receive_messages",
+        "class_name": "ClaudeSDKClient",
+        "is_async": True,
+        "is_streaming": True,
+    },
+    {
+        "package": "claude_agent_sdk.client",
+        "object": "ClaudeSDKClient",
+        "method": "receive_response",
+        "class_name": "ClaudeSDKClient",
+        "is_async": True,
+        "is_streaming": True,
+    },
+    {
+        "package": "claude_agent_sdk.client",
+        "object": "ClaudeSDKClient",
+        "method": "interrupt",
+        "class_name": "ClaudeSDKClient",
+        "is_async": True,
+        "is_streaming": False,
+    },
+    {
+        "package": "claude_agent_sdk.client",
+        "object": "ClaudeSDKClient",
+        "method": "disconnect",
+        "class_name": "ClaudeSDKClient",
+        "is_async": True,
+        "is_streaming": False,
+        # close proxy on a connection drop
+        "is_start_proxy": False,
+        "is_publish_span_context": False,
+        "is_release_proxy": True,
+    },
+    {
+        # No "object" and "class_name" fields for module-level functions
+        "package": "claude_agent_sdk",
+        "method": "query",
+        "is_async": True,
+        "is_streaming": True,
+        # start, send span to, and release proxy here as it is a module-level function doing all on its own
+        "is_start_proxy": True,
+        "is_publish_span_context": True,
+        "is_release_proxy": True,
+    },
+    {
+        # No "object" and "class_name" fields for module-level functions
+        "package": "claude_agent_sdk",
+        "method": "create_sdk_mcp_server",
+        "is_async": False,
+        "is_streaming": False,
+    },
+]
+
+_MODULE_FUNCTION_ORIGINALS: dict[tuple[str, str], Any] = {}
+
+
+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 _replace_function_aliases(original, wrapped):
+    for module in list(sys.modules.values()):
+        module_dict = getattr(module, "__dict__", None)
+        if not module_dict:
+            continue
+        for attr, value in list(module_dict.items()):
+            if value is original:
+                setattr(module, attr, wrapped)
+
+
+def _wrap_module_function(module_name: str, function_name: str, wrapper):
+    try:
+        module = sys.modules.get(module_name) or importlib.import_module(module_name)
+    except ModuleNotFoundError:
+        return
+
+    try:
+        original = getattr(module, function_name)
+    except AttributeError:
+        return
+
+    key = (module_name, function_name)
+    if key not in _MODULE_FUNCTION_ORIGINALS:
+        _MODULE_FUNCTION_ORIGINALS[key] = original
+
+    wrapped_function = FunctionWrapper(original, wrapper)
+    setattr(module, function_name, wrapped_function)
+    _replace_function_aliases(original, wrapped_function)
+
+
+def _unwrap_module_function(module_name: str, function_name: str):
+    key = (module_name, function_name)
+    original = _MODULE_FUNCTION_ORIGINALS.get(key)
+    if not original:
+        return
+
+    module = sys.modules.get(module_name)
+    if not module:
+        return
+
+    current = getattr(module, function_name, None)
+    setattr(module, function_name, original)
+    if current is not None:
+        _replace_function_aliases(current, original)
+    del _MODULE_FUNCTION_ORIGINALS[key]
+
+
+def _span_name(to_wrap: dict[str, str]) -> str:
+    class_name = to_wrap.get("class_name")
+    method = to_wrap.get("method")
+    return f"{class_name}.{method}" if class_name else method
+
+
+def _record_input(span, wrapped, args, kwargs):
+    try:
+        span.set_attribute(
+            "lmnr.span.input",
+            json_dumps(
+                get_input_from_func_args(
+                    wrapped,
+                    is_method=is_method(wrapped),
+                    func_args=list(args),
+                    func_kwargs=kwargs,
+                )
+            ),
+        )
+    except Exception:
+        pass
+
+
+def _record_output(span, to_wrap, value):
+    try:
+        span.set_attribute("lmnr.span.output", json_dumps(value))
+    except Exception:
+        pass
+
+
+def _get_span_context_payload() -> dict[str, str] | None:
+    current_span: ReadableSpan = trace.get_current_span(context=get_current_context())
+    if current_span is trace.INVALID_SPAN:
+        return None
+
+    span_context = current_span.get_span_context()
+    if span_context is None or not span_context.is_valid:
+        return None
+
+    span_ids_path = []
+    span_path = []
+    if hasattr(current_span, "attributes"):
+        readable_span: ReadableSpan = current_span
+        span_ids_path = list(readable_span.attributes.get(SPAN_IDS_PATH, tuple()))
+        span_path = list(readable_span.attributes.get(SPAN_PATH, tuple()))
+
+    project_api_key = Laminar.get_project_api_key()
+    laminar_url = Laminar.get_base_http_url()
+
+    return {
+        "trace_id": f"{span_context.trace_id:032x}",
+        "span_id": f"{span_context.span_id:016x}",
+        "project_api_key": project_api_key or "",
+        "span_ids_path": span_ids_path,
+        "span_path": span_path,
+        "laminar_url": laminar_url,
+    }
+
+
+def _publish_span_context() -> None:
+    payload = _get_span_context_payload()
+    if not payload:
+        return
+
+    set_trace_to_proxy(
+        trace_id=payload["trace_id"],
+        span_id=payload["span_id"],
+        project_api_key=payload["project_api_key"],
+        span_ids_path=payload["span_ids_path"],
+        span_path=payload["span_path"],
+        laminar_url=payload["laminar_url"],
+    )
+
+
+@_with_wrapper
+def _wrap_sync(to_wrap, wrapped, instance, args, kwargs):
+    with Laminar.start_as_current_span(
+        _span_name(to_wrap),
+        span_type=to_wrap.get("span_type", "DEFAULT"),
+    ) as span:
+        _record_input(span, wrapped, 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)
+            raise
+
+        _record_output(span, to_wrap, result)
+        return result
+
+
+@_with_wrapper
+async def _wrap_async(to_wrap, wrapped, instance, args, kwargs):
+    with Laminar.start_as_current_span(
+        _span_name(to_wrap),
+        span_type=to_wrap.get("span_type", "DEFAULT"),
+    ) as span:
+        _record_input(span, wrapped, args, kwargs)
+
+        original_base_url = None
+        if to_wrap.get("is_start_proxy"):
+            original_base_url = os.environ.get("ANTHROPIC_BASE_URL")
+            start_proxy()
+
+        if to_wrap.get("is_publish_span_context"):
+            proxy_base_url = get_cc_proxy_base_url()
+            if proxy_base_url:
+                _publish_span_context()
+            else:
+                logger.debug(
+                    "No claude proxy server found. Skipping span context publication."
+                )
+
+        try:
+            result = await wrapped(*args, **kwargs)
+        except Exception as e:  # pylint: disable=broad-except
+            span.set_status(Status(StatusCode.ERROR))
+            span.record_exception(e)
+            raise
+        finally:
+            if original_base_url is not None:
+                if original_base_url:
+                    os.environ["ANTHROPIC_BASE_URL"] = original_base_url
+                else:
+                    os.environ.pop("ANTHROPIC_BASE_URL", None)
+            if to_wrap.get("is_release_proxy"):
+                release_proxy()
+
+        _record_output(span, to_wrap, result)
+
+        return result
+
+
+@_with_wrapper
+def _wrap_async_gen(to_wrap, wrapped, instance, args, kwargs):
+    async def generator():
+        span = Laminar.start_span(
+            _span_name(to_wrap),
+            span_type=to_wrap.get("span_type", "DEFAULT"),
+        )
+        collected = []
+        async_iter = None
+
+        original_base_url = None
+        if to_wrap.get("is_start_proxy"):
+            original_base_url = os.environ.get("ANTHROPIC_BASE_URL")
+            start_proxy()
+
+        if to_wrap.get("is_publish_span_context"):
+            with Laminar.use_span(span):
+                proxy_base_url = get_cc_proxy_base_url()
+                if proxy_base_url:
+                    _publish_span_context()
+                else:
+                    logger.debug(
+                        "No claude proxy server found. Skipping span context publication."
+                    )
+
+        try:
+            with Laminar.use_span(span):
+                _record_input(span, wrapped, args, kwargs)
+                async_source = wrapped(*args, **kwargs)
+                async_iter = (
+                    async_source.__aiter__()
+                    if hasattr(async_source, "__aiter__")
+                    else async_source
+                )
+
+            while True:
+                try:
+                    with Laminar.use_span(
+                        span, record_exception=False, set_status_on_exception=False
+                    ):
+                        item = await async_iter.__anext__()
+                        collected.append(item)
+                except StopAsyncIteration:
+                    break
+                yield item
+        except Exception as e:  # pylint: disable=broad-except
+            with Laminar.use_span(span):
+                span.set_status(Status(StatusCode.ERROR))
+                span.record_exception(e)
+            raise
+        finally:
+            if original_base_url is not None:
+                if original_base_url:
+                    os.environ["ANTHROPIC_BASE_URL"] = original_base_url
+                else:
+                    os.environ.pop("ANTHROPIC_BASE_URL", None)
+            if async_iter and hasattr(async_iter, "aclose"):
+                try:
+                    with Laminar.use_span(span):
+                        await async_iter.aclose()
+                except Exception:  # pylint: disable=broad-except
+                    pass
+            with Laminar.use_span(span):
+                _record_output(span, to_wrap, collected)
+                span.end()
+
+            if to_wrap.get("is_release_proxy"):
+                release_proxy()
+
+    return generator()
+
+
+class ClaudeAgentInstrumentor(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")
+            is_streaming = wrapped_method.get("is_streaming", False)
+            is_async = wrapped_method.get("is_async", False)
+
+            if is_streaming:
+                wrapper_factory = _wrap_async_gen
+            elif is_async:
+                wrapper_factory = _wrap_async
+            else:
+                wrapper_factory = _wrap_sync
+
+            wrapper = wrapper_factory(wrapped_method)
+
+            if wrap_object:
+                target = f"{wrap_object}.{wrap_method}"
+                try:
+                    wrap_function_wrapper(
+                        wrap_package,
+                        target,
+                        wrapper,
+                    )
+                except (ModuleNotFoundError, AttributeError):
+                    pass  # that's ok, we don't want to fail if some methods do not exist
+            else:
+                try:
+                    _wrap_module_function(
+                        wrap_package,
+                        wrap_method,
+                        wrapper,
+                    )
+                except (ModuleNotFoundError, AttributeError):
+                    pass  # that's ok
+
+    def _uninstrument(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")
+
+            if wrap_object:
+                module_path = f"{wrap_package}.{wrap_object}"
+                try:
+                    unwrap(module_path, wrap_method)
+                except (ModuleNotFoundError, AttributeError):
+                    pass  # that's ok, we don't want to fail if some methods do not exist
+            else:
+                _unwrap_module_function(wrap_package, wrap_method)

lmnr/opentelemetry_lib/opentelemetry/instrumentation/claude_agent/proxy.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+import atexit
+import os
+import socket
+import threading
+import time
+from typing import Optional
+
+from lmnr.sdk.log import get_default_logger
+
+from lmnr_claude_code_proxy import run_server, set_current_trace, stop_server
+
+logger = get_default_logger(__name__)
+
+DEFAULT_ANTHROPIC_BASE_URL = "https://api.anthropic.com"
+DEFAULT_CC_PROXY_PORT = 45667
+CC_PROXY_PORT_ATTEMPTS = 5
+
+_CC_PROXY_LOCK = threading.Lock()
+_CC_PROXY_PORT: int | None = None
+_CC_PROXY_BASE_URL: str | None = None
+_CC_PROXY_TARGET_URL: str | None = None
+_CC_PROXY_SHUTDOWN_REGISTERED = False
+
+
+def _find_available_port(start_port: int, attempts: int) -> Optional[int]:
+    for offset in range(attempts):
+        candidate = start_port + offset
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            try:
+                sock.bind(("127.0.0.1", candidate))
+            except OSError:
+                continue
+        return candidate
+    return None
+
+
+def _wait_for_port(port: int, timeout: float = 5.0) -> bool:
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            sock.settimeout(0.2)
+            try:
+                sock.connect(("127.0.0.1", port))
+                return True
+            except OSError:
+                time.sleep(0.1)
+    return False
+
+
+def _stop_cc_proxy_locked():
+    global _CC_PROXY_PORT, _CC_PROXY_BASE_URL
+
+    try:
+        stop_server()
+    except Exception as e:
+        logger.debug("Unable to stop cc-proxy: %s", e)
+
+    if _CC_PROXY_TARGET_URL:
+        os.environ["ANTHROPIC_BASE_URL"] = _CC_PROXY_TARGET_URL
+
+    _CC_PROXY_PORT = None
+    _CC_PROXY_BASE_URL = None
+
+
+def _stop_cc_proxy():
+    with _CC_PROXY_LOCK:
+        _stop_cc_proxy_locked()
+
+
+def _register_proxy_shutdown():
+    global _CC_PROXY_SHUTDOWN_REGISTERED
+    if not _CC_PROXY_SHUTDOWN_REGISTERED:
+        atexit.register(_stop_cc_proxy)
+        _CC_PROXY_SHUTDOWN_REGISTERED = True
+
+
+def get_cc_proxy_base_url() -> str | None:
+    return _CC_PROXY_BASE_URL
+
+
+def start_proxy() -> Optional[str]:
+    with _CC_PROXY_LOCK:
+        global _CC_PROXY_PORT, _CC_PROXY_BASE_URL, _CC_PROXY_TARGET_URL
+
+        port = _find_available_port(DEFAULT_CC_PROXY_PORT, CC_PROXY_PORT_ATTEMPTS)
+        if port is None:
+            logger.warning("Unable to allocate port for cc-proxy.")
+            return None
+
+        target_url = (
+            _CC_PROXY_TARGET_URL
+            or os.environ.get("ANTHROPIC_ORIGINAL_BASE_URL")
+            or os.environ.get("ANTHROPIC_BASE_URL")
+            or DEFAULT_ANTHROPIC_BASE_URL
+        )
+        _CC_PROXY_TARGET_URL = target_url
+        os.environ.setdefault("ANTHROPIC_ORIGINAL_BASE_URL", target_url)
+
+        try:
+            run_server(target_url, port=port)
+        except OSError as exc:  # pragma: no cover
+            logger.warning("Unable to start cc-proxy: %s", exc)
+            return None
+
+        if not _wait_for_port(port):
+            logger.warning("cc-proxy failed to start on port %s", port)
+            stop_server()
+            return None
+
+        proxy_base_url = f"http://127.0.0.1:{port}"
+        _CC_PROXY_PORT = port
+        _CC_PROXY_BASE_URL = proxy_base_url
+        os.environ["ANTHROPIC_BASE_URL"] = proxy_base_url
+        _register_proxy_shutdown()
+
+        logger.info("Started claude proxy server on: " + str(proxy_base_url))
+        return proxy_base_url
+
+
+def release_proxy() -> None:
+    with _CC_PROXY_LOCK:
+        _stop_cc_proxy_locked()
+        logger.debug("Released claude proxy server")
+
+
+def set_trace_to_proxy(
+    trace_id: str,
+    span_id: str,
+    project_api_key: str,
+    span_path: list[str] = [],
+    span_ids_path: list[str] = [],
+    laminar_url: str = "https://api.lmnr.ai",
+):
+    set_current_trace(
+        trace_id=trace_id,
+        span_id=span_id,
+        project_api_key=project_api_key,
+        span_path=span_path,
+        span_ids_path=span_ids_path,
+        laminar_url=laminar_url,
+    )

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 import Laminar
+from lmnr.sdk.utils import json_dumps
+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