extend model feature flags

Author: Mickey Knox

app/modules/foundations/LLMFoundation/Sources/AIModel.swift

@@ -5,7 +5,9 @@ import Foundation
 
 // MARK: - LLMReasoning
 
-public struct LLMReasoning: Sendable, Hashable, Codable { }
+public struct LLMReasoning: Sendable, Hashable, Codable {
+  public init() { }
+}
 
 // MARK: - AIProviderModel
 
@@ -42,7 +44,10 @@ public struct AIModel: Hashable, Identifiable, Sendable, Codable {
     documentationURL: URL? = nil,
     reasoning: LLMReasoning? = nil,
     createdAt: TimeInterval,
-    rankForProgramming: Int)
+    rankForProgramming: Int,
+    supportsChat: Bool = true,
+    supportsTools: Bool = true,
+    supportsCompletion: Bool = false)
   {
     self.slug = slug
     self.name = name
@@ -54,6 +59,13 @@ public struct AIModel: Hashable, Identifiable, Sendable, Codable {
     self.reasoning = reasoning
     self.createdAt = createdAt
     self.rankForProgramming = rankForProgramming
+    self.supportsChat = supportsChat
+      // TODO: move to appropriate place
+      && AIModel.modelSupportsChat(id: slug)
+    self.supportsTools = supportsTools
+    self.supportsCompletion = supportsCompletion
+      // TODO: move to appropriate place
+      || AIModel.modelSupportsCompletion(id: slug)
   }
 
   /// A few models for debugging and providing default values.
@@ -122,6 +134,9 @@ public struct AIModel: Hashable, Identifiable, Sendable, Codable {
   public let reasoning: LLMReasoning?
   public let createdAt: TimeInterval
   public let rankForProgramming: Int
+  public let supportsChat: Bool
+  public let supportsTools: Bool
+  public let supportsCompletion: Bool
 
   public var id: String {
     slug
@@ -147,11 +162,11 @@ public struct AIModel: Hashable, Identifiable, Sendable, Codable {
 }
 
 extension AIModel {
-  public var supportsCompletion: Bool {
+  static func modelSupportsCompletion(id: String) -> Bool {
     ["inception/mercury-coder-small", "mistralai/codestral-latest"].contains(id)
   }
 
-  public var supportsChat: Bool {
+  static func modelSupportsChat(id: String) -> Bool {
     id != "inception/mercury-coder-small"
   }
 }

app/modules/serviceInterfaces/LocalServerServiceInterface/Sources/listModelsSchema.generated.swift

@@ -64,6 +64,10 @@ extension Schema {
     public let pricing: ModelPricing?
     public let createdAt: Double
     public let rankForProgramming: Int
+    public let supportsChat: Bool
+    public let supportsTools: Bool
+    public let supportsReasoning: Bool
+    public let supportsCompletion: Bool
 
     private enum CodingKeys: String, CodingKey {
       case providerId = "providerId"
@@ -77,6 +81,10 @@ extension Schema {
       case pricing = "pricing"
       case createdAt = "createdAt"
       case rankForProgramming = "rankForProgramming"
+      case supportsChat = "supportsChat"
+      case supportsTools = "supportsTools"
+      case supportsReasoning = "supportsReasoning"
+      case supportsCompletion = "supportsCompletion"
     }
 
     public init(
@@ -90,7 +98,11 @@ extension Schema {
         outputModalities: [ModelModality],
         pricing: ModelPricing? = nil,
         createdAt: Double,
-        rankForProgramming: Int
+        rankForProgramming: Int,
+        supportsChat: Bool,
+        supportsTools: Bool,
+        supportsReasoning: Bool,
+        supportsCompletion: Bool
     ) {
       self.providerId = providerId
       self.globalId = globalId
@@ -103,6 +115,10 @@ extension Schema {
       self.pricing = pricing
       self.createdAt = createdAt
       self.rankForProgramming = rankForProgramming
+      self.supportsChat = supportsChat
+      self.supportsTools = supportsTools
+      self.supportsReasoning = supportsReasoning
+      self.supportsCompletion = supportsCompletion
     }
 
     public init(from decoder: Decoder) throws {
@@ -118,6 +134,10 @@ extension Schema {
       pricing = try container.decodeIfPresent(ModelPricing?.self, forKey: .pricing)
       createdAt = try container.decode(Double.self, forKey: .createdAt)
       rankForProgramming = try container.decode(Int.self, forKey: .rankForProgramming)
+      supportsChat = try container.decode(Bool.self, forKey: .supportsChat)
+      supportsTools = try container.decode(Bool.self, forKey: .supportsTools)
+      supportsReasoning = try container.decode(Bool.self, forKey: .supportsReasoning)
+      supportsCompletion = try container.decode(Bool.self, forKey: .supportsCompletion)
     }
 
     public func encode(to encoder: Encoder) throws {
@@ -133,6 +153,10 @@ extension Schema {
       try container.encodeIfPresent(pricing, forKey: .pricing)
       try container.encode(createdAt, forKey: .createdAt)
       try container.encode(rankForProgramming, forKey: .rankForProgramming)
+      try container.encode(supportsChat, forKey: .supportsChat)
+      try container.encode(supportsTools, forKey: .supportsTools)
+      try container.encode(supportsReasoning, forKey: .supportsReasoning)
+      try container.encode(supportsCompletion, forKey: .supportsCompletion)
     }
   }
   public enum ModelModality: String, Codable, Sendable, CaseIterable {

app/modules/services/LLMService/Sources/LLMModelManager.swift

@@ -329,23 +329,31 @@ final class AIModelsManager: AIModelsManagerProtocol {
       name: apiProvider.name,
       settings: apiProvider.settings)))
     let response: Schema.ListModelsOutput = try await localServer.postRequest(path: "models", data: data)
-    return response.models.map { AIProviderModel(
-      providerId: $0.providerId,
-      provider: provider,
-      modelInfo: .init(
-        name: $0.name,
-        slug: $0.globalId,
-        contextSize: $0.contextLength,
-        maxOutputTokens: $0.maxCompletionTokens,
-        defaultPricing: $0.pricing != nil
-          ? .init(
-            input: $0.pricing?.prompt ?? 0,
-            output: $0.pricing?.completion ?? 0,
-            cacheWrite: $0.pricing?.inputCacheWrite ?? 0,
-            cachedInput: $0.pricing?.inputCacheRead ?? 0)
-          : nil,
-        createdAt: $0.createdAt,
-        rankForProgramming: $0.rankForProgramming)) }
+    return response.models.map { model in
+      let defaultPricing: ModelPricing? = {
+        guard let pricing = model.pricing else { return nil }
+        return .init(
+          input: pricing.prompt,
+          output: pricing.completion,
+          cacheWrite: pricing.inputCacheWrite ?? 0,
+          cachedInput: pricing.inputCacheRead ?? 0)
+      }()
+      return AIProviderModel(
+        providerId: model.providerId,
+        provider: provider,
+        modelInfo: .init(
+          name: model.name,
+          slug: model.globalId,
+          contextSize: model.contextLength,
+          maxOutputTokens: model.maxCompletionTokens,
+          defaultPricing: defaultPricing,
+          reasoning: model.supportsReasoning ? LLMReasoning() : nil,
+          createdAt: model.createdAt,
+          rankForProgramming: model.rankForProgramming,
+          supportsChat: model.supportsChat,
+          supportsTools: model.supportsTools,
+          supportsCompletion: model.supportsCompletion))
+    }
   }
 
   private func observerChangesToSettings() {

app/modules/services/LLMService/Tests/LLMModelManagerTests.swift

@@ -1094,6 +1094,10 @@ private func makeSchemaModel(
     maxCompletionTokens: 8_192,
     inputModalities: [.text],
     outputModalities: [.text],
+    supportsChat: true,
+    supportsTools: true,
+    supportsReasoning: false,
+    supportsCompletion: true,
     pricing: pricing ?? Schema.ModelPricing(
       prompt: 1.0,
       completion: 2.0),

local-server/build.sha256

@@ -1 +1 @@
-48ce67299157389bcd10431b9b1905d0fc47731f189f5993d547ec4706682062
+e73e92e58623897e84bfc944cf385d6089ea727d20db1a5ad64c5af486c6653b

local-server/src/server/endpoints/listModels.ts

@@ -71,6 +71,10 @@ export const registerEndpoint = (router: Router, aiProviders: AIProvider[]) => {
 					: undefined,
 				createdAt: model.created,
 				rankForProgramming: model.rankForProgramming,
+				supportsChat: model.supportsChat,
+				supportsTools: model.supportsTools,
+				supportsReasoning: model.supportsReasoning,
+				supportsCompletion: model.supportsCompletion,
 			})),
 		} satisfies ListModelsOutput)
 	})