Enable batch embedding and add copilot-plus-small (#969)

logancyang · Dec 25, 2024 · 62057c3 · 62057c3
1 parent 6e205cc
commit 62057c3
Show file tree

Hide file tree

Showing 3 changed files with 147 additions and 133 deletions.
diff --git a/src/LLMProviders/embeddingManager.ts b/src/LLMProviders/embeddingManager.ts
@@ -1,6 +1,6 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import { CustomModel } from "@/aiParams";
-import { EmbeddingModelProviders } from "@/constants";
+import { BREVILABS_API_BASE_URL, EmbeddingModelProviders } from "@/constants";
 import { getDecryptedKey } from "@/encryptionService";
 import { CustomError } from "@/error";
 import { getSettings, subscribeToSettingsChange } from "@/settings/model";
@@ -15,6 +15,7 @@ import { Notice } from "obsidian";
 type EmbeddingConstructorType = new (config: any) => Embeddings;
 
 const EMBEDDING_PROVIDER_CONSTRUCTORS = {
+  [EmbeddingModelProviders.COPILOT_PLUS]: OpenAIEmbeddings,
   [EmbeddingModelProviders.OPENAI]: OpenAIEmbeddings,
   [EmbeddingModelProviders.COHEREAI]: CohereEmbeddings,
   [EmbeddingModelProviders.GOOGLE]: GoogleGenerativeAIEmbeddings,
@@ -40,6 +41,7 @@ export default class EmbeddingManager {
   >;
 
   private readonly providerApiKeyMap: Record<EmbeddingModelProviders, () => string> = {
+    [EmbeddingModelProviders.COPILOT_PLUS]: () => getSettings().plusLicenseKey,
     [EmbeddingModelProviders.OPENAI]: () => getSettings().openAIApiKey,
     [EmbeddingModelProviders.COHEREAI]: () => getSettings().cohereApiKey,
     [EmbeddingModelProviders.GOOGLE]: () => getSettings().googleApiKey,
@@ -167,6 +169,15 @@ export default class EmbeddingManager {
         EmbeddingProviderConstructorMap[K]
       >[0] /*& Record<string, unknown>;*/;
     } = {
+      [EmbeddingModelProviders.COPILOT_PLUS]: {
+        modelName,
+        apiKey: getDecryptedKey(settings.plusLicenseKey),
+        timeout: 10000,
+        configuration: {
+          baseURL: BREVILABS_API_BASE_URL,
+          fetch: customModel.enableCors ? safeFetch : undefined,
+        },
+      },
       [EmbeddingModelProviders.OPENAI]: {
         modelName,
         openAIApiKey: getDecryptedKey(customModel.apiKey || settings.openAIApiKey),

diff --git a/src/constants.ts b/src/constants.ts
@@ -127,6 +127,7 @@ export enum EmbeddingModelProviders {
   OLLAMA = "ollama",
   LM_STUDIO = "lm-studio",
   OPENAI_FORMAT = "3rd party (openai-format)",
+  COPILOT_PLUS = "copilot-plus",
   // HUGGINGFACE = "huggingface",
   // VOYAGEAI = "voyageai",
 }
@@ -138,9 +139,18 @@ export enum EmbeddingModels {
   AZURE_OPENAI = "azure-openai",
   COHEREAI_EMBED_MULTILINGUAL_LIGHT_V3_0 = "embed-multilingual-light-v3.0",
   GOOGLE_ENG = "text-embedding-004",
+  COPILOT_PLUS_SMALL = "copilot-plus-small",
 }
 
 export const BUILTIN_EMBEDDING_MODELS: CustomModel[] = [
+  {
+    name: EmbeddingModels.COPILOT_PLUS_SMALL,
+    provider: EmbeddingModelProviders.COPILOT_PLUS,
+    enabled: true,
+    isBuiltIn: true,
+    isEmbeddingModel: true,
+    core: true,
+  },
   {
     name: EmbeddingModels.OPENAI_EMBEDDING_SMALL,
     provider: EmbeddingModelProviders.OPENAI,
@@ -238,7 +248,8 @@ export const DEFAULT_SETTINGS: CopilotSettings = {
   openRouterAiApiKey: "",
   defaultChainType: ChainType.LLM_CHAIN,
   defaultModelKey: ChatModels.GPT_4o + "|" + ChatModelProviders.OPENAI,
-  embeddingModelKey: EmbeddingModels.OPENAI_EMBEDDING_SMALL + "|" + ChatModelProviders.OPENAI,
+  embeddingModelKey:
+    EmbeddingModels.COPILOT_PLUS_SMALL + "|" + EmbeddingModelProviders.COPILOT_PLUS,
   temperature: 0.1,
   maxTokens: 1000,
   contextTurns: 15,