UNPKG

@traceloop/instrumentation-langchain

Version:
208 lines (203 loc) 10.1 kB
import { trace, context, SpanStatusCode } from '@opentelemetry/api'; import { safeExecuteInTheMiddle, InstrumentationBase, InstrumentationNodeModuleDefinition } from '@opentelemetry/instrumentation'; import { TraceloopSpanKindValues, SpanAttributes, CONTEXT_KEY_ALLOW_TRACE_CONTENT } from '@traceloop/ai-semantic-conventions'; function genericWrapper(tracer, shouldSendPrompts, spanKind, spanName) { // eslint-disable-next-line return (original) => { return function method(...args) { var _a, _b; const span = tracer().startSpan(spanName || `${this.constructor.name}.${spanKind}`); span.setAttribute(SpanAttributes.TRACELOOP_SPAN_KIND, spanKind); if (shouldSendPrompts) { try { if (args.length === 1 && typeof args[0] === "object" && !(args[0] instanceof Map)) { span.setAttribute(SpanAttributes.TRACELOOP_ENTITY_INPUT, JSON.stringify({ args: [], kwargs: args[0] })); } else { span.setAttribute(SpanAttributes.TRACELOOP_ENTITY_INPUT, JSON.stringify({ args: args.map((arg) => arg instanceof Map ? Array.from(arg.entries()) : arg), kwargs: {}, })); } } catch (e) { this._diag.debug(e); (_b = (_a = this._config).exceptionLogger) === null || _b === void 0 ? void 0 : _b.call(_a, e); } } 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 = execPromise .then((result) => { return new Promise((resolve) => { var _a, _b; span.setStatus({ code: SpanStatusCode.OK }); try { if (shouldSendPrompts) { if (result instanceof Map) { span.setAttribute(SpanAttributes.TRACELOOP_ENTITY_OUTPUT, JSON.stringify(Array.from(result.entries()))); } else { span.setAttribute(SpanAttributes.TRACELOOP_ENTITY_OUTPUT, JSON.stringify(result)); } } } catch (e) { this._diag.debug(e); (_b = (_a = this._config).exceptionLogger) === null || _b === void 0 ? void 0 : _b.call(_a, e); } finally { span.end(); resolve(result); } }); }) .catch((error) => { return new Promise((_, reject) => { span.setStatus({ code: SpanStatusCode.ERROR, message: error.message, }); span.end(); reject(error); }); }); return context.bind(execContext, wrappedPromise); }; }; } function taskWrapper(tracer, shouldSendPrompts, spanName) { return genericWrapper(tracer, shouldSendPrompts, TraceloopSpanKindValues.TASK, spanName); } function workflowWrapper(tracer, shouldSendPrompts, spanName) { return genericWrapper(tracer, shouldSendPrompts, TraceloopSpanKindValues.WORKFLOW, spanName); } var version = "0.12.0"; /* * 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 LangChainInstrumentation extends InstrumentationBase { constructor(config = {}) { super("@traceloop/instrumentation-langchain", version, config); } manuallyInstrument({ chainsModule, agentsModule, toolsModule, vectorStoreModule, runnablesModule, }) { if (chainsModule) { this._diag.debug("Manually instrumenting langchain chains"); this.patchChainModule(chainsModule); } if (agentsModule) { this._diag.debug("Manually instrumenting langchain agents"); this.patchAgentModule(agentsModule); } if (toolsModule) { this._diag.debug("Manually instrumenting langchain tools"); this.patchToolsModule(toolsModule); } if (vectorStoreModule) { this._diag.debug("Manually instrumenting langchain vector stores"); this.patchVectorStoreModule(vectorStoreModule); } if (runnablesModule) { this._diag.debug("Manually instrumenting @langchain/core/runnables"); this.patchRunnablesModule(runnablesModule); } } init() { const chainModule = new InstrumentationNodeModuleDefinition("langchain/chains.cjs", [">=0.3.0"], this.patchChainModule.bind(this), this.unpatchChainModule.bind(this)); const agentModule = new InstrumentationNodeModuleDefinition("langchain/agents.cjs", [">=0.3.0"], this.patchAgentModule.bind(this), this.unpatchAgentModule.bind(this)); const toolsModule = new InstrumentationNodeModuleDefinition("langchain/tools.cjs", [">=0.3.0"], this.patchToolsModule.bind(this), this.unpatchToolsModule.bind(this)); const vectorStoreModule = new InstrumentationNodeModuleDefinition("langchain/core/vectorstores.cjs", [">=0.3.0"], this.patchVectorStoreModule.bind(this), this.unpatchVectorStoreModule.bind(this)); const runnablesModule = new InstrumentationNodeModuleDefinition("@langchain/core/runnables.cjs", [">=0.3.0"], this.patchRunnablesModule.bind(this), this.unpatchRunnablesModule.bind(this)); return [ chainModule, agentModule, toolsModule, vectorStoreModule, runnablesModule, ]; } patchChainModule(moduleExports, moduleVersion) { this._diag.debug(`Patching langchain/chains.cjs@${moduleVersion}`); this._wrap(moduleExports.RetrievalQAChain.prototype, "_call", workflowWrapper(() => this.tracer, this._shouldSendPrompts(), "retrieval_qa.workflow")); this._wrap(moduleExports.BaseChain.prototype, "call", taskWrapper(() => this.tracer, this._shouldSendPrompts())); return moduleExports; } patchAgentModule(moduleExports, moduleVersion) { this._diag.debug(`Patching langchain/agents.cjs@${moduleVersion}`); this._wrap(moduleExports.AgentExecutor.prototype, "_call", workflowWrapper(() => this.tracer, this._shouldSendPrompts(), "langchain.agent")); return moduleExports; } patchToolsModule(moduleExports, moduleVersion) { this._diag.debug(`Patching langchain/tools.cjs@${moduleVersion}`); this._wrap(moduleExports.Tool.prototype, "call", taskWrapper(() => this.tracer, this._shouldSendPrompts())); return moduleExports; } patchVectorStoreModule(moduleExports, moduleVersion) { this._diag.debug(`Patching langchain/vectorstores.cjs@${moduleVersion}`); this._wrap(moduleExports.VectorStoreRetriever.prototype, "_getRelevantDocuments", taskWrapper(() => this.tracer, this._shouldSendPrompts())); return moduleExports; } patchRunnablesModule(moduleExports, moduleVersion) { this._diag.debug(`Patching @langchain/core/runnables@${moduleVersion}`); this._wrap(moduleExports.RunnableSequence.prototype, "invoke", taskWrapper(() => this.tracer, this._shouldSendPrompts())); return moduleExports; } unpatchChainModule(moduleExports, moduleVersion) { this._diag.debug(`Unpatching langchain/chains.cjs@${moduleVersion}`); this._unwrap(moduleExports.RetrievalQAChain.prototype, "_call"); this._unwrap(moduleExports.BaseChain.prototype, "call"); return moduleExports; } unpatchAgentModule(moduleExports, moduleVersion) { this._diag.debug(`Unpatching langchain/agents.cjs@${moduleVersion}`); this._unwrap(moduleExports.AgentExecutor.prototype, "_call"); return moduleExports; } unpatchToolsModule(moduleExports) { this._diag.debug(`Unpatching langchain/tools.cjs`); this._unwrap(moduleExports.Tool.prototype, "call"); return moduleExports; } unpatchVectorStoreModule(moduleExports) { this._diag.debug(`Unpatching langchain/vectorstores.cjs`); this._unwrap(moduleExports.VectorStoreRetriever.prototype, "_getRelevantDocuments"); return moduleExports; } unpatchRunnablesModule(moduleExports) { this._diag.debug(`Unpatching @langchain/core/runnables`); this._unwrap(moduleExports.Runnable.prototype, "invoke"); return moduleExports; } _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 { LangChainInstrumentation }; //# sourceMappingURL=index.mjs.map