some refacotor

Author: RulaKhaled

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

@@ -22,9 +22,7 @@ class MockGoogleGenAI {
               content: {
                 parts: [
                   {
-                    text: params.contents
-                      ? 'The capital of France is Paris.'
-                      : 'Mock response from Google GenAI!',
+                    text: params.contents ? 'The capital of France is Paris.' : 'Mock response from Google GenAI!',
                   },
                 ],
                 role: 'model',
@@ -43,9 +41,10 @@ class MockGoogleGenAI {
     };
 
     this.chats = {
-      create: () => {
-        // Return a chat instance with sendMessage method
+      create: (options) => {
+        // Return a chat instance with sendMessage method and model info
         return {
+          model: options?.model || 'unknown', // Include model from create options
           sendMessage: async () => {
             // Simulate processing time
             await new Promise(resolve => setTimeout(resolve, 10));

packages/core/src/index.ts

@@ -133,7 +133,12 @@ export { ANTHROPIC_AI_INTEGRATION_NAME } from './utils/anthropic-ai/constants';
 export { instrumentGoogleGenAIClient } from './utils/google-genai';
 export { GOOGLE_GENAI_INTEGRATION_NAME } from './utils/google-genai/constants';
 export type { OpenAiClient, OpenAiOptions, InstrumentedMethod } from './utils/openai/types';
-export type { AnthropicAiClient, AnthropicAiOptions, AnthropicAiInstrumentedMethod } from './utils/anthropic-ai/types';
+export type {
+  AnthropicAiClient,
+  AnthropicAiOptions,
+  AnthropicAiInstrumentedMethod,
+  AnthropicAiResponse,
+} from './utils/anthropic-ai/types';
 export type {
   GoogleGenAIClient,
   GoogleGenAIChat,

packages/core/src/utils/ai/gen-ai-attributes.ts

@@ -178,15 +178,3 @@ export const OPENAI_OPERATIONS = {
  * The response timestamp from Anthropic AI (ISO string)
  */
 export const ANTHROPIC_AI_RESPONSE_TIMESTAMP_ATTRIBUTE = 'anthropic.response.timestamp';
-
-// =============================================================================
-// GOOGLE GENAI OPERATIONS
-// =============================================================================
-
-/**
- * Google GenAI API operations
- */
-export const GOOGLE_GENAI_OPERATIONS = {
-  GENERATE_CONTENT: 'generateContent',
-  STREAM_GENERATE_CONTENT: 'streamGenerateContent',
-} as const;

packages/core/src/utils/google-genai/constants.ts

@@ -5,7 +5,7 @@ export const GOOGLE_GENAI_INTEGRATION_NAME = 'Google_GenAI';
 export const GOOGLE_GENAI_INSTRUMENTED_METHODS = ['models.generateContent', 'chats.create', 'sendMessage'] as const;
 
 // Constants for internal use
-export const GOOGLE_GENAI_MODEL_PROPERTY = '_sentryGoogleGenAIModel';
 export const GOOGLE_GENAI_SYSTEM_NAME = 'google_genai';
 export const CHATS_CREATE_METHOD = 'chats.create';
 export const CHAT_PATH = 'chat';
+

packages/core/src/utils/google-genai/index.ts

@@ -20,14 +20,7 @@ import {
   GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE,
 } from '../ai/gen-ai-attributes';
 import { buildMethodPath, getFinalOperationName, getSpanOperation } from '../ai/utils';
-import { isThenable } from '../is';
-import {
-  CHAT_PATH,
-  CHATS_CREATE_METHOD,
-  GOOGLE_GENAI_INTEGRATION_NAME,
-  GOOGLE_GENAI_MODEL_PROPERTY,
-  GOOGLE_GENAI_SYSTEM_NAME,
-} from './constants';
+import { CHAT_PATH, CHATS_CREATE_METHOD, GOOGLE_GENAI_INTEGRATION_NAME, GOOGLE_GENAI_SYSTEM_NAME } from './constants';
 import type {
   Candidate,
   ContentPart,
@@ -39,21 +32,26 @@ import type {
 import { shouldInstrument } from './utils';
 
 /**
- * Extract model from parameters or context
- * For chat instances, the model is stored during chat creation and retrieved from context
+ * Extract model from parameters or chat context object
+ * For chat instances, the model is available on the chat object as 'model' (older versions) or 'modelVersion' (newer versions)
  */
 export function extractModel(params: Record<string, unknown>, context?: unknown): string {
   if ('model' in params && typeof params.model === 'string') {
     return params.model;
   }
 
-  // For chat instances, try to get the model from the chat context
-  // This is because the model is set during chat creation
-  // and not passed as a parameter to the chat.sendMessage method
+  // Try to get model from chat context object (chat instance has model property)
   if (context && typeof context === 'object') {
-    const chatObj = context as Record<string, unknown>;
-    if (chatObj[GOOGLE_GENAI_MODEL_PROPERTY] && typeof chatObj[GOOGLE_GENAI_MODEL_PROPERTY] === 'string') {
-      return chatObj[GOOGLE_GENAI_MODEL_PROPERTY] as string;
+    const contextObj = context as Record<string, unknown>;
+
+    // Check for 'model' property (older versions, and streaming)
+    if ('model' in contextObj && typeof contextObj.model === 'string') {
+      return contextObj.model;
+    }
+
+    // Check for 'modelVersion' property (newer versions)
+    if ('modelVersion' in contextObj && typeof contextObj.modelVersion === 'string') {
+      return contextObj.modelVersion;
     }
   }
 
@@ -217,7 +215,7 @@ function instrumentMethod<T extends unknown[], R>(
   context: unknown,
   options?: GoogleGenAIOptions,
 ): (...args: T) => R | Promise<R> {
-  const isSyncCreate = !isThenable(originalMethod) && methodPath === CHATS_CREATE_METHOD;
+  const isSyncCreate = methodPath === CHATS_CREATE_METHOD 
 
   const run = (...args: T): R | Promise<R> => {
     const finalOptions = options || getRecordingOptionsFromIntegration();
@@ -238,13 +236,7 @@ function instrumentMethod<T extends unknown[], R>(
             if (finalOptions.recordInputs && args[0] && typeof args[0] === 'object') {
               addPrivateRequestAttributes(span, args[0] as Record<string, unknown>);
             }
-            const result = (originalMethod as (...args: T) => R).apply(context, args) as R;
-
-            if (typeof model === 'string' && model !== 'unknown' && typeof result === 'object') {
-              // We store the model in the result object so that it can be accessed later
-              // This is because the model is not passed as a parameter to the chat.sendMessage method
-              (result as Record<string, unknown>)[GOOGLE_GENAI_MODEL_PROPERTY] = model;
-            }
+            const result = (originalMethod as (...args: T) => R).apply(context, args);
 
             // No response attributes for create (returns object of chat instance, not generated content)
             return result;
@@ -255,7 +247,7 @@ function instrumentMethod<T extends unknown[], R>(
             throw error;
           }
         },
-      ) as R;
+      );
     }
 
     // Async/content-producing path
@@ -273,15 +265,15 @@ function instrumentMethod<T extends unknown[], R>(
 
           const result = await Promise.resolve((originalMethod as (...args: T) => Promise<R>).apply(context, args));
           addResponseAttributes(span, result as GoogleGenAIResponse, finalOptions.recordOutputs);
-          return result as R;
+          return result;
         } catch (error) {
           captureException(error, {
             mechanism: { handled: false, type: 'auto.ai.google_genai', data: { function: methodPath } },
           });
           throw error;
         }
       },
-    ) as Promise<R>;
+    );
   };
 
   return run;