From c39c4340aa8b5186d84682f834c185413db81619 Mon Sep 17 00:00:00 2001 From: Jacob Lee Date: Tue, 30 May 2023 12:13:46 -0700 Subject: [PATCH] Adds ZepMemory and ZepRetriever (#1450) * Zep Integration for Memory, Retriever, Docs + Tests * Adding tests for retriever, zep_memory impl * updated yarn.lock * updated * updated yarn.lock * updated * updated with PR changes * minor updates * Removed zep from docs/package, ^ver on langchain/package.json, updated example * few edits to make sure memory is properly added * reverting unnecessary package changes * added memory * bump js sdk ver to 0.3.0 * Updated Integration tests and bumped sdk ver * Updated retriever example, zepMemory/zepConvBufferMemory to zep_memory.ts * Adding 2 examples, removed unnecessary console.log * Removed Zep Memory Backed Chat from example & doc * removed option-2 implementation completely * Change entrypoint to memory/zep * Typo fix * Standardize parameters * Change entrypoint to fix docs * Fix Zep session loading * Reduce yarn.lock diff --------- Co-authored-by: Sharath Rajasekar --- .../indexes/retrievers/zep-retriever.mdx | 20 +++ .../modules/memory/examples/zep_memory.mdx | 26 ++++ examples/package.json | 1 + examples/src/memory/zep.ts | 39 +++++ examples/src/retrievers/zep.ts | 14 ++ langchain/.gitignore | 6 + langchain/package.json | 21 +++ langchain/scripts/create-entrypoints.js | 4 + .../src/memory/tests/zep_memory.int.test.ts | 44 ++++++ langchain/src/memory/zep.ts | 140 ++++++++++++++++++ .../src/retrievers/tests/zep.int.test.ts | 18 +++ langchain/src/retrievers/zep.ts | 57 +++++++ langchain/tsconfig.json | 2 + yarn.lock | 30 +++- 14 files changed, 420 insertions(+), 2 deletions(-) create mode 100644 docs/docs/modules/indexes/retrievers/zep-retriever.mdx create mode 100644 docs/docs/modules/memory/examples/zep_memory.mdx create mode 100644 examples/src/memory/zep.ts create mode 100644 examples/src/retrievers/zep.ts create mode 100644 langchain/src/memory/tests/zep_memory.int.test.ts create mode 100644 langchain/src/memory/zep.ts create mode 100644 langchain/src/retrievers/tests/zep.int.test.ts create mode 100644 langchain/src/retrievers/zep.ts diff --git a/docs/docs/modules/indexes/retrievers/zep-retriever.mdx b/docs/docs/modules/indexes/retrievers/zep-retriever.mdx new file mode 100644 index 000000000000..2fa90a2db90c --- /dev/null +++ b/docs/docs/modules/indexes/retrievers/zep-retriever.mdx @@ -0,0 +1,20 @@ +--- +hide_table_of_contents: true +--- + +# Zep Retriever + +This example shows how to use the Zep Retriever in a `RetrievalQAChain` to retrieve documents from Zep memory store. + +## Setup + +```bash npm2yarn +npm i @getzep/zep-js +``` + +## Usage + +import CodeBlock from "@theme/CodeBlock"; +import Example from "@examples/retrievers/zep.ts"; + +{Example} diff --git a/docs/docs/modules/memory/examples/zep_memory.mdx b/docs/docs/modules/memory/examples/zep_memory.mdx new file mode 100644 index 000000000000..1969e1a67dce --- /dev/null +++ b/docs/docs/modules/memory/examples/zep_memory.mdx @@ -0,0 +1,26 @@ +--- +hide_table_of_contents: true +--- + +# Zep Memory + +[Zep](https://github.com/getzep/zep) is a memory server that stores, summarizes, embeds, indexes, and enriches conversational AI chat histories, autonomous agent histories, document Q&A histories and exposes them via simple, low-latency APIs. + +Key Features: + +- Long-term memory persistence, with access to historical messages irrespective of your summarization strategy. +- Auto-summarization of memory messages based on a configurable message window. A series of summaries are stored, providing flexibility for future summarization strategies. +- Vector search over memories, with messages automatically embedded on creation. +- Auto-token counting of memories and summaries, allowing finer-grained control over prompt assembly. +- [Python](https://github.com/getzep/zep-python) and [JavaScript](https://github.com/getzep/zep-js) SDKs. + +## Setup + +See the instructions from [Zep](https://github.com/getzep/zep) for running the server locally or through an automated hosting provider. + +## Usage + +import CodeBlock from "@theme/CodeBlock"; +import Example from "@examples/memory/zep.ts"; + +{Example} diff --git a/examples/package.json b/examples/package.json index 4360e618543c..a2a7818dc280 100644 --- a/examples/package.json +++ b/examples/package.json @@ -24,6 +24,7 @@ "dependencies": { "@clickhouse/client": "^0.0.14", "@getmetal/metal-sdk": "^4.0.0", + "@getzep/zep-js": "0.3.1", "@gomomento/sdk": "^1.23.0", "@opensearch-project/opensearch": "^2.2.0", "@pinecone-database/pinecone": "^0.0.14", diff --git a/examples/src/memory/zep.ts b/examples/src/memory/zep.ts new file mode 100644 index 000000000000..4d1518e0f1db --- /dev/null +++ b/examples/src/memory/zep.ts @@ -0,0 +1,39 @@ +import { ChatOpenAI } from "langchain/chat_models/openai"; +import { ConversationChain } from "langchain/chains"; +import { ZepMemory } from "langchain/memory/zep"; + +const sessionId = "TestSession1234"; +const zepURL = "http://localhost:8000"; + +const memory = new ZepMemory({ + sessionId, + baseURL: zepURL, +}); + +const model = new ChatOpenAI({ + modelName: "gpt-3.5-turbo", + temperature: 0, +}); + +const chain = new ConversationChain({ llm: model, memory }); + +const res1 = await chain.call({ input: "Hi! I'm Jim." }); +console.log({ res1 }); +/* +{ + res1: { + text: "Hello Jim! It's nice to meet you. My name is AI. How may I assist you today?" + } +} +*/ + +const res2 = await chain.call({ input: "What did I just say my name was?" }); +console.log({ res2 }); + +/* +{ + res1: { + text: "You said your name was Jim." + } +} +*/ diff --git a/examples/src/retrievers/zep.ts b/examples/src/retrievers/zep.ts new file mode 100644 index 000000000000..c0fb60efb5fb --- /dev/null +++ b/examples/src/retrievers/zep.ts @@ -0,0 +1,14 @@ +import { ZepRetriever } from "langchain/retrievers/zep"; + +export const run = async () => { + const url = process.env.ZEP_URL || "http://localhost:8000"; + const sessionId = "TestSession1232"; + console.log(`Session ID: ${sessionId}, URL: ${url}`); + + const retriever = new ZepRetriever({ sessionId, url }); + + const query = "hello"; + const docs = await retriever.getRelevantDocuments(query); + + console.log(docs); +}; diff --git a/langchain/.gitignore b/langchain/.gitignore index a60d21610d8a..6112776675b0 100644 --- a/langchain/.gitignore +++ b/langchain/.gitignore @@ -145,6 +145,9 @@ text_splitter.d.ts memory.cjs memory.js memory.d.ts +memory/zep.cjs +memory/zep.js +memory/zep.d.ts document.cjs document.js document.d.ts @@ -274,6 +277,9 @@ retrievers/remote.d.ts retrievers/supabase.cjs retrievers/supabase.js retrievers/supabase.d.ts +retrievers/zep.cjs +retrievers/zep.js +retrievers/zep.d.ts retrievers/metal.cjs retrievers/metal.js retrievers/metal.d.ts diff --git a/langchain/package.json b/langchain/package.json index ea49039f60c1..daacd532c366 100644 --- a/langchain/package.json +++ b/langchain/package.json @@ -157,6 +157,9 @@ "memory.cjs", "memory.js", "memory.d.ts", + "memory/zep.cjs", + "memory/zep.js", + "memory/zep.d.ts", "document.cjs", "document.js", "document.d.ts", @@ -286,6 +289,9 @@ "retrievers/supabase.cjs", "retrievers/supabase.js", "retrievers/supabase.d.ts", + "retrievers/zep.cjs", + "retrievers/zep.js", + "retrievers/zep.d.ts", "retrievers/metal.cjs", "retrievers/metal.js", "retrievers/metal.d.ts", @@ -385,6 +391,7 @@ "@clickhouse/client": "^0.0.14", "@faker-js/faker": "^7.6.0", "@getmetal/metal-sdk": "^4.0.0", + "@getzep/zep-js": "^0.3.1", "@gomomento/sdk": "^1.23.0", "@huggingface/inference": "^1.5.1", "@jest/globals": "^29.5.0", @@ -456,6 +463,7 @@ "@aws-sdk/client-sagemaker-runtime": "^3.310.0", "@clickhouse/client": "^0.0.14", "@getmetal/metal-sdk": "*", + "@getzep/zep-js": "^0.3.1", "@gomomento/sdk": "^1.23.0", "@huggingface/inference": "^1.5.1", "@opensearch-project/opensearch": "*", @@ -512,6 +520,9 @@ "@getmetal/metal-sdk": { "optional": true }, + "@getzep/zep-js": { + "optional": true + }, "@gomomento/sdk": { "optional": true }, @@ -913,6 +924,11 @@ "import": "./memory.js", "require": "./memory.cjs" }, + "./memory/zep": { + "types": "./memory/zep.d.ts", + "import": "./memory/zep.js", + "require": "./memory/zep.cjs" + }, "./document": { "types": "./document.d.ts", "import": "./document.js", @@ -1134,6 +1150,11 @@ "import": "./retrievers/supabase.js", "require": "./retrievers/supabase.cjs" }, + "./retrievers/zep": { + "types": "./retrievers/zep.d.ts", + "import": "./retrievers/zep.js", + "require": "./retrievers/zep.cjs" + }, "./retrievers/metal": { "types": "./retrievers/metal.d.ts", "import": "./retrievers/metal.js", diff --git a/langchain/scripts/create-entrypoints.js b/langchain/scripts/create-entrypoints.js index 90f1e485f52e..5059f8d2a6f7 100644 --- a/langchain/scripts/create-entrypoints.js +++ b/langchain/scripts/create-entrypoints.js @@ -66,6 +66,7 @@ const entrypoints = { text_splitter: "text_splitter", // memory memory: "memory/index", + "memory/zep": "memory/zep", // document document: "document", // docstore @@ -119,6 +120,7 @@ const entrypoints = { retrievers: "retrievers/index", "retrievers/remote": "retrievers/remote/index", "retrievers/supabase": "retrievers/supabase", + "retrievers/zep": "retrievers/zep", "retrievers/metal": "retrievers/metal", "retrievers/databerry": "retrievers/databerry", "retrievers/contextual_compression": "retrievers/contextual_compression", @@ -190,6 +192,7 @@ const requiresOptionalDependency = [ "vectorstores/myscale", "vectorstores/redis", "vectorstores/tigris", + "memory/zep", "document_loaders/web/apify_dataset", "document_loaders/web/cheerio", "document_loaders/web/puppeteer", @@ -217,6 +220,7 @@ const requiresOptionalDependency = [ "chat_models/googlevertexai", "sql_db", "retrievers/supabase", + "retrievers/zep", "retrievers/metal", "retrievers/self_query", "output_parsers/expression", diff --git a/langchain/src/memory/tests/zep_memory.int.test.ts b/langchain/src/memory/tests/zep_memory.int.test.ts new file mode 100644 index 000000000000..f9684a70d82e --- /dev/null +++ b/langchain/src/memory/tests/zep_memory.int.test.ts @@ -0,0 +1,44 @@ +import { expect, test } from "@jest/globals"; +import { v4 as uuid } from "uuid"; +import { ZepMemory } from "../zep.js"; + +const sessionId = uuid(); +const baseURL = "http://localhost:8000"; +const zepMemory = new ZepMemory({ sessionId, baseURL }); + +beforeEach((done) => { + setTimeout(done, 1000); // 1-second delay before each test case +}); + +test("addMemory to Zep memory", async () => { + await zepMemory.saveContext( + { input: "Who was Octavia Butler?" }, + { + response: + "Octavia Estelle Butler (June 22, 1947 – " + + "February 24, 2006) was an American science fiction author.", + } + ); +}); + +test("getMessages from Zep memory", async () => { + const memoryVariables = await zepMemory.loadMemoryVariables({}); + console.log("memoryVariables", memoryVariables); + + // Check if memoryKey exists in the memoryVariables + expect(memoryVariables).toHaveProperty(zepMemory.memoryKey); + + const messages = memoryVariables[zepMemory.memoryKey]; + + // Check if messages is an array or string + if (typeof messages === "string") { + // In this case, we can at least expect a non-empty string. + expect(messages.length).toBeGreaterThanOrEqual(1); + } else if (Array.isArray(messages)) { + expect(messages.length).toBeGreaterThanOrEqual(1); + } else { + console.log("failed to get messages: ", messages); + // Fail the test because messages is neither string nor array + throw new Error("Returned messages is neither string nor array"); + } +}); diff --git a/langchain/src/memory/zep.ts b/langchain/src/memory/zep.ts new file mode 100644 index 000000000000..f3ef23c9c906 --- /dev/null +++ b/langchain/src/memory/zep.ts @@ -0,0 +1,140 @@ +import { ZepClient, Memory, Message } from "@getzep/zep-js"; +import { + InputValues, + OutputValues, + MemoryVariables, + getBufferString, + getInputValue, +} from "./base.js"; +import { BaseChatMemory, BaseChatMemoryInput } from "./chat_memory.js"; +import { + BaseChatMessage, + ChatMessage, + AIChatMessage, + HumanChatMessage, +} from "../schema/index.js"; + +export interface ZepMemoryInput extends BaseChatMemoryInput { + humanPrefix?: string; + + aiPrefix?: string; + + memoryKey?: string; + + baseURL: string; + + sessionId: string; +} + +export class ZepMemory extends BaseChatMemory implements ZepMemoryInput { + humanPrefix = "Human"; + + aiPrefix = "AI"; + + memoryKey = "history"; + + baseURL: string; + + sessionId: string; + + zepClient: ZepClient; + + constructor(fields: ZepMemoryInput) { + super({ + returnMessages: fields?.returnMessages ?? false, + inputKey: fields?.inputKey, + outputKey: fields?.outputKey, + }); + + this.humanPrefix = fields.humanPrefix ?? this.humanPrefix; + this.aiPrefix = fields.aiPrefix ?? this.aiPrefix; + this.memoryKey = fields.memoryKey ?? this.memoryKey; + this.baseURL = fields.baseURL; + this.sessionId = fields.sessionId; + this.zepClient = new ZepClient(this.baseURL); + } + + get memoryKeys() { + return [this.memoryKey]; + } + + async loadMemoryVariables(values: InputValues): Promise { + const lastN = values.lastN ?? 10; + const memory = await this.zepClient.getMemory(this.sessionId, lastN); + let messages: BaseChatMessage[] = []; + + if (memory) { + messages = memory.messages.map((message) => { + const { content, role } = message; + if (role === this.humanPrefix) { + return new HumanChatMessage(content); + } else if (role === this.aiPrefix) { + return new AIChatMessage(content); + } else { + // default to generic ChatMessage + return new ChatMessage(content, role); + } + }); + } + + if (this.returnMessages) { + const result = { + [this.memoryKey]: messages, + }; + return result; + } + const result = { + [this.memoryKey]: getBufferString( + messages, + this.humanPrefix, + this.aiPrefix + ), + }; + return result; + } + + async saveContext( + inputValues: InputValues, + outputValues: OutputValues + ): Promise { + const input = getInputValue(inputValues, this.inputKey); + const output = getInputValue(outputValues, this.outputKey); + + // Create new Memory and Message instances + const memory = new Memory({ + messages: [ + new Message({ + role: this.humanPrefix, + content: `${input}`, + }), + new Message({ + role: this.aiPrefix, + content: `${output}`, + }), + ], + }); + + // Add the new memory to the session using the ZepClient + if (this.sessionId) { + try { + await this.zepClient.addMemory(this.sessionId, memory); + } catch (error) { + console.error("Error adding memory: ", error); + } + } + + // Call the superclass's saveContext method + await super.saveContext(inputValues, outputValues); + } + + async clear(): Promise { + try { + await this.zepClient.deleteMemory(this.sessionId); + } catch (error) { + console.error("Error deleting session: ", error); + } + + // Clear the superclass's chat history + await super.clear(); + } +} diff --git a/langchain/src/retrievers/tests/zep.int.test.ts b/langchain/src/retrievers/tests/zep.int.test.ts new file mode 100644 index 000000000000..3c719d8154d2 --- /dev/null +++ b/langchain/src/retrievers/tests/zep.int.test.ts @@ -0,0 +1,18 @@ +/* eslint-disable no-process-env */ +/* eslint-disable @typescript-eslint/no-non-null-assertion */ +import { test, expect } from "@jest/globals"; + +import { ZepRetriever } from "../zep.js"; + +test.skip("ZepRetriever", async () => { + const baseURL = process.env.ZEP_API_URL || "http://localhost:8000"; + const sessionId = "your-session-id"; // Replace with the actual session ID + const topK = 3; // The number of documents to retrieve + const zepRetriever = new ZepRetriever({ sessionId, url: baseURL, topK }); + + const docs = await zepRetriever.getRelevantDocuments("hello"); + + expect(docs.length).toBeGreaterThan(0); + + console.log(docs); +}); diff --git a/langchain/src/retrievers/zep.ts b/langchain/src/retrievers/zep.ts new file mode 100644 index 000000000000..bc500deb1fec --- /dev/null +++ b/langchain/src/retrievers/zep.ts @@ -0,0 +1,57 @@ +import { ZepClient, SearchResult, SearchPayload } from "@getzep/zep-js"; +import { BaseRetriever } from "../schema/index.js"; +import { Document } from "../document.js"; + +export type ZepRetrieverConfig = { + sessionId: string; + url: string; + topK?: number; +}; + +export class ZepRetriever extends BaseRetriever { + private zepClient: ZepClient; + + private sessionId: string; + + private topK?: number; + + constructor(config: ZepRetrieverConfig) { + super(); + this.zepClient = new ZepClient(config.url); + this.sessionId = config.sessionId; + this.topK = config.topK; + } + + /** + * Converts an array of search results to an array of Document objects. + * @param {SearchResult[]} results - The array of search results. + * @returns {Document[]} An array of Document objects representing the search results. + */ + private searchResultToDoc(results: SearchResult[]): Document[] { + return results + .filter((r) => r.message) + .map( + ({ message: { content } = {}, ...metadata }, dist) => + new Document({ + pageContent: content ?? "", + metadata: { score: dist, ...metadata }, + }) + ); + } + + /** + * Retrieves the relevant documents based on the given query. + * @param {string} query - The query string. + * @returns {Promise} A promise that resolves to an array of relevant Document objects. + */ + async getRelevantDocuments(query: string): Promise { + const payload: SearchPayload = { text: query, meta: {} }; + const results: SearchResult[] = await this.zepClient.searchMemory( + this.sessionId, + payload, + this.topK + ); + + return this.searchResultToDoc(results); + } +} diff --git a/langchain/tsconfig.json b/langchain/tsconfig.json index 6e1ff6d209c0..2aceebd2e428 100644 --- a/langchain/tsconfig.json +++ b/langchain/tsconfig.json @@ -78,6 +78,7 @@ "src/vectorstores/tigris.ts", "src/text_splitter.ts", "src/memory/index.ts", + "src/memory/zep.ts", "src/document.ts", "src/docstore/index.ts", "src/document_loaders/base.ts", @@ -118,6 +119,7 @@ "src/output_parsers/expression.ts", "src/retrievers/remote/index.ts", "src/retrievers/supabase.ts", + "src/retrievers/zep.ts", "src/retrievers/metal.ts", "src/retrievers/databerry.ts", "src/retrievers/contextual_compression.ts", diff --git a/yarn.lock b/yarn.lock index b2f1930d83d7..950c4136cd57 100644 --- a/yarn.lock +++ b/yarn.lock @@ -5293,6 +5293,16 @@ __metadata: languageName: node linkType: hard +"@getzep/zep-js@npm:0.3.1, @getzep/zep-js@npm:^0.3.1": + version: 0.3.1 + resolution: "@getzep/zep-js@npm:0.3.1" + dependencies: + axios: ^1.4.0 + typescript: ^5.0.4 + checksum: 796df18baf3458d5e3887de9519b7d943e38d46f9f57b3515fa7fc5f8731271816b10d8868af6ef8e927345c409c5c865d8953781b88dd6889c159d6bfff9e48 + languageName: node + linkType: hard + "@gomomento/generated-types@npm:0.62.1": version: 0.62.1 resolution: "@gomomento/generated-types@npm:0.62.1" @@ -9689,6 +9699,17 @@ __metadata: languageName: node linkType: hard +"axios@npm:^1.4.0": + version: 1.4.0 + resolution: "axios@npm:1.4.0" + dependencies: + follow-redirects: ^1.15.0 + form-data: ^4.0.0 + proxy-from-env: ^1.1.0 + checksum: 7fb6a4313bae7f45e89d62c70a800913c303df653f19eafec88e56cea2e3821066b8409bc68be1930ecca80e861c52aa787659df0ffec6ad4d451c7816b9386b + languageName: node + linkType: hard + "axobject-query@npm:^3.1.1": version: 3.1.1 resolution: "axobject-query@npm:3.1.1" @@ -13952,6 +13973,7 @@ __metadata: dependencies: "@clickhouse/client": ^0.0.14 "@getmetal/metal-sdk": ^4.0.0 + "@getzep/zep-js": 0.3.1 "@gomomento/sdk": ^1.23.0 "@opensearch-project/opensearch": ^2.2.0 "@pinecone-database/pinecone": ^0.0.14 @@ -18252,6 +18274,7 @@ __metadata: "@clickhouse/client": ^0.0.14 "@faker-js/faker": ^7.6.0 "@getmetal/metal-sdk": ^4.0.0 + "@getzep/zep-js": ^0.3.1 "@gomomento/sdk": ^1.23.0 "@huggingface/inference": ^1.5.1 "@jest/globals": ^29.5.0 @@ -18337,6 +18360,7 @@ __metadata: "@aws-sdk/client-sagemaker-runtime": ^3.310.0 "@clickhouse/client": ^0.0.14 "@getmetal/metal-sdk": "*" + "@getzep/zep-js": ^0.3.1 "@gomomento/sdk": ^1.23.0 "@huggingface/inference": ^1.5.1 "@opensearch-project/opensearch": "*" @@ -18386,6 +18410,8 @@ __metadata: optional: true "@getmetal/metal-sdk": optional: true + "@getzep/zep-js": + optional: true "@gomomento/sdk": optional: true "@huggingface/inference": @@ -25881,7 +25907,7 @@ __metadata: languageName: node linkType: hard -"typescript@npm:^5.0.0": +"typescript@npm:^5.0.0, typescript@npm:^5.0.4": version: 5.0.4 resolution: "typescript@npm:5.0.4" bin: @@ -25911,7 +25937,7 @@ __metadata: languageName: node linkType: hard -"typescript@patch:typescript@^5.0.0#~builtin": +"typescript@patch:typescript@^5.0.0#~builtin, typescript@patch:typescript@^5.0.4#~builtin": version: 5.0.4 resolution: "typescript@patch:typescript@npm%3A5.0.4#~builtin::version=5.0.4&hash=1f5320" bin: