openinference-instrumentation-beeai 0.1.8__tar.gz → 0.1.10__tar.gz

{openinference_instrumentation_beeai-0.1.8 → openinference_instrumentation_beeai-0.1.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openinference-instrumentation-beeai
-Version: 0.1.8
+Version: 0.1.10
 Summary: OpenInference BeeAI Instrumentation
 Project-URL: Homepage, https://github.com/Arize-ai/openinference/tree/main/python/instrumentation/openinference-instrumentation-beeai
 Author: IBM Corp.

{openinference_instrumentation_beeai-0.1.8 → openinference_instrumentation_beeai-0.1.10}/src/openinference/instrumentation/beeai/__init__.py

@@ -1,6 +1,11 @@
+import contextlib
 import logging
 from importlib.metadata import PackageNotFoundError, version
-from typing import TYPE_CHECKING, Any, Callable, Collection
+from typing import TYPE_CHECKING, Any, Callable, Collection, Generator
+
+from opentelemetry.trace import StatusCode
+
+from openinference.instrumentation._spans import OpenInferenceSpan
 
 if TYPE_CHECKING:
     from beeai_framework.emitter import EventMeta
@@ -27,12 +32,13 @@ except PackageNotFoundError:
 
 
 class BeeAIInstrumentor(BaseInstrumentor):  # type: ignore
-    __slots__ = ("_tracer", "_cleanup", "_processes")
+    __slots__ = ("_tracer", "_cleanup", "_processes", "_processes_deps")
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self._cleanup: Callable[[], None] = lambda: None
         self._processes: dict[str, Processor] = {}
+        self._processes_deps: dict[str, list[Processor]] = {}
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
@@ -64,8 +70,16 @@ class BeeAIInstrumentor(BaseInstrumentor):  # type: ignore
     def _uninstrument(self, **kwargs: Any) -> None:
         self._cleanup()
         self._processes.clear()
+        self._processes_deps.clear()
 
-    def _build_tree(self, node: SpanWrapper) -> None:
+    def _build_tree(self, processor: Processor) -> None:
+        with self._build_tree_for_span(processor.span):
+            for child in self._processes_deps.pop(processor.run_id):
+                self._build_tree(child)
+        self._processes.pop(processor.run_id)
+
+    @contextlib.contextmanager
+    def _build_tree_for_span(self, node: SpanWrapper) -> Generator[OpenInferenceSpan, None, None]:
         with self._tracer.start_as_current_span(
             name=node.name,
             openinference_span_kind=node.kind,
@@ -73,23 +87,27 @@ class BeeAIInstrumentor(BaseInstrumentor):  # type: ignore
             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:
+            yield 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)
+                with self._build_tree_for_span(children):
+                    pass
 
             current_span.set_status(node.status)
-            if node.error is not None:
+            if node.error is not None and node.status == StatusCode.ERROR:
                 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 is None:
+            return
 
         if event.trace.run_id not in self._processes:
             parent = (
@@ -100,9 +118,10 @@ class BeeAIInstrumentor(BaseInstrumentor):  # type: ignore
             if event.trace.parent_run_id and not parent:
                 raise ValueError(f"Parent run with ID {event.trace.parent_run_id} was not found!")
 
+            self._processes_deps[event.trace.run_id] = []
             node = self._processes[event.trace.run_id] = ProcessorLocator.locate(data, event)
             if parent is not None:
-                parent.span.children.append(node.span)
+                self._processes_deps[parent.run_id].append(node)
         else:
             node = self._processes[event.trace.run_id]
 
@@ -110,8 +129,8 @@ class BeeAIInstrumentor(BaseInstrumentor):  # type: ignore
 
         if isinstance(data, RunContextFinishEvent):
             await node.end(data, event)
-            self._build_tree(node.span)
-            self._processes.pop(event.trace.run_id)
+            if event.trace.parent_run_id is None:
+                self._build_tree(node)
         else:
             if event.context.get("internal"):
                 return

{openinference_instrumentation_beeai-0.1.8 → openinference_instrumentation_beeai-0.1.10}/src/openinference/instrumentation/beeai/_span.py

@@ -77,5 +77,20 @@ class SpanWrapper:
     def set_status(self, status: StatusCode) -> None:
         self.status = status
 
+    def reset_exception(self) -> None:
+        self.error = None
+        self.set_status(StatusCode.OK)
+
     def record_exception(self, error: Exception) -> None:
+        from beeai_framework.errors import FrameworkError
+
         self.error = error
+        self.set_status(StatusCode.ERROR)
+        self.set_attributes(
+            {
+                SpanAttributes.OUTPUT_VALUE: error.explain()
+                if isinstance(error, FrameworkError)
+                else str(error),
+                SpanAttributes.OUTPUT_MIME_TYPE: OpenInferenceMimeTypeValues.TEXT.value,
+            }
+        )

{openinference_instrumentation_beeai-0.1.8 → openinference_instrumentation_beeai-0.1.10}/src/openinference/instrumentation/beeai/_utils.py

@@ -22,7 +22,7 @@ def _unpack_object(obj: dict[str, Any] | list[Any] | BaseModel, prefix: str = ""
         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})")
+            logger.debug(f"Cannot unpack object of type {type(obj_ref)} (prefix={prefix})")
             return {"value": str(obj)}
 
     if prefix and prefix.startswith("."):
@@ -34,7 +34,7 @@ def _unpack_object(obj: dict[str, Any] | list[Any] | BaseModel, prefix: str = ""
     for key, value in obj.items() if isinstance(obj, dict) else enumerate(obj):
         if value is None:
             continue
-        if is_primitive(value):
+        if is_primitive(value) or has_custom_str(value):
             output[f"{prefix}{key}"] = str(value)
         else:
             output.update(_unpack_object(value, prefix=f"{prefix}{key}"))
@@ -45,6 +45,13 @@ def is_primitive(value: Any) -> bool:
     return isinstance(value, str | bool | int | float | type(None))
 
 
+def has_custom_str(value: Any) -> bool:
+    if value.__class__.__module__ == "builtins":
+        return False
+
+    return value.__class__.__str__ is not object.__str__
+
+
 def stringify(value: Any, pretty: bool = False) -> str:
     if is_primitive(value):
         return str(value)

{openinference_instrumentation_beeai-0.1.8 → openinference_instrumentation_beeai-0.1.10}/src/openinference/instrumentation/beeai/processors/base.py

@@ -24,6 +24,9 @@ class Processor:
         assert isinstance(meta.creator, RunContext)
         target_cls = type(meta.creator.instance)
 
+        assert meta.trace is not None
+        self.run_id = meta.trace.run_id
+
         self.span = SpanWrapper(name=target_cls.__name__, kind=type(self).kind)
         self.span.started_at = meta.created_at
         self.span.attributes.update(
@@ -43,6 +46,9 @@ class Processor:
         pass
 
     async def end(self, event: "RunContextFinishEvent", meta: "EventMeta") -> None:
+        if event.error is not None:
+            self.span.record_exception(event.error)
+
         if event.output is not None:
             if SpanAttributes.OUTPUT_VALUE not in self.span.attributes:
                 self.span.attributes.update(
@@ -53,8 +59,4 @@ class Processor:
                 )
             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

{openinference_instrumentation_beeai-0.1.8 → openinference_instrumentation_beeai-0.1.10}/src/openinference/instrumentation/beeai/processors/tool.py

@@ -2,7 +2,7 @@ from typing import Any, ClassVar
 
 from beeai_framework.context import RunContext, RunContextStartEvent
 from beeai_framework.emitter import EventMeta
-from beeai_framework.tools import ToolErrorEvent, ToolRetryEvent, ToolSuccessEvent
+from beeai_framework.tools import ToolErrorEvent, ToolRetryEvent, ToolStartEvent, ToolSuccessEvent
 from beeai_framework.tools.tool import Tool
 from typing_extensions import override
 
@@ -48,9 +48,14 @@ class ToolProcessor(Processor):
         self.span.add_event(f"{meta.name} ({meta.path})", timestamp=meta.created_at)
 
         match event:
+            case ToolStartEvent():
+                pass
+            case None:  # finish event
+                pass
             case ToolSuccessEvent():
                 output_cls = type(event.output)
 
+                self.span.reset_exception()
                 self.span.set_attributes(
                     {
                         SpanAttributes.OUTPUT_VALUE: event.output.get_text_content(),
@@ -62,6 +67,7 @@ class ToolProcessor(Processor):
             case ToolErrorEvent():
                 span = self.span.child(meta.name, event=(event, meta))
                 span.record_exception(event.error)
+                self.span.record_exception(event.error)
             case ToolRetryEvent():
                 self.span.child(meta.name, event=(event, meta))
             case _:

openinference_instrumentation_beeai-0.1.10/src/openinference/instrumentation/beeai/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.10"

openinference_instrumentation_beeai-0.1.8/src/openinference/instrumentation/beeai/version.py

@@ -1 +0,0 @@
-__version__ = "0.1.8"