spb/zyquo-atlas Public License

1//2//  AIService.swift3//  Zyquo Atlas4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//8//  Orchestrates browser AI actions against the ported provider clients: builds9//  the grounded, injection-safe request, streams tokens into published state,10//  and cancels the in-flight request on navigation / Stop (docs/11//  AI-BROWSER-RESEARCH.md §5 — dropping the stream cancels the URLSession task).12//  One AIService drives one AI surface (e.g. a tab's chat-with-page sidebar).13//1415import Foundation16import Combine1718@MainActor19final class AIService: ObservableObject {20    /// Streamed answer text for the current action.21    @Published private(set) var output: String = ""22    @Published private(set) var isStreaming: Bool = false23    @Published private(set) var errorText: String?24    /// Progress note during map-reduce ("Reading section 3/9…").25    @Published private(set) var statusNote: String?2627    private var task: Task<Void, Never>?2829    /// Cancels any in-flight request (call on navigation, Stop, tab close).30    func cancel() {31        task?.cancel()32        task = nil33        if isStreaming { isStreaming = false }34    }3536    deinit { task?.cancel() }3738    /// Runs an action over a page context, streaming the answer into `output`.39    func run(_ action: AIAction,40             on context: PageContext,41             model: AIModel,42             apiKey: String,43             extra: String? = nil) {44        cancel()45        reset()46        isStreaming = true4748        task = Task { [weak self] in49            guard let self else { return }50            do {51                if action == .summarize {52                    try await self.summarize(context: context, model: model, apiKey: apiKey)53                } else {54                    let messages = [55                        Message(role: .user, text: action.userPrompt(for: context, extra: extra))56                    ]57                    try await self.stream(system: action.systemGuidance,58                                          messages: messages, model: model, apiKey: apiKey)59                }60                await MainActor.run { self.isStreaming = false }61            } catch is CancellationError {62                // Expected on navigation/Stop — leave partial output in place.63            } catch {64                await MainActor.run {65                    self.errorText = (error as? LocalizedError)?.errorDescription66                        ?? error.localizedDescription67                    self.isStreaming = false68                }69            }70        }71    }7273    /// Streams an arbitrary prompt (used by omnibox ask, ask-about-history,74    /// multi-tab compare). The caller builds the grounded, hygiene-wrapped text.75    func runRawPrompt(system: String?, userText: String, model: AIModel, apiKey: String) {76        cancel()77        reset()78        isStreaming = true79        task = Task { [weak self] in80            guard let self else { return }81            do {82                try await self.stream(system: system,83                                      messages: [Message(role: .user, text: userText)],84                                      model: model, apiKey: apiKey)85                await MainActor.run { self.isStreaming = false }86            } catch is CancellationError {87            } catch {88                await MainActor.run {89                    self.errorText = (error as? LocalizedError)?.errorDescription ?? error.localizedDescription90                    self.isStreaming = false91                }92            }93        }94    }9596    // MARK: - Summarize (stuff-first, map-reduce for long pages)9798    private func summarize(context: PageContext, model: AIModel, apiKey: String) async throws {99        switch Summarizer.plan(for: context, model: model) {100        case .stuff:101            let messages = [Message(role: .user, text: AIAction.summarize.userPrompt(for: context))]102            try await stream(system: AIAction.summarize.systemGuidance,103                             messages: messages, model: model, apiKey: apiKey)104105        case .mapReduce:106            let chunks = Summarizer.mapChunks(for: context)107            var partials: [String] = []108            for (i, chunk) in chunks.enumerated() {109                try Task.checkCancellation()110                await MainActor.run { self.statusNote = "Reading section \(i + 1)/\(chunks.count)…" }111                let partial = try await complete(112                    system: AIAction.summarize.systemGuidance,113                    userText: Summarizer.mapPrompt(chunk: chunk, title: context.title),114                    model: model, apiKey: apiKey115                )116                partials.append(partial)117            }118            await MainActor.run { self.statusNote = nil }119            let messages = [Message(role: .user,120                                    text: Summarizer.reducePrompt(partials: partials, title: context.title))]121            try await stream(system: AIAction.summarize.systemGuidance,122                             messages: messages, model: model, apiKey: apiKey)123        }124    }125126    // MARK: - Streaming primitive127128    /// Streams a chat request into `output`. Cancellation propagates to the129    /// provider stream by dropping the for-await loop.130    private func stream(system: String?, messages: [Message], model: AIModel, apiKey: String) async throws {131        let client = ProviderRegistry.client(for: model)132        // Reasoning models stream hidden thinking before the answer — budget for it.133        var params = ChatParameters(maxTokens: model.capabilities.reasoning ? 8192 : 1024)134        if model.parameterSupport.reasoningEffort { params.reasoningEffort = "low" }135        let request = ChatRequest(model: model, systemPrompt: system, messages: messages, parameters: params)136137        for try await event in client.streamChat(request, apiKey: apiKey) {138            try Task.checkCancellation()139            switch event {140            case .textDelta(let delta):141                await MainActor.run { self.output += delta }142            case .reasoningDelta, .citations, .usage, .finished:143                break144            }145        }146    }147148    /// Non-streaming completion used by the map phase.149    private func complete(system: String?, userText: String, model: AIModel, apiKey: String) async throws -> String {150        let client = ProviderRegistry.client(for: model)151        var params = ChatParameters(maxTokens: 512)152        if model.parameterSupport.reasoningEffort { params.reasoningEffort = "low" }153        let request = ChatRequest(model: model, systemPrompt: system,154                                  messages: [Message(role: .user, text: userText)],155                                  parameters: params, stream: false)156        let reply = try await client.complete(request, apiKey: apiKey)157        return reply.text158    }159160    private func reset() {161        output = ""162        errorText = nil163        statusNote = nil164    }165}166