[BUG] SmolagentsInstrumentor does not create model spans when using CodeAgent

[igor-gorohovsky]

Hello!

I'm trying to use opentelemetry-instrumentation-smolagents with this library, but unfortunately, it's not working as expected.

I noticed that maintainers from this repository have contributed to the instrumentation. Therefore, I wanted to bring this issue to your attention in case you could provide any insights.

More details and steps to reproduce can be found in the related issue in the OpenTelemetry repository:
Arize-ai/openinference#1636

Thank you!

[igor-gorohovsky]

If you can outline the vision about the fix, I’m happy to contribute. Right now, I’m not sure which side should be addressed.

[knutole]

This is a workaround that gets this result:

It involves using a custom instrumentor and wrapper, like so:

from phoenix.otel import register
from custom_instrumentor import SmolagentsInstrumentor

register()
SmolagentsInstrumentor().instrument()

The two files are edits of openinference.instrumentation.smolagents.SmolagentsInstrumentor and openinference.instrumentation.smolagents._wrappers.

# custom_instrumentor.py aka
# openinference.instrumentation.smolagents.SmolagentsInstrumentor

from typing import Any, Callable, Collection, Optional

from opentelemetry import trace as trace_api
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor  # type: ignore
from wrapt import wrap_function_wrapper

from smolagents.models import LiteLLMModel

from openinference.instrumentation import (
    OITracer,
    TraceConfig,
)
from custom_wrappers import (
    _ModelWrapper,
    _RunWrapper,
    _ToolCallWrapper,
    _StepWrapper,
)
from openinference.instrumentation.smolagents.version import __version__

_instruments = ("smolagents >= 1.2.2.dev0",)


class SmolagentsInstrumentor(BaseInstrumentor):  # type: ignore
    __slots__ = (
        "_original_run_method",
        "_original_step_methods",
        "_original_tool_call_method",
        "_original_model_call_methods",
        "_tracer",
    )

    def instrumentation_dependencies(self) -> Collection[str]:
        return _instruments

    def _instrument(self, **kwargs: Any) -> None:
        import smolagents
        from smolagents import MultiStepAgent, Tool, models

        if not (tracer_provider := kwargs.get("tracer_provider")):
            tracer_provider = trace_api.get_tracer_provider()
        if not (config := kwargs.get("config")):
            config = TraceConfig()
        else:
            assert isinstance(config, TraceConfig)

        self._tracer = OITracer(
            trace_api.get_tracer(__name__, __version__, tracer_provider),
            config=config,
        )

        self._original_run_method = getattr(MultiStepAgent, "run", None)
        wrap_function_wrapper(
            module="smolagents",
            name="MultiStepAgent.run",
            wrapper=_RunWrapper(tracer=self._tracer),
        )

        self._original_step_method = getattr(MultiStepAgent, "_execute_step", None)
        wrap_function_wrapper(
            module="smolagents",
            name="MultiStepAgent._execute_step",
            wrapper=_StepWrapper(tracer=self._tracer),
        )

        self._original_model_call_methods: Optional[dict[type, Callable[..., Any]]] = {}
        exported_model_subclasses = [
            attr
            for _, attr in vars(smolagents).items()
            if isinstance(attr, type) and issubclass(attr, models.Model)
        ]

        for model_subclass in exported_model_subclasses:
            method = "generate" if model_subclass is LiteLLMModel else "__call__"
            self._original_model_call_methods[model_subclass] = getattr(
                model_subclass, method
            )
            wrap_function_wrapper(
                module="smolagents",
                name=model_subclass.__name__ + "." + method,
                wrapper=_ModelWrapper(tracer=self._tracer),
            )

        self._original_tool_call_method = getattr(Tool, "__call__", None)
        wrap_function_wrapper(
            module="smolagents",
            name="Tool.__call__",
            wrapper=_ToolCallWrapper(tracer=self._tracer),
        )

    def _uninstrument(self, **kwargs: Any) -> None:
        from smolagents import MultiStepAgent, Tool

        if self._original_run_method is not None:
            MultiStepAgent.run = self._original_run_method
            self._original_run_method = None

        if self._original_step_method is not None:
            MultiStepAgent._execute_step = self._original_step_method
            self._original_step_method = None

        if self._original_model_call_methods is not None:
            for (
                model_subclass,
                original_model_call_method,
            ) in self._original_model_call_methods.items():
                method = "generate" if model_subclass is LiteLLMModel else "__call__"
                setattr(model_subclass, method, original_model_call_method)
            self._original_model_call_methods = None

        if self._original_tool_call_method is not None:
            Tool.__call__ = self._original_tool_call_method
            self._original_tool_call_method = None

# custom_wrappers.py aka
# openinference.instrumentation.smolagents._wrappers

from enum import Enum
from inspect import signature
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    Iterator,
    Mapping,
    Optional,
    Tuple,
)

from opentelemetry import context as context_api
from opentelemetry import trace as trace_api
from opentelemetry.util.types import AttributeValue

from openinference.instrumentation import get_attributes_from_context, safe_json_dumps
from openinference.semconv.trace import (
    MessageAttributes,
    OpenInferenceMimeTypeValues,
    OpenInferenceSpanKindValues,
    SpanAttributes,
    ToolAttributes,
    ToolCallAttributes,
)

import threading


# Add this thread-local storage for step context
_step_context: threading.local = threading.local()

if TYPE_CHECKING:
    from smolagents.tools import Tool  # type: ignore[import-untyped]


def _flatten(
    mapping: Optional[Mapping[str, Any]],
) -> Iterator[Tuple[str, AttributeValue]]:
    if not mapping:
        return
    for key, value in mapping.items():
        if value is None:
            continue
        if isinstance(value, Mapping):
            for sub_key, sub_value in _flatten(value):
                yield f"{key}.{sub_key}", sub_value
        elif isinstance(value, list) and any(
            isinstance(item, Mapping) for item in value
        ):
            for index, sub_mapping in enumerate(value):
                for sub_key, sub_value in _flatten(sub_mapping):
                    yield f"{key}.{index}.{sub_key}", sub_value
        else:
            if isinstance(value, Enum):
                value = value.value
            yield key, value


def _get_input_value(method: Callable[..., Any], *args: Any, **kwargs: Any) -> str:
    arguments = _bind_arguments(method, *args, **kwargs)
    arguments = _strip_method_args(arguments)
    return safe_json_dumps(arguments)


def _bind_arguments(
    method: Callable[..., Any], *args: Any, **kwargs: Any
) -> Dict[str, Any]:
    method_signature = signature(method)
    bound_args = method_signature.bind(*args, **kwargs)
    bound_args.apply_defaults()
    return bound_args.arguments


def _strip_method_args(arguments: Mapping[str, Any]) -> dict[str, Any]:
    return {
        key: value for key, value in arguments.items() if key not in ("self", "cls")
    }


def _smolagent_run_attributes(
    agent: Any, arguments: dict[str, Any]
) -> Iterator[Tuple[str, AttributeValue]]:
    if task := agent.task:
        yield "smolagents.task", task
    if additional_args := arguments.get("additional_args"):
        yield "smolagents.additional_args", safe_json_dumps(additional_args)
    yield "smolagents.max_steps", agent.max_steps
    yield "smolagents.tools_names", list(agent.tools.keys())
    for managed_agent_index, managed_agent in enumerate(agent.managed_agents.values()):
        yield f"smolagents.managed_agents.{managed_agent_index}.name", managed_agent.name
        yield (
            f"smolagents.managed_agents.{managed_agent_index}.description",
            managed_agent.description,
        )
        if getattr(managed_agent, "additional_prompting", None):
            yield (
                f"smolagents.managed_agents.{managed_agent_index}.additional_prompting",
                managed_agent.additional_prompting,
            )
        elif getattr(managed_agent, "managed_agent_prompt", None):
            yield (
                f"smolagents.managed_agents.{managed_agent_index}.managed_agent_prompt",
                managed_agent.managed_agent_prompt,
            )
        if getattr(managed_agent, "agent", None):
            yield (
                f"smolagents.managed_agents.{managed_agent_index}.max_steps",
                managed_agent.agent.max_steps,
            )
            yield (
                f"smolagents.managed_agents.{managed_agent_index}.tools_names",
                list(managed_agent.agent.tools.keys()),
            )


class _RunWrapper:
    def __init__(self, tracer: trace_api.Tracer) -> None:
        self._tracer = tracer

    def __call__(
        self,
        wrapped: Callable[..., Any],
        instance: Any,
        args: Tuple[Any, ...],
        kwargs: Mapping[str, Any],
    ) -> Any:
        if context_api.get_value(context_api._SUPPRESS_INSTRUMENTATION_KEY):
            return wrapped(*args, **kwargs)
        span_name = f"{instance.__class__.__name__}.run"
        agent = instance
        arguments = _bind_arguments(wrapped, *args, **kwargs)
        with self._tracer.start_as_current_span(
            span_name,
            attributes=dict(
                _flatten(
                    {
                        OPENINFERENCE_SPAN_KIND: AGENT,
                        INPUT_VALUE: _get_input_value(
                            wrapped,
                            *args,
                            **kwargs,
                        ),
                        **dict(_smolagent_run_attributes(agent, arguments)),
                        **dict(get_attributes_from_context()),
                    }
                )
            ),
        ) as span:
            agent_output = wrapped(*args, **kwargs)
            span.set_attribute(
                LLM_TOKEN_COUNT_PROMPT, agent.monitor.total_input_token_count
            )
            span.set_attribute(
                LLM_TOKEN_COUNT_COMPLETION, agent.monitor.total_output_token_count
            )
            span.set_attribute(
                LLM_TOKEN_COUNT_TOTAL,
                agent.monitor.total_input_token_count
                + agent.monitor.total_output_token_count,
            )
            span.set_status(trace_api.StatusCode.OK)
            span.set_attribute(OUTPUT_VALUE, str(agent_output))
        return agent_output


class _StepWrapper:
    def __init__(self, tracer: trace_api.Tracer) -> None:
        self._tracer = tracer

    def __call__(
        self,
        wrapped: Callable[..., Any],
        instance: Any,
        args: Tuple[Any, ...],
        kwargs: Mapping[str, Any],
    ) -> Any:
        if context_api.get_value(context_api._SUPPRESS_INSTRUMENTATION_KEY):
            return wrapped(*args, **kwargs)
        agent = instance
        span_name = f"Step {agent.step_number}"

        with self._tracer.start_as_current_span(
            span_name,
            attributes={
                OPENINFERENCE_SPAN_KIND: CHAIN,
                INPUT_VALUE: _get_input_value(wrapped, *args, **kwargs),
                **dict(get_attributes_from_context()),
            },
        ) as span:
            # Store the step span itself (not the context)
            _step_context.span = span

            try:
                result = wrapped(*args, **kwargs)
                step_log = args[0]  # ActionStep
                span.set_attribute(OUTPUT_VALUE, step_log.observations)
                if step_log.error is not None:
                    span.record_exception(step_log.error)
                span.set_status(trace_api.StatusCode.OK)

                return result
            finally:
                # Don't clear the span immediately - let Model/Tool calls use it
                pass  # We'll clear it when the next step starts


def _llm_input_messages(arguments: Mapping[str, Any]) -> Iterator[Tuple[str, Any]]:
    def process_message(idx: int, role: str, content: str) -> Iterator[Tuple[str, Any]]:
        yield f"{LLM_INPUT_MESSAGES}.{idx}.{MESSAGE_ROLE}", role
        yield f"{LLM_INPUT_MESSAGES}.{idx}.{MESSAGE_CONTENT}", content

    if isinstance(prompt := arguments.get("prompt"), str):
        yield from process_message(0, "user", prompt)
    elif isinstance(messages := arguments.get("messages"), list):
        for i, message in enumerate(messages):
            if not isinstance(message, dict):
                continue
            role, content = message.get("role"), message.get("content")
            if isinstance(content, list) and role:
                for subcontent in content:
                    if isinstance(subcontent, dict) and (
                        text := subcontent.get("text")
                    ):
                        yield from process_message(i, role, text)


def _llm_output_messages(output_message: Any) -> Iterator[Tuple[str, Any]]:
    if (role := getattr(output_message, "role", None)) is not None:
        yield (
            f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}",
            role,
        )
    if (content := getattr(output_message, "content", None)) is not None:
        yield (
            f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}",
            content,
        )
    if isinstance(tool_calls := getattr(output_message, "tool_calls", None), list):
        for tool_call_index, tool_call in enumerate(tool_calls):
            if (tool_call_id := getattr(tool_call, "id", None)) is not None:
                yield (
                    f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_ID}",
                    tool_call_id,
                )
            if (function := getattr(tool_call, "function", None)) is not None:
                if (name := getattr(function, "name", None)) is not None:
                    yield (
                        f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_NAME}",
                        name,
                    )
                if isinstance(arguments := getattr(function, "arguments", None), dict):
                    yield (
                        f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
                        safe_json_dumps(arguments),
                    )


def _output_value_and_mime_type(output: Any) -> Iterator[Tuple[str, Any]]:
    yield OUTPUT_MIME_TYPE, JSON
    yield OUTPUT_VALUE, output.model_dump_json()


def _llm_invocation_parameters(
    model: Any, arguments: Mapping[str, Any]
) -> Iterator[Tuple[str, Any]]:
    model_kwargs = _ if isinstance(_ := getattr(model, "kwargs", {}), dict) else {}
    kwargs = _ if isinstance(_ := arguments.get("kwargs"), dict) else {}
    yield LLM_INVOCATION_PARAMETERS, safe_json_dumps(model_kwargs | kwargs)


def _llm_tools(tools_to_call_from: list[Any]) -> Iterator[Tuple[str, Any]]:
    from smolagents import Tool
    from smolagents.models import get_tool_json_schema  # type: ignore[import-untyped]

    if not isinstance(tools_to_call_from, list):
        return
    for tool_index, tool in enumerate(tools_to_call_from):
        if isinstance(tool, Tool):
            yield (
                f"{LLM_TOOLS}.{tool_index}.{TOOL_JSON_SCHEMA}",
                safe_json_dumps(get_tool_json_schema(tool)),
            )


def _tools(tool: "Tool") -> Iterator[Tuple[str, Any]]:
    if tool_name := getattr(tool, "name", None):
        yield TOOL_NAME, tool_name
    if tool_description := getattr(tool, "description", None):
        yield TOOL_DESCRIPTION, tool_description
    yield TOOL_PARAMETERS, safe_json_dumps(tool.inputs)


def _input_value_and_mime_type(
    arguments: Mapping[str, Any],
) -> Iterator[Tuple[str, Any]]:
    yield INPUT_MIME_TYPE, JSON
    yield INPUT_VALUE, safe_json_dumps(arguments)


class _ModelWrapper:
    def __init__(self, tracer: trace_api.Tracer) -> None:
        self._tracer = tracer

    def __call__(
        self,
        wrapped: Callable[..., Any],
        instance: Any,
        args: Tuple[Any, ...],
        kwargs: Mapping[str, Any],
    ) -> Any:
        if context_api.get_value(context_api._SUPPRESS_INSTRUMENTATION_KEY):
            return wrapped(*args, **kwargs)
        arguments = _bind_arguments(wrapped, *args, **kwargs)
        span_name = f"{instance.__class__.__name__}.generate"
        model = instance

        step_span = getattr(_step_context, "span", None)
        if step_span:
            parent_context = trace_api.set_span_in_context(step_span)
        else:
            parent_context = context_api.get_current()

        with self._tracer.start_as_current_span(
            span_name,
            context=parent_context,
            attributes={
                OPENINFERENCE_SPAN_KIND: LLM,
                **dict(_input_value_and_mime_type(arguments)),
                **dict(_llm_invocation_parameters(instance, arguments)),
                **dict(_llm_input_messages(arguments)),
                **dict(get_attributes_from_context()),
            },
        ) as span:
            output_message = wrapped(*args, **kwargs)
            span.set_status(trace_api.StatusCode.OK)
            span.set_attribute(LLM_TOKEN_COUNT_PROMPT, model.last_input_token_count)
            span.set_attribute(
                LLM_TOKEN_COUNT_COMPLETION, model.last_output_token_count
            )
            span.set_attribute(LLM_MODEL_NAME, model.model_id)
            span.set_attribute(
                LLM_TOKEN_COUNT_TOTAL,
                model.last_input_token_count + model.last_output_token_count,
            )
            span.set_attributes(dict(_llm_output_messages(output_message)))
            span.set_attributes(
                dict(_llm_tools(arguments.get("tools_to_call_from", [])))
            )
            span.set_attributes(dict(_output_value_and_mime_type(output_message)))
        return output_message


class _ToolCallWrapper:
    def __init__(self, tracer: trace_api.Tracer) -> None:
        self._tracer = tracer

    def __call__(
        self,
        wrapped: Callable[..., Any],
        instance: Any,
        args: Tuple[Any, ...],
        kwargs: Mapping[str, Any],
    ) -> Any:
        if context_api.get_value(context_api._SUPPRESS_INSTRUMENTATION_KEY):
            return wrapped(*args, **kwargs)

        span_name = f"{instance.__class__.__name__}"
        step_span = getattr(_step_context, "span", None)
        if step_span:
            parent_context = trace_api.set_span_in_context(step_span)
        else:
            parent_context = context_api.get_current()

        with self._tracer.start_as_current_span(
            span_name,
            context=parent_context,
            attributes={
                OPENINFERENCE_SPAN_KIND: TOOL,
                INPUT_VALUE: _get_input_value(
                    wrapped,
                    *args,
                    **kwargs,
                ),
                **dict(_tools(instance)),
                **dict(get_attributes_from_context()),
            },
        ) as span:
            response = wrapped(*args, **kwargs)
            span.set_status(trace_api.StatusCode.OK)
            span.set_attributes(
                dict(
                    _output_value_and_mime_type_for_tool_span(
                        response=response,
                        output_type=instance.output_type,
                    )
                )
            )
        return response


def _output_value_and_mime_type_for_tool_span(
    response: Any, output_type: str
) -> Iterator[Tuple[str, Any]]:
    if output_type in (
        "string",
        "boolean",
        "integer",
        "number",
    ):
        yield OUTPUT_VALUE, response
        yield OUTPUT_MIME_TYPE, TEXT
    elif output_type == "object":
        yield OUTPUT_VALUE, safe_json_dumps(response)
        yield OUTPUT_MIME_TYPE, JSON

    # TODO: handle other types


# span attributes
INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE
INPUT_VALUE = SpanAttributes.INPUT_VALUE
LLM_INPUT_MESSAGES = SpanAttributes.LLM_INPUT_MESSAGES
LLM_INVOCATION_PARAMETERS = SpanAttributes.LLM_INVOCATION_PARAMETERS
LLM_MODEL_NAME = SpanAttributes.LLM_MODEL_NAME
LLM_OUTPUT_MESSAGES = SpanAttributes.LLM_OUTPUT_MESSAGES
LLM_PROMPTS = SpanAttributes.LLM_PROMPTS
LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION
LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT
LLM_TOKEN_COUNT_TOTAL = SpanAttributes.LLM_TOKEN_COUNT_TOTAL
LLM_TOOLS = SpanAttributes.LLM_TOOLS
OPENINFERENCE_SPAN_KIND = SpanAttributes.OPENINFERENCE_SPAN_KIND
OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE
OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE
TOOL_DESCRIPTION = SpanAttributes.TOOL_DESCRIPTION
TOOL_NAME = SpanAttributes.TOOL_NAME
TOOL_PARAMETERS = SpanAttributes.TOOL_PARAMETERS

# message attributes
MESSAGE_CONTENT = MessageAttributes.MESSAGE_CONTENT
MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON = (
    MessageAttributes.MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON
)
MESSAGE_FUNCTION_CALL_NAME = MessageAttributes.MESSAGE_FUNCTION_CALL_NAME
MESSAGE_NAME = MessageAttributes.MESSAGE_NAME
MESSAGE_ROLE = MessageAttributes.MESSAGE_ROLE
MESSAGE_TOOL_CALLS = MessageAttributes.MESSAGE_TOOL_CALLS

# mime types
JSON = OpenInferenceMimeTypeValues.JSON.value
TEXT = OpenInferenceMimeTypeValues.TEXT.value

# span kinds
AGENT = OpenInferenceSpanKindValues.AGENT.value
CHAIN = OpenInferenceSpanKindValues.CHAIN.value
LLM = OpenInferenceSpanKindValues.LLM.value
TOOL = OpenInferenceSpanKindValues.TOOL.value

# tool attributes
TOOL_JSON_SCHEMA = ToolAttributes.TOOL_JSON_SCHEMA

# tool call attributes
TOOL_CALL_FUNCTION_ARGUMENTS_JSON = ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON
TOOL_CALL_FUNCTION_NAME = ToolCallAttributes.TOOL_CALL_FUNCTION_NAME
TOOL_CALL_ID = ToolCallAttributes.TOOL_CALL_ID

[aymeric-roucher]

Sorry for leaving this onhold for so long!
I see you've added a fix here Arize-ai/openinference#1636 @igor-gorohovsky : thanks a lot for your help on this, it's really appreciated! 😃 Speaks a lot to the power of having involved community members.

If anyone else runs into the issue, upgrading the instrumentation package with pip install openinference-instrumentation-smolagents -U should fix it!

[albertvillanova]

@aymeric-roucher please note that the root cause here is a backward compatibility break: we no longer preserve the model's __call__ method and rely on generate instead.

I recall having suggested preserving __call__ in the original PR, but it seems this change was inadvertently introduced in a subsequent one. 😓

Maybe we should consider restoring __call__ to avoid breaking existing usages?

[aymeric-roucher]

@albertvillanova now the model usage in our main logic in agents.py is split between two options, generate and generate_stream, : we should keep it in the Model class for backwards compatibility (it's still there), but regarding the function called in agent logic, for me it makes sense to use generate or generate_stream directly depending on the need.
An alternative could be to reinstate __call__ usage in agents.py, but add to it a stream: bool argument that defaults to True? But we need to adapt the logic in the instrumentation package anyway.

[albertvillanova]

@aymeric-roucher, I see your point.

I think it is OK to keep things as they are now: internally we use 2 different methods; externally we allow users to keep calling model __call__ within their agent subclasses.