Add dimensions settings to embedding models

Author: intitni

Core/Sources/HostApp/AccountSettings/EmbeddingModelManagement/EmbeddingModelEdit.swift

@@ -15,6 +15,7 @@ struct EmbeddingModelEdit {
         var name: String
         var format: EmbeddingModel.Format
         var maxTokens: Int = 8191
+        var dimensions: Int = 1536
         var modelName: String = ""
         var ollamaKeepAlive: String = ""
         var apiKeyName: String { apiKeySelection.apiKeyName }
@@ -38,6 +39,7 @@ struct EmbeddingModelEdit {
         case testButtonClicked
         case testSucceeded(String)
         case testFailed(String)
+        case fixDimensions(Int)
         case checkSuggestedMaxTokens
         case apiKeySelection(APIKeySelection.Action)
         case baseURLSelection(BaseURLSelection.Action)
@@ -80,6 +82,7 @@ struct EmbeddingModelEdit {
             case .testButtonClicked:
                 guard !state.isTesting else { return .none }
                 state.isTesting = true
+                let dimensions = state.dimensions
                 let model = EmbeddingModel(
                     id: state.id,
                     name: state.name,
@@ -89,18 +92,33 @@ struct EmbeddingModelEdit {
                         baseURL: state.baseURL,
                         isFullURL: state.isFullURL,
                         maxTokens: state.maxTokens,
+                        dimensions: dimensions,
                         modelName: state.modelName
                     )
                 )
                 return .run { send in
                     do {
-                        _ = try await EmbeddingService(
+                        let result = try await EmbeddingService(
                             configuration: UserPreferenceEmbeddingConfiguration()
                                 .overriding {
                                     $0.model = model
                                 }
                         ).embed(text: "Hello")
-                        await send(.testSucceeded("Succeeded!"))
+                        if result.data.isEmpty {
+                            await send(.testFailed("No data returned"))
+                            return
+                        }
+                        let actualDimensions = result.data.first?.embedding.count ?? 0
+                        if actualDimensions != dimensions {
+                            await send(
+                                .testFailed("Invalid dimension, should be \(actualDimensions)")
+                            )
+                            await send(.fixDimensions(actualDimensions))
+                        } else {
+                            await send(
+                                .testSucceeded("Succeeded! (Dimensions: \(actualDimensions))")
+                            )
+                        }
                     } catch {
                         await send(.testFailed(error.localizedDescription))
                     }
@@ -131,6 +149,11 @@ struct EmbeddingModelEdit {
                     return .none
                 }
                 state.suggestedMaxTokens = knownModel.maxToken
+                state.dimensions = knownModel.dimensions
+                return .none
+
+            case let .fixDimensions(value):
+                state.dimensions = value
                 return .none
 
             case .apiKeySelection:

Core/Sources/HostApp/AccountSettings/EmbeddingModelManagement/EmbeddingModelEditView.swift

@@ -174,6 +174,51 @@ struct EmbeddingModelEditView: View {
             }
         }
     }
+    
+    struct DimensionsTextField: View {
+        @Perception.Bindable var store: StoreOf<EmbeddingModelEdit>
+
+        var body: some View {
+            WithPerceptionTracking {
+                HStack {
+                    let textFieldBinding = Binding(
+                        get: { String(store.dimensions) },
+                        set: {
+                            if let selectionDimensions = Int($0) {
+                                $store.dimensions.wrappedValue = selectionDimensions
+                            } else {
+                                $store.dimensions.wrappedValue = 0
+                            }
+                        }
+                    )
+
+                    TextField(text: textFieldBinding) {
+                        Text("Dimensions")
+                            .multilineTextAlignment(.trailing)
+                    }
+                    .overlay(alignment: .trailing) {
+                        Stepper(
+                            value: $store.dimensions,
+                            in: 0...Int.max,
+                            step: 100
+                        ) {
+                            EmptyView()
+                        }
+                    }
+                    .foregroundColor({
+                        if store.dimensions <= 0 {
+                            return .red
+                        }
+                        return .primary
+                    }() as Color)
+                }
+                
+                Text("If you are not sure, run test to get the correct value.")
+                    .font(.caption)
+                    .dynamicHeightTextInFormWorkaround()
+            }
+        }
+    }
 
     struct ApiKeyNamePicker: View {
         let store: StoreOf<EmbeddingModelEdit>
@@ -215,6 +260,7 @@ struct EmbeddingModelEditView: View {
                     }
 
                 MaxTokensTextField(store: store)
+                DimensionsTextField(store: store)
 
                 VStack(alignment: .leading, spacing: 8) {
                     Text(Image(systemName: "exclamationmark.triangle.fill")) + Text(
@@ -242,6 +288,7 @@ struct EmbeddingModelEditView: View {
                 TextField("Deployment Name", text: $store.modelName)
 
                 MaxTokensTextField(store: store)
+                DimensionsTextField(store: store)
             }
         }
     }
@@ -279,6 +326,7 @@ struct EmbeddingModelEditView: View {
                 TextField("Model Name", text: $store.modelName)
 
                 MaxTokensTextField(store: store)
+                DimensionsTextField(store: store)
 
                 Button("Custom Headers") {
                     isEditingCustomHeader.toggle()
@@ -299,6 +347,7 @@ struct EmbeddingModelEditView: View {
                 TextField("Model Name", text: $store.modelName)
 
                 MaxTokensTextField(store: store)
+                DimensionsTextField(store: store)
 
                 WithPerceptionTracking {
                     TextField(text: $store.ollamaKeepAlive, prompt: Text("Default Value")) {

Tool/Sources/Preferences/Types/OpenAIEmbeddingModel.swift

@@ -15,5 +15,16 @@ public extension OpenAIEmbeddingModel {
             return 8191
         }
     }
+    
+    var dimensions: Int {
+        switch self {
+        case .textEmbeddingAda002:
+            return 1536
+        case .textEmbedding3Small:
+            return 1536
+        case .textEmbedding3Large:
+            return 3072
+        }
+    }
 }