fixup! fixup! fix(tracing): add gen_ai.request.messages.original_length attributes (#18608)

Author: isaacs

packages/core/src/tracing/anthropic-ai/index.ts

@@ -76,11 +76,7 @@ function extractRequestAttributes(args: unknown[], methodPath: string): Record<s
  */
 function addPrivateRequestAttributes(span: Span, params: Record<string, unknown>): void {
   const messages = messagesFromParams(params);
-  if (messages.length) {
-    setMessagesAttribute(span, messages);
-  } else if ('input' in params) {
-    setMessagesAttribute(span, params.input);
-  }
+  setMessagesAttribute(span, messages);
 
   if ('prompt' in params) {
     span.setAttributes({ [GEN_AI_PROMPT_ATTRIBUTE]: JSON.stringify(params.prompt) });

packages/core/src/tracing/anthropic-ai/utils.ts

@@ -20,16 +20,11 @@ export function shouldInstrument(methodPath: string): methodPath is AnthropicAiI
  * Set the messages and messages original length attributes.
  */
 export function setMessagesAttribute(span: Span, messages: unknown): void {
-  if (Array.isArray(messages)) {
-    if (!messages.length) return;
+  const length = Array.isArray(messages) ? messages.length : undefined;
+  if (length !== 0) {
     span.setAttributes({
       [GEN_AI_REQUEST_MESSAGES_ATTRIBUTE]: getTruncatedJsonString(messages),
-      [GEN_AI_REQUEST_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE]: messages.length,
-    });
-  } else if (messages) {
-    span.setAttributes({
-      [GEN_AI_REQUEST_MESSAGES_ATTRIBUTE]: getTruncatedJsonString(messages),
-      [GEN_AI_REQUEST_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE]: undefined,
+      [GEN_AI_REQUEST_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE]: length,
     });
   }
 }
@@ -55,11 +50,15 @@ export function handleResponseError(span: Span, response: AnthropicAiResponse):
  * Include the system prompt in the messages list, if available
  */
 export function messagesFromParams(params: Record<string, unknown>): unknown[] {
-  const { system, messages } = params;
+  const { system, messages, input } = params;
 
   const systemMessages = typeof system === 'string' ? [{ role: 'system', content: params.system }] : [];
 
-  const userMessages = Array.isArray(messages) ? messages : messages != null ? [messages] : [];
+  const inputParamMessages = Array.isArray(input) ? input : input != null ? [input] : undefined;
+
+  const messagesParamMessages = Array.isArray(messages) ? messages : messages != null ? [messages] : [];
+
+  const userMessages = inputParamMessages ?? messagesParamMessages;
 
   return [...systemMessages, ...userMessages];
 }

packages/core/test/lib/utils/anthropic-utils.test.ts

@@ -1,5 +1,6 @@
 import { describe, expect, it } from 'vitest';
-import { messagesFromParams, shouldInstrument } from '../../../src/tracing/anthropic-ai/utils';
+import { messagesFromParams, setMessagesAttribute, shouldInstrument } from '../../../src/tracing/anthropic-ai/utils';
+import type { Span } from '../../../src/types-hoist/span';
 
 describe('anthropic-ai-utils', () => {
   describe('shouldInstrument', () => {
@@ -25,6 +26,19 @@ describe('anthropic-ai-utils', () => {
       ]);
     });
 
+    it('looks to params.input ahead of params.messages', () => {
+      expect(
+        messagesFromParams({
+          input: [{ role: 'user', content: 'input' }],
+          messages: [{ role: 'user', content: 'hello' }],
+          system: 'You are a friendly robot awaiting a greeting.',
+        }),
+      ).toStrictEqual([
+        { role: 'system', content: 'You are a friendly robot awaiting a greeting.' },
+        { role: 'user', content: 'input' },
+      ]);
+    });
+
     it('includes system message along with non-array messages', () => {
       expect(
         messagesFromParams({
@@ -53,4 +67,42 @@ describe('anthropic-ai-utils', () => {
       ).toStrictEqual([{ role: 'user', content: 'hello' }]);
     });
   });
+
+  describe('setMessagesAtribute', () => {
+    const mock = {
+      attributes: {} as Record<string, any>,
+      setAttributes(kv: Record<string, any>) {
+        for (const [key, val] of Object.entries(kv)) {
+          // eslint-disable-next-line @typescript-eslint/no-dynamic-delete
+          if (val === undefined) delete this.attributes[key];
+          else this.attributes[key] = val;
+        }
+      },
+    };
+    const span = mock as unknown as Span;
+
+    it('sets length along with truncated value', () => {
+      const content = 'A'.repeat(200_000);
+      setMessagesAttribute(span, [{ role: 'user', content }]);
+      const result = [{ role: 'user', content: 'A'.repeat(19972) }];
+      expect(mock.attributes).toStrictEqual({
+        "gen_ai.request.messages.original_length": 1,
+        'gen_ai.request.messages': JSON.stringify(result),
+      });
+    });
+
+    it('removes length when setting new value ', () => {
+      setMessagesAttribute(span, { content: 'hello, world' });
+      expect(mock.attributes).toStrictEqual({
+        'gen_ai.request.messages': '{"content":"hello, world"}',
+      });
+    });
+
+    it('ignores empty array', () => {
+      setMessagesAttribute(span, []);
+      expect(mock.attributes).toStrictEqual({
+        'gen_ai.request.messages': '{"content":"hello, world"}',
+      });
+    });
+  });
 });