openinference-instrumentation-beeai 0.1.6__py3-none-any.whl → 0.1.8__py3-none-any.whl

openinference/instrumentation/beeai/__init__.py

@@ -1,23 +1,10 @@
-# Copyright 2025 IBM Corp.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 import logging
-from importlib import import_module
 from importlib.metadata import PackageNotFoundError, version
-from typing import Any, Callable, Collection, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, Collection
+
+if TYPE_CHECKING:
+    from beeai_framework.emitter import EventMeta
 
-import wrapt
 from opentelemetry import trace as trace_api
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor  # type: ignore
 
@@ -25,13 +12,14 @@ from openinference.instrumentation import (
     OITracer,
     TraceConfig,
 )
-from openinference.semconv.trace import OpenInferenceSpanKindValues
-
-from .middleware import create_telemetry_middleware
+from openinference.instrumentation.beeai._span import SpanWrapper
+from openinference.instrumentation.beeai._utils import _datetime_to_span_time, exception_handler
+from openinference.instrumentation.beeai.processors.base import Processor
+from openinference.instrumentation.beeai.processors.locator import ProcessorLocator
 
 logger = logging.getLogger(__name__)
 
-_instruments = ("beeai-framework >= 0.1.10",)
+_instruments = ("beeai-framework >= 0.1.32",)
 try:
     __version__ = version("beeai-framework")
 except PackageNotFoundError:
@@ -39,14 +27,12 @@ except PackageNotFoundError:
 
 
 class BeeAIInstrumentor(BaseInstrumentor):  # type: ignore
-    __slots__ = (
-        "_original_react_agent_run",
-        "_original_tool_calling_agent_run",
-        "_original_chat_model_create",
-        "_original_chat_model_create_structure",
-        "_original_tool_run",
-        "_tracer",
-    )
+    __slots__ = ("_tracer", "_cleanup", "_processes")
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._cleanup: Callable[[], None] = lambda: None
+        self._processes: dict[str, Processor] = {}
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
@@ -60,89 +46,77 @@ class BeeAIInstrumentor(BaseInstrumentor):  # type: ignore
             else:
                 assert isinstance(config, TraceConfig)
 
-            from beeai_framework.agents.base import BaseAgent
-            from beeai_framework.agents.react.agent import ReActAgent
-            from beeai_framework.agents.tool_calling.agent import ToolCallingAgent
-            from beeai_framework.backend.chat import ChatModel
-            from beeai_framework.tools import Tool
-
             self._tracer = OITracer(
                 trace_api.get_tracer(__name__, __version__, tracer_provider),
                 config=config,
             )
 
-            F = TypeVar("F", bound=Callable[..., Any])
-
-            @wrapt.decorator  # type: ignore
-            def run_wrapper(
-                wrapped: F, instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
-            ) -> Any:
-                result = wrapped(*args, **kwargs)
-
-                span_kind = OpenInferenceSpanKindValues.UNKNOWN
-                if isinstance(instance, ChatModel):
-                    span_kind = OpenInferenceSpanKindValues.LLM
-                if isinstance(instance, BaseAgent):
-                    span_kind = OpenInferenceSpanKindValues.AGENT
-                if isinstance(instance, Tool):
-                    span_kind = OpenInferenceSpanKindValues.TOOL
-
-                if result.middleware:
-                    result.middleware(create_telemetry_middleware(self._tracer, span_kind.value))
+            from beeai_framework.emitter import Emitter, EmitterOptions
 
-                return result
-
-            ## Agent support
-            self._original_react_agent_run = getattr(
-                import_module("beeai_framework.agents.react.agent"), "run", None
-            )
-            setattr(ReActAgent, "run", run_wrapper(ReActAgent.run))
-            self._original_tool_calling_agent_run = getattr(
-                import_module("beeai_framework.agents.tool_calling.agent"), "run", None
-            )
-            setattr(ToolCallingAgent, "run", run_wrapper(ToolCallingAgent.run))
-            ## LLM support
-            self._original_chat_model_create = getattr(
-                import_module("beeai_framework.backend.chat"), "create", None
-            )
-            setattr(ChatModel, "create", run_wrapper(ChatModel.create))
-            self._original_chat_model_create_structure = getattr(
-                import_module("beeai_framework.backend.chat"), "create_structure", None
+            self._cleanup = Emitter.root().match(
+                "*.*",
+                self._handler,
+                EmitterOptions(match_nested=True, is_blocking=True),
             )
-            setattr(ChatModel, "create_structure", run_wrapper(ChatModel.create_structure))
-
-            ## Tool support
-            self._original_tool_run = getattr(import_module("beeai_framework.tools"), "run", None)
-            setattr(Tool, "run", run_wrapper(Tool.run))
-
-        except ImportError as e:
-            logger.error("ImportError during instrumentation", exc_info=e)
         except Exception as e:
             logger.error("Instrumentation error", exc_info=e)
 
     def _uninstrument(self, **kwargs: Any) -> None:
-        if self._original_react_agent_run is not None:
-            from beeai_framework.agents.react.agent import ReActAgent
-
-            setattr(ReActAgent, "run", self._original_react_agent_run)
-            self._original_react_agent_run = None
-        if self._original_tool_calling_agent_run is not None:
-            from beeai_framework.agents.tool_calling.agent import ToolCallingAgent
-
-            setattr(ToolCallingAgent, "run", self._original_tool_calling_agent_run)
-            self._original_tool_calling_agent_run = None
-        if self._original_chat_model_create is not None:
-            from beeai_framework.backend.chat import ChatModel
-
-            setattr(ChatModel, "create", self._original_chat_model_create)
-            self._original_chat_model_create = None
-        if self._original_chat_model_create_structure is not None:
-            from beeai_framework.backend.chat import ChatModel
-
-            setattr(ChatModel, "create_structure", self._original_chat_model_create_structure)
-            self._original_chat_model_create_structure = None
-        if self._original_tool_run is not None:
-            from beeai_framework.tools import Tool
-
-            setattr(Tool, "run", self._original_tool_run)
-            self._original_tool_run = None
+        self._cleanup()
+        self._processes.clear()
+
+    def _build_tree(self, node: SpanWrapper) -> None:
+        with self._tracer.start_as_current_span(
+            name=node.name,
+            openinference_span_kind=node.kind,
+            attributes=node.attributes,
+            start_time=_datetime_to_span_time(node.started_at) if node.started_at else None,
+            end_on_exit=False,  # we do it manually
+        ) as current_span:
+            for event in node.events:
+                current_span.add_event(
+                    name=event.name, attributes=event.attributes, timestamp=event.timestamp
+                )
+
+            for children in node.children:
+                self._build_tree(children)
+
+            current_span.set_status(node.status)
+            if node.error is not None:
+                current_span.record_exception(node.error)
+
+            current_span.end(_datetime_to_span_time(node.ended_at) if node.ended_at else None)
+
+    @exception_handler
+    async def _handler(self, data: Any, event: "EventMeta") -> None:
+        assert event.trace is not None, "Event must have a trace"
+
+        if event.trace.run_id not in self._processes:
+            parent = (
+                self._processes.get(event.trace.parent_run_id)
+                if event.trace.parent_run_id
+                else None
+            )
+            if event.trace.parent_run_id and not parent:
+                raise ValueError(f"Parent run with ID {event.trace.parent_run_id} was not found!")
+
+            node = self._processes[event.trace.run_id] = ProcessorLocator.locate(data, event)
+            if parent is not None:
+                parent.span.children.append(node.span)
+        else:
+            node = self._processes[event.trace.run_id]
+
+        from beeai_framework.context import RunContextFinishEvent
+
+        if isinstance(data, RunContextFinishEvent):
+            await node.end(data, event)
+            self._build_tree(node.span)
+            self._processes.pop(event.trace.run_id)
+        else:
+            if event.context.get("internal"):
+                return
+
+            await node.update(
+                data,
+                event,
+            )

openinference/instrumentation/beeai/_span.py

@@ -0,0 +1,81 @@
+from datetime import datetime
+from typing import TYPE_CHECKING, Any
+
+from opentelemetry.sdk.trace import Event
+from opentelemetry.trace import StatusCode
+
+from openinference.instrumentation.beeai._utils import _datetime_to_span_time, _unpack_object
+
+if TYPE_CHECKING:
+    from beeai_framework.emitter import EventMeta
+
+from openinference.semconv.trace import (
+    OpenInferenceMimeTypeValues,
+    OpenInferenceSpanKindValues,
+    SpanAttributes,
+)
+
+
+class SpanWrapper:
+    def __init__(self, *, name: str, kind: OpenInferenceSpanKindValues):
+        self.name = name
+        self.attributes: dict[str, Any] = {SpanAttributes.OPENINFERENCE_SPAN_KIND: kind}
+        self.events: list[Event] = []
+        self.status: StatusCode = StatusCode.OK
+        self.error: Exception | None = None
+        self.started_at: datetime | None = None
+        self.ended_at: datetime | None = None
+        self.children: list["SpanWrapper"] = []
+        self.kind = kind
+
+    def child(
+        self, name: str | None = None, event: tuple[Any, "EventMeta"] | None = None
+    ) -> "SpanWrapper":
+        child = SpanWrapper(name=name or f"{self.name}_child", kind=self.kind)
+        if event is not None:
+            from beeai_framework.utils.dicts import include_keys
+            from beeai_framework.utils.strings import to_json
+
+            value, meta = event
+
+            child.started_at = meta.created_at
+            child.ended_at = meta.created_at
+            child.attributes.update(
+                {
+                    SpanAttributes.INPUT_VALUE: to_json(value) if value is not None else value,
+                    SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value,
+                    **_unpack_object(
+                        include_keys(meta.model_dump(), {"id", "context", "path", "trace"}),
+                        prefix=SpanAttributes.METADATA,
+                    ),
+                }
+            )
+
+        self.children.append(child)
+        return child
+
+    def add_event(
+        self,
+        name: str,
+        attributes: dict[str, Any] | None = None,
+        timestamp: datetime | None = None,
+    ) -> None:
+        self.events.append(
+            Event(
+                name=name,
+                attributes=attributes or {},
+                timestamp=_datetime_to_span_time(timestamp) if timestamp else None,
+            )
+        )
+
+    def set_attribute(self, name: str, value: Any) -> None:
+        self.attributes[name] = value
+
+    def set_attributes(self, attributes: dict[str, Any]) -> None:
+        self.attributes.update(attributes)
+
+    def set_status(self, status: StatusCode) -> None:
+        self.status = status
+
+    def record_exception(self, error: Exception) -> None:
+        self.error = error

openinference/instrumentation/beeai/_utils.py

@@ -0,0 +1,77 @@
+import datetime
+import functools
+import json
+import logging
+from collections.abc import Awaitable
+from typing import Any, Callable, ParamSpec, TypeVar
+
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+def _datetime_to_span_time(dt: datetime.datetime) -> int:
+    if dt.tzinfo is None or dt.tzinfo.utcoffset(dt) is None:
+        dt = dt.replace(tzinfo=datetime.timezone.utc)
+
+    return int(dt.timestamp() * 1_000_000_000)
+
+
+def _unpack_object(obj: dict[str, Any] | list[Any] | BaseModel, prefix: str = "") -> dict[str, Any]:
+    if not isinstance(obj, dict) and not isinstance(obj, list):
+        obj_ref = obj
+        obj = json.loads(stringify(obj))
+        if not isinstance(obj, dict) and not isinstance(obj, list):
+            logger.warning(f"Cannot unpack object of type {type(obj_ref)} (prefix={prefix})")
+            return {"value": str(obj)}
+
+    if prefix and prefix.startswith("."):
+        prefix = prefix[1:]
+    if prefix and not prefix.endswith("."):
+        prefix += "."
+
+    output = {}
+    for key, value in obj.items() if isinstance(obj, dict) else enumerate(obj):
+        if value is None:
+            continue
+        if is_primitive(value):
+            output[f"{prefix}{key}"] = str(value)
+        else:
+            output.update(_unpack_object(value, prefix=f"{prefix}{key}"))
+    return output
+
+
+def is_primitive(value: Any) -> bool:
+    return isinstance(value, str | bool | int | float | type(None))
+
+
+def stringify(value: Any, pretty: bool = False) -> str:
+    if is_primitive(value):
+        return str(value)
+
+    from beeai_framework.utils.strings import to_json
+
+    return to_json(value, sort_keys=False, indent=4 if pretty else None)
+
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+
+def exception_handler(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T | None]]:
+    @functools.wraps(func)
+    async def wrapped(*args: P.args, **kwargs: P.kwargs) -> T | None:
+        try:
+            return await func(*args, **kwargs)
+        except Exception as e:
+            logger.error("Error has occurred in the telemetry package.", exc_info=e)
+            return None
+
+    return wrapped
+
+
+def safe_dump_model_schema(model: type[BaseModel]) -> dict[str, Any]:
+    try:
+        return model.model_json_schema(mode="serialization")
+    except:  # noqa: E722
+        return {}

openinference/instrumentation/beeai/processors/agents/base.py

@@ -0,0 +1,34 @@
+from typing import ClassVar
+
+from beeai_framework.agents import BaseAgent
+from beeai_framework.context import RunContext, RunContextStartEvent
+from beeai_framework.emitter import EventMeta
+
+from openinference.instrumentation.beeai._utils import _unpack_object
+from openinference.instrumentation.beeai.processors.base import Processor
+from openinference.semconv.trace import (
+    OpenInferenceSpanKindValues,
+    SpanAttributes,
+)
+
+
+class AgentProcessor(Processor):
+    kind: ClassVar[OpenInferenceSpanKindValues] = OpenInferenceSpanKindValues.AGENT
+
+    def __init__(self, event: RunContextStartEvent, meta: "EventMeta") -> None:
+        super().__init__(event, meta)
+
+        assert isinstance(meta.creator, RunContext)
+        assert isinstance(meta.creator.instance, BaseAgent)
+
+        agent = meta.creator.instance
+        self.span.name = agent.meta.name or self.span.name
+        self.span.set_attributes(
+            {
+                SpanAttributes.AGENT_NAME: agent.meta.name,
+                f"{SpanAttributes.METADATA}.agent_description": agent.meta.description,
+            }
+        )
+        self.span.set_attributes(
+            _unpack_object(agent.memory.to_json_safe(), prefix=f"{SpanAttributes.METADATA}.memory")
+        )

openinference/instrumentation/beeai/processors/agents/react.py

@@ -0,0 +1,77 @@
+from typing import Any
+
+from beeai_framework.context import RunContextStartEvent
+from beeai_framework.emitter import EventMeta
+from typing_extensions import override
+
+from openinference.instrumentation.beeai._utils import _unpack_object, stringify
+from openinference.instrumentation.beeai.processors.agents.base import AgentProcessor
+from openinference.semconv.trace import OpenInferenceMimeTypeValues, SpanAttributes
+
+
+class ReActAgentProcessor(AgentProcessor):
+    def __init__(self, event: "RunContextStartEvent", meta: "EventMeta"):
+        super().__init__(event, meta)
+        self._last_chunk_index = 0
+
+    @override
+    async def update(
+        self,
+        event: Any,
+        meta: "EventMeta",
+    ) -> None:
+        from beeai_framework.agents.react import (
+            ReActAgentErrorEvent,
+            ReActAgentRetryEvent,
+            ReActAgentStartEvent,
+            ReActAgentSuccessEvent,
+            ReActAgentUpdateEvent,
+        )
+        from beeai_framework.backend import UserMessage
+
+        await super().update(event, meta)
+
+        self.span.add_event(f"{meta.name} ({meta.path})", timestamp=meta.created_at)
+
+        match event:
+            case ReActAgentStartEvent():
+                last_user_message: UserMessage | None = next(
+                    (
+                        msg
+                        for msg in reversed(event.memory.messages)
+                        if isinstance(msg, UserMessage)
+                    ),
+                    None,
+                )
+                self.span.set_attributes(
+                    {
+                        SpanAttributes.INPUT_VALUE: last_user_message.text
+                        if last_user_message
+                        else "",
+                        SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.TEXT.value,
+                    }
+                )
+            case ReActAgentUpdateEvent():
+                value = event.update.parsed_value or event.update.value
+                self.span.set_attribute(
+                    f"{SpanAttributes.METADATA}.iteration", event.meta.iteration
+                )
+                self.span.set_attributes(
+                    _unpack_object(
+                        {event.update.key: stringify(value)},
+                        prefix=f"{SpanAttributes.METADATA}.iterations.{event.meta.iteration}",
+                    )
+                )
+
+            case ReActAgentErrorEvent():
+                span = self.span.child(meta.name, event=(event, meta))
+                span.record_exception(event.error)
+            case ReActAgentRetryEvent():
+                self.span.child(meta.name, event=(event, meta))
+            case ReActAgentSuccessEvent():
+                self.span.set_attribute(SpanAttributes.OUTPUT_VALUE, event.data.text)
+                self.span.set_attribute(
+                    SpanAttributes.OUTPUT_MIME_TYPE, OpenInferenceMimeTypeValues.TEXT.value
+                )
+            case _:
+                self.span.child(meta.name, event=(event, meta))

openinference/instrumentation/beeai/processors/agents/requirement_agent.py

@@ -0,0 +1,71 @@
+from typing import Any
+
+from beeai_framework.agents.experimental.events import (
+    RequirementAgentStartEvent,
+    RequirementAgentSuccessEvent,
+)
+from beeai_framework.agents.experimental.types import RequirementAgentRunStateStep
+from beeai_framework.context import RunContextStartEvent
+from beeai_framework.emitter import EventMeta
+from typing_extensions import override
+
+from openinference.instrumentation.beeai._utils import _unpack_object
+from openinference.instrumentation.beeai.processors.agents.base import AgentProcessor
+from openinference.semconv.trace import OpenInferenceMimeTypeValues, SpanAttributes
+
+
+class RequirementAgentProcessor(AgentProcessor):
+    def __init__(self, event: "RunContextStartEvent", meta: "EventMeta"):
+        super().__init__(event, meta)
+        self._steps: list[RequirementAgentRunStateStep] = []
+
+    @override
+    async def update(
+        self,
+        event: Any,
+        meta: "EventMeta",
+    ) -> None:
+        await super().update(event, meta)
+        self.span.add_event(f"{meta.name} ({meta.path})", timestamp=meta.created_at)
+
+        match event:
+            case RequirementAgentStartEvent():
+                self._sync_steps(event.state.steps)
+            case RequirementAgentSuccessEvent():
+                self._sync_steps(event.state.steps)
+
+                if event.state.answer is not None:
+                    self.span.set_attributes(
+                        {
+                            SpanAttributes.OUTPUT_VALUE: event.state.answer.text,
+                            SpanAttributes.OUTPUT_MIME_TYPE: OpenInferenceMimeTypeValues.TEXT.value,
+                        }
+                    )
+                if event.state.result is not None:
+                    self.span.set_attributes(
+                        _unpack_object(
+                            event.state.result, prefix=f"{SpanAttributes.METADATA}.result"
+                        )
+                    )
+            case _:
+                self.span.child(meta.name, event=(event, meta))
+
+    def _sync_steps(self, steps: list["RequirementAgentRunStateStep"]) -> None:
+        new_items = steps[len(self._steps) :]
+        self._steps.extend(new_items)
+
+        for idx, (new_step, old_step) in enumerate(zip(steps, self._steps)):
+            if new_step.id != old_step.id:
+                # TODO: cleanup old keys
+                self._steps.clear()
+                self._sync_steps(steps)
+                return
+
+            if new_step in new_items or new_step != old_step:
+                self._steps[idx] = new_step
+                self.span.set_attributes(
+                    _unpack_object(
+                        new_step,
+                        prefix=f"{SpanAttributes.METADATA}.steps.{idx}",
+                    )
+                )

openinference/instrumentation/beeai/processors/agents/tool_calling.py

@@ -0,0 +1,34 @@
+from typing import Any
+
+from beeai_framework.agents.tool_calling import (
+    ToolCallingAgentStartEvent,
+    ToolCallingAgentSuccessEvent,
+)
+from beeai_framework.emitter import EventMeta
+from typing_extensions import override
+
+from openinference.instrumentation.beeai.processors.agents.base import AgentProcessor
+from openinference.semconv.trace import OpenInferenceMimeTypeValues, SpanAttributes
+
+
+class ToolCallingAgentProcessor(AgentProcessor):
+    @override
+    async def update(
+        self,
+        event: Any,
+        meta: "EventMeta",
+    ) -> None:
+        await super().update(event, meta)
+        self.span.add_event(f"{meta.name} ({meta.path})", timestamp=meta.created_at)
+
+        match event:
+            case ToolCallingAgentStartEvent():
+                pass
+            case ToolCallingAgentSuccessEvent():
+                if event.state.result is not None:
+                    self.span.set_attribute(SpanAttributes.OUTPUT_VALUE, event.state.result.text)
+                    self.span.set_attribute(
+                        SpanAttributes.OUTPUT_MIME_TYPE, OpenInferenceMimeTypeValues.TEXT.value
+                    )
+            case _:
+                self.span.child(meta.name, event=(event, meta))

openinference/instrumentation/beeai/processors/base.py

@@ -0,0 +1,60 @@
+from typing import TYPE_CHECKING, Any, ClassVar
+
+if TYPE_CHECKING:
+    from beeai_framework.context import RunContextFinishEvent, RunContextStartEvent
+    from beeai_framework.emitter import EventMeta
+
+from opentelemetry.trace import StatusCode
+
+from openinference.instrumentation.beeai._span import SpanWrapper
+from openinference.instrumentation.beeai._utils import stringify
+from openinference.semconv.trace import (
+    OpenInferenceMimeTypeValues,
+    OpenInferenceSpanKindValues,
+    SpanAttributes,
+)
+
+
+class Processor:
+    kind: ClassVar[OpenInferenceSpanKindValues] = OpenInferenceSpanKindValues.UNKNOWN
+
+    def __init__(self, event: "RunContextStartEvent", meta: "EventMeta"):
+        from beeai_framework.context import RunContext
+
+        assert isinstance(meta.creator, RunContext)
+        target_cls = type(meta.creator.instance)
+
+        self.span = SpanWrapper(name=target_cls.__name__, kind=type(self).kind)
+        self.span.started_at = meta.created_at
+        self.span.attributes.update(
+            {
+                SpanAttributes.OPENINFERENCE_SPAN_KIND: type(self).kind,
+                SpanAttributes.INPUT_VALUE: stringify(event.input),
+                SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value,
+                f"{SpanAttributes.METADATA}.class_name": target_cls.__name__,
+            }
+        )
+
+    async def update(
+        self,
+        event: Any,
+        meta: "EventMeta",
+    ) -> None:
+        pass
+
+    async def end(self, event: "RunContextFinishEvent", meta: "EventMeta") -> None:
+        if event.output is not None:
+            if SpanAttributes.OUTPUT_VALUE not in self.span.attributes:
+                self.span.attributes.update(
+                    {
+                        SpanAttributes.OUTPUT_VALUE: stringify(event.output),
+                        SpanAttributes.OUTPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value,
+                    }
+                )
+            self.span.set_status(StatusCode.OK)
+
+        if event.error is not None:
+            self.span.set_status(StatusCode.ERROR)
+            self.span.record_exception(event.error)
+
+        self.span.ended_at = meta.created_at