@arizeai/openinference-instrumentation-beeai 1.1.1 → 1.2.0

package/src/middleware.ts

@@ -36,7 +36,7 @@ import {
   successLLMEventName,
 } from "./config";
 import { createFullPath } from "beeai-framework/emitter/utils";
-import type { ReActAgent } from "beeai-framework/agents/react/agent";
+import { ReActAgent } from "beeai-framework/agents/react/agent";
 import { BaseAgent } from "beeai-framework/agents/base";
 import { diag } from "@opentelemetry/api";
 import { FrameworkSpan, GeneratedResponse } from "./types";
@@ -45,6 +45,8 @@ import {
   OpenInferenceSpanKind,
   SemanticConventions,
 } from "@arizeai/openinference-semantic-conventions";
+import { ToolCallingAgent } from "beeai-framework/agents/toolCalling/agent";
+import type { ToolCallingAgentCallbacks } from "beeai-framework/agents/toolCalling/types";
 
 export const activeTracesMap = new Map<string, string>();
 
@@ -207,118 +209,158 @@ export function createTelemetryMiddleware(
         );
       }
 
-      /**
-       * 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,
-      });
+      try {
+        /**
+         * 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 serializedData = getSerializedObjectSafe(data, meta);
+        const { spanId, parentSpanId } = idNameManager.getIds({
+          path: meta.path,
+          id: meta.id,
+          runId: meta.trace.runId,
+          parentRunId: meta.trace.parentRunId,
+          groupId: meta.groupId,
+        });
 
-      // skip partialUpdate events with no data
-      if (meta.name === partialUpdateEventName && isEmpty(serializedData)) {
-        return;
-      }
+        const serializedData = getSerializedObjectSafe(data, meta);
 
-      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);
+        // skip partialUpdate events with no data
+        if (meta.name === partialUpdateEventName && isEmpty(serializedData)) {
+          return;
         }
-      }
 
-      // 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,
-        );
-      }
+        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),
+        });
 
-      // 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),
+        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),
+            );
+          }
+        }
+      } catch (e) {
+        diag.warn("Instrumentation error", e);
       }
     });
 
-    // The generated response and message history are collected from the `success` agent's event
+    // The generated response and message history are collected from the `success` event generated by ReActAgent
     emitter.match(
       (event) =>
         event.name === successLLMEventName &&
-        event.creator instanceof BaseAgent,
+        event.creator instanceof ReActAgent,
       (data: InferCallbackValue<ReActAgentCallbacks["success"]>) => {
-        const { data: dataObject, memory } = data;
-
-        generatedMessage = {
-          role: dataObject.role,
-          text: dataObject.text,
-        };
-        history = memory.messages.map((msg) => ({
-          text: msg.text,
-          role: msg.role,
-        }));
+        try {
+          const { data: dataObject, memory } = data;
+
+          generatedMessage = {
+            role: dataObject.role,
+            text: dataObject.text,
+          };
+          history = memory.messages.map((msg) => ({
+            text: msg.text,
+            role: msg.role,
+          }));
+        } catch (e) {
+          diag.warn(
+            "Instrumentation error. Unable to map messages to history for ReActAgent",
+            e,
+          );
+        }
+      },
+    );
+
+    // The generated response and message history are collected from the `success` event generated by ToolCallingAgentCallbacks
+    emitter.match(
+      (event) =>
+        event.name === successLLMEventName &&
+        event.creator instanceof ToolCallingAgent,
+      (data: InferCallbackValue<ToolCallingAgentCallbacks["success"]>) => {
+        try {
+          const { state } = data;
+
+          if (state.result) {
+            generatedMessage = {
+              role: state.result.role,
+              text: state.result.text,
+            };
+          }
+
+          history = state.memory.messages.map((msg) => ({
+            text: msg.text,
+            role: msg.role,
+          }));
+        } catch (e) {
+          diag.warn(
+            "Instrumentation error. Unable to map messages to history for ToolCallingAgent",
+            e,
+          );
+        }
       },
     );
   };

package/dist/examples/run-agent.d.ts

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

package/dist/examples/run-agent.d.ts.map

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

package/dist/examples/run-agent.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"run-agent.js","sourceRoot":"","sources":["../../examples/run-agent.ts"],"names":[],"mappings":";;;;;;;;;;;AAAA,6BAA2B;AAC3B,8DAAgE;AAChE,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,kBAAU,CAAC;IAC3B,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"}