feat(embeddingModel): add embedding model into mongodb

.env

@@ -176,4 +176,6 @@ HF_ORG_ADMIN=
 HF_ORG_EARLY_ACCESS=
 
 PUBLIC_SMOOTH_UPDATES=false
-COMMUNITY_TOOLS=false
+COMMUNITY_TOOLS=false
+
+ENCRYPTION_KEY=#your encryption key here

scripts/populate.ts

@@ -14,8 +14,7 @@ import type { User } from "../src/lib/types/User";
 import type { Assistant } from "../src/lib/types/Assistant";
 import type { Conversation } from "../src/lib/types/Conversation";
 import type { Settings } from "../src/lib/types/Settings";
-import type { CommunityToolDB, ToolLogoColor, ToolLogoIcon } from "../src/lib/types/Tool";
-import { defaultEmbeddingModel } from "../src/lib/server/embeddingModels.ts";
+import { getDefaultEmbeddingModel } from "../src/lib/server/embeddingModels.ts";
 import { Message } from "../src/lib/types/Message.ts";
 
 import { addChildren } from "../src/lib/utils/tree/addChildren.ts";
@@ -148,6 +147,7 @@ async function seed() {
 				updatedAt: faker.date.recent({ days: 30 }),
 				customPrompts: {},
 				assistants: [],
+				disableStream: false,
 			};
 			await collections.settings.updateOne(
 				{ userId: user._id },
@@ -216,7 +216,7 @@ async function seed() {
 								: faker.helpers.maybe(() => faker.hacker.phrase(), { probability: 0.5 })) ?? "";
 
 						const messages = await generateMessages(preprompt);
-
+						const defaultEmbeddingModel = await getDefaultEmbeddingModel();
 						const conv = {
 							_id: new ObjectId(),
 							userId: user._id,
@@ -226,7 +226,7 @@ async function seed() {
 							updatedAt: faker.date.recent({ days: 145 }),
 							model: faker.helpers.arrayElement(modelIds),
 							title: faker.internet.emoji() + " " + faker.hacker.phrase(),
-							embeddingModel: defaultEmbeddingModel.id,
+							embeddingModel: defaultEmbeddingModel.name,
 							messages,
 							rootMessageId: messages[0].id,
 						} satisfies Conversation;

src/hooks.server.ts

@@ -16,6 +16,7 @@ import { initExitHandler } from "$lib/server/exitHandler";
 import { ObjectId } from "mongodb";
 import { refreshAssistantsCounts } from "$lib/jobs/refresh-assistants-counts";
 import { refreshConversationStats } from "$lib/jobs/refresh-conversation-stats";
+import { pupulateEmbeddingModel } from "$lib/server/embeddingModels";
 
 // TODO: move this code on a started server hook, instead of using a "building" flag
 if (!building) {
@@ -26,6 +27,9 @@ if (!building) {
 	if (env.ENABLE_ASSISTANTS) {
 		refreshAssistantsCounts();
 	}
+
+	await pupulateEmbeddingModel();
+
 	refreshConversationStats();
 
 	// Init metrics server

src/lib/server/database.ts

@@ -19,6 +19,7 @@ import { logger } from "$lib/server/logger";
 import { building } from "$app/environment";
 import type { TokenCache } from "$lib/types/TokenCache";
 import { onExit } from "./exitHandler";
+import type { EmbeddingModel } from "$lib/types/EmbeddingModel";
 
 export const CONVERSATION_STATS_COLLECTION = "conversations.stats";
 
@@ -88,6 +89,7 @@ export class Database {
 		const semaphores = db.collection<Semaphore>("semaphores");
 		const tokenCaches = db.collection<TokenCache>("tokens");
 		const tools = db.collection<CommunityToolDB>("tools");
+		const embeddingModels = db.collection<EmbeddingModel>("embeddingModels");
 
 		return {
 			conversations,
@@ -106,6 +108,7 @@ export class Database {
 			semaphores,
 			tokenCaches,
 			tools,
+			embeddingModels,
 		};
 	}
 
@@ -129,6 +132,7 @@ export class Database {
 			semaphores,
 			tokenCaches,
 			tools,
+			embeddingModels,
 		} = this.getCollections();
 
 		conversations
@@ -228,6 +232,8 @@ export class Database {
 		tools.createIndex({ createdById: 1, userCount: -1 }).catch((e) => logger.error(e));
 		tools.createIndex({ userCount: 1 }).catch((e) => logger.error(e));
 		tools.createIndex({ last24HoursCount: 1 }).catch((e) => logger.error(e));
+
+		embeddingModels.createIndex({ name: 1 }, { unique: true }).catch((e) => logger.error(e));
 	}
 }
 

src/lib/server/embeddingEndpoints/embeddingEndpoints.ts

@@ -12,6 +12,7 @@ import {
 	embeddingEndpointOpenAIParametersSchema,
 } from "./openai/embeddingEndpoints";
 import { embeddingEndpointHfApi, embeddingEndpointHfApiSchema } from "./hfApi/embeddingHfApi";
+import type { EmbeddingModel } from "$lib/types/EmbeddingModel";
 
 // parameters passed when generating text
 interface EmbeddingEndpointParameters {
@@ -33,8 +34,8 @@ export const embeddingEndpointSchema = z.discriminatedUnion("type", [
 type EmbeddingEndpointTypeOptions = z.infer<typeof embeddingEndpointSchema>["type"];
 
 // generator function that takes in type discrimantor value for defining the endpoint and return the endpoint
-export type EmbeddingEndpointGenerator<T extends EmbeddingEndpointTypeOptions> = (
-	inputs: Extract<z.infer<typeof embeddingEndpointSchema>, { type: T }>
+type EmbeddingEndpointGenerator<T extends EmbeddingEndpointTypeOptions> = (
+	inputs: Extract<z.infer<typeof embeddingEndpointSchema>, { type: T }> & { model: EmbeddingModel }
 ) => EmbeddingEndpoint | Promise<EmbeddingEndpoint>;
 
 // list of all endpoint generators

src/lib/server/embeddingEndpoints/hfApi/embeddingHfApi.ts

@@ -3,22 +3,31 @@ import type { EmbeddingEndpoint, Embedding } from "../embeddingEndpoints";
 import { chunk } from "$lib/utils/chunk";
 import { env } from "$env/dynamic/private";
 import { logger } from "$lib/server/logger";
+import type { EmbeddingModel } from "$lib/types/EmbeddingModel";
+import { decrypt } from "$lib/utils/encryption";
 
 export const embeddingEndpointHfApiSchema = z.object({
 	weight: z.number().int().positive().default(1),
-	model: z.any(),
 	type: z.literal("hfapi"),
 	authorization: z
 		.string()
 		.optional()
 		.transform((v) => (!v && env.HF_TOKEN ? "Bearer " + env.HF_TOKEN : v)), // if the header is not set but HF_TOKEN is, use it as the authorization header
 });
 
+type EmbeddingEndpointHfApiInput = z.input<typeof embeddingEndpointHfApiSchema> & {
+	model: EmbeddingModel;
+};
+
 export async function embeddingEndpointHfApi(
-	input: z.input<typeof embeddingEndpointHfApiSchema>
+	input: EmbeddingEndpointHfApiInput
 ): Promise<EmbeddingEndpoint> {
-	const { model, authorization } = embeddingEndpointHfApiSchema.parse(input);
-	const url = "https://api-inference.huggingface.co/models/" + model.id;
+	const { model } = input;
+	const { authorization } = embeddingEndpointHfApiSchema.parse(input);
+
+	const decryptedAuthorization = authorization && decrypt(authorization);
+
+	const url = "https://api-inference.huggingface.co/models/" + model.name;
 
 	return async ({ inputs }) => {
 		const batchesInputs = chunk(inputs, 128);
@@ -30,7 +39,7 @@ export async function embeddingEndpointHfApi(
 					headers: {
 						Accept: "application/json",
 						"Content-Type": "application/json",
-						...(authorization ? { Authorization: authorization } : {}),
+						...(decryptedAuthorization ? { Authorization: decryptedAuthorization } : {}),
 					},
 					body: JSON.stringify({
 						inputs: {

src/lib/server/embeddingEndpoints/openai/embeddingEndpoints.ts

@@ -2,21 +2,28 @@ import { z } from "zod";
 import type { EmbeddingEndpoint, Embedding } from "../embeddingEndpoints";
 import { chunk } from "$lib/utils/chunk";
 import { env } from "$env/dynamic/private";
+import type { EmbeddingModel } from "$lib/types/EmbeddingModel";
+import { decrypt } from "$lib/utils/encryption";
 
 export const embeddingEndpointOpenAIParametersSchema = z.object({
 	weight: z.number().int().positive().default(1),
-	model: z.any(),
 	type: z.literal("openai"),
 	url: z.string().url().default("https://api.openai.com/v1/embeddings"),
 	apiKey: z.string().default(env.OPENAI_API_KEY),
 	defaultHeaders: z.record(z.string()).default({}),
 });
 
+type EmbeddingEndpointOpenAIInput = z.input<typeof embeddingEndpointOpenAIParametersSchema> & {
+	model: EmbeddingModel;
+};
+
 export async function embeddingEndpointOpenAI(
-	input: z.input<typeof embeddingEndpointOpenAIParametersSchema>
+	input: EmbeddingEndpointOpenAIInput
 ): Promise<EmbeddingEndpoint> {
-	const { url, model, apiKey, defaultHeaders } =
-		embeddingEndpointOpenAIParametersSchema.parse(input);
+	const { model } = input;
+	const { url, apiKey, defaultHeaders } = embeddingEndpointOpenAIParametersSchema.parse(input);
+
+	const decryptedApiKey = decrypt(apiKey);
 
 	const maxBatchSize = model.maxBatchSize || 100;
 
@@ -32,7 +39,7 @@ export async function embeddingEndpointOpenAI(
 					headers: {
 						Accept: "application/json",
 						"Content-Type": "application/json",
-						...(apiKey ? { Authorization: `Bearer ${apiKey}` } : {}),
+						...(decryptedApiKey ? { Authorization: `Bearer ${decryptedApiKey}` } : {}),
 						...defaultHeaders,
 					},
 					body: JSON.stringify({ input: batchInputs, model: model.name }),

src/lib/server/embeddingEndpoints/tei/embeddingEndpoints.ts

@@ -3,10 +3,11 @@ import type { EmbeddingEndpoint, Embedding } from "../embeddingEndpoints";
 import { chunk } from "$lib/utils/chunk";
 import { env } from "$env/dynamic/private";
 import { logger } from "$lib/server/logger";
+import type { EmbeddingModel } from "$lib/types/EmbeddingModel";
+import { decrypt } from "$lib/utils/encryption";
 
 export const embeddingEndpointTeiParametersSchema = z.object({
 	weight: z.number().int().positive().default(1),
-	model: z.any(),
 	type: z.literal("tei"),
 	url: z.string().url(),
 	authorization: z
@@ -35,10 +36,17 @@ const getModelInfoByUrl = async (url: string, authorization?: string) => {
 	}
 };
 
+type EmbeddingEndpointTeiInput = z.input<typeof embeddingEndpointTeiParametersSchema> & {
+	model: EmbeddingModel;
+};
+
 export async function embeddingEndpointTei(
-	input: z.input<typeof embeddingEndpointTeiParametersSchema>
+	input: EmbeddingEndpointTeiInput
 ): Promise<EmbeddingEndpoint> {
-	const { url, model, authorization } = embeddingEndpointTeiParametersSchema.parse(input);
+	const { model } = input;
+	const { url, authorization } = embeddingEndpointTeiParametersSchema.parse(input);
+
+	const decryptedAuthorization = authorization && decrypt(authorization);
 
 	const { max_client_batch_size, max_batch_tokens } = await getModelInfoByUrl(url);
 	const maxBatchSize = Math.min(
@@ -58,7 +66,7 @@ export async function embeddingEndpointTei(
 					headers: {
 						Accept: "application/json",
 						"Content-Type": "application/json",
-						...(authorization ? { Authorization: authorization } : {}),
+						...(decryptedAuthorization ? { Authorization: decryptedAuthorization } : {}),
 					},
 					body: JSON.stringify({ inputs: batchInputs, normalize: true, truncate: true }),
 				});

src/lib/server/embeddingEndpoints/transformersjs/embeddingEndpoints.ts

@@ -2,10 +2,10 @@ import { z } from "zod";
 import type { EmbeddingEndpoint } from "../embeddingEndpoints";
 import type { Tensor, FeatureExtractionPipeline } from "@huggingface/transformers";
 import { pipeline } from "@huggingface/transformers";
+import type { EmbeddingModel } from "$lib/types/EmbeddingModel";
 
 export const embeddingEndpointTransformersJSParametersSchema = z.object({
 	weight: z.number().int().positive().default(1),
-	model: z.any(),
 	type: z.literal("transformersjs"),
 });
 
@@ -36,10 +36,16 @@ export async function calculateEmbedding(modelName: string, inputs: string[]) {
 	return output.tolist();
 }
 
+type EmbeddingEndpointTransformersJSInput = z.input<
+	typeof embeddingEndpointTransformersJSParametersSchema
+> & {
+	model: EmbeddingModel;
+};
+
 export function embeddingEndpointTransformersJS(
-	input: z.input<typeof embeddingEndpointTransformersJSParametersSchema>
+	input: EmbeddingEndpointTransformersJSInput
 ): EmbeddingEndpoint {
-	const { model } = embeddingEndpointTransformersJSParametersSchema.parse(input);
+	const { model } = input;
 
 	return async ({ inputs }) => {
 		return calculateEmbedding(model.name, inputs);