From b2249d3c52434623e9241d33db79c86987e921cb Mon Sep 17 00:00:00 2001 From: Karthik Kalyanaraman Date: Fri, 22 Nov 2024 11:43:46 -0800 Subject: [PATCH 1/2] chromadb fixes --- package-lock.json | 13 +- package.json | 2 +- src/instrumentation/chroma/instrumentation.ts | 182 ++++++++++++++++-- 3 files changed, 170 insertions(+), 27 deletions(-) diff --git a/package-lock.json b/package-lock.json index 7a921ac2..cd894744 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "@langtrase/typescript-sdk", - "version": "6.1.0", + "version": "6.2.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@langtrase/typescript-sdk", - "version": "6.1.0", + "version": "6.2.0", "license": "Apache-2.0", "dependencies": { "@aws-sdk/client-bedrock-runtime": "^3.670.0", @@ -36,7 +36,7 @@ "@qdrant/js-client-rest": "^1.9.0", "@typescript-eslint/eslint-plugin": "^5.62.0", "ai": "^3.2.29", - "chromadb": "^1.8.1", + "chromadb": "^1.9.4", "cohere-ai": "^7.9.3", "dotenv": "^16.4.5", "eslint": "^8.34.0", @@ -7096,10 +7096,11 @@ "dev": true }, "node_modules/chromadb": { - "version": "1.8.1", - "resolved": "https://registry.npmjs.org/chromadb/-/chromadb-1.8.1.tgz", - "integrity": "sha512-NpbYydbg4Uqt/9BXKgkZXn0fqpsh2Z1yjhkhKH+rcHMoq0pwI18BFSU2QU7Fk/ZypwGefW2AvqyE/3ZJIgy4QA==", + "version": "1.9.4", + "resolved": "https://registry.npmjs.org/chromadb/-/chromadb-1.9.4.tgz", + "integrity": "sha512-KtBy3uvZWV5+B6tlSYxwi8YW+Dv+qBs2spffSaamLeAHLn+N4ss6xoscLWDp5J1ImfscT8NFW1lGZlTZbs8Huw==", "dev": true, + "license": "Apache-2.0", "dependencies": { "cliui": "^8.0.1", "isomorphic-fetch": "^3.0.0" diff --git a/package.json b/package.json index e6b08695..7ac83b36 100644 --- a/package.json +++ b/package.json @@ -48,7 +48,7 @@ "@qdrant/js-client-rest": "^1.9.0", "@typescript-eslint/eslint-plugin": "^5.62.0", "ai": "^3.2.29", - "chromadb": "^1.8.1", + "chromadb": "^1.9.4", "cohere-ai": "^7.9.3", "dotenv": "^16.4.5", "eslint": "^8.34.0", diff --git a/src/instrumentation/chroma/instrumentation.ts b/src/instrumentation/chroma/instrumentation.ts index 014a7061..1c74c999 100644 --- a/src/instrumentation/chroma/instrumentation.ts +++ b/src/instrumentation/chroma/instrumentation.ts @@ -22,6 +22,10 @@ import { version, name } from '../../../package.json' import { APIS } from '@langtrase/trace-attributes' class ChromaInstrumentation extends InstrumentationBase { + private originalGetOrCreateCollection: any + private originalCreateCollection: any + private originalGetCollection: any + constructor () { super(name, version) } @@ -32,18 +36,70 @@ class ChromaInstrumentation extends InstrumentationBase { } init (): Array> { + // eslint-disable-next-line @typescript-eslint/no-this-alias + const self = this const module = new InstrumentationNodeModuleDefinition( 'chromadb', ['>=1.8.1'], (moduleExports, moduleVersion) => { diag.debug(`Patching ChromaDB SDK version ${moduleVersion}`) - this._patch(moduleExports, moduleVersion as string) + + const ChromaClient = moduleExports.ChromaClient + + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (!ChromaClient?.prototype) { + diag.error('ChromaClient not found in exports') + return moduleExports + } + + // Store original methods + this.originalGetOrCreateCollection = ChromaClient.prototype.getOrCreateCollection + this.originalCreateCollection = ChromaClient.prototype.createCollection + this.originalGetCollection = ChromaClient.prototype.getCollection + + // Patch getOrCreateCollection + ChromaClient.prototype.getOrCreateCollection = async function (...args: any[]) { + // eslint-disable-next-line @typescript-eslint/no-unsafe-argument + const collection = await self.originalGetOrCreateCollection.apply(this, args) + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (collection) { + self._patchCollectionMethods(collection, moduleVersion) + } + return collection + } + + // Patch createCollection + ChromaClient.prototype.createCollection = async function (...args: any[]) { + // eslint-disable-next-line @typescript-eslint/no-unsafe-argument + const collection = await self.originalCreateCollection.apply(this, args) + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (collection) { + self._patchCollectionMethods(collection, moduleVersion) + } + return collection + } + + // Patch getCollection + ChromaClient.prototype.getCollection = async function (...args: any[]) { + // eslint-disable-next-line @typescript-eslint/no-unsafe-argument + const collection = await self.originalGetCollection.apply(this, args) + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (collection) { + self._patchCollectionMethods(collection, moduleVersion) + } + return collection + } + return moduleExports }, (moduleExports, moduleVersion) => { diag.debug(`Unpatching ChromaDB SDK version ${moduleVersion}`) - if (moduleExports !== undefined) { - this._unpatch(moduleExports) + const ChromaClient = moduleExports.ChromaClient + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (ChromaClient?.prototype) { + ChromaClient.prototype.getOrCreateCollection = this.originalGetOrCreateCollection + ChromaClient.prototype.createCollection = this.originalCreateCollection + ChromaClient.prototype.getCollection = this.originalGetCollection } } ) @@ -51,30 +107,116 @@ class ChromaInstrumentation extends InstrumentationBase { return [module] } - private _patch (chromadb: any, moduleVersion?: string): void { - if (isWrapped(chromadb.Collection.prototype)) { - Object.keys(APIS.chromadb).forEach((api) => { - this._unwrap(chromadb.Collection.prototype, APIS.chromadb[api as keyof typeof APIS.chromadb].OPERATION as string) - }) - } + private _patchCollectionMethods (collection: any, moduleVersion?: string): void { + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (!collection) return + // eslint-disable-next-line @typescript-eslint/no-this-alias + const self = this Object.keys(APIS.chromadb).forEach((api) => { - if (isWrapped(chromadb.Collection.prototype)) { - this._unpatch(chromadb) + const operation = APIS.chromadb[api as keyof typeof APIS.chromadb].OPERATION as string + const original = collection[operation] + + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (original && !isWrapped(original)) { + // eslint-disable-next-line @typescript-eslint/explicit-function-return-type + const wrapped = (...args: any[]) => { + const patchedMethod = collectionPatch( + // eslint-disable-next-line @typescript-eslint/no-unsafe-argument + function (...methodArgs: any[]) { + return original.apply(collection, methodArgs) + }, + api, + self.tracer, + self.instrumentationVersion, + moduleVersion + ) + return patchedMethod.apply(this, args) + } + + Object.assign(wrapped, original) + wrapped.__original = original + wrapped.__wrapped = true + + collection[operation] = wrapped } - this._wrap( - chromadb.Collection.prototype, - APIS.chromadb[api as keyof typeof APIS.chromadb].OPERATION, - (originalMethod: (...args: any[]) => any) => - collectionPatch(originalMethod, api, this.tracer, this.instrumentationVersion, moduleVersion) - ) }) } + private _patch (chromadb: any, moduleVersion?: string): void { + const ChromaClient = chromadb.ChromaClient + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (ChromaClient?.prototype) { + // eslint-disable-next-line @typescript-eslint/no-this-alias + const self = this + + // Store original methods if not already stored + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (!this.originalGetOrCreateCollection) { + this.originalGetOrCreateCollection = ChromaClient.prototype.getOrCreateCollection + } + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (!this.originalCreateCollection) { + this.originalCreateCollection = ChromaClient.prototype.createCollection + } + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (!this.originalGetCollection) { + this.originalGetCollection = ChromaClient.prototype.getCollection + } + + // Patch getOrCreateCollection + ChromaClient.prototype.getOrCreateCollection = async function (...args: any[]) { + // eslint-disable-next-line @typescript-eslint/no-unsafe-argument + const collection = await self.originalGetOrCreateCollection.apply(this, args) + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (collection) { + self._patchCollectionMethods(collection, moduleVersion) + } + return collection + } + + // Patch createCollection + ChromaClient.prototype.createCollection = async function (...args: any[]) { + // eslint-disable-next-line @typescript-eslint/no-unsafe-argument + const collection = await self.originalCreateCollection.apply(this, args) + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (collection) { + self._patchCollectionMethods(collection, moduleVersion) + } + return collection + } + + // Patch getCollection + ChromaClient.prototype.getCollection = async function (...args: any[]) { + // eslint-disable-next-line @typescript-eslint/no-unsafe-argument + const collection = await self.originalGetCollection.apply(this, args) + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (collection) { + self._patchCollectionMethods(collection, moduleVersion) + } + return collection + } + } + } + private _unpatch (chromadb: any): void { - Object.keys(APIS.chromadb).forEach((api) => { - this._unwrap(chromadb.Collection.prototype, APIS.chromadb[api as keyof typeof APIS.chromadb].OPERATION as string) - }) + const ChromaClient = chromadb.ChromaClient + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (ChromaClient?.prototype) { + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (this.originalGetOrCreateCollection) { + // eslint-disable-next-line @typescript-eslint/no-unsafe-argument + ChromaClient.prototype.getOrCreateCollection = this.originalGetOrCreateCollection + } + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (this.originalCreateCollection) { + ChromaClient.prototype.createCollection = this.originalCreateCollection + } + // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions + if (this.originalGetCollection) { + ChromaClient.prototype.getCollection = this.originalGetCollection + } + } } } From 1dfa9e0f858fa7f134979e8927d506dc1e07b534 Mon Sep 17 00:00:00 2001 From: Karthik Kalyanaraman Date: Fri, 22 Nov 2024 11:49:03 -0800 Subject: [PATCH 2/2] chromadb fixes --- package.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/package.json b/package.json index 7ac83b36..50d4e1a6 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@langtrase/typescript-sdk", - "version": "6.2.0", + "version": "6.3.0", "description": "A typescript SDK for Langtrace", "main": "dist/index.js", "types": "dist/index.d.ts",