-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
zep.ts
169 lines (153 loc) Β· 5.16 KB
/
zep.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
import {
MemorySearchPayload,
MemorySearchResult,
NotFoundError,
ZepClient,
} from "@getzep/zep-js";
import { BaseRetriever, BaseRetrieverInput } from "@langchain/core/retrievers";
import { Document } from "@langchain/core/documents";
/**
* Configuration interface for the ZepRetriever class. Extends the
* BaseRetrieverInput interface.
*
* @argument {string} sessionId - The ID of the Zep session.
* @argument {string} url - The URL of the Zep API.
* @argument {number} [topK] - The number of results to return.
* @argument {string} [apiKey] - The API key for the Zep API.
* @argument [searchScope] [searchScope] - The scope of the search: "messages" or "summary".
* @argument [searchType] [searchType] - The type of search to perform: "similarity" or "mmr".
* @argument {number} [mmrLambda] - The lambda value for the MMR search.
* @argument {Record<string, unknown>} [filter] - The metadata filter to apply to the search.
*/
export interface ZepRetrieverConfig extends BaseRetrieverInput {
sessionId: string;
url: string;
topK?: number;
apiKey?: string;
searchScope?: "messages" | "summary";
searchType?: "similarity" | "mmr";
mmrLambda?: number;
filter?: Record<string, unknown>;
}
/**
* Class for retrieving information from a Zep long-term memory store.
* Extends the BaseRetriever class.
* @example
* ```typescript
* const retriever = new ZepRetriever({
* url: "http:
* sessionId: "session_exampleUUID",
* topK: 3,
* });
* const query = "Can I drive red cars in France?";
* const docs = await retriever.getRelevantDocuments(query);
* ```
*/
export class ZepRetriever extends BaseRetriever {
static lc_name() {
return "ZepRetriever";
}
lc_namespace = ["langchain", "retrievers", "zep"];
get lc_secrets(): { [key: string]: string } | undefined {
return {
apiKey: "ZEP_API_KEY",
url: "ZEP_API_URL",
};
}
get lc_aliases(): { [key: string]: string } | undefined {
return { apiKey: "api_key" };
}
zepClientPromise: Promise<ZepClient>;
private sessionId: string;
private topK?: number;
private searchScope?: "messages" | "summary";
private searchType?: "similarity" | "mmr";
private mmrLambda?: number;
private filter?: Record<string, unknown>;
constructor(config: ZepRetrieverConfig) {
super(config);
this.sessionId = config.sessionId;
this.topK = config.topK;
this.searchScope = config.searchScope;
this.searchType = config.searchType;
this.mmrLambda = config.mmrLambda;
this.filter = config.filter;
this.zepClientPromise = ZepClient.init(config.url, config.apiKey);
}
/**
* Converts an array of message search results to an array of Document objects.
* @param {MemorySearchResult[]} results - The array of search results.
* @returns {Document[]} An array of Document objects representing the search results.
*/
private searchMessageResultToDoc(results: MemorySearchResult[]): Document[] {
return results
.filter((r) => r.message)
.map(
({
message: { content, metadata: messageMetadata } = {},
dist,
...rest
}) =>
new Document({
pageContent: content ?? "",
metadata: { score: dist, ...messageMetadata, ...rest },
})
);
}
/**
* Converts an array of summary search results to an array of Document objects.
* @param {MemorySearchResult[]} results - The array of search results.
* @returns {Document[]} An array of Document objects representing the search results.
*/
private searchSummaryResultToDoc(results: MemorySearchResult[]): Document[] {
return results
.filter((r) => r.summary)
.map(
({
summary: { content, metadata: summaryMetadata } = {},
dist,
...rest
}) =>
new Document({
pageContent: content ?? "",
metadata: { score: dist, ...summaryMetadata, ...rest },
})
);
}
/**
* Retrieves the relevant documents based on the given query.
* @param {string} query - The query string.
* @returns {Promise<Document[]>} A promise that resolves to an array of relevant Document objects.
*/
async _getRelevantDocuments(query: string): Promise<Document[]> {
const payload: MemorySearchPayload = {
text: query,
metadata: this.filter,
search_scope: this.searchScope,
search_type: this.searchType,
mmr_lambda: this.mmrLambda,
};
// Wait for ZepClient to be initialized
const zepClient = await this.zepClientPromise;
if (!zepClient) {
throw new Error("ZepClient is not initialized");
}
try {
const results: MemorySearchResult[] = await zepClient.memory.searchMemory(
this.sessionId,
payload,
this.topK
);
return this.searchScope === "summary"
? this.searchSummaryResultToDoc(results)
: this.searchMessageResultToDoc(results);
} catch (error) {
// eslint-disable-next-line no-instanceof/no-instanceof
if (error instanceof NotFoundError) {
return Promise.resolve([]); // Return an empty Document array
}
// If it's not a NotFoundError, throw the error again
throw error;
}
}
}