From 36884118365da01b2ae871e8fc27cce2c289c310 Mon Sep 17 00:00:00 2001 From: crisjc Date: Tue, 3 Dec 2024 09:42:24 +0000 Subject: [PATCH 1/5] Add caches for Azure CosmosDB --- libs/langchain-azure-cosmosdb/src/caches.ts | 146 ++++++++++++++++++++ 1 file changed, 146 insertions(+) diff --git a/libs/langchain-azure-cosmosdb/src/caches.ts b/libs/langchain-azure-cosmosdb/src/caches.ts index da7619c5ff96..4ee67379ad66 100644 --- a/libs/langchain-azure-cosmosdb/src/caches.ts +++ b/libs/langchain-azure-cosmosdb/src/caches.ts @@ -10,6 +10,12 @@ import { EmbeddingsInterface } from "@langchain/core/embeddings"; import { CosmosClient, CosmosClientOptions } from "@azure/cosmos"; import { DefaultAzureCredential } from "@azure/identity"; import { getEnvironmentVariable } from "@langchain/core/utils/env"; +import { MongoClient } from "mongodb"; +import { + AzureCosmosDBMongoDBConfig, + AzureCosmosDBMongoDBVectorStore, + AzureCosmosDBMongoDBSimilarityType + } from "./azure_cosmosdb_mongodb.js" import { AzureCosmosDBNoSQLConfig, AzureCosmosDBNoSQLVectorStore, @@ -189,3 +195,143 @@ export class AzureCosmosDBNoSQLSemanticCache extends BaseCache { } } } + +/** + * Represents a Semantic Cache that uses CosmosDB NoSQL backend as the underlying + * storage system. + * + * @example + * ```typescript + * const embeddings = new OpenAIEmbeddings(); + * const cache = new AzureCosmosDBNoSQLSemanticCache(embeddings, { + * connectionString: string + * }); + * const model = new ChatOpenAI({cache}); + * + * // Invoke the model to perform an action + * const response = await model.invoke("Do something random!"); + * console.log(response); + * ``` + */ +export class AzureCosmosDBMongoDBSemanticCache extends BaseCache { + private embeddings: EmbeddingsInterface; + + private config: AzureCosmosDBMongoDBConfig; + + private similarityScoreThreshold: number; + + private cacheDict: { [key: string]: AzureCosmosDBMongoDBVectorStore } = {}; + + private readonly client: MongoClient | undefined; + + private vectorDistanceFunction: string; + + constructor( + embeddings: EmbeddingsInterface, + dbConfig: AzureCosmosDBMongoDBConfig, + similarityScoreThreshold: number = 0.6, + ) { + super() + + const connectionString = + dbConfig.connectionString ?? + getEnvironmentVariable("AZURE_COSMOSDB_MONGODB_CONNECTION_STRING"); + + if (!dbConfig.client && !connectionString) { + throw new Error( + "AzureCosmosDBMongoDBVectorStore client or connection string must be set." + ); + } + + if (!dbConfig.client) { + this.client = new MongoClient (connectionString!, { + appName: "langchainjs", + }); + } + + this.config = { + ...dbConfig, + client: this.client + } + this.similarityScoreThreshold = similarityScoreThreshold; + this.embeddings = embeddings; + this.vectorDistanceFunction = + dbConfig?.indexOptions?.similarity ?? AzureCosmosDBMongoDBSimilarityType.COS; + + } + + private getLlmCache(llmKey: string) { + const key = getCacheKey(llmKey) + if (!this.cacheDict[key]) { + this.cacheDict[key] = new AzureCosmosDBMongoDBVectorStore( + this.embeddings, + this.config + ) + } + return this.cacheDict[key] + } + + /** + * Retrieves data from the cache. + * + * @param prompt The prompt for lookup. + * @param llmKey The LLM key used to construct the cache key. + * @returns An array of Generations if found, null otherwise. + */ + async lookup(prompt: string, llmKey: string): Promise { + const llmCache = this.getLlmCache(llmKey) + + const queryEmbedding = await this.embeddings.embedQuery(prompt) + const results = await llmCache.similaritySearchVectorWithScore(queryEmbedding, 1, this.config.indexOptions?.indexType) + if (!results.length) return null; + + const generations = results + .flatMap(([document, score]) => { + const isSimilar = + (this.vectorDistanceFunction === "COS" && + score <= this.similarityScoreThreshold) || + (this.vectorDistanceFunction !== "COS" && + score >= this.similarityScoreThreshold); + + if (!isSimilar) return undefined; + + return document.metadata.return_value.map((gen: string) => + deserializeStoredGeneration(JSON.parse(gen)) + ); + }) + .filter((gen) => gen !== undefined); + + return generations.length > 0 ? generations : null; + } + + /** + * Updates the cache with new data. + * + * @param prompt The prompt for update. + * @param llmKey The LLM key used to construct the cache key. + * @param value The value to be stored in the cache. + */ + public async update(prompt: string, llmKey: string, returnValue: Generation[]): Promise { + const serializedGenerations = returnValue.map((generation) => + JSON.stringify(serializeGeneration(generation)) + ); + const llmCache = this.getLlmCache(llmKey); + const metadata = { + llm_string: llmKey, + prompt, + return_value: serializedGenerations, + }; + const doc = new Document({ + pageContent: prompt, + metadata, + }); + await llmCache.addDocuments([doc]); + } + + public async clear(llmKey: string) { + const key = getCacheKey(llmKey); + if (this.cacheDict[key]) { + await this.cacheDict[key].delete(); + } + } +} \ No newline at end of file From f39c3db62b13a3ed495cea5a36122619ca9ecd65 Mon Sep 17 00:00:00 2001 From: crisjc Date: Wed, 4 Dec 2024 09:27:17 +0000 Subject: [PATCH 2/5] change threshold for similarity --- libs/langchain-azure-cosmosdb/src/caches.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langchain-azure-cosmosdb/src/caches.ts b/libs/langchain-azure-cosmosdb/src/caches.ts index 4ee67379ad66..f6485a58db72 100644 --- a/libs/langchain-azure-cosmosdb/src/caches.ts +++ b/libs/langchain-azure-cosmosdb/src/caches.ts @@ -229,7 +229,7 @@ export class AzureCosmosDBMongoDBSemanticCache extends BaseCache { constructor( embeddings: EmbeddingsInterface, dbConfig: AzureCosmosDBMongoDBConfig, - similarityScoreThreshold: number = 0.6, + similarityScoreThreshold: number = 0.0, ) { super() From f698add16ab56e931198ee52bc1659386cc62a32 Mon Sep 17 00:00:00 2001 From: jacoblee93 Date: Mon, 9 Dec 2024 16:19:57 -0800 Subject: [PATCH 3/5] Format --- libs/langchain-azure-cosmosdb/src/caches.ts | 52 ++++++++++++--------- 1 file changed, 30 insertions(+), 22 deletions(-) diff --git a/libs/langchain-azure-cosmosdb/src/caches.ts b/libs/langchain-azure-cosmosdb/src/caches.ts index f6485a58db72..175da1557c76 100644 --- a/libs/langchain-azure-cosmosdb/src/caches.ts +++ b/libs/langchain-azure-cosmosdb/src/caches.ts @@ -14,8 +14,8 @@ import { MongoClient } from "mongodb"; import { AzureCosmosDBMongoDBConfig, AzureCosmosDBMongoDBVectorStore, - AzureCosmosDBMongoDBSimilarityType - } from "./azure_cosmosdb_mongodb.js" + AzureCosmosDBMongoDBSimilarityType, +} from "./azure_cosmosdb_mongodb.js"; import { AzureCosmosDBNoSQLConfig, AzureCosmosDBNoSQLVectorStore, @@ -215,7 +215,7 @@ export class AzureCosmosDBNoSQLSemanticCache extends BaseCache { */ export class AzureCosmosDBMongoDBSemanticCache extends BaseCache { private embeddings: EmbeddingsInterface; - + private config: AzureCosmosDBMongoDBConfig; private similarityScoreThreshold: number; @@ -229,11 +229,11 @@ export class AzureCosmosDBMongoDBSemanticCache extends BaseCache { constructor( embeddings: EmbeddingsInterface, dbConfig: AzureCosmosDBMongoDBConfig, - similarityScoreThreshold: number = 0.0, + similarityScoreThreshold: number = 0.0 ) { - super() + super(); - const connectionString = + const connectionString = dbConfig.connectionString ?? getEnvironmentVariable("AZURE_COSMOSDB_MONGODB_CONNECTION_STRING"); @@ -244,31 +244,31 @@ export class AzureCosmosDBMongoDBSemanticCache extends BaseCache { } if (!dbConfig.client) { - this.client = new MongoClient (connectionString!, { + this.client = new MongoClient(connectionString!, { appName: "langchainjs", }); } this.config = { ...dbConfig, - client: this.client - } + client: this.client, + }; this.similarityScoreThreshold = similarityScoreThreshold; this.embeddings = embeddings; - this.vectorDistanceFunction = - dbConfig?.indexOptions?.similarity ?? AzureCosmosDBMongoDBSimilarityType.COS; - + this.vectorDistanceFunction = + dbConfig?.indexOptions?.similarity ?? + AzureCosmosDBMongoDBSimilarityType.COS; } private getLlmCache(llmKey: string) { - const key = getCacheKey(llmKey) + const key = getCacheKey(llmKey); if (!this.cacheDict[key]) { - this.cacheDict[key] = new AzureCosmosDBMongoDBVectorStore( + this.cacheDict[key] = new AzureCosmosDBMongoDBVectorStore( this.embeddings, this.config - ) + ); } - return this.cacheDict[key] + return this.cacheDict[key]; } /** @@ -279,10 +279,14 @@ export class AzureCosmosDBMongoDBSemanticCache extends BaseCache { * @returns An array of Generations if found, null otherwise. */ async lookup(prompt: string, llmKey: string): Promise { - const llmCache = this.getLlmCache(llmKey) + const llmCache = this.getLlmCache(llmKey); - const queryEmbedding = await this.embeddings.embedQuery(prompt) - const results = await llmCache.similaritySearchVectorWithScore(queryEmbedding, 1, this.config.indexOptions?.indexType) + const queryEmbedding = await this.embeddings.embedQuery(prompt); + const results = await llmCache.similaritySearchVectorWithScore( + queryEmbedding, + 1, + this.config.indexOptions?.indexType + ); if (!results.length) return null; const generations = results @@ -301,7 +305,7 @@ export class AzureCosmosDBMongoDBSemanticCache extends BaseCache { }) .filter((gen) => gen !== undefined); - return generations.length > 0 ? generations : null; + return generations.length > 0 ? generations : null; } /** @@ -311,7 +315,11 @@ export class AzureCosmosDBMongoDBSemanticCache extends BaseCache { * @param llmKey The LLM key used to construct the cache key. * @param value The value to be stored in the cache. */ - public async update(prompt: string, llmKey: string, returnValue: Generation[]): Promise { + public async update( + prompt: string, + llmKey: string, + returnValue: Generation[] + ): Promise { const serializedGenerations = returnValue.map((generation) => JSON.stringify(serializeGeneration(generation)) ); @@ -334,4 +342,4 @@ export class AzureCosmosDBMongoDBSemanticCache extends BaseCache { await this.cacheDict[key].delete(); } } -} \ No newline at end of file +} From 0bf7cab781e6c59235a801b08930e3283dc11236 Mon Sep 17 00:00:00 2001 From: crisjc Date: Mon, 16 Dec 2024 05:24:36 +0000 Subject: [PATCH 4/5] resolve comments --- libs/langchain-azure-cosmosdb/src/caches.ts | 345 ------------------ .../src/caches/caches_mongodb.ts | 149 ++++++++ .../src/caches/caches_nosql.ts | 192 ++++++++++ libs/langchain-azure-cosmosdb/src/index.ts | 3 +- .../src/tests/caches.int.test.ts | 244 ------------- .../tests/caches/caches_nosql.init.test.ts | 245 +++++++++++++ .../caches_nosql.test.ts} | 2 +- 7 files changed, 589 insertions(+), 591 deletions(-) delete mode 100644 libs/langchain-azure-cosmosdb/src/caches.ts create mode 100644 libs/langchain-azure-cosmosdb/src/caches/caches_mongodb.ts create mode 100644 libs/langchain-azure-cosmosdb/src/caches/caches_nosql.ts delete mode 100644 libs/langchain-azure-cosmosdb/src/tests/caches.int.test.ts create mode 100644 libs/langchain-azure-cosmosdb/src/tests/caches/caches_nosql.init.test.ts rename libs/langchain-azure-cosmosdb/src/tests/{caches.test.ts => caches/caches_nosql.test.ts} (96%) diff --git a/libs/langchain-azure-cosmosdb/src/caches.ts b/libs/langchain-azure-cosmosdb/src/caches.ts deleted file mode 100644 index 175da1557c76..000000000000 --- a/libs/langchain-azure-cosmosdb/src/caches.ts +++ /dev/null @@ -1,345 +0,0 @@ -import { - BaseCache, - deserializeStoredGeneration, - getCacheKey, - serializeGeneration, -} from "@langchain/core/caches"; -import { Generation } from "@langchain/core/outputs"; -import { Document } from "@langchain/core/documents"; -import { EmbeddingsInterface } from "@langchain/core/embeddings"; -import { CosmosClient, CosmosClientOptions } from "@azure/cosmos"; -import { DefaultAzureCredential } from "@azure/identity"; -import { getEnvironmentVariable } from "@langchain/core/utils/env"; -import { MongoClient } from "mongodb"; -import { - AzureCosmosDBMongoDBConfig, - AzureCosmosDBMongoDBVectorStore, - AzureCosmosDBMongoDBSimilarityType, -} from "./azure_cosmosdb_mongodb.js"; -import { - AzureCosmosDBNoSQLConfig, - AzureCosmosDBNoSQLVectorStore, -} from "./azure_cosmosdb_nosql.js"; - -const USER_AGENT_SUFFIX = "langchainjs-cdbnosql-semanticcache-javascript"; -const DEFAULT_CONTAINER_NAME = "semanticCacheContainer"; - -/** - * Represents a Semantic Cache that uses CosmosDB NoSQL backend as the underlying - * storage system. - * - * @example - * ```typescript - * const embeddings = new OpenAIEmbeddings(); - * const cache = new AzureCosmosDBNoSQLSemanticCache(embeddings, { - * databaseName: DATABASE_NAME, - * containerName: CONTAINER_NAME - * }); - * const model = new ChatOpenAI({cache}); - * - * // Invoke the model to perform an action - * const response = await model.invoke("Do something random!"); - * console.log(response); - * ``` - */ -export class AzureCosmosDBNoSQLSemanticCache extends BaseCache { - private embeddings: EmbeddingsInterface; - - private config: AzureCosmosDBNoSQLConfig; - - private similarityScoreThreshold: number; - - private cacheDict: { [key: string]: AzureCosmosDBNoSQLVectorStore } = {}; - - private vectorDistanceFunction: string; - - constructor( - embeddings: EmbeddingsInterface, - dbConfig: AzureCosmosDBNoSQLConfig, - similarityScoreThreshold: number = 0.6 - ) { - super(); - let client: CosmosClient; - - const connectionString = - dbConfig.connectionString ?? - getEnvironmentVariable("AZURE_COSMOSDB_NOSQL_CONNECTION_STRING"); - - const endpoint = - dbConfig.endpoint ?? - getEnvironmentVariable("AZURE_COSMOSDB_NOSQL_ENDPOINT"); - - if (!dbConfig.client && !connectionString && !endpoint) { - throw new Error( - "AzureCosmosDBNoSQLSemanticCache client, connection string or endpoint must be set." - ); - } - - if (!dbConfig.client) { - if (connectionString) { - // eslint-disable-next-line @typescript-eslint/no-non-null-assertion - let [endpoint, key] = connectionString!.split(";"); - [, endpoint] = endpoint.split("="); - [, key] = key.split("="); - - client = new CosmosClient({ - endpoint, - key, - userAgentSuffix: USER_AGENT_SUFFIX, - }); - } else { - // Use managed identity - client = new CosmosClient({ - endpoint, - aadCredentials: dbConfig.credentials ?? new DefaultAzureCredential(), - userAgentSuffix: USER_AGENT_SUFFIX, - } as CosmosClientOptions); - } - } else { - client = dbConfig.client; - } - - this.vectorDistanceFunction = - dbConfig.vectorEmbeddingPolicy?.vectorEmbeddings[0].distanceFunction ?? - "cosine"; - - this.config = { - ...dbConfig, - client, - databaseName: dbConfig.databaseName, - containerName: dbConfig.containerName ?? DEFAULT_CONTAINER_NAME, - }; - this.embeddings = embeddings; - this.similarityScoreThreshold = similarityScoreThreshold; - } - - private getLlmCache(llmKey: string) { - const key = getCacheKey(llmKey); - if (!this.cacheDict[key]) { - this.cacheDict[key] = new AzureCosmosDBNoSQLVectorStore( - this.embeddings, - this.config - ); - } - return this.cacheDict[key]; - } - - /** - * Retrieves data from the cache. - * - * @param prompt The prompt for lookup. - * @param llmKey The LLM key used to construct the cache key. - * @returns An array of Generations if found, null otherwise. - */ - public async lookup(prompt: string, llmKey: string) { - const llmCache = this.getLlmCache(llmKey); - - const results = await llmCache.similaritySearchWithScore(prompt, 1); - if (!results.length) return null; - - const generations = results - .flatMap(([document, score]) => { - const isSimilar = - (this.vectorDistanceFunction === "euclidean" && - score <= this.similarityScoreThreshold) || - (this.vectorDistanceFunction !== "euclidean" && - score >= this.similarityScoreThreshold); - - if (!isSimilar) return undefined; - - return document.metadata.return_value.map((gen: string) => - deserializeStoredGeneration(JSON.parse(gen)) - ); - }) - .filter((gen) => gen !== undefined); - - return generations.length > 0 ? generations : null; - } - - /** - * Updates the cache with new data. - * - * @param prompt The prompt for update. - * @param llmKey The LLM key used to construct the cache key. - * @param value The value to be stored in the cache. - */ - public async update( - prompt: string, - llmKey: string, - returnValue: Generation[] - ) { - const serializedGenerations = returnValue.map((generation) => - JSON.stringify(serializeGeneration(generation)) - ); - const llmCache = this.getLlmCache(llmKey); - const metadata = { - llm_string: llmKey, - prompt, - return_value: serializedGenerations, - }; - const doc = new Document({ - pageContent: prompt, - metadata, - }); - await llmCache.addDocuments([doc]); - } - - /** - * deletes the semantic cache for a given llmKey - * @param llmKey - */ - public async clear(llmKey: string) { - const key = getCacheKey(llmKey); - if (this.cacheDict[key]) { - await this.cacheDict[key].delete(); - } - } -} - -/** - * Represents a Semantic Cache that uses CosmosDB NoSQL backend as the underlying - * storage system. - * - * @example - * ```typescript - * const embeddings = new OpenAIEmbeddings(); - * const cache = new AzureCosmosDBNoSQLSemanticCache(embeddings, { - * connectionString: string - * }); - * const model = new ChatOpenAI({cache}); - * - * // Invoke the model to perform an action - * const response = await model.invoke("Do something random!"); - * console.log(response); - * ``` - */ -export class AzureCosmosDBMongoDBSemanticCache extends BaseCache { - private embeddings: EmbeddingsInterface; - - private config: AzureCosmosDBMongoDBConfig; - - private similarityScoreThreshold: number; - - private cacheDict: { [key: string]: AzureCosmosDBMongoDBVectorStore } = {}; - - private readonly client: MongoClient | undefined; - - private vectorDistanceFunction: string; - - constructor( - embeddings: EmbeddingsInterface, - dbConfig: AzureCosmosDBMongoDBConfig, - similarityScoreThreshold: number = 0.0 - ) { - super(); - - const connectionString = - dbConfig.connectionString ?? - getEnvironmentVariable("AZURE_COSMOSDB_MONGODB_CONNECTION_STRING"); - - if (!dbConfig.client && !connectionString) { - throw new Error( - "AzureCosmosDBMongoDBVectorStore client or connection string must be set." - ); - } - - if (!dbConfig.client) { - this.client = new MongoClient(connectionString!, { - appName: "langchainjs", - }); - } - - this.config = { - ...dbConfig, - client: this.client, - }; - this.similarityScoreThreshold = similarityScoreThreshold; - this.embeddings = embeddings; - this.vectorDistanceFunction = - dbConfig?.indexOptions?.similarity ?? - AzureCosmosDBMongoDBSimilarityType.COS; - } - - private getLlmCache(llmKey: string) { - const key = getCacheKey(llmKey); - if (!this.cacheDict[key]) { - this.cacheDict[key] = new AzureCosmosDBMongoDBVectorStore( - this.embeddings, - this.config - ); - } - return this.cacheDict[key]; - } - - /** - * Retrieves data from the cache. - * - * @param prompt The prompt for lookup. - * @param llmKey The LLM key used to construct the cache key. - * @returns An array of Generations if found, null otherwise. - */ - async lookup(prompt: string, llmKey: string): Promise { - const llmCache = this.getLlmCache(llmKey); - - const queryEmbedding = await this.embeddings.embedQuery(prompt); - const results = await llmCache.similaritySearchVectorWithScore( - queryEmbedding, - 1, - this.config.indexOptions?.indexType - ); - if (!results.length) return null; - - const generations = results - .flatMap(([document, score]) => { - const isSimilar = - (this.vectorDistanceFunction === "COS" && - score <= this.similarityScoreThreshold) || - (this.vectorDistanceFunction !== "COS" && - score >= this.similarityScoreThreshold); - - if (!isSimilar) return undefined; - - return document.metadata.return_value.map((gen: string) => - deserializeStoredGeneration(JSON.parse(gen)) - ); - }) - .filter((gen) => gen !== undefined); - - return generations.length > 0 ? generations : null; - } - - /** - * Updates the cache with new data. - * - * @param prompt The prompt for update. - * @param llmKey The LLM key used to construct the cache key. - * @param value The value to be stored in the cache. - */ - public async update( - prompt: string, - llmKey: string, - returnValue: Generation[] - ): Promise { - const serializedGenerations = returnValue.map((generation) => - JSON.stringify(serializeGeneration(generation)) - ); - const llmCache = this.getLlmCache(llmKey); - const metadata = { - llm_string: llmKey, - prompt, - return_value: serializedGenerations, - }; - const doc = new Document({ - pageContent: prompt, - metadata, - }); - await llmCache.addDocuments([doc]); - } - - public async clear(llmKey: string) { - const key = getCacheKey(llmKey); - if (this.cacheDict[key]) { - await this.cacheDict[key].delete(); - } - } -} diff --git a/libs/langchain-azure-cosmosdb/src/caches/caches_mongodb.ts b/libs/langchain-azure-cosmosdb/src/caches/caches_mongodb.ts new file mode 100644 index 000000000000..4ce56a72ecd0 --- /dev/null +++ b/libs/langchain-azure-cosmosdb/src/caches/caches_mongodb.ts @@ -0,0 +1,149 @@ +import { + BaseCache, + deserializeStoredGeneration, + getCacheKey, + serializeGeneration, +} from "@langchain/core/caches"; +import { Generation } from "@langchain/core/outputs"; +import { Document } from "@langchain/core/documents"; +import { EmbeddingsInterface } from "@langchain/core/embeddings"; +import { getEnvironmentVariable } from "@langchain/core/utils/env"; +import { MongoClient } from "mongodb"; +import { + AzureCosmosDBMongoDBConfig, + AzureCosmosDBMongoDBVectorStore, + AzureCosmosDBMongoDBSimilarityType, +} from "../azure_cosmosdb_mongodb.js"; + + +export class AzureCosmosDBMongoDBSemanticCache extends BaseCache { + private embeddings: EmbeddingsInterface; + + private config: AzureCosmosDBMongoDBConfig; + + private similarityScoreThreshold: number; + + private cacheDict: { [key: string]: AzureCosmosDBMongoDBVectorStore } = {}; + + private readonly client: MongoClient | undefined; + + private vectorDistanceFunction: string; + + constructor( + embeddings: EmbeddingsInterface, + dbConfig: AzureCosmosDBMongoDBConfig, + similarityScoreThreshold: number = 0.6 + ) { + super(); + + const connectionString = + dbConfig.connectionString ?? + getEnvironmentVariable("AZURE_COSMOSDB_MONGODB_CONNECTION_STRING"); + + if (!dbConfig.client && !connectionString) { + throw new Error( + "AzureCosmosDBMongoDBSemanticCache client or connection string must be set." + ); + } + + if (!dbConfig.client) { + this.client = new MongoClient(connectionString!, { + appName: "langchainjs", + }); + } + + this.config = { + ...dbConfig, + client: this.client, + collectionName: dbConfig.collectionName ?? "semanticCacheContainer" + }; + this.similarityScoreThreshold = similarityScoreThreshold; + this.embeddings = embeddings; + this.vectorDistanceFunction = + dbConfig?.indexOptions?.similarity ?? + AzureCosmosDBMongoDBSimilarityType.COS; + } + + private getLlmCache(llmKey: string) { + const key = getCacheKey(llmKey); + if (!this.cacheDict[key]) { + this.cacheDict[key] = new AzureCosmosDBMongoDBVectorStore( + this.embeddings, + this.config + ); + } + return this.cacheDict[key]; + } + + /** + * Retrieves data from the cache. + * + * @param prompt The prompt for lookup. + * @param llmKey The LLM key used to construct the cache key. + * @returns An array of Generations if found, null otherwise. + */ + async lookup(prompt: string, llmKey: string): Promise { + const llmCache = this.getLlmCache(llmKey); + + const queryEmbedding = await this.embeddings.embedQuery(prompt); + const results = await llmCache.similaritySearchVectorWithScore( + queryEmbedding, + 1, + this.config.indexOptions?.indexType + ); + if (!results.length) return null; + + const generations = results + .flatMap(([document, score]) => { + const isSimilar = + (this.vectorDistanceFunction === AzureCosmosDBMongoDBSimilarityType.L2 && + score <= this.similarityScoreThreshold) || + (this.vectorDistanceFunction !== AzureCosmosDBMongoDBSimilarityType.L2 && + score >= this.similarityScoreThreshold); + + if (!isSimilar) return undefined; + + return document.metadata.return_value.map((gen: string) => + deserializeStoredGeneration(JSON.parse(gen)) + ); + }) + .filter((gen) => gen !== undefined); + + return generations.length > 0 ? generations : null; + } + + /** + * Updates the cache with new data. + * + * @param prompt The prompt for update. + * @param llmKey The LLM key used to construct the cache key. + * @param value The value to be stored in the cache. + */ + public async update( + prompt: string, + llmKey: string, + returnValue: Generation[] + ): Promise { + const serializedGenerations = returnValue.map((generation) => + JSON.stringify(serializeGeneration(generation)) + ); + const llmCache = this.getLlmCache(llmKey); + const metadata = { + llm_string: llmKey, + prompt, + return_value: serializedGenerations, + }; + const doc = new Document({ + pageContent: prompt, + metadata, + }); + await llmCache.addDocuments([doc]); + } + + public async clear(llmKey: string) { + const key = getCacheKey(llmKey); + if (this.cacheDict[key]) { + await this.cacheDict[key].delete(); + } + } +} diff --git a/libs/langchain-azure-cosmosdb/src/caches/caches_nosql.ts b/libs/langchain-azure-cosmosdb/src/caches/caches_nosql.ts new file mode 100644 index 000000000000..905eedf2ff17 --- /dev/null +++ b/libs/langchain-azure-cosmosdb/src/caches/caches_nosql.ts @@ -0,0 +1,192 @@ +import { + BaseCache, + deserializeStoredGeneration, + getCacheKey, + serializeGeneration, + } from "@langchain/core/caches"; + import { Generation } from "@langchain/core/outputs"; + import { Document } from "@langchain/core/documents"; + import { EmbeddingsInterface } from "@langchain/core/embeddings"; + import { CosmosClient, CosmosClientOptions } from "@azure/cosmos"; + import { DefaultAzureCredential } from "@azure/identity"; + import { getEnvironmentVariable } from "@langchain/core/utils/env"; + import { + AzureCosmosDBNoSQLConfig, + AzureCosmosDBNoSQLVectorStore, + } from "../azure_cosmosdb_nosql.js"; + + const USER_AGENT_SUFFIX = "langchainjs-cdbnosql-semanticcache-javascript"; + const DEFAULT_CONTAINER_NAME = "semanticCacheContainer"; + + /** + * Represents a Semantic Cache that uses CosmosDB NoSQL backend as the underlying + * storage system. + * + * @example + * ```typescript + * const embeddings = new OpenAIEmbeddings(); + * const cache = new AzureCosmosDBNoSQLSemanticCache(embeddings, { + * databaseName: DATABASE_NAME, + * containerName: CONTAINER_NAME + * }); + * const model = new ChatOpenAI({cache}); + * + * // Invoke the model to perform an action + * const response = await model.invoke("Do something random!"); + * console.log(response); + * ``` + */ + export class AzureCosmosDBNoSQLSemanticCache extends BaseCache { + private embeddings: EmbeddingsInterface; + + private config: AzureCosmosDBNoSQLConfig; + + private similarityScoreThreshold: number; + + private cacheDict: { [key: string]: AzureCosmosDBNoSQLVectorStore } = {}; + + private vectorDistanceFunction: string; + + constructor( + embeddings: EmbeddingsInterface, + dbConfig: AzureCosmosDBNoSQLConfig, + similarityScoreThreshold: number = 0.6 + ) { + super(); + let client: CosmosClient; + + const connectionString = + dbConfig.connectionString ?? + getEnvironmentVariable("AZURE_COSMOSDB_NOSQL_CONNECTION_STRING"); + + const endpoint = + dbConfig.endpoint ?? + getEnvironmentVariable("AZURE_COSMOSDB_NOSQL_ENDPOINT"); + + if (!dbConfig.client && !connectionString && !endpoint) { + throw new Error( + "AzureCosmosDBNoSQLSemanticCache client, connection string or endpoint must be set." + ); + } + + if (!dbConfig.client) { + if (connectionString) { + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + let [endpoint, key] = connectionString!.split(";"); + [, endpoint] = endpoint.split("="); + [, key] = key.split("="); + + client = new CosmosClient({ + endpoint, + key, + userAgentSuffix: USER_AGENT_SUFFIX, + }); + } else { + // Use managed identity + client = new CosmosClient({ + endpoint, + aadCredentials: dbConfig.credentials ?? new DefaultAzureCredential(), + userAgentSuffix: USER_AGENT_SUFFIX, + } as CosmosClientOptions); + } + } else { + client = dbConfig.client; + } + + this.vectorDistanceFunction = + dbConfig.vectorEmbeddingPolicy?.vectorEmbeddings[0].distanceFunction ?? + "cosine"; + + this.config = { + ...dbConfig, + client, + databaseName: dbConfig.databaseName, + containerName: dbConfig.containerName ?? DEFAULT_CONTAINER_NAME, + }; + this.embeddings = embeddings; + this.similarityScoreThreshold = similarityScoreThreshold; + } + + private getLlmCache(llmKey: string) { + const key = getCacheKey(llmKey); + if (!this.cacheDict[key]) { + this.cacheDict[key] = new AzureCosmosDBNoSQLVectorStore( + this.embeddings, + this.config + ); + } + return this.cacheDict[key]; + } + + /** + * Retrieves data from the cache. + * + * @param prompt The prompt for lookup. + * @param llmKey The LLM key used to construct the cache key. + * @returns An array of Generations if found, null otherwise. + */ + public async lookup(prompt: string, llmKey: string) { + const llmCache = this.getLlmCache(llmKey); + + const results = await llmCache.similaritySearchWithScore(prompt, 1); + if (!results.length) return null; + + const generations = results + .flatMap(([document, score]) => { + const isSimilar = + (this.vectorDistanceFunction === "euclidean" && + score <= this.similarityScoreThreshold) || + (this.vectorDistanceFunction !== "euclidean" && + score >= this.similarityScoreThreshold); + + if (!isSimilar) return undefined; + + return document.metadata.return_value.map((gen: string) => + deserializeStoredGeneration(JSON.parse(gen)) + ); + }) + .filter((gen) => gen !== undefined); + + return generations.length > 0 ? generations : null; + } + + /** + * Updates the cache with new data. + * + * @param prompt The prompt for update. + * @param llmKey The LLM key used to construct the cache key. + * @param value The value to be stored in the cache. + */ + public async update( + prompt: string, + llmKey: string, + returnValue: Generation[] + ) { + const serializedGenerations = returnValue.map((generation) => + JSON.stringify(serializeGeneration(generation)) + ); + const llmCache = this.getLlmCache(llmKey); + const metadata = { + llm_string: llmKey, + prompt, + return_value: serializedGenerations, + }; + const doc = new Document({ + pageContent: prompt, + metadata, + }); + await llmCache.addDocuments([doc]); + } + + /** + * deletes the semantic cache for a given llmKey + * @param llmKey + */ + public async clear(llmKey: string) { + const key = getCacheKey(llmKey); + if (this.cacheDict[key]) { + await this.cacheDict[key].delete(); + } + } + } + \ No newline at end of file diff --git a/libs/langchain-azure-cosmosdb/src/index.ts b/libs/langchain-azure-cosmosdb/src/index.ts index c5160397b474..f0f040afb337 100644 --- a/libs/langchain-azure-cosmosdb/src/index.ts +++ b/libs/langchain-azure-cosmosdb/src/index.ts @@ -1,4 +1,5 @@ export * from "./azure_cosmosdb_mongodb.js"; export * from "./azure_cosmosdb_nosql.js"; -export * from "./caches.js"; export * from "./chat_histories.js"; +export * from "./caches/caches_nosql.js" +export * from "./caches/caches_mongodb.js" diff --git a/libs/langchain-azure-cosmosdb/src/tests/caches.int.test.ts b/libs/langchain-azure-cosmosdb/src/tests/caches.int.test.ts deleted file mode 100644 index d6b66ddaac05..000000000000 --- a/libs/langchain-azure-cosmosdb/src/tests/caches.int.test.ts +++ /dev/null @@ -1,244 +0,0 @@ -/* eslint-disable no-process-env */ -/* eslint-disable @typescript-eslint/no-explicit-any */ - -import { - CosmosClient, - IndexingMode, - VectorEmbeddingPolicy, -} from "@azure/cosmos"; -import { DefaultAzureCredential } from "@azure/identity"; -import { ChatOpenAI, OpenAIEmbeddings } from "@langchain/openai"; -import { AzureCosmosDBNoSQLSemanticCache } from "../caches.js"; - -const DATABASE_NAME = "langchainTestCacheDB"; -const CONTAINER_NAME = "testContainer"; - -function indexingPolicy(indexType: any) { - return { - indexingMode: IndexingMode.consistent, - includedPaths: [{ path: "/*" }], - excludedPaths: [{ path: '/"_etag"/?' }], - vectorIndexes: [{ path: "/embedding", type: indexType }], - }; -} - -function vectorEmbeddingPolicy( - distanceFunction: "euclidean" | "cosine" | "dotproduct", - dimension: number -): VectorEmbeddingPolicy { - return { - vectorEmbeddings: [ - { - path: "/embedding", - dataType: "float32", - distanceFunction, - dimensions: dimension, - }, - ], - }; -} - -async function initializeCache( - indexType: any, - distanceFunction: any, - similarityThreshold?: number -): Promise { - let cache: AzureCosmosDBNoSQLSemanticCache; - const embeddingModel = new OpenAIEmbeddings(); - const testEmbedding = await embeddingModel.embedDocuments(["sample text"]); - const dimension = Math.min( - testEmbedding[0].length, - indexType === "flat" ? 505 : 4096 - ); - if (process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING) { - cache = new AzureCosmosDBNoSQLSemanticCache( - new OpenAIEmbeddings(), - { - databaseName: DATABASE_NAME, - containerName: CONTAINER_NAME, - connectionString: process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING, - indexingPolicy: indexingPolicy(indexType), - vectorEmbeddingPolicy: vectorEmbeddingPolicy( - distanceFunction, - dimension - ), - }, - similarityThreshold - ); - } else if (process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT) { - cache = new AzureCosmosDBNoSQLSemanticCache( - new OpenAIEmbeddings(), - { - databaseName: DATABASE_NAME, - containerName: CONTAINER_NAME, - endpoint: process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT, - indexingPolicy: indexingPolicy(indexType), - vectorEmbeddingPolicy: vectorEmbeddingPolicy( - distanceFunction, - dimension - ), - }, - similarityThreshold - ); - } else { - throw new Error( - "Please set the environment variable AZURE_COSMOSDB_NOSQL_CONNECTION_STRING or AZURE_COSMOSDB_NOSQL_ENDPOINT" - ); - } - return cache; -} - -/* - * To run this test, you need have an Azure Cosmos DB for NoSQL instance - * running. You can deploy a free version on Azure Portal without any cost, - * following this guide: - * https://learn.microsoft.com/azure/cosmos-db/nosql/vector-search - * - * You do not need to create a database or collection, it will be created - * automatically by the test. - * - * Once you have the instance running, you need to set the following environment - * variables before running the test: - * - AZURE_COSMOSDB_NOSQL_CONNECTION_STRING or AZURE_COSMOSDB_NOSQL_ENDPOINT - * - AZURE_OPENAI_API_KEY - * - AZURE_OPENAI_API_INSTANCE_NAME - * - AZURE_OPENAI_API_EMBEDDINGS_DEPLOYMENT_NAME - * - AZURE_OPENAI_API_VERSION - */ -describe("Azure CosmosDB NoSQL Semantic Cache", () => { - beforeEach(async () => { - let client: CosmosClient; - - if (process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING) { - client = new CosmosClient( - process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING - ); - } else if (process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT) { - client = new CosmosClient({ - endpoint: process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT, - aadCredentials: new DefaultAzureCredential(), - }); - } else { - throw new Error( - "Please set the environment variable AZURE_COSMOSDB_NOSQL_CONNECTION_STRING or AZURE_COSMOSDB_NOSQL_ENDPOINT" - ); - } - - // Make sure the database does not exists - try { - await client.database(DATABASE_NAME).delete(); - } catch { - // Ignore error if the database does not exist - } - }); - - it("test AzureCosmosDBNoSqlSemanticCache with cosine quantizedFlat", async () => { - const cache = await initializeCache("quantizedFlat", "cosine"); - const model = new ChatOpenAI({ cache }); - const llmString = JSON.stringify(model._identifyingParams); - await cache.update("foo", llmString, [{ text: "fizz" }]); - - let cacheOutput = await cache.lookup("foo", llmString); - expect(cacheOutput).toEqual([{ text: "fizz" }]); - - cacheOutput = await cache.lookup("bar", llmString); - expect(cacheOutput).toEqual(null); - - await cache.clear(llmString); - }); - - it("test AzureCosmosDBNoSqlSemanticCache with cosine flat", async () => { - const cache = await initializeCache("flat", "cosine"); - const model = new ChatOpenAI({ cache }); - const llmString = JSON.stringify(model._identifyingParams); - await cache.update("foo", llmString, [{ text: "fizz" }]); - - let cacheOutput = await cache.lookup("foo", llmString); - expect(cacheOutput).toEqual([{ text: "fizz" }]); - - cacheOutput = await cache.lookup("bar", llmString); - expect(cacheOutput).toEqual(null); - - await cache.clear(llmString); - }); - - it("test AzureCosmosDBNoSqlSemanticCache with dotProduct quantizedFlat", async () => { - const cache = await initializeCache("quantizedFlat", "dotproduct"); - const model = new ChatOpenAI({ cache }); - const llmString = JSON.stringify(model._identifyingParams); - await cache.update("foo", llmString, [{ text: "fizz" }]); - - let cacheOutput = await cache.lookup("foo", llmString); - expect(cacheOutput).toEqual([{ text: "fizz" }]); - - cacheOutput = await cache.lookup("bar", llmString); - expect(cacheOutput).toEqual(null); - - await cache.clear(llmString); - }); - - it("test AzureCosmosDBNoSqlSemanticCache with dotProduct flat", async () => { - const cache = await initializeCache("flat", "cosine"); - const model = new ChatOpenAI({ cache }); - const llmString = JSON.stringify(model._identifyingParams); - await cache.update("foo", llmString, [{ text: "fizz" }]); - - let cacheOutput = await cache.lookup("foo", llmString); - expect(cacheOutput).toEqual([{ text: "fizz" }]); - - cacheOutput = await cache.lookup("bar", llmString); - expect(cacheOutput).toEqual(null); - - await cache.clear(llmString); - }); - - it("test AzureCosmosDBNoSqlSemanticCache with euclidean quantizedFlat", async () => { - const cache = await initializeCache("quantizedFlat", "euclidean"); - const model = new ChatOpenAI({ cache }); - const llmString = JSON.stringify(model._identifyingParams); - await cache.update("foo", llmString, [{ text: "fizz" }]); - - let cacheOutput = await cache.lookup("foo", llmString); - expect(cacheOutput).toEqual([{ text: "fizz" }]); - - cacheOutput = await cache.lookup("bar", llmString); - expect(cacheOutput).toEqual(null); - - await cache.clear(llmString); - }); - - it("test AzureCosmosDBNoSqlSemanticCache with euclidean flat", async () => { - const cache = await initializeCache("flat", "euclidean"); - const model = new ChatOpenAI({ cache }); - const llmString = JSON.stringify(model._identifyingParams); - await cache.update("foo", llmString, [{ text: "fizz" }]); - - let cacheOutput = await cache.lookup("foo", llmString); - expect(cacheOutput).toEqual([{ text: "fizz" }]); - - cacheOutput = await cache.lookup("bar", llmString); - expect(cacheOutput).toEqual(null); - - await cache.clear(llmString); - }); - - it("test AzureCosmosDBNoSqlSemanticCache response according to similarity score", async () => { - const cache = await initializeCache("quantizedFlat", "cosine"); - const model = new ChatOpenAI({ cache }); - const response1 = await model.invoke( - "Where is the headquarter of Microsoft?" - ); - console.log(response1.content); - // gives similarity score of 0.56 which is less than the threshold of 0.6. The cache - // will retun null which will allow the model to generate result. - const response2 = await model.invoke( - "List all Microsoft offices in India." - ); - expect(response2.content).not.toEqual(response1.content); - console.log(response2.content); - // gives similarity score of .63 > 0.6 - const response3 = await model.invoke("Tell me something about Microsoft"); - expect(response3.content).toEqual(response1.content); - console.log(response3.content); - }); -}); diff --git a/libs/langchain-azure-cosmosdb/src/tests/caches/caches_nosql.init.test.ts b/libs/langchain-azure-cosmosdb/src/tests/caches/caches_nosql.init.test.ts new file mode 100644 index 000000000000..88bab773e4fc --- /dev/null +++ b/libs/langchain-azure-cosmosdb/src/tests/caches/caches_nosql.init.test.ts @@ -0,0 +1,245 @@ +/* eslint-disable no-process-env */ +/* eslint-disable @typescript-eslint/no-explicit-any */ + +import { + CosmosClient, + IndexingMode, + VectorEmbeddingPolicy, + } from "@azure/cosmos"; + import { DefaultAzureCredential } from "@azure/identity"; + import { ChatOpenAI, OpenAIEmbeddings } from "@langchain/openai"; + import { AzureCosmosDBNoSQLSemanticCache } from "../../caches/caches_nosql.js" + + const DATABASE_NAME = "langchainTestCacheDB"; + const CONTAINER_NAME = "testContainer"; + + function indexingPolicy(indexType: any) { + return { + indexingMode: IndexingMode.consistent, + includedPaths: [{ path: "/*" }], + excludedPaths: [{ path: '/"_etag"/?' }], + vectorIndexes: [{ path: "/embedding", type: indexType }], + }; + } + + function vectorEmbeddingPolicy( + distanceFunction: "euclidean" | "cosine" | "dotproduct", + dimension: number + ): VectorEmbeddingPolicy { + return { + vectorEmbeddings: [ + { + path: "/embedding", + dataType: "float32", + distanceFunction, + dimensions: dimension, + }, + ], + }; + } + + async function initializeCache( + indexType: any, + distanceFunction: any, + similarityThreshold?: number + ): Promise { + let cache: AzureCosmosDBNoSQLSemanticCache; + const embeddingModel = new OpenAIEmbeddings(); + const testEmbedding = await embeddingModel.embedDocuments(["sample text"]); + const dimension = Math.min( + testEmbedding[0].length, + indexType === "flat" ? 505 : 4096 + ); + if (process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING) { + cache = new AzureCosmosDBNoSQLSemanticCache( + new OpenAIEmbeddings(), + { + databaseName: DATABASE_NAME, + containerName: CONTAINER_NAME, + connectionString: process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING, + indexingPolicy: indexingPolicy(indexType), + vectorEmbeddingPolicy: vectorEmbeddingPolicy( + distanceFunction, + dimension + ), + }, + similarityThreshold + ); + } else if (process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT) { + cache = new AzureCosmosDBNoSQLSemanticCache( + new OpenAIEmbeddings(), + { + databaseName: DATABASE_NAME, + containerName: CONTAINER_NAME, + endpoint: process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT, + indexingPolicy: indexingPolicy(indexType), + vectorEmbeddingPolicy: vectorEmbeddingPolicy( + distanceFunction, + dimension + ), + }, + similarityThreshold + ); + } else { + throw new Error( + "Please set the environment variable AZURE_COSMOSDB_NOSQL_CONNECTION_STRING or AZURE_COSMOSDB_NOSQL_ENDPOINT" + ); + } + return cache; + } + + /* + * To run this test, you need have an Azure Cosmos DB for NoSQL instance + * running. You can deploy a free version on Azure Portal without any cost, + * following this guide: + * https://learn.microsoft.com/azure/cosmos-db/nosql/vector-search + * + * You do not need to create a database or collection, it will be created + * automatically by the test. + * + * Once you have the instance running, you need to set the following environment + * variables before running the test: + * - AZURE_COSMOSDB_NOSQL_CONNECTION_STRING or AZURE_COSMOSDB_NOSQL_ENDPOINT + * - AZURE_OPENAI_API_KEY + * - AZURE_OPENAI_API_INSTANCE_NAME + * - AZURE_OPENAI_API_EMBEDDINGS_DEPLOYMENT_NAME + * - AZURE_OPENAI_API_VERSION + */ + describe("Azure CosmosDB NoSQL Semantic Cache", () => { + beforeEach(async () => { + let client: CosmosClient; + + if (process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING) { + client = new CosmosClient( + process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING + ); + } else if (process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT) { + client = new CosmosClient({ + endpoint: process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT, + aadCredentials: new DefaultAzureCredential(), + }); + } else { + throw new Error( + "Please set the environment variable AZURE_COSMOSDB_NOSQL_CONNECTION_STRING or AZURE_COSMOSDB_NOSQL_ENDPOINT" + ); + } + + // Make sure the database does not exists + try { + await client.database(DATABASE_NAME).delete(); + } catch { + // Ignore error if the database does not exist + } + }); + + it("test AzureCosmosDBNoSqlSemanticCache with cosine quantizedFlat", async () => { + const cache = await initializeCache("quantizedFlat", "cosine"); + const model = new ChatOpenAI({ cache }); + const llmString = JSON.stringify(model._identifyingParams); + await cache.update("foo", llmString, [{ text: "fizz" }]); + + let cacheOutput = await cache.lookup("foo", llmString); + expect(cacheOutput).toEqual([{ text: "fizz" }]); + + cacheOutput = await cache.lookup("bar", llmString); + expect(cacheOutput).toEqual(null); + + await cache.clear(llmString); + }); + + it("test AzureCosmosDBNoSqlSemanticCache with cosine flat", async () => { + const cache = await initializeCache("flat", "cosine"); + const model = new ChatOpenAI({ cache }); + const llmString = JSON.stringify(model._identifyingParams); + await cache.update("foo", llmString, [{ text: "fizz" }]); + + let cacheOutput = await cache.lookup("foo", llmString); + expect(cacheOutput).toEqual([{ text: "fizz" }]); + + cacheOutput = await cache.lookup("bar", llmString); + expect(cacheOutput).toEqual(null); + + await cache.clear(llmString); + }); + + it("test AzureCosmosDBNoSqlSemanticCache with dotProduct quantizedFlat", async () => { + const cache = await initializeCache("quantizedFlat", "dotproduct"); + const model = new ChatOpenAI({ cache }); + const llmString = JSON.stringify(model._identifyingParams); + await cache.update("foo", llmString, [{ text: "fizz" }]); + + let cacheOutput = await cache.lookup("foo", llmString); + expect(cacheOutput).toEqual([{ text: "fizz" }]); + + cacheOutput = await cache.lookup("bar", llmString); + expect(cacheOutput).toEqual(null); + + await cache.clear(llmString); + }); + + it("test AzureCosmosDBNoSqlSemanticCache with dotProduct flat", async () => { + const cache = await initializeCache("flat", "cosine"); + const model = new ChatOpenAI({ cache }); + const llmString = JSON.stringify(model._identifyingParams); + await cache.update("foo", llmString, [{ text: "fizz" }]); + + let cacheOutput = await cache.lookup("foo", llmString); + expect(cacheOutput).toEqual([{ text: "fizz" }]); + + cacheOutput = await cache.lookup("bar", llmString); + expect(cacheOutput).toEqual(null); + + await cache.clear(llmString); + }); + + it("test AzureCosmosDBNoSqlSemanticCache with euclidean quantizedFlat", async () => { + const cache = await initializeCache("quantizedFlat", "euclidean"); + const model = new ChatOpenAI({ cache }); + const llmString = JSON.stringify(model._identifyingParams); + await cache.update("foo", llmString, [{ text: "fizz" }]); + + let cacheOutput = await cache.lookup("foo", llmString); + expect(cacheOutput).toEqual([{ text: "fizz" }]); + + cacheOutput = await cache.lookup("bar", llmString); + expect(cacheOutput).toEqual(null); + + await cache.clear(llmString); + }); + + it("test AzureCosmosDBNoSqlSemanticCache with euclidean flat", async () => { + const cache = await initializeCache("flat", "euclidean"); + const model = new ChatOpenAI({ cache }); + const llmString = JSON.stringify(model._identifyingParams); + await cache.update("foo", llmString, [{ text: "fizz" }]); + + let cacheOutput = await cache.lookup("foo", llmString); + expect(cacheOutput).toEqual([{ text: "fizz" }]); + + cacheOutput = await cache.lookup("bar", llmString); + expect(cacheOutput).toEqual(null); + + await cache.clear(llmString); + }); + + it("test AzureCosmosDBNoSqlSemanticCache response according to similarity score", async () => { + const cache = await initializeCache("quantizedFlat", "cosine"); + const model = new ChatOpenAI({ cache }); + const response1 = await model.invoke( + "Where is the headquarter of Microsoft?" + ); + console.log(response1.content); + // gives similarity score of 0.56 which is less than the threshold of 0.6. The cache + // will retun null which will allow the model to generate result. + const response2 = await model.invoke( + "List all Microsoft offices in India." + ); + expect(response2.content).not.toEqual(response1.content); + console.log(response2.content); + // gives similarity score of .63 > 0.6 + const response3 = await model.invoke("Tell me something about Microsoft"); + expect(response3.content).toEqual(response1.content); + console.log(response3.content); + }); + }); + \ No newline at end of file diff --git a/libs/langchain-azure-cosmosdb/src/tests/caches.test.ts b/libs/langchain-azure-cosmosdb/src/tests/caches/caches_nosql.test.ts similarity index 96% rename from libs/langchain-azure-cosmosdb/src/tests/caches.test.ts rename to libs/langchain-azure-cosmosdb/src/tests/caches/caches_nosql.test.ts index 9de3f507acc0..3a7a253f22bc 100644 --- a/libs/langchain-azure-cosmosdb/src/tests/caches.test.ts +++ b/libs/langchain-azure-cosmosdb/src/tests/caches/caches_nosql.test.ts @@ -1,7 +1,7 @@ /* eslint-disable @typescript-eslint/no-explicit-any */ import { jest } from "@jest/globals"; import { FakeEmbeddings, FakeLLM } from "@langchain/core/utils/testing"; -import { AzureCosmosDBNoSQLSemanticCache } from "../index.js"; +import { AzureCosmosDBNoSQLSemanticCache } from "../../index.js"; // Create the mock Cosmos DB client const createMockClient = () => { From 262623e8186554881bf72c62e30d79c7773b11ab Mon Sep 17 00:00:00 2001 From: crisjc Date: Wed, 18 Dec 2024 10:34:59 +0000 Subject: [PATCH 5/5] add mongodb cache test and resolve comment --- .../src/caches/caches_mongodb.ts | 273 +++++++------ .../src/caches/caches_nosql.ts | 367 +++++++++--------- libs/langchain-azure-cosmosdb/src/index.ts | 4 +- .../tests/caches/caches_mongodb.int.test.ts | 136 +++++++ .../src/tests/caches/caches_mongodb.test.ts | 72 ++++ .../tests/caches/caches_nosql.init.test.ts | 245 ------------ .../src/tests/caches/caches_nosql.int.test.ts | 244 ++++++++++++ 7 files changed, 788 insertions(+), 553 deletions(-) create mode 100644 libs/langchain-azure-cosmosdb/src/tests/caches/caches_mongodb.int.test.ts create mode 100644 libs/langchain-azure-cosmosdb/src/tests/caches/caches_mongodb.test.ts delete mode 100644 libs/langchain-azure-cosmosdb/src/tests/caches/caches_nosql.init.test.ts create mode 100644 libs/langchain-azure-cosmosdb/src/tests/caches/caches_nosql.int.test.ts diff --git a/libs/langchain-azure-cosmosdb/src/caches/caches_mongodb.ts b/libs/langchain-azure-cosmosdb/src/caches/caches_mongodb.ts index 4ce56a72ecd0..b33589a4095d 100644 --- a/libs/langchain-azure-cosmosdb/src/caches/caches_mongodb.ts +++ b/libs/langchain-azure-cosmosdb/src/caches/caches_mongodb.ts @@ -1,8 +1,8 @@ import { - BaseCache, - deserializeStoredGeneration, - getCacheKey, - serializeGeneration, + BaseCache, + deserializeStoredGeneration, + getCacheKey, + serializeGeneration, } from "@langchain/core/caches"; import { Generation } from "@langchain/core/outputs"; import { Document } from "@langchain/core/documents"; @@ -10,140 +10,169 @@ import { EmbeddingsInterface } from "@langchain/core/embeddings"; import { getEnvironmentVariable } from "@langchain/core/utils/env"; import { MongoClient } from "mongodb"; import { - AzureCosmosDBMongoDBConfig, - AzureCosmosDBMongoDBVectorStore, - AzureCosmosDBMongoDBSimilarityType, + AzureCosmosDBMongoDBConfig, + AzureCosmosDBMongoDBVectorStore, + AzureCosmosDBMongoDBSimilarityType, } from "../azure_cosmosdb_mongodb.js"; - +/** + * Represents a Semantic Cache that uses CosmosDB MongoDB backend as the underlying + * storage system. + * + * @example + * ```typescript + * const embeddings = new OpenAIEmbeddings(); + * const cache = new AzureCosmosDBMongoDBSemanticCache(embeddings, { + * client?: MongoClient + * }); + * const model = new ChatOpenAI({cache}); + * + * // Invoke the model to perform an action + * const response = await model.invoke("Do something random!"); + * console.log(response); + * ``` + */ export class AzureCosmosDBMongoDBSemanticCache extends BaseCache { - private embeddings: EmbeddingsInterface; + private embeddings: EmbeddingsInterface; - private config: AzureCosmosDBMongoDBConfig; + private config: AzureCosmosDBMongoDBConfig; - private similarityScoreThreshold: number; + private similarityScoreThreshold: number; - private cacheDict: { [key: string]: AzureCosmosDBMongoDBVectorStore } = {}; + private cacheDict: { [key: string]: AzureCosmosDBMongoDBVectorStore } = {}; - private readonly client: MongoClient | undefined; + private readonly client: MongoClient | undefined; - private vectorDistanceFunction: string; + private vectorDistanceFunction: string; - constructor( - embeddings: EmbeddingsInterface, - dbConfig: AzureCosmosDBMongoDBConfig, - similarityScoreThreshold: number = 0.6 - ) { - super(); + constructor( + embeddings: EmbeddingsInterface, + dbConfig: AzureCosmosDBMongoDBConfig, + similarityScoreThreshold: number = 0.6 + ) { + super(); - const connectionString = - dbConfig.connectionString ?? - getEnvironmentVariable("AZURE_COSMOSDB_MONGODB_CONNECTION_STRING"); + const connectionString = + dbConfig.connectionString ?? + getEnvironmentVariable("AZURE_COSMOSDB_MONGODB_CONNECTION_STRING"); - if (!dbConfig.client && !connectionString) { - throw new Error( - "AzureCosmosDBMongoDBSemanticCache client or connection string must be set." - ); - } - - if (!dbConfig.client) { - this.client = new MongoClient(connectionString!, { - appName: "langchainjs", - }); - } - - this.config = { - ...dbConfig, - client: this.client, - collectionName: dbConfig.collectionName ?? "semanticCacheContainer" - }; - this.similarityScoreThreshold = similarityScoreThreshold; - this.embeddings = embeddings; - this.vectorDistanceFunction = - dbConfig?.indexOptions?.similarity ?? - AzureCosmosDBMongoDBSimilarityType.COS; + if (!dbConfig.client && !connectionString) { + throw new Error( + "AzureCosmosDBMongoDBSemanticCache client or connection string must be set." + ); } - private getLlmCache(llmKey: string) { - const key = getCacheKey(llmKey); - if (!this.cacheDict[key]) { - this.cacheDict[key] = new AzureCosmosDBMongoDBVectorStore( - this.embeddings, - this.config - ); - } - return this.cacheDict[key]; + if (!dbConfig.client) { + this.client = new MongoClient(connectionString!, { + appName: "langchainjs", + }); + } else { + this.client = dbConfig.client; } - /** - * Retrieves data from the cache. - * - * @param prompt The prompt for lookup. - * @param llmKey The LLM key used to construct the cache key. - * @returns An array of Generations if found, null otherwise. - */ - async lookup(prompt: string, llmKey: string): Promise { - const llmCache = this.getLlmCache(llmKey); - - const queryEmbedding = await this.embeddings.embedQuery(prompt); - const results = await llmCache.similaritySearchVectorWithScore( - queryEmbedding, - 1, - this.config.indexOptions?.indexType - ); - if (!results.length) return null; - - const generations = results - .flatMap(([document, score]) => { - const isSimilar = - (this.vectorDistanceFunction === AzureCosmosDBMongoDBSimilarityType.L2 && - score <= this.similarityScoreThreshold) || - (this.vectorDistanceFunction !== AzureCosmosDBMongoDBSimilarityType.L2 && - score >= this.similarityScoreThreshold); - - if (!isSimilar) return undefined; - - return document.metadata.return_value.map((gen: string) => - deserializeStoredGeneration(JSON.parse(gen)) - ); - }) - .filter((gen) => gen !== undefined); - - return generations.length > 0 ? generations : null; + this.config = { + ...dbConfig, + client: this.client, + collectionName: dbConfig.collectionName ?? "semanticCacheContainer", + }; + + this.similarityScoreThreshold = similarityScoreThreshold; + this.embeddings = embeddings; + this.vectorDistanceFunction = + dbConfig?.indexOptions?.similarity ?? + AzureCosmosDBMongoDBSimilarityType.COS; + } + + private getLlmCache(llmKey: string) { + const key = getCacheKey(llmKey); + if (!this.cacheDict[key]) { + this.cacheDict[key] = new AzureCosmosDBMongoDBVectorStore( + this.embeddings, + this.config + ); } - - /** - * Updates the cache with new data. - * - * @param prompt The prompt for update. - * @param llmKey The LLM key used to construct the cache key. - * @param value The value to be stored in the cache. - */ - public async update( - prompt: string, - llmKey: string, - returnValue: Generation[] - ): Promise { - const serializedGenerations = returnValue.map((generation) => - JSON.stringify(serializeGeneration(generation)) + return this.cacheDict[key]; + } + + /** + * Retrieves data from the cache. + * + * @param prompt The prompt for lookup. + * @param llmKey The LLM key used to construct the cache key. + * @returns An array of Generations if found, null otherwise. + */ + async lookup(prompt: string, llmKey: string): Promise { + const llmCache = this.getLlmCache(llmKey); + + const queryEmbedding = await this.embeddings.embedQuery(prompt); + const results = await llmCache.similaritySearchVectorWithScore( + queryEmbedding, + 1, + this.config.indexOptions?.indexType + ); + if (!results.length) return null; + + const generations = results + .flatMap(([document, score]) => { + const isSimilar = + (this.vectorDistanceFunction === + AzureCosmosDBMongoDBSimilarityType.L2 && + score <= this.similarityScoreThreshold) || + (this.vectorDistanceFunction !== + AzureCosmosDBMongoDBSimilarityType.L2 && + score >= this.similarityScoreThreshold); + + if (!isSimilar) return undefined; + + return document.metadata.return_value.map((gen: string) => + deserializeStoredGeneration(JSON.parse(gen)) ); - const llmCache = this.getLlmCache(llmKey); - const metadata = { - llm_string: llmKey, - prompt, - return_value: serializedGenerations, - }; - const doc = new Document({ - pageContent: prompt, - metadata, - }); - await llmCache.addDocuments([doc]); - } - - public async clear(llmKey: string) { - const key = getCacheKey(llmKey); - if (this.cacheDict[key]) { - await this.cacheDict[key].delete(); - } + }) + .filter((gen) => gen !== undefined); + + return generations.length > 0 ? generations : null; + } + + /** + * Updates the cache with new data. + * + * @param prompt The prompt for update. + * @param llmKey The LLM key used to construct the cache key. + * @param value The value to be stored in the cache. + */ + public async update( + prompt: string, + llmKey: string, + returnValue: Generation[] + ): Promise { + const serializedGenerations = returnValue.map((generation) => + JSON.stringify(serializeGeneration(generation)) + ); + + const llmCache = this.getLlmCache(llmKey); + + const metadata = { + llm_string: llmKey, + prompt, + return_value: serializedGenerations, + }; + + const doc = new Document({ + pageContent: prompt, + metadata, + }); + + await llmCache.addDocuments([doc]); + } + + /** + * deletes the semantic cache for a given llmKey + * @param llmKey + */ + public async clear(llmKey: string) { + const key = getCacheKey(llmKey); + if (this.cacheDict[key]) { + await this.cacheDict[key].delete(); } + } } diff --git a/libs/langchain-azure-cosmosdb/src/caches/caches_nosql.ts b/libs/langchain-azure-cosmosdb/src/caches/caches_nosql.ts index 905eedf2ff17..e110f848702b 100644 --- a/libs/langchain-azure-cosmosdb/src/caches/caches_nosql.ts +++ b/libs/langchain-azure-cosmosdb/src/caches/caches_nosql.ts @@ -1,192 +1,191 @@ import { - BaseCache, - deserializeStoredGeneration, - getCacheKey, - serializeGeneration, - } from "@langchain/core/caches"; - import { Generation } from "@langchain/core/outputs"; - import { Document } from "@langchain/core/documents"; - import { EmbeddingsInterface } from "@langchain/core/embeddings"; - import { CosmosClient, CosmosClientOptions } from "@azure/cosmos"; - import { DefaultAzureCredential } from "@azure/identity"; - import { getEnvironmentVariable } from "@langchain/core/utils/env"; - import { - AzureCosmosDBNoSQLConfig, - AzureCosmosDBNoSQLVectorStore, - } from "../azure_cosmosdb_nosql.js"; - - const USER_AGENT_SUFFIX = "langchainjs-cdbnosql-semanticcache-javascript"; - const DEFAULT_CONTAINER_NAME = "semanticCacheContainer"; - - /** - * Represents a Semantic Cache that uses CosmosDB NoSQL backend as the underlying - * storage system. - * - * @example - * ```typescript - * const embeddings = new OpenAIEmbeddings(); - * const cache = new AzureCosmosDBNoSQLSemanticCache(embeddings, { - * databaseName: DATABASE_NAME, - * containerName: CONTAINER_NAME - * }); - * const model = new ChatOpenAI({cache}); - * - * // Invoke the model to perform an action - * const response = await model.invoke("Do something random!"); - * console.log(response); - * ``` - */ - export class AzureCosmosDBNoSQLSemanticCache extends BaseCache { - private embeddings: EmbeddingsInterface; - - private config: AzureCosmosDBNoSQLConfig; - - private similarityScoreThreshold: number; - - private cacheDict: { [key: string]: AzureCosmosDBNoSQLVectorStore } = {}; - - private vectorDistanceFunction: string; - - constructor( - embeddings: EmbeddingsInterface, - dbConfig: AzureCosmosDBNoSQLConfig, - similarityScoreThreshold: number = 0.6 - ) { - super(); - let client: CosmosClient; - - const connectionString = - dbConfig.connectionString ?? - getEnvironmentVariable("AZURE_COSMOSDB_NOSQL_CONNECTION_STRING"); - - const endpoint = - dbConfig.endpoint ?? - getEnvironmentVariable("AZURE_COSMOSDB_NOSQL_ENDPOINT"); - - if (!dbConfig.client && !connectionString && !endpoint) { - throw new Error( - "AzureCosmosDBNoSQLSemanticCache client, connection string or endpoint must be set." - ); - } - - if (!dbConfig.client) { - if (connectionString) { - // eslint-disable-next-line @typescript-eslint/no-non-null-assertion - let [endpoint, key] = connectionString!.split(";"); - [, endpoint] = endpoint.split("="); - [, key] = key.split("="); - - client = new CosmosClient({ - endpoint, - key, - userAgentSuffix: USER_AGENT_SUFFIX, - }); - } else { - // Use managed identity - client = new CosmosClient({ - endpoint, - aadCredentials: dbConfig.credentials ?? new DefaultAzureCredential(), - userAgentSuffix: USER_AGENT_SUFFIX, - } as CosmosClientOptions); - } - } else { - client = dbConfig.client; - } - - this.vectorDistanceFunction = - dbConfig.vectorEmbeddingPolicy?.vectorEmbeddings[0].distanceFunction ?? - "cosine"; - - this.config = { - ...dbConfig, - client, - databaseName: dbConfig.databaseName, - containerName: dbConfig.containerName ?? DEFAULT_CONTAINER_NAME, - }; - this.embeddings = embeddings; - this.similarityScoreThreshold = similarityScoreThreshold; + BaseCache, + deserializeStoredGeneration, + getCacheKey, + serializeGeneration, +} from "@langchain/core/caches"; +import { Generation } from "@langchain/core/outputs"; +import { Document } from "@langchain/core/documents"; +import { EmbeddingsInterface } from "@langchain/core/embeddings"; +import { CosmosClient, CosmosClientOptions } from "@azure/cosmos"; +import { DefaultAzureCredential } from "@azure/identity"; +import { getEnvironmentVariable } from "@langchain/core/utils/env"; +import { + AzureCosmosDBNoSQLConfig, + AzureCosmosDBNoSQLVectorStore, +} from "../azure_cosmosdb_nosql.js"; + +const USER_AGENT_SUFFIX = "langchainjs-cdbnosql-semanticcache-javascript"; +const DEFAULT_CONTAINER_NAME = "semanticCacheContainer"; + +/** + * Represents a Semantic Cache that uses CosmosDB NoSQL backend as the underlying + * storage system. + * + * @example + * ```typescript + * const embeddings = new OpenAIEmbeddings(); + * const cache = new AzureCosmosDBNoSQLSemanticCache(embeddings, { + * databaseName: DATABASE_NAME, + * containerName: CONTAINER_NAME + * }); + * const model = new ChatOpenAI({cache}); + * + * // Invoke the model to perform an action + * const response = await model.invoke("Do something random!"); + * console.log(response); + * ``` + */ +export class AzureCosmosDBNoSQLSemanticCache extends BaseCache { + private embeddings: EmbeddingsInterface; + + private config: AzureCosmosDBNoSQLConfig; + + private similarityScoreThreshold: number; + + private cacheDict: { [key: string]: AzureCosmosDBNoSQLVectorStore } = {}; + + private vectorDistanceFunction: string; + + constructor( + embeddings: EmbeddingsInterface, + dbConfig: AzureCosmosDBNoSQLConfig, + similarityScoreThreshold: number = 0.6 + ) { + super(); + let client: CosmosClient; + + const connectionString = + dbConfig.connectionString ?? + getEnvironmentVariable("AZURE_COSMOSDB_NOSQL_CONNECTION_STRING"); + + const endpoint = + dbConfig.endpoint ?? + getEnvironmentVariable("AZURE_COSMOSDB_NOSQL_ENDPOINT"); + + if (!dbConfig.client && !connectionString && !endpoint) { + throw new Error( + "AzureCosmosDBNoSQLSemanticCache client, connection string or endpoint must be set." + ); } - - private getLlmCache(llmKey: string) { - const key = getCacheKey(llmKey); - if (!this.cacheDict[key]) { - this.cacheDict[key] = new AzureCosmosDBNoSQLVectorStore( - this.embeddings, - this.config - ); + + if (!dbConfig.client) { + if (connectionString) { + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + let [endpoint, key] = connectionString!.split(";"); + [, endpoint] = endpoint.split("="); + [, key] = key.split("="); + + client = new CosmosClient({ + endpoint, + key, + userAgentSuffix: USER_AGENT_SUFFIX, + }); + } else { + // Use managed identity + client = new CosmosClient({ + endpoint, + aadCredentials: dbConfig.credentials ?? new DefaultAzureCredential(), + userAgentSuffix: USER_AGENT_SUFFIX, + } as CosmosClientOptions); } - return this.cacheDict[key]; - } - - /** - * Retrieves data from the cache. - * - * @param prompt The prompt for lookup. - * @param llmKey The LLM key used to construct the cache key. - * @returns An array of Generations if found, null otherwise. - */ - public async lookup(prompt: string, llmKey: string) { - const llmCache = this.getLlmCache(llmKey); - - const results = await llmCache.similaritySearchWithScore(prompt, 1); - if (!results.length) return null; - - const generations = results - .flatMap(([document, score]) => { - const isSimilar = - (this.vectorDistanceFunction === "euclidean" && - score <= this.similarityScoreThreshold) || - (this.vectorDistanceFunction !== "euclidean" && - score >= this.similarityScoreThreshold); - - if (!isSimilar) return undefined; - - return document.metadata.return_value.map((gen: string) => - deserializeStoredGeneration(JSON.parse(gen)) - ); - }) - .filter((gen) => gen !== undefined); - - return generations.length > 0 ? generations : null; + } else { + client = dbConfig.client; } - - /** - * Updates the cache with new data. - * - * @param prompt The prompt for update. - * @param llmKey The LLM key used to construct the cache key. - * @param value The value to be stored in the cache. - */ - public async update( - prompt: string, - llmKey: string, - returnValue: Generation[] - ) { - const serializedGenerations = returnValue.map((generation) => - JSON.stringify(serializeGeneration(generation)) + + this.vectorDistanceFunction = + dbConfig.vectorEmbeddingPolicy?.vectorEmbeddings[0].distanceFunction ?? + "cosine"; + + this.config = { + ...dbConfig, + client, + databaseName: dbConfig.databaseName, + containerName: dbConfig.containerName ?? DEFAULT_CONTAINER_NAME, + }; + this.embeddings = embeddings; + this.similarityScoreThreshold = similarityScoreThreshold; + } + + private getLlmCache(llmKey: string) { + const key = getCacheKey(llmKey); + if (!this.cacheDict[key]) { + this.cacheDict[key] = new AzureCosmosDBNoSQLVectorStore( + this.embeddings, + this.config ); - const llmCache = this.getLlmCache(llmKey); - const metadata = { - llm_string: llmKey, - prompt, - return_value: serializedGenerations, - }; - const doc = new Document({ - pageContent: prompt, - metadata, - }); - await llmCache.addDocuments([doc]); } - - /** - * deletes the semantic cache for a given llmKey - * @param llmKey - */ - public async clear(llmKey: string) { - const key = getCacheKey(llmKey); - if (this.cacheDict[key]) { - await this.cacheDict[key].delete(); - } + return this.cacheDict[key]; + } + + /** + * Retrieves data from the cache. + * + * @param prompt The prompt for lookup. + * @param llmKey The LLM key used to construct the cache key. + * @returns An array of Generations if found, null otherwise. + */ + public async lookup(prompt: string, llmKey: string) { + const llmCache = this.getLlmCache(llmKey); + + const results = await llmCache.similaritySearchWithScore(prompt, 1); + if (!results.length) return null; + + const generations = results + .flatMap(([document, score]) => { + const isSimilar = + (this.vectorDistanceFunction === "euclidean" && + score <= this.similarityScoreThreshold) || + (this.vectorDistanceFunction !== "euclidean" && + score >= this.similarityScoreThreshold); + + if (!isSimilar) return undefined; + + return document.metadata.return_value.map((gen: string) => + deserializeStoredGeneration(JSON.parse(gen)) + ); + }) + .filter((gen) => gen !== undefined); + + return generations.length > 0 ? generations : null; + } + + /** + * Updates the cache with new data. + * + * @param prompt The prompt for update. + * @param llmKey The LLM key used to construct the cache key. + * @param value The value to be stored in the cache. + */ + public async update( + prompt: string, + llmKey: string, + returnValue: Generation[] + ) { + const serializedGenerations = returnValue.map((generation) => + JSON.stringify(serializeGeneration(generation)) + ); + const llmCache = this.getLlmCache(llmKey); + const metadata = { + llm_string: llmKey, + prompt, + return_value: serializedGenerations, + }; + const doc = new Document({ + pageContent: prompt, + metadata, + }); + await llmCache.addDocuments([doc]); + } + + /** + * deletes the semantic cache for a given llmKey + * @param llmKey + */ + public async clear(llmKey: string) { + const key = getCacheKey(llmKey); + if (this.cacheDict[key]) { + await this.cacheDict[key].delete(); } } - \ No newline at end of file +} diff --git a/libs/langchain-azure-cosmosdb/src/index.ts b/libs/langchain-azure-cosmosdb/src/index.ts index f0f040afb337..00498b06c53c 100644 --- a/libs/langchain-azure-cosmosdb/src/index.ts +++ b/libs/langchain-azure-cosmosdb/src/index.ts @@ -1,5 +1,5 @@ export * from "./azure_cosmosdb_mongodb.js"; export * from "./azure_cosmosdb_nosql.js"; export * from "./chat_histories.js"; -export * from "./caches/caches_nosql.js" -export * from "./caches/caches_mongodb.js" +export * from "./caches/caches_nosql.js"; +export * from "./caches/caches_mongodb.js"; diff --git a/libs/langchain-azure-cosmosdb/src/tests/caches/caches_mongodb.int.test.ts b/libs/langchain-azure-cosmosdb/src/tests/caches/caches_mongodb.int.test.ts new file mode 100644 index 000000000000..7902bdc06e79 --- /dev/null +++ b/libs/langchain-azure-cosmosdb/src/tests/caches/caches_mongodb.int.test.ts @@ -0,0 +1,136 @@ +/* eslint-disable no-nested-ternary */ +/* eslint-disable @typescript-eslint/no-explicit-any */ +/* eslint-disable no-process-env */ +import { ChatOpenAI, OpenAIEmbeddings } from "@langchain/openai"; +import { MongoClient } from "mongodb"; +import { AzureCosmosDBMongoDBSemanticCache } from "../../caches/caches_mongodb.js"; +import { + AzureCosmosDBMongoDBIndexOptions, + AzureCosmosDBMongoDBSimilarityType, +} from "../../azure_cosmosdb_mongodb.js"; + +const DATABASE_NAME = "langchain"; +const COLLECTION_NAME = "test"; + +async function initializeCache( + indexType: any, + distanceFunction: any, + similarityThreshold: number = 0.6 +): Promise { + const embeddingModel = new OpenAIEmbeddings(); + const testEmbedding = await embeddingModel.embedDocuments(["sample text"]); + const dimension = testEmbedding[0].length; + + const indexOptions: AzureCosmosDBMongoDBIndexOptions = { + indexType, + // eslint-disable-next-line no-nested-ternary + similarity: + distanceFunction === "cosine" + ? AzureCosmosDBMongoDBSimilarityType.COS + : distanceFunction === "euclidean" + ? AzureCosmosDBMongoDBSimilarityType.L2 + : AzureCosmosDBMongoDBSimilarityType.IP, + dimensions: dimension, + }; + + let cache: AzureCosmosDBMongoDBSemanticCache; + + const connectionString = process.env.AZURE_COSMOSDB_MONGODB_CONNECTION_STRING; + if (connectionString) { + cache = new AzureCosmosDBMongoDBSemanticCache( + embeddingModel, + { + databaseName: DATABASE_NAME, + collectionName: COLLECTION_NAME, + connectionString, + indexOptions, + }, + similarityThreshold + ); + } else { + throw new Error( + "Please set the environment variable AZURE_COSMOSDB_MONGODB_CONNECTION_STRING" + ); + } + + return cache; +} + +describe("AzureCosmosDBMongoDBSemanticCache", () => { + beforeEach(async () => { + const connectionString = + process.env.AZURE_COSMOSDB_MONGODB_CONNECTION_STRING; + const client = new MongoClient(connectionString!); + + try { + await client.db(DATABASE_NAME).collection(COLLECTION_NAME).drop(); + } catch (error) { + throw new Error("Please set collection name here"); + } + }); + + it("should store and retrieve cache using cosine similarity with ivf index", async () => { + const cache = await initializeCache("ivf", "cosine"); + const model = new ChatOpenAI({ cache }); + const llmString = JSON.stringify(model._identifyingParams); + await cache.update("foo", llmString, [{ text: "fizz" }]); + + let cacheOutput = await cache.lookup("foo", llmString); + expect(cacheOutput).toEqual([{ text: "fizz" }]); + + cacheOutput = await cache.lookup("bar", llmString); + expect(cacheOutput).toEqual(null); + + await cache.clear(llmString); + }); + + it("should store and retrieve cache using euclidean similarity with hnsw index", async () => { + const cache = await initializeCache("hnsw", "euclidean"); + const model = new ChatOpenAI({ cache }); + const llmString = JSON.stringify(model._identifyingParams); + await cache.update("foo", llmString, [{ text: "fizz" }]); + + let cacheOutput = await cache.lookup("foo", llmString); + expect(cacheOutput).toEqual([{ text: "fizz" }]); + + cacheOutput = await cache.lookup("bar", llmString); + expect(cacheOutput).toEqual(null); + + await cache.clear(llmString); + }); + + it("should return null if similarity score is below threshold (cosine similarity with ivf index)", async () => { + const cache = await initializeCache("ivf", "cosine", 0.8); + const model = new ChatOpenAI({ cache }); + const llmString = JSON.stringify(model._identifyingParams); + await cache.update("foo", llmString, [{ text: "fizz" }]); + + const cacheOutput = await cache.lookup("foo", llmString); + expect(cacheOutput).toEqual([{ text: "fizz" }]); + + const resultBelowThreshold = await cache.lookup("bar", llmString); + expect(resultBelowThreshold).toEqual(null); + + await cache.clear(llmString); + }); + + it("should handle a variety of cache updates and lookups", async () => { + const cache = await initializeCache("ivf", "cosine", 0.7); + const model = new ChatOpenAI({ cache }); + const llmString = JSON.stringify(model._identifyingParams); + + await cache.update("test1", llmString, [{ text: "response 1" }]); + await cache.update("test2", llmString, [{ text: "response 2" }]); + + let cacheOutput = await cache.lookup("test1", llmString); + expect(cacheOutput).toEqual([{ text: "response 1" }]); + + cacheOutput = await cache.lookup("test2", llmString); + expect(cacheOutput).toEqual([{ text: "response 2" }]); + + cacheOutput = await cache.lookup("test3", llmString); + expect(cacheOutput).toEqual(null); + + await cache.clear(llmString); + }); +}); diff --git a/libs/langchain-azure-cosmosdb/src/tests/caches/caches_mongodb.test.ts b/libs/langchain-azure-cosmosdb/src/tests/caches/caches_mongodb.test.ts new file mode 100644 index 000000000000..4bd9cf0ce996 --- /dev/null +++ b/libs/langchain-azure-cosmosdb/src/tests/caches/caches_mongodb.test.ts @@ -0,0 +1,72 @@ +/* eslint-disable @typescript-eslint/no-explicit-any */ +// eslint-disable-next-line import/no-extraneous-dependencies +import { jest } from "@jest/globals"; +import { FakeEmbeddings, FakeLLM } from "@langchain/core/utils/testing"; +import { Document } from "@langchain/core/documents"; +import { MongoClient } from "mongodb"; +import { AzureCosmosDBMongoDBSemanticCache } from "../../index.js"; + +const createMockClient = () => ({ + db: jest.fn().mockReturnValue({ + collectionName: "documents", + collection: jest.fn().mockReturnValue({ + listIndexes: jest.fn().mockReturnValue({ + toArray: jest.fn().mockReturnValue([ + { + name: "vectorSearchIndex", + }, + ]), + }), + findOne: jest.fn().mockReturnValue({ + metadata: { + return_value: ['{"text": "fizz"}'], + }, + similarityScore: 0.8, + }), + insertMany: jest.fn().mockImplementation((docs: any) => ({ + insertedIds: docs.map((_: any, i: any) => `id${i}`), + })), + aggregate: jest.fn().mockReturnValue({ + map: jest.fn().mockReturnValue({ + toArray: jest.fn().mockReturnValue([ + [ + new Document({ + pageContent: "test", + metadata: { return_value: ['{"text": "fizz"}'] }, + }), + 0.8, + ], + ]), + }), + }), + }), + command: jest.fn(), + }), + connect: jest.fn(), + close: jest.fn(), +}); + +describe("AzureCosmosDBMongoDBSemanticCache", () => { + it("should store, retrieve, and clear cache in MongoDB", async () => { + const mockClient = createMockClient() as any; + const embeddings = new FakeEmbeddings(); + const cache = new AzureCosmosDBMongoDBSemanticCache( + embeddings, + { + client: mockClient as MongoClient, + }, + 0.8 + ); + + expect(cache).toBeDefined(); + + const llm = new FakeLLM({}); + const llmString = JSON.stringify(llm._identifyingParams()); + + await cache.update("foo", llmString, [{ text: "fizz" }]); + expect(mockClient.db().collection().insertMany).toHaveBeenCalled(); + + const result = await cache.lookup("foo", llmString); + expect(result).toEqual([{ text: "fizz" }]); + }); +}); diff --git a/libs/langchain-azure-cosmosdb/src/tests/caches/caches_nosql.init.test.ts b/libs/langchain-azure-cosmosdb/src/tests/caches/caches_nosql.init.test.ts deleted file mode 100644 index 88bab773e4fc..000000000000 --- a/libs/langchain-azure-cosmosdb/src/tests/caches/caches_nosql.init.test.ts +++ /dev/null @@ -1,245 +0,0 @@ -/* eslint-disable no-process-env */ -/* eslint-disable @typescript-eslint/no-explicit-any */ - -import { - CosmosClient, - IndexingMode, - VectorEmbeddingPolicy, - } from "@azure/cosmos"; - import { DefaultAzureCredential } from "@azure/identity"; - import { ChatOpenAI, OpenAIEmbeddings } from "@langchain/openai"; - import { AzureCosmosDBNoSQLSemanticCache } from "../../caches/caches_nosql.js" - - const DATABASE_NAME = "langchainTestCacheDB"; - const CONTAINER_NAME = "testContainer"; - - function indexingPolicy(indexType: any) { - return { - indexingMode: IndexingMode.consistent, - includedPaths: [{ path: "/*" }], - excludedPaths: [{ path: '/"_etag"/?' }], - vectorIndexes: [{ path: "/embedding", type: indexType }], - }; - } - - function vectorEmbeddingPolicy( - distanceFunction: "euclidean" | "cosine" | "dotproduct", - dimension: number - ): VectorEmbeddingPolicy { - return { - vectorEmbeddings: [ - { - path: "/embedding", - dataType: "float32", - distanceFunction, - dimensions: dimension, - }, - ], - }; - } - - async function initializeCache( - indexType: any, - distanceFunction: any, - similarityThreshold?: number - ): Promise { - let cache: AzureCosmosDBNoSQLSemanticCache; - const embeddingModel = new OpenAIEmbeddings(); - const testEmbedding = await embeddingModel.embedDocuments(["sample text"]); - const dimension = Math.min( - testEmbedding[0].length, - indexType === "flat" ? 505 : 4096 - ); - if (process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING) { - cache = new AzureCosmosDBNoSQLSemanticCache( - new OpenAIEmbeddings(), - { - databaseName: DATABASE_NAME, - containerName: CONTAINER_NAME, - connectionString: process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING, - indexingPolicy: indexingPolicy(indexType), - vectorEmbeddingPolicy: vectorEmbeddingPolicy( - distanceFunction, - dimension - ), - }, - similarityThreshold - ); - } else if (process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT) { - cache = new AzureCosmosDBNoSQLSemanticCache( - new OpenAIEmbeddings(), - { - databaseName: DATABASE_NAME, - containerName: CONTAINER_NAME, - endpoint: process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT, - indexingPolicy: indexingPolicy(indexType), - vectorEmbeddingPolicy: vectorEmbeddingPolicy( - distanceFunction, - dimension - ), - }, - similarityThreshold - ); - } else { - throw new Error( - "Please set the environment variable AZURE_COSMOSDB_NOSQL_CONNECTION_STRING or AZURE_COSMOSDB_NOSQL_ENDPOINT" - ); - } - return cache; - } - - /* - * To run this test, you need have an Azure Cosmos DB for NoSQL instance - * running. You can deploy a free version on Azure Portal without any cost, - * following this guide: - * https://learn.microsoft.com/azure/cosmos-db/nosql/vector-search - * - * You do not need to create a database or collection, it will be created - * automatically by the test. - * - * Once you have the instance running, you need to set the following environment - * variables before running the test: - * - AZURE_COSMOSDB_NOSQL_CONNECTION_STRING or AZURE_COSMOSDB_NOSQL_ENDPOINT - * - AZURE_OPENAI_API_KEY - * - AZURE_OPENAI_API_INSTANCE_NAME - * - AZURE_OPENAI_API_EMBEDDINGS_DEPLOYMENT_NAME - * - AZURE_OPENAI_API_VERSION - */ - describe("Azure CosmosDB NoSQL Semantic Cache", () => { - beforeEach(async () => { - let client: CosmosClient; - - if (process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING) { - client = new CosmosClient( - process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING - ); - } else if (process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT) { - client = new CosmosClient({ - endpoint: process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT, - aadCredentials: new DefaultAzureCredential(), - }); - } else { - throw new Error( - "Please set the environment variable AZURE_COSMOSDB_NOSQL_CONNECTION_STRING or AZURE_COSMOSDB_NOSQL_ENDPOINT" - ); - } - - // Make sure the database does not exists - try { - await client.database(DATABASE_NAME).delete(); - } catch { - // Ignore error if the database does not exist - } - }); - - it("test AzureCosmosDBNoSqlSemanticCache with cosine quantizedFlat", async () => { - const cache = await initializeCache("quantizedFlat", "cosine"); - const model = new ChatOpenAI({ cache }); - const llmString = JSON.stringify(model._identifyingParams); - await cache.update("foo", llmString, [{ text: "fizz" }]); - - let cacheOutput = await cache.lookup("foo", llmString); - expect(cacheOutput).toEqual([{ text: "fizz" }]); - - cacheOutput = await cache.lookup("bar", llmString); - expect(cacheOutput).toEqual(null); - - await cache.clear(llmString); - }); - - it("test AzureCosmosDBNoSqlSemanticCache with cosine flat", async () => { - const cache = await initializeCache("flat", "cosine"); - const model = new ChatOpenAI({ cache }); - const llmString = JSON.stringify(model._identifyingParams); - await cache.update("foo", llmString, [{ text: "fizz" }]); - - let cacheOutput = await cache.lookup("foo", llmString); - expect(cacheOutput).toEqual([{ text: "fizz" }]); - - cacheOutput = await cache.lookup("bar", llmString); - expect(cacheOutput).toEqual(null); - - await cache.clear(llmString); - }); - - it("test AzureCosmosDBNoSqlSemanticCache with dotProduct quantizedFlat", async () => { - const cache = await initializeCache("quantizedFlat", "dotproduct"); - const model = new ChatOpenAI({ cache }); - const llmString = JSON.stringify(model._identifyingParams); - await cache.update("foo", llmString, [{ text: "fizz" }]); - - let cacheOutput = await cache.lookup("foo", llmString); - expect(cacheOutput).toEqual([{ text: "fizz" }]); - - cacheOutput = await cache.lookup("bar", llmString); - expect(cacheOutput).toEqual(null); - - await cache.clear(llmString); - }); - - it("test AzureCosmosDBNoSqlSemanticCache with dotProduct flat", async () => { - const cache = await initializeCache("flat", "cosine"); - const model = new ChatOpenAI({ cache }); - const llmString = JSON.stringify(model._identifyingParams); - await cache.update("foo", llmString, [{ text: "fizz" }]); - - let cacheOutput = await cache.lookup("foo", llmString); - expect(cacheOutput).toEqual([{ text: "fizz" }]); - - cacheOutput = await cache.lookup("bar", llmString); - expect(cacheOutput).toEqual(null); - - await cache.clear(llmString); - }); - - it("test AzureCosmosDBNoSqlSemanticCache with euclidean quantizedFlat", async () => { - const cache = await initializeCache("quantizedFlat", "euclidean"); - const model = new ChatOpenAI({ cache }); - const llmString = JSON.stringify(model._identifyingParams); - await cache.update("foo", llmString, [{ text: "fizz" }]); - - let cacheOutput = await cache.lookup("foo", llmString); - expect(cacheOutput).toEqual([{ text: "fizz" }]); - - cacheOutput = await cache.lookup("bar", llmString); - expect(cacheOutput).toEqual(null); - - await cache.clear(llmString); - }); - - it("test AzureCosmosDBNoSqlSemanticCache with euclidean flat", async () => { - const cache = await initializeCache("flat", "euclidean"); - const model = new ChatOpenAI({ cache }); - const llmString = JSON.stringify(model._identifyingParams); - await cache.update("foo", llmString, [{ text: "fizz" }]); - - let cacheOutput = await cache.lookup("foo", llmString); - expect(cacheOutput).toEqual([{ text: "fizz" }]); - - cacheOutput = await cache.lookup("bar", llmString); - expect(cacheOutput).toEqual(null); - - await cache.clear(llmString); - }); - - it("test AzureCosmosDBNoSqlSemanticCache response according to similarity score", async () => { - const cache = await initializeCache("quantizedFlat", "cosine"); - const model = new ChatOpenAI({ cache }); - const response1 = await model.invoke( - "Where is the headquarter of Microsoft?" - ); - console.log(response1.content); - // gives similarity score of 0.56 which is less than the threshold of 0.6. The cache - // will retun null which will allow the model to generate result. - const response2 = await model.invoke( - "List all Microsoft offices in India." - ); - expect(response2.content).not.toEqual(response1.content); - console.log(response2.content); - // gives similarity score of .63 > 0.6 - const response3 = await model.invoke("Tell me something about Microsoft"); - expect(response3.content).toEqual(response1.content); - console.log(response3.content); - }); - }); - \ No newline at end of file diff --git a/libs/langchain-azure-cosmosdb/src/tests/caches/caches_nosql.int.test.ts b/libs/langchain-azure-cosmosdb/src/tests/caches/caches_nosql.int.test.ts new file mode 100644 index 000000000000..8d05abfdfe6c --- /dev/null +++ b/libs/langchain-azure-cosmosdb/src/tests/caches/caches_nosql.int.test.ts @@ -0,0 +1,244 @@ +/* eslint-disable no-process-env */ +/* eslint-disable @typescript-eslint/no-explicit-any */ + +import { + CosmosClient, + IndexingMode, + VectorEmbeddingPolicy, +} from "@azure/cosmos"; +import { DefaultAzureCredential } from "@azure/identity"; +import { ChatOpenAI, OpenAIEmbeddings } from "@langchain/openai"; +import { AzureCosmosDBNoSQLSemanticCache } from "../../caches/caches_nosql.js"; + +const DATABASE_NAME = "langchainTestCacheDB"; +const CONTAINER_NAME = "testContainer"; + +function indexingPolicy(indexType: any) { + return { + indexingMode: IndexingMode.consistent, + includedPaths: [{ path: "/*" }], + excludedPaths: [{ path: '/"_etag"/?' }], + vectorIndexes: [{ path: "/embedding", type: indexType }], + }; +} + +function vectorEmbeddingPolicy( + distanceFunction: "euclidean" | "cosine" | "dotproduct", + dimension: number +): VectorEmbeddingPolicy { + return { + vectorEmbeddings: [ + { + path: "/embedding", + dataType: "float32", + distanceFunction, + dimensions: dimension, + }, + ], + }; +} + +async function initializeCache( + indexType: any, + distanceFunction: any, + similarityThreshold?: number +): Promise { + let cache: AzureCosmosDBNoSQLSemanticCache; + const embeddingModel = new OpenAIEmbeddings(); + const testEmbedding = await embeddingModel.embedDocuments(["sample text"]); + const dimension = Math.min( + testEmbedding[0].length, + indexType === "flat" ? 505 : 4096 + ); + if (process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING) { + cache = new AzureCosmosDBNoSQLSemanticCache( + new OpenAIEmbeddings(), + { + databaseName: DATABASE_NAME, + containerName: CONTAINER_NAME, + connectionString: process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING, + indexingPolicy: indexingPolicy(indexType), + vectorEmbeddingPolicy: vectorEmbeddingPolicy( + distanceFunction, + dimension + ), + }, + similarityThreshold + ); + } else if (process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT) { + cache = new AzureCosmosDBNoSQLSemanticCache( + new OpenAIEmbeddings(), + { + databaseName: DATABASE_NAME, + containerName: CONTAINER_NAME, + endpoint: process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT, + indexingPolicy: indexingPolicy(indexType), + vectorEmbeddingPolicy: vectorEmbeddingPolicy( + distanceFunction, + dimension + ), + }, + similarityThreshold + ); + } else { + throw new Error( + "Please set the environment variable AZURE_COSMOSDB_NOSQL_CONNECTION_STRING or AZURE_COSMOSDB_NOSQL_ENDPOINT" + ); + } + return cache; +} + +/* + * To run this test, you need have an Azure Cosmos DB for NoSQL instance + * running. You can deploy a free version on Azure Portal without any cost, + * following this guide: + * https://learn.microsoft.com/azure/cosmos-db/nosql/vector-search + * + * You do not need to create a database or collection, it will be created + * automatically by the test. + * + * Once you have the instance running, you need to set the following environment + * variables before running the test: + * - AZURE_COSMOSDB_NOSQL_CONNECTION_STRING or AZURE_COSMOSDB_NOSQL_ENDPOINT + * - AZURE_OPENAI_API_KEY + * - AZURE_OPENAI_API_INSTANCE_NAME + * - AZURE_OPENAI_API_EMBEDDINGS_DEPLOYMENT_NAME + * - AZURE_OPENAI_API_VERSION + */ +describe("Azure CosmosDB NoSQL Semantic Cache", () => { + beforeEach(async () => { + let client: CosmosClient; + + if (process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING) { + client = new CosmosClient( + process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING + ); + } else if (process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT) { + client = new CosmosClient({ + endpoint: process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT, + aadCredentials: new DefaultAzureCredential(), + }); + } else { + throw new Error( + "Please set the environment variable AZURE_COSMOSDB_NOSQL_CONNECTION_STRING or AZURE_COSMOSDB_NOSQL_ENDPOINT" + ); + } + + // Make sure the database does not exists + try { + await client.database(DATABASE_NAME).delete(); + } catch { + // Ignore error if the database does not exist + } + }); + + it("test AzureCosmosDBNoSqlSemanticCache with cosine quantizedFlat", async () => { + const cache = await initializeCache("quantizedFlat", "cosine"); + const model = new ChatOpenAI({ cache }); + const llmString = JSON.stringify(model._identifyingParams); + await cache.update("foo", llmString, [{ text: "fizz" }]); + + let cacheOutput = await cache.lookup("foo", llmString); + expect(cacheOutput).toEqual([{ text: "fizz" }]); + + cacheOutput = await cache.lookup("bar", llmString); + expect(cacheOutput).toEqual(null); + + await cache.clear(llmString); + }); + + it("test AzureCosmosDBNoSqlSemanticCache with cosine flat", async () => { + const cache = await initializeCache("flat", "cosine"); + const model = new ChatOpenAI({ cache }); + const llmString = JSON.stringify(model._identifyingParams); + await cache.update("foo", llmString, [{ text: "fizz" }]); + + let cacheOutput = await cache.lookup("foo", llmString); + expect(cacheOutput).toEqual([{ text: "fizz" }]); + + cacheOutput = await cache.lookup("bar", llmString); + expect(cacheOutput).toEqual(null); + + await cache.clear(llmString); + }); + + it("test AzureCosmosDBNoSqlSemanticCache with dotProduct quantizedFlat", async () => { + const cache = await initializeCache("quantizedFlat", "dotproduct"); + const model = new ChatOpenAI({ cache }); + const llmString = JSON.stringify(model._identifyingParams); + await cache.update("foo", llmString, [{ text: "fizz" }]); + + let cacheOutput = await cache.lookup("foo", llmString); + expect(cacheOutput).toEqual([{ text: "fizz" }]); + + cacheOutput = await cache.lookup("bar", llmString); + expect(cacheOutput).toEqual(null); + + await cache.clear(llmString); + }); + + it("test AzureCosmosDBNoSqlSemanticCache with dotProduct flat", async () => { + const cache = await initializeCache("flat", "cosine"); + const model = new ChatOpenAI({ cache }); + const llmString = JSON.stringify(model._identifyingParams); + await cache.update("foo", llmString, [{ text: "fizz" }]); + + let cacheOutput = await cache.lookup("foo", llmString); + expect(cacheOutput).toEqual([{ text: "fizz" }]); + + cacheOutput = await cache.lookup("bar", llmString); + expect(cacheOutput).toEqual(null); + + await cache.clear(llmString); + }); + + it("test AzureCosmosDBNoSqlSemanticCache with euclidean quantizedFlat", async () => { + const cache = await initializeCache("quantizedFlat", "euclidean"); + const model = new ChatOpenAI({ cache }); + const llmString = JSON.stringify(model._identifyingParams); + await cache.update("foo", llmString, [{ text: "fizz" }]); + + let cacheOutput = await cache.lookup("foo", llmString); + expect(cacheOutput).toEqual([{ text: "fizz" }]); + + cacheOutput = await cache.lookup("bar", llmString); + expect(cacheOutput).toEqual(null); + + await cache.clear(llmString); + }); + + it("test AzureCosmosDBNoSqlSemanticCache with euclidean flat", async () => { + const cache = await initializeCache("flat", "euclidean"); + const model = new ChatOpenAI({ cache }); + const llmString = JSON.stringify(model._identifyingParams); + await cache.update("foo", llmString, [{ text: "fizz" }]); + + let cacheOutput = await cache.lookup("foo", llmString); + expect(cacheOutput).toEqual([{ text: "fizz" }]); + + cacheOutput = await cache.lookup("bar", llmString); + expect(cacheOutput).toEqual(null); + + await cache.clear(llmString); + }); + + it("test AzureCosmosDBNoSqlSemanticCache response according to similarity score", async () => { + const cache = await initializeCache("quantizedFlat", "cosine"); + const model = new ChatOpenAI({ cache }); + const response1 = await model.invoke( + "Where is the headquarter of Microsoft?" + ); + console.log(response1.content); + // gives similarity score of 0.56 which is less than the threshold of 0.6. The cache + // will retun null which will allow the model to generate result. + const response2 = await model.invoke( + "List all Microsoft offices in India." + ); + expect(response2.content).not.toEqual(response1.content); + console.log(response2.content); + // gives similarity score of .63 > 0.6 + const response3 = await model.invoke("Tell me something about Microsoft"); + expect(response3.content).toEqual(response1.content); + console.log(response3.content); + }); +});