diff --git a/pyproject.toml b/pyproject.toml index 2fd42bf0..0587d44e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ [project] name = "lmnr" -version = "0.7.20" +version = "0.7.21" description = "Python SDK for Laminar" authors = [ { name = "lmnr.ai", email = "founders@lmnr.ai" } @@ -50,6 +50,7 @@ classifiers = [ "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", ] [project.scripts] @@ -136,10 +137,11 @@ dev = [ "groq>=0.30.0", "anthropic[bedrock]>=0.60.0", "langchain-openai>=0.3.32", + "kernel>=0.18.0", ] [build-system] -requires = ["uv_build>=0.7.21,<0.8"] +requires = ["uv_build>=0.9.7,<0.10"] build-backend = "uv_build" [tool.uv.workspace] diff --git a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/kernel/__init__.py b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/kernel/__init__.py new file mode 100644 index 00000000..55829f0e --- /dev/null +++ b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/kernel/__init__.py @@ -0,0 +1,374 @@ +"""OpenTelemetry Kernel instrumentation""" + +import functools +from typing import Collection + +from lmnr.opentelemetry_lib.decorators import json_dumps +from lmnr.opentelemetry_lib.opentelemetry.instrumentation.kernel.utils import ( + process_tool_output_formatter, + screenshot_tool_output_formatter, +) +from lmnr.sdk.decorators import observe +from lmnr.sdk.utils import get_input_from_func_args, is_async +from lmnr import Laminar +from opentelemetry.instrumentation.instrumentor import BaseInstrumentor +from opentelemetry.instrumentation.utils import unwrap + +from opentelemetry.trace.status import Status, StatusCode +from wrapt import wrap_function_wrapper + +_instruments = ("kernel >= 0.2.0",) + + +WRAPPED_METHODS = [ + { + "package": "kernel.resources.browsers", + "object": "BrowsersResource", + "method": "create", + "class_name": "Browser", + }, + { + "package": "kernel.resources.browsers", + "object": "BrowsersResource", + "method": "retrieve", + "class_name": "Browser", + }, + { + "package": "kernel.resources.browsers", + "object": "BrowsersResource", + "method": "list", + "class_name": "Browser", + }, + { + "package": "kernel.resources.browsers", + "object": "BrowsersResource", + "method": "delete", + "class_name": "Browser", + }, + { + "package": "kernel.resources.browsers", + "object": "BrowsersResource", + "method": "delete_by_id", + "class_name": "Browser", + }, + { + "package": "kernel.resources.browsers", + "object": "BrowsersResource", + "method": "load_extensions", + "class_name": "Browser", + }, + { + "package": "kernel.resources.browsers.computer", + "object": "ComputerResource", + "method": "capture_screenshot", + "class_name": "Computer", + "span_type": "TOOL", + "output_formatter": screenshot_tool_output_formatter, + }, + { + "package": "kernel.resources.browsers.computer", + "object": "ComputerResource", + "method": "click_mouse", + "class_name": "Computer", + "span_type": "TOOL", + }, + { + "package": "kernel.resources.browsers.computer", + "object": "ComputerResource", + "method": "drag_mouse", + "class_name": "Computer", + "span_type": "TOOL", + }, + { + "package": "kernel.resources.browsers.computer", + "object": "ComputerResource", + "method": "move_mouse", + "class_name": "Computer", + "span_type": "TOOL", + }, + { + "package": "kernel.resources.browsers.computer", + "object": "ComputerResource", + "method": "press_key", + "class_name": "Computer", + "span_type": "TOOL", + }, + { + "package": "kernel.resources.browsers.computer", + "object": "ComputerResource", + "method": "scroll", + "class_name": "Computer", + "span_type": "TOOL", + }, + { + "package": "kernel.resources.browsers.computer", + "object": "ComputerResource", + "method": "type_text", + "class_name": "Computer", + "span_type": "TOOL", + }, + { + "package": "kernel.resources.browsers.playwright", + "object": "PlaywrightResource", + "method": "execute", + "class_name": "Playwright", + }, + { + "package": "kernel.resources.browsers.process", + "object": "ProcessResource", + "method": "exec", + "class_name": "Process", + "span_type": "TOOL", + "output_formatter": process_tool_output_formatter, + }, + { + "package": "kernel.resources.browsers.process", + "object": "ProcessResource", + "method": "kill", + "class_name": "Process", + "span_type": "TOOL", + "output_formatter": process_tool_output_formatter, + }, + { + "package": "kernel.resources.browsers.process", + "object": "ProcessResource", + "method": "spawn", + "class_name": "Process", + "span_type": "TOOL", + "output_formatter": process_tool_output_formatter, + }, + { + "package": "kernel.resources.browsers.process", + "object": "ProcessResource", + "method": "status", + "class_name": "Process", + "span_type": "TOOL", + "output_formatter": process_tool_output_formatter, + }, + { + "package": "kernel.resources.browsers.process", + "object": "ProcessResource", + "method": "stdin", + "class_name": "Process", + "span_type": "TOOL", + "output_formatter": process_tool_output_formatter, + }, + { + "package": "kernel.resources.browsers.process", + "object": "ProcessResource", + "method": "stdout_stream", + "class_name": "Process", + "span_type": "TOOL", + "output_formatter": process_tool_output_formatter, + }, +] + + +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 + + +@_with_wrapper +def _wrap( + to_wrap, + wrapped, + instance, + args, + kwargs, +): + with Laminar.start_as_current_span( + f"{to_wrap.get('class_name')}.{to_wrap.get('method')}", + span_type=to_wrap.get("span_type", "DEFAULT"), + ) 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) + 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, +): + with Laminar.start_as_current_span( + f"{to_wrap.get('class_name')}.{to_wrap.get('method')}", + span_type=to_wrap.get("span_type", "DEFAULT"), + ) 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) + 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 +def _wrap_app_action( + to_wrap, + wrapped, + instance, + args, + kwargs, +): + """ + Wraps app.action() decorator factory to add tracing to action handlers. + + wrapped: the original `action` method + args: (name,) - the action name + kwargs: potentially {'name': ...} + + Returns a decorator that wraps handlers with tracing before registering them. + """ + + # Call the original action method to get the decorator + original_decorator = wrapped(*args, **kwargs) + + # Get the action name from args + action_name = args[0] if args else kwargs.get("name", "unknown") + + # Create a wrapper for the decorator that intercepts the handler + def tracing_decorator(handler): + # Apply the observe decorator to add tracing + observed_handler = observe( + name=f"action.{action_name}", + span_type="DEFAULT", + )(handler) + + # Create an additional wrapper to add post-execution logic + if is_async(handler): + + @functools.wraps(handler) + async def async_wrapper_with_flush(*handler_args, **handler_kwargs): + # Execute the observed handler (tracing happens here) + result = await observed_handler(*handler_args, **handler_kwargs) + + Laminar.flush() + + return result + + # Register the wrapper with the original decorator + return original_decorator(async_wrapper_with_flush) + else: + + @functools.wraps(handler) + def sync_wrapper_with_flush(*handler_args, **handler_kwargs): + # Execute the observed handler (tracing happens here) + result = observed_handler(*handler_args, **handler_kwargs) + + Laminar.flush() + + return result + + # Register the wrapper with the original decorator + return original_decorator(sync_wrapper_with_flush) + + return tracing_decorator + + +class KernelInstrumentor(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, AttributeError): + pass # that's ok, we don't want to fail if some methods do not exist + + for wrapped_method in WRAPPED_METHODS: + wrap_package = wrapped_method.get("package") + wrap_object = f"Async{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, AttributeError): + pass # that's ok, we don't want to fail if some methods do not exist + + try: + wrap_function_wrapper( + "kernel.app_framework", + "KernelApp.action", + _wrap_app_action({}), + ) + except (ModuleNotFoundError, AttributeError): + pass + + 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, AttributeError): + pass # that's ok, we don't want to fail if some methods do not exist + + for wrapped_method in WRAPPED_METHODS: + wrap_package = wrapped_method.get("package") + wrap_object = f"Async{wrapped_method.get('object')}" + try: + unwrap( + f"{wrap_package}.{wrap_object}", + wrapped_method.get("method"), + ) + except (ModuleNotFoundError, AttributeError): + pass # that's ok, we don't want to fail if some methods do not exist + + try: + unwrap("kernel.app_framework.KernelApp", "action") + except (ModuleNotFoundError, AttributeError): + pass diff --git a/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/kernel/utils.py b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/kernel/utils.py new file mode 100644 index 00000000..69c888a1 --- /dev/null +++ b/src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/kernel/utils.py @@ -0,0 +1,36 @@ +# import base64 +import base64 +from copy import deepcopy +from typing import Any + +from lmnr.opentelemetry_lib.decorators import json_dumps +from pydantic import BaseModel + + +def screenshot_tool_output_formatter(output: Any) -> str: + # output is of type BinaryAPIResponse, which implements + # the iter_bytes method from httpx.Response + + return "" + # The below implementation works, but it may consume the entire iterator, + # making the response unusable after the formatter is called. + # This is UNLESS somewhere in code output.read() (httpx.Response.read()) + # is called. + # We cannot rely on that now, so we return a placeholder. + # response_bytes = [] + # for chunk in output.iter_bytes(): + # response_bytes.append(chunk) + # response_base64 = base64.b64encode(response_bytes).decode("utf-8") + # return f"data:image/png;base64,{response_base64}" + + +def process_tool_output_formatter(output: Any) -> str: + if not isinstance(output, (dict, BaseModel)): + return json_dumps(output) + + output = output.model_dump() if isinstance(output, BaseModel) else deepcopy(output) + if "stderr_b64" in output: + output["stderr"] = base64.b64decode(output["stderr_b64"]).decode("utf-8") + if "stdout_b64" in output: + output["stdout"] = base64.b64decode(output["stdout_b64"]).decode("utf-8") + return json_dumps(output) diff --git a/src/lmnr/opentelemetry_lib/tracing/_instrument_initializers.py b/src/lmnr/opentelemetry_lib/tracing/_instrument_initializers.py index e5dfffcc..9ac99103 100644 --- a/src/lmnr/opentelemetry_lib/tracing/_instrument_initializers.py +++ b/src/lmnr/opentelemetry_lib/tracing/_instrument_initializers.py @@ -199,6 +199,16 @@ def init_instrumentor(self, *args, **kwargs) -> BaseInstrumentor | None: return HaystackInstrumentor() +class KernelInstrumentorInitializer(InstrumentorInitializer): + def init_instrumentor(self, *args, **kwargs) -> BaseInstrumentor | None: + if not is_package_installed("kernel"): + return None + + from ..opentelemetry.instrumentation.kernel import KernelInstrumentor + + return KernelInstrumentor() + + class LanceDBInstrumentorInitializer(InstrumentorInitializer): def init_instrumentor(self, *args, **kwargs) -> BaseInstrumentor | None: if not is_package_installed("lancedb"): diff --git a/src/lmnr/opentelemetry_lib/tracing/instruments.py b/src/lmnr/opentelemetry_lib/tracing/instruments.py index ea64fa2e..c17ed40a 100644 --- a/src/lmnr/opentelemetry_lib/tracing/instruments.py +++ b/src/lmnr/opentelemetry_lib/tracing/instruments.py @@ -27,6 +27,7 @@ class Instruments(Enum): GOOGLE_GENAI = "google_genai" GROQ = "groq" HAYSTACK = "haystack" + KERNEL = "kernel" LANCEDB = "lancedb" LANGCHAIN = "langchain" LANGGRAPH = "langgraph" @@ -74,6 +75,7 @@ class Instruments(Enum): Instruments.GOOGLE_GENAI: initializers.GoogleGenAIInstrumentorInitializer(), Instruments.GROQ: initializers.GroqInstrumentorInitializer(), Instruments.HAYSTACK: initializers.HaystackInstrumentorInitializer(), + Instruments.KERNEL: initializers.KernelInstrumentorInitializer(), Instruments.LANCEDB: initializers.LanceDBInstrumentorInitializer(), Instruments.LANGCHAIN: initializers.LangchainInstrumentorInitializer(), Instruments.LANGGRAPH: initializers.LanggraphInstrumentorInitializer(), diff --git a/src/lmnr/version.py b/src/lmnr/version.py index 9ecf8c75..62871ea0 100644 --- a/src/lmnr/version.py +++ b/src/lmnr/version.py @@ -3,7 +3,7 @@ from packaging import version -__version__ = "0.7.20" +__version__ = "0.7.21" PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}"