openinference-instrumentation-beeai 0.1.9__py3-none-any.whl → 0.1.11__py3-none-any.whl

openinference/instrumentation/beeai/__init__.py

@@ -1,9 +1,12 @@
+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
 
@@ -29,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
@@ -66,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, 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)
 
-    def _build_tree(self, node: SpanWrapper) -> None:
+    @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,
@@ -75,13 +87,16 @@ 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 and node.status == StatusCode.ERROR:
@@ -91,7 +106,8 @@ class BeeAIInstrumentor(BaseInstrumentor):  # type: ignore
 
     @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 = (
@@ -102,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]
 
@@ -112,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/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(

openinference/instrumentation/beeai/processors/chat.py

@@ -99,7 +99,7 @@ class ChatModelProcessor(Processor):
 
             case ChatModelSuccessEvent():
                 if not self._messages:  # only when no streaming
-                    self._add_new_messages(event.value.messages)
+                    self._add_new_messages(event.value.output)
 
                 usage = event.value.usage
                 if usage:
@@ -126,7 +126,7 @@ class ChatModelProcessor(Processor):
                         SpanAttributes.OPENINFERENCE_SPAN_KIND: type(self).kind,
                         SpanAttributes.OUTPUT_VALUE: event.value.get_text_content(),
                         SpanAttributes.OUTPUT_MIME_TYPE: OpenInferenceMimeTypeValues.TEXT.value,
-                        f"{SpanAttributes.METADATA}.chunks_count": len(event.value.messages),
+                        f"{SpanAttributes.METADATA}.chunks_count": len(event.value.output),
                         **_unpack_object(
                             usage.model_dump(exclude_none=True) if usage else {},
                             prefix=f"{SpanAttributes.METADATA}.usage",

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,6 +48,10 @@ 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)
 

openinference/instrumentation/beeai/version.py

@@ -1 +1 @@
-__version__ = "0.1.9"
+__version__ = "0.1.11"

{openinference_instrumentation_beeai-0.1.9.dist-info → openinference_instrumentation_beeai-0.1.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openinference-instrumentation-beeai
-Version: 0.1.9
+Version: 0.1.11
 Summary: OpenInference BeeAI Instrumentation
 Project-URL: Homepage, https://github.com/Arize-ai/openinference/tree/main/python/instrumentation/openinference-instrumentation-beeai
 Author: IBM Corp.
@@ -15,7 +15,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Python: <3.14,>=3.11
-Requires-Dist: beeai-framework<0.2.0,>=0.1.36
+Requires-Dist: beeai-framework<0.2.0,>=0.1.49
 Requires-Dist: openinference-instrumentation>=0.1.37
 Requires-Dist: openinference-semantic-conventions>=0.1.21
 Requires-Dist: opentelemetry-api>=1.36.0
@@ -23,9 +23,9 @@ Requires-Dist: opentelemetry-instrumentation>=0.57b0
 Requires-Dist: opentelemetry-semantic-conventions>=0.57b0
 Requires-Dist: wrapt>=1.17.2
 Provides-Extra: instruments
-Requires-Dist: beeai-framework>=0.1.36; extra == 'instruments'
+Requires-Dist: beeai-framework>=0.1.49; extra == 'instruments'
 Provides-Extra: test
-Requires-Dist: beeai-framework>=0.1.36; extra == 'test'
+Requires-Dist: beeai-framework>=0.1.49; extra == 'test'
 Requires-Dist: opentelemetry-exporter-otlp; extra == 'test'
 Requires-Dist: opentelemetry-sdk; extra == 'test'
 Description-Content-Type: text/markdown

{openinference_instrumentation_beeai-0.1.9.dist-info → openinference_instrumentation_beeai-0.1.11.dist-info}/RECORD

@@ -1,21 +1,21 @@
-openinference/instrumentation/beeai/__init__.py,sha256=tEM7-1EhCKIFsoLeMEa7GcRYPFfjuGeXi2w4CEKA1R4,4547
+openinference/instrumentation/beeai/__init__.py,sha256=yTygxcKT4VHgegEQndDrUrhjKZ9ifLsJtwZjF6A3R2M,5229
 openinference/instrumentation/beeai/_span.py,sha256=iVlYou4vnNKtDxpypMdZuD2AKeaDiG1Cu5PXVzgQ8w4,3259
 openinference/instrumentation/beeai/_utils.py,sha256=tfQsQEcevyLJno8WmLTOe936GVTIS2etnAFVbAPyztc,2521
-openinference/instrumentation/beeai/version.py,sha256=XIaxbMbyiP-L3kguR1GhxirFblTXiHR1lMfDVITvHUI,22
+openinference/instrumentation/beeai/version.py,sha256=nllDrH0jyChMuuYrK0CC55iTBKUNTUjejtcwxyUF2EQ,23
 openinference/instrumentation/beeai/processors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-openinference/instrumentation/beeai/processors/base.py,sha256=eAcKq6ip23sp1pjCa6efjDEcfoPJ_Jc3GwaWk5oZU7k,2125
-openinference/instrumentation/beeai/processors/chat.py,sha256=a4Ps6opRb2POrQ8Nla1JHUhkdVFZF5rmht8Xca3F_xA,9285
+openinference/instrumentation/beeai/processors/base.py,sha256=-h8tx1moKVi-1t3RpN-c7klGRJyHwoYe16rKWcZ_bhs,2204
+openinference/instrumentation/beeai/processors/chat.py,sha256=Nl0JWThA7f7Oeet6136fARdpNmakMH0SYym3T9kCGFg,9281
 openinference/instrumentation/beeai/processors/embedding.py,sha256=T9fZs2M7qEs4SnLYbSXRbhe3P7rCNch-snRQBDSC9Es,2598
 openinference/instrumentation/beeai/processors/locator.py,sha256=G9TFW_HgXM1TrOVdvtRU2Eq3D-atHLAET4fSo4F02X8,3635
 openinference/instrumentation/beeai/processors/requirement.py,sha256=Q9DgHDd-5rmP88Fe00d7cNKQg5zQ7ILuRzVAej5xfmk,2666
-openinference/instrumentation/beeai/processors/tool.py,sha256=ddkJP0Y7-pZFQNlKSB-SOOwquAbWL9ram8VS46SU8-Q,2633
+openinference/instrumentation/beeai/processors/tool.py,sha256=o5aKAqEHZIk9bmK2rWUqtACiyMltIFCnzaunCKxRtu8,2765
 openinference/instrumentation/beeai/processors/workflow.py,sha256=OMwFFHv3mp4M4hFvH7utYd_fiSkTnBcl2oUVaMEdy-A,3815
 openinference/instrumentation/beeai/processors/agents/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 openinference/instrumentation/beeai/processors/agents/base.py,sha256=3fidrUoU9pVixq9YN_y1jkNMozNMX-YG2CMuh855cLk,1244
 openinference/instrumentation/beeai/processors/agents/react.py,sha256=rS3xlvgyZ5G6MyDMeSh4xFLT_66h7GVAYEYlwZCpIdY,3022
 openinference/instrumentation/beeai/processors/agents/requirement_agent.py,sha256=HpleY8pNWojuUqcae2PZgat7Xq2edU9C-uz0YjZQUyc,2774
 openinference/instrumentation/beeai/processors/agents/tool_calling.py,sha256=yaWP5JmGuvZIha9iUSKgv0MJgI0QSbuiJLLQFnbqUZw,1223
-openinference_instrumentation_beeai-0.1.9.dist-info/METADATA,sha256=q2qwIddIyMKySNHBL6DRFSCCQi5NqoK9xHuSn1fkEm0,5491
-openinference_instrumentation_beeai-0.1.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-openinference_instrumentation_beeai-0.1.9.dist-info/entry_points.txt,sha256=ee7EUhbWv-XK1dxhPXuFVy9qstzj-lc-265Phe2Ml9s,183
-openinference_instrumentation_beeai-0.1.9.dist-info/RECORD,,
+openinference_instrumentation_beeai-0.1.11.dist-info/METADATA,sha256=TS92bnYQ0L02tUQBftmfxYhZl_8UUWmiKfAnZ_i83JU,5492
+openinference_instrumentation_beeai-0.1.11.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+openinference_instrumentation_beeai-0.1.11.dist-info/entry_points.txt,sha256=ee7EUhbWv-XK1dxhPXuFVy9qstzj-lc-265Phe2Ml9s,183
+openinference_instrumentation_beeai-0.1.11.dist-info/RECORD,,