Skip to content

Commit 1b4e32c

Browse files
nforronforro
authored
Merge pull request #59 from nforro/misc
BeeAI: few improvements
2 parents 57f89fe + 06e332b commit 1b4e32c

File tree

5 files changed

+142
-1
lines changed

5 files changed

+142
-1
lines changed

‎beeai/Containerfile‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,15 +14,23 @@ RUN dnf -y install \
1414
spectool \
1515
&& dnf -y clean all
1616

17+
# https://github.com/i-am-bee/beeai-framework/pull/939
18+
# already merged, can be removed after next upstream release
1719
COPY beeai-gemini.patch /tmp
20+
# https://github.com/i-am-bee/beeai-framework/issues/959
21+
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
1824

1925
# Install BeeAI Framework and Phoenix
2026
RUN pip3 install --no-cache-dir \
2127
beeai-framework[mcp,duckduckgo]==0.1.31 \
2228
openinference-instrumentation-beeai \
2329
arize-phoenix-otel \
2430
&& cd /usr/local/lib/python3.13/site-packages \
25-
&& patch -p2 -i /tmp/beeai-gemini.patch
31+
&& patch -p2 -i /tmp/beeai-gemini.patch \
32+
&& patch -p2 -i /tmp/beeai-gemini-malformed-function-call.patch \
33+
&& patch -p5 -i /tmp/beeai-instrumentation.patch
2634

2735
# Create user
2836
RUN useradd -m -G wheel beeai

‎beeai/agents/base_agent.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,3 +50,12 @@ async def _run_with_schema(self, input: TInputSchema) -> TOutputSchema:
5050

5151
async def run_with_schema(self, input: TInputSchema) -> TOutputSchema:
5252
return await self._run_with_schema(input)
53+
54+
55+
if os.getenv("LITELLM_DEBUG"):
56+
# the following two modules call `litellm_debug(False)` on import
57+
# import them explicitly now to ensure our call to `litellm_debug()` is not negated later
58+
import beeai_framework.adapters.litellm.chat
59+
import beeai_framework.adapters.litellm.embedding
60+
from beeai_framework.adapters.litellm.utils import litellm_debug
61+
litellm_debug(True)

‎beeai/beeai-gemini-malformed-function-call.patch‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
diff --git a/python/beeai_framework/adapters/litellm/chat.py b/python/beeai_framework/adapters/litellm/chat.py
2+
index b35acde3..8d01b0c9 100644
3+
--- a/python/beeai_framework/adapters/litellm/chat.py
4+
+++ b/python/beeai_framework/adapters/litellm/chat.py
5+
@@ -231,6 +231,13 @@ class LiteLLMChatModel(ChatModel, ABC):
6+
)
7+
8+
def _transform_output(self, chunk: ModelResponse | ModelResponseStream) -> ChatModelOutput:
9+
+ if not chunk.choices:
10+
+ usage = chunk.get("usage") # type: ignore
11+
+ return ChatModelOutput(
12+
+ messages=[],
13+
+ finish_reason="stop",
14+
+ usage=ChatModelUsage(**usage.model_dump()) if usage else None,
15+
+ )
16+
choice = chunk.choices[0]
17+
finish_reason = choice.finish_reason
18+
usage = chunk.get("usage") # type: ignore

‎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,

‎beeai/templates/beeai-agent.env‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,5 +9,6 @@ CHAT_MODEL=
99
BEEAI_MAX_ITERATIONS=100
1010
BEEAI_MAX_RETRIES_PER_STEP=5
1111
BEEAI_TOTAL_MAX_RETRIES=10
12+
LITELLM_DEBUG=
1213

1314
GITLAB_USER=

0 commit comments

Comments
 (0)