mongodb-rag-core
Version:
Common elements used by MongoDB Chatbot Framework components.
323 lines • 13.1 kB
JavaScript
"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.makeMongoDbEmbeddedContentStore = exports.listDataSourcesCache = void 0;
const _1 = require(".");
const MongoDbDatabaseConnection_1 = require("../MongoDbDatabaseConnection");
const assert_1 = require("assert");
const mongodb_1 = require("mongodb");
function makeMatchQuery({ sourceNames, chunkAlgoHash }) {
const operator = chunkAlgoHash.operation === "equals" ? "$eq" : "$ne";
return {
chunkAlgoHash: { [operator]: chunkAlgoHash.hashValue },
// run on specific source names if specified, run on all if not
...(sourceNames
? {
sourceName: {
$in: sourceNames,
},
}
: undefined),
};
}
/**
24-hour cache of listDataSources aggregation as query is a full scan of all documents in collection
*/
exports.listDataSourcesCache = {
data: null,
expiresAt: 0,
isRefreshing: false,
};
const CACHE_STALE_AGE = 24 * 60 * 60 * 1000; // 24 hours
const CACHE_MAX_AGE = 1000 * 60 * 60 * 24 * 7; // 7 days
function makeMongoDbEmbeddedContentStore({ connectionUri, databaseName, searchIndex: { embeddingName, numDimensions = 1536, filters = [
{
type: "filter",
path: "sourceName",
},
{
type: "filter",
path: "metadata.version.label",
},
{
type: "filter",
path: "metadata.version.isCurrent",
},
{
type: "filter",
path: "sourceType",
},
{
type: "filter",
path: "url",
},
], name = "vector_index", }, collectionName = "embedded_content", }) {
const { mongoClient, db, drop, close } = (0, MongoDbDatabaseConnection_1.makeMongoDbDatabaseConnection)({
connectionUri,
databaseName,
});
const embeddedContentCollection = db.collection(collectionName);
const embeddingPath = `embeddings.${embeddingName}`;
async function fetchFreshListDataSources() {
const freshData = await embeddedContentCollection
.aggregate([
{
$group: {
_id: "$sourceName",
versions: {
$addToSet: {
$cond: [
{ $ifNull: ["$metadata.version.label", false] },
{
label: "$metadata.version.label",
isCurrent: "$metadata.version.isCurrent",
},
"$$REMOVE",
],
},
},
sourceType: { $addToSet: "$sourceType" },
},
},
{
$project: {
_id: 0,
id: "$_id",
versions: {
$map: {
input: {
$filter: {
input: "$versions",
as: "v",
cond: { $ne: ["$$v.label", null] },
},
},
as: "v",
in: {
label: "$$v.label",
isCurrent: { $ifNull: ["$$v.isCurrent", false] },
},
},
},
type: {
$arrayElemAt: [
{
$filter: {
input: "$sourceType",
as: "t",
cond: { $ne: ["$$t", null] },
},
},
0,
],
},
},
},
])
.toArray();
exports.listDataSourcesCache.data = freshData;
exports.listDataSourcesCache.expiresAt = Date.now() + CACHE_STALE_AGE;
exports.listDataSourcesCache.isRefreshing = false;
return freshData;
}
return {
drop,
close,
metadata: {
databaseName,
collectionName,
embeddingName,
embeddingPath,
},
async loadEmbeddedContent({ page }) {
return await embeddedContentCollection.find((0, _1.pageIdentity)(page)).toArray();
},
async deleteEmbeddedContent({ page, dataSources, inverseDataSources = false, }) {
const deleteResult = await embeddedContentCollection.deleteMany({
...(page ? (0, _1.pageIdentity)(page) : undefined),
...(dataSources
? {
sourceName: {
[inverseDataSources ? "$nin" : "$in"]: dataSources,
},
}
: undefined),
});
if (!deleteResult.acknowledged) {
throw new Error("EmbeddedContent deletion not acknowledged!");
}
},
async updateEmbeddedContent({ page, embeddedContent }) {
(0, assert_1.strict)(embeddedContent.length !== 0);
embeddedContent.forEach((embeddedContent) => {
(0, assert_1.strict)(embeddedContent.sourceName === page.sourceName &&
embeddedContent.url === page.url, `EmbeddedContent source/url (${embeddedContent.sourceName} / ${embeddedContent.url}) must match give page source/url (${page.sourceName} / ${page.url})!`);
});
await mongoClient.withSession(async (session) => {
await session.withTransaction(async () => {
// First delete all the embeddedContent for the given page
const deleteResult = await embeddedContentCollection.deleteMany((0, _1.pageIdentity)(page), { session });
if (!deleteResult.acknowledged) {
throw new Error("EmbeddedContent deletion not acknowledged!");
}
// Insert the embedded content for the page
const insertResult = await embeddedContentCollection.insertMany([...embeddedContent], {
session,
});
if (!insertResult.acknowledged) {
throw new Error("EmbeddedContent insertion not acknowledged!");
}
const { insertedCount } = insertResult;
if (insertedCount !== embeddedContent.length) {
throw new Error(`Expected ${embeddedContent.length} inserted, got ${insertedCount}`);
}
});
});
},
/**
@param vector - The vector to search for nearest neighbors to.
@param options - Options for performing a nearest-neighbor search.
*/
async findNearestNeighbors(vector, options) {
const { indexName, path, k, minScore, filter = {}, numCandidates, } = {
// Default options
indexName: name,
path: embeddingPath,
k: 3,
minScore: 0,
// User options override
...(options ?? {}),
};
return embeddedContentCollection
.aggregate([
{
$vectorSearch: {
index: indexName,
queryVector: vector,
path,
limit: k,
numCandidates: numCandidates ?? k * 15,
filter: handleFilters(filter),
},
},
{
$addFields: {
score: {
$meta: "vectorSearchScore",
},
},
},
{ $match: { score: { $gte: minScore } } },
])
.toArray();
},
async init() {
await embeddedContentCollection.createIndex({ sourceName: 1 });
await embeddedContentCollection.createIndex({ url: 1 });
await embeddedContentCollection.createIndex({
"metadata.version.isCurrent": 1,
});
await embeddedContentCollection.createIndex({
"metadata.version.label": 1,
});
await embeddedContentCollection.createIndex({
sourceType: 1,
});
try {
const searchIndex = {
name,
type: "vectorSearch",
definition: {
fields: [
{
numDimensions,
path: embeddingPath,
similarity: "cosine",
type: "vector",
},
...filters,
],
},
};
await embeddedContentCollection.createSearchIndex(searchIndex);
}
catch (error) {
if (error instanceof mongodb_1.MongoServerError) {
(0, assert_1.strict)(error.codeName === "IndexAlreadyExists", `An unexpected MongoError occurred: ${error.name}`);
}
else {
throw error;
}
}
},
async listDataSources() {
const now = Date.now();
// If cache is fresh (< 24h), return it immediately
if (exports.listDataSourcesCache.data && now < exports.listDataSourcesCache.expiresAt) {
return exports.listDataSourcesCache.data;
}
// If cache exists but is stale (< 7 days), return it and refresh in background
if (exports.listDataSourcesCache.data &&
now - exports.listDataSourcesCache.expiresAt < CACHE_MAX_AGE) {
if (!exports.listDataSourcesCache.isRefreshing) {
exports.listDataSourcesCache.isRefreshing = true;
void fetchFreshListDataSources().catch((err) => {
exports.listDataSourcesCache.isRefreshing = false;
console.error("Error refreshing listDataSources cache:", err);
});
}
return exports.listDataSourcesCache.data;
}
// Cache is too old (>= 7 days) — fetch fresh and set cache
return await fetchFreshListDataSources();
},
async getDataSources(matchQuery) {
const result = await embeddedContentCollection
.aggregate([
{ $match: makeMatchQuery(matchQuery) },
{
$group: {
_id: null,
uniqueSources: { $addToSet: "$sourceName" },
},
},
{ $project: { _id: 0, uniqueSources: 1 } },
])
.toArray();
const uniqueSources = result.length > 0 ? result[0].uniqueSources : [];
return uniqueSources;
},
};
}
exports.makeMongoDbEmbeddedContentStore = makeMongoDbEmbeddedContentStore;
const handleFilters = (filter) => {
const vectorSearchFilter = {};
if (filter.sourceName) {
vectorSearchFilter["sourceName"] = Array.isArray(filter.sourceName)
? { $in: filter.sourceName }
: filter.sourceName;
}
if (filter.sourceType) {
vectorSearchFilter["sourceType"] = Array.isArray(filter.sourceType)
? { $in: filter.sourceType }
: filter.sourceType;
}
// Handle version filter. Note: unversioned embeddings (isCurrent: null) are treated as current
const { current, label } = filter.version ?? {};
if (label) {
vectorSearchFilter["metadata.version.label"] = Array.isArray(label)
? { $in: label }
: label;
}
// Return current embeddings if either:
// 1. current=true was explicitly requested, or
// 2. [Default] no version filters were specified (current and label are undefined)
else if (current === true || current === undefined) {
vectorSearchFilter["metadata.version.isCurrent"] = { $ne: false }; // Include unversioned embeddings
}
// Only find embeddings that are explicitly marked as non-current (isCurrent: false)
else if (current === false) {
vectorSearchFilter["metadata.version.isCurrent"] = false;
}
return vectorSearchFilter;
};
//# sourceMappingURL=MongoDbEmbeddedContentStore.js.map