Skip to content

Commit

Permalink
fix(vertex): support v1.2.0 (#290)
Browse files Browse the repository at this point in the history
  • Loading branch information
@nirga
nirga authored May 31, 2024
1 parent f30c012 commit e62c9b4
Show file tree
Hide file tree
Showing 7 changed files with 70 additions and 121 deletions.
12 changes: 6 additions & 6 deletions package-lock.json
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion packages/instrumentation-vertexai/package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@
},
"devDependencies": {
"@google-cloud/aiplatform": "^3.10.0",
"@google-cloud/vertexai": "^0.2.1"
"@google-cloud/vertexai": "^1.2.0"
},
"homepage": "https://github.com/traceloop/openllmetry-js/tree/main/packages/instrumentation-openai",
"gitHead": "ef1e70d6037f7b5c061056ef2be16e3f55f02ed5"
Expand Down
124 changes: 44 additions & 80 deletions packages/instrumentation-vertexai/src/vertexai-instrumentation.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,47 +49,41 @@ export class VertexAIInstrumentation extends InstrumentationBase<any> {
protected init(): InstrumentationModuleDefinition<any> {
const vertexAIModule = new InstrumentationNodeModuleDefinition<any>(
"@google-cloud/vertexai",
[">=0.2.1"],
[">=1.1.0"],
this.wrap.bind(this),
this.unwrap.bind(this),
);

return vertexAIModule;
}

private modelConfig: vertexAI.ModelParams = { model: "" };

private setModel(newValue: vertexAI.ModelParams) {
this.modelConfig = { ...newValue };
}

public manuallyInstrument(module: typeof vertexAI) {
this._diag.debug("Manually instrumenting @google-cloud/vertexai");

this._wrap(
module.VertexAI_Preview.prototype,
"getGenerativeModel",
this.wrapperMethod("getGenerativeModel"),
module.GenerativeModel.prototype,
"generateContentStream",
this.wrapperMethod(),
);
this._wrap(
module.GenerativeModel.prototype,
"generateContentStream",
this.wrapperMethod("generateContentStream"),
"generateContent",
this.wrapperMethod(),
);
}

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

this._wrap(
module.VertexAI_Preview.prototype,
"getGenerativeModel",
this.wrapperMethod("getGenerativeModel"),
module.GenerativeModel.prototype,
"generateContentStream",
this.wrapperMethod(),
);
this._wrap(
module.GenerativeModel.prototype,
"generateContentStream",
this.wrapperMethod("generateContentStream"),
"generateContent",
this.wrapperMethod(),
);

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

this._unwrap(module.VertexAI_Preview.prototype, "getGenerativeModel");
this._unwrap(module.GenerativeModel.prototype, "generateContentStream");
this._unwrap(module.GenerativeModel.prototype, "generateContent");
}

private wrapperMethod(
wrappedMethodName: "getGenerativeModel" | "generateContentStream",
) {
private wrapperMethod() {
// eslint-disable-next-line @typescript-eslint/no-this-alias
const plugin = this;
// eslint-disable-next-line @typescript-eslint/ban-types
return (original: Function) => {
return function method(
this: any,
this: vertexAI.GenerativeModel,
...args: (vertexAI.GenerateContentRequest & vertexAI.ModelParams)[]
) {
if (wrappedMethodName === "getGenerativeModel") {
plugin.setModel(args[0]);

return context.bind(
context.active(),
safeExecuteInTheMiddle(
() => {
return context.with(context.active(), () => {
return original.apply(this, args);
});
},
(e) => {
if (e) {
plugin._diag.error("Error in VertexAI Instrumentation", e);
}
},
),
);
}

const span = plugin._startSpan({
instance: this,
params: args[0],
});

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

private _startSpan({
instance,
params,
}: {
instance: vertexAI.GenerativeModel;
params: vertexAI.GenerateContentRequest;
}): Span {
const attributes: Attributes = {
Expand All @@ -167,28 +142,18 @@ export class VertexAIInstrumentation extends InstrumentationBase<any> {
};

try {
attributes[SpanAttributes.LLM_REQUEST_MODEL] = this.modelConfig.model;

if (
this.modelConfig.generation_config !== undefined &&
typeof this.modelConfig.generation_config === "object"
) {
if (this.modelConfig.generation_config.max_output_tokens) {
attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] =
this.modelConfig.generation_config.max_output_tokens;
}
if (this.modelConfig.generation_config.temperature) {
attributes[SpanAttributes.LLM_REQUEST_TEMPERATURE] =
this.modelConfig.generation_config.temperature;
}
if (this.modelConfig.generation_config.top_p) {
attributes[SpanAttributes.LLM_REQUEST_TOP_P] =
this.modelConfig.generation_config.top_p;
}
if (this.modelConfig.generation_config.top_k) {
attributes[SpanAttributes.LLM_TOP_K] =
this.modelConfig.generation_config.top_k;
}
attributes[SpanAttributes.LLM_REQUEST_MODEL] = instance["model"];
attributes[SpanAttributes.LLM_RESPONSE_MODEL] = instance["model"];

if (instance["generationConfig"]) {
attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] =
instance["generationConfig"].max_output_tokens;
attributes[SpanAttributes.LLM_REQUEST_TEMPERATURE] =
instance["generationConfig"].temperature;
attributes[SpanAttributes.LLM_REQUEST_TOP_P] =
instance["generationConfig"].top_p;
attributes[SpanAttributes.LLM_TOP_K] =
instance["generationConfig"].top_k;
}

if (this._shouldSendPrompts() && "contents" in params) {
Expand All @@ -213,7 +178,9 @@ export class VertexAIInstrumentation extends InstrumentationBase<any> {
.then(async (result) => {
await this._endSpan({
span,
result: result as vertexAI.StreamGenerateContentResult,
result: result as
| vertexAI.StreamGenerateContentResult
| vertexAI.GenerateContentResult,
});
return new Promise<T>((resolve) => resolve(result));
})
Expand All @@ -236,14 +203,11 @@ export class VertexAIInstrumentation extends InstrumentationBase<any> {
result,
}: {
span: Span;
result: vertexAI.StreamGenerateContentResult;
result:
| vertexAI.StreamGenerateContentResult
| vertexAI.GenerateContentResult;
}) {
try {
span.setAttribute(
SpanAttributes.LLM_RESPONSE_MODEL,
this.modelConfig.model,
);

const streamResponse = await result.response;

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

if (streamResponse.usageMetadata?.candidates_token_count)
if (streamResponse.usageMetadata?.candidatesTokenCount)
span.setAttribute(
SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
streamResponse.usageMetadata.candidates_token_count,
streamResponse.usageMetadata.candidatesTokenCount,
);

if (streamResponse.usageMetadata?.prompt_token_count)
if (streamResponse.usageMetadata?.promptTokenCount)
span.setAttribute(
SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
streamResponse.usageMetadata.prompt_token_count,
streamResponse.usageMetadata.promptTokenCount,
);

if (this._shouldSendPrompts()) {
streamResponse.candidates.forEach((candidate, index) => {
streamResponse.candidates?.forEach((candidate, index) => {
if (candidate.finishReason)
span.setAttribute(
`${SpanAttributes.LLM_COMPLETIONS}.${index}.finish_reason`,
Expand Down Expand Up @@ -298,10 +262,10 @@ export class VertexAIInstrumentation extends InstrumentationBase<any> {
const result = parts
.map((part) => {
if (part.text) return part.text;
else if (part.file_data)
return part.file_data.file_uri + "-" + part.file_data.mime_type;
else if (part.inline_data)
return part.inline_data.data + "-" + part.inline_data.mime_type;
else if (part.fileData)
return part.fileData.fileUri + "-" + part.fileData.mimeType;
else if (part.inlineData)
return part.inlineData.data + "-" + part.inlineData.mimeType;
else return "";
})
.filter(Boolean);
Expand Down
31 changes: 11 additions & 20 deletions packages/instrumentation-vertexai/tests/gemini.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,9 +60,9 @@ describe.skip("Test Gemini GenerativeModel Instrumentation", () => {

const generativeModel = vertexAI.preview.getGenerativeModel({
model,
generation_config: {
top_p: 0.9,
max_output_tokens: 256,
generationConfig: {
topP: 0.9,
maxOutputTokens: 256,
},
});
const prompt = "What is Node.js?";
Expand All @@ -78,7 +78,7 @@ describe.skip("Test Gemini GenerativeModel Instrumentation", () => {
const aggregatedResponse = await responseStream.response;

const fullTextResponse =
aggregatedResponse.candidates[0].content.parts[0].text;
aggregatedResponse.candidates![0].content.parts[0].text;

const spans = memoryExporter.getFinishedSpans();

Expand All @@ -87,10 +87,7 @@ describe.skip("Test Gemini GenerativeModel Instrumentation", () => {
assert.strictEqual(attributes["gen_ai.system"], "VertexAI");
assert.strictEqual(attributes["llm.request.type"], "completion");
assert.strictEqual(attributes["gen_ai.request.model"], model);
assert.strictEqual(
attributes["gen_ai.request.top_p"],
generativeModel.generation_config?.top_p,
);
assert.strictEqual(attributes["gen_ai.request.top_p"], 0.9);
assert.strictEqual(attributes["gen_ai.prompt.0.content"], prompt);
assert.strictEqual(attributes["gen_ai.prompt.0.role"], "user");
assert.strictEqual(attributes["gen_ai.response.model"], model);
Expand All @@ -111,9 +108,9 @@ describe.skip("Test Gemini GenerativeModel Instrumentation", () => {

const generativeModel = vertexAI.preview.getGenerativeModel({
model,
generation_config: {
top_p: 0.9,
max_output_tokens: 256,
generationConfig: {
topP: 0.9,
maxOutputTokens: 256,
},
});
const prompt = "What are the 4 cardinal directions?";
Expand All @@ -129,7 +126,7 @@ describe.skip("Test Gemini GenerativeModel Instrumentation", () => {

const fullTextResponse = [];
for await (const item of responseStream.stream) {
fullTextResponse.push(item.candidates[0].content.parts[0].text);
fullTextResponse.push(item.candidates![0].content.parts[0].text);
}

assert.ok(fullTextResponse);
Expand All @@ -143,14 +140,8 @@ describe.skip("Test Gemini GenerativeModel Instrumentation", () => {
assert.strictEqual(attributes["gen_ai.system"], "VertexAI");
assert.strictEqual(attributes["llm.request.type"], "completion");
assert.strictEqual(attributes["gen_ai.request.model"], model);
assert.strictEqual(
attributes["gen_ai.request.top_p"],
generativeModel.generation_config?.top_p,
);
assert.strictEqual(
attributes["gen_ai.request.max_tokens"],
generativeModel.generation_config?.max_output_tokens,
);
assert.strictEqual(attributes["gen_ai.request.top_p"], 0.9);
assert.strictEqual(attributes["gen_ai.request.max_tokens"], 256);
assert.strictEqual(attributes["gen_ai.prompt.0.content"], prompt);
assert.strictEqual(attributes["gen_ai.prompt.0.role"], "user");
assert.strictEqual(attributes["gen_ai.response.model"], model);
Expand Down
2 changes: 1 addition & 1 deletion packages/sample-app/package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@
"@aws-sdk/client-bedrock-runtime": "^3.499.0",
"@azure/openai": "^1.0.0-beta.11",
"@google-cloud/aiplatform": "^3.10.0",
"@google-cloud/vertexai": "^0.2.1",
"@google-cloud/vertexai": "^1.2.0",
"@langchain/community": "^0.0.34",
"@pinecone-database/pinecone": "^2.0.1",
"@traceloop/node-server-sdk": "*",
Expand Down
Loading

0 comments on commit e62c9b4

Please sign in to comment.