feat(mcp): Support local embeddings via environment variable

1 files changed, 25 insertions, 6 deletions

diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts
index 5fc6a7e..9102c62 100644
--- a/packages/mcp/src/index.ts
+++ b/packages/mcp/src/index.ts
@@ -6,8 +6,11 @@ import { z } from "zod";
 import {
 	createSupabaseClient,
 	EmbeddingService,
+	LocalEmbeddingProvider,
 	SupabaseStore,
 	SupabaseProjectStore,
+	type EmbeddingProvider,
+	type LocalEmbeddingModel,
 } from "@imemio/sdk";
 
 function getRequiredEnvironmentVariable(name: string): string {
@@ -24,16 +27,32 @@ function getOptionalEnvironmentVariable(name: string): string | undefined {
 	return process.env[name];
 }
 
+function createEmbeddingProvider(): EmbeddingProvider | null {
+	const embeddingType = getOptionalEnvironmentVariable("IMEMIO_EMBEDDING_TYPE");
+	const openaiApiKey = getOptionalEnvironmentVariable("OPENAI_API_KEY");
+
+	if (embeddingType === "local") {
+		const model = getOptionalEnvironmentVariable(
+			"IMEMIO_LOCAL_EMBEDDING_MODEL",
+		) as LocalEmbeddingModel | undefined;
+
+		return new LocalEmbeddingProvider({ model });
+	}
+
+	if (openaiApiKey) {
+		return new EmbeddingService({ apiKey: openaiApiKey });
+	}
+
+	return null;
+}
+
 const supabaseUrl = getRequiredEnvironmentVariable("SUPABASE_URL");
 const supabaseKey = getRequiredEnvironmentVariable("SUPABASE_ANON_KEY");
 const userId = getRequiredEnvironmentVariable("IMEMIO_USER_ID");
-const openaiApiKey = getOptionalEnvironmentVariable("OPENAI_API_KEY");
 const client = createSupabaseClient(supabaseUrl, supabaseKey);
 const memoryStore = new SupabaseStore(client, userId);
 const projectStore = new SupabaseProjectStore(client, userId);
-const embeddingService = openaiApiKey
-	? new EmbeddingService({ apiKey: openaiApiKey })
-	: null;
+const embeddingService = createEmbeddingProvider();
 const tagSchema = z.object({
 	id: z.string(),
 	name: z.string(),
@@ -560,7 +579,7 @@ server.tool(
 				content: [
 					{
 						type: "text" as const,
-						text: "Embedding service not configured. Set OPENAI_API_KEY environment variable.",
+						text: "Embedding service not configured. Set OPENAI_API_KEY or IMEMIO_EMBEDDING_TYPE=local.",
 					},
 				],
 				isError: true,
@@ -595,7 +614,7 @@ server.tool(
 				content: [
 					{
 						type: "text" as const,
-						text: "Embedding service not configured. Set OPENAI_API_KEY environment variable.",
+						text: "Embedding service not configured. Set OPENAI_API_KEY or IMEMIO_EMBEDDING_TYPE=local.",
 					},
 				],
 				isError: true,