@arizeai/openinference-instrumentation-beeai 1.5.22 → 1.5.24

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@arizeai/openinference-instrumentation-beeai",
-  "version": "1.5.22",
+  "version": "1.5.24",
   "private": false,
   "description": "OpenInference Instrumentation for BeeAI framework",
   "keywords": [
@@ -33,8 +33,8 @@
     "@opentelemetry/instrumentation": "^0.57.1",
     "remeda": "^2.20.2",
     "semver": "^7.7.0",
-    "@arizeai/openinference-core": "2.5.2",
-    "@arizeai/openinference-semantic-conventions": "2.7.0"
+    "@arizeai/openinference-core": "2.5.4",
+    "@arizeai/openinference-semantic-conventions": "2.8.0"
   },
   "devDependencies": {
     "@opentelemetry/exporter-trace-otlp-proto": "^0.50.0",

package/src/helpers/buildTraceTree.ts

@@ -111,26 +111,25 @@ function buildLLMMainSpanData(data: BuiltTraceTreeProps["data"]) {
 
   if (!startBeeaiSpan && !successBeeeaiSpan) return {};
 
+  const startData = startBeeaiSpan?.attributes.data;
+  const successData = successBeeeaiSpan?.attributes.data;
+
   const provider =
-    startBeeaiSpan?.attributes.data?.[SemanticConventions.LLM_PROVIDER] ||
-    successBeeeaiSpan?.attributes.data?.[SemanticConventions.LLM_PROVIDER];
+    startData?.[SemanticConventions.LLM_PROVIDER] ||
+    successData?.[SemanticConventions.LLM_PROVIDER];
 
   const modelName =
-    startBeeaiSpan?.attributes.data?.[SemanticConventions.LLM_MODEL_NAME] ||
-    successBeeeaiSpan?.attributes.data?.[SemanticConventions.LLM_MODEL_NAME];
+    startData?.[SemanticConventions.LLM_MODEL_NAME] ||
+    successData?.[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],
+      [SemanticConventions.INPUT_VALUE]: startData?.[SemanticConventions.INPUT_VALUE],
+      [SemanticConventions.INPUT_MIME_TYPE]: startData?.[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],
+      [SemanticConventions.OUTPUT_MIME_TYPE]: successData?.[SemanticConventions.OUTPUT_MIME_TYPE],
+      [SemanticConventions.OUTPUT_VALUE]: successData?.[SemanticConventions.OUTPUT_VALUE],
     }),
     ...(provider && { [SemanticConventions.LLM_PROVIDER]: provider }),
     ...(modelName && { [SemanticConventions.LLM_MODEL_NAME]: modelName }),

package/src/helpers/getErrorSafe.ts

@@ -23,4 +23,5 @@ export function getErrorSafe(data: unknown): string | undefined {
   if (error instanceof FrameworkError) {
     return FrameworkError.ensure(error).explain();
   }
+  return undefined;
 }

package/src/helpers/getSerializedObjectSafe.ts

@@ -98,6 +98,196 @@ function parseLLMOutputMessages(messages: readonly unknown[]) {
 
 const matchesEvent = (name: string, events: readonly string[]): boolean => events.includes(name);
 
+const NATIVE_TOOL_EVENT_NAMES = [
+  startToolEventName,
+  successToolEventName,
+  finishToolEventName,
+  errorToolEventName,
+  retryToolEventName,
+] as const;
+
+/** Maps an error onto the OpenTelemetry exception attributes. */
+function getExceptionAttributes(error: Error | undefined) {
+  if (!error) {
+    return {};
+  }
+  return {
+    "exception.message": error.message,
+    "exception.stacktrace": error.stack,
+    "exception.type": error.name,
+  };
+}
+
+/** Extracts the attributes for an agent start/success/error/retry event. */
+function getAgentEventAttributes(dataObject: unknown) {
+  const event = isObjectWithStringKeys(dataObject) ? dataObject : {};
+  const agentMeta = isObjectWithStringKeys(event.meta) ? event.meta : undefined;
+  const tools = Array.isArray(event.tools) ? event.tools.filter(isObjectWithStringKeys) : [];
+  const memory = isObjectWithStringKeys(event.memory) ? event.memory : undefined;
+  const messages = memory && Array.isArray(memory.messages) ? memory.messages : [];
+  const error = event.error instanceof Error ? event.error : undefined;
+  const data = event.data;
+  return {
+    [SemanticConventions.OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.AGENT,
+    ...(typeof agentMeta?.iteration === "number" && { iteration: agentMeta.iteration }),
+    ...(tools?.length > 0 && {
+      [SemanticConventions.LLM_TOOLS]: tools.map((tool) => ({
+        [SemanticConventions.TOOL_NAME]: typeof tool.name === "string" ? tool.name : undefined,
+        [SemanticConventions.TOOL_DESCRIPTION]:
+          typeof tool.description === "string" ? tool.description : undefined,
+        "tool.options": tool.options,
+      })),
+    }),
+    ...(messages.length > 0 && {
+      [SemanticConventions.INPUT_MIME_TYPE]: MimeType.JSON,
+      [SemanticConventions.INPUT_VALUE]: JSON.stringify(messages),
+    }),
+    ...getExceptionAttributes(error),
+    ...(data != null
+      ? {
+          [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.JSON,
+          [SemanticConventions.OUTPUT_VALUE]: JSON.stringify(data),
+        }
+      : {}),
+  };
+}
+
+/** Extracts the attributes for an agent update/partial update event. */
+function getUpdateEventAttributes(dataObject: unknown) {
+  const event = isObjectWithStringKeys(dataObject) ? dataObject : {};
+  const data = isObjectWithStringKeys(event.data) ? event.data : {};
+
+  const output = data.final_answer || data.tool_output;
+  return {
+    [SemanticConventions.OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.AGENT,
+    ...(typeof data.thought === "string" && { thought: data.thought }),
+    ...(typeof data.tool_name === "string" && {
+      [SemanticConventions.TOOL_NAME]: data.tool_name,
+    }),
+    ...(data.tool_input != null
+      ? { [SemanticConventions.TOOL_PARAMETERS]: JSON.stringify(data.tool_input) }
+      : {}),
+    ...(typeof output === "string" && {
+      [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.JSON,
+      [SemanticConventions.OUTPUT_VALUE]: output,
+    }),
+  };
+}
+
+/** Extracts the attributes for a tool event emitted by an agent. */
+function getAgentToolEventAttributes(dataObject: unknown) {
+  const event = isObjectWithStringKeys(dataObject) ? dataObject : {};
+  const data = isObjectWithStringKeys(event.data) ? event.data : {};
+  const iteration = isObjectWithStringKeys(data.iteration) ? data.iteration : undefined;
+  const tool = isObjectWithStringKeys(data.tool) ? data.tool : undefined;
+  const error = data.error instanceof Error ? data.error : undefined;
+  const output = data.result instanceof Serializable ? data.result.createSnapshot() : undefined;
+
+  return {
+    [SemanticConventions.OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.TOOL,
+    ...(typeof iteration?.thought === "string" && { thought: iteration.thought }),
+    ...(data?.input
+      ? {
+          [SemanticConventions.TOOL_PARAMETERS]: JSON.stringify(data.input),
+        }
+      : {}),
+    ...(typeof tool?.description === "string" && {
+      [SemanticConventions.TOOL_DESCRIPTION]: tool.description,
+    }),
+    ...(typeof tool?.name === "string" && {
+      [SemanticConventions.TOOL_NAME]: tool.name,
+    }),
+    ...getExceptionAttributes(error),
+    ...(output != null ? { [SemanticConventions.OUTPUT_VALUE]: JSON.stringify(output) } : {}),
+  };
+}
+
+/** Extracts the attributes for a tool event emitted by a Tool itself. */
+function getNativeToolEventAttributes(dataObject: unknown) {
+  if (!dataObject) {
+    return {
+      [SemanticConventions.OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.TOOL,
+    };
+  }
+  const event = isObjectWithStringKeys(dataObject) ? dataObject : {};
+  const { input, output } = event;
+  const error = event.error instanceof Error ? event.error : undefined;
+
+  return {
+    [SemanticConventions.OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.TOOL,
+    ...(input != null ? { [SemanticConventions.TOOL_PARAMETERS]: JSON.stringify(input) } : {}),
+    ...(output != null ? { [SemanticConventions.OUTPUT_VALUE]: JSON.stringify(output) } : {}),
+    ...getExceptionAttributes(error),
+  };
+}
+
+/** Extracts the model/provider attributes from a ChatModel snapshot. */
+function getChatModelCreatorAttributes(creator: ChatModel) {
+  const creatorSnapshot = creator.createSnapshot();
+  const snapshot: Record<string, unknown> = isObjectWithStringKeys(creatorSnapshot)
+    ? creatorSnapshot
+    : {};
+  return {
+    ...(typeof snapshot.providerId === "string" && {
+      [SemanticConventions.LLM_PROVIDER]: snapshot.providerId,
+      [`${SemanticAttributePrefixes.metadata}.${LLMAttributePostfixes.provider}`]:
+        snapshot.providerId,
+    }),
+    ...(typeof snapshot.modelId === "string" && {
+      [SemanticConventions.LLM_MODEL_NAME]: snapshot.modelId,
+      [`${SemanticAttributePrefixes.metadata}.${LLMAttributePostfixes.model_name}`]:
+        snapshot.modelId,
+    }),
+    ...(snapshot.parameters != null
+      ? {
+          [SemanticConventions.LLM_INVOCATION_PARAMETERS]: JSON.stringify(snapshot.parameters),
+        }
+      : {}),
+  };
+}
+
+/** Extracts the attributes for an LLM start/success/error/new-token event. */
+function getLLMEventAttributes({
+  dataObject,
+  creator,
+}: {
+  dataObject: unknown;
+  creator: ChatModel;
+}) {
+  const event = isObjectWithStringKeys(dataObject) ? dataObject : {};
+  const value = isObjectWithStringKeys(event.value) ? event.value : undefined;
+  const input = isObjectWithStringKeys(event.input) ? event.input : undefined;
+  const usage = value && isObjectWithStringKeys(value.usage) ? value.usage : undefined;
+  const inputMessages = input && Array.isArray(input.messages) ? input.messages : [];
+  const outputMessages = value && Array.isArray(value.messages) ? value.messages : [];
+  const error = event.error instanceof Error ? event.error : undefined;
+
+  return {
+    [SemanticConventions.OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.LLM,
+    ...(typeof usage?.completionTokens === "number" && {
+      [SemanticConventions.LLM_TOKEN_COUNT_COMPLETION]: usage.completionTokens,
+    }),
+    ...(typeof usage?.promptTokens === "number" && {
+      [SemanticConventions.LLM_TOKEN_COUNT_PROMPT]: usage.promptTokens,
+    }),
+    ...(typeof usage?.totalTokens === "number" && {
+      [SemanticConventions.LLM_TOKEN_COUNT_TOTAL]: usage.totalTokens,
+    }),
+    ...(inputMessages.length > 0 && {
+      [SemanticConventions.INPUT_MIME_TYPE]: MimeType.JSON,
+      [SemanticConventions.INPUT_VALUE]: JSON.stringify(inputMessages),
+      ...parserLLMInputMessages(inputMessages),
+    }),
+    ...(outputMessages.length > 0 && {
+      [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.JSON,
+      [SemanticConventions.OUTPUT_VALUE]: JSON.stringify(outputMessages),
+      ...parseLLMOutputMessages(outputMessages),
+    }),
+    ...getExceptionAttributes(error),
+    ...getChatModelCreatorAttributes(creator),
+  };
+}
+
 export function getSerializedObjectSafe(dataObject: unknown, meta: EventMeta<unknown>) {
   try {
     // agent events
@@ -105,125 +295,22 @@ export function getSerializedObjectSafe(dataObject: unknown, meta: EventMeta<unk
       matchesEvent(meta.name, [startEventName, successEventName, errorEventName, retryEventName]) &&
       meta.creator instanceof BaseAgent
     ) {
-      const event = isObjectWithStringKeys(dataObject) ? dataObject : {};
-      const agentMeta = isObjectWithStringKeys(event.meta) ? event.meta : undefined;
-      const tools = Array.isArray(event.tools) ? event.tools.filter(isObjectWithStringKeys) : [];
-      const memory = isObjectWithStringKeys(event.memory) ? event.memory : undefined;
-      const messages = memory && Array.isArray(memory.messages) ? memory.messages : [];
-      const error = event.error instanceof Error ? event.error : undefined;
-      const data = event.data;
-      return {
-        [SemanticConventions.OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.AGENT,
-        ...(typeof agentMeta?.iteration === "number" && { iteration: agentMeta.iteration }),
-        ...(tools?.length > 0 && {
-          [SemanticConventions.LLM_TOOLS]: tools.map((tool) => ({
-            [SemanticConventions.TOOL_NAME]: typeof tool.name === "string" ? tool.name : undefined,
-            [SemanticConventions.TOOL_DESCRIPTION]:
-              typeof tool.description === "string" ? tool.description : undefined,
-            "tool.options": tool.options,
-          })),
-        }),
-        ...(messages.length > 0 && {
-          [SemanticConventions.INPUT_MIME_TYPE]: MimeType.JSON,
-          [SemanticConventions.INPUT_VALUE]: JSON.stringify(messages),
-        }),
-        ...(error && {
-          "exception.message": error.message,
-          "exception.stacktrace": error.stack,
-          "exception.type": error.name,
-        }),
-        ...(data != null
-          ? {
-              [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.JSON,
-              [SemanticConventions.OUTPUT_VALUE]: JSON.stringify(data),
-            }
-          : {}),
-      };
+      return getAgentEventAttributes(dataObject);
     }
 
     // update events
     if (matchesEvent(meta.name, [updateEventName, partialUpdateEventName])) {
-      const event = isObjectWithStringKeys(dataObject) ? dataObject : {};
-      const data = isObjectWithStringKeys(event.data) ? event.data : {};
-
-      const output = data.final_answer || data.tool_output;
-      return {
-        [SemanticConventions.OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.AGENT,
-        ...(typeof data.thought === "string" && { thought: data.thought }),
-        ...(typeof data.tool_name === "string" && {
-          [SemanticConventions.TOOL_NAME]: data.tool_name,
-        }),
-        ...(data.tool_input != null
-          ? { [SemanticConventions.TOOL_PARAMETERS]: JSON.stringify(data.tool_input) }
-          : {}),
-        ...(typeof output === "string" && {
-          [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.JSON,
-          [SemanticConventions.OUTPUT_VALUE]: output,
-        }),
-      };
+      return getUpdateEventAttributes(dataObject);
     }
 
     // tool events (from agent)
     if (matchesEvent(meta.name, [toolErrorEventName, toolStartEventName, toolSuccessEventName])) {
-      const event = isObjectWithStringKeys(dataObject) ? dataObject : {};
-      const data = isObjectWithStringKeys(event.data) ? event.data : {};
-      const iteration = isObjectWithStringKeys(data.iteration) ? data.iteration : undefined;
-      const tool = isObjectWithStringKeys(data.tool) ? data.tool : undefined;
-      const error = data.error instanceof Error ? data.error : undefined;
-      const output = data.result instanceof Serializable ? data.result.createSnapshot() : undefined;
-
-      return {
-        [SemanticConventions.OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.TOOL,
-        ...(typeof iteration?.thought === "string" && { thought: iteration.thought }),
-        ...(data?.input
-          ? {
-              [SemanticConventions.TOOL_PARAMETERS]: JSON.stringify(data.input),
-            }
-          : {}),
-        ...(typeof tool?.description === "string" && {
-          [SemanticConventions.TOOL_DESCRIPTION]: tool.description,
-        }),
-        ...(typeof tool?.name === "string" && {
-          [SemanticConventions.TOOL_NAME]: tool.name,
-        }),
-        ...(error && {
-          "exception.message": error.message,
-          "exception.stacktrace": error.stack,
-          "exception.type": error.name,
-        }),
-        ...(output != null ? { [SemanticConventions.OUTPUT_VALUE]: JSON.stringify(output) } : {}),
-      };
+      return getAgentToolEventAttributes(dataObject);
     }
-    // tool events native
-    if (
-      [
-        startToolEventName,
-        successToolEventName,
-        finishToolEventName,
-        errorToolEventName,
-        retryToolEventName,
-      ].some((eventName) => eventName === meta.name) &&
-      meta.creator instanceof Tool
-    ) {
-      if (!dataObject) {
-        return {
-          [SemanticConventions.OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.TOOL,
-        };
-      }
-      const event = isObjectWithStringKeys(dataObject) ? dataObject : {};
-      const { input, output } = event;
-      const error = event.error instanceof Error ? event.error : undefined;
 
-      return {
-        [SemanticConventions.OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.TOOL,
-        ...(input != null ? { [SemanticConventions.TOOL_PARAMETERS]: JSON.stringify(input) } : {}),
-        ...(output != null ? { [SemanticConventions.OUTPUT_VALUE]: JSON.stringify(output) } : {}),
-        ...(error && {
-          "exception.message": error.message,
-          "exception.stacktrace": error.stack,
-          "exception.type": error.name,
-        }),
-      };
+    // tool events native
+    if (matchesEvent(meta.name, NATIVE_TOOL_EVENT_NAMES) && meta.creator instanceof Tool) {
+      return getNativeToolEventAttributes(dataObject);
     }
 
     // llm events
@@ -236,60 +323,7 @@ export function getSerializedObjectSafe(dataObject: unknown, meta: EventMeta<unk
       ]) &&
       meta.creator instanceof ChatModel
     ) {
-      const event = isObjectWithStringKeys(dataObject) ? dataObject : {};
-      const value = isObjectWithStringKeys(event.value) ? event.value : undefined;
-      const input = isObjectWithStringKeys(event.input) ? event.input : undefined;
-      const usage = value && isObjectWithStringKeys(value.usage) ? value.usage : undefined;
-      const inputMessages = input && Array.isArray(input.messages) ? input.messages : [];
-      const outputMessages = value && Array.isArray(value.messages) ? value.messages : [];
-      const error = event.error instanceof Error ? event.error : undefined;
-
-      const creatorSnapshot = meta.creator.createSnapshot();
-      const creator: Record<string, unknown> = isObjectWithStringKeys(creatorSnapshot)
-        ? creatorSnapshot
-        : {};
-      return {
-        [SemanticConventions.OPENINFERENCE_SPAN_KIND]: OpenInferenceSpanKind.LLM,
-        ...(typeof usage?.completionTokens === "number" && {
-          [SemanticConventions.LLM_TOKEN_COUNT_COMPLETION]: usage.completionTokens,
-        }),
-        ...(typeof usage?.promptTokens === "number" && {
-          [SemanticConventions.LLM_TOKEN_COUNT_PROMPT]: usage.promptTokens,
-        }),
-        ...(typeof usage?.totalTokens === "number" && {
-          [SemanticConventions.LLM_TOKEN_COUNT_TOTAL]: usage.totalTokens,
-        }),
-        ...(inputMessages.length > 0 && {
-          [SemanticConventions.INPUT_MIME_TYPE]: MimeType.JSON,
-          [SemanticConventions.INPUT_VALUE]: JSON.stringify(inputMessages),
-          ...parserLLMInputMessages(inputMessages),
-        }),
-        ...(outputMessages.length > 0 && {
-          [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.JSON,
-          [SemanticConventions.OUTPUT_VALUE]: JSON.stringify(outputMessages),
-          ...parseLLMOutputMessages(outputMessages),
-        }),
-        ...(error && {
-          "exception.message": error.message,
-          "exception.stacktrace": error.stack,
-          "exception.type": error.name,
-        }),
-        ...(typeof creator.providerId === "string" && {
-          [SemanticConventions.LLM_PROVIDER]: creator.providerId,
-          [`${SemanticAttributePrefixes.metadata}.${LLMAttributePostfixes.provider}`]:
-            creator.providerId,
-        }),
-        ...(typeof creator.modelId === "string" && {
-          [SemanticConventions.LLM_MODEL_NAME]: creator.modelId,
-          [`${SemanticAttributePrefixes.metadata}.${LLMAttributePostfixes.model_name}`]:
-            creator.modelId,
-        }),
-        ...(creator.parameters != null
-          ? {
-              [SemanticConventions.LLM_INVOCATION_PARAMETERS]: JSON.stringify(creator.parameters),
-            }
-          : {}),
-      };
+      return getLLMEventAttributes({ dataObject, creator: meta.creator });
     }
     if (meta.name === finishLLMEventName && meta.creator instanceof ChatModel) {
       return {

package/src/helpers/traceSerializer.ts

@@ -26,7 +26,7 @@ export function traceSerializer({ ignored_keys = [] }: { ignored_keys?: string[]
         return (key, value) => {
           // Ignore specific keys, all owned entities and keys starting with underscore
           if (mergedIgnoreKeys.has(key) || key.startsWith("_")) {
-            return;
+            return undefined;
           }
 
           return value;

package/src/middleware.ts

@@ -178,6 +178,90 @@ export function createTelemetryMiddleware(tracer: OITracer, mainSpanKind: OpenIn
       },
     );
 
+    /**
+     * Creates the artificial "iteration" tree level for a top-level groupId the
+     * first time it is seen. Nested groups (like tokens) are skipped because
+     * they would introduce unuseful complexity.
+     */
+    function createGroupSpanIfNeeded({
+      groupId,
+      parentRunId,
+      createdAt,
+    }: {
+      groupId: string | undefined;
+      parentRunId: string | undefined;
+      createdAt: Date;
+    }) {
+      if (!groupId || parentRunId || groupIterations.includes(groupId)) {
+        return;
+      }
+      spansMap.set(
+        groupId,
+        createSpan({
+          id: groupId,
+          name: groupId,
+          target: "groupId",
+          data: {
+            [SemanticConventions.OPENINFERENCE_SPAN_KIND]: "Chain",
+          },
+          startedAt: convertDateToPerformance(createdAt),
+        }),
+      );
+      groupIterations.push(groupId);
+    }
+
+    /**
+     * Drops the previous `partialUpdate` span for this iteration once a newer one
+     * arrives, unless it has nested spans, in which case it is only marked for
+     * deletion.
+     */
+    function dropSupersededPartialUpdateSpan({
+      eventName,
+      lastIteration,
+    }: {
+      eventName: string;
+      lastIteration: string | undefined;
+    }) {
+      if (lastIteration == null || partialUpdateEventName !== eventName) {
+        return;
+      }
+      const lastIterationEventSpanId = eventsIterationsMap.get(lastIteration)?.get(eventName);
+      if (!lastIterationEventSpanId || !spansMap.has(lastIterationEventSpanId)) {
+        return;
+      }
+      const { context: spanContext } = spansMap.get(lastIterationEventSpanId)!;
+      if (parentIdsMap.has(spanContext.span_id)) {
+        spansToDeleteMap.set(lastIterationEventSpanId, undefined);
+        return;
+      }
+      // delete span
+      cleanSpanSources({ spanId: lastIterationEventSpanId });
+      spansMap.delete(lastIterationEventSpanId);
+    }
+
+    /**
+     * Saves the last event span for each iteration.
+     */
+    function recordIterationEventSpan({
+      eventName,
+      lastIteration,
+      spanId,
+    }: {
+      eventName: string;
+      lastIteration: string | undefined;
+      spanId: string;
+    }) {
+      if (lastIteration == null) {
+        return;
+      }
+      const iterationEvents = eventsIterationsMap.get(lastIteration);
+      if (iterationEvents) {
+        iterationEvents.set(eventName, spanId);
+      } else {
+        eventsIterationsMap.set(lastIteration, new Map<string, string>([[eventName, spanId]]));
+      }
+    }
+
     /**
      * 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
@@ -204,21 +288,11 @@ export function createTelemetryMiddleware(tracer: OITracer, mainSpanKind: OpenIn
          * 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);
-        }
+        createGroupSpanIfNeeded({
+          groupId: meta.groupId,
+          parentRunId: meta.trace.parentRunId,
+          createdAt: meta.createdAt,
+        });
 
         const { spanId, parentSpanId } = idNameManager.getIds({
           path: meta.path,
@@ -252,21 +326,7 @@ export function createTelemetryMiddleware(tracer: OITracer, mainSpanKind: OpenIn
         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);
-          }
-        }
+        dropSupersededPartialUpdateSpan({ eventName: meta.name, lastIteration });
 
         // create new span
         spansMap.set(span.context.span_id, span);
@@ -276,13 +336,11 @@ export function createTelemetryMiddleware(tracer: OITracer, mainSpanKind: OpenIn
         }
 
         // 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));
-          }
-        }
+        recordIterationEventSpan({
+          eventName: meta.name,
+          lastIteration,
+          spanId: span.context.span_id,
+        });
       } catch (e) {
         diag.warn("Instrumentation error", e);
       }