144 lines
5 KiB
JavaScript
144 lines
5 KiB
JavaScript
|
import { ZepClient } from "@getzep/zep-cloud";
|
||
|
import { NotFoundError, } from "@getzep/zep-cloud/api";
|
||
|
import { BaseRetriever } from "@langchain/core/retrievers";
|
||
|
import { Document } from "@langchain/core/documents";
|
||
|
/**
|
||
|
* 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";
|
||
|
}
|
||
|
get lc_secrets() {
|
||
|
return {
|
||
|
apiKey: "ZEP_API_KEY",
|
||
|
};
|
||
|
}
|
||
|
get lc_aliases() {
|
||
|
return { apiKey: "api_key" };
|
||
|
}
|
||
|
constructor(config) {
|
||
|
super(config);
|
||
|
Object.defineProperty(this, "lc_namespace", {
|
||
|
enumerable: true,
|
||
|
configurable: true,
|
||
|
writable: true,
|
||
|
value: ["langchain", "retrievers", "zep"]
|
||
|
});
|
||
|
Object.defineProperty(this, "client", {
|
||
|
enumerable: true,
|
||
|
configurable: true,
|
||
|
writable: true,
|
||
|
value: void 0
|
||
|
});
|
||
|
Object.defineProperty(this, "sessionId", {
|
||
|
enumerable: true,
|
||
|
configurable: true,
|
||
|
writable: true,
|
||
|
value: void 0
|
||
|
});
|
||
|
Object.defineProperty(this, "topK", {
|
||
|
enumerable: true,
|
||
|
configurable: true,
|
||
|
writable: true,
|
||
|
value: void 0
|
||
|
});
|
||
|
Object.defineProperty(this, "searchScope", {
|
||
|
enumerable: true,
|
||
|
configurable: true,
|
||
|
writable: true,
|
||
|
value: void 0
|
||
|
});
|
||
|
Object.defineProperty(this, "searchType", {
|
||
|
enumerable: true,
|
||
|
configurable: true,
|
||
|
writable: true,
|
||
|
value: void 0
|
||
|
});
|
||
|
Object.defineProperty(this, "mmrLambda", {
|
||
|
enumerable: true,
|
||
|
configurable: true,
|
||
|
writable: true,
|
||
|
value: void 0
|
||
|
});
|
||
|
Object.defineProperty(this, "filter", {
|
||
|
enumerable: true,
|
||
|
configurable: true,
|
||
|
writable: true,
|
||
|
value: void 0
|
||
|
});
|
||
|
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.
|
||
|
*/
|
||
|
searchMessageResultToDoc(results) {
|
||
|
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.
|
||
|
*/
|
||
|
searchSummaryResultToDoc(results) {
|
||
|
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) {
|
||
|
try {
|
||
|
const results = 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;
|
||
|
}
|
||
|
}
|
||
|
}
|