1  //
  2  //  ChatViewModel.swift
  3  //  MLXChatExample
  4  //
  5  //  Created by İbrahim Çetin on 20.04.2025.
  6  //
  7  
  8  import Foundation
  9  import MLXLMCommon
 10  import UniformTypeIdentifiers
 11  import Observation
 12  
 13  /// ViewModel that manages the chat interface and coordinates with MLXService for text generation.
 14  /// Handles user input, message history, media attachments, and generation state.
 15  @Observable
 16  @MainActor
 17  class ChatViewModel {
 18      /// Service responsible for ML model operations
 19      private let mlxService: MLXService
 20  
 21      init(mlxService: MLXService) {
 22          self.mlxService = mlxService
 23      }
 24  
 25      /// Current user input text
 26      var prompt: String = ""
 27  
 28      var messages: [Message] = [
 29          .system("You are Jade, created by Awful Security.")
 30      ]
 31  
 32      /// Currently selected language model for generation
 33      var selectedModel: LMModel = MLXService.availableModels.first!
 34  
 35      /// Manages image and video attachments for the current message
 36      var mediaSelection = MediaSelection()
 37  
 38      /// Indicates if text generation is in progress
 39      var isGenerating = false
 40  
 41      /// Current generation task, used for cancellation
 42      private var generateTask: Task<Void, any Error>?
 43  
 44      /// Stores performance metrics from the current generation
 45      private var generateCompletionInfo: GenerateCompletionInfo?
 46  
 47      /// Current generation speed in tokens per second
 48      var tokensPerSecond: Double {
 49          generateCompletionInfo?.tokensPerSecond ?? 0
 50      }
 51  
 52      /// Progress of the current model download, if any
 53      var modelDownloadProgress: Progress? {
 54          mlxService.modelDownloadProgress
 55      }
 56  
 57      /// Most recent error message, if any
 58      var errorMessage: String?
 59      
 60      var scratchpadText: String = ""
 61  
 62      /// Generates response for the current prompt and media attachments
 63      func generate() async {
 64          // Cancel any existing generation task
 65          if let existingTask = generateTask {
 66              existingTask.cancel()
 67              generateTask = nil
 68          }
 69  
 70          isGenerating = true
 71  
 72          // Add user message with any media attachments
 73          messages.append(.user(prompt.trimmingCharacters(in: .whitespacesAndNewlines), images: mediaSelection.images, videos: mediaSelection.videos))
 74          // Add empty assistant message that will be filled during generation
 75          messages.append(.assistant(""))
 76  
 77          // Clear the input after sending
 78          clear(.prompt)
 79  
 80          generateTask = Task {
 81              // Process generation chunks and update UI
 82              for await generation in try await mlxService.generate(
 83                  messages: messages, model: selectedModel)
 84              {
 85                  switch generation {
 86                  case .chunk(let chunk):
 87                      // Append new text to the current assistant message
 88                      if let assistantMessage = messages.last {
 89                          assistantMessage.content += chunk
 90                      }
 91                  case .info(let info):
 92                      // Update performance metrics
 93                      generateCompletionInfo = info
 94                  }
 95              }
 96          }
 97  
 98          do {
 99              // Handle task completion and cancellation
100              try await withTaskCancellationHandler {
101                  try await generateTask?.value
102                  messages.last?.isComplete = true
103              } onCancel: {
104                  Task { @MainActor in
105                      generateTask?.cancel()
106  
107                      // Mark message as cancelled
108                      if let assistantMessage = messages.last {
109                          assistantMessage.content += "\n[Cancelled]"
110                      }
111                  }
112              }
113          } catch {
114              errorMessage = error.localizedDescription
115          }
116  
117          isGenerating = false
118          generateTask = nil
119      }
120  
121      /// Processes and adds media attachments to the current message
122      func addMedia(_ result: Result<URL, any Error>) {
123          do {
124              let url = try result.get()
125  
126              // Determine media type and add to appropriate collection
127              if let mediaType = UTType(filenameExtension: url.pathExtension) {
128                  if mediaType.conforms(to: .image) {
129                      mediaSelection.images = [url]
130                  } else if mediaType.conforms(to: .movie) {
131                      mediaSelection.videos = [url]
132                  }
133              }
134          } catch {
135              errorMessage = "Failed to load media item.\n\nError: \(error)"
136          }
137      }
138  
139      /// Clears various aspects of the chat state based on provided options
140      func clear(_ options: ClearOption) {
141          if options.contains(.prompt) {
142              prompt = ""
143              mediaSelection = .init()
144          }
145  
146          if options.contains(.chat) {
147              messages = []
148              generateTask?.cancel()
149          }
150  
151          if options.contains(.meta) {
152              generateCompletionInfo = nil
153          }
154  
155          errorMessage = nil
156      }
157  }
158  
159  /// Manages the state of media attachments in the chat
160  @Observable
161  class MediaSelection {
162      /// Controls visibility of media selection UI
163      var isShowing = false
164  
165      /// Currently selected image URLs
166      var images: [URL] = []
167  
168      /// Currently selected video URLs
169      var videos: [URL] = []
170  
171      /// Whether any media is currently selected
172      var isEmpty: Bool {
173          images.isEmpty && videos.isEmpty
174      }
175  }
176  
177  /// Options for clearing different aspects of the chat state
178  struct ClearOption: RawRepresentable, OptionSet {
179      let rawValue: Int
180  
181      /// Clears current prompt and media selection
182      static let prompt = ClearOption(rawValue: 1 << 0)
183      /// Clears chat history and cancels generation
184      static let chat = ClearOption(rawValue: 1 << 1)
185      /// Clears generation metadata
186      static let meta = ClearOption(rawValue: 1 << 2)
187  }