@traceloop/instrumentation-pinecone
Version:
OpenTelemetry instrumentation for pinecone vector DB
165 lines (160 loc) • 9.61 kB
JavaScript
;
var api = require('@opentelemetry/api');
var instrumentation = require('@opentelemetry/instrumentation');
var aiSemanticConventions = require('@traceloop/ai-semantic-conventions');
var version = "0.13.0";
class PineconeInstrumentation extends instrumentation.InstrumentationBase {
constructor(config = {}) {
super("@traceloop/instrumentation-pinecone", version, config);
}
setConfig(config = {}) {
super.setConfig(config);
}
manuallyInstrument(module) {
this.patch(module);
}
init() {
const module = new instrumentation.InstrumentationNodeModuleDefinition("@pinecone-database/pinecone", [">=2.0.1"], this.patch.bind(this), this.unpatch.bind(this));
return module;
}
patch(moduleExports, moduleVersion) {
this._diag.debug(`Patching @pinecone-database/pinecone@${moduleVersion}`);
this._wrap(moduleExports.Index.prototype, "query", this.queryWrapper(this.tracer));
this._wrap(moduleExports.Index.prototype, "upsert", this.genericWrapper("upsert", this.tracer));
this._wrap(moduleExports.Index.prototype, "deleteAll", this.genericWrapper("delete", this.tracer));
this._wrap(moduleExports.Index.prototype, "deleteMany", this.genericWrapper("delete", this.tracer));
this._wrap(moduleExports.Index.prototype, "deleteOne", this.genericWrapper("delete", this.tracer));
return moduleExports;
}
unpatch(moduleExports, moduleVersion) {
this._diag.debug(`Unpatching @pinecone-database/pinecone@${moduleVersion}`);
this._unwrap(moduleExports.Index.prototype, "query");
this._unwrap(moduleExports.Index.prototype, "upsert");
this._unwrap(moduleExports.Index.prototype, "deleteAll");
this._unwrap(moduleExports.Index.prototype, "deleteMany");
this._unwrap(moduleExports.Index.prototype, "deleteOne");
}
genericWrapper(methodName, tracer) {
// eslint-disable-next-line @typescript-eslint/no-this-alias
const plugin = this;
// eslint-disable-next-line
return (original) => {
return function method(...args) {
const span = tracer.startSpan(`pinecone.${methodName}`);
span.setAttribute(aiSemanticConventions.SpanAttributes.VECTOR_DB_VENDOR, "Pinecone");
const execContext = api.trace.setSpan(api.context.active(), span);
const execPromise = instrumentation.safeExecuteInTheMiddle(() => {
return api.context.with(execContext, () => {
return original.apply(this, args);
});
}, (e) => {
plugin._diag.error(`Error in Pinecone instrumentation`, e);
});
const wrappedPromise = execPromise
.then((result) => {
return new Promise((resolve) => {
span.setStatus({ code: api.SpanStatusCode.OK });
span.end();
resolve(result);
});
})
.catch((error) => {
return new Promise((_, reject) => {
span.setStatus({
code: api.SpanStatusCode.ERROR,
message: error.message,
});
span.end();
reject(error);
});
});
return api.context.bind(execContext, wrappedPromise);
};
};
}
queryWrapper(tracer) {
// eslint-disable-next-line @typescript-eslint/no-this-alias
const plugin = this;
// eslint-disable-next-line
return (original) => {
return function method(...args) {
var _a, _b;
const span = tracer.startSpan(`pinecone.query`);
const execContext = api.trace.setSpan(api.context.active(), span);
try {
const options = args[0];
span.setAttribute(aiSemanticConventions.SpanAttributes.VECTOR_DB_VENDOR, "Pinecone");
const query_request_event = span.addEvent("pinecone.query.request");
query_request_event.setAttribute(aiSemanticConventions.EventAttributes.VECTOR_DB_QUERY_TOP_K, options.topK);
query_request_event.setAttribute(aiSemanticConventions.EventAttributes.VECTOR_DB_QUERY_INCLUDE_VALUES, options.includeValues || false);
query_request_event.setAttribute(aiSemanticConventions.EventAttributes.VECTOR_DB_QUERY_INCLUDE_METADATA, options.includeMetadata || false);
query_request_event.setAttribute(aiSemanticConventions.EventAttributes.VECTOR_DB_QUERY_ID, options.id);
query_request_event.setAttribute(aiSemanticConventions.EventAttributes.VECTOR_DB_QUERY_EMBEDDINGS_VECTOR, options.vector);
query_request_event.setAttribute(aiSemanticConventions.EventAttributes.VECTOR_DB_QUERY_METADATA_FILTER, JSON.stringify(options.filter ? options.filter : {}));
}
catch (e) {
this._diag.debug(e);
(_b = (_a = this._config).exceptionLogger) === null || _b === void 0 ? void 0 : _b.call(_a, e);
}
const execPromise = instrumentation.safeExecuteInTheMiddle(() => {
return api.context.with(execContext, () => {
return original.apply(this, args);
});
}, (e) => {
if (e) {
plugin._diag.error(`Error in Pinecone instrumentation`, e);
}
});
const wrappedPromise = execPromise
.then((result) => {
return new Promise((resolve) => {
var _a, _b, _c, _d, _e, _f;
span.setStatus({ code: api.SpanStatusCode.OK });
try {
const result_obj = result;
const query_result_event = span.addEvent("pinecone.query.result");
query_result_event.setAttribute(aiSemanticConventions.EventAttributes.VECTOR_DB_QUERY_RESULT_NAMESPACE, result_obj.namespace);
if (((_a = result_obj.usage) === null || _a === void 0 ? void 0 : _a.readUnits) !== undefined) {
query_result_event.setAttribute(aiSemanticConventions.EventAttributes.VECTOR_DB_QUERY_RESULT_READ_UNITS_CONSUMED, (_b = result_obj.usage) === null || _b === void 0 ? void 0 : _b.readUnits);
}
query_result_event.setAttribute(aiSemanticConventions.EventAttributes.VECTOR_DB_QUERY_RESULT_MATCHES_LENGTH, result_obj.matches.length);
for (let i = 0; i < result_obj.matches.length; i++) {
const match = result_obj.matches[i];
const query_result_match_event = query_result_event.addEvent(`pinecone.query.result.${i}`);
if (match.score !== undefined) {
query_result_match_event.setAttribute(aiSemanticConventions.EventAttributes.VECTOR_DB_QUERY_RESULT_SCORE.replace("{i}", i.toString()), match.score);
}
if (match.sparseValues !== undefined) {
query_result_match_event.setAttribute(aiSemanticConventions.EventAttributes.VECTOR_DB_QUERY_RESULT_SPARSE_INDICES.replace("{i}", i.toString()), (_c = match.sparseValues) === null || _c === void 0 ? void 0 : _c.indices);
query_result_match_event.setAttribute(aiSemanticConventions.EventAttributes.VECTOR_DB_QUERY_RESULT_SPARSE_VALUES.replace("{i}", i.toString()), (_d = match.sparseValues) === null || _d === void 0 ? void 0 : _d.values);
}
query_result_match_event.setAttribute(aiSemanticConventions.EventAttributes.VECTOR_DB_QUERY_RESULT_ID.replace("{i}", i.toString()), match.id);
query_result_match_event.setAttribute(aiSemanticConventions.EventAttributes.VECTOR_DB_QUERY_RESULT_VALUES.replace("{i}", i.toString()), match.values);
query_result_match_event.addEvent(`pinecone.query.result.${i}.metadata`, match.metadata);
}
}
catch (e) {
this._diag.debug(e);
(_f = (_e = this._config).exceptionLogger) === null || _f === void 0 ? void 0 : _f.call(_e, e);
}
span.end();
resolve(result);
});
})
.catch((error) => {
return new Promise((_, reject) => {
span.setStatus({
code: api.SpanStatusCode.ERROR,
message: error.message,
});
span.end();
reject(error);
});
});
return api.context.bind(execContext, wrappedPromise);
};
};
}
}
exports.PineconeInstrumentation = PineconeInstrumentation;
//# sourceMappingURL=index.js.map