@sigit / siti
Code Commits Issues Pulls Cloud Agent CI/CD
release/windows-microsoft-store
development fix/asc-key-path main release/windows-microsoft-store
swift 472 lines 15.6 KB
Raw
1 //
2 // ChatViewModel.swift
3 // Siti AI visionOS
4 //
5
6 import Foundation
7 import Onde
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: ChatMessage {
40 switch role {
41 case .user:
42 return userMessage(content: text)
43 case .assistant:
44 return assistantMessage(content: text)
45 }
46 }
47 }
48
49 // ─────────────────────────────────────────────────────────────────────────────
50 // MARK: - Stream bridge
51 // ─────────────────────────────────────────────────────────────────────────────
52
53 /// Thread-safe bridge between the Rust callback thread and Swift concurrency.
54 private final class StreamBridge: @unchecked Sendable {
55 private let lock = NSLock()
56 private var continuation: AsyncStream<String>.Continuation?
57 private var cancelled = false
58
59 func attach(_ continuation: AsyncStream<String>.Continuation) {
60 lock.lock()
61 self.continuation = continuation
62 let shouldFinish = cancelled
63 lock.unlock()
64
65 if shouldFinish {
66 continuation.finish()
67 }
68 }
69
70 func yield(_ delta: String) {
71 lock.lock()
72 guard !cancelled else {
73 lock.unlock()
74 return
75 }
76 let continuation = continuation
77 lock.unlock()
78 continuation?.yield(delta)
79 }
80
81 func finish() {
82 lock.lock()
83 let continuation = continuation
84 self.continuation = nil
85 lock.unlock()
86 continuation?.finish()
87 }
88
89 func cancel() {
90 lock.lock()
91 cancelled = true
92 let continuation = continuation
93 self.continuation = nil
94 lock.unlock()
95 continuation?.finish()
96 }
97
98 var isCancelled: Bool {
99 lock.lock()
100 let cancelled = cancelled
101 lock.unlock()
102 return cancelled
103 }
104 }
105
106 /// Bridges the UniFFI callback interface to an `AsyncStream<String>`.
107 private final class ChunkCollector: StreamChunkListener {
108 private let bridge: StreamBridge
109
110 init(bridge: StreamBridge) {
111 self.bridge = bridge
112 }
113
114 /// Called by the Rust runtime for each streamed token.
115 /// - Returns: `true` to keep streaming, `false` to cancel.
116 func onChunk(chunk: StreamChunk) -> Bool {
117 if bridge.isCancelled {
118 return false
119 }
120
121 if chunk.done {
122 bridge.finish()
123 return false
124 }
125
126 bridge.yield(chunk.delta)
127 return !bridge.isCancelled
128 }
129 }
130
131 // ─────────────────────────────────────────────────────────────────────────────
132 // MARK: - ViewModel
133 // ─────────────────────────────────────────────────────────────────────────────
134
135 @MainActor
136 final class ChatViewModel: ObservableObject {
137
138 // ── Published state ───────────────────────────────────────────────────────
139
140 @Published private(set) var messages: [Message]
141 @Published private(set) var isModelLoading: Bool = false
142 @Published private(set) var isModelReady: Bool = false
143 @Published private(set) var isSending: Bool = false
144 @Published private(set) var loadingProgress: String = "Preparing model…"
145 @Published private(set) var engineInfo: EngineInfo = EngineInfo(
146 status: .unloaded,
147 modelName: nil,
148 approxMemory: nil,
149 historyLength: 0
150 )
151
152 @Published var alertError: IdentifiableError? = nil
153
154 // ── Private state ─────────────────────────────────────────────────────────
155
156 // Lazy so the Rust/UniFFI runtime isn't initialized during
157 // @StateObject construction on the UIKit event-fetch thread.
158 private var engine: OndeChatEngine?
159 private var streamingTask: Task<Void, Never>?
160 private var streamBridge: StreamBridge?
161 private var didHydrateEngineHistory = false
162
163 private static let persistedMessagesKey = "ai.siti.vision.messages"
164
165 init() {
166 self.messages = Self.loadPersistedMessages()
167 }
168
169 private func getOrCreateEngine() -> OndeChatEngine {
170 if let existing = engine { return existing }
171 let new = OndeChatEngine()
172 engine = new
173 return new
174 }
175
176 var hasMessages: Bool {
177 !messages.isEmpty
178 }
179
180 var isBusy: Bool {
181 isModelLoading || isSending
182 }
183
184 // ─────────────────────────────────────────────────────────────────────────
185 // MARK: Model lifecycle
186 // ─────────────────────────────────────────────────────────────────────────
187
188 func loadModelIfNeeded(systemPrompt: String, samplingPreset: SamplingPreset) async {
189 if isModelReady {
190 await applySettings(systemPrompt: systemPrompt, samplingPreset: samplingPreset)
191 return
192 }
193
194 await loadModel(
195 systemPrompt: systemPrompt,
196 samplingPreset: samplingPreset,
197 forceReload: false
198 )
199 }
200
201 func retryLoadingModel(systemPrompt: String, samplingPreset: SamplingPreset) async {
202 await loadModel(
203 systemPrompt: systemPrompt,
204 samplingPreset: samplingPreset,
205 forceReload: false
206 )
207 }
208
209 func reloadModel(systemPrompt: String, samplingPreset: SamplingPreset) async {
210 await loadModel(
211 systemPrompt: systemPrompt,
212 samplingPreset: samplingPreset,
213 forceReload: true
214 )
215 }
216
217 func applySettings(systemPrompt: String, samplingPreset: SamplingPreset) async {
218 guard isModelReady else { return }
219
220 let engine = getOrCreateEngine()
221 let normalizedPrompt = normalizedSystemPrompt(systemPrompt)
222
223 if let normalizedPrompt {
224 await engine.setSystemPrompt(prompt: normalizedPrompt)
225 } else {
226 await engine.clearSystemPrompt()
227 }
228
229 await engine.setSampling(sampling: samplingPreset.samplingConfig)
230 await refreshEngineInfo()
231 }
232
233 func clearConversation() async {
234 cancelStreaming()
235 messages.removeAll()
236 persistMessages()
237
238 if let engine {
239 _ = await engine.clearHistory()
240 await refreshEngineInfo()
241 } else {
242 engineInfo = EngineInfo(
243 status: .unloaded,
244 modelName: nil,
245 approxMemory: nil,
246 historyLength: 0
247 )
248 }
249 }
250
251 // ─────────────────────────────────────────────────────────────────────────
252 // MARK: Sending messages
253 // ─────────────────────────────────────────────────────────────────────────
254
255 func send(text: String) {
256 let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
257 guard !trimmed.isEmpty, !isSending, isModelReady else { return }
258
259 alertError = nil
260 messages.append(Message(role: .user, text: trimmed))
261 persistMessages()
262
263 streamingTask = Task {
264 await streamResponse(for: trimmed)
265 }
266 }
267
268 func sendSuggestion(_ suggestion: PromptSuggestion) {
269 send(text: suggestion.prompt)
270 }
271
272 func cancelStreaming() {
273 streamBridge?.cancel()
274 streamingTask?.cancel()
275 streamingTask = nil
276 streamBridge = nil
277
278 if let idx = messages.indices.last, messages[idx].isStreaming {
279 if messages[idx].text.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty {
280 messages.remove(at: idx)
281 } else {
282 messages[idx].isStreaming = false
283 }
284 persistMessages()
285 }
286
287 isSending = false
288 if engineInfo.status == .generating {
289 engineInfo = EngineInfo(
290 status: .ready,
291 modelName: engineInfo.modelName,
292 approxMemory: engineInfo.approxMemory,
293 historyLength: engineInfo.historyLength
294 )
295 }
296 }
297
298 // ─────────────────────────────────────────────────────────────────────────
299 // MARK: Private helpers
300 // ─────────────────────────────────────────────────────────────────────────
301
302 private func loadModel(
303 systemPrompt: String,
304 samplingPreset: SamplingPreset,
305 forceReload: Bool
306 ) async {
307 guard !isModelLoading else { return }
308
309 if forceReload {
310 cancelStreaming()
311 didHydrateEngineHistory = false
312 }
313
314 isModelLoading = true
315 isModelReady = false
316 alertError = nil
317 loadingProgress = forceReload ? "Reloading model…" : "Loading model…"
318
319 defer {
320 isModelLoading = false
321 }
322
323 do {
324 let elapsed = try await getOrCreateEngine().loadDefaultModel(
325 systemPrompt: normalizedSystemPrompt(systemPrompt),
326 sampling: samplingPreset.samplingConfig
327 )
328
329 loadingProgress = String(format: "Model ready in %.1f s", elapsed)
330 isModelReady = true
331 didHydrateEngineHistory = false
332 await restorePersistedConversationIfNeeded()
333 await refreshEngineInfo()
334 } catch {
335 isModelReady = false
336 await refreshEngineInfo()
337 alertError = IdentifiableError(error)
338 loadingProgress = "Unable to load model"
339 }
340 }
341
342 private func streamResponse(for userText: String) async {
343 isSending = true
344 engineInfo = EngineInfo(
345 status: .generating,
346 modelName: engineInfo.modelName,
347 approxMemory: engineInfo.approxMemory,
348 historyLength: engineInfo.historyLength
349 )
350
351 let assistantIndex = messages.count
352 messages.append(Message(role: .assistant, text: "", isStreaming: true))
353
354 let bridge = StreamBridge()
355 streamBridge = bridge
356 let listener = ChunkCollector(bridge: bridge)
357
358 let stream = AsyncStream<String> { continuation in
359 bridge.attach(continuation)
360 }
361
362 let capturedEngine = getOrCreateEngine()
363 let producer = Task.detached(priority: .userInitiated) {
364 do {
365 try await streamChatMessage(
366 engine: capturedEngine,
367 message: userText,
368 listener: listener
369 )
370 bridge.finish()
371 } catch {
372 bridge.finish()
373 await MainActor.run {
374 if !bridge.isCancelled {
375 self.alertError = IdentifiableError(error)
376 }
377 }
378 }
379 }
380
381 for await delta in stream {
382 guard !Task.isCancelled else { break }
383 if assistantIndex < messages.count {
384 messages[assistantIndex].text += delta
385 }
386 }
387
388 _ = await producer.result
389
390 let assistantHasText = assistantIndex < messages.count &&
391 !messages[assistantIndex].text.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
392
393 if assistantIndex < messages.count {
394 if assistantHasText {
395 messages[assistantIndex].isStreaming = false
396 } else {
397 messages.remove(at: assistantIndex)
398 }
399 }
400
401 persistMessages()
402 streamBridge = nil
403 streamingTask = nil
404 isSending = false
405 await refreshEngineInfo()
406 }
407
408 private func restorePersistedConversationIfNeeded() async {
409 guard !didHydrateEngineHistory else { return }
410 guard !messages.isEmpty else {
411 didHydrateEngineHistory = true
412 return
413 }
414
415 loadingProgress = "Restoring conversation…"
416 let engine = getOrCreateEngine()
417
418 for message in messages {
419 await engine.pushHistory(message: message.asChatMessage)
420 }
421
422 didHydrateEngineHistory = true
423 }
424
425 private func refreshEngineInfo() async {
426 guard let engine else {
427 engineInfo = EngineInfo(
428 status: .unloaded,
429 modelName: nil,
430 approxMemory: nil,
431 historyLength: 0
432 )
433 return
434 }
435
436 let info = await engine.info()
437 engineInfo = info
438 isModelReady = info.status == .ready || info.status == .generating
439 }
440
441 private func normalizedSystemPrompt(_ prompt: String) -> String? {
442 let trimmed = prompt.trimmingCharacters(in: .whitespacesAndNewlines)
443 return trimmed.isEmpty ? nil : trimmed
444 }
445
446 private func persistMessages() {
447 do {
448 let data = try JSONEncoder().encode(messages)
449 UserDefaults.standard.set(data, forKey: Self.persistedMessagesKey)
450 } catch {
451 logPersistenceFailure(error)
452 }
453 }
454
455 private static func loadPersistedMessages() -> [Message] {
456 guard let data = UserDefaults.standard.data(forKey: persistedMessagesKey) else {
457 return []
458 }
459
460 do {
461 return try JSONDecoder().decode([Message].self, from: data).map {
462 Message(id: $0.id, role: $0.role, text: $0.text, isStreaming: false)
463 }
464 } catch {
465 return []
466 }
467 }
468
469 private func logPersistenceFailure(_ error: Error) {
470 NSLog("SitiVision: failed to persist messages: %@", String(describing: error))
471 }
472 }