Spaces:
Running
Running
File size: 6,546 Bytes
1b9697f 0e5c445 5da61b4 2772555 82fcab7 2e28042 922b1b2 e4a770a 992a8de 634bd69 4dbcbb6 a4c3fca ddbe7d2 cad3e14 8bb1351 4dbcbb6 8bb1351 cad3e14 1b9697f cad3e14 ddbe7d2 cad3e14 4dbcbb6 cad3e14 4dbcbb6 a4c3fca 4dbcbb6 a4c3fca 4dbcbb6 634bd69 4dbcbb6 cad3e14 4dbcbb6 2e28042 634bd69 992a8de a4c3fca 992a8de 2e28042 e4a770a 922b1b2 4dbcbb6 cad3e14 5da61b4 2e28042 ddbe7d2 2e28042 ddbe7d2 e6addfc ddbe7d2 634bd69 ddbe7d2 634bd69 ddbe7d2 634bd69 ddbe7d2 634bd69 ddbe7d2 634bd69 ddbe7d2 881070b ddbe7d2 a4c3fca ddbe7d2 4dbcbb6 ddbe7d2 cad3e14 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 |
import { MONGODB_URL, MONGODB_DB_NAME, MONGODB_DIRECT_CONNECTION } from "$env/static/private";
import { GridFSBucket, MongoClient } from "mongodb";
import type { Conversation } from "$lib/types/Conversation";
import type { SharedConversation } from "$lib/types/SharedConversation";
import type { AbortedGeneration } from "$lib/types/AbortedGeneration";
import type { Settings } from "$lib/types/Settings";
import type { User } from "$lib/types/User";
import type { MessageEvent } from "$lib/types/MessageEvent";
import type { Session } from "$lib/types/Session";
import type { Assistant } from "$lib/types/Assistant";
import type { Report } from "$lib/types/Report";
import type { ConversationStats } from "$lib/types/ConversationStats";
import type { MigrationResult } from "$lib/types/MigrationResult";
import type { Semaphore } from "$lib/types/Semaphore";
import type { AssistantStats } from "$lib/types/AssistantStats";
import { logger } from "$lib/server/logger";
if (!MONGODB_URL) {
throw new Error(
"Please specify the MONGODB_URL environment variable inside .env.local. Set it to mongodb://localhost:27017 if you are running MongoDB locally, or to a MongoDB Atlas free instance for example."
);
}
export const CONVERSATION_STATS_COLLECTION = "conversations.stats";
const client = new MongoClient(MONGODB_URL, {
directConnection: MONGODB_DIRECT_CONNECTION === "true",
});
export const connectPromise = client.connect().catch(logger.error);
export function getCollections(mongoClient: MongoClient) {
const db = mongoClient.db(MONGODB_DB_NAME + (import.meta.env.MODE === "test" ? "-test" : ""));
const conversations = db.collection<Conversation>("conversations");
const conversationStats = db.collection<ConversationStats>(CONVERSATION_STATS_COLLECTION);
const assistants = db.collection<Assistant>("assistants");
const assistantStats = db.collection<AssistantStats>("assistants.stats");
const reports = db.collection<Report>("reports");
const sharedConversations = db.collection<SharedConversation>("sharedConversations");
const abortedGenerations = db.collection<AbortedGeneration>("abortedGenerations");
const settings = db.collection<Settings>("settings");
const users = db.collection<User>("users");
const sessions = db.collection<Session>("sessions");
const messageEvents = db.collection<MessageEvent>("messageEvents");
const bucket = new GridFSBucket(db, { bucketName: "files" });
const migrationResults = db.collection<MigrationResult>("migrationResults");
const semaphores = db.collection<Semaphore>("semaphores");
return {
conversations,
conversationStats,
assistants,
assistantStats,
reports,
sharedConversations,
abortedGenerations,
settings,
users,
sessions,
messageEvents,
bucket,
migrationResults,
semaphores,
};
}
const db = client.db(MONGODB_DB_NAME + (import.meta.env.MODE === "test" ? "-test" : ""));
const collections = getCollections(client);
const {
conversations,
conversationStats,
assistants,
assistantStats,
reports,
sharedConversations,
abortedGenerations,
settings,
users,
sessions,
messageEvents,
semaphores,
} = collections;
export { client, db, collections };
client.on("open", () => {
conversations
.createIndex(
{ sessionId: 1, updatedAt: -1 },
{ partialFilterExpression: { sessionId: { $exists: true } } }
)
.catch(logger.error);
conversations
.createIndex(
{ userId: 1, updatedAt: -1 },
{ partialFilterExpression: { userId: { $exists: true } } }
)
.catch(logger.error);
conversations
.createIndex(
{ "message.id": 1, "message.ancestors": 1 },
{ partialFilterExpression: { userId: { $exists: true } } }
)
.catch(logger.error);
// To do stats on conversations
conversations.createIndex({ updatedAt: 1 }).catch(logger.error);
// Not strictly necessary, could use _id, but more convenient. Also for stats
conversations.createIndex({ createdAt: 1 }).catch(logger.error);
// To do stats on conversation messages
conversations.createIndex({ "messages.createdAt": 1 }, { sparse: true }).catch(logger.error);
// Unique index for stats
conversationStats
.createIndex(
{
type: 1,
"date.field": 1,
"date.span": 1,
"date.at": 1,
distinct: 1,
},
{ unique: true }
)
.catch(logger.error);
// Allow easy check of last computed stat for given type/dateField
conversationStats
.createIndex({
type: 1,
"date.field": 1,
"date.at": 1,
})
.catch(logger.error);
abortedGenerations.createIndex({ updatedAt: 1 }, { expireAfterSeconds: 30 }).catch(logger.error);
abortedGenerations.createIndex({ conversationId: 1 }, { unique: true }).catch(logger.error);
sharedConversations.createIndex({ hash: 1 }, { unique: true }).catch(logger.error);
settings.createIndex({ sessionId: 1 }, { unique: true, sparse: true }).catch(logger.error);
settings.createIndex({ userId: 1 }, { unique: true, sparse: true }).catch(logger.error);
settings.createIndex({ assistants: 1 }).catch(logger.error);
users.createIndex({ hfUserId: 1 }, { unique: true }).catch(logger.error);
users.createIndex({ sessionId: 1 }, { unique: true, sparse: true }).catch(logger.error);
// No unicity because due to renames & outdated info from oauth provider, there may be the same username on different users
users.createIndex({ username: 1 }).catch(logger.error);
messageEvents.createIndex({ createdAt: 1 }, { expireAfterSeconds: 60 }).catch(logger.error);
sessions.createIndex({ expiresAt: 1 }, { expireAfterSeconds: 0 }).catch(logger.error);
sessions.createIndex({ sessionId: 1 }, { unique: true }).catch(logger.error);
assistants.createIndex({ createdById: 1, userCount: -1 }).catch(logger.error);
assistants.createIndex({ userCount: 1 }).catch(logger.error);
assistants.createIndex({ featured: 1, userCount: -1 }).catch(logger.error);
assistants.createIndex({ modelId: 1, userCount: -1 }).catch(logger.error);
assistants.createIndex({ searchTokens: 1 }).catch(logger.error);
assistants.createIndex({ last24HoursCount: 1 }).catch(logger.error);
assistantStats
// Order of keys is important for the queries
.createIndex({ "date.span": 1, "date.at": 1, assistantId: 1 }, { unique: true })
.catch(logger.error);
reports.createIndex({ assistantId: 1 }).catch(logger.error);
reports.createIndex({ createdBy: 1, assistantId: 1 }).catch(logger.error);
// Unique index for semaphore and migration results
semaphores.createIndex({ key: 1 }, { unique: true }).catch(logger.error);
semaphores.createIndex({ createdAt: 1 }, { expireAfterSeconds: 60 }).catch(logger.error);
});
|