Skip to content

Commit e1c2fcf

Browse files
nicohrubecnicohrubec
committed
move functions
1 parent 440affd commit e1c2fcf

File tree

2 files changed

+118
-107
lines changed

2 files changed

+118
-107
lines changed

‎packages/core/src/tracing/google-genai/index.ts‎

Lines changed: 9 additions & 106 deletions
Original file line numberDiff line numberDiff line change
@@ -1,35 +1,25 @@
11
import { getClient } from '../../currentScopes';
22
import { captureException } from '../../exports';
3-
import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '../../semanticAttributes';
43
import { SPAN_STATUS_ERROR } from '../../tracing';
54
import { startSpan, startSpanManual } from '../../tracing/trace';
6-
import type { Span, SpanAttributeValue } from '../../types-hoist/span';
5+
import type { Span } from '../../types-hoist/span';
76
import { handleCallbackErrors } from '../../utils/handleCallbackErrors';
87
import {
98
GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE,
109
GEN_AI_INPUT_MESSAGES_ATTRIBUTE,
1110
GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE,
12-
GEN_AI_OPERATION_NAME_ATTRIBUTE,
13-
GEN_AI_REQUEST_AVAILABLE_TOOLS_ATTRIBUTE,
14-
GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE,
15-
GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE,
1611
GEN_AI_REQUEST_MODEL_ATTRIBUTE,
17-
GEN_AI_REQUEST_PRESENCE_PENALTY_ATTRIBUTE,
18-
GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE,
19-
GEN_AI_REQUEST_TOP_K_ATTRIBUTE,
20-
GEN_AI_REQUEST_TOP_P_ATTRIBUTE,
2112
GEN_AI_RESPONSE_MODEL_ATTRIBUTE,
2213
GEN_AI_RESPONSE_TEXT_ATTRIBUTE,
2314
GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE,
24-
GEN_AI_SYSTEM_ATTRIBUTE,
2515
GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE,
2616
GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE,
2717
GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE,
2818
GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE,
2919
} from '../ai/gen-ai-attributes';
3020
import { truncateGenAiMessages } from '../ai/messageTruncation';
3121
import { buildMethodPath, extractSystemInstructions, getFinalOperationName, getSpanOperation } from '../ai/utils';
32-
import { CHAT_PATH, CHATS_CREATE_METHOD, GOOGLE_GENAI_SYSTEM_NAME } from './constants';
22+
import { CHAT_PATH, CHATS_CREATE_METHOD } from './constants';
3323
import { instrumentStream } from './streaming';
3424
import type {
3525
Candidate,
@@ -39,100 +29,13 @@ import type {
3929
GoogleGenAIResponse,
4030
} from './types';
4131
import type { ContentListUnion, ContentUnion, Message, PartListUnion } from './utils';
42-
import { contentUnionToMessages, isEmbeddingsMethod, isStreamingMethod, shouldInstrument } from './utils';
43-
44-
/**
45-
* Extract model from parameters or chat context object
46-
* For chat instances, the model is available on the chat object as 'model' (older versions) or 'modelVersion' (newer versions)
47-
*/
48-
export function extractModel(params: Record<string, unknown>, context?: unknown): string {
49-
if ('model' in params && typeof params.model === 'string') {
50-
return params.model;
51-
}
52-
53-
// Try to get model from chat context object (chat instance has model property)
54-
if (context && typeof context === 'object') {
55-
const contextObj = context as Record<string, unknown>;
56-
57-
// Check for 'model' property (older versions, and streaming)
58-
if ('model' in contextObj && typeof contextObj.model === 'string') {
59-
return contextObj.model;
60-
}
61-
62-
// Check for 'modelVersion' property (newer versions)
63-
if ('modelVersion' in contextObj && typeof contextObj.modelVersion === 'string') {
64-
return contextObj.modelVersion;
65-
}
66-
}
67-
68-
return 'unknown';
69-
}
70-
71-
/**
72-
* Extract generation config parameters
73-
*/
74-
function extractConfigAttributes(config: Record<string, unknown>): Record<string, SpanAttributeValue> {
75-
const attributes: Record<string, SpanAttributeValue> = {};
76-
77-
if ('temperature' in config && typeof config.temperature === 'number') {
78-
attributes[GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE] = config.temperature;
79-
}
80-
if ('topP' in config && typeof config.topP === 'number') {
81-
attributes[GEN_AI_REQUEST_TOP_P_ATTRIBUTE] = config.topP;
82-
}
83-
if ('topK' in config && typeof config.topK === 'number') {
84-
attributes[GEN_AI_REQUEST_TOP_K_ATTRIBUTE] = config.topK;
85-
}
86-
if ('maxOutputTokens' in config && typeof config.maxOutputTokens === 'number') {
87-
attributes[GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE] = config.maxOutputTokens;
88-
}
89-
if ('frequencyPenalty' in config && typeof config.frequencyPenalty === 'number') {
90-
attributes[GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE] = config.frequencyPenalty;
91-
}
92-
if ('presencePenalty' in config && typeof config.presencePenalty === 'number') {
93-
attributes[GEN_AI_REQUEST_PRESENCE_PENALTY_ATTRIBUTE] = config.presencePenalty;
94-
}
95-
96-
return attributes;
97-
}
98-
99-
/**
100-
* Extract request attributes from method arguments
101-
* Builds the base attributes for span creation including system info, model, and config
102-
*/
103-
function extractRequestAttributes(
104-
methodPath: string,
105-
params?: Record<string, unknown>,
106-
context?: unknown,
107-
): Record<string, SpanAttributeValue> {
108-
const attributes: Record<string, SpanAttributeValue> = {
109-
[GEN_AI_SYSTEM_ATTRIBUTE]: GOOGLE_GENAI_SYSTEM_NAME,
110-
[GEN_AI_OPERATION_NAME_ATTRIBUTE]: getFinalOperationName(methodPath),
111-
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.ai.google_genai',
112-
};
113-
114-
if (params) {
115-
attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] = extractModel(params, context);
116-
117-
// Extract generation config parameters
118-
if ('config' in params && typeof params.config === 'object' && params.config) {
119-
const config = params.config as Record<string, unknown>;
120-
Object.assign(attributes, extractConfigAttributes(config));
121-
122-
// Extract available tools from config
123-
if ('tools' in config && Array.isArray(config.tools)) {
124-
const functionDeclarations = config.tools.flatMap(
125-
(tool: { functionDeclarations: unknown[] }) => tool.functionDeclarations,
126-
);
127-
attributes[GEN_AI_REQUEST_AVAILABLE_TOOLS_ATTRIBUTE] = JSON.stringify(functionDeclarations);
128-
}
129-
}
130-
} else {
131-
attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] = extractModel({}, context);
132-
}
133-
134-
return attributes;
135-
}
32+
import {
33+
contentUnionToMessages,
34+
extractRequestAttributes,
35+
isEmbeddingsMethod,
36+
isStreamingMethod,
37+
shouldInstrument,
38+
} from './utils';
13639

13740
/**
13841
* Add private request attributes to spans.

‎packages/core/src/tracing/google-genai/utils.ts‎

Lines changed: 109 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,19 @@
1-
import { GOOGLE_GENAI_INSTRUMENTED_METHODS } from './constants';
1+
import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '../../semanticAttributes';
2+
import type { SpanAttributeValue } from '../../types-hoist/span';
3+
import {
4+
GEN_AI_OPERATION_NAME_ATTRIBUTE,
5+
GEN_AI_REQUEST_AVAILABLE_TOOLS_ATTRIBUTE,
6+
GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE,
7+
GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE,
8+
GEN_AI_REQUEST_MODEL_ATTRIBUTE,
9+
GEN_AI_REQUEST_PRESENCE_PENALTY_ATTRIBUTE,
10+
GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE,
11+
GEN_AI_REQUEST_TOP_K_ATTRIBUTE,
12+
GEN_AI_REQUEST_TOP_P_ATTRIBUTE,
13+
GEN_AI_SYSTEM_ATTRIBUTE,
14+
} from '../ai/gen-ai-attributes';
15+
import { getFinalOperationName } from '../ai/utils';
16+
import { GOOGLE_GENAI_INSTRUMENTED_METHODS, GOOGLE_GENAI_SYSTEM_NAME } from './constants';
217
import type { GoogleGenAIIstrumentedMethod } from './types';
318

419
/**
@@ -29,6 +44,99 @@ export function isEmbeddingsMethod(methodPath: string): boolean {
2944
return methodPath.includes('embedContent');
3045
}
3146

47+
/**
48+
* Extract model from parameters or chat context object
49+
* For chat instances, the model is available on the chat object as 'model' (older versions) or 'modelVersion' (newer versions)
50+
*/
51+
export function extractModel(params: Record<string, unknown>, context?: unknown): string {
52+
if ('model' in params && typeof params.model === 'string') {
53+
return params.model;
54+
}
55+
56+
// Try to get model from chat context object (chat instance has model property)
57+
if (context && typeof context === 'object') {
58+
const contextObj = context as Record<string, unknown>;
59+
60+
// Check for 'model' property (older versions, and streaming)
61+
if ('model' in contextObj && typeof contextObj.model === 'string') {
62+
return contextObj.model;
63+
}
64+
65+
// Check for 'modelVersion' property (newer versions)
66+
if ('modelVersion' in contextObj && typeof contextObj.modelVersion === 'string') {
67+
return contextObj.modelVersion;
68+
}
69+
}
70+
71+
return 'unknown';
72+
}
73+
74+
/**
75+
* Extract generation config parameters
76+
*/
77+
function extractConfigAttributes(config: Record<string, unknown>): Record<string, SpanAttributeValue> {
78+
const attributes: Record<string, SpanAttributeValue> = {};
79+
80+
if ('temperature' in config && typeof config.temperature === 'number') {
81+
attributes[GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE] = config.temperature;
82+
}
83+
if ('topP' in config && typeof config.topP === 'number') {
84+
attributes[GEN_AI_REQUEST_TOP_P_ATTRIBUTE] = config.topP;
85+
}
86+
if ('topK' in config && typeof config.topK === 'number') {
87+
attributes[GEN_AI_REQUEST_TOP_K_ATTRIBUTE] = config.topK;
88+
}
89+
if ('maxOutputTokens' in config && typeof config.maxOutputTokens === 'number') {
90+
attributes[GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE] = config.maxOutputTokens;
91+
}
92+
if ('frequencyPenalty' in config && typeof config.frequencyPenalty === 'number') {
93+
attributes[GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE] = config.frequencyPenalty;
94+
}
95+
if ('presencePenalty' in config && typeof config.presencePenalty === 'number') {
96+
attributes[GEN_AI_REQUEST_PRESENCE_PENALTY_ATTRIBUTE] = config.presencePenalty;
97+
}
98+
99+
return attributes;
100+
}
101+
102+
/**
103+
* Extract request attributes from method arguments
104+
* Builds the base attributes for span creation including system info, model, and config
105+
*/
106+
export function extractRequestAttributes(
107+
methodPath: string,
108+
params?: Record<string, unknown>,
109+
context?: unknown,
110+
): Record<string, SpanAttributeValue> {
111+
const attributes: Record<string, SpanAttributeValue> = {
112+
[GEN_AI_SYSTEM_ATTRIBUTE]: GOOGLE_GENAI_SYSTEM_NAME,
113+
[GEN_AI_OPERATION_NAME_ATTRIBUTE]: getFinalOperationName(methodPath),
114+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.ai.google_genai',
115+
};
116+
117+
if (params) {
118+
attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] = extractModel(params, context);
119+
120+
// Extract generation config parameters
121+
if ('config' in params && typeof params.config === 'object' && params.config) {
122+
const config = params.config as Record<string, unknown>;
123+
Object.assign(attributes, extractConfigAttributes(config));
124+
125+
// Extract available tools from config
126+
if ('tools' in config && Array.isArray(config.tools)) {
127+
const functionDeclarations = config.tools.flatMap(
128+
(tool: { functionDeclarations: unknown[] }) => tool.functionDeclarations,
129+
);
130+
attributes[GEN_AI_REQUEST_AVAILABLE_TOOLS_ATTRIBUTE] = JSON.stringify(functionDeclarations);
131+
}
132+
}
133+
} else {
134+
attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] = extractModel({}, context);
135+
}
136+
137+
return attributes;
138+
}
139+
32140
// Copied from https://googleapis.github.io/js-genai/release_docs/index.html
33141
export type ContentListUnion = Content | Content[] | PartListUnion;
34142
export type ContentUnion = Content | PartUnion[] | PartUnion;

0 commit comments

Comments
 (0)