feat: support output dimensions for embedding models (#490) (#522)

Add optional outputDimension field for embedding models to request
specific dimensions from providers that support Matryoshka
Representation Learning (MRL).

Supported providers:
- OpenAI text-embedding-3-* (via 'dimensions' parameter)
- Google gemini-embedding-001 (via 'config.outputDimensionality')

Changes:
- Add outputDimension field to EmbeddingModel type
- Update all provider getEmbedding methods to accept dimensions option
- Add Output Dimensions input field in Add Embedding Model modal
- Validate that returned dimension matches requested outputDimension

Author: kevin-on

src/components/settings/modals/AddEmbeddingModelModal.tsx

@@ -43,7 +43,9 @@ function AddEmbeddingModelModalComponent({
     providerType: DEFAULT_PROVIDERS[0].type,
     id: '',
     model: '',
+    outputDimension: undefined,
   })
+  const [outputDimensionInput, setOutputDimensionInput] = useState('')
 
   const handleSubmit = async () => {
     try {
@@ -69,6 +71,7 @@ function AddEmbeddingModelModalComponent({
       const embeddingResult = await providerClient.getEmbedding(
         formData.model,
         'test',
+        { dimensions: formData.outputDimension },
       )
 
       if (!Array.isArray(embeddingResult) || embeddingResult.length === 0) {
@@ -77,6 +80,18 @@ function AddEmbeddingModelModalComponent({
 
       const dimension = embeddingResult.length
 
+      // Validate that the model respected the requested output dimension
+      if (
+        formData.outputDimension !== undefined &&
+        dimension !== formData.outputDimension
+      ) {
+        throw new Error(
+          `Requested output dimension ${formData.outputDimension}, but the model returned ${dimension} dimensions. ` +
+            `This model may not support custom output dimensions (Matryoshka Representation Learning). ` +
+            `Leave the "Output Dimensions" field empty to use the model's default dimension.`,
+        )
+      }
+
       if (!supportedDimensionsForIndex.includes(dimension)) {
         const confirmed = await new Promise<boolean>((resolve) => {
           new ConfirmModal(plugin.app, {
@@ -174,6 +189,24 @@ function AddEmbeddingModelModalComponent({
         />
       </ObsidianSetting>
 
+      <ObsidianSetting
+        name="Output Dimensions"
+        desc="Optional. Request a specific output dimension from models that support Matryoshka Representation Learning (MRL), such as OpenAI's text-embedding-3-* or Google's gemini-embedding-001. Leave empty to use the model's default dimension."
+      >
+        <ObsidianTextInput
+          value={outputDimensionInput}
+          placeholder="e.g., 768"
+          onChange={(value: string) => {
+            setOutputDimensionInput(value)
+            const parsed = parseInt(value, 10)
+            setFormData((prev) => ({
+              ...prev,
+              outputDimension: isNaN(parsed) ? undefined : parsed,
+            }))
+          }}
+        />
+      </ObsidianSetting>
+
       <ObsidianSetting>
         <ObsidianButton text="Add" onClick={handleSubmit} cta />
         <ObsidianButton text="Cancel" onClick={onClose} />

src/core/llm/anthropic.ts

@@ -603,7 +603,11 @@ https://github.com/glowingjade/obsidian-smart-composer/issues/286`,
     throw new Error(`Unsupported tool choice: ${JSON.stringify(toolChoice)}`)
   }
 
-  async getEmbedding(_model: string, _text: string): Promise<number[]> {
+  async getEmbedding(
+    _model: string,
+    _text: string,
+    _options?: { dimensions?: number },
+  ): Promise<number[]> {
     throw new Error(
       `Provider ${this.provider.id} does not support embeddings. Please use a different provider.`,
     )

src/core/llm/anthropicClaudeCodeProvider.ts

@@ -77,7 +77,11 @@ export class AnthropicClaudeCodeProvider extends BaseLLMProvider<
     )
   }
 
-  async getEmbedding(_model: string, _text: string): Promise<number[]> {
+  async getEmbedding(
+    _model: string,
+    _text: string,
+    _options?: { dimensions?: number },
+  ): Promise<number[]> {
     throw new Error(
       `Provider ${this.provider.id} does not support embeddings. Please use a different provider.`,
     )

src/core/llm/azureOpenaiProvider.ts

@@ -57,7 +57,11 @@ export class AzureOpenAIProvider extends BaseLLMProvider<
     return this.adapter.streamResponse(this.client, request, options)
   }
 
-  async getEmbedding(_model: string, _text: string): Promise<number[]> {
+  async getEmbedding(
+    _model: string,
+    _text: string,
+    _options?: { dimensions?: number },
+  ): Promise<number[]> {
     throw new Error(
       `Provider ${this.provider.id} does not support embeddings. Please use a different provider.`,
     )

src/core/llm/base.ts

@@ -29,5 +29,9 @@ export abstract class BaseLLMProvider<P extends LLMProvider> {
     options?: LLMOptions,
   ): Promise<AsyncIterable<LLMResponseStreaming>>
 
-  abstract getEmbedding(model: string, text: string): Promise<number[]>
+  abstract getEmbedding(
+    model: string,
+    text: string,
+    options?: { dimensions?: number },
+  ): Promise<number[]>
 }

src/core/llm/deepseekStudioProvider.ts

@@ -80,7 +80,11 @@ export class DeepSeekStudioProvider extends BaseLLMProvider<
     return this.adapter.streamResponse(this.client, formattedRequest, options)
   }
 
-  async getEmbedding(_model: string, _text: string): Promise<number[]> {
+  async getEmbedding(
+    _model: string,
+    _text: string,
+    _options?: { dimensions?: number },
+  ): Promise<number[]> {
     throw new Error(
       `Provider ${this.provider.id} does not support embeddings. Please use a different provider.`,
     )

src/core/llm/gemini.ts

@@ -543,7 +543,11 @@ export class GeminiProvider extends BaseLLMProvider<
     return config
   }
 
-  async getEmbedding(model: string, text: string): Promise<number[]> {
+  async getEmbedding(
+    model: string,
+    text: string,
+    options?: { dimensions?: number },
+  ): Promise<number[]> {
     if (!this.apiKey) {
       throw new LLMAPIKeyNotSetException(
         `Provider ${this.provider.id} API key is missing. Please set it in settings menu.`,
@@ -554,6 +558,9 @@ export class GeminiProvider extends BaseLLMProvider<
       const response = await this.client.models.embedContent({
         model: model,
         contents: text,
+        ...(options?.dimensions && {
+          config: { outputDimensionality: options.dimensions },
+        }),
       })
       return response.embeddings?.[0]?.values ?? []
     } catch (error) {

src/core/llm/geminiPlanProvider.ts

@@ -79,7 +79,11 @@ export class GeminiPlanProvider extends BaseLLMProvider<
     )
   }
 
-  async getEmbedding(_model: string, _text: string): Promise<number[]> {
+  async getEmbedding(
+    _model: string,
+    _text: string,
+    _options?: { dimensions?: number },
+  ): Promise<number[]> {
     throw new Error(
       `Provider ${this.provider.id} does not support embeddings. Please use a different provider.`,
     )

src/core/llm/lmStudioProvider.ts

@@ -55,11 +55,16 @@ export class LmStudioProvider extends BaseLLMProvider<
     return this.adapter.streamResponse(this.client, request, options)
   }
 
-  async getEmbedding(model: string, text: string): Promise<number[]> {
+  async getEmbedding(
+    model: string,
+    text: string,
+    options?: { dimensions?: number },
+  ): Promise<number[]> {
     const embedding = await this.client.embeddings.create({
       model: model,
       input: text,
       encoding_format: 'float',
+      ...(options?.dimensions && { dimensions: options.dimensions }),
     })
     return embedding.data[0].embedding
   }

src/core/llm/mistralProvider.ts

@@ -56,7 +56,11 @@ export class MistralProvider extends BaseLLMProvider<
     return this.adapter.streamResponse(this.client, request, options)
   }
 
-  async getEmbedding(_model: string, _text: string): Promise<number[]> {
+  async getEmbedding(
+    _model: string,
+    _text: string,
+    _options?: { dimensions?: number },
+  ): Promise<number[]> {
     throw new Error(
       `Provider ${this.provider.id} does not support embeddings. Please use a different provider.`,
     )