update with tests

Author: RulaKhaled

dev-packages/node-integration-tests/suites/tracing/anthropic/scenario.mjs

@@ -10,6 +10,10 @@ class MockAnthropic {
       create: this._messagesCreate.bind(this),
       countTokens: this._messagesCountTokens.bind(this)
     };
+
+    this.models = {
+      retrieve: this._modelsRetrieve.bind(this),
+    };
   }
 
   /**
@@ -46,22 +50,28 @@ class MockAnthropic {
     };
   }
 
-  async _messagesCountTokens(params) {
+  async _messagesCountTokens() {
     // Simulate processing time
     await new Promise(resolve => setTimeout(resolve, 10));
 
-    if (params.model === 'error-model') {
-      const error = new Error('Model not found');
-      error.status = 404;
-      error.headers = { 'x-request-id': 'mock-request-123' };
-      throw error;
-    }
-
-    // For countTokens, just return input_tokens to match implementation
+    // For countTokens, just return input_tokens
     return {
         input_tokens: 15
       }
   }
+
+  async _modelsRetrieve(modelId) {
+    // Simulate processing time
+    await new Promise(resolve => setTimeout(resolve, 10));
+
+    // Match what the actual implementation would return
+    return {
+      id: modelId,
+      name: modelId,
+      created_at: 1715145600,
+      model: modelId,  // Add model field to match the check in addResponseAttributes
+    };
+  }
 }
 
 async function run() {
@@ -100,6 +110,9 @@ async function run() {
         { role: 'user', content: 'What is the capital of France?' },
       ],
     });
+
+    // Fourth test: models.retrieve
+    await client.models.retrieve('claude-3-haiku-20240307');
   });
 }
 

dev-packages/node-integration-tests/suites/tracing/anthropic/test.ts

@@ -58,6 +58,22 @@ describe('Anthropic integration', () => {
         origin: 'auto.ai.anthropic',
         status: 'ok',
       }),
+      // Fourth span - models.retrieve
+      expect.objectContaining({
+        data: {
+          'gen_ai.operation.name': 'retrieve',
+          'sentry.op': 'gen_ai.retrieve',
+          'sentry.origin': 'auto.ai.anthropic',
+          'gen_ai.system': 'anthropic',
+          'gen_ai.request.model': 'claude-3-haiku-20240307',
+          'gen_ai.response.id': 'claude-3-haiku-20240307',
+          'gen_ai.response.model': 'claude-3-haiku-20240307',
+        },
+        description: 'retrieve claude-3-haiku-20240307',
+        op: 'gen_ai.retrieve',
+        origin: 'auto.ai.anthropic',
+        status: 'ok',
+      }),
     ]),
   };
 
@@ -119,6 +135,22 @@ describe('Anthropic integration', () => {
         origin: 'auto.ai.anthropic',
         status: 'ok',
       }),
+      // Fourth span - models.retrieve with PII
+      expect.objectContaining({
+        data: {
+          'gen_ai.operation.name': 'retrieve',
+          'sentry.op': 'gen_ai.retrieve',
+          'sentry.origin': 'auto.ai.anthropic',
+          'gen_ai.system': 'anthropic',
+          'gen_ai.request.model': 'claude-3-haiku-20240307',
+          'gen_ai.response.id': 'claude-3-haiku-20240307',
+          'gen_ai.response.model': 'claude-3-haiku-20240307',
+        },
+        description: 'retrieve claude-3-haiku-20240307',
+        op: 'gen_ai.retrieve',
+        origin: 'auto.ai.anthropic',
+        status: 'ok',
+      }),
     ]),
   };
 
@@ -141,6 +173,18 @@ describe('Anthropic integration', () => {
         }),
         op: 'gen_ai.messages.countTokens',
       }),
+      // Check models.retrieve with options
+      expect.objectContaining({
+        data: expect.objectContaining({
+          'gen_ai.operation.name': 'retrieve',
+          'gen_ai.system': 'anthropic',
+          'gen_ai.request.model': 'claude-3-haiku-20240307',
+          'gen_ai.response.id': 'claude-3-haiku-20240307',
+          'gen_ai.response.model': 'claude-3-haiku-20240307',
+        }),
+        op: 'gen_ai.retrieve',
+        description: 'retrieve claude-3-haiku-20240307',
+      }),
     ]),
   };
 

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

@@ -5,7 +5,7 @@ export const ANTHROPIC_AI_INTEGRATION_NAME = 'Anthropic_AI';
 export const ANTHROPIC_AI_INSTRUMENTED_METHODS = [
   'messages.create',
   'messages.countTokens',
-  'models.list',
   'models.get',
   'completions.create',
+  'models.retrieve',
 ] as const;

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

@@ -51,7 +51,12 @@ function extractRequestAttributes(args: unknown[], methodPath: string): Record<s
     attributes[GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE] = params.frequency_penalty;
     attributes[GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE] = params.max_tokens;
   } else {
-    attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] = 'unknown';
+    if (methodPath === 'models.retrieve' || methodPath === 'models.get') {
+      // models.retrieve(model-id) and models.get(model-id)
+      attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] = args[0];
+    } else {
+      attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] = 'unknown';
+    }
   }
 
   return attributes;
@@ -109,6 +114,11 @@ function addResponseAttributes(span: Span, response: AnthropicAiResponse, record
     span.setAttributes({
       [ANTHROPIC_AI_RESPONSE_TIMESTAMP_ATTRIBUTE]: new Date(response.created * 1000).toISOString(),
     });
+    if ('created_at' in response && typeof response.created_at === 'number') {
+      span.setAttributes({
+        [ANTHROPIC_AI_RESPONSE_TIMESTAMP_ATTRIBUTE]: new Date(response.created_at * 1000).toISOString(),
+      });
+    }
   }
 
   if (response.usage) {
@@ -174,10 +184,10 @@ function instrumentMethod<T extends unknown[], R>(
           captureException(error, {
             mechanism: {
               handled: false,
-            },
-            type: 'auto.ai.anthropic',
-            data: {
-              function: methodPath,
+              type: 'auto.ai.anthropic',
+              data: {
+                function: methodPath,
+              },
             },
           });
           throw error;

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

@@ -20,7 +20,8 @@ export type AnthropicAiResponse = {
   [key: string]: unknown; // Allow for additional unknown properties
   id: string;
   model: string;
-  created: number;
+  created?: number;
+  created_at?: number; // Available for Models.retrieve
   messages?: Array<Message>;
   content?: string; // Available for Messages.create
   completion?: string; // Available for Completions.create