diff --git a/package-lock.json b/package-lock.json index 65101865..5608513f 100644 --- a/package-lock.json +++ b/package-lock.json @@ -11,6 +11,9 @@ "workspaces": [ "packages/*" ], + "dependencies": { + "@traceloop/node-server-sdk": "^0.5.25" + }, "devDependencies": { "@commitlint/cli": "^18.4.4", "@commitlint/config-conventional": "^18.4.4", @@ -29,7 +32,7 @@ "eslint-plugin-prettier": "^5.0.1", "husky": "^8.0.3", "lerna": "^8.1.2", - "nx": "^17.2.8", + "nx": "^17.3.2", "prettier": "^3.0.3", "rollup": "^4.13.0", "rollup-plugin-dts": "^6.1.0", @@ -6595,6 +6598,50 @@ "node": ">=10" } }, + "node_modules/@sinonjs/commons": { + "version": "3.0.1", + "resolved": "https://registry.npmjs.org/@sinonjs/commons/-/commons-3.0.1.tgz", + "integrity": "sha512-K3mCHKQ9sVh8o1C9cxkwxaOmXoAMlDxC1mYyHrjqOWEcBjYr76t96zL2zlj5dUGZ3HSw240X1qgH3Mjf1yJWpQ==", + "dev": true, + "dependencies": { + "type-detect": "4.0.8" + } + }, + "node_modules/@sinonjs/fake-timers": { + "version": "11.2.2", + "resolved": "https://registry.npmjs.org/@sinonjs/fake-timers/-/fake-timers-11.2.2.tgz", + "integrity": "sha512-G2piCSxQ7oWOxwGSAyFHfPIsyeJGXYtc6mFbnFA+kRXkiEnTl8c/8jul2S329iFBnDI9HGoeWWAZvuvOkZccgw==", + "dev": true, + "dependencies": { + "@sinonjs/commons": "^3.0.0" + } + }, + "node_modules/@sinonjs/samsam": { + "version": "8.0.0", + "resolved": "https://registry.npmjs.org/@sinonjs/samsam/-/samsam-8.0.0.tgz", + "integrity": "sha512-Bp8KUVlLp8ibJZrnvq2foVhP0IVX2CIprMJPK0vqGqgrDa0OHVKeZyBykqskkrdxV6yKBPmGasO8LVjAKR3Gew==", + "dev": true, + "dependencies": { + "@sinonjs/commons": "^2.0.0", + "lodash.get": "^4.4.2", + "type-detect": "^4.0.8" + } + }, + "node_modules/@sinonjs/samsam/node_modules/@sinonjs/commons": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/@sinonjs/commons/-/commons-2.0.0.tgz", + "integrity": "sha512-uLa0j859mMrg2slwQYdO/AkrOfmH+X6LTVmNTS9CqexuE2IvVORIkSpJLqePAbEnKJ77aMmCwr1NUZ57120Xcg==", + "dev": true, + "dependencies": { + "type-detect": "4.0.8" + } + }, + "node_modules/@sinonjs/text-encoding": { + "version": "0.7.2", + "resolved": "https://registry.npmjs.org/@sinonjs/text-encoding/-/text-encoding-0.7.2.tgz", + "integrity": "sha512-sXXKG+uL9IrKqViTtao2Ws6dy0znu9sOaP1di/jKGW1M6VssO8vlpXCQcpZ+jisQ1tTFAC5Jo/EOzFbggBagFQ==", + "dev": true + }, "node_modules/@smithy/abort-controller": { "version": "2.1.3", "resolved": "https://registry.npmjs.org/@smithy/abort-controller/-/abort-controller-2.1.3.tgz", @@ -7714,6 +7761,21 @@ "resolved": "https://registry.npmjs.org/@types/shimmer/-/shimmer-1.0.5.tgz", "integrity": "sha512-9Hp0ObzwwO57DpLFF0InUjUm/II8GmKAvzbefxQTihCb7KI6yc9yzf0nLc4mVdby5N4DRCgQM2wCup9KTieeww==" }, + "node_modules/@types/sinon": { + "version": "17.0.3", + "resolved": "https://registry.npmjs.org/@types/sinon/-/sinon-17.0.3.tgz", + "integrity": "sha512-j3uovdn8ewky9kRBG19bOwaZbexJu/XjtkHyjvUgt4xfPFz18dcORIMqnYh66Fx3Powhcr85NT5+er3+oViapw==", + "dev": true, + "dependencies": { + "@types/sinonjs__fake-timers": "*" + } + }, + "node_modules/@types/sinonjs__fake-timers": { + "version": "8.1.5", + "resolved": "https://registry.npmjs.org/@types/sinonjs__fake-timers/-/sinonjs__fake-timers-8.1.5.tgz", + "integrity": "sha512-mQkU2jY8jJEF7YHjHvsQO8+3ughTL1mcnn96igfhONmR+fUPSKIkefQYpSe8bsly2Ep7oQbn/6VG5/9/0qcArQ==", + "dev": true + }, "node_modules/@types/tough-cookie": { "version": "4.0.5", "resolved": "https://registry.npmjs.org/@types/tough-cookie/-/tough-cookie-4.0.5.tgz", @@ -13573,6 +13635,12 @@ "safe-buffer": "~5.1.0" } }, + "node_modules/just-extend": { + "version": "6.2.0", + "resolved": "https://registry.npmjs.org/just-extend/-/just-extend-6.2.0.tgz", + "integrity": "sha512-cYofQu2Xpom82S6qD778jBDpwvvy39s1l/hrYij2u9AMdQcGRpaBu6kY4mVhuno5kJVi1DAz4aiphA2WI1/OAw==", + "dev": true + }, "node_modules/jwa": { "version": "2.0.0", "resolved": "https://registry.npmjs.org/jwa/-/jwa-2.0.0.tgz", @@ -14971,6 +15039,12 @@ "integrity": "sha512-FT1yDzDYEoYWhnSGnpE/4Kj1fLZkDFyqRb7fNt6FdYOSxlUWAtp42Eh6Wb0rGIv/m9Bgo7x4GhQbm5Ys4SG5ow==", "dev": true }, + "node_modules/lodash.get": { + "version": "4.4.2", + "resolved": "https://registry.npmjs.org/lodash.get/-/lodash.get-4.4.2.tgz", + "integrity": "sha512-z+Uw/vLuy6gQe8cfaFWD7p0wVv8fJl3mbzXh33RS+0oW2wvUqiRXiQ69gLWSLpgB5/6sU+r6BlQR0MBILadqTQ==", + "dev": true + }, "node_modules/lodash.isfunction": { "version": "3.0.9", "resolved": "https://registry.npmjs.org/lodash.isfunction/-/lodash.isfunction-3.0.9.tgz", @@ -15858,6 +15932,12 @@ "node": ">= 14.0.0" } }, + "node_modules/mocha-param": { + "version": "2.0.1", + "resolved": "https://registry.npmjs.org/mocha-param/-/mocha-param-2.0.1.tgz", + "integrity": "sha512-TDrDAChx9XtkGmRKWGOzMoQefwHsfYUxyjNWgkfAze+EFRIRT28yJVcpcNhw9iWg2NvfeMQZnSwWCNMYwPxZew==", + "dev": true + }, "node_modules/mocha/node_modules/ansi-colors": { "version": "4.1.1", "resolved": "https://registry.npmjs.org/ansi-colors/-/ansi-colors-4.1.1.tgz", @@ -16222,6 +16302,25 @@ "integrity": "sha512-Yd3UES5mWCSqR+qNT93S3UoYUkqAZ9lLg8a7g9rimsWmYGK8cVToA4/sF3RrshdyV3sAGMXVUmpMYOw+dLpOuw==", "devOptional": true }, + "node_modules/nise": { + "version": "5.1.9", + "resolved": "https://registry.npmjs.org/nise/-/nise-5.1.9.tgz", + "integrity": "sha512-qOnoujW4SV6e40dYxJOb3uvuoPHtmLzIk4TFo+j0jPJoC+5Z9xja5qH5JZobEPsa8+YYphMrOSwnrshEhG2qww==", + "dev": true, + "dependencies": { + "@sinonjs/commons": "^3.0.0", + "@sinonjs/fake-timers": "^11.2.2", + "@sinonjs/text-encoding": "^0.7.2", + "just-extend": "^6.2.0", + "path-to-regexp": "^6.2.1" + } + }, + "node_modules/nise/node_modules/path-to-regexp": { + "version": "6.2.1", + "resolved": "https://registry.npmjs.org/path-to-regexp/-/path-to-regexp-6.2.1.tgz", + "integrity": "sha512-JLyh7xT1kizaEvcaXOQwOc2/Yhw6KZOvPf1S8401UyLk86CU79LN3vl7ztXGm/pZ+YjoyAJ4rxmHwbkBXJX+yw==", + "dev": true + }, "node_modules/nocache": { "version": "3.0.4", "resolved": "https://registry.npmjs.org/nocache/-/nocache-3.0.4.tgz", @@ -19826,6 +19925,33 @@ "resolved": "https://registry.npmjs.org/is-arrayish/-/is-arrayish-0.3.2.tgz", "integrity": "sha512-eVRqCvVlZbuw3GrM63ovNSNAeA1K16kaR/LRY/92w0zxQ5/1YzwblUX652i4Xs9RwAGjW9d9y6X88t8OaAJfWQ==" }, + "node_modules/sinon": { + "version": "17.0.1", + "resolved": "https://registry.npmjs.org/sinon/-/sinon-17.0.1.tgz", + "integrity": "sha512-wmwE19Lie0MLT+ZYNpDymasPHUKTaZHUH/pKEubRXIzySv9Atnlw+BUMGCzWgV7b7wO+Hw6f1TEOr0IUnmU8/g==", + "dev": true, + "dependencies": { + "@sinonjs/commons": "^3.0.0", + "@sinonjs/fake-timers": "^11.2.2", + "@sinonjs/samsam": "^8.0.0", + "diff": "^5.1.0", + "nise": "^5.1.5", + "supports-color": "^7.2.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/sinon" + } + }, + "node_modules/sinon/node_modules/diff": { + "version": "5.2.0", + "resolved": "https://registry.npmjs.org/diff/-/diff-5.2.0.tgz", + "integrity": "sha512-uIFDxqpRZGZ6ThOk84hEfqWoHx2devRFvpTZcTHur85vImfaxUbTW9Ryh4CpCuDnToOP1CEtXKIgytHBPVff5A==", + "dev": true, + "engines": { + "node": ">=0.3.1" + } + }, "node_modules/slash": { "version": "3.0.0", "resolved": "https://registry.npmjs.org/slash/-/slash-3.0.0.tgz", @@ -21033,6 +21159,15 @@ "node": ">= 0.8.0" } }, + "node_modules/type-detect": { + "version": "4.0.8", + "resolved": "https://registry.npmjs.org/type-detect/-/type-detect-4.0.8.tgz", + "integrity": "sha512-0fr/mIH1dlO+x7TlcMy+bIDqKPsw/70tVyeHW787goQjhmqaZe10uwLujubK9q9Lg6Fiho1KUKDYz0Z7k7g5/g==", + "dev": true, + "engines": { + "node": ">=4" + } + }, "node_modules/type-fest": { "version": "0.21.3", "resolved": "https://registry.npmjs.org/type-fest/-/type-fest-0.21.3.tgz", @@ -21919,9 +22054,12 @@ "@pollyjs/core": "^6.0.6", "@pollyjs/persister-fs": "^6.0.6", "@types/mocha": "^10.0.6", + "@types/sinon": "^17.0.3", "hnswlib-node": "^1.4.2", "langchain": "^0.1.25", "mocha": "^10.2.0", + "mocha-param": "^2.0.1", + "sinon": "^17.0.1", "ts-mocha": "^10.0.0" }, "engines": { @@ -22875,6 +23013,7 @@ "@google-cloud/aiplatform": "^3.10.0", "@google-cloud/vertexai": "^0.2.1", "@langchain/community": "^0.0.34", + "@opentelemetry/api": "^1.8.0", "@pinecone-database/pinecone": "^2.0.1", "@traceloop/node-server-sdk": "*", "cohere-ai": "^7.7.5", diff --git a/package.json b/package.json index 674c23db..b81330a4 100644 --- a/package.json +++ b/package.json @@ -25,7 +25,7 @@ "eslint-plugin-prettier": "^5.0.1", "husky": "^8.0.3", "lerna": "^8.1.2", - "nx": "^17.2.8", + "nx": "^17.3.2", "prettier": "^3.0.3", "rollup": "^4.13.0", "rollup-plugin-dts": "^6.1.0", @@ -38,5 +38,8 @@ "commitizen": { "path": "./node_modules/cz-conventional-changelog" } + }, + "dependencies": { + "@traceloop/node-server-sdk": "^0.5.25" } } diff --git a/packages/instrumentation-langchain/package.json b/packages/instrumentation-langchain/package.json index b4df37fd..737ca117 100644 --- a/packages/instrumentation-langchain/package.json +++ b/packages/instrumentation-langchain/package.json @@ -49,10 +49,13 @@ "@pollyjs/core": "^6.0.6", "@pollyjs/persister-fs": "^6.0.6", "@types/mocha": "^10.0.6", + "@types/sinon": "^17.0.3", "hnswlib-node": "^1.4.2", "langchain": "^0.1.25", "mocha": "^10.2.0", - "ts-mocha": "^10.0.0" + "ts-mocha": "^10.0.0", + "mocha-param": "^2.0.1", + "sinon": "^17.0.1" }, "homepage": "https://github.com/traceloop/openllmetry-js/tree/main/packages/instrumentation-langchain" } diff --git a/packages/instrumentation-langchain/src/callback-handler.ts b/packages/instrumentation-langchain/src/callback-handler.ts new file mode 100644 index 00000000..be0ef74e --- /dev/null +++ b/packages/instrumentation-langchain/src/callback-handler.ts @@ -0,0 +1,171 @@ +/* + * 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. + */ +import { Serialized } from "@langchain/core/load/serializable"; +import { BaseCallbackHandler, BaseCallbackHandlerInput } from "@langchain/core/callbacks/base"; +import { AgentAction, AgentFinish } from "@langchain/core/agents"; +import { ChainValues } from "@langchain/core/utils/types"; +import { DiagLogger, Tracer, Span, SpanStatus, SpanStatusCode } from "@opentelemetry/api"; +import { LLMResult } from "@langchain/core/outputs"; +import { DocumentInterface } from "@langchain/core/documents"; +import { SpanAttributes, TraceloopSpanKindValues } from "@traceloop/ai-semantic-conventions"; + +// exported for tests +export enum CallbackTrigger { + LLM, Chain, Tool, Agent, Retriever +} + +interface TraceloopCallbackHandlerInput extends BaseCallbackHandlerInput { + tracer: Tracer; + logger: DiagLogger; + shouldSendPrompts: boolean; +} + +export class TraceloopCallbackHandler extends BaseCallbackHandler { + name = "TraceloopCallbackHandler"; + + private readonly ID_SEPARATOR = '.'; + + private readonly tracer: Tracer; + private readonly logger: DiagLogger; + private readonly shouldSendPrompts: boolean; + private readonly activeSpans: Map; + + constructor(input: TraceloopCallbackHandlerInput) { + super(input); + this.tracer = input.tracer; + this.logger = input.logger; + this.shouldSendPrompts = input.shouldSendPrompts; + this.activeSpans = new Map(); + + this.logger.debug(`Built CallbackHandler, it will${this.shouldSendPrompts ? '' : ' not'} send prompts.`); + } + + override handleLLMStart(llm: Serialized, prompts: string[], runId: string, parentRunId?: string, extraParams?: Record, tags?: string[], metadata?: Record, name?: string) { + this.logger.debug(`Inside 'handleLLMStart': {llm: ${JSON.stringify(llm)}, prompts: ${prompts}, runId: ${runId}, parentRunId: ${parentRunId}}`); + + this.pushSpan(CallbackTrigger.LLM, llm.id.join(this.ID_SEPARATOR), TraceloopSpanKindValues.TASK, JSON.stringify(prompts)); + } + + override handleLLMEnd(output: LLMResult, runId: string, parentRunId?: string, tags?: string[]) { + this.logger.debug(`Inside 'handleLLMEnd': {output: ${JSON.stringify(output)}, runId: ${runId}, parentRunId: ${parentRunId}}`); + + this.popSpan(CallbackTrigger.LLM, { code: SpanStatusCode.OK }, JSON.stringify(output)); + } + + override handleLLMError(err: any, runId: string, parentRunId?: string, tags?: string[]) { + this.logger.debug(`Inside 'handleLLMError': {err: ${JSON.stringify(err)}, runId: ${runId}, parentRunId: ${parentRunId}}`); + + this.popSpan(CallbackTrigger.LLM, { code: SpanStatusCode.ERROR, message: JSON.stringify(err) }); + } + + override handleChainStart(chain: Serialized, inputs: ChainValues, runId: string, parentRunId?: string, tags?: string[], metadata?: Record, runType?: string, name?: string) { + this.logger.debug(`Inside 'handleChainStart': {chain: ${JSON.stringify(chain)}, inputs: ${JSON.stringify(inputs)}, runId: ${runId}, parentRunId: ${parentRunId}}`); + + this.pushSpan(CallbackTrigger.Chain, chain.id.join(this.ID_SEPARATOR), TraceloopSpanKindValues.WORKFLOW, JSON.stringify(inputs)); + } + + override handleChainEnd(outputs: ChainValues, runId: string, parentRunId?: string, tags?: string[], kwargs?: { inputs?: Record; }) { + this.logger.debug(`Inside 'handleChainEnd': {outputs: ${JSON.stringify(outputs)}, runId: ${runId}, parentRunId: ${parentRunId}}`); + + this.popSpan(CallbackTrigger.Chain, { code: SpanStatusCode.OK }, JSON.stringify(outputs)); + } + + override handleChainError(err: any, runId: string, parentRunId?: string, tags?: string[], kwargs?: { inputs?: Record; }) { + this.logger.debug(`Inside 'handleChainError': {err: ${JSON.stringify(err)}, runId: ${runId}, parentRunId: ${parentRunId}}`); + + this.popSpan(CallbackTrigger.Chain, { code: SpanStatusCode.ERROR, message: JSON.stringify(err) }); + } + + override handleToolStart(tool: Serialized, input: string, runId: string, parentRunId?: string, tags?: string[], metadata?: Record, name?: string) { + this.logger.debug(`Inside 'handleToolStart': {tool: ${JSON.stringify(tool)}, input: ${JSON.stringify(input)}, runId: ${runId}, parentRunId: ${parentRunId}}`); + + this.pushSpan(CallbackTrigger.Tool, tool.id.join(this.ID_SEPARATOR), TraceloopSpanKindValues.TOOL, input); + } + + override handleToolEnd(output: string, runId: string, parentRunId?: string, tags?: string[]) { + this.logger.debug(`Inside 'handleToolEnd': {output: ${output}, runId: ${runId}, parentRunId: ${parentRunId}}`); + + this.popSpan(CallbackTrigger.Tool, { code: SpanStatusCode.OK }, output); + } + + override handleToolError(err: any, runId: string, parentRunId?: string, tags?: string[]) { + this.logger.debug(`Inside 'handleToolError': {err: ${JSON.stringify(err)}, runId: ${runId}, parentRunId: ${parentRunId}}`); + + this.popSpan(CallbackTrigger.Tool, { code: SpanStatusCode.ERROR, message: JSON.stringify(err) }); + } + + override handleAgentAction(action: AgentAction, runId: string, parentRunId?: string, tags?: string[]): void | Promise { + this.logger.debug(`Inside 'handleAgentAction': {action: ${JSON.stringify(action)}, runId: ${runId}, parentRunId: ${parentRunId}}`); + + this.pushSpan(CallbackTrigger.Agent, action.tool, TraceloopSpanKindValues.AGENT, action.toolInput); + } + + override handleAgentEnd(action: AgentFinish, runId: string, parentRunId?: string | undefined, tags?: string[] | undefined): void | Promise { + this.logger.debug(`Inside 'handleAgentEnd': {action: ${JSON.stringify(action)}, runId: ${runId}, parentRunId: ${parentRunId}}`); + + this.popSpan(CallbackTrigger.Agent, { code: SpanStatusCode.OK }, JSON.stringify(action.returnValues)); + } + + override handleRetrieverStart(retriever: Serialized, query: string, runId: string, parentRunId?: string, tags?: string[], metadata?: Record, name?: string) { + this.logger.debug(`Inside 'handleRetrieverStart': {retriever: ${JSON.stringify(retriever)}, query: ${query}, runId: ${runId}, parentRunId: ${parentRunId}}`); + + this.pushSpan(CallbackTrigger.Retriever, retriever.id.join(this.ID_SEPARATOR), TraceloopSpanKindValues.TASK, query); + } + + override handleRetrieverEnd(documents: DocumentInterface>[], runId: string, parentRunId?: string, tags?: string[]) { + this.logger.debug(`Inside 'handleRetrieverEnd': {documents: ${JSON.stringify(documents)}, runId: ${runId}, parentRunId: ${parentRunId}}`); + + this.popSpan(CallbackTrigger.Retriever, { code: SpanStatusCode.OK }, JSON.stringify(documents)); + } + + override handleRetrieverError(err: any, runId: string, parentRunId?: string, tags?: string[]) { + this.logger.debug(`Inside 'handleRetrieverError': {err: ${JSON.stringify(err)}, runId: ${runId}, parentRunId: ${parentRunId}}`); + + this.popSpan(CallbackTrigger.Retriever, { code: SpanStatusCode.ERROR, message: JSON.stringify(err) }); + } + + private pushSpan(trigger: CallbackTrigger, id: string, spanKind: TraceloopSpanKindValues, inputs?: string): void { + const span = this.tracer.startSpan(id); + span.setAttribute(SpanAttributes.TRACELOOP_SPAN_KIND, spanKind); + + if (this.shouldSendPrompts && !!inputs) { + span.setAttribute(SpanAttributes.TRACELOOP_ENTITY_INPUT, inputs); + } + + if (!this.activeSpans.has(trigger)) { + this.activeSpans.set(trigger, []); + } + + this.activeSpans.get(trigger)!.push(span); + } + + private popSpan(trigger: CallbackTrigger, status: SpanStatus, outputs?: string): void { + if (!this.activeSpans.has(trigger) || this.activeSpans.get(trigger)!.length === 0) { + const message = `Tried to pop span for CallbackTrigger.${CallbackTrigger[trigger]}, but there was none. Current active spans: ${JSON.stringify(this.activeSpans)}`; + this.logger.error(message); + throw new Error(message); + } + + const span = this.activeSpans.get(trigger)!.pop()!; + span.setStatus(status); + + if (this.shouldSendPrompts && !!outputs) { + span.setAttribute(SpanAttributes.TRACELOOP_ENTITY_OUTPUT, outputs); + } + + span.end(); + } +} diff --git a/packages/instrumentation-langchain/src/index.ts b/packages/instrumentation-langchain/src/index.ts index ec28d1d9..e9cabf7d 100644 --- a/packages/instrumentation-langchain/src/index.ts +++ b/packages/instrumentation-langchain/src/index.ts @@ -15,3 +15,4 @@ */ export * from "./instrumentation"; export * from "./types"; +export * from "./callback-handler"; diff --git a/packages/instrumentation-langchain/test/callback-handler.test.ts b/packages/instrumentation-langchain/test/callback-handler.test.ts new file mode 100644 index 00000000..ac20a200 --- /dev/null +++ b/packages/instrumentation-langchain/test/callback-handler.test.ts @@ -0,0 +1,207 @@ +/* + * 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. + */ + +import * as assert from "assert"; +import { stub, resetHistory } from "sinon"; +import itParam from "mocha-param"; + +import { trace, DiagConsoleLogger, SpanStatusCode } from "@opentelemetry/api"; +import { SpanAttributes, TraceloopSpanKindValues } from "@traceloop/ai-semantic-conventions"; + +import { CallbackTrigger, TraceloopCallbackHandler } from "../src"; + +import { name, version } from "../package.json"; + +describe("Test TraceloopCallbackHandler", () => { + + const serialized = { lc: 42, type: undefined, id: [ "unit", "tests" ] }; + const runId = "run-uuid"; + const error = new Error("handcrafted error"); + + const tracer = trace.getTracer(name, version); + const stubSpan = tracer.startSpan("stubSpan"); + + const startSpanStub = stub(tracer, "startSpan").returns(stubSpan); + + const endStub = stub(stubSpan, "end"); + const setAttributeStub = stub(stubSpan, "setAttribute"); + const setStatusStub = stub(stubSpan, "setStatus"); + const callbackHandler = new TraceloopCallbackHandler({ tracer, logger: new DiagConsoleLogger(), shouldSendPrompts: false }); + + interface TestArguments { + startFn: Function; + startFnArgs: any[]; + endFn: Function; + endFnArgs: any[]; + errorFn?: Function; + errorFnArgs?: any[]; + spanKind: TraceloopSpanKindValues; + } + + const TEST_ARGUMENTS_PER_CALLBACK_TRIGGER = { + [CallbackTrigger.LLM]: { + startFn: callbackHandler.handleLLMStart, + startFnArgs: [ serialized, [ "first prompt", "second prompt" ], runId ], + endFn: callbackHandler.handleLLMEnd, + endFnArgs: [ [ [ { text: "output" } ] ], runId ], + errorFn: callbackHandler.handleLLMError, + errorFnArgs: [ error, runId ], + spanKind: TraceloopSpanKindValues.TASK, + }, + [CallbackTrigger.Chain]: { + startFn: callbackHandler.handleChainStart, + startFnArgs: [ serialized, { "key": "value" }, runId ], + endFn: callbackHandler.handleChainEnd, + endFnArgs: [ { "key": "value" }, runId ], + errorFn: callbackHandler.handleChainError, + errorFnArgs: [ error, runId ], + spanKind: TraceloopSpanKindValues.WORKFLOW, + }, + [CallbackTrigger.Tool]: { + startFn: callbackHandler.handleToolStart, + startFnArgs: [ serialized, "input", runId ], + endFn: callbackHandler.handleToolEnd, + endFnArgs: [ "output", runId ], + errorFn: callbackHandler.handleToolError, + errorFnArgs: [ error, runId ], + spanKind: TraceloopSpanKindValues.TOOL, + }, + [CallbackTrigger.Agent]: { + startFn: callbackHandler.handleAgentAction, + startFnArgs: [ { tool: "tool", toolInput: "toolInput", log: "log" }, runId ], + endFn: callbackHandler.handleAgentEnd, + endFnArgs: [ { returnValues: { "key": "value" } }, runId ], + spanKind: TraceloopSpanKindValues.AGENT, + }, + [CallbackTrigger.Retriever]: { + startFn: callbackHandler.handleRetrieverStart, + startFnArgs: [ serialized, "query", runId ], + endFn: callbackHandler.handleRetrieverEnd, + endFnArgs: [ [ { pageContent: "content", metadata: { "key": "value" } } ], runId ], + errorFn: callbackHandler.handleRetrieverError, + errorFnArgs: [ error, runId ], + spanKind: TraceloopSpanKindValues.TASK, + }, + }; + + const TEST_ARGUMENTS_ARRAY = Object + .values(CallbackTrigger) + .filter(trigger => typeof trigger === "number") + .map(trigger => trigger as CallbackTrigger) + .map(trigger => TEST_ARGUMENTS_PER_CALLBACK_TRIGGER[trigger]); + + afterEach(() => { + resetHistory(); + }); + + itParam( + "should handle spans correctly when handle start is called before handle end", + TEST_ARGUMENTS_ARRAY, + ({ startFn, startFnArgs, endFn, endFnArgs, spanKind }: TestArguments) => { + startFn.call(callbackHandler, ...startFnArgs); + assert.ok(startSpanStub.calledOnce); + assert.ok(setAttributeStub.calledOnce); + assert.deepStrictEqual(setAttributeStub.lastCall.firstArg, SpanAttributes.TRACELOOP_SPAN_KIND); + assert.deepStrictEqual(setAttributeStub.lastCall.lastArg, spanKind); + + endFn.call(callbackHandler, ...endFnArgs); + assert.ok(setStatusStub.calledOnce); + assert.deepStrictEqual(setStatusStub.lastCall.firstArg, { code: SpanStatusCode.OK }); + assert.ok(endStub.calledOnce); + } + ); + + itParam( + "should handle spans correctly when there are processing errors", + TEST_ARGUMENTS_ARRAY.filter(({ errorFn }: TestArguments) => !!errorFn), + ({ startFn, startFnArgs, errorFn, errorFnArgs }: TestArguments) => { + startFn.call(callbackHandler, ...startFnArgs); + resetHistory(); + + errorFn!.call(callbackHandler, ...errorFnArgs!); + assert.ok(setStatusStub.calledOnce); + assert.deepStrictEqual(setStatusStub.lastCall.firstArg["code"], SpanStatusCode.ERROR); + assert.ok(endStub.calledOnce); + } + ); + + it("should handle spans correctly on nested handle start and handle end calls", () => { + const nestedCallAmount = 10; + const { startFn, startFnArgs, endFn, endFnArgs, spanKind }: TestArguments = TEST_ARGUMENTS_PER_CALLBACK_TRIGGER[CallbackTrigger.Chain]; + + [...Array(nestedCallAmount).keys()].forEach(_ => { + startFn.call(callbackHandler, ...startFnArgs); + assert.ok(startSpanStub.calledOnce); + assert.ok(setAttributeStub.calledOnce); + assert.deepStrictEqual(setAttributeStub.lastCall.firstArg, SpanAttributes.TRACELOOP_SPAN_KIND); + assert.deepStrictEqual(setAttributeStub.lastCall.lastArg, spanKind); + resetHistory(); + }); + + [...Array(nestedCallAmount).keys()].forEach(_ => { + endFn.call(callbackHandler, ...endFnArgs); + assert.ok(setStatusStub.calledOnce); + assert.deepStrictEqual(setStatusStub.lastCall.firstArg, { code: SpanStatusCode.OK }); + assert.ok(endStub.calledOnce); + resetHistory(); + }); + + assert.throws(() => endFn.call(callbackHandler, ...endFnArgs), Error); + }); + + it( "should send inputs and outputs when shouldSendPrompts is set", () => { + const [ inputIndex, outputIndex ] = [ 1, 0 ]; + const handler = new TraceloopCallbackHandler({ tracer, logger: new DiagConsoleLogger(), shouldSendPrompts: true }); + const { startFnArgs, endFnArgs, spanKind }: TestArguments = TEST_ARGUMENTS_PER_CALLBACK_TRIGGER[CallbackTrigger.Tool]; + + // funtion signature does not allow for spread operator :( + handler.handleToolStart(startFnArgs[0], startFnArgs[1], startFnArgs[2]); + assert.ok(startSpanStub.calledOnce); + assert.ok(setAttributeStub.calledTwice); + assert.deepStrictEqual(setAttributeStub.firstCall.firstArg, SpanAttributes.TRACELOOP_SPAN_KIND); + assert.deepStrictEqual(setAttributeStub.firstCall.lastArg, spanKind); + assert.deepStrictEqual(setAttributeStub.lastCall.firstArg, SpanAttributes.TRACELOOP_ENTITY_INPUT); + assert.deepStrictEqual(setAttributeStub.lastCall.lastArg, startFnArgs[inputIndex]); + resetHistory(); + + // funtion signature does not allow for spread operator :( + handler.handleToolEnd(endFnArgs[0], endFnArgs[1]); + assert.ok(setStatusStub.calledOnce); + assert.deepStrictEqual(setStatusStub.lastCall.firstArg, { code: SpanStatusCode.OK }); + assert.ok(setAttributeStub.calledOnce); + assert.deepStrictEqual(setAttributeStub.firstCall.firstArg, SpanAttributes.TRACELOOP_ENTITY_OUTPUT); + assert.deepStrictEqual(setAttributeStub.firstCall.lastArg, endFnArgs[outputIndex]); + assert.ok(endStub.calledOnce); + }); + + itParam( + "should throw if handle end is called before handle start", + TEST_ARGUMENTS_ARRAY, + ({ endFn, endFnArgs }: TestArguments) => { + assert.throws(() => endFn.call(callbackHandler, ...endFnArgs), Error); + } + ); + + itParam( + "should throw if handle end is called twice after handle start", + TEST_ARGUMENTS_ARRAY, + ({ startFn, startFnArgs, endFn, endFnArgs }: TestArguments) => { + startFn.call(callbackHandler, ...startFnArgs); + endFn.call(callbackHandler, ...endFnArgs); + assert.throws(() => endFn.call(callbackHandler, ...endFnArgs), Error); + } + ); +}); diff --git a/packages/sample-app/package.json b/packages/sample-app/package.json index a9bdadcf..c5f20087 100644 --- a/packages/sample-app/package.json +++ b/packages/sample-app/package.json @@ -35,6 +35,7 @@ "@google-cloud/aiplatform": "^3.10.0", "@google-cloud/vertexai": "^0.2.1", "@langchain/community": "^0.0.34", + "@opentelemetry/api": "^1.8.0", "@pinecone-database/pinecone": "^2.0.1", "@traceloop/node-server-sdk": "*", "cohere-ai": "^7.7.5", diff --git a/packages/sample-app/src/sample_langchain.ts b/packages/sample-app/src/sample_langchain.ts index 410d6c9e..ff0d7fd1 100644 --- a/packages/sample-app/src/sample_langchain.ts +++ b/packages/sample-app/src/sample_langchain.ts @@ -1,80 +1,33 @@ import * as fs from "fs"; import * as traceloop from "@traceloop/node-server-sdk"; +import { TraceloopCallbackHandler } from "@traceloop/instrumentation-langchain"; import { HNSWLib } from "@langchain/community/vectorstores/hnswlib"; -import { OpenAIEmbeddings, OpenAI } from "@langchain/openai"; import { RecursiveCharacterTextSplitter } from "langchain/text_splitter"; import { RetrievalQAChain, loadQAStuffChain } from "langchain/chains"; -import { createOpenAIToolsAgent, AgentExecutor } from "langchain/agents"; -import { SerpAPI } from "@langchain/community/tools/serpapi"; -import { Calculator } from "langchain/tools/calculator"; -import { ChatPromptTemplate, PromptTemplate } from "@langchain/core/prompts"; -import { ChatOpenAI } from "@langchain/openai"; -import { pull } from "langchain/hub"; +import { PromptTemplate } from "@langchain/core/prompts"; +import { OllamaEmbeddings } from "langchain/embeddings/ollama"; +import { Ollama } from "@langchain/community/llms/ollama"; +// import * as ChainsModule from "langchain/chains"; + +import { DiagConsoleLogger, ProxyTracerProvider } from "@opentelemetry/api"; traceloop.initialize({ appName: "sample_langchain", apiKey: process.env.TRACELOOP_API_KEY, disableBatch: true, + logLevel: 'debug', + /*instrumentModules: { langchain: { chainsModule: ChainsModule } }, + traceContent: true,*/ }); class SampleLangchain { - @traceloop.workflow({ name: "sample_retrieval_qa_example" }) - async retrievalQAChainExample() { - // Initialize the LLM to use to answer the question. - const model = new ChatOpenAI({}); - - const text = fs.readFileSync( - "packages/sample-app/data/paul_graham/paul_graham_essay.txt", - "utf8", - ); - const textSplitter = new RecursiveCharacterTextSplitter({ - chunkSize: 1000, - }); - const docs = await textSplitter.createDocuments([text]); - - // Create a vector store from the documents. - const vectorStore = await HNSWLib.fromDocuments( - docs, - new OpenAIEmbeddings(), - ); - - // Initialize a retriever wrapper around the vector store - const vectorStoreRetriever = vectorStore.asRetriever(); - - const chain = RetrievalQAChain.fromLLM(model, vectorStoreRetriever); - const answer = await chain.invoke({ - query: "What did the president say about Justice Breyer?", - }); - - return answer; - } - - @traceloop.workflow({ name: "sample_tools_example" }) - async toolsExample() { - const llm = new ChatOpenAI({}); - const tools = [new Calculator(), new SerpAPI()]; - const prompt = await pull( - "hwchase17/openai-tools-agent", - ); - const agent = await createOpenAIToolsAgent({ llm, tools, prompt }); - const agentExecutor = new AgentExecutor({ - agent, - tools, - }); - const result = await agentExecutor.invoke({ - input: - "By searching the Internet, find how many albums has Boldy James dropped since 2010 and how many albums has Nas dropped since 2010? Find who dropped more albums and show the difference in percent.", - }); - return result; - } - @traceloop.workflow({ name: "sample_qa_stuff_chain" }) async qaStuffChainExample() { - const slowerModel = new OpenAI({ - modelName: "gpt-3.5-turbo-instruct", - temperature: 0.0, + const slowerModel = new Ollama({ + baseUrl: "http://localhost:11434", + model: "llama2", }); const SYSTEM_TEMPLATE = `Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. @@ -82,7 +35,7 @@ class SampleLangchain { {context}`; const prompt = PromptTemplate.fromTemplate(SYSTEM_TEMPLATE); const text = fs.readFileSync( - "packages/sample-app/data/paul_graham/paul_graham_essay.txt", + "data/paul_graham/paul_graham_essay.txt", "utf8", ); const textSplitter = new RecursiveCharacterTextSplitter({ @@ -91,12 +44,22 @@ class SampleLangchain { const docs = await textSplitter.createDocuments([text]); const vectorStore = await HNSWLib.fromDocuments( docs, - new OpenAIEmbeddings(), + new OllamaEmbeddings({ + baseUrl: "http://localhost:11434", + model: "llama2", + }), ); const chain = new RetrievalQAChain({ combineDocumentsChain: loadQAStuffChain(slowerModel, { prompt }), retriever: vectorStore.asRetriever(2), returnSourceDocuments: true, + callbacks: [ + new TraceloopCallbackHandler({ + tracer: new ProxyTracerProvider().getTracer("@traceloop/instrumentation-langchain", "0.5.24"), + logger: new DiagConsoleLogger(), + shouldSendPrompts: true, + }), + ], }); const result = await chain.call({ query: "What did the author do growing up?", @@ -111,11 +74,7 @@ traceloop.withAssociationProperties( async () => { const sampleLangchain = new SampleLangchain(); const retrievalQAChainResult = - await sampleLangchain.retrievalQAChainExample(); + await sampleLangchain.qaStuffChainExample(); console.log(retrievalQAChainResult); - const toolsResult = await sampleLangchain.toolsExample(); - console.log(toolsResult); - const qaStuffChainResult = await sampleLangchain.qaStuffChainExample(); - console.log(qaStuffChainResult); }, );