From cd0329e0cf6ef34cc8a3b1b304cc727e28dd87d5 Mon Sep 17 00:00:00 2001
From: SuYao <sy20010504@gmail.com>
Date: Sun, 25 May 2025 13:07:21 +0800
Subject: [PATCH] fix: update dimensions handling in KnowledgeBaseParams
 (#6417)

fix: update dimensions handling in KnowledgeBaseParams and add supported dimension providers
---
 src/main/embeddings/EmbeddingsFactory.ts      | 4 ++--
 src/renderer/src/config/providers.ts          | 1 +
 src/renderer/src/services/KnowledgeService.ts | 3 ++-
 src/renderer/src/types/index.ts               | 2 +-
 4 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/main/embeddings/EmbeddingsFactory.ts b/src/main/embeddings/EmbeddingsFactory.ts
index 69de15171e..5924d00d7d 100644
--- a/src/main/embeddings/EmbeddingsFactory.ts
+++ b/src/main/embeddings/EmbeddingsFactory.ts
@@ -23,14 +23,14 @@ export default class EmbeddingsFactory {
         azureOpenAIApiVersion: apiVersion,
         azureOpenAIApiDeploymentName: model,
         azureOpenAIApiInstanceName: getInstanceName(baseURL),
-        // dimensions,
+        dimensions,
         batchSize
       })
     }
     return new OpenAiEmbeddings({
       model,
       apiKey,
-      // dimensions,
+      dimensions,
       batchSize,
       configuration: { baseURL }
     })
diff --git a/src/renderer/src/config/providers.ts b/src/renderer/src/config/providers.ts
index 506eea3ad1..9096c1d5ce 100644
--- a/src/renderer/src/config/providers.ts
+++ b/src/renderer/src/config/providers.ts
@@ -103,6 +103,7 @@ export function getProviderLogo(providerId: string) {
 
 // export const SUPPORTED_REANK_PROVIDERS = ['silicon', 'jina', 'voyageai', 'dashscope', 'aihubmix']
 export const NOT_SUPPORTED_REANK_PROVIDERS = ['ollama']
+export const ONLY_SUPPORTED_DIMENSION_PROVIDERS = ['ollama', 'infini']
 
 export const PROVIDER_CONFIG = {
   openai: {
diff --git a/src/renderer/src/services/KnowledgeService.ts b/src/renderer/src/services/KnowledgeService.ts
index 4200085b36..626788a4e2 100644
--- a/src/renderer/src/services/KnowledgeService.ts
+++ b/src/renderer/src/services/KnowledgeService.ts
@@ -2,6 +2,7 @@ import type { ExtractChunkData } from '@cherrystudio/embedjs-interfaces'
 import { DEFAULT_KNOWLEDGE_DOCUMENT_COUNT, DEFAULT_KNOWLEDGE_THRESHOLD } from '@renderer/config/constant'
 import { getEmbeddingMaxContext } from '@renderer/config/embedings'
 import Logger from '@renderer/config/logger'
+import { ONLY_SUPPORTED_DIMENSION_PROVIDERS } from '@renderer/config/providers'
 import AiProvider from '@renderer/providers/AiProvider'
 import store from '@renderer/store'
 import { FileType, KnowledgeBase, KnowledgeBaseParams, KnowledgeReference } from '@renderer/types'
@@ -38,7 +39,7 @@ export const getKnowledgeBaseParams = (base: KnowledgeBase): KnowledgeBaseParams
   return {
     id: base.id,
     model: base.model.id,
-    dimensions: base.dimensions,
+    dimensions: ONLY_SUPPORTED_DIMENSION_PROVIDERS.includes(base.model.provider) ? base.dimensions : undefined,
     apiKey: aiProvider.getApiKey() || 'secret',
     apiVersion: provider.apiVersion,
     baseURL: host,
diff --git a/src/renderer/src/types/index.ts b/src/renderer/src/types/index.ts
index db8323e322..4b29fb5caa 100644
--- a/src/renderer/src/types/index.ts
+++ b/src/renderer/src/types/index.ts
@@ -393,7 +393,7 @@ export interface KnowledgeBase {
 export type KnowledgeBaseParams = {
   id: string
   model: string
-  dimensions: number
+  dimensions?: number
   apiKey: string
   apiVersion?: string
   baseURL: string