spb/zyquo-local Public MIT

1//2//  ChatController.swift3//  Zyquo Local4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//89import Foundation10import Observation1112/// Drives one streaming generation at a time for the selected conversation:13/// send/stop/regenerate/edit-resend, live `<think>` parsing, throttled14/// tokens-per-second ticker (≤4 Hz), stats capture, and auto-titling.15@MainActor16@Observable17final class ChatController {18    private(set) var isGenerating = false19    /// Visible (non-thinking) streamed text of the in-flight response.20    private(set) var streamingText = ""21    /// Streamed `<think>` content of the in-flight response.22    private(set) var streamingThinking = ""23    /// True while the stream is inside a `<think>` block.24    private(set) var isThinking = false25    /// Live generation speed, updated at most 4 Hz.26    private(set) var liveTokensPerSecond: Double = 027    /// Conversation currently streaming (may differ from the selection).28    private(set) var streamingConversationID: UUID?2930    private weak var app: AppModel?31    private var generationTask: Task<Void, Never>?3233    func bind(to app: AppModel) {34        self.app = app35    }3637    // MARK: - Actions3839    /// Sends `prompt` in the given conversation, streaming the response.40    func send(prompt: String, in conversation: Conversation) {41        guard let app, !isGenerating else { return }42        var conversation = conversation43        let userMessage = Message(role: .user, content: prompt)44        conversation.messages.append(userMessage)45        app.update(conversation)46        stream(prompt: prompt, conversation: conversation)47    }4849    /// Regenerates the last assistant response (optionally after the user50    /// switched models).51    func regenerate(in conversation: Conversation) {52        guard !isGenerating else { return }53        var conversation = conversation54        guard let lastUser = conversation.messages.last(where: { $0.role == .user }) else { return }55        // Drop trailing assistant message(s) after the last user turn.56        while let last = conversation.messages.last, last.role == .assistant {57            conversation.messages.removeLast()58        }59        app?.update(conversation)60        stream(prompt: lastUser.content, conversation: conversation, replayingLastUser: true)61    }6263    /// Edits a previous user message and resends from that point.64    func editAndResend(messageID: UUID, newText: String, in conversation: Conversation) {65        guard !isGenerating else { return }66        var conversation = conversation67        guard let index = conversation.messages.firstIndex(where: { $0.id == messageID }) else { return }68        conversation.messages[index].content = newText69        conversation.messages.removeSubrange((index + 1)...)70        app?.update(conversation)71        stream(prompt: newText, conversation: conversation, replayingLastUser: true)72    }7374    func stop() {75        generationTask?.cancel()76        Task { await app?.engine.stopGeneration() }77    }7879    // MARK: - Streaming core8081    /// `replayingLastUser`: the prompt is already the last user message in82    /// `conversation.messages`; the engine session must be rebuilt so its83    /// history excludes it (it is re-sent as the new turn).84    private func stream(prompt: String, conversation: Conversation, replayingLastUser: Bool = false) {85        guard let app else { return }86        let conversationID = conversation.id8788        streamingText = ""89        streamingThinking = ""90        isThinking = false91        liveTokensPerSecond = 092        isGenerating = true93        streamingConversationID = conversationID9495        generationTask = Task {96            var parser = ThinkTagParser()97            var stats: GenerationStats?98            var finish: GenerationFinishReason = .stop99            let started = Date()100            var tokenCount = 0101            var lastTick = Date.distantPast102103            do {104                // The engine session's history must exclude the new prompt:105                // strip the trailing user message before (re)building.106                var sessionConversation = conversation107                if let last = sessionConversation.messages.last, last.role == .user {108                    sessionConversation.messages.removeLast()109                }110                if replayingLastUser {111                    try await app.engine.startSession(conversation: sessionConversation)112                } else {113                    try await app.engine.ensureSession(conversation: sessionConversation)114                }115116                let events = try await app.engine.generate(prompt: prompt, params: conversation.params)117                for try await event in events {118                    switch event {119                    case .token(let text):120                        tokenCount += 1121                        let (visible, thinking, inThink) = parser.consume(text)122                        if !visible.isEmpty { streamingText += visible }123                        if !thinking.isEmpty { streamingThinking += thinking }124                        isThinking = inThink125                        // ≤4 Hz ticker to avoid flicker.126                        let now = Date()127                        if now.timeIntervalSince(lastTick) >= 0.25 {128                            lastTick = now129                            let elapsed = now.timeIntervalSince(started)130                            if elapsed > 0.5 {131                                liveTokensPerSecond = Double(tokenCount) / elapsed132                            }133                        }134                    case .stats(let s):135                        stats = s136                    case .finished(let reason):137                        finish = reason138                    }139                }140            } catch {141                app.lastError = error.localizedDescription142            }143            finalize(conversationID: conversationID, stats: stats, finish: finish)144        }145    }146147    private func finalize(conversationID: UUID, stats: GenerationStats?, finish: GenerationFinishReason) {148        defer {149            isGenerating = false150            streamingConversationID = nil151            streamingText = ""152            streamingThinking = ""153            isThinking = false154            liveTokensPerSecond = 0155            generationTask = nil156        }157        guard let app, var conversation = app.conversations.first(where: { $0.id == conversationID })158        else { return }159        let content = streamingText.trimmingCharacters(in: .whitespacesAndNewlines)160        let thinking = streamingThinking.trimmingCharacters(in: .whitespacesAndNewlines)161        guard !content.isEmpty || !thinking.isEmpty else { return }162        let message = Message(163            role: .assistant,164            content: content,165            thinking: thinking.isEmpty ? nil : thinking,166            stats: stats.map {167                MessageStats(168                    timeToFirstToken: $0.timeToFirstToken,169                    tokensPerSecond: $0.tokensPerSecond,170                    promptTokenCount: $0.promptTokenCount,171                    generationTokenCount: $0.generationTokenCount,172                    peakMemoryBytes: $0.peakMemoryBytes173                )174            }175        )176        conversation.messages.append(message)177        app.update(conversation)178        _ = finish  // reason currently not surfaced beyond stats179180        if conversation.title == "New Chat" {181            autoTitle(conversation: conversation)182        }183    }184185    /// Short, cheap title generation after the first exchange, using the186    /// loaded model itself. Falls back to a truncated first prompt.187    private func autoTitle(conversation: Conversation) {188        guard let app else { return }189        guard let firstUser = conversation.messages.first(where: { $0.role == .user }) else { return }190        let fallback = String(firstUser.content.prefix(48))191192        Task {193            var title = fallback194            do {195                let prompt = """196                    Reply with a title of at most 5 words for a conversation that starts with \197                    this message, and nothing else — no quotes, no punctuation at the end:198                    \(String(firstUser.content.prefix(500)))199                    """200                let temp = Conversation(params: GenerationParams(temperature: 0.1, maxTokens: 24))201                try await app.engine.startSession(conversation: temp)202                var generated = ""203                let events = try await app.engine.generate(prompt: prompt, params: temp.params)204                for try await event in events {205                    if case .token(let t) = event { generated += t }206                }207                var parser = ThinkTagParser()208                let (visible, _, _) = parser.consume(generated)209                let cleaned = visible210                    .trimmingCharacters(in: .whitespacesAndNewlines)211                    .trimmingCharacters(in: CharacterSet(charactersIn: "\"'.“”"))212                if !cleaned.isEmpty { title = String(cleaned.prefix(60)) }213                // Rebind the engine session to the real conversation.214                try? await app.engine.startSession(conversation: conversation)215            } catch {216                // Fallback title already set.217            }218            if var c = app.conversations.first(where: { $0.id == conversation.id }) {219                c.title = title220                app.update(c, touch: false)221            }222        }223    }224}225226/// Incremental parser splitting a token stream into visible text and227/// `<think>…</think>` content, robust to tags split across chunks.228struct ThinkTagParser {229    private var inThink = false230    private var pending = ""231232    /// Returns (visibleDelta, thinkingDelta, isInsideThink).233    mutating func consume(_ chunk: String) -> (String, String, Bool) {234        pending += chunk235        var visible = ""236        var thinking = ""237238        while true {239            if inThink {240                if let range = pending.range(of: "</think>") {241                    thinking += pending[..<range.lowerBound]242                    pending = String(pending[range.upperBound...])243                    inThink = false244                } else {245                    // Keep a possible partial closing tag in the buffer.246                    let safe = safeEmitLength(of: pending, partial: "</think>")247                    thinking += pending.prefix(safe)248                    pending = String(pending.dropFirst(safe))249                    break250                }251            } else {252                if let range = pending.range(of: "<think>") {253                    visible += pending[..<range.lowerBound]254                    pending = String(pending[range.upperBound...])255                    inThink = true256                } else {257                    let safe = safeEmitLength(of: pending, partial: "<think>")258                    visible += pending.prefix(safe)259                    pending = String(pending.dropFirst(safe))260                    break261                }262            }263        }264        return (visible, thinking, inThink)265    }266267    /// Length of `text` that can be emitted without cutting a partial `tag`268    /// suffix that might complete in the next chunk.269    private func safeEmitLength(of text: String, partial tag: String) -> Int {270        let maxKeep = min(tag.count - 1, text.count)271        for keep in stride(from: maxKeep, through: 1, by: -1) {272            if text.hasSuffix(String(tag.prefix(keep))) {273                return text.count - keep274            }275        }276        return text.count277    }278}279