@sigit / siti
Code Commits Issues Pulls Cloud Agent CI/CD
development
development fix/asc-key-path main release/windows-microsoft-store
swift 350 lines 12.1 KB
Raw
1 //
2 // ChatViewModel.swift
3 // Siti AI visionOS
4 //
5
6 import Ed
7 import Foundation
8
9 // ─────────────────────────────────────────────────────────────────────────────
10 // MARK: - Message model
11 // ─────────────────────────────────────────────────────────────────────────────
12
13 struct Message: Identifiable, Codable, Hashable {
14 enum Role: String, Codable {
15 case user
16 case assistant
17 }
18
19 let id: UUID
20 let role: Role
21 var text: String
22 /// True while the assistant is still streaming tokens for this message.
23 var isStreaming: Bool
24
25 init(
26 id: UUID = UUID(),
27 role: Role,
28 text: String,
29 isStreaming: Bool = false
30 ) {
31 self.id = id
32 self.role = role
33 self.text = text
34 self.isStreaming = isStreaming
35 }
36 }
37
38 private extension Message {
39 var asChatMessage: EdChatMessage {
40 switch role {
41 case .user:
42 return EdChatMessage(role: .user, content: text)
43 case .assistant:
44 return EdChatMessage(role: .assistant, content: text)
45 }
46 }
47 }
48
49 private extension EdEngineInfo {
50 static let unloaded = EdEngineInfo(status: .unloaded)
51 }
52
53 // ─────────────────────────────────────────────────────────────────────────────
54 // MARK: - ViewModel
55 // ─────────────────────────────────────────────────────────────────────────────
56
57 @MainActor
58 final class ChatViewModel: ObservableObject {
59
60 // ── Published state ───────────────────────────────────────────────────────
61
62 @Published private(set) var messages: [Message]
63 @Published private(set) var isModelLoading: Bool = false
64 @Published private(set) var isModelReady: Bool = false
65 @Published private(set) var isSending: Bool = false
66 @Published private(set) var loadingProgress: String = "Preparing model…"
67 @Published private(set) var engineInfo: EdEngineInfo = .unloaded
68
69 @Published var alertError: IdentifiableError? = nil
70
71 // ── Private state ─────────────────────────────────────────────────────────
72
73 // Lazy so the Rust/UniFFI runtime isn't initialized during
74 // @StateObject construction on the UIKit event-fetch thread.
75 private var engine: EdAgent?
76 private var streamingTask: Task<Void, Never>?
77 private var didHydrateEngineHistory = false
78
79 private static let persistedMessagesKey = "ai.siti.vision.messages"
80
81 init() {
82 self.messages = Self.loadPersistedMessages()
83 }
84
85 private func getOrCreateEngine() -> EdAgent {
86 if let existing = engine { return existing }
87 let new = EdAgent(appID: nil)
88 engine = new
89 return new
90 }
91
92 var hasMessages: Bool {
93 !messages.isEmpty
94 }
95
96 var isBusy: Bool {
97 isModelLoading || isSending
98 }
99
100 // ─────────────────────────────────────────────────────────────────────────
101 // MARK: Model lifecycle
102 // ─────────────────────────────────────────────────────────────────────────
103
104 func loadModelIfNeeded(systemPrompt: String, samplingPreset: SamplingPreset) async {
105 if isModelReady {
106 await applySettings(systemPrompt: systemPrompt, samplingPreset: samplingPreset)
107 return
108 }
109
110 await loadModel(
111 systemPrompt: systemPrompt,
112 samplingPreset: samplingPreset,
113 forceReload: false
114 )
115 }
116
117 func retryLoadingModel(systemPrompt: String, samplingPreset: SamplingPreset) async {
118 await loadModel(
119 systemPrompt: systemPrompt,
120 samplingPreset: samplingPreset,
121 forceReload: false
122 )
123 }
124
125 func reloadModel(systemPrompt: String, samplingPreset: SamplingPreset) async {
126 await loadModel(
127 systemPrompt: systemPrompt,
128 samplingPreset: samplingPreset,
129 forceReload: true
130 )
131 }
132
133 func applySettings(systemPrompt: String, samplingPreset: SamplingPreset) async {
134 guard isModelReady else { return }
135
136 let engine = getOrCreateEngine()
137 let normalizedPrompt = normalizedSystemPrompt(systemPrompt)
138
139 if let normalizedPrompt {
140 await engine.setSystemPrompt(normalizedPrompt)
141 } else {
142 await engine.clearSystemPrompt()
143 }
144
145 await engine.setSampling(samplingPreset.samplingConfig)
146 await refreshEngineInfo()
147 }
148
149 func clearConversation() async {
150 cancelStreaming()
151 messages.removeAll()
152 persistMessages()
153
154 if let engine {
155 _ = await engine.clearHistory()
156 await refreshEngineInfo()
157 } else {
158 engineInfo = .unloaded
159 }
160 }
161
162 // ─────────────────────────────────────────────────────────────────────────
163 // MARK: Sending messages
164 // ─────────────────────────────────────────────────────────────────────────
165
166 func send(text: String) {
167 let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
168 guard !trimmed.isEmpty, !isSending, isModelReady else { return }
169
170 alertError = nil
171 messages.append(Message(role: .user, text: trimmed))
172 persistMessages()
173
174 streamingTask = Task {
175 await streamResponse(for: trimmed)
176 }
177 }
178
179 func sendSuggestion(_ suggestion: PromptSuggestion) {
180 send(text: suggestion.prompt)
181 }
182
183 func cancelStreaming() {
184 // Cancelling the task ends the `for try await` in `streamResponse`,
185 // which stops generation on the Rust side.
186 streamingTask?.cancel()
187 streamingTask = nil
188
189 if let idx = messages.indices.last, messages[idx].isStreaming {
190 if messages[idx].text.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty {
191 messages.remove(at: idx)
192 } else {
193 messages[idx].isStreaming = false
194 }
195 persistMessages()
196 }
197
198 isSending = false
199 if engineInfo.status == .generating {
200 engineInfo = EdEngineInfo(
201 status: .ready,
202 modelName: engineInfo.modelName,
203 approximateMemory: engineInfo.approximateMemory,
204 historyLength: engineInfo.historyLength
205 )
206 }
207 }
208
209 // ─────────────────────────────────────────────────────────────────────────
210 // MARK: Private helpers
211 // ─────────────────────────────────────────────────────────────────────────
212
213 private func loadModel(
214 systemPrompt: String,
215 samplingPreset: SamplingPreset,
216 forceReload: Bool
217 ) async {
218 guard !isModelLoading else { return }
219
220 if forceReload {
221 cancelStreaming()
222 didHydrateEngineHistory = false
223 }
224
225 isModelLoading = true
226 isModelReady = false
227 alertError = nil
228 loadingProgress = forceReload ? "Reloading model…" : "Loading model…"
229
230 defer {
231 isModelLoading = false
232 }
233
234 do {
235 let elapsed = try await getOrCreateEngine().load(
236 gguf: SitiDefaults.defaultModel,
237 systemPrompt: normalizedSystemPrompt(systemPrompt),
238 sampling: samplingPreset.samplingConfig
239 )
240
241 loadingProgress = String(format: "Model ready in %.1f s", elapsed)
242 isModelReady = true
243 didHydrateEngineHistory = false
244 await restorePersistedConversationIfNeeded()
245 await refreshEngineInfo()
246 } catch {
247 isModelReady = false
248 await refreshEngineInfo()
249 alertError = IdentifiableError(error)
250 loadingProgress = "Unable to load model"
251 }
252 }
253
254 private func streamResponse(for userText: String) async {
255 isSending = true
256 engineInfo = EdEngineInfo(
257 status: .generating,
258 modelName: engineInfo.modelName,
259 approximateMemory: engineInfo.approximateMemory,
260 historyLength: engineInfo.historyLength
261 )
262
263 let assistantIndex = messages.count
264 messages.append(Message(role: .assistant, text: "", isStreaming: true))
265
266 do {
267 for try await delta in getOrCreateEngine().stream(userText) {
268 if assistantIndex < messages.count {
269 messages[assistantIndex].text += delta
270 }
271 }
272 } catch {
273 if !Task.isCancelled {
274 alertError = IdentifiableError(error)
275 }
276 }
277
278 let assistantHasText = assistantIndex < messages.count &&
279 !messages[assistantIndex].text.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
280
281 if assistantIndex < messages.count {
282 if assistantHasText {
283 messages[assistantIndex].isStreaming = false
284 } else {
285 messages.remove(at: assistantIndex)
286 }
287 }
288
289 persistMessages()
290 streamingTask = nil
291 isSending = false
292 await refreshEngineInfo()
293 }
294
295 private func restorePersistedConversationIfNeeded() async {
296 guard !didHydrateEngineHistory else { return }
297 guard !messages.isEmpty else {
298 didHydrateEngineHistory = true
299 return
300 }
301
302 loadingProgress = "Restoring conversation…"
303 await getOrCreateEngine().restoreHistory(messages.map(\.asChatMessage))
304
305 didHydrateEngineHistory = true
306 }
307
308 private func refreshEngineInfo() async {
309 guard let engine else {
310 engineInfo = .unloaded
311 return
312 }
313
314 let info = await engine.info()
315 engineInfo = info
316 isModelReady = info.status == .ready || info.status == .generating
317 }
318
319 private func normalizedSystemPrompt(_ prompt: String) -> String? {
320 let trimmed = prompt.trimmingCharacters(in: .whitespacesAndNewlines)
321 return trimmed.isEmpty ? nil : trimmed
322 }
323
324 private func persistMessages() {
325 do {
326 let data = try JSONEncoder().encode(messages)
327 UserDefaults.standard.set(data, forKey: Self.persistedMessagesKey)
328 } catch {
329 logPersistenceFailure(error)
330 }
331 }
332
333 private static func loadPersistedMessages() -> [Message] {
334 guard let data = UserDefaults.standard.data(forKey: persistedMessagesKey) else {
335 return []
336 }
337
338 do {
339 return try JSONDecoder().decode([Message].self, from: data).map {
340 Message(id: $0.id, role: $0.role, text: $0.text, isStreaming: false)
341 }
342 } catch {
343 return []
344 }
345 }
346
347 private func logPersistenceFailure(_ error: Error) {
348 NSLog("SitiVision: failed to persist messages: %@", String(describing: error))
349 }
350 }