From 1fc50ef742f229d4ef0a8dca045a7ff8039d297d Mon Sep 17 00:00:00 2001 From: Jordan Koch Date: Tue, 18 Aug 2026 17:12:17 -0700 Subject: [PATCH] feat: multi-model LLM load balancer (local + frontier + optional Nova) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fan out the shared multi-model load balancer to MLX Code. Optionally spread each chat across all installed local models (native MLX + Ollama), all frontier models via OpenRouter, and an optional Nova Gateway — health-gated and load-balanced — instead of a single pinned MLX model. Off by default; existing single-model behavior is preserved when the toggles are off. - Port pure, network-free building blocks from AIStudio: ModelRegistry (discovery + pool composition), LoadBalancer (round-robin / least-busy with health gating), OpenRouterProvider / OpenAICompatibleRequest, KeychainStore. - Add LLMBalancer service (discoverEnabledPool / healthMap / generateBalanced / dispatchBalanced) wiring discovery → health → balanced dispatch. Local MLX runs in-process via MLXService; Ollama / OpenRouter / Nova ride the generic OpenAI-compatible HTTP path. MLX discovery uses MLXService.discoverModels(). - Wire balanced dispatch into ChatViewModel with clean fallback to the single pinned MLX model when the pool is empty. - Three persisted toggles in AppSettings + a new Settings → Balancer tab (gateway URL, Keychain-stored OpenRouter key). - Add network-free LoadBalancerTests (24 tests). - Update README (Multi-model load balancing section + Mermaid flow) and CHANGELOG. Nova is never a hard requirement: works with zero Nova present; the Nova Gateway is one optional, health-probed entry that drops out if unavailable. Co-Authored-By: Claude Opus 4.8 --- CHANGELOG.md | 15 + MLX Code Tests/LoadBalancerTests.swift | 276 +++++++++++++++ MLX Code.xcodeproj/project.pbxproj | 20 ++ MLX Code/Models/AppSettings.swift | 65 ++++ MLX Code/Services/KeychainStore.swift | 80 +++++ MLX Code/Services/LLMBalancer.swift | 380 +++++++++++++++++++++ MLX Code/Services/ModelRegistry.swift | 241 +++++++++++++ MLX Code/Services/OpenRouterProvider.swift | 158 +++++++++ MLX Code/ViewModels/ChatViewModel.swift | 133 ++++---- MLX Code/Views/SettingsView.swift | 100 ++++++ README.md | 51 +++ 11 files changed, 1461 insertions(+), 58 deletions(-) create mode 100644 MLX Code Tests/LoadBalancerTests.swift create mode 100644 MLX Code/Services/KeychainStore.swift create mode 100644 MLX Code/Services/LLMBalancer.swift create mode 100644 MLX Code/Services/ModelRegistry.swift create mode 100644 MLX Code/Services/OpenRouterProvider.swift diff --git a/CHANGELOG.md b/CHANGELOG.md index b6f8426..c782333 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,21 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added +- Multi-model load balancing: optionally spread each chat across all installed local models (native MLX + Ollama), all frontier models (OpenRouter), and an optional Nova Gateway — health-gated and load-balanced — instead of a single pinned model +- Three independent, persisted toggles in **Settings → Balancer**: All local models, All frontier models (OpenRouter), Nova Gateway (with a configurable gateway URL and Keychain-stored OpenRouter key) +- `ModelRegistry` (model discovery + pool composition), `LoadBalancer` (round-robin / least-busy with health gating), `OpenRouterProvider` / `OpenAICompatibleRequest`, and `KeychainStore` — pure, network-free, unit-tested building blocks ported from AIStudio +- `LLMBalancer` service wiring discovery → health map → balanced dispatch into `ChatViewModel` +- Network-free `LoadBalancerTests` suite (24 tests) covering parsing, pool composition, and selection policies + +### Changed +- `ChatViewModel` routes generation through the balancer when any toggle is on; falls back cleanly to the single pinned MLX model when the pool is empty or all toggles are off (existing behavior preserved) + +### Notes +- Nova is never a hard requirement — the balancer works with zero Nova present; the Nova Gateway is one optional, health-probed entry that drops out if unavailable + ## [5.0.0] - 2026-02-19 ### Removed diff --git a/MLX Code Tests/LoadBalancerTests.swift b/MLX Code Tests/LoadBalancerTests.swift new file mode 100644 index 0000000..07b1439 --- /dev/null +++ b/MLX Code Tests/LoadBalancerTests.swift @@ -0,0 +1,276 @@ +// +// LoadBalancerTests.swift +// MLX Code Tests +// +// Deterministic (no-network) tests for model discovery parsing, pool +// composition, and the load-balancer selection policies. Ported from AIStudio +// and adapted to MLX Code's backend set and MLX discovery source. +// +// Copyright © 2026 Jordan Koch. All rights reserved. +// + +import XCTest +@testable import MLX_Code + +final class LoadBalancerTests: XCTestCase { + + // MARK: - parseOllamaTags + + func testParseOllamaTagsMapsModels() { + let json = """ + {"models": [ + {"name": "mistral:latest", "size": 123}, + {"name": "llama3.2:3b"}, + {"size": 5} + ]} + """ + let models = ModelRegistry.parseOllamaTags(Data(json.utf8)) + XCTAssertEqual(models.count, 2) + XCTAssertEqual(models.map { $0.modelName }, ["mistral:latest", "llama3.2:3b"]) + XCTAssertTrue(models.allSatisfy { $0.backend == .ollama }) + XCTAssertEqual(models[0].id, "ollama|mistral:latest") + XCTAssertEqual(models[0].endpoint, "http://localhost:11434/api/chat") + } + + func testParseOllamaTagsEmptyAndGarbage() { + XCTAssertTrue(ModelRegistry.parseOllamaTags(Data("nonsense".utf8)).isEmpty) + XCTAssertTrue(ModelRegistry.parseOllamaTags(Data("{}".utf8)).isEmpty) + XCTAssertTrue(ModelRegistry.parseOllamaTags(Data(#"{"models": []}"#.utf8)).isEmpty) + XCTAssertTrue(ModelRegistry.parseOllamaTags(Data()).isEmpty) + } + + // MARK: - parseMLXModels (HF hub cache form) + + func testParseMLXModelsFromHubDirs() { + let dirs = [ + "models--mlx-community--Llama-3.2-3B-Instruct-4bit", + "models--meta-llama--Llama-3.1-8B", // not MLX → excluded + "models--mlx-community--Qwen2.5-7B-4bit", + "blobs", // not a model dir → excluded + "models--" // empty repo → excluded + ] + let models = ModelRegistry.parseMLXModels(hubDirectoryNames: dirs) + XCTAssertEqual(models.count, 2) + XCTAssertEqual(models.map { $0.modelName }, + ["mlx-community/Llama-3.2-3B-Instruct-4bit", "mlx-community/Qwen2.5-7B-4bit"]) + XCTAssertTrue(models.allSatisfy { $0.backend == .mlx }) + XCTAssertEqual(models[0].displayName, "Llama-3.2-3B-Instruct-4bit") + } + + func testParseMLXModelsEmpty() { + XCTAssertTrue(ModelRegistry.parseMLXModels(hubDirectoryNames: []).isEmpty) + XCTAssertTrue(ModelRegistry.parseMLXModels(hubDirectoryNames: ["random", "stuff"]).isEmpty) + } + + // MARK: - mlxModels(fromLocalNames:) — MLX Code's own discovery source + + func testMLXModelsFromLocalNames() { + let models = ModelRegistry.mlxModels(fromLocalNames: [ + "Qwen2.5-7B-Instruct-4bit", + "mlx-community/Llama-3.2-3B-Instruct-4bit", + " ", // blank → excluded + "" // empty → excluded + ]) + XCTAssertEqual(models.count, 2) + XCTAssertTrue(models.allSatisfy { $0.backend == .mlx }) + XCTAssertEqual(models[0].modelName, "Qwen2.5-7B-Instruct-4bit") + XCTAssertEqual(models[0].id, "mlx|Qwen2.5-7B-Instruct-4bit") + // display name is the short (last path component) form + XCTAssertEqual(models[1].displayName, "Llama-3.2-3B-Instruct-4bit") + XCTAssertTrue(models.allSatisfy { $0.endpoint.isEmpty }) + } + + // MARK: - frontier + nova mapping + + func testFrontierModelsMapping() { + let frontier = ModelRegistry.frontierModels(from: ["openai/gpt-4o", "anthropic/claude-sonnet-4.5", ""]) + XCTAssertEqual(frontier.count, 2) + XCTAssertTrue(frontier.allSatisfy { $0.backend == .openRouter }) + XCTAssertEqual(frontier[0].endpoint, OpenRouterProvider.chatCompletionsURL) + } + + func testNovaGatewayModel() { + let nova = ModelRegistry.novaGatewayModel() + XCTAssertEqual(nova.backend, .novaGateway) + XCTAssertEqual(nova.endpoint, "http://127.0.0.1:18792/v1/chat/completions") + } + + // MARK: - Pool composition (toggles) + + private func samplePool() -> (local: [DiscoveredModel], frontier: [DiscoveredModel], nova: DiscoveredModel) { + let ollama = [DiscoveredModel(modelName: "mistral:latest", backend: .ollama, endpoint: "e")] + let mlx = [DiscoveredModel(modelName: "mlx-community/Qwen", backend: .mlx, endpoint: "")] + let frontier = ModelRegistry.frontierModels(from: ["openai/gpt-4o"]) + let nova = ModelRegistry.novaGatewayModel() + return (ollama + mlx, frontier, nova) + } + + func testAssemblePoolLocalOnly() { + let s = samplePool() + let pool = ModelRegistry.assemblePool( + ollama: [s.local[0]], mlx: [s.local[1]], frontier: s.frontier, novaGateway: s.nova, + useAllLocalModels: true, enableAllFrontierModels: false, useNovaGateway: false) + XCTAssertEqual(pool.count, 2) + XCTAssertTrue(pool.allSatisfy { $0.backend == .ollama || $0.backend == .mlx }) + } + + func testAssemblePoolFrontierOnly() { + let s = samplePool() + let pool = ModelRegistry.assemblePool( + ollama: [s.local[0]], mlx: [s.local[1]], frontier: s.frontier, novaGateway: s.nova, + useAllLocalModels: false, enableAllFrontierModels: true, useNovaGateway: false) + XCTAssertEqual(pool.count, 1) + XCTAssertEqual(pool[0].backend, .openRouter) + } + + func testAssemblePoolBothPlusNova() { + let s = samplePool() + let pool = ModelRegistry.assemblePool( + ollama: [s.local[0]], mlx: [s.local[1]], frontier: s.frontier, novaGateway: s.nova, + useAllLocalModels: true, enableAllFrontierModels: true, useNovaGateway: true) + XCTAssertEqual(pool.count, 4) + XCTAssertTrue(pool.contains { $0.backend == .novaGateway }) + } + + func testAssemblePoolNovaAbsentWhenToggleOff() { + let s = samplePool() + let pool = ModelRegistry.assemblePool( + ollama: [s.local[0]], mlx: [s.local[1]], frontier: s.frontier, novaGateway: s.nova, + useAllLocalModels: true, enableAllFrontierModels: false, useNovaGateway: false) + XCTAssertFalse(pool.contains { $0.backend == .novaGateway }) + } + + func testAssemblePoolAllOff() { + let s = samplePool() + let pool = ModelRegistry.assemblePool( + ollama: [s.local[0]], mlx: [s.local[1]], frontier: s.frontier, novaGateway: s.nova, + useAllLocalModels: false, enableAllFrontierModels: false, useNovaGateway: false) + XCTAssertTrue(pool.isEmpty) + } + + func testAssemblePoolDeduplicates() { + let dup = DiscoveredModel(modelName: "mistral:latest", backend: .ollama, endpoint: "e") + let pool = ModelRegistry.assemblePool( + ollama: [dup, dup], mlx: [], frontier: [], novaGateway: nil, + useAllLocalModels: true, enableAllFrontierModels: false, useNovaGateway: false) + XCTAssertEqual(pool.count, 1) + } + + // MARK: - Round-robin policy + + func testRoundRobinCyclesAndWraps() { + let pool = ["a", "b", "c"].map { DiscoveredModel(modelName: $0, backend: .ollama, endpoint: "e") } + let health = Dictionary(uniqueKeysWithValues: pool.map { ($0.id, true) }) + let lb = LoadBalancer() + + var picks: [String] = [] + for _ in 0..<7 { + picks.append(lb.next(pool: pool, health: health, policy: .roundRobin)!.modelName) + } + XCTAssertEqual(picks, ["a", "b", "c", "a", "b", "c", "a"]) + } + + // MARK: - Least-busy policy + + func testLeastBusyPicksLowestInFlight() { + let pool = ["a", "b", "c"].map { DiscoveredModel(modelName: $0, backend: .ollama, endpoint: "e") } + let health = Dictionary(uniqueKeysWithValues: pool.map { ($0.id, true) }) + let lb = LoadBalancer() + + // a: 2 in-flight, b: 0, c: 1 → b is least busy. + lb.checkOut(pool[0].id); lb.checkOut(pool[0].id) + lb.checkOut(pool[2].id) + XCTAssertEqual(lb.next(pool: pool, health: health, policy: .leastBusy)?.modelName, "b") + } + + func testLeastBusyTieBreaksByPoolOrder() { + let pool = ["a", "b", "c"].map { DiscoveredModel(modelName: $0, backend: .ollama, endpoint: "e") } + let health = Dictionary(uniqueKeysWithValues: pool.map { ($0.id, true) }) + let lb = LoadBalancer() + // All zero in-flight → first in pool order wins, deterministically. + XCTAssertEqual(lb.next(pool: pool, health: health, policy: .leastBusy)?.modelName, "a") + } + + func testCheckInNeverGoesNegative() { + let lb = LoadBalancer() + lb.checkIn("x") + XCTAssertEqual(lb.inFlight["x"], 0) + lb.checkOut("x"); lb.checkIn("x"); lb.checkIn("x") + XCTAssertEqual(lb.inFlight["x"], 0) + } + + // MARK: - Health gating + + func testHealthMapExcludesUnhealthy() { + let pool = ["a", "b", "c"].map { DiscoveredModel(modelName: $0, backend: .ollama, endpoint: "e") } + let lb = LoadBalancer() + // b marked unhealthy; a & c absent-from-map default to healthy. + let health = [pool[1].id: false] + let picked = (0..<4).map { _ in lb.next(pool: pool, health: health, policy: .roundRobin)!.modelName } + XCTAssertFalse(picked.contains("b")) + XCTAssertEqual(Set(picked), ["a", "c"]) + } + + func testAllUnhealthyReturnsNil() { + let pool = ["a", "b"].map { DiscoveredModel(modelName: $0, backend: .ollama, endpoint: "e") } + let health = Dictionary(uniqueKeysWithValues: pool.map { ($0.id, false) }) + let lb = LoadBalancer() + XCTAssertNil(lb.next(pool: pool, health: health, policy: .roundRobin)) + XCTAssertNil(lb.next(pool: pool, health: health, policy: .leastBusy)) + } + + func testEmptyPoolReturnsNil() { + let lb = LoadBalancer() + XCTAssertNil(lb.next(pool: [], health: [:], policy: .roundRobin)) + } + + // MARK: - OpenAI-compatible request building (network-free) + + func testOpenAICompatibleRequestBuild() throws { + let request = try OpenAICompatibleRequest.build( + endpoint: "https://openrouter.ai/api/v1/chat/completions", + model: "openai/gpt-4o", + messages: [["role": "user", "content": "hi"]], + temperature: 0.5, + maxTokens: 128, + stream: false, + headers: ["Authorization": "Bearer test"] + ) + XCTAssertEqual(request.httpMethod, "POST") + XCTAssertEqual(request.value(forHTTPHeaderField: "Authorization"), "Bearer test") + XCTAssertEqual(request.value(forHTTPHeaderField: "Content-Type"), "application/json") + let body = try XCTUnwrap(request.httpBody) + let json = try XCTUnwrap(try JSONSerialization.jsonObject(with: body) as? [String: Any]) + XCTAssertEqual(json["model"] as? String, "openai/gpt-4o") + XCTAssertEqual(json["stream"] as? Bool, false) + } + + func testOpenAICompatibleRequestInvalidURLThrows() { + XCTAssertThrowsError(try OpenAICompatibleRequest.build( + endpoint: "", + model: "m", + messages: [], + temperature: 0.5, + maxTokens: 10, + stream: false + )) + } + + func testParseOpenRouterModels() { + let json = #"{"data":[{"id":"openai/gpt-4o"},{"id":"anthropic/claude-sonnet-4.5"},{"foo":"bar"}]}"# + XCTAssertEqual(OpenRouterProvider.parseModels(Data(json.utf8)), + ["openai/gpt-4o", "anthropic/claude-sonnet-4.5"]) + XCTAssertTrue(OpenRouterProvider.parseModels(Data("garbage".utf8)).isEmpty) + } + + // MARK: - Backend type + + func testNovaGatewayBackendType() { + XCTAssertEqual(LLMBackendType.novaGateway.rawValue, "novagateway") + XCTAssertEqual(LLMBackendType.novaGateway.displayName, "Nova Gateway") + XCTAssertEqual(LLMBackendType.novaGateway.defaultURL, "http://127.0.0.1:18792") + XCTAssertFalse(LLMBackendType.novaGateway.icon.isEmpty) + // MLX Code's balancer backend set: ollama, mlx, openRouter, novaGateway. + XCTAssertEqual(LLMBackendType.allCases.count, 4) + } +} diff --git a/MLX Code.xcodeproj/project.pbxproj b/MLX Code.xcodeproj/project.pbxproj index d36af77..f56b280 100644 --- a/MLX Code.xcodeproj/project.pbxproj +++ b/MLX Code.xcodeproj/project.pbxproj @@ -7,6 +7,11 @@ objects = { /* Begin PBXBuildFile section */ + AABB01000000000000000002 /* ModelRegistry.swift in Sources */ = {isa = PBXBuildFile; fileRef = AABB01000000000000000001 /* ModelRegistry.swift */; }; + AABB02000000000000000002 /* OpenRouterProvider.swift in Sources */ = {isa = PBXBuildFile; fileRef = AABB02000000000000000001 /* OpenRouterProvider.swift */; }; + AABB03000000000000000002 /* KeychainStore.swift in Sources */ = {isa = PBXBuildFile; fileRef = AABB03000000000000000001 /* KeychainStore.swift */; }; + AABB04000000000000000002 /* LLMBalancer.swift in Sources */ = {isa = PBXBuildFile; fileRef = AABB04000000000000000001 /* LLMBalancer.swift */; }; + AABB05000000000000000002 /* LoadBalancerTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = AABB05000000000000000001 /* LoadBalancerTests.swift */; }; 043C4E7F3FCA17ECD6EB1A51 /* PromptTemplatesView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 942A3F3878CD789BDF616A78 /* PromptTemplatesView.swift */; }; 0475B2DD296EECD045126C17 /* PrerequisitesView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 22BF0FCAFA31E5516E7B5EC6 /* PrerequisitesView.swift */; }; 054C45D18234B5AE6A2A3233 /* GettingStarted.md in Resources */ = {isa = PBXBuildFile; fileRef = 3072558342FD102B940FF3DB /* GettingStarted.md */; }; @@ -178,6 +183,11 @@ /* End PBXCopyFilesBuildPhase section */ /* Begin PBXFileReference section */ + AABB01000000000000000001 /* ModelRegistry.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ModelRegistry.swift; sourceTree = ""; }; + AABB02000000000000000001 /* OpenRouterProvider.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = OpenRouterProvider.swift; sourceTree = ""; }; + AABB03000000000000000001 /* KeychainStore.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = KeychainStore.swift; sourceTree = ""; }; + AABB04000000000000000001 /* LLMBalancer.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = LLMBalancer.swift; sourceTree = ""; }; + AABB05000000000000000001 /* LoadBalancerTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = LoadBalancerTests.swift; sourceTree = ""; }; 0095C89C8F1E4539863AE74D96B702F4 /* MLXService.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MLXService.swift; sourceTree = ""; }; 01A056387C5D6883F0932FA4 /* GitService.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; name = GitService.swift; path = "MLX Code/Services/GitService.swift"; sourceTree = ""; }; 0674F2D9E84C4AC3AC8ED170619E99EA /* FileService.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FileService.swift; sourceTree = ""; }; @@ -487,6 +497,10 @@ 8D7B12981853422D971FCAD3B063E827 /* Services */ = { isa = PBXGroup; children = ( + AABB01000000000000000001 /* ModelRegistry.swift */, + AABB02000000000000000001 /* OpenRouterProvider.swift */, + AABB03000000000000000001 /* KeychainStore.swift */, + AABB04000000000000000001 /* LLMBalancer.swift */, 0095C89C8F1E4539863AE74D96B702F4 /* MLXService.swift */, BB001CDEF123456789ABCDE000022222 /* XcodeActionHandler.swift */, 2113E8E02AFC404C917BDECE45AD0C31 /* XcodeService.swift */, @@ -523,6 +537,7 @@ CE911C9235A55156197C996A /* MLX Code Tests */ = { isa = PBXGroup; children = ( + AABB05000000000000000001 /* LoadBalancerTests.swift */, 1B364AD048BC9F2567D45CC7 /* AppSettingsTests.swift */, 8C37CB4BC161157E297B1C70 /* CommandValidatorTests.swift */, F975089CB3BD79BDBC9F5CD3 /* ContextManagerTests.swift */, @@ -859,6 +874,7 @@ isa = PBXSourcesBuildPhase; buildActionMask = 2147483647; files = ( + AABB05000000000000000002 /* LoadBalancerTests.swift in Sources */, 678CCB7627496E9731B7D734 /* AppSettingsTests.swift in Sources */, AA74C582F22930DF98A0CB7B /* CommandValidatorTests.swift in Sources */, 2CAAE35717EEF0F0BB6889D1 /* ContextManagerTests.swift in Sources */, @@ -899,6 +915,10 @@ isa = PBXSourcesBuildPhase; buildActionMask = 2147483647; files = ( + AABB01000000000000000002 /* ModelRegistry.swift in Sources */, + AABB02000000000000000002 /* OpenRouterProvider.swift in Sources */, + AABB03000000000000000002 /* KeychainStore.swift in Sources */, + AABB04000000000000000002 /* LLMBalancer.swift in Sources */, E2FCF04057F99E48BF46CCF3 /* CommandValidator.swift in Sources */, A0013163E03EA758655F244E /* ModelSecurityValidator.swift in Sources */, 4A14B6380E9E4A44949059C9DE2D703A /* MLXCodeApp.swift in Sources */, diff --git a/MLX Code/Models/AppSettings.swift b/MLX Code/Models/AppSettings.swift index a4f2c58..3a66a66 100644 --- a/MLX Code/Models/AppSettings.swift +++ b/MLX Code/Models/AppSettings.swift @@ -108,6 +108,22 @@ class AppSettings: ObservableObject { /// Enable user memories in LLM system prompt @Published var enableMemories: Bool = true + // MARK: - Multi-Model Load Balancer Settings + + /// Balance across ALL installed local models (native MLX + Ollama), not just + /// the single pinned model. Off by default — preserves existing behavior. + @Published var useAllLocalModels: Bool = false + + /// Include all frontier models (via an OpenRouter API key) in the balancer pool. + @Published var enableAllFrontierModels: Bool = false + + /// Include the optional Nova Gateway as one balancer entry. Nova is never a + /// hard requirement — if its health probe fails it simply drops from the pool. + @Published var useNovaGateway: Bool = false + + /// Nova Gateway base URL (OpenAI-compatible; health via `/v1/models`). + @Published var novaGatewayURL: String = ModelRegistry.novaGatewayDefaultURL + // MARK: - Private Properties private let userDefaults = UserDefaults.standard @@ -141,6 +157,10 @@ class AppSettings: ObservableObject { static let credentialScanOnPush = "credentialScanOnPush" static let authorName = "authorName" static let enableMemories = "enableMemories" + static let useAllLocalModels = "useAllLocalModels" + static let enableAllFrontierModels = "enableAllFrontierModels" + static let useNovaGateway = "useNovaGateway" + static let novaGatewayURL = "novaGatewayURL" } // MARK: - Initialization @@ -324,6 +344,20 @@ class AppSettings: ObservableObject { enableMemories = userDefaults.bool(forKey: Keys.enableMemories) } + // Load load-balancer settings + if userDefaults.object(forKey: Keys.useAllLocalModels) != nil { + useAllLocalModels = userDefaults.bool(forKey: Keys.useAllLocalModels) + } + if userDefaults.object(forKey: Keys.enableAllFrontierModels) != nil { + enableAllFrontierModels = userDefaults.bool(forKey: Keys.enableAllFrontierModels) + } + if userDefaults.object(forKey: Keys.useNovaGateway) != nil { + useNovaGateway = userDefaults.bool(forKey: Keys.useNovaGateway) + } + if let url = userDefaults.string(forKey: Keys.novaGatewayURL), !url.isEmpty { + novaGatewayURL = url + } + // Load available models if let modelsData = userDefaults.data(forKey: Keys.availableModels), let models = try? JSONDecoder().decode([MLXModel].self, from: modelsData) { @@ -372,6 +406,12 @@ class AppSettings: ObservableObject { userDefaults.set(authorName, forKey: Keys.authorName) userDefaults.set(enableMemories, forKey: Keys.enableMemories) + // Save load-balancer settings + userDefaults.set(useAllLocalModels, forKey: Keys.useAllLocalModels) + userDefaults.set(enableAllFrontierModels, forKey: Keys.enableAllFrontierModels) + userDefaults.set(useNovaGateway, forKey: Keys.useNovaGateway) + userDefaults.set(novaGatewayURL, forKey: Keys.novaGatewayURL) + // Save selected model ID if let modelId = selectedModel?.id.uuidString { userDefaults.set(modelId, forKey: Keys.selectedModelId) @@ -635,6 +675,31 @@ class AppSettings: ObservableObject { .debounce(for: .seconds(1.0), scheduler: DispatchQueue.main) .sink { [weak self] _ in self?.saveSettings() } .store(in: &cancellables) + + // Observe load-balancer settings + $useAllLocalModels + .dropFirst() + .debounce(for: .seconds(0.5), scheduler: DispatchQueue.main) + .sink { [weak self] _ in self?.saveSettings() } + .store(in: &cancellables) + + $enableAllFrontierModels + .dropFirst() + .debounce(for: .seconds(0.5), scheduler: DispatchQueue.main) + .sink { [weak self] _ in self?.saveSettings() } + .store(in: &cancellables) + + $useNovaGateway + .dropFirst() + .debounce(for: .seconds(0.5), scheduler: DispatchQueue.main) + .sink { [weak self] _ in self?.saveSettings() } + .store(in: &cancellables) + + $novaGatewayURL + .dropFirst() + .debounce(for: .seconds(1.0), scheduler: DispatchQueue.main) + .sink { [weak self] _ in self?.saveSettings() } + .store(in: &cancellables) } } diff --git a/MLX Code/Services/KeychainStore.swift b/MLX Code/Services/KeychainStore.swift new file mode 100644 index 0000000..7cc2f4b --- /dev/null +++ b/MLX Code/Services/KeychainStore.swift @@ -0,0 +1,80 @@ +// +// KeychainStore.swift +// MLX Code +// +// Ported from AIStudio by Jordan Koch on 2026-08-18. +// Copyright © 2026 Jordan Koch. All rights reserved. +// + +import Foundation +import Security + +/// Minimal macOS Keychain wrapper for storing secrets (e.g. API keys). +/// +/// Secrets are stored as classic generic-password items in the login keychain so +/// that the round-trip works in unsigned/unsandboxed unit-test processes (no +/// data-protection keychain, which would require entitlements). Never store +/// secrets in UserDefaults. +struct KeychainStore { + let service: String + let account: String + + /// - Parameters: + /// - service: Keychain service identifier. Defaults to the OpenRouter service. + /// - account: Account/key name within the service. + init(service: String = OpenRouterProvider.keychainService, account: String = "apiKey") { + self.service = service + self.account = account + } + + private var baseQuery: [String: Any] { + [ + kSecClass as String: kSecClassGenericPassword, + kSecAttrService as String: service, + kSecAttrAccount as String: account + ] + } + + /// Store (or replace) a secret. Returns true on success. + @discardableResult + func set(_ value: String) -> Bool { + guard let data = value.data(using: .utf8) else { return false } + + // Remove any existing item first so we can cleanly re-add. + SecItemDelete(baseQuery as CFDictionary) + + var attributes = baseQuery + attributes[kSecValueData as String] = data + let status = SecItemAdd(attributes as CFDictionary, nil) + return status == errSecSuccess + } + + /// Retrieve a secret, or nil if none is stored. + func get() -> String? { + var query = baseQuery + query[kSecReturnData as String] = true + query[kSecMatchLimit as String] = kSecMatchLimitOne + + var result: AnyObject? + let status = SecItemCopyMatching(query as CFDictionary, &result) + guard status == errSecSuccess, + let data = result as? Data, + let value = String(data: data, encoding: .utf8) else { + return nil + } + return value + } + + /// Delete the stored secret. Returns true if an item was removed or none existed. + @discardableResult + func delete() -> Bool { + let status = SecItemDelete(baseQuery as CFDictionary) + return status == errSecSuccess || status == errSecItemNotFound + } + + /// True if a non-empty secret is stored. + var hasValue: Bool { + guard let value = get() else { return false } + return !value.isEmpty + } +} diff --git a/MLX Code/Services/LLMBalancer.swift b/MLX Code/Services/LLMBalancer.swift new file mode 100644 index 0000000..19ba03a --- /dev/null +++ b/MLX Code/Services/LLMBalancer.swift @@ -0,0 +1,380 @@ +// +// LLMBalancer.swift +// MLX Code +// +// Created by Jordan Koch on 2026-08-18. +// Copyright © 2026 Jordan Koch. All rights reserved. +// +// Multi-model load balancer for MLX Code. Wires the reusable, network-free +// `ModelRegistry` / `LoadBalancer` pieces (ported from AIStudio) into MLX Code's +// own model call path so that — when the settings toggles are on — chat can be +// spread across ALL installed local models (native MLX + Ollama), all frontier +// models (via an OpenRouter key), and the optional Nova Gateway, load-balanced +// and health-gated, instead of a single pinned MLX model. +// +// Nova is NEVER a hard requirement: with zero Nova present the balancer still +// works over local MLX/Ollama models and/or OpenRouter. The Nova Gateway is just +// one optional entry whose health probe failing simply removes it from the pool. +// + +import Foundation + +// MARK: - Backend type + +/// LLM backend type identifier for the balancer pool. MLX Code is MLX-native, but +/// the balancer can additionally spread work to Ollama, OpenRouter frontier +/// models, and the optional Nova Gateway. +enum LLMBackendType: String, CaseIterable, Codable, Sendable { + case ollama = "ollama" + case mlx = "mlx" + case openRouter = "openrouter" + case novaGateway = "novagateway" + + var displayName: String { + switch self { + case .ollama: return "Ollama" + case .mlx: return "MLX Native" + case .openRouter: return "OpenRouter (Frontier Models)" + case .novaGateway: return "Nova Gateway" + } + } + + var icon: String { + switch self { + case .ollama: return "network" + case .mlx: return "cpu" + case .openRouter: return "cloud" + case .novaGateway: return "sparkle.magnifyingglass" + } + } + + var defaultURL: String { + switch self { + case .ollama: return ModelRegistry.ollamaBaseURL + case .mlx: return "" + case .openRouter: return OpenRouterProvider.baseURL + case .novaGateway: return ModelRegistry.novaGatewayDefaultURL + } + } +} + +// MARK: - Errors + +/// Errors thrown by the balanced dispatch path. +enum LLMError: LocalizedError { + case invalidURL + case noBackendAvailable + + var errorDescription: String? { + switch self { + case .invalidURL: return "The backend URL is invalid" + case .noBackendAvailable: return "No LLM backend is currently available" + } + } +} + +/// The balancer reuses MLX Code's own `Message` type as its chat message shape. +typealias ChatMessage = Message + +// MARK: - Balancer + +/// Load-balanced multi-model dispatch. Mirrors AIStudio's `LLMBackendManager` +/// balancer surface (`discoverEnabledPool` / `healthMap` / `generateBalanced` / +/// `dispatchBalanced`) but routed through MLX Code's in-process `MLXService` for +/// local MLX inference and OpenAI-compatible HTTP for the other backends. +actor LLMBalancer { + static let shared = LLMBalancer() + + /// Pure round-robin / least-busy selector over the healthy pool. + private let balancer = LoadBalancer() + + /// Selection policy — least-busy by default (spreads concurrent load). + var policy: BalancerPolicy = .leastBusy + + /// The most recently discovered pool (for diagnostics / UI). + private(set) var discoveredModels: [DiscoveredModel] = [] + + /// Cached OpenRouter model ids (fetched once, falls back to the popular set). + private var cachedFrontierIDs: [String]? + + private let session: URLSession = .shared + + private init() {} + + // MARK: Toggle snapshot + + /// A minimal, `Sendable` snapshot of the three balancer toggles + gateway URL, + /// read from the `@MainActor` `AppSettings`. + private struct ToggleSnapshot: Sendable { + let useAllLocalModels: Bool + let enableAllFrontierModels: Bool + let useNovaGateway: Bool + let novaGatewayURL: String + } + + private func snapshot() async -> ToggleSnapshot { + await MainActor.run { + let s = AppSettings.shared + return ToggleSnapshot( + useAllLocalModels: s.useAllLocalModels, + enableAllFrontierModels: s.enableAllFrontierModels, + useNovaGateway: s.useNovaGateway, + novaGatewayURL: s.novaGatewayURL + ) + } + } + + /// True when any balancing toggle is on. When false, callers must preserve the + /// existing single-model behavior. + func isBalancingEnabled() async -> Bool { + let s = await snapshot() + return s.useAllLocalModels || s.enableAllFrontierModels || s.useNovaGateway + } + + /// Set the selection policy. + func setPolicy(_ newPolicy: BalancerPolicy) { + policy = newPolicy + } + + // MARK: Discovery + + /// The OpenRouter API key from the Keychain (empty if none stored). + private func openRouterKey() -> String { + KeychainStore(service: OpenRouterProvider.keychainService).get() ?? "" + } + + /// Discover the enabled balancer pool honoring the three toggles. Resilient: + /// any unreachable source simply contributes zero models. + func discoverEnabledPool() async -> [DiscoveredModel] { + let s = await snapshot() + + var ollama: [DiscoveredModel] = [] + var mlx: [DiscoveredModel] = [] + var frontier: [DiscoveredModel] = [] + + if s.useAllLocalModels { + ollama = await ModelRegistry.discoverOllama(baseURL: ModelRegistry.ollamaBaseURL, session: session) + // MLX Code discovers SafeTensors models from its own configured paths. + let localNames = ((try? await MLXService.shared.discoverModels()) ?? []).map { $0.name } + mlx = ModelRegistry.mlxModels(fromLocalNames: localNames) + } + if s.enableAllFrontierModels { + frontier = ModelRegistry.frontierModels(from: await frontierModelIDs()) + } + let nova = s.useNovaGateway ? ModelRegistry.novaGatewayModel(url: s.novaGatewayURL) : nil + + let pool = ModelRegistry.assemblePool( + ollama: ollama, + mlx: mlx, + frontier: frontier, + novaGateway: nova, + useAllLocalModels: s.useAllLocalModels, + enableAllFrontierModels: s.enableAllFrontierModels, + useNovaGateway: s.useNovaGateway + ) + discoveredModels = pool + return pool + } + + /// Fetch the live OpenRouter model list (cached), falling back to the popular + /// set on any failure or when no API key is present. + private func frontierModelIDs() async -> [String] { + if let cached = cachedFrontierIDs { return cached } + let key = openRouterKey() + guard !key.isEmpty, let url = URL(string: OpenRouterProvider.modelsURL) else { + return OpenRouterProvider.fallbackModels + } + var request = URLRequest(url: url) + for (k, v) in OpenRouterProvider.authHeaders(apiKey: key) { + request.setValue(v, forHTTPHeaderField: k) + } + do { + let (data, response) = try await session.data(for: request) + guard (response as? HTTPURLResponse)?.statusCode == 200 else { + cachedFrontierIDs = OpenRouterProvider.fallbackModels + return OpenRouterProvider.fallbackModels + } + let ids = OpenRouterProvider.parseModels(data) + let result = ids.isEmpty ? OpenRouterProvider.fallbackModels : ids + cachedFrontierIDs = result + return result + } catch { + return OpenRouterProvider.fallbackModels + } + } + + // MARK: Health + + /// Build a `[modelId: Bool]` health map for `pool` by probing each distinct + /// backend once. This is the health-gating that lets an unavailable backend + /// (e.g. the Nova Gateway) drop out while everything else keeps working. + private func healthMap(for pool: [DiscoveredModel], novaURL: String) async -> [String: Bool] { + var backendHealth: [LLMBackendType: Bool] = [:] + for backend in Set(pool.map { $0.backend }) { + backendHealth[backend] = await checkAvailability(backend, novaURL: novaURL) + } + var map: [String: Bool] = [:] + for model in pool { + map[model.id] = backendHealth[model.backend] ?? false + } + return map + } + + /// Live availability probe for a single backend. Never throws. + private func checkAvailability(_ backend: LLMBackendType, novaURL: String) async -> Bool { + switch backend { + case .mlx: + // In-process; available only when a model is actually loaded. + return await MLXService.shared.isLoaded() + case .ollama: + return await httpOK("\(ModelRegistry.ollamaBaseURL)/api/tags") + case .openRouter: + return !openRouterKey().isEmpty + case .novaGateway: + return await httpOK("\(novaURL)/v1/models") + } + } + + /// GET the URL and report whether it returned HTTP 200. Any failure → false. + private func httpOK(_ urlString: String) async -> Bool { + guard let url = URL(string: urlString) else { return false } + var request = URLRequest(url: url) + request.timeoutInterval = 3 + do { + let (_, response) = try await session.data(for: request) + return (response as? HTTPURLResponse)?.statusCode == 200 + } catch { + return false + } + } + + // MARK: Balanced dispatch + + /// Balanced chat completion over the healthy enabled pool. Returns `nil` when + /// the pool is empty (caller should fall back to the single-model path). + /// Throws only when every healthy candidate failed mid-flight. + /// + /// For the in-process MLX backend the real `streamHandler` is passed straight + /// through (true token streaming). For HTTP backends the completed response is + /// delivered to `streamHandler` in one call. + func chatCompletionBalanced( + messages: [Message], + parameters: ModelParameters?, + streamHandler: ((String) -> Void)? = nil + ) async throws -> String? { + let s = await snapshot() + let pool = await discoverEnabledPool() + guard !pool.isEmpty else { return nil } + + let health = await healthMap(for: pool, novaURL: s.novaGatewayURL) + var remaining = pool + var lastError: Error? + + while let choice = balancer.next(pool: remaining, health: health, policy: policy) { + balancer.checkOut(choice.id) + do { + let result = try await dispatchBalanced( + model: choice, + messages: messages, + parameters: parameters, + streamHandler: streamHandler + ) + balancer.checkIn(choice.id) + return result + } catch { + balancer.checkIn(choice.id) + lastError = error + remaining.removeAll { $0.id == choice.id } + continue + } + } + + if let lastError = lastError { throw lastError } + return nil + } + + /// Route a single balancer-selected model to its backend implementation. + private func dispatchBalanced( + model: DiscoveredModel, + messages: [Message], + parameters: ModelParameters?, + streamHandler: ((String) -> Void)? + ) async throws -> String { + switch model.backend { + case .mlx: + // In-process native inference through the existing MLX path. + return try await MLXService.shared.chatCompletion( + messages: messages, + parameters: parameters, + streamHandler: streamHandler + ) + case .ollama: + let content = try await postOpenAICompatible( + endpoint: "\(ModelRegistry.ollamaBaseURL)/v1/chat/completions", + model: model.modelName, + headers: [:], + messages: messages, + parameters: parameters + ) + streamHandler?(content) + return content + case .openRouter: + let key = openRouterKey() + guard !key.isEmpty else { throw LLMError.noBackendAvailable } + let content = try await postOpenAICompatible( + endpoint: model.endpoint, + model: model.modelName, + headers: OpenRouterProvider.authHeaders(apiKey: key), + messages: messages, + parameters: parameters + ) + streamHandler?(content) + return content + case .novaGateway: + let content = try await postOpenAICompatible( + endpoint: model.endpoint, + model: model.modelName, + headers: [:], + messages: messages, + parameters: parameters + ) + streamHandler?(content) + return content + } + } + + /// POST an OpenAI-compatible chat-completions request and return the assistant + /// message content. Non-streaming. + private func postOpenAICompatible( + endpoint: String, + model: String, + headers: [String: String], + messages: [Message], + parameters: ModelParameters? + ) async throws -> String { + let params = parameters ?? ModelParameters() + let payloadMessages: [[String: String]] = messages.map { + ["role": $0.role.rawValue, "content": $0.content] + } + let request = try OpenAICompatibleRequest.build( + endpoint: endpoint, + model: model, + messages: payloadMessages, + temperature: Float(params.temperature), + maxTokens: params.maxTokens, + stream: false, + headers: headers + ) + let (data, response) = try await session.data(for: request) + guard (response as? HTTPURLResponse)?.statusCode == 200 else { + throw LLMError.noBackendAvailable + } + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let choices = json["choices"] as? [[String: Any]], + let message = choices.first?["message"] as? [String: Any], + let content = message["content"] as? String else { + throw LLMError.noBackendAvailable + } + return content + } +} diff --git a/MLX Code/Services/ModelRegistry.swift b/MLX Code/Services/ModelRegistry.swift new file mode 100644 index 0000000..97ffc05 --- /dev/null +++ b/MLX Code/Services/ModelRegistry.swift @@ -0,0 +1,241 @@ +// +// ModelRegistry.swift +// MLX Code +// +// Ported from AIStudio by Jordan Koch on 2026-08-18. +// Copyright © 2026 Jordan Koch. All rights reserved. +// +// Discovers every model available on the machine and normalizes them into a +// flat `[DiscoveredModel]` pool that the load balancer can spread work across — +// the single-user version of how Nova's gateway balances load. +// +// The parsing/composition logic is factored into pure, network-free functions so +// it is fully unit-testable; the thin discovery wrappers do the actual I/O and +// never throw to the caller (an unreachable backend simply contributes zero +// models). +// + +import Foundation + +// MARK: - Discovered model + +/// A single model discovered on the machine, normalized across backends. +struct DiscoveredModel: Identifiable, Hashable, Sendable { + /// Stable, pool-unique identifier (`"|"`). + let id: String + /// The raw model name/id passed to the backend API (e.g. `mistral:latest`). + let modelName: String + /// Human-friendly label for pickers. + let displayName: String + /// Which backend serves this model. + let backend: LLMBackendType + /// Base or chat-completions endpoint the model is reachable at (informational + /// for MLX, which runs in-process). + let endpoint: String + + init(modelName: String, displayName: String? = nil, backend: LLMBackendType, endpoint: String) { + self.id = "\(backend.rawValue)|\(modelName)" + self.modelName = modelName + self.displayName = displayName ?? modelName + self.backend = backend + self.endpoint = endpoint + } +} + +// MARK: - Model registry + +/// Discovers and normalizes the models available on this machine. +/// +/// All the JSON/structured-input → `[DiscoveredModel]` mapping lives in pure, +/// network-free functions (`parseOllamaTags`, `parseMLXModels`, `frontierModels`, +/// `assemblePool`) so they can be unit-tested without hitting the network. The +/// `discover*` wrappers add the thin I/O layer and swallow every error. +enum ModelRegistry { + /// Default local Ollama base URL. + static let ollamaBaseURL = "http://localhost:11434" + /// Default Nova Gateway base URL (OpenAI-compatible, inherits Nova's routing). + static let novaGatewayDefaultURL = "http://127.0.0.1:18792" + + // MARK: Pure parsing (network-free, unit-tested) + + /// Map an Ollama `/api/tags` response body to `[DiscoveredModel]`. + /// Returns `[]` for empty/garbage input — never throws. + static func parseOllamaTags(_ data: Data, baseURL: String = ollamaBaseURL) -> [DiscoveredModel] { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let models = json["models"] as? [[String: Any]] else { + return [] + } + let endpoint = "\(baseURL)/api/chat" + return models.compactMap { entry -> DiscoveredModel? in + guard let name = entry["name"] as? String, !name.isEmpty else { return nil } + return DiscoveredModel(modelName: name, backend: .ollama, endpoint: endpoint) + } + } + + /// Map a set of Hugging Face hub cache directory names to locally-installed + /// MLX models. A hub directory is named `models----`; this converts + /// it back to the `/` model id and keeps only MLX models. Pure and + /// network-free; the discovery wrapper feeds it the real directory listing. + static func parseMLXModels(hubDirectoryNames names: [String]) -> [DiscoveredModel] { + names.compactMap { raw -> DiscoveredModel? in + guard raw.hasPrefix("models--") else { return nil } + let repo = raw.dropFirst("models--".count).replacingOccurrences(of: "--", with: "/") + guard !repo.isEmpty, repo.lowercased().contains("mlx") else { return nil } + let short = repo.split(separator: "/").last.map(String.init) ?? repo + // MLX runs in-process — no HTTP endpoint. + return DiscoveredModel(modelName: repo, displayName: short, backend: .mlx, endpoint: "") + } + } + + /// Map already-discovered local MLX model names into the registry. MLX Code + /// discovers SafeTensors models from its own configured models directory + /// (`MLXService.discoverModels()`) rather than the HF hub cache, so this thin + /// mapper is the app-specific MLX discovery source. Pure and network-free. + static func mlxModels(fromLocalNames names: [String]) -> [DiscoveredModel] { + names.compactMap { name -> DiscoveredModel? in + let trimmed = name.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return nil } + let short = trimmed.split(separator: "/").last.map(String.init) ?? trimmed + return DiscoveredModel(modelName: trimmed, displayName: short, backend: .mlx, endpoint: "") + } + } + + /// Map OpenRouter model ids (from `OpenRouterProvider.parseModels`) into the + /// registry as frontier models. Pure and network-free. + static func frontierModels(from openRouterModelIds: [String]) -> [DiscoveredModel] { + openRouterModelIds.compactMap { id -> DiscoveredModel? in + guard !id.isEmpty else { return nil } + return DiscoveredModel(modelName: id, backend: .openRouter, endpoint: OpenRouterProvider.chatCompletionsURL) + } + } + + /// The single Nova Gateway "model" — the app routes to Nova and inherits her + /// own internal routing, so it presents as one balancer entry. + static func novaGatewayModel(url: String = novaGatewayDefaultURL) -> DiscoveredModel { + DiscoveredModel( + modelName: "nova", + displayName: "Nova Gateway", + backend: .novaGateway, + endpoint: "\(url)/v1/chat/completions" + ) + } + + /// Compose the enabled balancer pool from per-source model lists and the three + /// toggles. Pure and network-free — this is the toggle-composition contract the + /// load balancer runs over. + static func assemblePool( + ollama: [DiscoveredModel] = [], + mlx: [DiscoveredModel] = [], + frontier: [DiscoveredModel] = [], + novaGateway: DiscoveredModel? = nil, + useAllLocalModels: Bool, + enableAllFrontierModels: Bool, + useNovaGateway: Bool + ) -> [DiscoveredModel] { + var pool: [DiscoveredModel] = [] + if useAllLocalModels { + pool.append(contentsOf: ollama) + pool.append(contentsOf: mlx) + } + if enableAllFrontierModels { + pool.append(contentsOf: frontier) + } + if useNovaGateway, let nova = novaGateway { + pool.append(nova) + } + // De-duplicate by id while preserving first-seen order. + var seen = Set() + return pool.filter { seen.insert($0.id).inserted } + } + + // MARK: Discovery I/O (thin, resilient — never throws) + + /// Discover local Ollama models. Any failure → `[]`. + static func discoverOllama(baseURL: String = ollamaBaseURL, session: URLSession = .shared) async -> [DiscoveredModel] { + guard let url = URL(string: "\(baseURL)/api/tags") else { return [] } + do { + let (data, response) = try await session.data(from: url) + guard (response as? HTTPURLResponse)?.statusCode == 200 else { return [] } + return parseOllamaTags(data, baseURL: baseURL) + } catch { + return [] + } + } + + /// Discover locally-installed MLX models from the Hugging Face hub cache. Any + /// failure (no cache dir, unreadable) → `[]`. + static func discoverMLX(hubPath: String? = nil) -> [DiscoveredModel] { + let path = hubPath ?? (NSHomeDirectory() as NSString).appendingPathComponent(".cache/huggingface/hub") + guard let entries = try? FileManager.default.contentsOfDirectory(atPath: path) else { return [] } + return parseMLXModels(hubDirectoryNames: entries) + } +} + +// MARK: - Load balancer + +/// Balancer selection policy. +enum BalancerPolicy: String, CaseIterable, Sendable { + /// Cycle through the pool in order, wrapping around. + case roundRobin + /// Prefer the model with the fewest in-flight requests. + case leastBusy +} + +/// Pure, network-free load balancer over a `[DiscoveredModel]` pool. Given the +/// pool, a per-model health map and a policy, it returns the next model to use — +/// no network, so it is fully unit-testable. In-flight counts (for `.leastBusy`) +/// are tracked internally via `checkOut`/`checkIn`. +/// +/// Composes with `FailoverPlanner`: the manager gates the pool to healthy backends +/// first (skip unhealthy, fall through), then the balancer spreads load across +/// what remains. A model is eligible unless the health map marks it `false`. +final class LoadBalancer { + /// In-flight request count per model id. + private(set) var inFlight: [String: Int] = [:] + /// Rolling cursor for round-robin. + private var cursor: Int = 0 + + init() {} + + /// The healthy subset of `pool`, preserving pool order. A model is healthy + /// unless the map explicitly marks it `false`. + func healthy(in pool: [DiscoveredModel], health: [String: Bool]) -> [DiscoveredModel] { + pool.filter { health[$0.id] != false } + } + + /// Select the next model from the healthy subset of `pool` under `policy`. + /// Returns `nil` when nothing is healthy (the manager then falls back cleanly). + func next(pool: [DiscoveredModel], health: [String: Bool] = [:], policy: BalancerPolicy) -> DiscoveredModel? { + let candidates = healthy(in: pool, health: health) + guard !candidates.isEmpty else { return nil } + + switch policy { + case .roundRobin: + let choice = candidates[cursor % candidates.count] + cursor += 1 + return choice + case .leastBusy: + // Lowest in-flight count wins; ties broken by pool order (first). + return candidates.min { lhs, rhs in + (inFlight[lhs.id] ?? 0) < (inFlight[rhs.id] ?? 0) + } + } + } + + /// Mark a request as started against `modelId`. + func checkOut(_ modelId: String) { + inFlight[modelId, default: 0] += 1 + } + + /// Mark a request against `modelId` as finished. + func checkIn(_ modelId: String) { + let current = inFlight[modelId] ?? 0 + inFlight[modelId] = max(0, current - 1) + } + + /// Reset all balancer state (cursor + in-flight counts). + func reset() { + inFlight.removeAll() + cursor = 0 + } +} diff --git a/MLX Code/Services/OpenRouterProvider.swift b/MLX Code/Services/OpenRouterProvider.swift new file mode 100644 index 0000000..be4ae10 --- /dev/null +++ b/MLX Code/Services/OpenRouterProvider.swift @@ -0,0 +1,158 @@ +// +// OpenRouterProvider.swift +// MLX Code +// +// Ported from AIStudio by Jordan Koch on 2026-08-18. +// Copyright © 2026 Jordan Koch. All rights reserved. +// +// OpenRouter frontier-model access + the deterministic (network-free) pieces of +// the OpenAI-compatible request path and the automatic-failover selection logic. +// These are factored out as pure helpers so they can be unit-tested without +// hitting the network. +// + +import Foundation + +// MARK: - OpenRouter constants & helpers + +/// Static configuration and pure helpers for the OpenRouter provider. +enum OpenRouterProvider { + /// OpenAI-compatible base URL (already includes `/v1`). + static let baseURL = "https://openrouter.ai/api/v1" + /// Full chat-completions endpoint. + static var chatCompletionsURL: String { "\(baseURL)/chat/completions" } + /// Models listing endpoint. + static var modelsURL: String { "\(baseURL)/models" } + + /// Attribution headers required/recommended by OpenRouter. + static let referer = "https://github.com/kochj23/MLXCode" + static let title = "MLX Code" + + /// macOS Keychain service used to store the OpenRouter API key. + static let keychainService = "com.jordankoch.mlxcode.openrouter" + + /// Hardcoded fallback model list used when the live `/models` fetch fails. + /// A few popular current models spanning providers. + static let fallbackModels: [String] = [ + "anthropic/claude-sonnet-4.5", + "openai/gpt-4o", + "google/gemini-2.0-flash-001", + "meta-llama/llama-3.3-70b-instruct", + "deepseek/deepseek-chat" + ] + + /// Default model selected when none has been chosen yet. + static var defaultModel: String { fallbackModels.first ?? "openai/gpt-4o" } + + /// Auth + attribution headers for OpenRouter requests. + static func authHeaders(apiKey: String) -> [String: String] { + [ + "Authorization": "Bearer \(apiKey)", + "HTTP-Referer": referer, + "X-Title": title + ] + } + + /// Parse the model ids out of an OpenRouter `/models` response body. + /// Returns an empty array if the payload can't be parsed. + static func parseModels(_ data: Data) -> [String] { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let models = json["data"] as? [[String: Any]] else { + return [] + } + return models.compactMap { $0["id"] as? String } + } +} + +// MARK: - OpenAI-compatible request construction + +/// Pure builders for the generic OpenAI-compatible `/v1/chat/completions` request. +/// No network, no shared state — fully unit-testable. +enum OpenAICompatibleRequest { + /// Map a prompt + optional system prompt + prior history into the OpenAI + /// `messages` array shape. + static func chatMessages( + prompt: String, + systemPrompt: String?, + history: [ChatMessage] + ) -> [[String: String]] { + var messages: [[String: String]] = [] + + if let system = systemPrompt, !system.isEmpty { + messages.append(["role": "system", "content": system]) + } + + for msg in history where msg.role != .system { + messages.append(["role": msg.role.rawValue, "content": msg.content]) + } + + // Add the new user prompt unless it's already the trailing message. + if history.last?.role != .user || history.last?.content != prompt { + messages.append(["role": "user", "content": prompt]) + } + + return messages + } + + /// Build a POST `URLRequest` for a full chat-completions endpoint URL. + /// - Parameter endpoint: the complete URL (e.g. OpenRouter's + /// `https://openrouter.ai/api/v1/chat/completions`, or a local + /// `http://host/v1/chat/completions`). + static func build( + endpoint: String, + model: String, + messages: [[String: String]], + temperature: Float, + maxTokens: Int, + stream: Bool, + headers: [String: String] = [:] + ) throws -> URLRequest { + guard let url = URL(string: endpoint) else { + throw LLMError.invalidURL + } + + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + for (key, value) in headers { + request.setValue(value, forHTTPHeaderField: key) + } + + let body: [String: Any] = [ + "model": model, + "messages": messages, + "temperature": temperature, + "max_tokens": maxTokens, + "stream": stream + ] + request.httpBody = try JSONSerialization.data(withJSONObject: body) + return request + } +} + +// MARK: - Automatic failover selection + +/// Pure selection logic for health-checked automatic failover — the single-user +/// version of Nova's load balancing. Given an ordered preference chain and a map +/// of per-backend availability, pick which backend to use. No network here so it +/// is fully unit-testable; the manager injects live availability. +enum FailoverPlanner { + /// Default ordered preference chain: local-first, then frontier. + static let defaultChain: [LLMBackendType] = [.ollama, .mlx, .openRouter] + + /// The ordered subset of `chain` that is currently available. + static func orderedHealthy( + chain: [LLMBackendType], + availability: [LLMBackendType: Bool] + ) -> [LLMBackendType] { + chain.filter { availability[$0] == true } + } + + /// The first backend in `chain` that is available, or nil if none are. + static func firstHealthy( + chain: [LLMBackendType], + availability: [LLMBackendType: Bool] + ) -> LLMBackendType? { + chain.first { availability[$0] == true } + } +} diff --git a/MLX Code/ViewModels/ChatViewModel.swift b/MLX Code/ViewModels/ChatViewModel.swift index 7841947..9594f9b 100644 --- a/MLX Code/ViewModels/ChatViewModel.swift +++ b/MLX Code/ViewModels/ChatViewModel.swift @@ -358,80 +358,97 @@ class ChatViewModel: ObservableObject { budget: budget ) - // Get response from MLX service with streaming - let response = try await MLXService.shared.chatCompletion( - messages: optimizedMessages, - parameters: selectedModel?.parameters, - streamHandler: { [weak self] token in - Task { @MainActor [weak self] in - guard let self = self else { return } - - // Check if we should stop generation - guard !shouldStopGeneration else { - logWarning("⚠️ Stopping generation due to repetition/length limit", category: "ChatViewModel") - return - } + // Shared streaming token handler. Used by both the single-model MLX + // path and the multi-model load-balanced path so behavior is identical. + let streamHandler: (String) -> Void = { [weak self] token in + Task { @MainActor [weak self] in + guard let self = self else { return } + + // Check if we should stop generation + guard !shouldStopGeneration else { + logWarning("⚠️ Stopping generation due to repetition/length limit", category: "ChatViewModel") + return + } - // First token received - update status - if self.isWaitingForFirstToken { - self.isWaitingForFirstToken = false - self.statusMessage = "Generating..." - } + // First token received - update status + if self.isWaitingForFirstToken { + self.isWaitingForFirstToken = false + self.statusMessage = "Generating..." + } - accumulatedResponse += token + accumulatedResponse += token - // Update token count - self.tokenCount += 1 - self.currentTokenCount = self.tokenCount + // Update token count + self.tokenCount += 1 + self.currentTokenCount = self.tokenCount - // Calculate tokens per second - if let startTime = self.generationStartTime { - let elapsed = Date().timeIntervalSince(startTime) - if elapsed > 0 { - self.tokensPerSecond = Double(self.tokenCount) / elapsed - } + // Calculate tokens per second + if let startTime = self.generationStartTime { + let elapsed = Date().timeIntervalSince(startTime) + if elapsed > 0 { + self.tokensPerSecond = Double(self.tokenCount) / elapsed } + } - // Tool call detection is handled in MLXService — it breaks the stream - // loop as soon as appears, so we just update UI here. + // Tool call detection is handled in MLXService — it breaks the stream + // loop as soon as appears, so we just update UI here. - // Check for repetition - if let detector = self.repetitionDetector { - let hasRepetition = detector.addToken(token) - let hasExcessiveRepetition = detector.detectExcessiveRepetition() + // Check for repetition + if let detector = self.repetitionDetector { + let hasRepetition = detector.addToken(token) + let hasExcessiveRepetition = detector.detectExcessiveRepetition() - if hasRepetition || hasExcessiveRepetition { - shouldStopGeneration = true + if hasRepetition || hasExcessiveRepetition { + shouldStopGeneration = true - if accumulatedResponse.count > 500 { - let keepLength = Int(Double(accumulatedResponse.count) * 0.8) - let truncateIndex = accumulatedResponse.index(accumulatedResponse.startIndex, offsetBy: keepLength) - accumulatedResponse = String(accumulatedResponse[.. 500 { + let keepLength = Int(Double(accumulatedResponse.count) * 0.8) + let truncateIndex = accumulatedResponse.index(accumulatedResponse.startIndex, offsetBy: keepLength) + accumulatedResponse = String(accumulatedResponse[.. ChatViewModel.maxResponseLength { - shouldStopGeneration = true - accumulatedResponse += "\n\n[Response truncated: maximum length reached]" - } + // Check for maximum length + if accumulatedResponse.count > ChatViewModel.maxResponseLength { + shouldStopGeneration = true + accumulatedResponse += "\n\n[Response truncated: maximum length reached]" + } - // Check for maximum token count - if self.tokenCount > ChatViewModel.maxResponseTokens { - shouldStopGeneration = true - accumulatedResponse += "\n\n[Response truncated: maximum tokens reached]" - } + // Check for maximum token count + if self.tokenCount > ChatViewModel.maxResponseTokens { + shouldStopGeneration = true + accumulatedResponse += "\n\n[Response truncated: maximum tokens reached]" + } - // Update the message content - if let messageId = self.streamingMessageId, - let index = self.currentConversation?.messages.firstIndex(where: { $0.id == messageId }) { - self.currentConversation?.messages[index].content = accumulatedResponse - } + // Update the message content + if let messageId = self.streamingMessageId, + let index = self.currentConversation?.messages.firstIndex(where: { $0.id == messageId }) { + self.currentConversation?.messages[index].content = accumulatedResponse } } - ) + } + + // Get response. When any load-balancer toggle is on, spread work across + // the healthy enabled pool (all local models + frontier + Nova); the + // balancer returns nil when its pool is empty so we fall back cleanly to + // the single pinned MLX model — preserving existing behavior when off. + let response: String + if await LLMBalancer.shared.isBalancingEnabled(), + let balanced = try await LLMBalancer.shared.chatCompletionBalanced( + messages: optimizedMessages, + parameters: selectedModel?.parameters, + streamHandler: streamHandler + ) { + response = balanced + } else { + response = try await MLXService.shared.chatCompletion( + messages: optimizedMessages, + parameters: selectedModel?.parameters, + streamHandler: streamHandler + ) + } // Update final message if let messageId = streamingMessageId, diff --git a/MLX Code/Views/SettingsView.swift b/MLX Code/Views/SettingsView.swift index dd2215d..8b71567 100644 --- a/MLX Code/Views/SettingsView.swift +++ b/MLX Code/Views/SettingsView.swift @@ -31,6 +31,12 @@ struct SettingsView: View { /// Download status messages @State private var downloadStatus: [UUID: String] = [:] + /// OpenRouter API key (loaded from / saved to the Keychain). + @State private var openRouterKey: String = "" + + /// Whether the OpenRouter key was just saved (for UI feedback). + @State private var openRouterKeySaved = false + var body: some View { ZStack { @@ -91,6 +97,12 @@ struct SettingsView: View { Label("Model", systemImage: "cpu") } + // Load-balancer settings + balancerSettings + .tabItem { + Label("Balancer", systemImage: "arrow.triangle.branch") + } + // Appearance settings appearanceSettings .tabItem { @@ -124,6 +136,94 @@ struct SettingsView: View { .frame(width: 700, height: 600) } + // MARK: - Load Balancer Settings + + private var balancerSettings: some View { + ScrollView { + VStack(alignment: .leading, spacing: 20) { + // Intro + VStack(alignment: .leading, spacing: 6) { + Text("Multi-Model Load Balancing") + .font(.headline) + .foregroundColor(.primary) + Text("Spread chat across ALL installed local models, frontier models, and the optional Nova Gateway — health-gated and load-balanced — instead of a single pinned model. When every toggle is off, MLX Code uses only the selected model as before.") + .font(.caption) + .foregroundColor(.secondary) + } + + Divider() + + // Three toggles + VStack(alignment: .leading, spacing: 12) { + Toggle("All local models (native MLX + Ollama)", isOn: $settings.useAllLocalModels) + .toggleStyle(.switch) + Text("Discovers every installed SafeTensors MLX model plus any models served by a local Ollama instance.") + .font(.caption) + .foregroundColor(.secondary) + + Toggle("All frontier models (OpenRouter)", isOn: $settings.enableAllFrontierModels) + .toggleStyle(.switch) + Text("Requires an OpenRouter API key below. Frontier models are unavailable (skipped) until a key is set.") + .font(.caption) + .foregroundColor(.secondary) + + Toggle("Nova Gateway (optional)", isOn: $settings.useNovaGateway) + .toggleStyle(.switch) + Text("Optional. If the gateway health check fails, this entry shows as unavailable and everything else keeps working — Nova is never required.") + .font(.caption) + .foregroundColor(.secondary) + } + + Divider() + + // Nova Gateway URL + VStack(alignment: .leading, spacing: 8) { + Text("Nova Gateway URL") + .font(.subheadline) + .foregroundColor(.primary) + TextField("http://127.0.0.1:18792", text: $settings.novaGatewayURL) + .textFieldStyle(.roundedBorder) + .disableAutocorrection(true) + Text("OpenAI-compatible base URL. Health is probed via \(settings.novaGatewayURL)/v1/models.") + .font(.caption) + .foregroundColor(.secondary) + } + + Divider() + + // OpenRouter API key + VStack(alignment: .leading, spacing: 8) { + Text("OpenRouter API Key") + .font(.subheadline) + .foregroundColor(.primary) + HStack { + SecureField("sk-or-...", text: $openRouterKey) + .textFieldStyle(.roundedBorder) + Button("Save") { + let store = KeychainStore(service: OpenRouterProvider.keychainService) + if openRouterKey.isEmpty { + store.delete() + } else { + store.set(openRouterKey) + } + openRouterKeySaved = true + } + .buttonStyle(.borderedProminent) + } + Text(openRouterKeySaved + ? "Saved to the macOS Keychain." + : "Stored securely in the macOS Keychain — never in UserDefaults.") + .font(.caption) + .foregroundColor(openRouterKeySaved ? .green : .secondary) + } + } + .padding() + } + .onAppear { + openRouterKey = KeychainStore(service: OpenRouterProvider.keychainService).get() ?? "" + } + } + // MARK: - General Settings private var generalSettings: some View { diff --git a/README.md b/README.md index bf78284..8fa2f31 100644 --- a/README.md +++ b/README.md @@ -198,6 +198,57 @@ Models download automatically via native Hub Swift API. Custom models from any m --- +## Multi-model load balancing + +By default MLX Code talks to a single pinned MLX model. Optionally, it can spread each chat across **every** model available to the machine -- all installed local models (native MLX + a local Ollama, if present), all frontier models via an OpenRouter key, and an optional Nova Gateway -- health-gated and load-balanced. This is off by default; when every toggle is off, behavior is exactly as before. + +Three independent toggles live in **Settings → Balancer**: + +| Toggle | Adds to the pool | Requires | +|---|---|---| +| **All local models** | Every installed SafeTensors MLX model + any models served by a local Ollama | Nothing (Ollama optional) | +| **All frontier models** | Frontier models via OpenRouter | An OpenRouter API key (stored in the Keychain) | +| **Nova Gateway** | One entry that routes to Nova's OpenAI-compatible gateway | Nothing -- **entirely optional** | + +**Nova is never a hard requirement.** With zero Nova present the balancer still works over local MLX/Ollama models and/or OpenRouter. The Nova Gateway is just one optional entry: if its health probe (`GET /v1/models`) fails, that entry is marked unavailable and every other model keeps working. + +Selection is **health-gated then balanced**: each distinct backend is probed once (MLX = a model is loaded; Ollama = `/api/tags` 200; OpenRouter = a key is present; Nova = `/v1/models` 200), unhealthy models are dropped, and the `LoadBalancer` picks the next model by **least-busy** (fewest in-flight requests) or **round-robin**. A model that fails mid-request is removed and the next healthy one is tried; if the pool is empty it falls back cleanly to the single pinned MLX model. + +```mermaid +flowchart TB + subgraph Discovery["Discovery (ModelRegistry, network-free parsing)"] + MLXd["Local MLX models
MLXService.discoverModels()"] + Ollama["Ollama /api/tags"] + OR["OpenRouter /models"] + Nova["Nova Gateway entry"] + end + + subgraph Toggles["AppSettings toggles"] + T1["useAllLocalModels"] + T2["enableAllFrontierModels"] + T3["useNovaGateway"] + end + + MLXd --> Pool + Ollama --> Pool + OR --> Pool + Nova --> Pool + T1 -.gates.-> Pool + T2 -.gates.-> Pool + T3 -.gates.-> Pool + + Pool["assemblePool()
[DiscoveredModel] (deduped)"] --> Health["healthMap()
probe each backend once"] + Health --> Balancer["LoadBalancer.next()
least-busy / round-robin"] + Balancer --> Dispatch{"dispatchBalanced"} + Dispatch -->|mlx| InProc["MLXService.chatCompletion
(in-process, streamed)"] + Dispatch -->|ollama / openRouter / novaGateway| HTTP["OpenAI-compatible POST"] + Dispatch -->|all fail / empty pool| Fallback["Single pinned MLX model"] +``` + +Implemented by `ModelRegistry` / `LoadBalancer` / `OpenRouterProvider` / `KeychainStore` (the pure, network-free, unit-tested pieces) wired together by `LLMBalancer` and invoked from `ChatViewModel`. Covered by the network-free `LoadBalancerTests` suite. + +--- + ## Nova API Server Local HTTP API on port **37422** (loopback only).