spb/zyquo-local Public MIT

1//2//  ChatSession.swift3//  Zyquo Local4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//89import Foundation10import MLXLMCommon1112/// Zyquo Local's multi-turn session: converts persisted conversation history13/// into the model's chat-template form, truncates oldest turns when the14/// context window would overflow (always keeping the system prompt), and15/// reuses the KV cache across turns via the underlying MLX session.16///17/// Not thread-safe by design — owned and accessed only by `InferenceEngine`.18final class ChatSession {19    /// Conversation this session is bound to.20    let conversationID: UUID21    private let contextWindow: Int22    private let mlxSession: MLXLMCommon.ChatSession2324    /// Rough chars-per-token estimate used for truncation budgeting.25    private static let charsPerToken = 3.526    /// Fraction of the context window budgeted for history (rest is reserved27    /// for the generated response).28    private static let historyBudgetFraction = 0.72930    /// Number of tokens estimated to be used by history + system prompt.31    private(set) var estimatedHistoryTokens: Int3233    init(34        container: ModelContainer,35        conversation: Conversation,36        contextWindow: Int,37        params: GenerationParams38    ) {39        self.conversationID = conversation.id40        self.contextWindow = contextWindow4142        let truncated = Self.truncatedHistory(43            messages: conversation.messages,44            contextWindow: contextWindow45        )46        self.estimatedHistoryTokens = Self.estimateTokens(47            truncated.map(\.content).joined() + (conversation.systemPrompt ?? "")48        )4950        let history: [Chat.Message] = truncated.map { message in51            switch message.role {52            case .system: .system(message.content)53            case .user: .user(message.content)54            case .assistant: .assistant(message.content)55            }56        }57        self.mlxSession = MLXLMCommon.ChatSession(58            container,59            instructions: conversation.systemPrompt,60            history: history,61            generateParameters: params.toMLX()62        )63    }6465    /// Streams a response to a new user prompt, honoring per-call parameters66    /// without losing the KV cache.67    func stream(prompt: String, params: GenerationParams) -> AsyncThrowingStream<Generation, Error> {68        mlxSession.generateParameters = params.toMLX()69        estimatedHistoryTokens += Self.estimateTokens(prompt)70        return mlxSession.streamDetails(to: prompt)71    }7273    /// Estimated context usage in tokens (for the UI context bar).74    var estimatedContextUsage: (used: Int, window: Int) {75        (estimatedHistoryTokens, contextWindow)76    }7778    func noteResponse(_ text: String) {79        estimatedHistoryTokens += Self.estimateTokens(text)80    }8182    // MARK: - Truncation8384    static func estimateTokens(_ text: String) -> Int {85        Int(Double(text.count) / charsPerToken) + 186    }8788    /// Drops oldest non-system turns until the estimated history fits the89    /// budgeted share of the context window. The system prompt (a leading90    /// system message, if any) is always kept.91    static func truncatedHistory(messages: [Message], contextWindow: Int) -> [Message] {92        let budget = Int(Double(contextWindow) * historyBudgetFraction)93        var system: [Message] = []94        var turns: [Message] = []95        for m in messages {96            if m.role == .system && turns.isEmpty {97                system.append(m)98            } else {99                turns.append(m)100            }101        }102        func total(_ list: [Message]) -> Int {103            list.reduce(0) { $0 + estimateTokens($1.content) }104        }105        while turns.count > 1, total(system) + total(turns) > budget {106            turns.removeFirst()107        }108        return system + turns109    }110}111