cleaner way of handling missing model

Author: RulaKhaled

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

@@ -5,7 +5,9 @@ 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';
+
+// 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,6 +3,7 @@ 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,
@@ -25,7 +26,7 @@ import {
   CHAT_PATH,
   CHATS_CREATE_METHOD,
   GOOGLE_GENAI_INTEGRATION_NAME,
-  GOOGLE_GENAI_MODEL_PROPERTY,
+  GOOGLE_GENAI_SCOPE_MODEL_KEY,
   GOOGLE_GENAI_SYSTEM_NAME,
 } from './constants';
 import type {
@@ -39,21 +40,35 @@ 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
+ * Store model information on the root span for later retrieval by chat.sendMessage
  */
-export function extractModel(params: Record<string, unknown>, context?: unknown): string {
+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
+ */
+export function extractModel(params: Record<string, unknown>, _context?: unknown): string {
   if ('model' in params && typeof params.model === 'string') {
     return params.model;
   }
 
-  // We retrieve the model from the chat context
-  // because the model is defined when the chat is created,
-  // not provided as a parameter in chat.sendMessage
-  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];
+  // 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;
     }
   }
 
@@ -240,11 +255,9 @@ function instrumentMethod<T extends unknown[], R>(
             }
             const result = (originalMethod as (...args: T) => R).apply(context, args);
 
-            if (typeof model === 'string' && model !== 'unknown' && typeof result === 'object') {
-              // For chat instances, retrieve the model from the chat context.
-              // The model is defined when the chat is created,
-              // not provided as a parameter in chat.sendMessage.
-              (result as Record<string, unknown>)[GOOGLE_GENAI_MODEL_PROPERTY] = model;
+            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)