-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
zep_cloud.ts
161 lines (145 loc) Β· 4.89 KB
/
zep_cloud.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
import { ZepClient } from "@getzep/zep-cloud";
import {
SearchScope,
SearchType,
MemorySearchResult,
NotFoundError,
} from "@getzep/zep-cloud/api";
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} [apiKey] - The Zep Cloud Project Key.
* @argument {number} [topK] - The number of results to return.
* @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 ZepCloudRetrieverConfig extends BaseRetrieverInput {
sessionId: string;
topK?: number;
apiKey: string;
searchScope?: SearchScope;
searchType?: SearchType;
mmrLambda?: number;
filter?: Record<string, unknown>;
}
/**
* Class for retrieving information from a Zep Cloud long-term memory store.
* Extends the BaseRetriever class.
* @example
* ```typescript
* const retriever = new ZepCloudRetriever({
* apiKey: "<zep cloud project api key>",
* sessionId: "session_exampleUUID",
* topK: 3,
* });
* const query = "Can I drive red cars in France?";
* const docs = await retriever.getRelevantDocuments(query);
* ```
*/
export class ZepCloudRetriever extends BaseRetriever {
static lc_name() {
return "ZepRetriever";
}
lc_namespace = ["langchain", "retrievers", "zep"];
get lc_secrets(): { [key: string]: string } | undefined {
return {
apiKey: "ZEP_API_KEY",
};
}
get lc_aliases(): { [key: string]: string } | undefined {
return { apiKey: "api_key" };
}
client: ZepClient;
private sessionId: string;
private topK?: number;
private searchScope?: SearchScope;
private searchType?: SearchType;
private mmrLambda?: number;
private filter?: Record<string, unknown>;
constructor(config: ZepCloudRetrieverConfig) {
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.client = new ZepClient({ apiKey: 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 } = {},
score,
...rest
}) =>
new Document({
pageContent: content ?? "",
metadata: { score, ...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 } = {},
score,
...rest
}) =>
new Document({
pageContent: content ?? "",
metadata: { score, ...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[]> {
try {
const results: MemorySearchResult[] = await this.client.memory.search(
this.sessionId,
{
text: query,
metadata: this.filter,
searchScope: this.searchScope,
searchType: this.searchType,
mmrLambda: this.mmrLambda,
limit: 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;
}
}
}