UNPKG

@traceloop/instrumentation-vertexai

Version:
389 lines (383 loc) 26.5 kB
'use strict'; var api = require('@opentelemetry/api'); var instrumentation = require('@opentelemetry/instrumentation'); var aiSemanticConventions = require('@traceloop/ai-semantic-conventions'); var tslib = require('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 instrumentation.InstrumentationBase { constructor(config = {}) { super("@traceloop/instrumentation-vertexai", version, config); } setConfig(config = {}) { super.setConfig(config); } init() { const aiPlatformModule = new instrumentation.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 = api.trace.setSpan(api.context.active(), span); const execPromise = instrumentation.safeExecuteInTheMiddle(() => { return api.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 api.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 = { [aiSemanticConventions.SpanAttributes.LLM_SYSTEM]: "Google", [aiSemanticConventions.SpanAttributes.LLM_REQUEST_TYPE]: "completion", }; try { if (params !== undefined) { if (params.endpoint) { const model = params.endpoint.split("/").pop(); attributes[aiSemanticConventions.SpanAttributes.LLM_REQUEST_MODEL] = model; attributes[aiSemanticConventions.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[aiSemanticConventions.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[aiSemanticConventions.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[aiSemanticConventions.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[aiSemanticConventions.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[`${aiSemanticConventions.SpanAttributes.LLM_PROMPTS}.0.role`] = "user"; attributes[`${aiSemanticConventions.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[`${aiSemanticConventions.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[`${aiSemanticConventions.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: api.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: api.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(aiSemanticConventions.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(aiSemanticConventions.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(aiSemanticConventions.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(aiSemanticConventions.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(`${aiSemanticConventions.SpanAttributes.LLM_COMPLETIONS}.${index}.role`, "assistant"); span.setAttribute(`${aiSemanticConventions.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(`${aiSemanticConventions.SpanAttributes.LLM_COMPLETIONS}.${index}.role`, "assistant"); span.setAttribute(`${aiSemanticConventions.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: api.SpanStatusCode.OK }); span.end(); } _shouldSendPrompts() { const contextShouldSendPrompts = api.context .active() .getValue(aiSemanticConventions.CONTEXT_KEY_ALLOW_TRACE_CONTENT); if (contextShouldSendPrompts !== undefined) { return contextShouldSendPrompts; } return this._config.traceContent !== undefined ? this._config.traceContent : true; } } class VertexAIInstrumentation extends instrumentation.InstrumentationBase { constructor(config = {}) { super("@traceloop/instrumentation-vertexai", version, config); } setConfig(config = {}) { super.setConfig(config); } init() { const vertexAIModule = new instrumentation.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 = api.trace.setSpan(api.context.active(), span); const execPromise = instrumentation.safeExecuteInTheMiddle(() => { return api.context.with(execContext, () => { return original.apply(this, args); }); }, // eslint-disable-next-line @typescript-eslint/no-empty-function () => { }); const wrappedPromise = plugin._wrapPromise(span, execPromise); return api.context.bind(execContext, wrappedPromise); }; }; } _startSpan({ instance, params, }) { var _a, _b; const attributes = { [aiSemanticConventions.SpanAttributes.LLM_SYSTEM]: "Google", [aiSemanticConventions.SpanAttributes.LLM_REQUEST_TYPE]: "completion", }; try { attributes[aiSemanticConventions.SpanAttributes.LLM_REQUEST_MODEL] = instance["model"]; attributes[aiSemanticConventions.SpanAttributes.LLM_RESPONSE_MODEL] = instance["model"]; if (instance["generationConfig"]) { attributes[aiSemanticConventions.SpanAttributes.LLM_REQUEST_MAX_TOKENS] = instance["generationConfig"].max_output_tokens; attributes[aiSemanticConventions.SpanAttributes.LLM_REQUEST_TEMPERATURE] = instance["generationConfig"].temperature; attributes[aiSemanticConventions.SpanAttributes.LLM_REQUEST_TOP_P] = instance["generationConfig"].top_p; attributes[aiSemanticConventions.SpanAttributes.LLM_TOP_K] = instance["generationConfig"].top_k; } if (this._shouldSendPrompts() && "contents" in params) { let i = 0; if (instance["systemInstruction"]) { attributes[`${aiSemanticConventions.SpanAttributes.LLM_PROMPTS}.${i}.role`] = "system"; attributes[`${aiSemanticConventions.SpanAttributes.LLM_PROMPTS}.${i}.content`] = this._formatPartsData(instance["systemInstruction"].parts); i++; } params.contents.forEach((content, j) => { var _a; attributes[`${aiSemanticConventions.SpanAttributes.LLM_PROMPTS}.${i + j}.role`] = (_a = content.role) !== null && _a !== void 0 ? _a : "user"; attributes[`${aiSemanticConventions.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: api.SpanKind.CLIENT, attributes, }); } _wrapPromise(span, promise) { return promise .then((result) => tslib.__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: api.SpanStatusCode.ERROR, message: error.message, }); span.recordException(error); span.end(); reject(error); }); }); } _endSpan(_a) { return tslib.__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(aiSemanticConventions.SpanAttributes.LLM_USAGE_TOTAL_TOKENS, streamResponse.usageMetadata.totalTokenCount); if ((_c = streamResponse.usageMetadata) === null || _c === void 0 ? void 0 : _c.candidatesTokenCount) span.setAttribute(aiSemanticConventions.SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, streamResponse.usageMetadata.candidatesTokenCount); if ((_d = streamResponse.usageMetadata) === null || _d === void 0 ? void 0 : _d.promptTokenCount) span.setAttribute(aiSemanticConventions.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(`${aiSemanticConventions.SpanAttributes.LLM_COMPLETIONS}.${index}.finish_reason`, candidate.finishReason); if (candidate.content) { span.setAttribute(`${aiSemanticConventions.SpanAttributes.LLM_COMPLETIONS}.${index}.role`, (_a = candidate.content.role) !== null && _a !== void 0 ? _a : "assistant"); span.setAttribute(`${aiSemanticConventions.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: api.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 = api.context .active() .getValue(aiSemanticConventions.CONTEXT_KEY_ALLOW_TRACE_CONTENT); if (contextShouldSendPrompts !== undefined) { return contextShouldSendPrompts; } return this._config.traceContent !== undefined ? this._config.traceContent : true; } } exports.AIPlatformInstrumentation = AIPlatformInstrumentation; exports.VertexAIInstrumentation = VertexAIInstrumentation; //# sourceMappingURL=index.js.map