spb/zyquo-cloud Public MIT

1//2//  ConversationStore.swift3//  Zyquo Cloud4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//8//  Central chat state: the conversation list, selection, and the chat engine9//  that streams provider responses into messages. All mutation happens on the10//  main actor; streaming work runs in tasks that hop back for UI updates.11//1213import Foundation1415@MainActor16final class ConversationStore: ObservableObject {17    @Published var conversations: [Conversation] = []18    @Published var selectedID: Conversation.ID?19    @Published var searchText: String = ""20    /// Live streaming task per conversation (supports parallel streams in compare mode).21    @Published private(set) var streamingConversationIDs: Set<Conversation.ID> = []2223    let catalog: ModelCatalog24    let vault: KeyVaultStore25    private let persistence: PersistenceService26    private var streamTasks: [Conversation.ID: Task<Void, Never>] = [:]2728    /// Global default system prompt (Settings → Advanced).29    @Published var defaultSystemPrompt: String {30        didSet { persistence.save(defaultSystemPrompt, to: "default-system-prompt.json") }31    }3233    init(34        catalog: ModelCatalog,35        vault: KeyVaultStore,36        persistence: PersistenceService = .shared37    ) {38        self.catalog = catalog39        self.vault = vault40        self.persistence = persistence41        self.defaultSystemPrompt = persistence.load(String.self, from: "default-system-prompt.json") ?? ""42        conversations = persistence.loadConversations()43        selectedID = conversations.first?.id44    }4546    // MARK: - Selection & lookup4748    var selected: Conversation? {49        get { conversations.first { $0.id == selectedID } }50    }5152    func binding(for id: Conversation.ID) -> Int? {53        conversations.firstIndex { $0.id == id }54    }5556    func isStreaming(_ id: Conversation.ID) -> Bool {57        streamingConversationIDs.contains(id)58    }5960    // MARK: - CRUD6162    @discardableResult63    func newConversation(model: AIModel? = nil, persona: Persona? = nil) -> Conversation {64        let chosen = model ?? catalog.defaultModel65        var conversation = Conversation(66            modelID: chosen?.id ?? "",67            provider: chosen?.provider ?? .openai,68            systemPrompt: persona?.systemPrompt ?? (defaultSystemPrompt.isEmpty ? nil : defaultSystemPrompt)69        )70        if let persona {71            conversation.personaID = persona.id72            conversation.parameters = persona.parameters73            if let modelID = persona.modelID, let provider = persona.provider {74                conversation.modelID = modelID75                conversation.provider = provider76            }77        }78        conversations.insert(conversation, at: 0)79        selectedID = conversation.id80        persistence.save(conversation)81        return conversation82    }8384    func delete(_ id: Conversation.ID) {85        stopStreaming(id)86        if let conversation = conversations.first(where: { $0.id == id }) {87            persistence.delete(conversation)88        }89        conversations.removeAll { $0.id == id }90        if selectedID == id { selectedID = conversations.first?.id }91    }9293    func update(_ conversation: Conversation) {94        guard let index = conversations.firstIndex(where: { $0.id == conversation.id }) else { return }95        var updated = conversation96        updated.updatedAt = Date()97        conversations[index] = updated98        persistence.save(updated)99    }100101    func togglePin(_ id: Conversation.ID) {102        guard var conversation = conversations.first(where: { $0.id == id }) else { return }103        conversation.isPinned.toggle()104        update(conversation)105    }106107    func rename(_ id: Conversation.ID, to title: String) {108        guard var conversation = conversations.first(where: { $0.id == id }) else { return }109        conversation.title = title110        conversation.hasAutoTitle = false111        update(conversation)112    }113114    // MARK: - Sidebar grouping115116    struct SidebarGroup: Identifiable {117        let id: String118        let title: String119        let conversations: [Conversation]120    }121122    /// Pinned / Today / Yesterday / Previous 7 Days / Older, filtered by search.123    var sidebarGroups: [SidebarGroup] {124        let filtered = searchText.isEmpty125            ? conversations126            : conversations.filter { conversation in127                conversation.title.localizedCaseInsensitiveContains(searchText)128                    || conversation.messages.contains {129                        $0.text.localizedCaseInsensitiveContains(searchText)130                    }131            }132        let calendar = Calendar.current133        let now = Date()134        var pinned: [Conversation] = []135        var today: [Conversation] = []136        var yesterday: [Conversation] = []137        var week: [Conversation] = []138        var older: [Conversation] = []139        for conversation in filtered {140            if conversation.isPinned { pinned.append(conversation); continue }141            if calendar.isDateInToday(conversation.updatedAt) { today.append(conversation) }142            else if calendar.isDateInYesterday(conversation.updatedAt) { yesterday.append(conversation) }143            else if conversation.updatedAt > now.addingTimeInterval(-7 * 86_400) { week.append(conversation) }144            else { older.append(conversation) }145        }146        return [147            SidebarGroup(id: "pinned", title: "Pinned", conversations: pinned),148            SidebarGroup(id: "today", title: "Today", conversations: today),149            SidebarGroup(id: "yesterday", title: "Yesterday", conversations: yesterday),150            SidebarGroup(id: "week", title: "Previous 7 Days", conversations: week),151            SidebarGroup(id: "older", title: "Older", conversations: older),152        ].filter { !$0.conversations.isEmpty }153    }154155    // MARK: - Chat engine156157    /// Sends the user's text (and attachments) in a conversation and streams158    /// the assistant reply. Optionally targets a different model for this159    /// message only (per-message model switch).160    func send(161        text: String,162        attachments: [Attachment] = [],163        in conversationID: Conversation.ID,164        overrideModel: AIModel? = nil165    ) {166        guard var conversation = conversations.first(where: { $0.id == conversationID }) else { return }167        guard let model = overrideModel168            ?? catalog.model(id: conversation.modelID, provider: conversation.provider)169            ?? catalog.defaultModel170        else { return }171172        if let overrideModel {173            conversation.modelID = overrideModel.id174            conversation.provider = overrideModel.provider175        }176177        var userMessage = Message(role: .user, text: text, attachments: attachments)178        userMessage.modelID = model.id179        userMessage.provider = model.provider180        conversation.messages.append(userMessage)181        update(conversation)182183        generateReply(in: conversationID, model: model)184    }185186    /// Regenerates the last assistant message (optionally with another model).187    func regenerate(in conversationID: Conversation.ID, with model: AIModel? = nil) {188        guard var conversation = conversations.first(where: { $0.id == conversationID }) else { return }189        if conversation.messages.last?.role == .assistant {190            conversation.messages.removeLast()191            update(conversation)192        }193        guard let target = model194            ?? catalog.model(id: conversation.modelID, provider: conversation.provider)195        else { return }196        generateReply(in: conversationID, model: target)197    }198199    /// Replaces a user message's text and regenerates from that point.200    func editAndResend(messageID: Message.ID, newText: String, in conversationID: Conversation.ID) {201        guard var conversation = conversations.first(where: { $0.id == conversationID }),202              let index = conversation.messages.firstIndex(where: { $0.id == messageID })203        else { return }204        conversation.messages[index].text = newText205        conversation.messages.removeSubrange((index + 1)...)206        update(conversation)207        guard let model = catalog.model(id: conversation.modelID, provider: conversation.provider) else { return }208        generateReply(in: conversationID, model: model)209    }210211    func deleteMessage(_ messageID: Message.ID, in conversationID: Conversation.ID) {212        guard var conversation = conversations.first(where: { $0.id == conversationID }) else { return }213        conversation.messages.removeAll { $0.id == messageID }214        update(conversation)215    }216217    func stopStreaming(_ conversationID: Conversation.ID) {218        streamTasks[conversationID]?.cancel()219        streamTasks[conversationID] = nil220        streamingConversationIDs.remove(conversationID)221        finalizeStreamingMessage(in: conversationID)222    }223224    private func generateReply(in conversationID: Conversation.ID, model: AIModel) {225        guard let index = conversations.firstIndex(where: { $0.id == conversationID }) else { return }226227        var placeholder = Message(role: .assistant, text: "")228        placeholder.modelID = model.id229        placeholder.provider = model.provider230        placeholder.isStreaming = true231        conversations[index].messages.append(placeholder)232        let messageID = placeholder.id233234        let request = ChatRequest(235            model: model,236            systemPrompt: conversations[index].systemPrompt,237            messages: conversations[index].messages.filter { $0.id != messageID && $0.errorText == nil },238            parameters: conversations[index].parameters239        )240241        streamingConversationIDs.insert(conversationID)242        let task = Task { [weak self] in243            guard let self else { return }244            do {245                let apiKey = try self.vault.apiKey(for: model.provider)246                let client = ProviderRegistry.client(for: model)247                var usage: TokenUsage?248                for try await event in client.streamChat(request, apiKey: apiKey) {249                    if Task.isCancelled { break }250                    switch event {251                    case .textDelta(let delta):252                        self.mutateMessage(messageID, in: conversationID) { $0.text += delta }253                    case .reasoningDelta(let delta):254                        self.mutateMessage(messageID, in: conversationID) {255                            $0.reasoning = ($0.reasoning ?? "") + delta256                        }257                    case .citations(let citations):258                        self.mutateMessage(messageID, in: conversationID) { $0.citations = citations }259                    case .usage(let u):260                        usage = u261                    case .finished:262                        break263                    }264                }265                let finalUsage = usage266                self.mutateMessage(messageID, in: conversationID) { message in267                    message.isStreaming = false268                    if let u = finalUsage {269                        message.usage = u270                        message.estimatedCost = model.pricing?.cost(271                            inputTokens: u.inputTokens, outputTokens: u.outputTokens272                        )273                    }274                }275            } catch let error as ProviderError {276                if case .cancelled = error {277                    self.mutateMessage(messageID, in: conversationID) { $0.isStreaming = false }278                } else {279                    self.mutateMessage(messageID, in: conversationID) {280                        $0.isStreaming = false281                        $0.errorText = error.localizedDescription282                    }283                }284            } catch {285                self.mutateMessage(messageID, in: conversationID) {286                    $0.isStreaming = false287                    $0.errorText = error.localizedDescription288                }289            }290            self.streamingConversationIDs.remove(conversationID)291            self.streamTasks[conversationID] = nil292            self.persistSnapshot(of: conversationID)293            self.autoTitleIfNeeded(conversationID)294        }295        streamTasks[conversationID] = task296    }297298    private func mutateMessage(299        _ messageID: Message.ID,300        in conversationID: Conversation.ID,301        _ mutate: (inout Message) -> Void302    ) {303        guard let ci = conversations.firstIndex(where: { $0.id == conversationID }),304              let mi = conversations[ci].messages.firstIndex(where: { $0.id == messageID })305        else { return }306        mutate(&conversations[ci].messages[mi])307    }308309    private func finalizeStreamingMessage(in conversationID: Conversation.ID) {310        guard let ci = conversations.firstIndex(where: { $0.id == conversationID }) else { return }311        for mi in conversations[ci].messages.indices where conversations[ci].messages[mi].isStreaming {312            conversations[ci].messages[mi].isStreaming = false313        }314        persistSnapshot(of: conversationID)315    }316317    private func persistSnapshot(of conversationID: Conversation.ID) {318        guard var conversation = conversations.first(where: { $0.id == conversationID }) else { return }319        conversation.updatedAt = Date()320        if let index = conversations.firstIndex(where: { $0.id == conversationID }) {321            conversations[index] = conversation322        }323        persistence.save(conversation)324    }325326    // MARK: - Auto titles327328    /// After the first completed exchange, asks the provider's cheapest model329    /// for a short title.330    private func autoTitleIfNeeded(_ conversationID: Conversation.ID) {331        guard let conversation = conversations.first(where: { $0.id == conversationID }),332              conversation.hasAutoTitle,333              conversation.messages.filter({ $0.role == .assistant && !$0.text.isEmpty }).count == 1,334              let cheap = catalog.cheapestModel(for: conversation.provider),335              let firstUser = conversation.messages.first(where: { $0.role == .user })336        else { return }337338        let assistantText = conversation.messages.last { $0.role == .assistant }?.text ?? ""339        let prompt = """340        Write a title of at most 5 words for this conversation. Reply with the title only, no quotes.341342        User: \(firstUser.text.prefix(500))343        Assistant: \(assistantText.prefix(500))344        """345        Task { [weak self] in346            guard let self else { return }347            do {348                let apiKey = try self.vault.apiKey(for: cheap.provider)349                let client = ProviderRegistry.client(for: cheap)350                let request = ChatRequest(351                    model: cheap,352                    systemPrompt: nil,353                    messages: [Message(role: .user, text: prompt)],354                    parameters: ChatParameters(maxTokens: 24),355                    stream: false356                )357                let reply = try await client.complete(request, apiKey: apiKey)358                let title = reply.text359                    .trimmingCharacters(in: .whitespacesAndNewlines)360                    .trimmingCharacters(in: CharacterSet(charactersIn: "\"“”"))361                guard !title.isEmpty,362                      var current = self.conversations.first(where: { $0.id == conversationID }),363                      current.hasAutoTitle364                else { return }365                current.title = String(title.prefix(60))366                self.update(current)367            } catch {368                // Title generation is best-effort; keep "New Chat" on failure.369            }370        }371    }372}373