8000 fix(vertex): support v1.2.0 by nirga · Pull Request #290 · traceloop/openllmetry-js · GitHub
[go: up one dir, main page]
Skip to content

fix(vertex): support v1.2.0 #290

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
May 31, 2024
Merged

fix(vertex): support v1.2.0 #290

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions package-lock.json
View file
Open in desktop

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
Open in desktop
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
Open in desktop
8000
Original file line numberDiff 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
Open in desktop
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
Open in desktop
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
Loading
0