@arizeai/openinference-instrumentation-beeai 1.0.0 → 1.1.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@arizeai/openinference-instrumentation-beeai",
-  "version": "1.0.0",
+  "version": "1.1.0",
   "license": "Apache-2.0",
   "author": "gallas.milan@gmail.com",
   "description": "OpenInference Instrumentation for BeeAI framework",
@@ -36,12 +36,12 @@
     "@arizeai/openinference-semantic-conventions": "1.0.0"
   },
   "peerDependencies": {
-    "beeai-framework": "^0.1.5"
+    "beeai-framework": "^0.1.8"
   },
   "devDependencies": {
     "@types/jest": "^29.5.12",
     "@types/node": "^20.14.11",
-    "beeai-framework": "^0.1.5",
+    "beeai-framework": "^0.1.8",
     "import-in-the-middle": "^1.13.0",
     "@opentelemetry/exporter-trace-otlp-proto": "^0.50.0",
     "@opentelemetry/resources": "^1.30.1",

package/src/config.ts

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { BeeCallbacks } from "beeai-framework/agents/bee/types";
+import type { ReActAgentCallbacks } from "beeai-framework/agents/react/types";
 import { ChatModelEvents } from "beeai-framework/backend/chat";
 import { ToolEvents } from "beeai-framework/tools/base";
 
@@ -25,15 +25,16 @@ export const INSTRUMENTATION_IGNORED_KEYS = process.env
     )
   : [];
 
-export const partialUpdateEventName: keyof BeeCallbacks = "partialUpdate";
-export const updateEventName: keyof BeeCallbacks = "update";
-export const toolErrorEventName: keyof BeeCallbacks = "toolError";
-export const toolStartEventName: keyof BeeCallbacks = "toolStart";
-export const toolSuccessEventName: keyof BeeCallbacks = "toolSuccess";
-export const startEventName: keyof BeeCallbacks = "start";
-export const successEventName: keyof BeeCallbacks = "success";
-export const errorEventName: keyof BeeCallbacks = "error";
-export const retryEventName: keyof BeeCallbacks = "retry";
+export const partialUpdateEventName: keyof ReActAgentCallbacks =
+  "partialUpdate";
+export const updateEventName: keyof ReActAgentCallbacks = "update";
+export const toolErrorEventName: keyof ReActAgentCallbacks = "toolError";
+export const toolStartEventName: keyof ReActAgentCallbacks = "toolStart";
+export const toolSuccessEventName: keyof ReActAgentCallbacks = "toolSuccess";
+export const startEventName: keyof ReActAgentCallbacks = "start";
+export const successEventName: keyof ReActAgentCallbacks = "success";
+export const errorEventName: keyof ReActAgentCallbacks = "error";
+export const retryEventName: keyof ReActAgentCallbacks = "retry";
 
 export const newTokenLLMEventName: keyof ChatModelEvents = `newToken`;
 export const successLLMEventName: keyof ChatModelEvents = `success`;

package/src/helpers/buildTraceTree.ts

@@ -5,9 +5,11 @@ import {
   OpenInferenceSpanKind,
   SemanticConventions,
 } from "@arizeai/openinference-semantic-conventions";
+import { startEventName, successEventName } from "../config";
 
 interface BuiltTraceTreeProps {
   tracer: OITracer;
+  mainSpanKind: OpenInferenceSpanKind;
   data: {
     prompt?: string | null;
     history: GeneratedResponse[] | undefined;
@@ -72,7 +74,90 @@ function buildSpansForParent({ tracer, data }: BuildSpansForParentProps) {
     });
 }
 
-export function buildTraceTree({ tracer, data }: BuiltTraceTreeProps) {
+function buildAgentMainSpanData(data: BuiltTraceTreeProps["data"]) {
+  return {
+    ...(data.prompt && { [SemanticConventions.INPUT_VALUE]: data.prompt }),
+    ...(data.generatedMessage !== undefined && {
+      [SemanticConventions.OUTPUT_VALUE]: JSON.stringify(data.generatedMessage),
+    }),
+    ...(data.history && { history: JSON.stringify(data.history) }),
+  };
+}
+
+function buildToolMainSpanData(data: BuiltTraceTreeProps["data"]) {
+  const startBeeaiSpan = data.spans.find(
+    (span) => span.name === startEventName,
+  );
+  const successBeeeaiSpan = data.spans.find(
+    (span) => span.name === successEventName,
+  );
+
+  return {
+    ...(startBeeaiSpan && {
+      [SemanticConventions.INPUT_VALUE]:
+        startBeeaiSpan.attributes.data?.[SemanticConventions.TOOL_PARAMETERS],
+      [SemanticConventions.TOOL_PARAMETERS]:
+        startBeeaiSpan.attributes.data?.[SemanticConventions.TOOL_PARAMETERS],
+    }),
+    ...(successBeeeaiSpan && {
+      [SemanticConventions.OUTPUT_VALUE]:
+        successBeeeaiSpan.attributes.data?.[SemanticConventions.OUTPUT_VALUE],
+    }),
+  };
+}
+
+function buildLLMMainSpanData(data: BuiltTraceTreeProps["data"]) {
+  const startBeeaiSpan = data.spans.find(
+    (span) => span.name === startEventName,
+  );
+  const successBeeeaiSpan = data.spans.find(
+    (span) => span.name === successEventName,
+  );
+
+  if (!startBeeaiSpan && !successBeeeaiSpan) return {};
+
+  const provider =
+    startBeeaiSpan?.attributes.data?.[SemanticConventions.LLM_PROVIDER] ||
+    successBeeeaiSpan?.attributes.data?.[SemanticConventions.LLM_PROVIDER];
+
+  const modelName =
+    startBeeaiSpan?.attributes.data?.[SemanticConventions.LLM_MODEL_NAME] ||
+    successBeeeaiSpan?.attributes.data?.[SemanticConventions.LLM_MODEL_NAME];
+
+  return {
+    ...(startBeeaiSpan && {
+      [SemanticConventions.INPUT_VALUE]:
+        startBeeaiSpan?.attributes.data?.[SemanticConventions.INPUT_VALUE],
+      [SemanticConventions.INPUT_MIME_TYPE]:
+        startBeeaiSpan?.attributes.data?.[SemanticConventions.INPUT_MIME_TYPE],
+    }),
+    ...(successBeeeaiSpan && {
+      [SemanticConventions.OUTPUT_MIME_TYPE]:
+        successBeeeaiSpan.attributes.data?.[
+          SemanticConventions.OUTPUT_MIME_TYPE
+        ],
+      [SemanticConventions.OUTPUT_VALUE]:
+        successBeeeaiSpan.attributes.data?.[SemanticConventions.OUTPUT_VALUE],
+    }),
+    ...(provider && { [SemanticConventions.LLM_PROVIDER]: provider }),
+    ...(modelName && { [SemanticConventions.LLM_MODEL_NAME]: modelName }),
+  };
+}
+
+export function buildTraceTree({
+  tracer,
+  data,
+  mainSpanKind,
+}: BuiltTraceTreeProps) {
+  const computedData =
+    mainSpanKind === OpenInferenceSpanKind.AGENT
+      ? buildAgentMainSpanData(data)
+      : mainSpanKind === OpenInferenceSpanKind.TOOL
+        ? buildToolMainSpanData(data)
+        : mainSpanKind === OpenInferenceSpanKind.LLM
+          ? buildLLMMainSpanData(data)
+          : {};
+
   tracer.startActiveSpan(
     `beeai-framework-main`,
     {
@@ -82,16 +167,9 @@ export function buildTraceTree({ tracer, data }: BuiltTraceTreeProps) {
       attributes: {
         source: data.source,
         traceId: data.traceId,
-        [SemanticConventions.OPENINFERENCE_SPAN_KIND]:
-          OpenInferenceSpanKind.AGENT,
+        [SemanticConventions.OPENINFERENCE_SPAN_KIND]: mainSpanKind,
         ["beeai.version"]: data.version,
-        ...(data.prompt && { [SemanticConventions.INPUT_VALUE]: data.prompt }),
-        ...(data.generatedMessage !== undefined && {
-          [SemanticConventions.OUTPUT_VALUE]: JSON.stringify(
-            data.generatedMessage,
-          ),
-        }),
-        ...(data.history && { history: JSON.stringify(data.history) }),
+        ...computedData,
       },
     },
     (activeSpan) => {

package/src/helpers/getSerializedObjectSafe.ts

@@ -18,7 +18,7 @@ import { Serializable } from "beeai-framework/internals/serializable";
 import { instrumentationLogger } from "beeai-framework/instrumentation/logger";
 import { getProp } from "beeai-framework/internals/helpers/object";
 import { EventMeta, InferCallbackValue } from "beeai-framework/emitter/types";
-import { BeeCallbacks } from "beeai-framework/agents/bee/types";
+import type { ReActAgentCallbacks } from "beeai-framework/agents/react/types";
 import { ChatModel, ChatModelEvents } from "beeai-framework/backend/chat";
 import { BaseAgent } from "beeai-framework/agents/base";
 import { Message, MessageContentPart } from "beeai-framework/backend/message";
@@ -96,16 +96,16 @@ export function getSerializedObjectSafe(
   // agent events
   if (
     [startEventName, successEventName, errorEventName, retryEventName].includes(
-      meta.name as keyof BeeCallbacks,
+      meta.name as keyof ReActAgentCallbacks,
     ) &&
     meta.creator instanceof BaseAgent
   ) {
     const { meta, tools, memory, error, data } =
       dataObject as InferCallbackValue<
-        | BeeCallbacks["start"]
-        | BeeCallbacks["error"]
-        | BeeCallbacks["success"]
-        | BeeCallbacks["retry"]
+        | ReActAgentCallbacks["start"]
+        | ReActAgentCallbacks["error"]
+        | ReActAgentCallbacks["success"]
+        | ReActAgentCallbacks["retry"]
       >;
     return {
       [SemanticConventions.OPENINFERENCE_SPAN_KIND]:
@@ -137,11 +137,11 @@ export function getSerializedObjectSafe(
   // update events
   if (
     [updateEventName, partialUpdateEventName].includes(
-      meta.name as keyof BeeCallbacks,
+      meta.name as keyof ReActAgentCallbacks,
     )
   ) {
     const { data } = dataObject as InferCallbackValue<
-      BeeCallbacks["partialUpdate"] | BeeCallbacks["update"]
+      ReActAgentCallbacks["partialUpdate"] | ReActAgentCallbacks["update"]
     >;
 
     const output = data.final_answer || data.tool_output;
@@ -165,13 +165,13 @@ export function getSerializedObjectSafe(
   // tool events (from agent)
   if (
     [toolErrorEventName, toolStartEventName, toolSuccessEventName].includes(
-      meta.name as keyof BeeCallbacks,
+      meta.name as keyof ReActAgentCallbacks,
     )
   ) {
     const { data } = dataObject as InferCallbackValue<
-      | BeeCallbacks["toolError"]
-      | BeeCallbacks["toolStart"]
-      | BeeCallbacks["toolSuccess"]
+      | ReActAgentCallbacks["toolError"]
+      | ReActAgentCallbacks["toolStart"]
+      | ReActAgentCallbacks["toolSuccess"]
     >;
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     const output: any = data.result && data.result.createSnapshot();

package/src/instrumentation.ts

@@ -9,6 +9,7 @@ import * as bee from "beeai-framework";
 
 import { OITracer, TraceConfigOptions } from "@arizeai/openinference-core";
 import { createTelemetryMiddleware } from "./middleware";
+import { OpenInferenceSpanKind } from "@arizeai/openinference-semantic-conventions";
 
 const MODULE_NAME = "beeai-framework";
 
@@ -83,21 +84,71 @@ export class BeeAIInstrumentation extends InstrumentationBase {
       return module;
     }
 
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    const createWrappedMethod = <T extends (...args: any[]) => any>(
+      instrumentation: BeeAIInstrumentation,
+      original: T,
+      mainSpanKind: OpenInferenceSpanKind,
+    ) => {
+      return function wrappedMethod(this: unknown, ...args: unknown[]) {
+        const returned = original.apply(this, args);
+        if (returned?.middleware) {
+          returned.middleware(
+            createTelemetryMiddleware(instrumentation.oiTracer, mainSpanKind),
+          );
+        }
+        return returned;
+      };
+    };
     diag.info("[BeeaiInstrumentation] Patching BaseAgent.run");
 
     // eslint-disable-next-line @typescript-eslint/no-this-alias
     const instrumentation: BeeAIInstrumentation = this;
 
-    this._wrap(module.BaseAgent.prototype, "run", (original) => {
-      return function wrappedRun(this: unknown, ...args) {
-        const returned = original.apply(this, args);
-        returned.middleware(
-          createTelemetryMiddleware(instrumentation.oiTracer),
-        );
+    this._wrap(module.BaseAgent.prototype, "run", (original) =>
+      createWrappedMethod(
+        instrumentation,
+        original,
+        OpenInferenceSpanKind.AGENT,
+      ),
+    );
 
-        return returned;
-      };
-    });
+    // tool instrumentation support
+    if (module.Tool) {
+      this._wrap(module.Tool.prototype, "run", (original) =>
+        createWrappedMethod(
+          instrumentation,
+          original,
+          OpenInferenceSpanKind.TOOL,
+        ),
+      );
+    }
+    // model instrumentation support
+    if (module.ChatModel) {
+      this._wrap(module.ChatModel.prototype, "create", (original) =>
+        createWrappedMethod(
+          instrumentation,
+          original,
+          OpenInferenceSpanKind.LLM,
+        ),
+      );
+      this._wrap(module.ChatModel.prototype, "createStructure", (original) =>
+        createWrappedMethod(
+          instrumentation,
+          original,
+          OpenInferenceSpanKind.LLM,
+        ),
+      );
+    }
+    if (module.EmbeddingModel) {
+      this._wrap(module.EmbeddingModel.prototype, "create", (original) =>
+        createWrappedMethod(
+          instrumentation,
+          original,
+          OpenInferenceSpanKind.EMBEDDING,
+        ),
+      );
+    }
 
     return module;
   }

package/src/middleware.ts

@@ -20,7 +20,7 @@ import { createSpan } from "./helpers/create-span";
 import { IdNameManager } from "./helpers/idNameManager";
 import { getErrorSafe } from "./helpers/getErrorSafe";
 import { findLast, isEmpty } from "remeda";
-import type { BeeCallbacks } from "beeai-framework/agents/bee/types";
+import type { ReActAgentCallbacks } from "beeai-framework/agents/react/types";
 import type { InferCallbackValue } from "beeai-framework/emitter/types";
 import { FrameworkError } from "beeai-framework/errors";
 import { Version } from "beeai-framework/version";
@@ -36,12 +36,15 @@ import {
   successLLMEventName,
 } from "./config";
 import { createFullPath } from "beeai-framework/emitter/utils";
-import type { BeeAgent } from "beeai-framework/agents/bee/agent";
+import type { ReActAgent } from "beeai-framework/agents/react/agent";
 import { BaseAgent } from "beeai-framework/agents/base";
 import { diag } from "@opentelemetry/api";
 import { FrameworkSpan, GeneratedResponse } from "./types";
 import { buildTraceTree } from "./helpers/buildTraceTree";
-import { SemanticConventions } from "@arizeai/openinference-semantic-conventions";
+import {
+  OpenInferenceSpanKind,
+  SemanticConventions,
+} from "@arizeai/openinference-semantic-conventions";
 
 export const activeTracesMap = new Map<string, string>();
 
@@ -52,7 +55,10 @@ export const activeTracesMap = new Map<string, string>();
  * Then we create the open telemetry spans when all data are collected. We are retroactively deleting some unnecessary internal spans
  * see "emitter.match((event) => event.path === `${basePath}.run.${finishEventName}`" section
  */
-export function createTelemetryMiddleware(tracer: OITracer) {
+export function createTelemetryMiddleware(
+  tracer: OITracer,
+  mainSpanKind: OpenInferenceSpanKind,
+) {
   return (context: GetRunContext<RunInstance, unknown>) => {
     if (!context.emitter?.trace?.id) {
       throw new FrameworkError(`Fatal error. Missing traceId`, [], { context });
@@ -73,7 +79,7 @@ export function createTelemetryMiddleware(tracer: OITracer) {
 
     let prompt: string | undefined | null = null;
     if (instance instanceof BaseAgent) {
-      prompt = (runParams as Parameters<BeeAgent["run"]>)[0].prompt;
+      prompt = (runParams as Parameters<ReActAgent["run"]>)[0].prompt;
     }
 
     const spansMap = new Map<string, FrameworkSpan>();
@@ -151,6 +157,7 @@ export function createTelemetryMiddleware(tracer: OITracer) {
           // create tracer spans from collected data
           buildTraceTree({
             tracer,
+            mainSpanKind,
             data: {
               prompt: prompt,
               history,
@@ -301,7 +308,7 @@ export function createTelemetryMiddleware(tracer: OITracer) {
       (event) =>
         event.name === successLLMEventName &&
         event.creator instanceof BaseAgent,
-      (data: InferCallbackValue<BeeCallbacks["success"]>) => {
+      (data: InferCallbackValue<ReActAgentCallbacks["success"]>) => {
         const { data: dataObject, memory } = data;
 
         generatedMessage = {

package/dist/examples/bee.d.ts

@@ -1,2 +0,0 @@
-import "./instrumentation";
-//# sourceMappingURL=bee.d.ts.map

package/dist/examples/bee.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"bee.d.ts","sourceRoot":"","sources":["../../examples/bee.ts"],"names":[],"mappings":"AAAA,OAAO,mBAAmB,CAAC"}

package/dist/examples/bee.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"bee.js","sourceRoot":"","sources":["../../examples/bee.ts"],"names":[],"mappings":";;;;;;;;;;;AAAA,6BAA2B;AAC3B,4DAA4D;AAC5D,uEAA+E;AAC/E,oEAAiE;AACjE,oFAAqF;AACrF,uEAAwE;AAExE,MAAM,GAAG,GAAG,IAAI,sBAAe,CAAC,UAAU,CAAC,CAAC,CAAC,+DAA+D;AAC5G,MAAM,KAAK,GAAG,IAAI,gBAAQ,CAAC;IACzB,GAAG,EAAE,8CAA8C;IACnD,MAAM,EAAE,IAAI,yBAAW,EAAE,EAAE,4CAA4C;IACvE,KAAK,EAAE,CAAC,IAAI,uCAAoB,EAAE,EAAE,IAAI,yBAAa,EAAE,CAAC,EAAE,2CAA2C;CACtG,CAAC,CAAC;AAEH,SAAe,IAAI;;QACjB,MAAM,QAAQ,GAAG,MAAM,KAAK,CAAC,GAAG,CAAC;YAC/B,MAAM,EAAE,0CAA0C;SACnD,CAAC,CAAC;QACH,sCAAsC;QACtC,OAAO,CAAC,GAAG,CAAC,aAAa,EAAE,QAAQ,CAAC,MAAM,CAAC,IAAI,CAAC,CAAC;IACnD,CAAC;CAAA;AAED,sCAAsC;AACtC,IAAI,EAAE,CAAC,KAAK,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC"}