openinference-instrumentation-beeai 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

openinference/instrumentation/beeai/processors/chat.py

@@ -0,0 +1,239 @@
+import json
+from datetime import datetime
+from typing import Any, ClassVar
+
+from beeai_framework.backend import (
+    AnyMessage,
+    ChatModel,
+    MessageImageContent,
+    MessageTextContent,
+    MessageToolCallContent,
+    MessageToolResultContent,
+)
+from beeai_framework.backend.events import (
+    ChatModelStartEvent,
+    ChatModelSuccessEvent,
+)
+from beeai_framework.context import RunContext, RunContextStartEvent
+from beeai_framework.emitter import EventMeta
+from beeai_framework.tools import AnyTool
+from beeai_framework.utils.lists import remove_falsy
+from typing_extensions import override
+
+from openinference.instrumentation.beeai._utils import (
+    _unpack_object,
+    safe_dump_model_schema,
+    stringify,
+)
+from openinference.instrumentation.beeai.processors.base import Processor
+from openinference.semconv.trace import (
+    MessageAttributes,
+    MessageContentAttributes,
+    OpenInferenceMimeTypeValues,
+    OpenInferenceSpanKindValues,
+    SpanAttributes,
+    ToolAttributes,
+    ToolCallAttributes,
+)
+
+
+class ChatModelProcessor(Processor):
+    kind: ClassVar[OpenInferenceSpanKindValues] = OpenInferenceSpanKindValues.LLM
+
+    def __init__(self, event: RunContextStartEvent, meta: "EventMeta"):
+        super().__init__(event, meta)
+
+        self._last_updated_at = datetime.now()
+        self._messages: dict[str, "AnyMessage"] = {}
+
+        assert isinstance(meta.creator, RunContext)
+        assert isinstance(meta.creator.instance, ChatModel)
+        llm = meta.creator.instance
+        self.span.set_attributes(
+            {
+                SpanAttributes.LLM_MODEL_NAME: llm.model_id,
+                SpanAttributes.LLM_PROVIDER: llm.provider_id,
+            }
+        )
+
+    @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)
+
+        if meta.name == "finish":
+            return
+
+        match event:
+            case ChatModelStartEvent():
+                assert isinstance(meta.creator, ChatModel)
+                self._last_updated_at = meta.created_at
+                self.span.set_attributes(
+                    _process_messages(
+                        event.input.messages,
+                        prefix=SpanAttributes.LLM_INPUT_MESSAGES,
+                    ),
+                )
+                self.span.set_attributes(
+                    {
+                        SpanAttributes.LLM_TOOLS: json.loads(stringify(event.input.tools or [])),
+                        SpanAttributes.LLM_INVOCATION_PARAMETERS: stringify(
+                            meta.creator.parameters.model_dump(
+                                exclude_none=True, exclude_unset=True
+                            )
+                            | event.input.model_dump(
+                                exclude_none=True,
+                                exclude_unset=True,
+                                exclude={
+                                    "tools",
+                                    "messages",
+                                },
+                            ),
+                        ),
+                    }
+                )
+
+            case ChatModelSuccessEvent():
+                if not self._messages:  # only when no streaming
+                    self._add_new_messages(event.value.messages)
+
+                usage = event.value.usage
+
+                self.span.set_attributes(
+                    {
+                        **(
+                            {
+                                SpanAttributes.LLM_TOKEN_COUNT_TOTAL: usage.total_tokens,
+                                SpanAttributes.LLM_TOKEN_COUNT_PROMPT: usage.prompt_tokens,
+                                SpanAttributes.LLM_TOKEN_COUNT_COMPLETION: usage.completion_tokens,
+                            }
+                            if usage
+                            else {}
+                        ),
+                        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),
+                        **_unpack_object(
+                            usage.model_dump(exclude_none=True) if usage else {},
+                            prefix=f"{SpanAttributes.METADATA}.usage",
+                        ),
+                    }
+                )
+            case _:
+                self.span.child(meta.name, event=(event, meta))
+
+    def _add_new_messages(self, messages: list["AnyMessage"]) -> None:
+        for new_msg in messages:
+            msg_id = new_msg.meta.get("id") or f"latest_{new_msg.role}"
+            if msg_id in self._messages:
+                self._messages[msg_id].merge(new_msg)
+            else:
+                self._messages[msg_id] = new_msg.clone()
+
+            _aggregate_msg_content(self._messages[msg_id])
+            self.span.set_attributes(
+                _process_messages(
+                    [self._messages[msg_id]], prefix=SpanAttributes.LLM_OUTPUT_MESSAGES
+                )
+            )
+
+
+def _process_tools(tools: list[AnyTool]) -> list[dict[str, str | Any]]:
+    return [
+        {
+            SpanAttributes.TOOL_NAME: t.name,
+            SpanAttributes.TOOL_DESCRIPTION: t.description,
+            # TODO: difference between TOOL_PARAMETERS and TOOL_JSON_SCHEMA is not obvious
+            SpanAttributes.TOOL_PARAMETERS: safe_dump_model_schema(t.input_schema),
+            ToolAttributes.TOOL_JSON_SCHEMA: safe_dump_model_schema(t.input_schema),
+        }
+        for t in tools
+    ]
+
+
+def _process_messages(
+    messages: list["AnyMessage"], prefix: str = "", offset: int = 0
+) -> dict[str, Any]:
+    if prefix and not prefix.endswith("."):
+        prefix += "."
+
+    output = {}
+    for _i, msg in enumerate(messages):
+        i = _i + offset
+        output[f"{prefix}{i}.{MessageAttributes.MESSAGE_ROLE}"] = str(msg.role)
+
+        output.update(
+            _unpack_object(
+                remove_falsy(
+                    [
+                        (
+                            {
+                                MessageContentAttributes.MESSAGE_CONTENT_TYPE: "text",
+                                MessageContentAttributes.MESSAGE_CONTENT_TEXT: content.text,
+                            }
+                            if isinstance(content, MessageTextContent)
+                            else {
+                                MessageContentAttributes.MESSAGE_CONTENT_TYPE: "image",
+                                MessageContentAttributes.MESSAGE_CONTENT_IMAGE: content.image_url[
+                                    "url"
+                                ],
+                            }
+                            if isinstance(content, MessageImageContent)
+                            else {
+                                MessageAttributes.MESSAGE_TOOL_CALL_ID: content.tool_call_id,
+                                MessageContentAttributes.MESSAGE_CONTENT_TYPE: "text",
+                                MessageContentAttributes.MESSAGE_CONTENT_TEXT: stringify(
+                                    content.result, pretty=True
+                                ),
+                            }
+                            if isinstance(content, MessageToolResultContent)
+                            else None
+                        )
+                        for content in msg.content
+                    ],
+                ),
+                prefix=f"{prefix}{i}.{MessageAttributes.MESSAGE_CONTENTS}",
+            )
+        )
+        # )
+
+        tool_calls: list[dict[str, Any]] = [
+            {
+                ToolCallAttributes.TOOL_CALL_ID: msg_call.id,
+                ToolCallAttributes.TOOL_CALL_FUNCTION_NAME: msg_call.tool_name,
+                ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON: msg_call.args,
+            }
+            for msg_call in msg.get_by_type(MessageToolCallContent)
+        ]
+        output.update(
+            _unpack_object(tool_calls, prefix=f"{prefix}{i}.{MessageAttributes.MESSAGE_TOOL_CALLS}")
+        )
+
+    return output
+
+
+def _aggregate_msg_content(message: "AnyMessage") -> None:
+    from beeai_framework.backend import MessageTextContent, MessageToolCallContent
+
+    contents = message.content.copy()
+    aggregated_content: list[Any] = []
+
+    for content in contents:
+        last_content = aggregated_content[-1] if aggregated_content else None
+        if isinstance(last_content, MessageTextContent) and isinstance(content, MessageTextContent):
+            last_content.text += content.text
+        elif isinstance(last_content, MessageToolCallContent) and isinstance(
+            content, MessageToolCallContent
+        ):
+            last_content.args += content.args
+        else:
+            aggregated_content.append(content)
+
+    message.content.clear()
+    message.content.extend(aggregated_content)

openinference/instrumentation/beeai/processors/embedding.py

@@ -0,0 +1,71 @@
+from typing import TYPE_CHECKING, Any, ClassVar
+
+if TYPE_CHECKING:
+    from beeai_framework.context import RunContextStartEvent
+    from beeai_framework.emitter import EventMeta
+
+from beeai_framework.backend import EmbeddingModel
+from beeai_framework.backend.events import (
+    EmbeddingModelStartEvent,
+    EmbeddingModelSuccessEvent,
+)
+from beeai_framework.context import RunContext
+from typing_extensions import override
+
+from openinference.instrumentation.beeai.processors.base import Processor
+from openinference.semconv.trace import (
+    EmbeddingAttributes,
+    OpenInferenceSpanKindValues,
+    SpanAttributes,
+)
+
+
+class EmbeddingModelProcessor(Processor):
+    kind: ClassVar[OpenInferenceSpanKindValues] = OpenInferenceSpanKindValues.EMBEDDING
+
+    def __init__(self, event: "RunContextStartEvent", meta: "EventMeta"):
+        super().__init__(event, meta)
+
+        assert isinstance(meta.creator, RunContext)
+        assert isinstance(meta.creator.instance, EmbeddingModel)
+
+        llm = meta.creator.instance
+        self.span.set_attributes(
+            {
+                SpanAttributes.EMBEDDING_MODEL_NAME: llm.model_id,
+                SpanAttributes.LLM_PROVIDER: llm.provider_id,
+            }
+        )
+
+    @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)
+        self.span.child(meta.name, event=(event, meta))
+
+        if isinstance(event, EmbeddingModelStartEvent):
+            for idx, txt in enumerate(event.input.values):
+                self.span.set_attribute(
+                    f"{SpanAttributes.EMBEDDING_EMBEDDINGS}.{idx}.{EmbeddingAttributes.EMBEDDING_TEXT}",
+                    txt,
+                )
+        elif isinstance(event, EmbeddingModelSuccessEvent):
+            for idx, embedding in enumerate(event.value.embeddings):
+                self.span.set_attribute(
+                    f"{SpanAttributes.EMBEDDING_EMBEDDINGS}.{idx}.{EmbeddingAttributes.EMBEDDING_VECTOR}",
+                    embedding,
+                )
+
+            if event.value.usage:
+                self.span.set_attributes(
+                    {
+                        SpanAttributes.LLM_TOKEN_COUNT_TOTAL: event.value.usage.total_tokens,
+                        SpanAttributes.LLM_TOKEN_COUNT_PROMPT: event.value.usage.prompt_tokens,
+                        SpanAttributes.LLM_TOKEN_COUNT_COMPLETION: event.value.usage.completion_tokens,  # noqa: E501
+                    }
+                )

openinference/instrumentation/beeai/processors/locator.py

@@ -0,0 +1,106 @@
+import contextlib
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from beeai_framework.emitter import EventMeta
+
+from openinference.instrumentation.beeai.processors.base import Processor
+
+
+class ProcessorLocator:
+    entries: dict[type, type] = {}
+    _loaded: bool = False
+
+    @staticmethod
+    def _load() -> None:
+        if ProcessorLocator._loaded:
+            return
+
+        ProcessorLocator._loaded = True
+
+        with contextlib.suppress(ImportError):
+            from beeai_framework.backend.chat import ChatModel
+
+            from .chat import ChatModelProcessor
+
+            ProcessorLocator.entries[ChatModel] = ChatModelProcessor
+
+        with contextlib.suppress(ImportError):
+            from beeai_framework.backend.embedding import EmbeddingModel
+
+            from .embedding import EmbeddingModelProcessor
+
+            ProcessorLocator.entries[EmbeddingModel] = EmbeddingModelProcessor
+
+        with contextlib.suppress(ImportError):
+            from beeai_framework.agents.react.agent import ReActAgent
+
+            from .agents.react import ReActAgentProcessor
+
+            ProcessorLocator.entries[ReActAgent] = ReActAgentProcessor
+
+        with contextlib.suppress(ImportError):
+            from beeai_framework.agents.tool_calling.agent import ToolCallingAgent
+
+            from .agents.tool_calling import ToolCallingAgentProcessor
+
+            ProcessorLocator.entries[ToolCallingAgent] = ToolCallingAgentProcessor
+
+        with contextlib.suppress(ImportError):
+            from beeai_framework.agents.experimental.agent import RequirementAgent
+
+            from .agents.requirement_agent import RequirementAgentProcessor
+
+            ProcessorLocator.entries[RequirementAgent] = RequirementAgentProcessor
+
+        with contextlib.suppress(ImportError):
+            from beeai_framework.agents.experimental.requirements.requirement import Requirement
+
+            from openinference.instrumentation.beeai.processors.requirement import (
+                RequirementProcessor,
+            )
+
+            ProcessorLocator.entries[Requirement] = RequirementProcessor
+
+        with contextlib.suppress(ImportError):
+            from beeai_framework.agents.base import BaseAgent
+
+            from openinference.instrumentation.beeai.processors.agents.base import AgentProcessor
+
+            ProcessorLocator.entries[BaseAgent] = AgentProcessor
+
+        with contextlib.suppress(ImportError):
+            from beeai_framework.tools.tool import Tool
+
+            from openinference.instrumentation.beeai.processors.tool import ToolProcessor
+
+            ProcessorLocator.entries[Tool] = ToolProcessor
+
+        with contextlib.suppress(ImportError):
+            from beeai_framework.workflows.workflow import Workflow
+
+            from .workflow import WorkflowProcessor
+
+            ProcessorLocator.entries[Workflow] = WorkflowProcessor
+
+    @staticmethod
+    def locate(data: Any, event: "EventMeta") -> Processor:
+        ProcessorLocator._load()
+
+        from beeai_framework.context import RunContext, RunContextStartEvent
+
+        assert isinstance(data, RunContextStartEvent)
+        assert isinstance(event.creator, RunContext)
+
+        instance_cls = type(event.creator.instance)
+
+        if instance_cls in ProcessorLocator.entries:
+            cls_processor = ProcessorLocator.entries[instance_cls]
+        else:
+            for cls, cls_processor in ProcessorLocator.entries.items():
+                if isinstance(event.creator.instance, cls):
+                    break
+            else:
+                cls_processor = Processor
+
+        return cls_processor(data, event)  # type: ignore

openinference/instrumentation/beeai/processors/requirement.py

@@ -0,0 +1,67 @@
+from typing import Any, ClassVar
+
+from beeai_framework.agents.experimental.requirements.ask_permission import AskPermissionRequirement
+from beeai_framework.agents.experimental.requirements.conditional import ConditionalRequirement
+from beeai_framework.agents.experimental.requirements.events import RequirementInitEvent
+from beeai_framework.agents.experimental.requirements.requirement import Requirement
+from beeai_framework.context import RunContext, 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.base import Processor
+from openinference.semconv.trace import OpenInferenceSpanKindValues, SpanAttributes
+
+
+class RequirementProcessor(Processor):
+    kind: ClassVar[OpenInferenceSpanKindValues] = OpenInferenceSpanKindValues.UNKNOWN
+
+    def __init__(self, event: "RunContextStartEvent", meta: "EventMeta") -> None:
+        super().__init__(event, meta)
+
+        assert isinstance(meta.creator, RunContext)
+        assert isinstance(meta.creator.instance, Requirement)
+
+        requirement = meta.creator.instance
+        self.span.name = requirement.name or self.span.name
+        self._sync_state(meta.creator.instance)
+
+    def _sync_state(self, instance: "Requirement[Any]") -> None:
+        attributes = {
+            "name": instance.name,
+            "enabled": instance.enabled,
+            "priority": instance.priority,
+            "state": instance.state,
+            "middlewares": [str(m) for m in instance.middlewares],
+        }
+
+        if isinstance(instance, ConditionalRequirement):
+            attributes["target_name"] = str(instance.source)
+        elif isinstance(instance, AskPermissionRequirement):
+            attributes["includes"] = [str(t) for t in instance._include]
+            attributes["excludes"] = [str(t) for t in instance._exclude]
+            attributes["remember_choices"] = instance._remember_choices
+            attributes["_always_allow"] = instance._always_allow
+
+        self.span.set_attributes(
+            _unpack_object(
+                attributes,
+                prefix=f"{SpanAttributes.METADATA}",
+            )
+        )
+
+    @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 RequirementInitEvent():
+                pass
+            case _:
+                self.span.child(meta.name, event=(event, meta))

openinference/instrumentation/beeai/processors/tool.py

@@ -0,0 +1,72 @@
+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.tool import Tool
+from typing_extensions import override
+
+from openinference.instrumentation.beeai._utils import safe_dump_model_schema, stringify
+from openinference.instrumentation.beeai.processors.base import Processor
+from openinference.semconv.trace import (
+    OpenInferenceMimeTypeValues,
+    OpenInferenceSpanKindValues,
+    SpanAttributes,
+    ToolAttributes,
+)
+
+
+class ToolProcessor(Processor):
+    kind: ClassVar[OpenInferenceSpanKindValues] = OpenInferenceSpanKindValues.TOOL
+
+    def __init__(self, event: "RunContextStartEvent", meta: "EventMeta") -> None:
+        super().__init__(event, meta)
+
+        assert isinstance(meta.creator, RunContext)
+        assert isinstance(meta.creator.instance, Tool)
+
+        tool = meta.creator.instance
+        self.span.name = tool.name
+        self.span.set_attributes(
+            {
+                SpanAttributes.TOOL_NAME: tool.name,
+                SpanAttributes.TOOL_DESCRIPTION: tool.description,
+                # TODO: what's the difference?
+                SpanAttributes.TOOL_PARAMETERS: stringify(
+                    safe_dump_model_schema(tool.input_schema)
+                ),
+                ToolAttributes.TOOL_JSON_SCHEMA: stringify(
+                    safe_dump_model_schema(tool.input_schema)
+                ),
+            }
+        )
+
+    @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 ToolSuccessEvent():
+                output_cls = type(event.output)
+
+                self.span.set_attributes(
+                    {
+                        SpanAttributes.OUTPUT_VALUE: event.output.get_text_content(),
+                        SpanAttributes.OUTPUT_MIME_TYPE: OpenInferenceMimeTypeValues.TEXT.value,
+                        f"{SpanAttributes.METADATA}.output_class": output_cls.__name__,
+                        f"{SpanAttributes.METADATA}.is_empty": event.output.is_empty(),
+                    }
+                )
+            case ToolErrorEvent():
+                span = self.span.child(meta.name, event=(event, meta))
+                span.record_exception(event.error)
+            case ToolRetryEvent():
+                self.span.child(meta.name, event=(event, meta))
+            case _:
+                self.span.child(meta.name, event=(event, meta))

openinference/instrumentation/beeai/processors/workflow.py

@@ -0,0 +1,108 @@
+from typing import Any, ClassVar
+
+from beeai_framework.context import RunContext, RunContextStartEvent
+from beeai_framework.emitter import EventMeta
+from beeai_framework.workflows import (
+    Workflow,
+    WorkflowErrorEvent,
+    WorkflowStartEvent,
+    WorkflowSuccessEvent,
+)
+from beeai_framework.workflows.agent.agent import Schema as AgentWorkflowSchema
+from pydantic import BaseModel
+from typing_extensions import override
+
+from openinference.instrumentation.beeai._utils import (
+    _unpack_object,
+    safe_dump_model_schema,
+    stringify,
+)
+from openinference.instrumentation.beeai.processors.base import Processor
+from openinference.semconv.trace import (
+    OpenInferenceMimeTypeValues,
+    OpenInferenceSpanKindValues,
+    SpanAttributes,
+)
+
+
+class WorkflowProcessor(Processor):
+    kind: ClassVar[OpenInferenceSpanKindValues] = OpenInferenceSpanKindValues.CHAIN
+
+    def __init__(self, event: "RunContextStartEvent", meta: "EventMeta"):
+        super().__init__(event, meta)
+
+        assert isinstance(meta.creator, RunContext)
+        assert isinstance(meta.creator.instance, Workflow)
+
+        self._last_step = 0
+
+        workflow = meta.creator.instance
+        self.span.name = workflow.name or self.span.name
+        self.span.set_attributes(
+            {
+                f"{SpanAttributes.METADATA}.name": workflow.name,
+                f"{SpanAttributes.METADATA}.all_steps": workflow.step_names,
+                f"{SpanAttributes.METADATA}.start_step": workflow.start_step,
+                **_unpack_object(
+                    safe_dump_model_schema(workflow.schema),
+                    prefix=f"{SpanAttributes.METADATA}.schema",
+                ),
+            }
+        )
+
+    @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 WorkflowStartEvent() | WorkflowSuccessEvent() | WorkflowErrorEvent():
+                self._update_state(event)
+            case _:
+                self.span.child(meta.name, event=(event, meta))
+
+    def _update_state(
+        self, event: "WorkflowStartEvent[Any] | WorkflowSuccessEvent[Any] | WorkflowErrorEvent[Any]"
+    ) -> None:
+        self.span.set_attribute(f"{SpanAttributes.METADATA}.current_step", event.step)
+        self.span.set_attributes(
+            _unpack_object(
+                _serialize_state(event.run.state), prefix=f"{SpanAttributes.METADATA}.state"
+            )
+        )
+
+        # update steps
+        for idx, step in enumerate(event.run.steps):
+            if idx < self._last_step:
+                continue
+
+            self.span.set_attributes(
+                _unpack_object(
+                    {
+                        "name": step.name,
+                        "state": _serialize_state(step.state),
+                    },
+                    prefix=f"{SpanAttributes.METADATA}.steps.{idx}",
+                )
+            )
+            self._last_step = idx
+
+        if isinstance(event, WorkflowSuccessEvent):
+            if event.next == Workflow.END or event.run.result is not None:
+                result = event.run.result if event.run.result is not None else event.state
+                self.span.attributes.update(
+                    {
+                        SpanAttributes.OUTPUT_VALUE: stringify(_serialize_state(result)),
+                        SpanAttributes.OUTPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON.value,
+                    }
+                )
+
+
+def _serialize_state(result: BaseModel) -> dict[str, Any]:
+    exclude = {"new_messages", "inputs"} if isinstance(result, AgentWorkflowSchema) else set()
+    return result.model_dump(exclude=exclude, exclude_none=True)

openinference/instrumentation/beeai/version.py

@@ -1 +1 @@
-__version__ = "0.1.5"
+__version__ = "0.1.7"