@traceloop/instrumentation-llamaindex
Version:
Llamaindex Instrumentation
399 lines (390 loc) • 18.5 kB
JavaScript
;
var instrumentation = require('@opentelemetry/instrumentation');
var lodash = require('lodash');
var api = require('@opentelemetry/api');
var aiSemanticConventions = require('@traceloop/ai-semantic-conventions');
var tslib = require('tslib');
function _interopNamespaceDefault(e) {
var n = Object.create(null);
if (e) {
Object.keys(e).forEach(function (k) {
if (k !== 'default') {
var d = Object.getOwnPropertyDescriptor(e, k);
Object.defineProperty(n, k, d.get ? d : {
enumerable: true,
get: function () { return e[k]; }
});
}
});
}
n.default = e;
return Object.freeze(n);
}
var lodash__namespace = /*#__PURE__*/_interopNamespaceDefault(lodash);
const shouldSendPrompts = (config) => {
const contextShouldSendPrompts = api.context
.active()
.getValue(aiSemanticConventions.CONTEXT_KEY_ALLOW_TRACE_CONTENT);
if (contextShouldSendPrompts !== undefined) {
return !!contextShouldSendPrompts;
}
return config.traceContent !== undefined ? config.traceContent : true;
};
// Adopted from https://github.com/open-telemetry/opentelemetry-js/issues/2951#issuecomment-1214587378
function bindAsyncGenerator(ctx, generator) {
var _a;
return {
next: api.context.bind(ctx, generator.next.bind(generator)),
return: api.context.bind(ctx, generator.return.bind(generator)),
throw: api.context.bind(ctx, generator.throw.bind(generator)),
[Symbol.asyncIterator]() {
return bindAsyncGenerator(ctx, generator[Symbol.asyncIterator]());
},
[Symbol.asyncDispose]: ((_a = generator[Symbol.asyncDispose]) === null || _a === void 0 ? void 0 : _a.bind(generator)) ||
(() => Promise.resolve()),
};
}
function generatorWrapper(streamingResult, ctx, fn) {
return tslib.__asyncGenerator(this, arguments, function* generatorWrapper_1() {
var _a, e_1, _b, _c;
try {
for (var _d = true, _e = tslib.__asyncValues(bindAsyncGenerator(ctx, streamingResult)), _f; _f = yield tslib.__await(_e.next()), _a = _f.done, !_a; _d = true) {
_c = _f.value;
_d = false;
const chunk = _c;
yield yield tslib.__await(chunk);
}
}
catch (e_1_1) { e_1 = { error: e_1_1 }; }
finally {
try {
if (!_d && !_a && (_b = _e.return)) yield tslib.__await(_b.call(_e));
}
finally { if (e_1) throw e_1.error; }
}
fn();
});
}
function llmGeneratorWrapper(streamingResult, ctx, fn) {
return tslib.__asyncGenerator(this, arguments, function* llmGeneratorWrapper_1() {
var _a, e_2, _b, _c;
let message = "";
try {
for (var _d = true, _e = tslib.__asyncValues(bindAsyncGenerator(ctx, streamingResult)), _f; _f = yield tslib.__await(_e.next()), _a = _f.done, !_a; _d = true) {
_c = _f.value;
_d = false;
const messageChunk = _c;
if (messageChunk.delta) {
message += messageChunk.delta;
}
if (messageChunk.text) {
message += messageChunk.text;
}
yield yield tslib.__await(messageChunk);
}
}
catch (e_2_1) { e_2 = { error: e_2_1 }; }
finally {
try {
if (!_d && !_a && (_b = _e.return)) yield tslib.__await(_b.call(_e));
}
finally { if (e_2) throw e_2.error; }
}
fn(message);
});
}
function genericWrapper(className, methodName, kind, tracer, shouldSendPrompts) {
// eslint-disable-next-line
return (original) => {
return function method(...args) {
const params = args[0];
const streaming = params && params.stream;
const name = `${lodash__namespace.snakeCase(className)}.${lodash__namespace.snakeCase(methodName)}`;
const span = tracer().startSpan(`${name}`, {}, api.context.active());
span.setAttribute(aiSemanticConventions.SpanAttributes.TRACELOOP_SPAN_KIND, kind);
if (kind === aiSemanticConventions.TraceloopSpanKindValues.WORKFLOW) {
span.setAttribute(aiSemanticConventions.SpanAttributes.TRACELOOP_WORKFLOW_NAME, name);
}
if (shouldSendPrompts) {
try {
if (args.length === 1 &&
typeof args[0] === "object" &&
!(args[0] instanceof Map)) {
span.setAttribute(aiSemanticConventions.SpanAttributes.TRACELOOP_ENTITY_INPUT, JSON.stringify({ args: [], kwargs: args[0] }));
}
else {
span.setAttribute(aiSemanticConventions.SpanAttributes.TRACELOOP_ENTITY_INPUT, JSON.stringify({
args: args.map((arg) => arg instanceof Map ? Array.from(arg.entries()) : arg),
kwargs: {},
}));
}
}
catch (_a) {
/* empty */
}
}
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 = execPromise
.then((result) => {
return new Promise((resolve) => {
if (streaming) {
result = generatorWrapper(result, execContext, () => {
span.setStatus({ code: api.SpanStatusCode.OK });
span.end();
});
resolve(result);
}
else {
span.setStatus({ code: api.SpanStatusCode.OK });
try {
if (shouldSendPrompts) {
if (result instanceof Map) {
span.setAttribute(aiSemanticConventions.SpanAttributes.TRACELOOP_ENTITY_OUTPUT, JSON.stringify(Array.from(result.entries())));
}
else {
span.setAttribute(aiSemanticConventions.SpanAttributes.TRACELOOP_ENTITY_OUTPUT, JSON.stringify(result));
}
}
}
finally {
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);
};
};
}
class CustomLLMInstrumentation {
constructor(config, diag, tracer) {
this.config = config;
this.diag = diag;
this.tracer = tracer;
}
chatWrapper({ className }) {
// eslint-disable-next-line @typescript-eslint/no-this-alias
const plugin = this;
return (original) => {
return function method(...args) {
var _a, _b;
const params = args[0];
const messages = params === null || params === void 0 ? void 0 : params.messages;
const streaming = params === null || params === void 0 ? void 0 : params.stream;
const span = plugin
.tracer()
.startSpan(`llamaindex.${lodash__namespace.snakeCase(className)}.chat`, {
kind: api.SpanKind.CLIENT,
});
try {
span.setAttribute(aiSemanticConventions.SpanAttributes.LLM_SYSTEM, className);
span.setAttribute(aiSemanticConventions.SpanAttributes.LLM_REQUEST_MODEL, this.metadata.model);
span.setAttribute(aiSemanticConventions.SpanAttributes.LLM_REQUEST_TYPE, "chat");
span.setAttribute(aiSemanticConventions.SpanAttributes.LLM_REQUEST_TOP_P, this.metadata.topP);
if (shouldSendPrompts(plugin.config)) {
for (const messageIdx in messages) {
const content = messages[messageIdx].content;
if (typeof content === "string") {
span.setAttribute(`${aiSemanticConventions.SpanAttributes.LLM_PROMPTS}.${messageIdx}.content`, content);
}
else if (content[0].type ===
"text") {
span.setAttribute(`${aiSemanticConventions.SpanAttributes.LLM_PROMPTS}.${messageIdx}.content`, content[0].text);
}
span.setAttribute(`${aiSemanticConventions.SpanAttributes.LLM_PROMPTS}.${messageIdx}.role`, messages[messageIdx].role);
}
}
}
catch (e) {
plugin.diag.warn(e);
(_b = (_a = plugin.config).exceptionLogger) === null || _b === void 0 ? void 0 : _b.call(_a, e);
}
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 = execPromise
.then((result) => {
return new Promise((resolve) => {
if (streaming) {
result = plugin.handleStreamingResponse(result, span, execContext, this.metadata);
}
else {
result = plugin.handleResponse(result, span, this.metadata);
}
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);
};
};
}
handleResponse(result, span, metadata) {
var _a, _b;
span.setAttribute(aiSemanticConventions.SpanAttributes.LLM_RESPONSE_MODEL, metadata.model);
if (!shouldSendPrompts(this.config)) {
span.setStatus({ code: api.SpanStatusCode.OK });
span.end();
return result;
}
try {
if (result.message) {
span.setAttribute(`${aiSemanticConventions.SpanAttributes.LLM_COMPLETIONS}.0.role`, result.message.role);
const content = result.message.content;
if (typeof content === "string") {
span.setAttribute(`${aiSemanticConventions.SpanAttributes.LLM_COMPLETIONS}.0.content`, content);
}
else if (content[0].type === "text") {
span.setAttribute(`${aiSemanticConventions.SpanAttributes.LLM_COMPLETIONS}.0.content`, content[0].text);
}
span.setStatus({ code: api.SpanStatusCode.OK });
}
}
catch (e) {
this.diag.warn(e);
(_b = (_a = this.config).exceptionLogger) === null || _b === void 0 ? void 0 : _b.call(_a, e);
}
span.end();
return result;
}
handleStreamingResponse(result, span, execContext, metadata) {
span.setAttribute(aiSemanticConventions.SpanAttributes.LLM_RESPONSE_MODEL, metadata.model);
if (!shouldSendPrompts(this.config)) {
span.setStatus({ code: api.SpanStatusCode.OK });
span.end();
return result;
}
return llmGeneratorWrapper(result, execContext, (message) => {
span.setAttribute(`${aiSemanticConventions.SpanAttributes.LLM_COMPLETIONS}.0.content`, message);
span.setStatus({ code: api.SpanStatusCode.OK });
span.end();
});
}
}
var version = "0.19.0";
class LlamaIndexInstrumentation extends instrumentation.InstrumentationBase {
constructor(config = {}) {
super("@traceloop/instrumentation-llamaindex", version, config);
}
setConfig(config = {}) {
super.setConfig(config);
}
manuallyInstrument(module) {
this._diag.debug("Manually instrumenting llamaindex");
this.patch(module);
}
init() {
const llamaindexModule = new instrumentation.InstrumentationNodeModuleDefinition("llamaindex", [">=0.1.0"], this.patch.bind(this), this.unpatch.bind(this));
const openaiModule = new instrumentation.InstrumentationNodeModuleDefinition("@llamaindex/openai", [">=0.1.0"], this.patchOpenAI.bind(this), this.unpatchOpenAI.bind(this));
return [llamaindexModule, openaiModule];
}
isLLM(llm) {
return (llm &&
llm.complete !== undefined &&
llm.chat !== undefined);
}
isEmbedding(embedding) {
return !!(embedding === null || embedding === void 0 ? void 0 : embedding.getQueryEmbedding);
}
isSynthesizer(synthesizer) {
return (synthesizer && synthesizer.synthesize !== undefined);
}
isRetriever(retriever) {
return retriever && retriever.retrieve !== undefined;
}
patch(moduleExports, moduleVersion) {
this._diag.debug(`Patching llamaindex@${moduleVersion}`);
const customLLMInstrumentation = new CustomLLMInstrumentation(this._config, this._diag, () => this.tracer);
this._wrap(moduleExports.RetrieverQueryEngine.prototype, "query", genericWrapper(moduleExports.RetrieverQueryEngine.name, "query", aiSemanticConventions.TraceloopSpanKindValues.WORKFLOW, () => this.tracer, shouldSendPrompts(this._config)));
this._wrap(moduleExports.ContextChatEngine.prototype, "chat", genericWrapper(moduleExports.ContextChatEngine.name, "chat", aiSemanticConventions.TraceloopSpanKindValues.WORKFLOW, () => this.tracer, shouldSendPrompts(this._config)));
// OpenAIAgent has been moved to @llamaindex/openai package in newer versions
// This instrumentation is handled separately
for (const key in moduleExports) {
const cls = moduleExports[key];
if (this.isLLM(cls.prototype)) {
this._wrap(cls.prototype, "chat", customLLMInstrumentation.chatWrapper({ className: cls.name }));
}
else if (this.isEmbedding(cls.prototype)) {
this._wrap(cls.prototype, "getQueryEmbedding", genericWrapper(cls.name, "getQueryEmbedding", aiSemanticConventions.TraceloopSpanKindValues.TASK, () => this.tracer, shouldSendPrompts(this._config)));
}
else if (this.isSynthesizer(cls.prototype)) {
this._wrap(cls.prototype, "synthesize", genericWrapper(cls.name, "synthesize", aiSemanticConventions.TraceloopSpanKindValues.TASK, () => this.tracer, shouldSendPrompts(this._config)));
}
else if (this.isRetriever(cls.prototype)) {
this._wrap(cls.prototype, "retrieve", genericWrapper(cls.name, "retrieve", aiSemanticConventions.TraceloopSpanKindValues.TASK, () => this.tracer, shouldSendPrompts(this._config)));
}
}
return moduleExports;
}
unpatch(moduleExports, moduleVersion) {
this._diag.debug(`Unpatching llamaindex@${moduleVersion}`);
this._unwrap(moduleExports.RetrieverQueryEngine.prototype, "query");
for (const key in moduleExports) {
const cls = moduleExports[key];
if (this.isLLM(cls.prototype)) {
this._unwrap(cls.prototype, "complete");
this._unwrap(cls.prototype, "chat");
}
else if (this.isEmbedding(cls.prototype)) {
this._unwrap(cls.prototype, "getQueryEmbedding");
}
else if (this.isSynthesizer(cls.prototype)) {
this._unwrap(cls.prototype, "synthesize");
}
else if (this.isRetriever(cls.prototype)) {
this._unwrap(cls.prototype, "retrieve");
}
}
return moduleExports;
}
patchOpenAI(moduleExports, moduleVersion) {
this._diag.debug(`Patching @llamaindex/openai@${moduleVersion}`);
// Instrument OpenAIAgent if it exists
if (moduleExports.OpenAIAgent && moduleExports.OpenAIAgent.prototype) {
this._wrap(moduleExports.OpenAIAgent.prototype, "chat", genericWrapper(moduleExports.OpenAIAgent.name, "agent", aiSemanticConventions.TraceloopSpanKindValues.AGENT, () => this.tracer, shouldSendPrompts(this._config)));
}
return moduleExports;
}
unpatchOpenAI(moduleExports, moduleVersion) {
this._diag.debug(`Unpatching @llamaindex/openai@${moduleVersion}`);
// Unwrap OpenAIAgent if it exists
if (moduleExports.OpenAIAgent && moduleExports.OpenAIAgent.prototype) {
this._unwrap(moduleExports.OpenAIAgent.prototype, "chat");
}
return moduleExports;
}
}
exports.LlamaIndexInstrumentation = LlamaIndexInstrumentation;
//# sourceMappingURL=index.js.map