@traceloop/instrumentation-vertexai
Version:
Google's VertexAI Instrumentation
386 lines (381 loc) • 25.5 kB
JavaScript
import { trace, context, SpanKind, SpanStatusCode } from '@opentelemetry/api';
import { InstrumentationBase, InstrumentationNodeModuleDefinition, safeExecuteInTheMiddle } from '@opentelemetry/instrumentation';
import { SpanAttributes, CONTEXT_KEY_ALLOW_TRACE_CONTENT } from '@traceloop/ai-semantic-conventions';
import { __awaiter } from 'tslib';
var version = "0.22.2";
/*
* Copyright Traceloop
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
class AIPlatformInstrumentation extends InstrumentationBase {
constructor(config = {}) {
super("@traceloop/instrumentation-vertexai", version, config);
}
setConfig(config = {}) {
super.setConfig(config);
}
init() {
const aiPlatformModule = new InstrumentationNodeModuleDefinition("@google-cloud/aiplatform", [">=3.10.0"], this.wrap.bind(this), this.unwrap.bind(this));
return aiPlatformModule;
}
manuallyInstrument(module) {
this._diag.debug(`Manually instrumenting @google-cloud/aiplatform`);
this._wrap(module.PredictionServiceClient.prototype, "predict", this.wrapperMethod());
}
wrap(module, moduleVersion) {
this._diag.debug(`Patching @google-cloud/aiplatform@${moduleVersion}`);
this._wrap(module.PredictionServiceClient.prototype, "predict", this.wrapperMethod());
return module;
}
unwrap(module, moduleVersion) {
this._diag.debug(`Unpatching @google-cloud/aiplatform@${moduleVersion}`);
this._unwrap(module.PredictionServiceClient.prototype, "predict");
}
wrapperMethod() {
// eslint-disable-next-line @typescript-eslint/no-this-alias
const plugin = this;
// eslint-disable-next-line
return (original) => {
return function method(...args) {
const span = plugin._startSpan({
params: args[0],
});
const execContext = trace.setSpan(context.active(), span);
const execPromise = safeExecuteInTheMiddle(() => {
return context.with(execContext, () => {
return original.apply(this, args);
});
}, (e) => {
if (e) {
plugin._diag.error("Error in VertexAIPlatform instrumentation", e);
}
});
const wrappedPromise = plugin._wrapPromise(span, execPromise);
return context.bind(execContext, wrappedPromise);
};
};
}
_startSpan({ params, }) {
var _a, _b, _c, _d, _e, _f, _g, _h, _j, _k, _l, _m, _o, _p, _q, _r, _s, _t, _u, _v, _w, _x, _y, _z, _0, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15;
const attributes = {
[SpanAttributes.LLM_SYSTEM]: "Google",
[SpanAttributes.LLM_REQUEST_TYPE]: "completion",
};
try {
if (params !== undefined) {
if (params.endpoint) {
const model = params.endpoint.split("/").pop();
attributes[SpanAttributes.LLM_REQUEST_MODEL] = model;
attributes[SpanAttributes.LLM_RESPONSE_MODEL] = model;
}
if (params === null || params === void 0 ? void 0 : params.parameters) {
if ((_b = (_a = params === null || params === void 0 ? void 0 : params.parameters.structValue) === null || _a === void 0 ? void 0 : _a.fields) === null || _b === void 0 ? void 0 : _b.maxOutputTokens.numberValue) {
attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] =
(_d = (_c = params === null || params === void 0 ? void 0 : params.parameters.structValue) === null || _c === void 0 ? void 0 : _c.fields) === null || _d === void 0 ? void 0 : _d.maxOutputTokens.numberValue;
}
if ((_f = (_e = params === null || params === void 0 ? void 0 : params.parameters.structValue) === null || _e === void 0 ? void 0 : _e.fields) === null || _f === void 0 ? void 0 : _f.temperature.numberValue) {
attributes[SpanAttributes.LLM_REQUEST_TEMPERATURE] =
(_h = (_g = params === null || params === void 0 ? void 0 : params.parameters.structValue) === null || _g === void 0 ? void 0 : _g.fields) === null || _h === void 0 ? void 0 : _h.temperature.numberValue;
}
if ((_k = (_j = params === null || params === void 0 ? void 0 : params.parameters.structValue) === null || _j === void 0 ? void 0 : _j.fields) === null || _k === void 0 ? void 0 : _k.topP.numberValue) {
attributes[SpanAttributes.LLM_REQUEST_TOP_P] =
(_m = (_l = params === null || params === void 0 ? void 0 : params.parameters.structValue) === null || _l === void 0 ? void 0 : _l.fields) === null || _m === void 0 ? void 0 : _m.topP.numberValue;
}
if ((_p = (_o = params === null || params === void 0 ? void 0 : params.parameters.structValue) === null || _o === void 0 ? void 0 : _o.fields) === null || _p === void 0 ? void 0 : _p.topK.numberValue) {
attributes[SpanAttributes.LLM_TOP_K] =
(_r = (_q = params === null || params === void 0 ? void 0 : params.parameters.structValue) === null || _q === void 0 ? void 0 : _q.fields) === null || _r === void 0 ? void 0 : _r.topK.numberValue;
}
}
if (this._shouldSendPrompts() &&
params.instances &&
((_s = params.instances) === null || _s === void 0 ? void 0 : _s.length) !== 0) {
if (((_t = params.instances[0].structValue) === null || _t === void 0 ? void 0 : _t.fields) &&
"prompt" in params.instances[0].structValue.fields &&
((_v = (_u = params.instances[0].structValue) === null || _u === void 0 ? void 0 : _u.fields) === null || _v === void 0 ? void 0 : _v.prompt.stringValue)) {
attributes[`${SpanAttributes.LLM_PROMPTS}.0.role`] = "user";
attributes[`${SpanAttributes.LLM_PROMPTS}.0.content`] =
(_x = (_w = params.instances[0].structValue) === null || _w === void 0 ? void 0 : _w.fields) === null || _x === void 0 ? void 0 : _x.prompt.stringValue;
}
else if (params.instances[0].structValue &&
((_2 = (_1 = (_0 = (_z = (_y = params.instances[0].structValue.fields) === null || _y === void 0 ? void 0 : _y.messages.listValue) === null || _z === void 0 ? void 0 : _z.values) === null || _0 === void 0 ? void 0 : _0[0].structValue) === null || _1 === void 0 ? void 0 : _1.fields) === null || _2 === void 0 ? void 0 : _2.content.stringValue)) {
attributes[`${SpanAttributes.LLM_PROMPTS}.0.role`] =
(_8 = (_7 = (_6 = (_5 = (_4 = (_3 = params.instances[0].structValue.fields) === null || _3 === void 0 ? void 0 : _3.messages.listValue) === null || _4 === void 0 ? void 0 : _4.values) === null || _5 === void 0 ? void 0 : _5[0].structValue) === null || _6 === void 0 ? void 0 : _6.fields) === null || _7 === void 0 ? void 0 : _7.author.stringValue) !== null && _8 !== void 0 ? _8 : "user";
attributes[`${SpanAttributes.LLM_PROMPTS}.0.content`] =
(_13 = (_12 = (_11 = (_10 = (_9 = params.instances[0].structValue.fields) === null || _9 === void 0 ? void 0 : _9.messages.listValue) === null || _10 === void 0 ? void 0 : _10.values) === null || _11 === void 0 ? void 0 : _11[0].structValue) === null || _12 === void 0 ? void 0 : _12.fields) === null || _13 === void 0 ? void 0 : _13.content.stringValue;
}
}
}
}
catch (e) {
this._diag.debug(e);
(_15 = (_14 = this._config).exceptionLogger) === null || _15 === void 0 ? void 0 : _15.call(_14, e);
}
return this.tracer.startSpan(`vertexai.completion`, {
kind: SpanKind.CLIENT,
attributes,
});
}
_wrapPromise(span, promise) {
return promise
.then((result) => {
return new Promise((resolve) => {
this._endSpan({
span,
result: result,
});
resolve(result);
});
})
.catch((error) => {
return new Promise((_, reject) => {
span.setStatus({
code: SpanStatusCode.ERROR,
message: error.message,
});
span.recordException(error);
span.end();
reject(error);
});
});
}
_endSpan({ span, result, }) {
var _a, _b, _c, _d, _e, _f, _g, _h, _j, _k, _l, _m, _o, _p, _q, _r, _s, _t, _u, _v, _w, _x, _y, _z, _0, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _30, _31, _32, _33, _34;
try {
if (result[0].model)
span.setAttribute(SpanAttributes.LLM_RESPONSE_MODEL, result[0].model);
if (result) {
if (result[0].metadata) {
if (typeof ((_g = (_f = (_e = (_d = (_c = (_b = (_a = result[0].metadata) === null || _a === void 0 ? void 0 : _a.structValue) === null || _b === void 0 ? void 0 : _b.fields) === null || _c === void 0 ? void 0 : _c.tokenMetadata.structValue) === null || _d === void 0 ? void 0 : _d.fields) === null || _e === void 0 ? void 0 : _e.outputTokenCount.structValue) === null || _f === void 0 ? void 0 : _f.fields) === null || _g === void 0 ? void 0 : _g.totalTokens.numberValue) === "number")
span.setAttribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, (_p = (_o = (_m = (_l = (_k = (_j = (_h = result[0].metadata) === null || _h === void 0 ? void 0 : _h.structValue) === null || _j === void 0 ? void 0 : _j.fields) === null || _k === void 0 ? void 0 : _k.tokenMetadata.structValue) === null || _l === void 0 ? void 0 : _l.fields) === null || _m === void 0 ? void 0 : _m.outputTokenCount.structValue) === null || _o === void 0 ? void 0 : _o.fields) === null || _p === void 0 ? void 0 : _p.totalTokens.numberValue);
if (typeof ((_w = (_v = (_u = (_t = (_s = (_r = (_q = result[0].metadata) === null || _q === void 0 ? void 0 : _q.structValue) === null || _r === void 0 ? void 0 : _r.fields) === null || _s === void 0 ? void 0 : _s.tokenMetadata.structValue) === null || _t === void 0 ? void 0 : _t.fields) === null || _u === void 0 ? void 0 : _u.inputTokenCount.structValue) === null || _v === void 0 ? void 0 : _v.fields) === null || _w === void 0 ? void 0 : _w.totalTokens.numberValue) === "number")
span.setAttribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, (_3 = (_2 = (_1 = (_0 = (_z = (_y = (_x = result[0].metadata) === null || _x === void 0 ? void 0 : _x.structValue) === null || _y === void 0 ? void 0 : _y.fields) === null || _z === void 0 ? void 0 : _z.tokenMetadata.structValue) === null || _0 === void 0 ? void 0 : _0.fields) === null || _1 === void 0 ? void 0 : _1.inputTokenCount.structValue) === null || _2 === void 0 ? void 0 : _2.fields) === null || _3 === void 0 ? void 0 : _3.totalTokens.numberValue);
if (typeof ((_10 = (_9 = (_8 = (_7 = (_6 = (_5 = (_4 = result[0].metadata) === null || _4 === void 0 ? void 0 : _4.structValue) === null || _5 === void 0 ? void 0 : _5.fields) === null || _6 === void 0 ? void 0 : _6.tokenMetadata.structValue) === null || _7 === void 0 ? void 0 : _7.fields) === null || _8 === void 0 ? void 0 : _8.inputTokenCount.structValue) === null || _9 === void 0 ? void 0 : _9.fields) === null || _10 === void 0 ? void 0 : _10.totalTokens.numberValue) === "number" &&
typeof ((_17 = (_16 = (_15 = (_14 = (_13 = (_12 = (_11 = result[0].metadata) === null || _11 === void 0 ? void 0 : _11.structValue) === null || _12 === void 0 ? void 0 : _12.fields) === null || _13 === void 0 ? void 0 : _13.tokenMetadata.structValue) === null || _14 === void 0 ? void 0 : _14.fields) === null || _15 === void 0 ? void 0 : _15.outputTokenCount.structValue) === null || _16 === void 0 ? void 0 : _16.fields) === null || _17 === void 0 ? void 0 : _17.totalTokens.numberValue) === "number")
span.setAttribute(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, ((_24 = (_23 = (_22 = (_21 = (_20 = (_19 = (_18 = result[0].metadata) === null || _18 === void 0 ? void 0 : _18.structValue) === null || _19 === void 0 ? void 0 : _19.fields) === null || _20 === void 0 ? void 0 : _20.tokenMetadata.structValue) === null || _21 === void 0 ? void 0 : _21.fields) === null || _22 === void 0 ? void 0 : _22.inputTokenCount.structValue) === null || _23 === void 0 ? void 0 : _23.fields) === null || _24 === void 0 ? void 0 : _24.totalTokens.numberValue) +
((_31 = (_30 = (_29 = (_28 = (_27 = (_26 = (_25 = result[0].metadata) === null || _25 === void 0 ? void 0 : _25.structValue) === null || _26 === void 0 ? void 0 : _26.fields) === null || _27 === void 0 ? void 0 : _27.tokenMetadata.structValue) === null || _28 === void 0 ? void 0 : _28.fields) === null || _29 === void 0 ? void 0 : _29.outputTokenCount.structValue) === null || _30 === void 0 ? void 0 : _30.fields) === null || _31 === void 0 ? void 0 : _31.totalTokens.numberValue));
}
if (this._shouldSendPrompts()) {
(_32 = result[0].predictions) === null || _32 === void 0 ? void 0 : _32.forEach((prediction, index) => {
var _a, _b, _c, _d, _e, _f, _g, _h, _j, _k, _l, _m, _o, _p, _q, _r, _s, _t, _u, _v;
if (((_a = prediction.structValue) === null || _a === void 0 ? void 0 : _a.fields) &&
"content" in prediction.structValue.fields &&
!!((_c = (_b = prediction.structValue) === null || _b === void 0 ? void 0 : _b.fields) === null || _c === void 0 ? void 0 : _c.content.stringValue)) {
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.role`, "assistant");
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.content`, (_e = (_d = prediction.structValue) === null || _d === void 0 ? void 0 : _d.fields) === null || _e === void 0 ? void 0 : _e.content.stringValue);
}
else if (((_f = prediction.structValue) === null || _f === void 0 ? void 0 : _f.fields) &&
"candidates" in prediction.structValue.fields &&
!!((_o = (_m = (_l = (_k = (_j = (_h = (_g = prediction.structValue) === null || _g === void 0 ? void 0 : _g.fields) === null || _h === void 0 ? void 0 : _h.candidates.listValue) === null || _j === void 0 ? void 0 : _j.values) === null || _k === void 0 ? void 0 : _k[0]) === null || _l === void 0 ? void 0 : _l.structValue) === null || _m === void 0 ? void 0 : _m.fields) === null || _o === void 0 ? void 0 : _o.content.stringValue)) {
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.role`, "assistant");
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.content`, (_v = (_u = (_t = (_s = (_r = (_q = (_p = prediction.structValue) === null || _p === void 0 ? void 0 : _p.fields) === null || _q === void 0 ? void 0 : _q.candidates.listValue) === null || _r === void 0 ? void 0 : _r.values) === null || _s === void 0 ? void 0 : _s[0]) === null || _t === void 0 ? void 0 : _t.structValue) === null || _u === void 0 ? void 0 : _u.fields) === null || _v === void 0 ? void 0 : _v.content.stringValue);
}
});
}
}
}
catch (e) {
this._diag.debug(e);
(_34 = (_33 = this._config).exceptionLogger) === null || _34 === void 0 ? void 0 : _34.call(_33, e);
}
span.setStatus({ code: SpanStatusCode.OK });
span.end();
}
_shouldSendPrompts() {
const contextShouldSendPrompts = context
.active()
.getValue(CONTEXT_KEY_ALLOW_TRACE_CONTENT);
if (contextShouldSendPrompts !== undefined) {
return contextShouldSendPrompts;
}
return this._config.traceContent !== undefined
? this._config.traceContent
: true;
}
}
class VertexAIInstrumentation extends InstrumentationBase {
constructor(config = {}) {
super("@traceloop/instrumentation-vertexai", version, config);
}
setConfig(config = {}) {
super.setConfig(config);
}
init() {
const vertexAIModule = new InstrumentationNodeModuleDefinition("@google-cloud/vertexai", [">=1.1.0"], this.wrap.bind(this), this.unwrap.bind(this));
return vertexAIModule;
}
manuallyInstrument(module) {
this._diag.debug("Manually instrumenting @google-cloud/vertexai");
this._wrap(module.GenerativeModel.prototype, "generateContentStream", this.wrapperMethod());
this._wrap(module.GenerativeModel.prototype, "generateContent", this.wrapperMethod());
}
wrap(module, moduleVersion) {
this._diag.debug(`Patching @google-cloud/vertexai@${moduleVersion}`);
this._wrap(module.GenerativeModel.prototype, "generateContentStream", this.wrapperMethod());
this._wrap(module.GenerativeModel.prototype, "generateContent", this.wrapperMethod());
return module;
}
unwrap(module, moduleVersion) {
this._diag.debug(`Unpatching @google-cloud/vertexai@${moduleVersion}`);
this._unwrap(module.GenerativeModel.prototype, "generateContentStream");
this._unwrap(module.GenerativeModel.prototype, "generateContent");
}
wrapperMethod() {
// eslint-disable-next-line @typescript-eslint/no-this-alias
const plugin = this;
// eslint-disable-next-line
return (original) => {
return function method(...args) {
const span = plugin._startSpan({
instance: this,
params: args[0],
});
const execContext = trace.setSpan(context.active(), span);
const execPromise = safeExecuteInTheMiddle(() => {
return context.with(execContext, () => {
return original.apply(this, args);
});
},
// eslint-disable-next-line @typescript-eslint/no-empty-function
() => { });
const wrappedPromise = plugin._wrapPromise(span, execPromise);
return context.bind(execContext, wrappedPromise);
};
};
}
_startSpan({ instance, params, }) {
var _a, _b;
const attributes = {
[SpanAttributes.LLM_SYSTEM]: "Google",
[SpanAttributes.LLM_REQUEST_TYPE]: "completion",
};
try {
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) {
let i = 0;
if (instance["systemInstruction"]) {
attributes[`${SpanAttributes.LLM_PROMPTS}.${i}.role`] = "system";
attributes[`${SpanAttributes.LLM_PROMPTS}.${i}.content`] =
this._formatPartsData(instance["systemInstruction"].parts);
i++;
}
params.contents.forEach((content, j) => {
var _a;
attributes[`${SpanAttributes.LLM_PROMPTS}.${i + j}.role`] =
(_a = content.role) !== null && _a !== void 0 ? _a : "user";
attributes[`${SpanAttributes.LLM_PROMPTS}.${i + j}.content`] =
this._formatPartsData(content.parts);
});
}
}
catch (e) {
this._diag.debug(e);
(_b = (_a = this._config).exceptionLogger) === null || _b === void 0 ? void 0 : _b.call(_a, e);
}
return this.tracer.startSpan(`vertexai.completion`, {
kind: SpanKind.CLIENT,
attributes,
});
}
_wrapPromise(span, promise) {
return promise
.then((result) => __awaiter(this, void 0, void 0, function* () {
yield this._endSpan({
span,
result: result,
});
return new Promise((resolve) => resolve(result));
}))
.catch((error) => {
return new Promise((_, reject) => {
span.setStatus({
code: SpanStatusCode.ERROR,
message: error.message,
});
span.recordException(error);
span.end();
reject(error);
});
});
}
_endSpan(_a) {
return __awaiter(this, arguments, void 0, function* ({ span, result, }) {
var _b, _c, _d, _e, _f, _g;
try {
const streamResponse = yield result.response;
if (((_b = streamResponse.usageMetadata) === null || _b === void 0 ? void 0 : _b.totalTokenCount) !== undefined)
span.setAttribute(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, streamResponse.usageMetadata.totalTokenCount);
if ((_c = streamResponse.usageMetadata) === null || _c === void 0 ? void 0 : _c.candidatesTokenCount)
span.setAttribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, streamResponse.usageMetadata.candidatesTokenCount);
if ((_d = streamResponse.usageMetadata) === null || _d === void 0 ? void 0 : _d.promptTokenCount)
span.setAttribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, streamResponse.usageMetadata.promptTokenCount);
if (this._shouldSendPrompts()) {
(_e = streamResponse.candidates) === null || _e === void 0 ? void 0 : _e.forEach((candidate, index) => {
var _a;
if (candidate.finishReason)
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.finish_reason`, candidate.finishReason);
if (candidate.content) {
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.role`, (_a = candidate.content.role) !== null && _a !== void 0 ? _a : "assistant");
span.setAttribute(`${SpanAttributes.LLM_COMPLETIONS}.${index}.content`, this._formatPartsData(candidate.content.parts));
}
});
}
}
catch (e) {
this._diag.debug(e);
(_g = (_f = this._config).exceptionLogger) === null || _g === void 0 ? void 0 : _g.call(_f, e);
}
span.setStatus({ code: SpanStatusCode.OK });
span.end();
});
}
_formatPartsData(parts) {
const result = parts
.map((part) => {
if (part.text)
return part.text;
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);
return result.join("\n");
}
_shouldSendPrompts() {
const contextShouldSendPrompts = context
.active()
.getValue(CONTEXT_KEY_ALLOW_TRACE_CONTENT);
if (contextShouldSendPrompts !== undefined) {
return contextShouldSendPrompts;
}
return this._config.traceContent !== undefined
? this._config.traceContent
: true;
}
}
export { AIPlatformInstrumentation, VertexAIInstrumentation };
//# sourceMappingURL=index.mjs.map