blaxel 0.1.22rc70__py3-none-any.whl → 0.2.0rc2__py3-none-any.whl

blaxel/telemetry/__init__.py

@@ -0,0 +1,6 @@
+"""Blaxel telemetry module."""
+
+from .exporters import *  # noqa: F403, F401
+from .instrumentation import *  # noqa: F403, F401
+from .log import *  # noqa: F403, F401
+from .manager import *  # noqa: F403, F401

blaxel/telemetry/instrumentation/blaxel_core.py

@@ -0,0 +1,124 @@
+import json
+import logging
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import Tracer, get_tracer
+
+import blaxel.core.agents
+import blaxel.core.jobs
+from blaxel.telemetry.span import SpanManager
+
+logger = logging.getLogger(__name__)
+
+class BlaxelCoreInstrumentor(BaseInstrumentor):
+    def instrumentation_dependencies(self):
+        return ["blaxel.core"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "0.2.0", tracer_provider)
+        self._patch_agent(tracer)
+        self._patch_job(tracer)
+        self._patch_tools(tracer)
+
+    def _uninstrument(self, **kwargs):
+        # Optionally restore original methods
+        pass
+
+    def _patch_agent(self, tracer: Tracer):
+        orig_run = blaxel.core.agents.BlAgent.run
+
+        def traced_run(self, *args, **kwargs):
+            attributes = {
+                "agent.name": self.name,
+                "agent.args": str(args),
+                "span.type": "agent.run",
+                **SpanManager.get_default_attributes(),
+            }
+            with tracer.start_span(self.name, attributes=attributes) as span:
+                try:
+                    result = orig_run(self, *args, **kwargs)
+                    span.set_attribute("agent.run.result", result)
+                    return result
+                except Exception as e:
+                    span.set_attribute("agent.run.error", str(e))
+                    raise
+
+        blaxel.core.agents.BlAgent.run = traced_run
+
+    def _patch_job(self, tracer: Tracer):
+        orig_run = blaxel.core.jobs.BlJob.run
+
+        def traced_run(self, *args, **kwargs):
+            attributes = {
+                "job.name": self.name,
+                "span.type": "job.run",
+                **SpanManager.get_default_attributes(),
+            }
+            with tracer.start_span(self.name, attributes=attributes) as span:
+                try:
+                    result = orig_run(self, *args, **kwargs)
+                    span.set_attribute("job.run.result", result)
+                    return result
+                except Exception as e:
+                    span.set_attribute("job.run.error", str(e))
+                    raise
+
+        blaxel.core.jobs.BlJob.run = traced_run
+
+    def _patch_tools(self, tracer: Tracer):
+        # Patch PersistentWebSocket.list_tools
+        orig_list_tools = blaxel.core.tools.PersistentWebSocket.list_tools
+
+        async def traced_list_tools(self, *args, **kwargs):
+            span_attributes = {
+                "tool.server": self.url,
+                "tool.server_name": self.name,
+                "span.type": "tool.list",
+                **SpanManager.get_default_attributes(),
+            }
+            with tracer.start_span(self.name, attributes=span_attributes):
+                result = await orig_list_tools(self, *args, **kwargs)
+                # Optionally: span.set_attribute("tool.list.result", str(result))
+                return result
+
+        blaxel.core.tools.PersistentWebSocket.list_tools = traced_list_tools
+
+        # Patch convert_mcp_tool_to_blaxel_tool to wrap tool calls
+        orig_convert = blaxel.core.tools.convert_mcp_tool_to_blaxel_tool
+
+        def traced_convert_mcp_tool_to_blaxel_tool(websocket_client, name, url, tool):
+            Tool = orig_convert(websocket_client, name, url, tool)
+            orig_coroutine = Tool.coroutine
+            orig_sync_coroutine = Tool.sync_coroutine
+
+            async def traced_coroutine(*args, **kwargs):
+                span_attributes = {
+                    "tool.name": tool.name,
+                    "tool.args": json.dumps(kwargs),
+                    "tool.server": url,
+                    "tool.server_name": name,
+                    "span.type": "tool.call",
+                    **SpanManager.get_default_attributes(),
+                }
+                logger.info(span_attributes)
+                with tracer.start_span("blaxel-tool-call", attributes=span_attributes):
+                    return await orig_coroutine(*args, **kwargs)
+
+            def traced_sync_coroutine(*args, **kwargs):
+                span_attributes = {
+                    "tool.name": tool.name,
+                    "tool.args": json.dumps(kwargs),
+                    "tool.server": url,
+                    "tool.server_name": name,
+                    "span.type": "tool.call",
+                    **SpanManager.get_default_attributes(),
+                }
+                with tracer.start_span("blaxel-tool-call", attributes=span_attributes):
+                    return orig_sync_coroutine(*args, **kwargs)
+
+            Tool.coroutine = traced_coroutine
+            Tool.sync_coroutine = traced_sync_coroutine
+            return Tool
+
+        blaxel.core.tools.convert_mcp_tool_to_blaxel_tool = traced_convert_mcp_tool_to_blaxel_tool

blaxel/telemetry/instrumentation/blaxel_langgraph.py

@@ -0,0 +1,74 @@
+import logging
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import Tracer, get_tracer
+from wrapt import wrap_function_wrapper
+
+from blaxel.telemetry.span import SpanManager
+
+logger = logging.getLogger(__name__)
+
+class BlaxelLanggraphInstrumentor(BaseInstrumentor):
+    def instrumentation_dependencies(self):
+        return ["blaxel_langgraph", "blaxel.core"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "0.2.0", tracer_provider)
+        # self._patch_gemini(tracer)
+        self._patch_tool(tracer)
+
+    def _uninstrument(self, **kwargs):
+        import blaxel.langgraph.custom.gemini
+
+        # Restore original methods
+        if hasattr(blaxel.langgraph.custom.gemini.GeminiRestClient, "_blaxel_original_generate_content"):
+            blaxel.langgraph.custom.gemini.GeminiRestClient.generate_content = blaxel.langgraph.custom.gemini.GeminiRestClient._blaxel_original_generate_content
+            delattr(blaxel.langgraph.custom.gemini.GeminiRestClient, "_blaxel_original_generate_content")
+
+        if hasattr(blaxel.langgraph.custom.gemini.GeminiRestClient, "_blaxel_original_generate_content_async"):
+            blaxel.langgraph.custom.gemini.GeminiRestClient.generate_content_async = blaxel.langgraph.custom.gemini.GeminiRestClient._blaxel_original_generate_content_async
+            delattr(blaxel.langgraph.custom.gemini.GeminiRestClient, "_blaxel_original_generate_content_async")
+
+        if hasattr(blaxel.langgraph.custom.gemini.GeminiRestClient, "_blaxel_original_stream_generate_content"):
+            blaxel.langgraph.custom.gemini.GeminiRestClient.stream_generate_content = blaxel.langgraph.custom.gemini.GeminiRestClient._blaxel_original_stream_generate_content
+            delattr(blaxel.langgraph.custom.gemini.GeminiRestClient, "_blaxel_original_stream_generate_content")
+
+        if hasattr(blaxel.langgraph.custom.gemini.GeminiRestClient, "_blaxel_original_stream_generate_content_async"):
+            blaxel.langgraph.custom.gemini.GeminiRestClient.stream_generate_content_async = blaxel.langgraph.custom.gemini.GeminiRestClient._blaxel_original_stream_generate_content_async
+            delattr(blaxel.langgraph.custom.gemini.GeminiRestClient, "_blaxel_original_stream_generate_content_async")
+
+    def _patch_tool(self, tracer: Tracer):
+        def _traced_run(original, instance, args, kwargs):
+            attributes = {
+                "tool.type": "langgraph",
+                "span.type": "tool.call",
+                **SpanManager.get_default_attributes(),
+            }
+            with tracer.start_span("langgraph.tool", attributes=attributes) as span:
+                try:
+                    result = original(*args, **kwargs)
+                    span.set_attribute("tool.name", result.name)
+                    return result
+                except Exception as e:
+                    span.set_attribute("tool.error", str(e))
+                    raise
+
+        async def _traced_arun(original, instance, args, kwargs):
+            attributes = {
+                "tool.type": "langgraph",
+                "span.type": "tool.call",
+                **SpanManager.get_default_attributes(),
+            }
+            with tracer.start_span("langgraph.tool", attributes=attributes) as span:
+                try:
+                    result = await original(*args, **kwargs)
+                    span.set_attribute("tool.name", result.name)
+                    return result
+                except Exception as e:
+                    span.set_attribute("tool.error", str(e))
+                    raise
+
+        wrap_function_wrapper('langchain_core.tools', 'BaseTool.run', _traced_run)
+        wrap_function_wrapper('langchain_core.tools', 'BaseTool.arun', _traced_arun)
+        wrap_function_wrapper('langchain_core.tools', 'BaseTool.arun', _traced_arun)

blaxel/telemetry/instrumentation/blaxel_langgraph_gemini.py

@@ -0,0 +1,360 @@
+"""OpenTelemetry Google Generative AI API instrumentation"""
+
+import logging
+import os
+import types
+from typing import Collection
+
+from opentelemetry import context as context_api
+from opentelemetry.instrumentation.google_generativeai.config import Config
+from opentelemetry.instrumentation.google_generativeai.utils import dont_throw
+from opentelemetry.instrumentation.google_generativeai.version import __version__
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY, unwrap
+from opentelemetry.semconv_ai import (
+    SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY,
+    LLMRequestTypeValues,
+    SpanAttributes,
+)
+from opentelemetry.trace import SpanKind, get_tracer
+from opentelemetry.trace.status import Status, StatusCode
+from wrapt import wrap_function_wrapper
+
+logger = logging.getLogger(__name__)
+
+WRAPPED_METHODS = [
+    {
+        "package": "blaxel_langgraph.custom.gemini",
+        "object": "GeminiRestClient",
+        "method": "generate_content",
+        "span_name": "gemini.generate_content",
+    },
+    {
+        "package": "blaxel_langgraph.custom.gemini",
+        "object": "GeminiRestClient",
+        "method": "generate_content_async",
+        "span_name": "gemini.generate_content_async",
+    },
+]
+
+
+def should_send_prompts():
+    return (
+        os.getenv("TRACELOOP_TRACE_CONTENT") or "true"
+    ).lower() == "true" or context_api.get_value("override_enable_content_tracing")
+
+
+def is_streaming_response(response):
+    return isinstance(response, types.GeneratorType)
+
+
+def is_async_streaming_response(response):
+    return isinstance(response, types.AsyncGeneratorType)
+
+
+def _set_span_attribute(span, name, value):
+    if value is not None:
+        if value != "":
+            span.set_attribute(name, value)
+    return
+
+
+def _set_input_attributes(span, args, kwargs, llm_model):
+    if not should_send_prompts():
+        return
+
+    if "contents" in kwargs:
+        contents = kwargs["contents"]
+        if isinstance(contents, list):
+            for i, content in enumerate(contents):
+                if hasattr(content, "parts"):
+                    for part in content.parts:
+                        if hasattr(part, "text"):
+                            _set_span_attribute(
+                                span,
+                                f"{SpanAttributes.LLM_PROMPTS}.{i}.content",
+                                part.text,
+                            )
+                            _set_span_attribute(
+                                span,
+                                f"{SpanAttributes.LLM_PROMPTS}.{i}.role",
+                                getattr(content, "role", "user"),
+                            )
+    elif args and len(args) > 0:
+        prompt = ""
+        for arg in args:
+            if isinstance(arg, str):
+                prompt = f"{prompt}{arg}\n"
+            elif isinstance(arg, list):
+                for subarg in arg:
+                    prompt = f"{prompt}{subarg}\n"
+        if prompt:
+            _set_span_attribute(
+                span,
+                f"{SpanAttributes.LLM_PROMPTS}.0.content",
+                prompt,
+            )
+            _set_span_attribute(
+                span,
+                f"{SpanAttributes.LLM_PROMPTS}.0.role",
+                "user",
+            )
+    elif "prompt" in kwargs:
+        _set_span_attribute(
+            span,
+            f"{SpanAttributes.LLM_PROMPTS}.0.content",
+            kwargs["prompt"],
+        )
+        _set_span_attribute(
+            span,
+            f"{SpanAttributes.LLM_PROMPTS}.0.role",
+            "user",
+        )
+
+    _set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, llm_model)
+    _set_span_attribute(
+        span, SpanAttributes.LLM_REQUEST_TEMPERATURE, kwargs.get("temperature")
+    )
+    _set_span_attribute(
+        span, SpanAttributes.LLM_REQUEST_MAX_TOKENS, kwargs.get("max_output_tokens")
+    )
+    _set_span_attribute(span, SpanAttributes.LLM_REQUEST_TOP_P, kwargs.get("top_p"))
+    _set_span_attribute(span, SpanAttributes.LLM_TOP_K, kwargs.get("top_k"))
+    _set_span_attribute(
+        span, SpanAttributes.LLM_PRESENCE_PENALTY, kwargs.get("presence_penalty")
+    )
+    _set_span_attribute(
+        span, SpanAttributes.LLM_FREQUENCY_PENALTY, kwargs.get("frequency_penalty")
+    )
+
+    return
+
+
+@dont_throw
+def _set_response_attributes(span, response, llm_model):
+    _set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, llm_model)
+
+    if hasattr(response, "usage_metadata"):
+        _set_span_attribute(
+            span,
+            SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
+            response.usage_metadata.total_token_count,
+        )
+        _set_span_attribute(
+            span,
+            SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
+            response.usage_metadata.candidates_token_count,
+        )
+        _set_span_attribute(
+            span,
+            SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
+            response.usage_metadata.prompt_token_count,
+        )
+
+        if isinstance(response.text, list):
+            for index, item in enumerate(response):
+                prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{index}"
+                _set_span_attribute(span, f"{prefix}.content", item.text)
+                _set_span_attribute(span, f"{prefix}.role", "assistant")
+        elif isinstance(response.text, str):
+            _set_span_attribute(
+                span, f"{SpanAttributes.LLM_COMPLETIONS}.0.content", response.text
+            )
+            _set_span_attribute(
+                span, f"{SpanAttributes.LLM_COMPLETIONS}.0.role", "assistant"
+            )
+    else:
+        if isinstance(response, list):
+            for index, item in enumerate(response):
+                prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{index}"
+                _set_span_attribute(span, f"{prefix}.content", item)
+                _set_span_attribute(span, f"{prefix}.role", "assistant")
+        elif isinstance(response, str):
+            _set_span_attribute(
+                span, f"{SpanAttributes.LLM_COMPLETIONS}.0.content", response
+            )
+            _set_span_attribute(
+                span, f"{SpanAttributes.LLM_COMPLETIONS}.0.role", "assistant"
+            )
+
+    return
+
+
+def _build_from_streaming_response(span, response, llm_model):
+    complete_response = ""
+    for item in response:
+        item_to_yield = item
+        complete_response += str(item.text)
+
+        yield item_to_yield
+
+    _set_response_attributes(span, complete_response, llm_model)
+
+    span.set_status(Status(StatusCode.OK))
+    span.end()
+
+
+async def _abuild_from_streaming_response(span, response, llm_model):
+    complete_response = ""
+    async for item in response:
+        item_to_yield = item
+        complete_response += str(item.text)
+
+        yield item_to_yield
+
+    _set_response_attributes(span, complete_response, llm_model)
+
+    span.set_status(Status(StatusCode.OK))
+    span.end()
+
+
+# @dont_throw
+def _handle_request(span, args, kwargs, llm_model):
+    if span.is_recording():
+        _set_input_attributes(span, args, kwargs, llm_model)
+
+
+@dont_throw
+def _handle_response(span, response, llm_model):
+    if span.is_recording():
+        _set_response_attributes(span, response, llm_model)
+
+        span.set_status(Status(StatusCode.OK))
+
+
+def _with_tracer_wrapper(func):
+    """Helper for providing tracer for wrapper functions."""
+
+    def _with_tracer(tracer, to_wrap):
+        def wrapper(wrapped, instance, args, kwargs):
+            return func(tracer, to_wrap, wrapped, instance, args, kwargs)
+
+        return wrapper
+
+    return _with_tracer
+
+
+@_with_tracer_wrapper
+async def _awrap(tracer, to_wrap, wrapped, instance, args, kwargs):
+    """Instruments and calls every function defined in TO_WRAP."""
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY) or context_api.get_value(
+        SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY
+    ):
+        return await wrapped(*args, **kwargs)
+
+    llm_model = "unknown"
+    if hasattr(instance, "_model_id"):
+        llm_model = instance._model_id.replace("models/", "")
+    if hasattr(instance, "_model_name"):
+        llm_model = instance._model_name.replace(
+            "publishers/google/models/", ""
+        ).replace("models/", "")
+    if hasattr(instance, "model") and hasattr(instance.model, "model_name"):
+        llm_model = instance.model.model_name.replace("models/", "")
+
+    name = to_wrap.get("span_name")
+    span = tracer.start_span(
+        name,
+        kind=SpanKind.CLIENT,
+        attributes={
+            SpanAttributes.LLM_SYSTEM: "Gemini",
+            SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
+        },
+    )
+
+    _handle_request(span, args, kwargs, llm_model)
+
+    response = await wrapped(*args, **kwargs)
+
+    if response:
+        if is_streaming_response(response):
+            return _build_from_streaming_response(span, response, llm_model)
+        elif is_async_streaming_response(response):
+            return _abuild_from_streaming_response(span, response, llm_model)
+        else:
+            _handle_response(span, response, llm_model)
+
+    span.end()
+    return response
+
+
+@_with_tracer_wrapper
+def _wrap(tracer, to_wrap, wrapped, instance, args, kwargs):
+    """Instruments and calls every function defined in TO_WRAP."""
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY) or context_api.get_value(
+        SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY
+    ):
+        return wrapped(*args, **kwargs)
+
+    llm_model = "unknown"
+    if hasattr(instance, "_model_id"):
+        llm_model = instance._model_id.replace("models/", "")
+    if hasattr(instance, "_model_name"):
+        llm_model = instance._model_name.replace(
+            "publishers/google/models/", ""
+        ).replace("models/", "")
+    if hasattr(instance, "model") and hasattr(instance.model, "model_name"):
+        llm_model = instance.model.model_name.replace("models/", "")
+
+    name = to_wrap.get("span_name")
+    span = tracer.start_span(
+        name,
+        kind=SpanKind.CLIENT,
+        attributes={
+            SpanAttributes.LLM_SYSTEM: "Gemini",
+            SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
+        },
+    )
+
+    _handle_request(span, args, kwargs, llm_model)
+
+    response = wrapped(*args, **kwargs)
+
+    if response:
+        if is_streaming_response(response):
+            return _build_from_streaming_response(span, response, llm_model)
+        elif is_async_streaming_response(response):
+            return _abuild_from_streaming_response(span, response, llm_model)
+        else:
+            _handle_response(span, response, llm_model)
+
+    span.end()
+    return response
+
+
+class BlaxelLanggraphGeminiInstrumentor(BaseInstrumentor):
+    """An instrumentor for Google Generative AI's client library."""
+
+    def __init__(self, exception_logger=None):
+        super().__init__()
+        Config.exception_logger = exception_logger
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["blaxel_langgraph"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, __version__, tracer_provider)
+        for wrapped_method in WRAPPED_METHODS:
+            wrap_package = wrapped_method.get("package")
+            wrap_object = wrapped_method.get("object")
+            wrap_method = wrapped_method.get("method")
+
+            wrap_function_wrapper(
+                wrap_package,
+                f"{wrap_object}.{wrap_method}",
+                (
+                    _awrap(tracer, wrapped_method)
+                    if wrap_method == "generate_content_async"
+                    else _wrap(tracer, wrapped_method)
+                ),
+            )
+
+    def _uninstrument(self, **kwargs):
+        for wrapped_method in WRAPPED_METHODS:
+            wrap_package = wrapped_method.get("package")
+            wrap_object = wrapped_method.get("object")
+            unwrap(
+                f"{wrap_package}.{wrap_object}",
+                wrapped_method.get("method", ""),
+            )

blaxel/telemetry/instrumentation/blaxel_llamaindex.py

@@ -0,0 +1,89 @@
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as package_version
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.llamaindex.utils import (
+    _with_tracer_wrapper,
+    process_request,
+    process_response,
+    start_as_current_span_async,
+)
+from opentelemetry.semconv_ai import SpanAttributes, TraceloopSpanKindValues
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+
+TO_INSTRUMENT = [
+    {
+        "class": "FunctionTool",
+        "v9_module": "llama_index.tools.function_tool",
+        "v10_module": "llama_index.core.tools.function_tool",
+        "v10_legacy_module": "llama_index.legacy.tools.function_tool",
+    },
+    {
+        "class": "QueryEngineTool",
+        "v9_module": "llama_index.tools.query_engine",
+        "v10_module": "llama_index.core.tools.query_engine",
+        "v10_legacy_module": "llama_index.legacy.tools.query_engine",
+    },
+]
+
+
+class BlaxelLlamaIndexInstrumentor(BaseInstrumentor):
+    def instrumentation_dependencies(self):
+        return ["blaxel_llamaindex", "blaxel.core"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "0.2.0", tracer_provider)
+
+        for module in TO_INSTRUMENT:
+            try:
+                package_version("llama-index-core")
+                self._instrument_module(module["v10_module"], module["class"], tracer)
+                self._instrument_module(module["v10_legacy_module"], module["class"], tracer)
+
+            except PackageNotFoundError:
+                self._instrument_module(module["v9_module"], module["class"], tracer)
+
+    def _uninstrument(self, **kwargs):
+        pass
+
+    def _instrument_module(self, module_name, class_name, tracer):
+        wrap_function_wrapper(
+            module_name, f"{class_name}.call", query_wrapper(tracer)
+        )
+        wrap_function_wrapper(
+            module_name, f"{class_name}.acall", aquery_wrapper(tracer)
+        )
+
+
+@_with_tracer_wrapper
+def query_wrapper(tracer, wrapped, instance, args, kwargs):
+    name = instance.__class__.__name__
+    with tracer.start_as_current_span(f"{name}.tool") as span:
+        span.set_attribute(
+            SpanAttributes.TRACELOOP_SPAN_KIND,
+            TraceloopSpanKindValues.TOOL.value,
+        )
+        span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, name)
+
+        process_request(span, args, kwargs)
+        res = wrapped(*args, **kwargs)
+        process_response(span, res)
+        return res
+
+
+@_with_tracer_wrapper
+async def aquery_wrapper(tracer, wrapped, instance, args, kwargs):
+    name = instance.__class__.__name__
+    async with start_as_current_span_async(tracer=tracer, name=f"{name}.tool") as span:
+        span.set_attribute(
+            SpanAttributes.TRACELOOP_SPAN_KIND,
+            TraceloopSpanKindValues.TOOL.value,
+        )
+        span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, name)
+
+        process_request(span, args, kwargs)
+        res = await wrapped(*args, **kwargs)
+        process_response(span, res)
+        return res