Skip to content

Commit e62c9b4

Browse files
nirganirga
authored
fix(vertex): support v1.2.0 (#290)
1 parent f30c012 commit e62c9b4

File tree

7 files changed

+70
-121
lines changed

7 files changed

+70
-121
lines changed

‎package-lock.json‎

Lines changed: 6 additions & 6 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/instrumentation-vertexai/package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@
4343
},
4444
"devDependencies": {
4545
"@google-cloud/aiplatform": "^3.10.0",
46-
"@google-cloud/vertexai": "^0.2.1"
46+
"@google-cloud/vertexai": "^1.2.0"
4747
},
4848
"homepage": "https://github.com/traceloop/openllmetry-js/tree/main/packages/instrumentation-openai",
4949
"gitHead": "ef1e70d6037f7b5c061056ef2be16e3f55f02ed5"

‎packages/instrumentation-vertexai/src/vertexai-instrumentation.ts‎

Lines changed: 44 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -49,47 +49,41 @@ export class VertexAIInstrumentation extends InstrumentationBase<any> {
4949
protected init(): InstrumentationModuleDefinition<any> {
5050
const vertexAIModule = new InstrumentationNodeModuleDefinition<any>(
5151
"@google-cloud/vertexai",
52-
[">=0.2.1"],
52+
[">=1.1.0"],
5353
this.wrap.bind(this),
5454
this.unwrap.bind(this),
5555
);
5656

5757
return vertexAIModule;
5858
}
5959

60-
private modelConfig: vertexAI.ModelParams = { model: "" };
61-
62-
private setModel(newValue: vertexAI.ModelParams) {
63-
this.modelConfig = { ...newValue };
64-
}
65-
6660
public manuallyInstrument(module: typeof vertexAI) {
6761
this._diag.debug("Manually instrumenting @google-cloud/vertexai");
6862

6963
this._wrap(
70-
module.VertexAI_Preview.prototype,
71-
"getGenerativeModel",
72-
this.wrapperMethod("getGenerativeModel"),
64+
module.GenerativeModel.prototype,
65+
"generateContentStream",
66+
this.wrapperMethod(),
7367
);
7468
this._wrap(
7569
module.GenerativeModel.prototype,
76-
"generateContentStream",
77-
this.wrapperMethod("generateContentStream"),
70+
"generateContent",
71+
this.wrapperMethod(),
7872
);
7973
}
8074

8175
private wrap(module: typeof vertexAI, moduleVersion?: string) {
8276
this._diag.debug(`Patching @google-cloud/vertexai@${moduleVersion}`);
8377

8478
this._wrap(
85-
module.VertexAI_Preview.prototype,
86-
"getGenerativeModel",
87-
this.wrapperMethod("getGenerativeModel"),
79+
module.GenerativeModel.prototype,
80+
"generateContentStream",
81+
this.wrapperMethod(),
8882
);
8983
this._wrap(
9084
module.GenerativeModel.prototype,
91-
"generateContentStream",
92-
this.wrapperMethod("generateContentStream"),
85+
"generateContent",
86+
this.wrapperMethod(),
9387
);
9488

9589
return module;
@@ -98,42 +92,21 @@ export class VertexAIInstrumentation extends InstrumentationBase<any> {
9892
private unwrap(module: typeof vertexAI, moduleVersion?: string): void {
9993
this._diag.debug(`Unpatching @google-cloud/vertexai@${moduleVersion}`);
10094

101-
this._unwrap(module.VertexAI_Preview.prototype, "getGenerativeModel");
10295
this._unwrap(module.GenerativeModel.prototype, "generateContentStream");
96+
this._unwrap(module.GenerativeModel.prototype, "generateContent");
10397
}
10498

105-
private wrapperMethod(
106-
wrappedMethodName: "getGenerativeModel" | "generateContentStream",
107-
) {
99+
private wrapperMethod() {
108100
// eslint-disable-next-line @typescript-eslint/no-this-alias
109101
const plugin = this;
110102
// eslint-disable-next-line @typescript-eslint/ban-types
111103
return (original: Function) => {
112104
return function method(
113-
this: any,
105+
this: vertexAI.GenerativeModel,
114106
...args: (vertexAI.GenerateContentRequest & vertexAI.ModelParams)[]
115107
) {
116-
if (wrappedMethodName === "getGenerativeModel") {
117-
plugin.setModel(args[0]);
118-
119-
return context.bind(
120-
context.active(),
121-
safeExecuteInTheMiddle(
122-
() => {
123-
return context.with(context.active(), () => {
124-
return original.apply(this, args);
125-
});
126-
},
127-
(e) => {
128-
if (e) {
129-
plugin._diag.error("Error in VertexAI Instrumentation", e);
130-
}
131-
},
132-
),
133-
);
134-
}
135-
136108
const span = plugin._startSpan({
109+
instance: this,
137110
params: args[0],
138111
});
139112

@@ -157,8 +130,10 @@ export class VertexAIInstrumentation extends InstrumentationBase<any> {
157130
}
158131

159132
private _startSpan({
133+
instance,
160134
params,
161135
}: {
136+
instance: vertexAI.GenerativeModel;
162137
params: vertexAI.GenerateContentRequest;
163138
}): Span {
164139
const attributes: Attributes = {
@@ -167,28 +142,18 @@ export class VertexAIInstrumentation extends InstrumentationBase<any> {
167142
};
168143

169144
try {
170-
attributes[SpanAttributes.LLM_REQUEST_MODEL] = this.modelConfig.model;
171-
172-
if (
173-
this.modelConfig.generation_config !== undefined &&
174-
typeof this.modelConfig.generation_config === "object"
175-
) {
176-
if (this.modelConfig.generation_config.max_output_tokens) {
177-
attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] =
178-
this.modelConfig.generation_config.max_output_tokens;
179-
}
180-
if (this.modelConfig.generation_config.temperature) {
181-
attributes[SpanAttributes.LLM_REQUEST_TEMPERATURE] =
182-
this.modelConfig.generation_config.temperature;
183-
}
184-
if (this.modelConfig.generation_config.top_p) {
185-
attributes[SpanAttributes.LLM_REQUEST_TOP_P] =
186-
this.modelConfig.generation_config.top_p;
187-
}
188-
if (this.modelConfig.generation_config.top_k) {
189-
attributes[SpanAttributes.LLM_TOP_K] =
190-
this.modelConfig.generation_config.top_k;
191-
}
145+
attributes[SpanAttributes.LLM_REQUEST_MODEL] = instance["model"];
146+
attributes[SpanAttributes.LLM_RESPONSE_MODEL] = instance["model"];
147+
148+
if (instance["generationConfig"]) {
149+
attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] =
150+
instance["generationConfig"].max_output_tokens;
151+
attributes[SpanAttributes.LLM_REQUEST_TEMPERATURE] =
152+
instance["generationConfig"].temperature;
153+
attributes[SpanAttributes.LLM_REQUEST_TOP_P] =
154+
instance["generationConfig"].top_p;
155+
attributes[SpanAttributes.LLM_TOP_K] =
156+
instance["generationConfig"].top_k;
192157
}
193158

194159
if (this._shouldSendPrompts() && "contents" in params) {
@@ -213,7 +178,9 @@ export class VertexAIInstrumentation extends InstrumentationBase<any> {
213178
.then(async (result) => {
214179
await this._endSpan({
215180
span,
216-
result: result as vertexAI.StreamGenerateContentResult,
181+
result: result as
182+
| vertexAI.StreamGenerateContentResult
183+
| vertexAI.GenerateContentResult,
217184
});
218185
return new Promise<T>((resolve) => resolve(result));
219186
})
@@ -236,14 +203,11 @@ export class VertexAIInstrumentation extends InstrumentationBase<any> {
236203
result,
237204
}: {
238205
span: Span;
239-
result: vertexAI.StreamGenerateContentResult;
206+
result:
207+
| vertexAI.StreamGenerateContentResult
208+
| vertexAI.GenerateContentResult;
240209
}) {
241210
try {
242-
span.setAttribute(
243-
SpanAttributes.LLM_RESPONSE_MODEL,
244-
this.modelConfig.model,
245-
);
246-
247211
const streamResponse = await result.response;
248212

249213
if (streamResponse.usageMetadata?.totalTokenCount !== undefined)
@@ -252,20 +216,20 @@ export class VertexAIInstrumentation extends InstrumentationBase<any> {
252216
streamResponse.usageMetadata.totalTokenCount,
253217
);
254218

255-
if (streamResponse.usageMetadata?.candidates_token_count)
219+
if (streamResponse.usageMetadata?.candidatesTokenCount)
256220
span.setAttribute(
257221
SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
258-
streamResponse.usageMetadata.candidates_token_count,
222+
streamResponse.usageMetadata.candidatesTokenCount,
259223
);
260224

261-
if (streamResponse.usageMetadata?.prompt_token_count)
225+
if (streamResponse.usageMetadata?.promptTokenCount)
262226
span.setAttribute(
263227
SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
264-
streamResponse.usageMetadata.prompt_token_count,
228+
streamResponse.usageMetadata.promptTokenCount,
265229
);
266230

267231
if (this._shouldSendPrompts()) {
268-
streamResponse.candidates.forEach((candidate, index) => {
232+
streamResponse.candidates?.forEach((candidate, index) => {
269233
if (candidate.finishReason)
270234
span.setAttribute(
271235
`${SpanAttributes.LLM_COMPLETIONS}.${index}.finish_reason`,
@@ -298,10 +262,10 @@ export class VertexAIInstrumentation extends InstrumentationBase<any> {
298262
const result = parts
299263
.map((part) => {
300264
if (part.text) return part.text;
301-
else if (part.file_data)
302-
return part.file_data.file_uri + "-" + part.file_data.mime_type;
303-
else if (part.inline_data)
304-
return part.inline_data.data + "-" + part.inline_data.mime_type;
265+
else if (part.fileData)
266+
return part.fileData.fileUri + "-" + part.fileData.mimeType;
267+
else if (part.inlineData)
268+
return part.inlineData.data + "-" + part.inlineData.mimeType;
305269
else return "";
306270
})
307271
.filter(Boolean);

‎packages/instrumentation-vertexai/tests/gemini.test.ts‎

Lines changed: 11 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -60,9 +60,9 @@ describe.skip("Test Gemini GenerativeModel Instrumentation", () => {
6060

6161
const generativeModel = vertexAI.preview.getGenerativeModel({
6262
model,
63-
generation_config: {
64-
top_p: 0.9,
65-
max_output_tokens: 256,
63+
generationConfig: {
64+
topP: 0.9,
65+
maxOutputTokens: 256,
6666
},
6767
});
6868
const prompt = "What is Node.js?";
@@ -78,7 +78,7 @@ describe.skip("Test Gemini GenerativeModel Instrumentation", () => {
7878
const aggregatedResponse = await responseStream.response;
7979

8080
const fullTextResponse =
81-
aggregatedResponse.candidates[0].content.parts[0].text;
81+
aggregatedResponse.candidates![0].content.parts[0].text;
8282

8383
const spans = memoryExporter.getFinishedSpans();
8484

@@ -87,10 +87,7 @@ describe.skip("Test Gemini GenerativeModel Instrumentation", () => {
8787
assert.strictEqual(attributes["gen_ai.system"], "VertexAI");
8888
assert.strictEqual(attributes["llm.request.type"], "completion");
8989
assert.strictEqual(attributes["gen_ai.request.model"], model);
90-
assert.strictEqual(
91-
attributes["gen_ai.request.top_p"],
92-
generativeModel.generation_config?.top_p,
93-
);
90+
assert.strictEqual(attributes["gen_ai.request.top_p"], 0.9);
9491
assert.strictEqual(attributes["gen_ai.prompt.0.content"], prompt);
9592
assert.strictEqual(attributes["gen_ai.prompt.0.role"], "user");
9693
assert.strictEqual(attributes["gen_ai.response.model"], model);
@@ -111,9 +108,9 @@ describe.skip("Test Gemini GenerativeModel Instrumentation", () => {
111108

112109
const generativeModel = vertexAI.preview.getGenerativeModel({
113110
model,
114-
generation_config: {
115-
top_p: 0.9,
116-
max_output_tokens: 256,
111+
generationConfig: {
112+
topP: 0.9,
113+
maxOutputTokens: 256,
117114
},
118115
});
119116
const prompt = "What are the 4 cardinal directions?";
@@ -129,7 +126,7 @@ describe.skip("Test Gemini GenerativeModel Instrumentation", () => {
129126

130127
const fullTextResponse = [];
131128
for await (const item of responseStream.stream) {
132-
fullTextResponse.push(item.candidates[0].content.parts[0].text);
129+
fullTextResponse.push(item.candidates![0].content.parts[0].text);
133130
}
134131

135132
assert.ok(fullTextResponse);
@@ -143,14 +140,8 @@ describe.skip("Test Gemini GenerativeModel Instrumentation", () => {
143140
assert.strictEqual(attributes["gen_ai.system"], "VertexAI");
144141
assert.strictEqual(attributes["llm.request.type"], "completion");
145142
assert.strictEqual(attributes["gen_ai.request.model"], model);
146-
assert.strictEqual(
147-
attributes["gen_ai.request.top_p"],
148-
generativeModel.generation_config?.top_p,
149-
);
150-
assert.strictEqual(
151-
attributes["gen_ai.request.max_tokens"],
152-
generativeModel.generation_config?.max_output_tokens,
153-
);
143+
assert.strictEqual(attributes["gen_ai.request.top_p"], 0.9);
144+
assert.strictEqual(attributes["gen_ai.request.max_tokens"], 256);
154145
assert.strictEqual(attributes["gen_ai.prompt.0.content"], prompt);
155146
assert.strictEqual(attributes["gen_ai.prompt.0.role"], "user");
156147
assert.strictEqual(attributes["gen_ai.response.model"], model);

‎packages/sample-app/package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@
3636
"@aws-sdk/client-bedrock-runtime": "^3.499.0",
3737
"@azure/openai": "^1.0.0-beta.11",
3838
"@google-cloud/aiplatform": "^3.10.0",
39-
"@google-cloud/vertexai": "^0.2.1",
39+
"@google-cloud/vertexai": "^1.2.0",
4040
"@langchain/community": "^0.0.34",
4141
"@pinecone-database/pinecone": "^2.0.1",
4242
"@traceloop/node-server-sdk": "*",

0 commit comments

Comments
 (0)