Skip to content

Commit 06e332b

Browse files
nforronforro
committed
Group spans/traces based on agent run and add metadata
Signed-off-by: Nikola Forró <[email protected]>
1 parent 46e967e commit 06e332b

File tree

2 files changed

+109
-1
lines changed

2 files changed

+109
-1
lines changed

‎beeai/Containerfile‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ RUN dnf -y install \
1919
COPY beeai-gemini.patch /tmp
2020
# https://github.com/i-am-bee/beeai-framework/issues/959
2121
COPY beeai-gemini-malformed-function-call.patch /tmp
22+
# https://github.com/i-am-bee/beeai-framework/issues/906
23+
COPY beeai-instrumentation.patch /tmp
2224

2325
# Install BeeAI Framework and Phoenix
2426
RUN pip3 install --no-cache-dir \
@@ -27,7 +29,8 @@ RUN pip3 install --no-cache-dir \
2729
arize-phoenix-otel \
2830
&& cd /usr/local/lib/python3.13/site-packages \
2931
&& patch -p2 -i /tmp/beeai-gemini.patch \
30-
&& patch -p2 -i /tmp/beeai-gemini-malformed-function-call.patch
32+
&& patch -p2 -i /tmp/beeai-gemini-malformed-function-call.patch \
33+
&& patch -p5 -i /tmp/beeai-instrumentation.patch
3134

3235
# Create user
3336
RUN useradd -m -G wheel beeai

‎beeai/beeai-instrumentation.patch‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
diff --git a/python/instrumentation/openinference-instrumentation-beeai/src/openinference/instrumentation/beeai/__init__.py b/python/instrumentation/openinference-instrumentation-beeai/src/openinference/instrumentation/beeai/__init__.py
2+
index 21da0aa2..c32ae843 100644
3+
--- a/python/instrumentation/openinference-instrumentation-beeai/src/openinference/instrumentation/beeai/__init__.py
4+
+++ b/python/instrumentation/openinference-instrumentation-beeai/src/openinference/instrumentation/beeai/__init__.py
5+
@@ -40,6 +40,7 @@ except PackageNotFoundError:
6+
7+
class BeeAIInstrumentor(BaseInstrumentor): # type: ignore
8+
__slots__ = (
9+
+ "_original_requirement_agent_run",
10+
"_original_react_agent_run",
11+
"_original_tool_calling_agent_run",
12+
"_original_chat_model_create",
13+
@@ -61,6 +62,7 @@ class BeeAIInstrumentor(BaseInstrumentor): # type: ignore
14+
assert isinstance(config, TraceConfig)
15+
16+
from beeai_framework.agents.base import BaseAgent
17+
+ from beeai_framework.agents.experimental.agent import RequirementAgent
18+
from beeai_framework.agents.react.agent import ReActAgent
19+
from beeai_framework.agents.tool_calling.agent import ToolCallingAgent
20+
from beeai_framework.backend.chat import ChatModel
21+
@@ -93,6 +95,10 @@ class BeeAIInstrumentor(BaseInstrumentor): # type: ignore
22+
return result
23+
24+
## Agent support
25+
+ self._original_requirement_agent_run = getattr(
26+
+ import_module("beeai_framework.agents.experimental.agent"), "run", None
27+
+ )
28+
+ setattr(RequirementAgent, "run", run_wrapper(RequirementAgent.run))
29+
self._original_react_agent_run = getattr(
30+
import_module("beeai_framework.agents.react.agent"), "run", None
31+
)
32+
@@ -121,6 +127,11 @@ class BeeAIInstrumentor(BaseInstrumentor): # type: ignore
33+
logger.error("Instrumentation error", exc_info=e)
34+
35+
def _uninstrument(self, **kwargs: Any) -> None:
36+
+ if self._original_requirement_agent_run is not None:
37+
+ from beeai_framework.agents.experimental.agent import RequirementAgent
38+
+
39+
+ setattr(RequirementAgent, "run", self._original_requirement_agent_run)
40+
+ self._original_requirement_agent_run = None
41+
if self._original_react_agent_run is not None:
42+
from beeai_framework.agents.react.agent import ReActAgent
43+
44+
diff --git a/python/instrumentation/openinference-instrumentation-beeai/src/openinference/instrumentation/beeai/middleware.py b/python/instrumentation/openinference-instrumentation-beeai/src/openinference/instrumentation/beeai/middleware.py
45+
index f8150f04..24a0781e 100644
46+
--- a/python/instrumentation/openinference-instrumentation-beeai/src/openinference/instrumentation/beeai/middleware.py
47+
+++ b/python/instrumentation/openinference-instrumentation-beeai/src/openinference/instrumentation/beeai/middleware.py
48+
@@ -19,6 +19,8 @@ from importlib.metadata import PackageNotFoundError, version
49+
from typing import Any, Callable, Dict, Optional, cast
50+
51+
from beeai_framework.agents.base import BaseAgent
52+
+from beeai_framework.agents.experimental.agent import RequirementAgent
53+
+from beeai_framework.agents.experimental.events import RequirementAgentSuccessEvent
54+
from beeai_framework.agents.react.agent import ReActAgent
55+
from beeai_framework.agents.react.events import ReActAgentSuccessEvent
56+
from beeai_framework.agents.tool_calling.agent import ToolCallingAgent
57+
@@ -283,6 +285,26 @@ def create_telemetry_middleware(
58+
}
59+
for m in react_agent_typed_data.memory.messages
60+
]
61+
+ if isinstance(meta.creator, RequirementAgent):
62+
+ requirement_agent_typed_data = cast(RequirementAgentSuccessEvent, data)
63+
+ history = [
64+
+ {
65+
+ "text": m.text,
66+
+ "role": m.role.value if hasattr(m.role, "value") else m.role,
67+
+ }
68+
+ for m in requirement_agent_typed_data.state.memory.messages
69+
+ ]
70+
+ if (
71+
+ hasattr(requirement_agent_typed_data.state, "result")
72+
+ and requirement_agent_typed_data.state.result is not None
73+
+ ):
74+
+ result_role = requirement_agent_typed_data.state.result.role
75+
+ generated_message = {
76+
+ "role": result_role.value
77+
+ if hasattr(result_role, "value")
78+
+ else result_role,
79+
+ "text": requirement_agent_typed_data.state.result.text,
80+
+ }
81+
except Exception as e:
82+
logger.error("Instrumentation error: failed to extract success message", exc_info=e)
83+
84+
diff --git a/python/instrumentation/openinference-instrumentation-beeai/src/openinference/instrumentation/beeai/utils/build_trace_tree.py b/python/instrumentation/openinference-instrumentation-beeai/src/openinference/instrumentation/beeai/utils/build_trace_tree.py
85+
index 7f9e0f6a..eff19555 100644
86+
--- a/python/instrumentation/openinference-instrumentation-beeai/src/openinference/instrumentation/beeai/utils/build_trace_tree.py
87+
+++ b/python/instrumentation/openinference-instrumentation-beeai/src/openinference/instrumentation/beeai/utils/build_trace_tree.py
88+
@@ -16,7 +16,7 @@ from typing import Any, List, Optional, TypedDict
89+
90+
from opentelemetry import trace
91+
92+
-from openinference.instrumentation import OITracer
93+
+from openinference.instrumentation import OITracer, using_metadata
94+
from openinference.semconv.trace import OpenInferenceSpanKindValues, SpanAttributes
95+
96+
from .create_span import FrameworkSpan
97+
@@ -153,7 +153,7 @@ def build_trace_tree(tracer: OITracer, main_span_kind: str, data: BuildTraceTree
98+
**computed_data,
99+
}
100+
101+
- with tracer.start_as_current_span(
102+
+ with using_metadata({"source": data["source"]}), tracer.start_as_current_span(
103+
name="beeai-framework-main",
104+
start_time=data["startTime"],
105+
attributes=attributes,

0 commit comments

Comments
 (0)