this is much cleaner, get model from context

Author: RulaKhaled

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

@@ -41,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/utils/google-genai/constants.ts

@@ -9,5 +9,3 @@ export const GOOGLE_GENAI_SYSTEM_NAME = 'google_genai';
 export const CHATS_CREATE_METHOD = 'chats.create';
 export const CHAT_PATH = 'chat';
 
-// Span attribute key for storing chat context
-export const GOOGLE_GENAI_SCOPE_MODEL_KEY = 'google_genai.chat.model';

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

@@ -3,7 +3,6 @@ import { captureException } from '../../exports';
 import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '../../semanticAttributes';
 import { startSpan } from '../../tracing/trace';
 import type { Span, SpanAttributeValue } from '../../types-hoist/span';
-import { getActiveSpan, getRootSpan, spanToJSON } from '../../utils/spanUtils';
 import {
   GEN_AI_OPERATION_NAME_ATTRIBUTE,
   GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE,
@@ -22,13 +21,7 @@ import {
 } 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_SCOPE_MODEL_KEY,
-  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,
@@ -40,35 +33,26 @@ import type {
 import { shouldInstrument } from './utils';
 
 /**
- * Store model information on the root span for later retrieval by chat.sendMessage
- */
-function storeModelOnRootSpan(model: string): void {
-  const activeSpan = getActiveSpan();
-  if (activeSpan) {
-    const rootSpan = getRootSpan(activeSpan);
-    rootSpan.setAttributes({
-      [GOOGLE_GENAI_SCOPE_MODEL_KEY]: model,
-    });
-  }
-}
-
-/**
- * Extract model from parameters or root span attributes
- * For chat instances, the model is stored on the root span during chat creation
+ * 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 {
+export function extractModel(params: Record<string, unknown>, context?: unknown): string {
   if ('model' in params && typeof params.model === 'string') {
     return params.model;
   }
 
-  // Try to get model from root span attributes
-  const activeSpan = getActiveSpan();
-  if (activeSpan) {
-    const rootSpan = getRootSpan(activeSpan);
-    const spanData = spanToJSON(rootSpan);
-    const modelFromSpan = spanData.data?.[GOOGLE_GENAI_SCOPE_MODEL_KEY];
-    if (typeof modelFromSpan === 'string') {
-      return modelFromSpan;
+  // Try to get model from chat context object (chat instance has model property)
+  if (context && typeof context === 'object') {
+    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;
     }
   }
 
@@ -255,11 +239,6 @@ function instrumentMethod<T extends unknown[], R>(
             }
             const result = (originalMethod as (...args: T) => R).apply(context, args);
 
-            if (typeof model === 'string' && model !== 'unknown') {
-              // Store model information on root span for later retrieval by chat.sendMessage
-              storeModelOnRootSpan(model);
-            }
-
             // No response attributes for create (returns object of chat instance, not generated content)
             return result;
           } catch (error) {