@aksolab/recall
Version:
A memory management package for AI SDK memory functionality
309 lines (308 loc) • 12.5 kB
JavaScript
import { SchemaFieldTypes, VectorAlgorithms } from 'redis';
import { ArchiveProvider } from '../base';
import { embed } from 'ai';
import { openai } from '@ai-sdk/openai';
/**
* Utility function to set up Redis schema for archive storage.
* This should be called once during application initialization.
*/
export async function setupRedisSchema(client, indexName = 'idx:archive', collectionName = 'recall:memory:archive:', dimensions = 1536) {
try {
// Drop existing index if it exists
try {
await client.ft.dropIndex(indexName);
}
catch (e) {
if (!e.message.includes('Unknown Index')) {
throw e;
}
}
// Create index for archive entries
await client.ft.create(indexName, {
'$.name': {
type: SchemaFieldTypes.TEXT,
AS: 'name'
},
'$.content': {
type: SchemaFieldTypes.TEXT,
AS: 'content'
},
'$.timestamp': {
type: SchemaFieldTypes.NUMERIC,
SORTABLE: true,
AS: 'timestamp'
},
'$.metadata': {
type: SchemaFieldTypes.TEXT,
AS: 'metadata'
},
'$.embeddings': {
type: SchemaFieldTypes.VECTOR,
ALGORITHM: VectorAlgorithms.HNSW,
TYPE: 'FLOAT32',
DIM: dimensions,
DISTANCE_METRIC: 'COSINE',
AS: 'embeddings'
}
}, {
ON: 'JSON',
PREFIX: collectionName
});
}
catch (error) {
console.error('Error setting up Redis schema:', error);
throw error;
}
}
export class RedisArchiveProvider extends ArchiveProvider {
client;
indexName;
collectionName;
constructor(config) {
super(config);
this.client = config.client;
this.indexName = config.indexName || 'idx:archive';
this.collectionName = config.collectionName || 'recall:memory:archive:';
}
async initialize() {
try {
// Check if Redis Search is available
//const rawModules = await this.client.moduleList();
//const modules = (rawModules as unknown) as Array<{ name: string }>;
//const hasSearch = modules.some(module =>
// module?.name?.toLowerCase() === 'search' ||
// module?.name?.toLowerCase() === 'redisearch'
//);
//if (!hasSearch) {
// throw new Error(
// 'Redis Search module is not available. This provider requires Redis Stack or Redis with RediSearch module installed. ' +
// 'Please ensure you have the correct Redis version with Search capability enabled.'
// );
//}
// Check if index exists
try {
await this.client.ft.info(this.indexName);
}
catch (e) {
await setupRedisSchema(this.client, this.indexName, this.collectionName, this.config.dimensions);
if (e.message.includes('Unknown Index')) {
throw new Error(`Redis Search index '${this.indexName}' not found. Please run setupRedisSchema() before using the provider.`);
}
throw e;
}
}
catch (error) {
if (error instanceof Error &&
(error.message.includes('Redis Search module is not available') ||
error.message.includes('Redis Search index'))) {
console.error('\x1b[31mError:\x1b[0m', error.message);
}
else {
console.error('Error initializing Redis provider:', error);
}
throw error;
}
}
async cleanup() {
// No need to disconnect since we don't own the connection
return;
}
getKey(id) {
return `${this.collectionName}${id}`;
}
async generateEmbeddings(text) {
const { embedding } = await embed({
model: openai.embedding(this.config.embeddingModel || 'text-embedding-3-small'),
value: text,
});
return embedding;
}
async addEntry(entry) {
const timestamp = Date.now();
const id = `entry-${timestamp}-${Math.random().toString(36).slice(2)}`;
const embeddings = await this.generateEmbeddings(entry.content);
const fullEntry = {
...entry,
id,
embeddings,
timestamp
};
const key = this.getKey(id);
await this.client.json.set(key, '$', fullEntry);
return fullEntry;
}
async addEntries(entries) {
return Promise.all(entries.map(entry => this.addEntry(entry)));
}
async updateEntry(id, entry) {
const key = this.getKey(id);
const existing = await this.getEntry(id);
if (!existing) {
throw new Error(`Entry with id ${id} not found`);
}
const updatedEntry = {
...existing,
...entry
};
// If content was updated, regenerate embeddings
if (entry.content && entry.content !== existing.content) {
updatedEntry.embeddings = await this.generateEmbeddings(entry.content);
}
await this.client.json.set(key, '$', updatedEntry);
return updatedEntry;
}
async deleteEntry(id) {
await this.client.del(this.getKey(id));
}
async deleteEntriesByName(name) {
const results = await this.client.ft.search(this.indexName, `@name:(${name})`, { RETURN: ['name'] });
if (results.documents.length === 0) {
return 0;
}
await Promise.all(results.documents.map(doc => this.client.del(doc.id)));
return results.documents.length;
}
async searchByText(query, options = {}) {
const { limit = 20 } = options;
// Prepare search terms for text search - exact phrase match and individual terms
const terms = query.split(/\s+/).filter(term => term.length > 0);
const exactPhrase = `"${query}"`;
const fuzzyTerms = terms.map(term => `%${term}%`).join('|');
// Combine exact phrase and fuzzy term matching
const textQuery = `(@name|content:(${exactPhrase})) | (@name|content:(${fuzzyTerms}))`;
const results = await this.client.ft.search(this.indexName, textQuery, {
LIMIT: { from: options.offset || 0, size: limit },
RETURN: ['name', 'content', 'metadata', 'timestamp'],
SORTBY: 'timestamp',
DIALECT: 2,
});
if (results.documents.length === 0) {
return [];
}
return results.documents.map((doc) => {
const entry = doc.value;
const content = entry.content.toLowerCase();
const exactPhraseMatch = content.includes(query.toLowerCase());
const termMatches = terms.filter(term => content.includes(term.toLowerCase()));
const score = exactPhraseMatch ? 1.0 : termMatches.length / terms.length;
return {
entry: { ...entry, id: doc.id.replace(this.collectionName, '') },
score: score * 100,
matches: {
exactPhrase: exactPhraseMatch,
terms: termMatches
}
};
});
}
async searchBySimilarity(query, options = {}) {
const { limit = 20 } = options;
const embedding = await this.generateEmbeddings(query);
const vectorQuery = Buffer.from(new Float32Array(embedding).buffer);
const results = await this.client.ft.search(this.indexName, '*=>[KNN 20 @embeddings $vec_query AS vector_score]', {
PARAMS: { vec_query: vectorQuery },
RETURN: ['name', 'content', 'metadata', 'timestamp', 'vector_score'],
SORTBY: 'vector_score',
DIALECT: 2,
LIMIT: { from: options.offset || 0, size: limit }
});
if (results.documents.length === 0) {
return [];
}
// Find min and max scores for normalization
const scores = results.documents.map((doc) => Number(doc.value.vector_score || 0));
const minScore = Math.min(...scores);
const maxScore = Math.max(...scores);
const range = maxScore - minScore;
return results.documents.map((doc) => {
const entry = doc.value;
const rawScore = Number(doc.value.vector_score || 0);
const normalizedScore = range === 0 ? 1 : (maxScore - rawScore) / range;
const score = normalizedScore * 100;
const { vector_score, ...entryWithoutScore } = entry;
return {
entry: { ...entryWithoutScore, id: doc.id.replace(this.collectionName, '') },
score
};
});
}
async hybridSearch(query, options = {}) {
const { limit = 20, vectorWeight = 0.7, textWeight = 0.3 } = options;
const embedding = await this.generateEmbeddings(query);
const vectorQuery = Buffer.from(new Float32Array(embedding).buffer);
const results = await this.client.ft.search(this.indexName, '*=>[KNN 20 @embeddings $vec_query AS vector_score]', {
PARAMS: { vec_query: vectorQuery },
RETURN: ['name', 'content', 'metadata', 'timestamp', 'vector_score'],
DIALECT: 2,
LIMIT: { from: options.offset || 0, size: limit }
});
if (results.documents.length === 0) {
return [];
}
// Find min and max scores for normalization
const scores = results.documents.map((doc) => Number(doc.value.vector_score || 0));
const minScore = Math.min(...scores);
const maxScore = Math.max(...scores);
const range = maxScore - minScore;
const processedResults = results.documents.map((doc) => {
const entry = doc.value;
const rawVectorScore = Number(doc.value.vector_score || 0);
const normalizedVectorScore = range === 0 ? 1 : (maxScore - rawVectorScore) / range;
const content = entry.content.toLowerCase();
const searchTerms = query.toLowerCase().split(/\s+/);
const exactPhraseMatch = content.includes(query.toLowerCase());
const termMatches = searchTerms.filter(term => content.includes(term));
let textScore = 0;
if (exactPhraseMatch) {
textScore = 1.0;
}
else if (termMatches.length > 0) {
textScore = termMatches.length / searchTerms.length;
}
const { vector_score, ...entryWithoutScore } = entry;
const finalScore = (normalizedVectorScore * vectorWeight + textScore * textWeight) * 100;
return {
entry: { ...entryWithoutScore, id: doc.id.replace(this.collectionName, '') },
score: finalScore,
matches: {
exactPhrase: exactPhraseMatch,
terms: termMatches
}
};
});
return processedResults
.filter(result => !options.minScore || result.score >= options.minScore)
.sort((a, b) => b.score - a.score)
.slice(0, limit);
}
async listEntries(options = {}) {
const { limit = 100, offset = 0 } = options;
const results = await this.client.ft.search(this.indexName, '*', {
LIMIT: { from: offset, size: limit },
SORTBY: 'timestamp',
DESC: true,
});
return results.documents.map(doc => {
const value = doc.value;
return {
...value,
id: doc.id.replace(this.collectionName, '')
};
});
}
async getEntry(id) {
const key = this.getKey(id);
const entry = await this.client.json.get(key);
return entry ? { ...entry, id } : null;
}
async clear() {
const keys = await this.client.keys(`${this.collectionName}*`);
if (keys.length > 0) {
await this.client.del(keys);
}
}
async count() {
const info = await this.client.ft.info(this.indexName);
return Number(info.numDocs);
}
}