@arizeai/openinference-instrumentation-beeai 1.0.0

package/src/helpers/traceSerializer.ts

@@ -0,0 +1,46 @@
+/**
+ * Copyright 2025 IBM Corp.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+const DEFAULT_IGNORE_KEYS = [
+  "emitter",
+  "logger",
+  "tokens",
+  "createdBy",
+  "client",
+];
+
+export function traceSerializer({
+  ignored_keys = [],
+}: {
+  ignored_keys?: string[];
+}) {
+  const mergedIgnoreKeys = new Set([...DEFAULT_IGNORE_KEYS, ...ignored_keys]);
+
+  return (body: object) =>
+    JSON.stringify(
+      body,
+      (() => {
+        return (key, value) => {
+          // Ignore specific keys, all owned entities and keys starting with underscore
+          if (mergedIgnoreKeys.has(key) || key.startsWith("_")) {
+            return;
+          }
+
+          return value;
+        };
+      })(),
+    );
+}

package/src/index.ts

@@ -0,0 +1 @@
+export * from "./instrumentation";

package/src/instrumentation.ts

@@ -0,0 +1,123 @@
+import {
+  InstrumentationBase,
+  InstrumentationNodeModuleDefinition,
+  InstrumentationConfig,
+} from "@opentelemetry/instrumentation";
+import { diag } from "@opentelemetry/api";
+import { Version } from "beeai-framework";
+import * as bee from "beeai-framework";
+
+import { OITracer, TraceConfigOptions } from "@arizeai/openinference-core";
+import { createTelemetryMiddleware } from "./middleware";
+
+const MODULE_NAME = "beeai-framework";
+
+/**
+ * Flag to check if the openai module has been patched
+ * Note: This is a fallback in case the module is made immutable (e.x. Deno, webpack, etc.)
+ */
+let _isOpenInferencePatched = false;
+
+/**
+ * function to check if instrumentation is enabled / disabled
+ */
+export function isPatched() {
+  return _isOpenInferencePatched;
+}
+
+export class BeeAIInstrumentation extends InstrumentationBase {
+  private oiTracer: OITracer;
+
+  constructor({
+    instrumentationConfig,
+    traceConfig,
+  }: {
+    /**
+     * The config for the instrumentation
+     * @see {@link InstrumentationConfig}
+     */
+    instrumentationConfig?: InstrumentationConfig;
+    /**
+     * The OpenInference trace configuration. Can be used to mask or redact sensitive information on spans.
+     * @see {@link TraceConfigOptions}
+     */
+    traceConfig?: TraceConfigOptions;
+  } = {}) {
+    super(
+      "@arizeai/openinference-instrumentation-beeai",
+      Version,
+      Object.assign({}, instrumentationConfig),
+    );
+    this.oiTracer = new OITracer({ tracer: this.tracer, traceConfig });
+  }
+
+  protected init() {
+    const module = new InstrumentationNodeModuleDefinition(
+      MODULE_NAME,
+      ["0.*"],
+      this.patch.bind(this),
+      this.unpatch.bind(this),
+    );
+
+    return module;
+  }
+
+  /**
+   * Manually instruments the BeeAI module. This is needed when the module is not loaded via require (commonjs)
+   * @param {openai} module
+   */
+  manuallyInstrument(module: typeof bee) {
+    diag.debug(`[BeeaiInstrumentation] Manually instrumenting ${MODULE_NAME}`);
+    this.patch(module);
+  }
+
+  private patch(module: typeof bee & { openInferencePatched?: boolean }) {
+    if (!module?.BaseAgent) {
+      diag.warn(
+        "[BeeaiInstrumentation] BaseAgent not found, skipping instrumentation.",
+      );
+      return module;
+    }
+
+    if (module?.openInferencePatched || _isOpenInferencePatched) {
+      return module;
+    }
+
+    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),
+        );
+
+        return returned;
+      };
+    });
+
+    return module;
+  }
+
+  /**
+   * Un-patches the OpenAI module's chat completions API
+   */
+  private unpatch(
+    moduleExports: typeof bee & { openInferencePatched?: boolean },
+    moduleVersion?: string,
+  ) {
+    diag.debug(`Removing patch for ${MODULE_NAME}@${moduleVersion}`);
+    this._unwrap(moduleExports.BaseAgent.prototype, "run");
+
+    _isOpenInferencePatched = false;
+    try {
+      // This can fail if the module is made immutable via the runtime or bundler
+      moduleExports.openInferencePatched = false;
+    } catch (e) {
+      diag.warn(`Failed to unset ${MODULE_NAME} patched flag on the module`, e);
+    }
+  }
+}

package/src/middleware.ts

@@ -0,0 +1,318 @@
+/**
+ * Copyright 2025 IBM Corp.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { OITracer } from "@arizeai/openinference-core";
+import { getSerializedObjectSafe } from "./helpers/getSerializedObjectSafe";
+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 { InferCallbackValue } from "beeai-framework/emitter/types";
+import { FrameworkError } from "beeai-framework/errors";
+import { Version } from "beeai-framework/version";
+import { Role } from "beeai-framework/backend/message";
+import type { GetRunContext, RunInstance } from "beeai-framework/context";
+import { traceSerializer } from "./helpers/traceSerializer";
+import {
+  errorLLMEventName,
+  finishLLMEventName,
+  INSTRUMENTATION_IGNORED_KEYS,
+  newTokenLLMEventName,
+  partialUpdateEventName,
+  successLLMEventName,
+} from "./config";
+import { createFullPath } from "beeai-framework/emitter/utils";
+import type { BeeAgent } from "beeai-framework/agents/bee/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";
+
+export const activeTracesMap = new Map<string, string>();
+
+/**
+ * First we collect, filter and transform events from the beeai emitter to our internal framework spans in the `emitter.match("*.*"` block
+ * Some special data are collected in other `emitter.match` blocks like (history, response, rawPromt)
+ *
+ * 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) {
+  return (context: GetRunContext<RunInstance, unknown>) => {
+    if (!context.emitter?.trace?.id) {
+      throw new FrameworkError(`Fatal error. Missing traceId`, [], { context });
+    }
+
+    const traceId = context.emitter?.trace?.id;
+    if (activeTracesMap.has(traceId)) {
+      return;
+    }
+    activeTracesMap.set(traceId, context.instance.constructor.name);
+
+    diag.debug("createTelemetryMiddleware", {
+      source: context.instance.constructor.name,
+      traceId: traceId,
+    });
+    const { emitter, runParams, instance } = context;
+    const basePath = createFullPath(emitter.namespace, "");
+
+    let prompt: string | undefined | null = null;
+    if (instance instanceof BaseAgent) {
+      prompt = (runParams as Parameters<BeeAgent["run"]>)[0].prompt;
+    }
+
+    const spansMap = new Map<string, FrameworkSpan>();
+    const parentIdsMap = new Map<string, number>();
+    const spansToDeleteMap = new Map<string, undefined>();
+
+    let generatedMessage: GeneratedResponse | undefined = undefined;
+    let history: GeneratedResponse[] | undefined = undefined;
+    const groupIterations: string[] = [];
+
+    const idNameManager = new IdNameManager();
+
+    const eventsIterationsMap = new Map<string, Map<string, string>>();
+
+    const startTimeDate = new Date().getTime();
+    const startTimePerf = performance.now();
+    function convertDateToPerformance(date: Date) {
+      return date.getTime() - startTimeDate + startTimePerf;
+    }
+
+    const serializer = traceSerializer({
+      ignored_keys: INSTRUMENTATION_IGNORED_KEYS,
+    });
+
+    function cleanSpanSources({ spanId }: { spanId: string }) {
+      const parentId = spansMap.get(spanId)?.parent_id;
+      if (!parentId) {
+        return;
+      }
+
+      const spanCount = parentIdsMap.get(parentId);
+      if (!spanCount) {
+        return;
+      }
+
+      if (spanCount > 1) {
+        // increase the span count for the parentId
+        parentIdsMap.set(parentId, spanCount - 1);
+      } else if (spanCount === 1) {
+        parentIdsMap.delete(parentId);
+        // check the `spansToDelete` if the span should be deleted when it has no children's
+        if (spansToDeleteMap.has(parentId)) {
+          spansMap.delete(parentId);
+          spansToDeleteMap.delete(parentId);
+        }
+      }
+    }
+
+    /**
+     * Create OpenTelemetry spans from collected data
+     */
+    emitter.match(
+      (event) => event.path === `${basePath}.run.${finishLLMEventName}`,
+      async (_, meta) => {
+        try {
+          diag.debug("run finish event", { path: meta.path, traceId: traceId });
+
+          if (!prompt && instance instanceof BaseAgent) {
+            prompt = findLast(
+              instance.memory.messages,
+              (message) => message.role === Role.USER,
+            )?.text;
+
+            if (!prompt) {
+              throw new FrameworkError(
+                "The prompt must be defined for the Agent's run",
+                [],
+                {
+                  context,
+                },
+              );
+            }
+          }
+
+          // create tracer spans from collected data
+          buildTraceTree({
+            tracer,
+            data: {
+              prompt: prompt,
+              history,
+              generatedMessage,
+              spans: JSON.parse(serializer(Array.from(spansMap.values()))),
+              traceId,
+              version: Version,
+              runErrorSpanKey: `${basePath}.run.${errorLLMEventName}`,
+              startTime: startTimePerf,
+              endTime: performance.now(),
+              source: activeTracesMap.get(traceId)!,
+            },
+          });
+        } catch (e) {
+          diag.warn("Instrumentation error", e);
+        } finally {
+          activeTracesMap.delete(traceId);
+        }
+      },
+    );
+
+    /**
+     * This block collects all "not run category" events with their data and prepares spans for the OpenTelemetry.
+     * The huge number of `newToken` events are skipped and only the last one for each parent event is saved because of `generated_token_count` information
+     * The framework event tree structure is different from the open-telemetry tree structure and must be transformed from groupId and parentGroupId pattern via idNameManager
+     * The artificial "iteration" main tree level is computed from the `meta.groupId`
+     */
+    emitter.match("*.*", (data, meta) => {
+      // allow `run.error` event due to the runtime error information
+      if (
+        meta.path.includes(".run.") &&
+        meta.path !== `${basePath}.run.${errorLLMEventName}`
+      ) {
+        return;
+      }
+      // skip all new token events
+      if (meta.name === newTokenLLMEventName) {
+        return;
+      }
+      if (!meta.trace?.runId) {
+        throw new FrameworkError(
+          `Fatal error. Missing runId for event: ${meta.path}`,
+          [],
+          {
+            context,
+          },
+        );
+      }
+
+      /**
+       * create groupId span level (id does not exist)
+       * I use only the top-level groups like iterations other nested groups like tokens would introduce unuseful complexity
+       */
+      if (
+        meta.groupId &&
+        !meta.trace.parentRunId &&
+        !groupIterations.includes(meta.groupId)
+      ) {
+        spansMap.set(
+          meta.groupId,
+          createSpan({
+            id: meta.groupId,
+            name: meta.groupId,
+            target: "groupId",
+            data: {
+              [SemanticConventions.OPENINFERENCE_SPAN_KIND]: "Chain",
+            },
+            startedAt: convertDateToPerformance(meta.createdAt),
+          }),
+        );
+        groupIterations.push(meta.groupId);
+      }
+
+      const { spanId, parentSpanId } = idNameManager.getIds({
+        path: meta.path,
+        id: meta.id,
+        runId: meta.trace.runId,
+        parentRunId: meta.trace.parentRunId,
+        groupId: meta.groupId,
+      });
+
+      const serializedData = getSerializedObjectSafe(data, meta);
+
+      // skip partialUpdate events with no data
+      if (meta.name === partialUpdateEventName && isEmpty(serializedData)) {
+        return;
+      }
+
+      const span = createSpan({
+        id: spanId,
+        name: meta.name,
+        target: meta.path,
+        ...(parentSpanId && { parent: { id: parentSpanId } }),
+        ctx: meta.context,
+        data: serializedData,
+        error: getErrorSafe(data),
+        startedAt: convertDateToPerformance(meta.createdAt),
+      });
+
+      const lastIteration = groupIterations[groupIterations.length - 1];
+
+      // delete the `partialUpdate` event if does not have nested spans
+      const lastIterationEventSpanId = eventsIterationsMap
+        .get(lastIteration)
+        ?.get(meta.name);
+      if (
+        lastIterationEventSpanId &&
+        partialUpdateEventName === meta.name &&
+        spansMap.has(lastIterationEventSpanId)
+      ) {
+        const { context } = spansMap.get(lastIterationEventSpanId)!;
+        if (parentIdsMap.has(context.span_id)) {
+          spansToDeleteMap.set(lastIterationEventSpanId, undefined);
+        } else {
+          // delete span
+          cleanSpanSources({ spanId: lastIterationEventSpanId });
+          spansMap.delete(lastIterationEventSpanId);
+        }
+      }
+
+      // create new span
+      spansMap.set(span.context.span_id, span);
+      // update number of nested spans for parent_id if exists
+      if (span.parent_id) {
+        parentIdsMap.set(
+          span.parent_id,
+          (parentIdsMap.get(span.parent_id) || 0) + 1,
+        );
+      }
+
+      // save the last event for each iteration
+      if (groupIterations.length > 0) {
+        if (eventsIterationsMap.has(lastIteration)) {
+          eventsIterationsMap
+            .get(lastIteration)!
+            .set(meta.name, span.context.span_id);
+        } else {
+          eventsIterationsMap.set(
+            lastIteration,
+            new Map().set(meta.name, span.context.span_id),
+          );
+        }
+      }
+    });
+
+    // The generated response and message history are collected from the `success` agent's event
+    emitter.match(
+      (event) =>
+        event.name === successLLMEventName &&
+        event.creator instanceof BaseAgent,
+      (data: InferCallbackValue<BeeCallbacks["success"]>) => {
+        const { data: dataObject, memory } = data;
+
+        generatedMessage = {
+          role: dataObject.role,
+          text: dataObject.text,
+        };
+        history = memory.messages.map((msg) => ({
+          text: msg.text,
+          role: msg.role,
+        }));
+      },
+    );
+  };
+}

package/src/types.ts

@@ -0,0 +1,38 @@
+/**
+ * Copyright 2025 IBM Corp.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { Attributes, SpanStatus, TimeInput } from "@opentelemetry/api";
+
+export interface FrameworkSpan {
+  attributes: {
+    metadata?: Attributes | null;
+    data?: Attributes | null;
+    target: string;
+  };
+  context: {
+    span_id: string;
+  };
+  end_time: TimeInput;
+  name: string;
+  parent_id?: string;
+  start_time: TimeInput;
+  status: SpanStatus;
+}
+
+export interface GeneratedResponse {
+  text: string;
+  role: string;
+}