import { BaseRetriever } from "@langchain/core/retrievers";
export const LAST_ACCESSED_AT_KEY = "last_accessed_at";
export const BUFFER_IDX = "buffer_idx";
/**
* TimeWeightedVectorStoreRetriever retrieves documents based on their time-weighted relevance.
* ref: https://github.com/langchain-ai/langchain/blob/master/libs/langchain/langchain/retrievers/time_weighted_retriever.py
* @example
* ```typescript
* const retriever = new TimeWeightedVectorStoreRetriever({
* vectorStore: new MemoryVectorStore(new OpenAIEmbeddings()),
* memoryStream: [],
* searchKwargs: 2,
* });
* await retriever.addDocuments([
* { pageContent: "My name is John.", metadata: {} },
* { pageContent: "My favourite food is pizza.", metadata: {} },
*
* ]);
* const results = await retriever.getRelevantDocuments(
* "What is my favourite food?",
* );
* ```
*/
export class TimeWeightedVectorStoreRetriever extends BaseRetriever {
static lc_name() {
return "TimeWeightedVectorStoreRetriever";
}
get lc_namespace() {
return ["langchain", "retrievers", "time_weighted"];
}
/**
* Constructor to initialize the required fields
* @param fields - The fields required for initializing the TimeWeightedVectorStoreRetriever
*/
constructor(fields) {
super(fields);
/**
* The vectorstore to store documents and determine salience.
*/
Object.defineProperty(this, "vectorStore", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
/**
* The number of top K most relevant documents to consider when searching.
*/
Object.defineProperty(this, "searchKwargs", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
/**
* The memory_stream of documents to search through.
*/
Object.defineProperty(this, "memoryStream", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
/**
* The exponential decay factor used as (1.0-decay_rate)**(hrs_passed).
*/
Object.defineProperty(this, "decayRate", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
/**
* The maximum number of documents to retrieve in a given call.
*/
Object.defineProperty(this, "k", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
/**
* Other keys in the metadata to factor into the score, e.g. 'importance'.
*/
Object.defineProperty(this, "otherScoreKeys", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
/**
* The salience to assign memories not retrieved from the vector store.
*/
Object.defineProperty(this, "defaultSalience", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
this.vectorStore = fields.vectorStore;
this.searchKwargs = fields.searchKwargs ?? 100;
this.memoryStream = fields.memoryStream ?? [];
this.decayRate = fields.decayRate ?? 0.01;
this.k = fields.k ?? 4;
this.otherScoreKeys = fields.otherScoreKeys ?? [];
this.defaultSalience = fields.defaultSalience ?? null;
}
/**
* Get the memory stream of documents.
* @returns The memory stream of documents.
*/
getMemoryStream() {
return this.memoryStream;
}
/**
* Set the memory stream of documents.
* @param memoryStream The new memory stream of documents.
*/
setMemoryStream(memoryStream) {
this.memoryStream = memoryStream;
}
/**
* Get relevant documents based on time-weighted relevance
* @param query - The query to search for
* @returns The relevant documents
*/
async _getRelevantDocuments(query, runManager) {
const now = Math.floor(Date.now() / 1000);
const memoryDocsAndScores = this.getMemoryDocsAndScores();
const salientDocsAndScores = await this.getSalientDocuments(query, runManager);
const docsAndScores = { ...memoryDocsAndScores, ...salientDocsAndScores };
return this.computeResults(docsAndScores, now);
}
/**
* NOTE: When adding documents to a vector store, use addDocuments
* via retriever instead of directly to the vector store.
* This is because it is necessary to process the document
* in prepareDocuments.
*
* @param docs - The documents to add to vector store in the retriever
*/
async addDocuments(docs) {
const now = Math.floor(Date.now() / 1000);
const savedDocs = this.prepareDocuments(docs, now);
this.memoryStream.push(...savedDocs);
await this.vectorStore.addDocuments(savedDocs);
}
/**
* Get memory documents and their scores
* @returns An object containing memory documents and their scores
*/
getMemoryDocsAndScores() {
const memoryDocsAndScores = {};
for (const doc of this.memoryStream.slice(-this.k)) {
const bufferIdx = doc.metadata[BUFFER_IDX];
if (bufferIdx === undefined) {
throw new Error(`Found a document in the vector store that is missing required metadata. This retriever only supports vector stores with documents that have been added through the "addDocuments" method on a TimeWeightedVectorStoreRetriever, not directly added or loaded into the backing vector store.`);
}
memoryDocsAndScores[bufferIdx] = {
doc,
score: this.defaultSalience ?? 0,
};
}
return memoryDocsAndScores;
}
/**
* Get salient documents and their scores based on the query
* @param query - The query to search for
* @returns An object containing salient documents and their scores
*/
async getSalientDocuments(query, runManager) {
const docAndScores = await this.vectorStore.similaritySearchWithScore(query, this.searchKwargs, undefined, runManager?.getChild());
const results = {};
for (const [fetchedDoc, score] of docAndScores) {
const bufferIdx = fetchedDoc.metadata[BUFFER_IDX];
if (bufferIdx === undefined) {
throw new Error(`Found a document in the vector store that is missing required metadata. This retriever only supports vector stores with documents that have been added through the "addDocuments" method on a TimeWeightedVectorStoreRetriever, not directly added or loaded into the backing vector store.`);
}
const doc = this.memoryStream[bufferIdx];
results[bufferIdx] = { doc, score };
}
return results;
}
/**
* Compute the final result set of documents based on the combined scores
* @param docsAndScores - An object containing documents and their scores
* @param now - The current timestamp
* @returns The final set of documents
*/
computeResults(docsAndScores, now) {
const recordedDocs = Object.values(docsAndScores)
.map(({ doc, score }) => ({
doc,
score: this.getCombinedScore(doc, score, now),
}))
.sort((a, b) => b.score - a.score);
const results = [];
for (const { doc } of recordedDocs) {
const bufferedDoc = this.memoryStream[doc.metadata[BUFFER_IDX]];
bufferedDoc.metadata[LAST_ACCESSED_AT_KEY] = now;
results.push(bufferedDoc);
if (results.length > this.k) {
break;
}
}
return results;
}
/**
* Prepare documents with necessary metadata before saving
* @param docs - The documents to prepare
* @param now - The current timestamp
* @returns The prepared documents
*/
prepareDocuments(docs, now) {
return docs.map((doc, i) => ({
...doc,
metadata: {
...doc.metadata,
[LAST_ACCESSED_AT_KEY]: doc.metadata[LAST_ACCESSED_AT_KEY] ?? now,
created_at: doc.metadata.created_at ?? now,
[BUFFER_IDX]: this.memoryStream.length + i,
},
}));
}
/**
* Calculate the combined score based on vector relevance and other factors
* @param doc - The document to calculate the score for
* @param vectorRelevance - The relevance score from the vector store
* @param nowMsec - The current timestamp in milliseconds
* @returns The combined score for the document
*/
getCombinedScore(doc, vectorRelevance, nowMsec) {
const hoursPassed = this.getHoursPassed(nowMsec, doc.metadata[LAST_ACCESSED_AT_KEY]);
let score = (1.0 - this.decayRate) ** hoursPassed;
for (const key of this.otherScoreKeys) {
score += doc.metadata[key];
}
if (vectorRelevance !== null) {
score += vectorRelevance;
}
return score;
}
/**
* Calculate the hours passed between two time points
* @param time - The current time in seconds
* @param refTime - The reference time in seconds
* @returns The number of hours passed between the two time points
*/
getHoursPassed(time, refTime) {
return (time - refTime) / 3600;
}
}