Function cosineSimilaritySearch

packages/baseai/src/utils/memory/db/lib.ts:215–241

({
	chunks,
	queryEmbedding,
	topK
}: {
	chunks: MemoryChunk[];
	queryEmbedding: number[];
	topK: number;
})

}

export function cosineSimilaritySearch({
	chunks,
	queryEmbedding,
	topK
}: {
	chunks: MemoryChunk[];
	queryEmbedding: number[];
	topK: number;
}): SimilarChunk[] {
	// Return empty array if topK is invalid or no chunks are provided
	if (topK <= 0 || chunks.length === 0) {
		return [];
	}

	// Calculate cosine similarity for each chunk
	const chunksWithCosineSimilarity: SimilarChunk[] = chunks.map(chunk => ({
		text: chunk.text,
		attributes: chunk.attributes,
		similarity: cosineSimilarity(chunk.embedding, queryEmbedding) || 0
	}));

	// Sort chunks by similarity (descending order)
	chunksWithCosineSimilarity.sort((a, b) => b.similarity - a.similarity);

	// Return top K results
	return chunksWithCosineSimilarity.slice(0, topK);
}

// Get the path of the db file for a memory
export function getMemoryDbPath(memoryName: string): string {