feat(core): Instrument tool calls for Anthropic AI

RulaKhaled · RulaKhaled · commit 74c18ee3c94e · 2025-08-27T17:25:44.000+02:00
diff --git a/packages/core/src/utils/anthropic-ai/constants.ts b/packages/core/src/utils/anthropic-ai/constants.ts
@@ -9,4 +9,5 @@ export const ANTHROPIC_AI_INSTRUMENTED_METHODS = [
   'models.get',
   'completions.create',
   'models.retrieve',
+  'beta.messages.create',
 ] as const;
diff --git a/packages/core/src/utils/anthropic-ai/index.ts b/packages/core/src/utils/anthropic-ai/index.ts
@@ -8,6 +8,7 @@ import {
   ANTHROPIC_AI_RESPONSE_TIMESTAMP_ATTRIBUTE,
   GEN_AI_OPERATION_NAME_ATTRIBUTE,
   GEN_AI_PROMPT_ATTRIBUTE,
+  GEN_AI_REQUEST_AVAILABLE_TOOLS_ATTRIBUTE,
   GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE,
   GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE,
   GEN_AI_REQUEST_MESSAGES_ATTRIBUTE,
@@ -19,6 +20,7 @@ import {
   GEN_AI_RESPONSE_ID_ATTRIBUTE,
   GEN_AI_RESPONSE_MODEL_ATTRIBUTE,
   GEN_AI_RESPONSE_TEXT_ATTRIBUTE,
+  GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE,
   GEN_AI_SYSTEM_ATTRIBUTE,
 } from '../ai/gen-ai-attributes';
 import { buildMethodPath, getFinalOperationName, getSpanOperation, setTokenUsageAttributes } from '../ai/utils';
@@ -31,6 +33,7 @@ import type {
   AnthropicAiOptions,
   AnthropicAiResponse,
   AnthropicAiStreamingEvent,
+  ContentBlock,
 } from './types';
 import { shouldInstrument } from './utils';
 
@@ -46,6 +49,9 @@ function extractRequestAttributes(args: unknown[], methodPath: string): Record<s
 
   if (args.length > 0 && typeof args[0] === 'object' && args[0] !== null) {
     const params = args[0] as Record<string, unknown>;
+    if (params.tools && Array.isArray(params.tools)) {
+      attributes[GEN_AI_REQUEST_AVAILABLE_TOOLS_ATTRIBUTE] = JSON.stringify(params.tools);
+    }
 
     attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] = params.model ?? 'unknown';
     if ('temperature' in params) attributes[GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE] = params.temperature;
@@ -96,10 +102,20 @@ function addResponseAttributes(span: Span, response: AnthropicAiResponse, record
       if (Array.isArray(response.content)) {
         span.setAttributes({
           [GEN_AI_RESPONSE_TEXT_ATTRIBUTE]: response.content
-            .map((item: { text: string | undefined }) => item.text)
+            .map((item: ContentBlock) => item.text)
             .filter((text): text is string => text !== undefined)
             .join(''),
         });
+
+        const toolCalls: Array<ContentBlock> = [];
+
+        for (const item of response.content) {
+          if (item.type === 'tool_use' || item.type === 'server_tool_use') {
+            toolCalls.push(item);
+          }
+        }
+
+        span.setAttributes({ [GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE]: JSON.stringify(toolCalls) });
       }
     }
     // Completions.create
diff --git a/packages/core/src/utils/anthropic-ai/streaming.ts b/packages/core/src/utils/anthropic-ai/streaming.ts
@@ -7,6 +7,7 @@ import {
   GEN_AI_RESPONSE_MODEL_ATTRIBUTE,
   GEN_AI_RESPONSE_STREAMING_ATTRIBUTE,
   GEN_AI_RESPONSE_TEXT_ATTRIBUTE,
+  GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE,
 } from '../ai/gen-ai-attributes';
 import { setTokenUsageAttributes } from '../ai/utils';
 import type { AnthropicAiStreamingEvent } from './types';
@@ -32,6 +33,17 @@ interface StreamingState {
   cacheCreationInputTokens: number | undefined;
   /** Number of cache read input tokens used. */
   cacheReadInputTokens: number | undefined;
+  /** Accumulated tool calls (finalized) */
+  toolCalls: Array<Record<string, unknown>>;
+  /** In-progress tool call blocks keyed by index */
+  activeToolBlocks: Record<
+    number,
+    {
+      id?: string;
+      name?: string;
+      inputJsonParts: string[];
+    }
+  >;
 }
 
 /**
@@ -43,12 +55,7 @@ interface StreamingState {
  * @returns Whether an error occurred
  */
 
-function isErrorEvent(
-  event: AnthropicAiStreamingEvent,
-  state: StreamingState,
-  recordOutputs: boolean,
-  span: Span,
-): boolean {
+function isErrorEvent(event: AnthropicAiStreamingEvent, span: Span): boolean {
   if ('type' in event && typeof event.type === 'string') {
     // If the event is an error, set the span status and capture the error
     // These error events are not rejected by the API by default, but are sent as metadata of the response
@@ -69,11 +76,6 @@ function isErrorEvent(
       });
       return true;
     }
-
-    if (recordOutputs && event.type === 'content_block_delta') {
-      const text = event.delta?.text;
-      if (text) state.responseTexts.push(text);
-    }
   }
   return false;
 }
@@ -110,6 +112,66 @@ function handleMessageMetadata(event: AnthropicAiStreamingEvent, state: Streamin
   }
 }
 
+/**
+ * Handle start of a content block (e.g., tool_use)
+ */
+function handleContentBlockStart(event: AnthropicAiStreamingEvent, state: StreamingState): void {
+  if (event.type !== 'content_block_start' || typeof event.index !== 'number' || !event.content_block) return;
+  if (event.content_block.type === 'tool_use') {
+    state.activeToolBlocks[event.index] = {
+      id: event.content_block.id,
+      name: event.content_block.name,
+      inputJsonParts: [],
+    };
+  }
+}
+
+/**
+ * Handle deltas of a content block, including input_json_delta for tool_use
+ */
+function handleContentBlockDelta(
+  event: AnthropicAiStreamingEvent,
+  state: StreamingState,
+  recordOutputs: boolean,
+): void {
+  if (event.type !== 'content_block_delta' || typeof event.index !== 'number' || !event.delta) return;
+
+  if ('partial_json' in event.delta && typeof event.delta.partial_json === 'string') {
+    const active = state.activeToolBlocks[event.index];
+    if (active) {
+      active.inputJsonParts.push(event.delta.partial_json);
+    }
+  }
+
+  if (recordOutputs && event.delta.text) {
+    state.responseTexts.push(event.delta.text);
+  }
+}
+
+/**
+ * Handle stop of a content block; finalize tool_use entries
+ */
+function handleContentBlockStop(event: AnthropicAiStreamingEvent, state: StreamingState): void {
+  if (event.type !== 'content_block_stop' || typeof event.index !== 'number') return;
+
+  const active = state.activeToolBlocks[event.index];
+  if (!active) return;
+
+  const raw = active.inputJsonParts.join('');
+  const parsedInput = raw ? JSON.parse(raw) : {};
+
+  state.toolCalls.push({
+    type: 'tool_use',
+    id: active.id,
+    name: active.name,
+    input: parsedInput,
+  });
+
+  // Avoid deleting a dynamic key; rebuild the map without this index
+  const remainingEntries = Object.entries(state.activeToolBlocks).filter(([key]) => key !== String(event.index));
+  state.activeToolBlocks = Object.fromEntries(remainingEntries);
+}
+
 /**
  * Processes an event
  * @param event - The event to process
@@ -128,10 +190,19 @@ function processEvent(
     return;
   }
 
-  const isError = isErrorEvent(event, state, recordOutputs, span);
+  const isError = isErrorEvent(event, span);
   if (isError) return;
 
   handleMessageMetadata(event, state);
+
+  // Tool call events are sent via 3 separate events:
+  // - content_block_start (start of the tool call)
+  // - content_block_delta (delta aka input of the tool call)
+  // - content_block_stop (end of the tool call)
+  // We need to handle them all to capture the full tool call.
+  handleContentBlockStart(event, state);
+  handleContentBlockDelta(event, state, recordOutputs);
+  handleContentBlockStop(event, state);
 }
 
 /**
@@ -153,6 +224,8 @@ export async function* instrumentStream(
     completionTokens: undefined,
     cacheCreationInputTokens: undefined,
     cacheReadInputTokens: undefined,
+    toolCalls: [],
+    activeToolBlocks: {},
   };
 
   try {
@@ -197,6 +270,13 @@ export async function* instrumentStream(
       });
     }
 
+    // Set tool calls if any were captured
+    if (recordOutputs && state.toolCalls.length > 0) {
+      span.setAttributes({
+        [GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE]: JSON.stringify(state.toolCalls),
+      });
+    }
+
     span.end();
   }
 }
diff --git a/packages/core/src/utils/anthropic-ai/types.ts b/packages/core/src/utils/anthropic-ai/types.ts
@@ -16,14 +16,25 @@ export type Message = {
   content: string | unknown[];
 };
 
+export type ContentBlock = {
+  type: 'tool_use' | 'server_tool_use' | string;
+  text?: string;
+  /** Tool name when type is tool_use */
+  name?: string;
+  /** Tool invocation id when type is tool_use */
+  id?: string;
+  input?: Record<string, unknown>;
+  tool_use_id?: string;
+};
+
 export type AnthropicAiResponse = {
   [key: string]: unknown; // Allow for additional unknown properties
   id: string;
   model: string;
   created?: number;
   created_at?: number; // Available for Models.retrieve
   messages?: Array<Message>;
-  content?: string; // Available for Messages.create
+  content?: string | Array<ContentBlock>; // Available for Messages.create
   completion?: string; // Available for Completions.create
   input_tokens?: number; // Available for Models.countTokens
   usage?: {
@@ -87,7 +98,14 @@ export type AnthropicAiMessage = {
  * Streaming event type for Anthropic AI
  */
 export type AnthropicAiStreamingEvent = {
-  type: 'message_delta' | 'content_block_start' | 'content_block_delta' | 'content_block_stop' | 'error';
+  type:
+    | 'message_start'
+    | 'message_delta'
+    | 'message_stop'
+    | 'content_block_start'
+    | 'content_block_delta'
+    | 'content_block_stop'
+    | 'error';
   error?: {
     type: string;
     message: string;
@@ -96,9 +114,15 @@ export type AnthropicAiStreamingEvent = {
   delta?: {
     type: unknown;
     text?: string;
+    /** Present for fine-grained tool streaming */
+    partial_json?: string;
+    stop_reason?: string;
+    stop_sequence?: number;
   };
   usage?: {
     output_tokens: number; // Final total output tokens; emitted on the last `message_delta` event
   };
   message?: AnthropicAiMessage;
+  /** Present for fine-grained tool streaming */
+  content_block?: ContentBlock;
 };
