diff --git a/Package.resolved b/Package.resolved index 23d12109..8d90fa77 100644 --- a/Package.resolved +++ b/Package.resolved @@ -1,5 +1,5 @@ { - "originHash" : "279e029257048c3eb542ea9222e1b700d481a3a6c37cdad50f82bb1d9faca5fb", + "originHash" : "50c257fc09167ff54a06ff15cfb91ce88e50e7e43303b3edf68e0c73c609108c", "pins" : [ { "identity" : "eventsource", @@ -19,24 +19,6 @@ "version" : "1.3.1" } }, - { - "identity" : "llama.swift", - "kind" : "remoteSourceControl", - "location" : "https://github.com/mattt/llama.swift", - "state" : { - "revision" : "716419d4d7aa542fce301e809cde7234c68ddbc6", - "version" : "2.10549.0" - } - }, - { - "identity" : "swift-asn1", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-asn1.git", - "state" : { - "revision" : "d9a5b37470adc940d22c3bcd5ca6953a516b727f", - "version" : "1.7.2" - } - }, { "identity" : "swift-atomics", "kind" : "remoteSourceControl", @@ -55,15 +37,6 @@ "version" : "1.6.0" } }, - { - "identity" : "swift-crypto", - "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-crypto.git", - "state" : { - "revision" : "da9d28d69ebe3894b18376c8f2395c2f37b8448f", - "version" : "4.5.2" - } - }, { "identity" : "swift-http-types", "kind" : "remoteSourceControl", @@ -73,15 +46,6 @@ "version" : "1.8.0" } }, - { - "identity" : "swift-huggingface", - "kind" : "remoteSourceControl", - "location" : "https://github.com/huggingface/swift-huggingface", - "state" : { - "revision" : "c30143a4616f3fa994ee3ab3e3464f752845c07c", - "version" : "1.0.0" - } - }, { "identity" : "swift-nio", "kind" : "remoteSourceControl", diff --git a/Package.swift b/Package.swift index a4c7f912..0592f42f 100644 --- a/Package.swift +++ b/Package.swift @@ -29,7 +29,8 @@ let package = Package( .default(enabledTraits: []), ], dependencies: [ - .package(url: "https://github.com/huggingface/swift-transformers", from: "1.0.0"), + // fork(main 分支含 underscore 输入名修复 c2cfb40,已 push) + .package(url: "https://github.com/foxterm/swift-transformers.git", branch: "main"), .package(url: "https://github.com/huggingface/swift-huggingface", from: "1.0.0"), .package( url: "https://github.com/mattt/EventSource", @@ -100,6 +101,13 @@ let package = Package( package: "async-http-client", condition: .when(traits: ["AsyncHTTPClient"]) ), + ], + swiftSettings: [ + // 将启用的 trait 同步为源文件编译条件(#if CoreML / #if MLX / #if Llama) + .define("CoreML", .when(traits: ["CoreML"])), + .define("MLX", .when(traits: ["MLX"])), + .define("Llama", .when(traits: ["Llama"])), + .define("AsyncHTTPClient", .when(traits: ["AsyncHTTPClient"])), ] ), .macro( diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index 4b57d92d..0682cd99 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -85,6 +85,20 @@ public final class LanguageModelSession: @unchecked Sendable { private let model: any LanguageModel +#if Llama + /// 底层模型的本地上下文窗口(仅 Llama 等可查询的本地模型返回非 nil; + /// 来自 GGUF 的 context_length)。 + public var modelContextSize: UInt32? { + (model as? LlamaLanguageModel)?.modelContextSize + } + + /// 本地 Llama 模型性能快照(权重/参数/内存占用/系统可用/CPU/速度)。 + /// 云端模型返回 nil。 + public var localModelPerformance: LlamaLanguageModel.LlamaPerformanceSnapshot? { + (model as? LlamaLanguageModel)?.performanceSnapshot() + } +#endif + /// The tools that the model can call during the session. /// /// For a session created with dynamic instructions, this is empty. @@ -97,6 +111,19 @@ public final class LanguageModelSession: @unchecked Sendable { /// can read the session's tools. public let tools: [any Tool] + /// 动态工具覆盖:非 nil 时优先用于每次请求(工具开关实时生效, + /// 无需重建会话/重新创建模型;本地 Llama 模型因此不必重载 GGUF)。 + @ObservationIgnored private var dynamicTools: [any Tool]? + @ObservationIgnored private let dynamicToolsLock = NSLock() + + /// 更新会话工具集(运行中实时生效,不重建模型,不打断上下文)。 + /// 下次模型请求(respond / stream)即使用新工具。 + nonisolated public func setDynamicTools(_ tools: [any Tool]) { + dynamicToolsLock.lock() + dynamicTools = tools + dynamicToolsLock.unlock() + } + /// The instructions for the session, if any. /// /// For a session created with dynamic instructions, this is `nil`. @@ -167,8 +194,13 @@ public final class LanguageModelSession: @unchecked Sendable { /// It's public so that language models outside this module /// can read the inputs for each request. nonisolated public func resolvedRequestContext() -> RequestContext { + // 动态工具覆盖优先;未设置时用创建时的 tools + dynamicToolsLock.lock() + let effectiveTools = dynamicTools ?? tools + dynamicToolsLock.unlock() + guard let dynamicInstructions else { - return RequestContext(transcript: transcript, instructions: instructions, tools: tools) + return RequestContext(transcript: transcript, instructions: instructions, tools: effectiveTools) } // Evaluate the body for this request. diff --git a/Sources/AnyLanguageModel/Models/CoreMLLanguageModel.swift b/Sources/AnyLanguageModel/Models/CoreMLLanguageModel.swift index 5c17b946..05ad5922 100644 --- a/Sources/AnyLanguageModel/Models/CoreMLLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/CoreMLLanguageModel.swift @@ -56,8 +56,26 @@ } do { - // Load the model with the specified compute units - self.model = try Models.LanguageModel.loadCompiled(url: url, computeUnits: computeUnits) + // 本地 tokenizer 优先:模型目录含 tokenizer.json 时直接本地加载,避免从 HuggingFace Hub 下载 + // (本地模型名通常是转换产物名,在 Hub 上不存在)。 + let modelFolder = url.deletingLastPathComponent() + var localTokenizer: (any Tokenizer)? = nil + if FileManager.default.fileExists( + atPath: modelFolder.appendingPathComponent("tokenizer.json").path + ) { + localTokenizer = try? await AutoTokenizer.from(modelFolder: modelFolder) + } + if let localTokenizer { + // Load the model with the specified compute units + self.model = try Models.LanguageModel.loadCompiled( + url: url, + computeUnits: computeUnits, + tokenizer: localTokenizer + ) + } else { + // Load the model with the specified compute units + self.model = try Models.LanguageModel.loadCompiled(url: url, computeUnits: computeUnits) + } } catch { // Map CoreML errors to our specific error cases throw CoreMLLanguageModelError.modelInvalid(url, underlyingError: error) diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index f0a1eb0e..baa0a08c 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -1,4 +1,8 @@ import Foundation +#if os(macOS) + import Darwin + import Darwin.Mach +#endif #if Llama import JSONSchema import LlamaSwift @@ -20,6 +24,140 @@ import Foundation } } + /// Incrementally parses thinking spans (all common tag pairs) out of raw + /// generated text, streaming the thinking content separately from the + /// visible body. Handles tags that span multiple tokens/chunks by + /// buffering only the trailing tag-prefix. + struct ThoughtParser { + /// 兜底标签对全集(模型模板检测失败时使用) + static let defaultPairs: [(open: String, close: String)] = [ + ("<|startofthink|>", ""), + ("<|begin_of_thought|>", ""), + ("", ""), + ("", ""), + ("", ""), + ] + + private let pairs: [(open: String, close: String)] + private let maxOpenTagLength: Int + + private(set) var thinking = "" + private var buffer = "" + private var activeClose: String? + + /// - Parameter pairs: 思考标签对;缺省使用兜底全集。 + init(pairs: [(open: String, close: String)] = ThoughtParser.defaultPairs) { + self.pairs = pairs + self.maxOpenTagLength = pairs.map { $0.open.count }.max() ?? 8 + } + + mutating func consume(_ raw: String) { + buffer += raw + while true { + if let close = activeClose { + if let found = buffer.range(of: close) { + thinking += buffer[.. keep { + let cut = buffer.index(buffer.endIndex, offsetBy: -keep) + thinking += buffer[.., close: String)? + for pair in self.pairs { + if let r = buffer.range(of: pair.open), + best == nil || r.lowerBound < best!.range.lowerBound { + best = (r, pair.close) + } + } + if let best { + buffer = String(buffer[best.range.upperBound...]) + activeClose = best.close + } else { + // 保留最长 open 标签前缀,等待下一 chunk 确认 + let keep = self.maxOpenTagLength - 1 + if buffer.count > keep { + let cut = buffer.index(buffer.endIndex, offsetBy: -keep) + buffer = String(buffer[cut...]) + } else { + break + } + } + } + } + } + + /// 生成结束:flush 剩余 buffer(未闭合的思考段也完整输出) + mutating func finalize() { + thinking += buffer + buffer = "" + activeClose = nil + } + } + +extension ThoughtParser { + /// 从 llama 模型元数据(tokenizer.chat_template)动态检测思考标签对, + /// 无需针对每个模型写死标签。查不到模板或无思考标签时返回 nil(用兜底全集)。 + static func detectThoughtSpans(model: OpaquePointer?) -> [(open: String, close: String)]? { + guard let model else { return nil } + var buffer = [CChar](repeating: 0, count: 65536) + let count = llama_model_meta_val_str( + model, "tokenizer.chat_template", &buffer, buffer.count + ) + guard count > 0 else { return nil } + let template = String(cString: buffer) + return thoughtSpanPairs(fromTemplate: template) + } + + /// 从 chat template 文本提取思考标签对: + /// 1) 已知标签集合匹配(模板里有哪个用哪个,全无则跳过); + /// 2) 通用成对发现:… 且 X 含 think/thought/reasoning 关键词。 + static func thoughtSpanPairs(fromTemplate template: String) -> [(open: String, close: String)]? { + let known: [(open: String, close: String)] = [ + ("<|startofthink|>", ""), + ("<|begin_of_thought|>", ""), + ("", ""), + ("", ""), + ("", ""), + ] + let matched = known.filter { template.contains($0.open) } + if !matched.isEmpty { return matched } + + // 通用成对发现: 与 同时出现在模板,且 X 含思考关键词 + let pattern = #"<(/?)\|?([A-Za-z_]*?(?:think|thought|reasoning)[A-Za-z_]*?)\|?>"# + guard let regex = try? NSRegularExpression(pattern: pattern) else { return nil } + let ns = template as NSString + var openTags: Set = [] + var closeTags: Set = [] + for match in regex.matches(in: template, range: NSRange(location: 0, length: ns.length)) { + if let range = Range(match.range, in: template) { + let tag = String(template[range]) + if tag.hasPrefix("] = [] + private var locked = false + + func acquire() async { + if !locked { + locked = true + return + } + await withCheckedContinuation { waiters.append($0) } + } + + func release() { + if waiters.isEmpty { + locked = false + } else { + waiters.removeFirst().resume() + } + } + } + + private let structuredGate = InferenceGate() + /// A context kept alive for one session so exchanges reuse its state. private struct CachedSessionContext { let sessionID: ObjectIdentifier @@ -635,6 +814,12 @@ import Foundation self.defaultOptions = ResolvedGenerationOptions() } + /// 模型原生上下文窗口长度(读取 GGUF 自带 context_length, + /// 经 llama_n_ctx_train 取得)。模型加载前为 nil。 + public var modelContextSize: UInt32? { + nativeContextSize + } + deinit { if let cached = cachedSessionContext { llama_free(cached.context) @@ -848,7 +1033,10 @@ import Foundation let imageMarker = mtmdContext != nil ? String(cString: mtmd_default_marker()) : nil if type == String.self { - let maxTokens = runtimeOptions.maximumResponseTokens ?? 100 + let maxTokens = Self.resolveMaxTokens( + requested: runtimeOptions.maximumResponseTokens, + contextSize: runtimeOptions.contextSize + ) let outputFormat = currentToolCallFormat() var toolRounds = ToolRoundLimit(provider: "Llama", options: options) var allEntries: [Transcript.Entry] = [] @@ -1002,7 +1190,10 @@ import Foundation } let context = try makeFreshContext(options: runtimeOptions) defer { llama_free(context) } - let maxTokens = structuredOptions.maximumResponseTokens ?? 512 + let maxTokens = Self.resolveMaxTokens( + requested: structuredOptions.maximumResponseTokens, + contextSize: structuredOptions.contextSize + ) let (jsonString, usage) = try await generateStructuredJSON( context: context, prompt: fullPrompt, @@ -1083,7 +1274,10 @@ import Foundation try ensureModelLoaded() let runtimeOptions = resolvedOptions(from: options) - let maxTokens = runtimeOptions.maximumResponseTokens ?? 100 + let maxTokens = Self.resolveMaxTokens( + requested: runtimeOptions.maximumResponseTokens, + contextSize: runtimeOptions.contextSize + ) let outputFormat = self.currentToolCallFormat() var toolRounds = ToolRoundLimit(provider: "Llama", options: options) var accumulatedEntries: [Transcript.Entry] = [] @@ -1129,11 +1323,36 @@ import Foundation var roundRaw = "" let terminator = toolContext?.format.callTerminator let withholdToolCalls = toolContext != nil + var thoughtParser = ThoughtParser( + pairs: self.detectedThoughtPairs ?? ThoughtParser.defaultPairs + ) + func updateReasoningEntry() { + guard !thoughtParser.thinking.isEmpty else { return } + let reasoningEntry = Transcript.Entry.reasoning( + Transcript.Reasoning( + origin: .llama, + segments: [.text( + Transcript.TextSegment(content: thoughtParser.thinking) + )] + ) + ) + if let lastIndex = accumulatedEntries.indices.last, + case .reasoning = accumulatedEntries[lastIndex] { + accumulatedEntries[lastIndex] = reasoningEntry + } else { + accumulatedEntries.append(reasoningEntry) + } + } let collectToken: (String) -> Bool = { tokenText in roundRaw += tokenText + // Stream thinking spans as reasoning transcript entries, + // separate from the visible body text. + thoughtParser.consume(tokenText) + updateReasoningEntry() let visible = outputFormat.streamingVisibleText( in: roundRaw, - withholdToolCalls: withholdToolCalls + withholdToolCalls: withholdToolCalls, + thoughtPairs: self.detectedThoughtPairs ) let total = emittedBase + visible if !total.isEmpty, total != lastYieldedText { @@ -1171,6 +1390,9 @@ import Foundation } usage.add(roundUsage) + // Flush any unclosed thinking span into the reasoning entry. + thoughtParser.finalize() + updateReasoningEntry() // Publish counts before tool execution, which may stop or fail. yieldSnapshot(lastYieldedText ?? emittedBase) @@ -1181,7 +1403,8 @@ import Foundation let roundVisible = outputFormat.streamingVisibleText( in: roundRaw, withholdToolCalls: withholdToolCalls, - holdPartialMarkers: false + holdPartialMarkers: false, + thoughtPairs: self.detectedThoughtPairs ) guard let format = toolContext?.format else { @@ -1259,6 +1482,10 @@ import Foundation throw LlamaLanguageModelError.invalidModelPath } + // 加载前内存预检:llama.cpp 在内存耗尽时会直接 abort()(进程级崩溃, + // 无异常可捕获),这里提前估算并拦截,避免把整机内存打满导致系统卡死。 + try ensureSufficientMemoryForLoading() + // Initialize backend lazily - must be done before loading model llama_backend_init() @@ -1300,9 +1527,212 @@ import Foundation self.model = loadedModel self.vocab = llama_model_get_vocab(loadedModel) self.tokenCache = StructuredGenerationTokenCache() + self.detectedThoughtPairs = ThoughtParser.detectThoughtSpans(model: loadedModel) + // 模型原生训练上下文:GGUF 自带 context_length, + // 未显式设置时用它作为上下文/输出上限的默认依据 + let native = llama_n_ctx_train(loadedModel) + self.nativeContextSize = native > 0 ? UInt32(native) : nil self.isModelLoaded = true } + /// 加载前内存预检:估算 GGUF 权重 + 计算缓冲所需内存, + /// 可用内存不足时提前抛 `insufficientMemory`,避免 llama.cpp 内部分配失败 + /// 调用 abort() 直接终止进程(该路径无异常可捕获,会让 App/系统直接卡死)。 + private func ensureSufficientMemoryForLoading() throws { + let fileSize: UInt64 + do { + let attrs = try FileManager.default.attributesOfItem(atPath: modelPath) + fileSize = (attrs[.size] as? NSNumber)?.uint64Value ?? 0 + } catch { + throw LlamaLanguageModelError.invalidModelPath + } + guard fileSize > 0 else { + throw LlamaLanguageModelError.invalidModelPath + } + + #if os(macOS) + // 权重 + 推理/计算缓冲保守按 1.5 倍估算;多模态投影器另计 + var estimate = Double(fileSize) * 1.5 + if let mmprojPath { + if let attrs = try? FileManager.default.attributesOfItem(atPath: mmprojPath), + let m = (attrs[.size] as? NSNumber)?.uint64Value { + estimate += Double(m) * 1.2 + } + } + let available = Double(Self.availableMemoryBytes()) + // 记录估算基线供运行时护栏使用 + memoryLoadEstimateBytes = UInt64(estimate) + // 保留 20% 余量,防止整机进入内存压力 + if estimate > available * 0.8 { + throw LlamaLanguageModelError.insufficientMemory( + requiredBytes: UInt64(estimate), + availableBytes: UInt64(available) + ) + } + #endif + } + + /// 运行时内存护栏:推理中 KV cache 随上下文持续增长、系统内存动态变化, + /// 每 N 个 token 检查一次系统可用内存,低于阈值立即中断生成。 + /// 这避免 llama.cpp 在内存耗尽时调用 abort() 把整机干死机。 + private func checkRuntimeMemoryPressure() throws { + #if os(macOS) + let free = Self.availableMemoryBytes() + // 阈值:模型估算的 1/4,保底 1GB;低于即认为内存吃紧并提前停止。 + // 留足余量,避免 macOS 内存压缩/换页导致整机卡顿(鼠标都动不了)。 + let lowerBound = max(memoryLoadEstimateBytes / 4, 1 * 1024 * 1024 * 1024) + if free < lowerBound { + throw LlamaLanguageModelError.insufficientMemory( + requiredBytes: memoryLoadEstimateBytes, + availableBytes: free + ) + } + #endif + } + + /// 推荐推理线程数:取物理核心的一半(保底 2)。 + /// llama.cpp 每个线程占一个 CPU 核心全速运算,默认用全部核心会让 + /// 推理占满整机 CPU,UI 与系统其他应用全部卡顿。 + private static func recommendedThreadCount() -> Int32 { + #if os(macOS) + var physical = 0 + var size = MemoryLayout.size + if sysctlbyname("hw.physicalcpu", &physical, &size, nil, 0) == 0, physical > 0 { + return Int32(max(2, physical / 2)) + } + #endif + let logical = ProcessInfo.processInfo.processorCount + return Int32(max(2, logical / 2)) + } + + /// macOS 当前可用内存(free + inactive 可回收页)。 + /// 注意:`os_proc_available_memory()` 是 iOS 专用 API,macOS 上不可用。 + #if os(macOS) + private static func availableMemoryBytes() -> UInt64 { + var stats = vm_statistics64() + var count = mach_msg_type_number_t( + MemoryLayout.size / MemoryLayout.size + ) + let result = withUnsafeMutablePointer(to: &stats) { + $0.withMemoryRebound(to: integer_t.self, capacity: Int(count)) { + host_statistics64(mach_host_self(), HOST_VM_INFO64, $0, &count) + } + } + guard result == KERN_SUCCESS else { + return ProcessInfo.processInfo.physicalMemory + } + let pageSize = UInt64(sysconf(Int32(_SC_PAGESIZE))) + return (UInt64(stats.free_count) + UInt64(stats.inactive_count)) * pageSize + } + + /// 本进程实际物理内存占用(phys_footprint,含 Metal/统一内存分配)。 + private static func processFootprintBytes() -> UInt64 { + var info = task_vm_info_data_t() + var count = mach_msg_type_number_t( + MemoryLayout.size / MemoryLayout.size + ) + let kr = withUnsafeMutablePointer(to: &info) { + $0.withMemoryRebound(to: integer_t.self, capacity: Int(count)) { + task_info(mach_task_self_, task_flavor_t(TASK_VM_INFO), $0, &count) + } + } + guard kr == KERN_SUCCESS else { return 0 } + return UInt64(info.phys_footprint) + } + + // CPU 采集暂不启用: + // 进程级采样(proc_pidinfo)本身正确(空闲 0%,忙时真实占用), + // 但当前无需在性能行展示 CPU,保留实现待后续需要时恢复。 + /* + private final class ProcessCPUSampleState: @unchecked Sendable { + let lock = NSLock() + var last: (cpuSeconds: Double, wall: CFAbsoluteTime)? + } + private static let processCPUSampleState = ProcessCPUSampleState() + + private static func processCPUPercent() -> Double { + var info = proc_taskinfo() + let size = MemoryLayout.size + let r = proc_pidinfo(getpid(), PROC_PIDTASKINFO, 0, &info, Int32(size)) + guard r == size else { return 0 } + // pti_total_user / pti_total_system 单位为微秒 + let cpuSeconds = Double(info.pti_total_user + info.pti_total_system) / 1_000_000 + processCPUSampleState.lock.lock() + defer { processCPUSampleState.lock.unlock() } + guard let last = processCPUSampleState.last, cpuSeconds >= last.cpuSeconds else { + processCPUSampleState.last = (cpuSeconds, CFAbsoluteTimeGetCurrent()) + return 0 + } + let now = CFAbsoluteTimeGetCurrent() + let dt = max(now - last.wall, 0.05) + let dCPU = cpuSeconds - last.cpuSeconds + processCPUSampleState.last = (cpuSeconds, now) + let ncpu = Double(ProcessInfo.processInfo.activeProcessorCount) + let pct = dCPU / dt * 100 / ncpu + return min(max(pct, 0), 100) + } + */ + #endif + + /// 本地模型性能快照:权重/参数/上下文/KV cache/进程内存/系统可用/CPU/速度。 + public struct LlamaPerformanceSnapshot: Sendable { + /// GGUF 模型权重字节数(llama_model_size) + public let modelSizeBytes: UInt64 + /// 模型参数量 + public let parameterCount: UInt64 + /// 当前上下文已用 token 数 + public let contextTokens: UInt32 + /// KV cache 估算字节(context × 层 × 维度 × 4,f16) + public let kvCacheBytes: UInt64 + /// 本进程实际物理内存占用(phys_footprint,统一内存含 GPU 分配) + public let processFootprintBytes: UInt64 + /// 系统当前可用内存 + public let availableSystemBytes: UInt64 + /// 本进程 CPU 使用率(归一化到单核,0-100%) + public let systemCPUPercent: Double + /// 最近一次生成速度(token/秒) + public let tokensPerSecond: Double + /// Metal 卸载层数(-1 = 全部卸载到 GPU) + public let gpuLayers: Int32 + } + + /// 本地模型性能快照(权重/参数/KV cache/进程内存/系统可用/CPU/速度)。 + /// 模型未加载或非 macOS 返回 nil。 + public func performanceSnapshot() -> LlamaPerformanceSnapshot? { + guard isModelLoaded, let model else { return nil } + #if os(macOS) + let size = llama_model_size(model) + let params = llama_model_n_params(model) + let nLayer = Int32(llama_model_n_layer(model)) + let nEmbd = Int32(llama_model_n_embd(model)) + var ctxTokens: UInt32 = 0 + var kvBytes: UInt64 = 0 + sessionContextLock.lock() + if let cached = cachedSessionContext?.context { + ctxTokens = llama_n_ctx(cached) + // KV cache ≈ context tokens × 层数 × 维度 × 2(K+V) × 2(f16) + kvBytes = UInt64(ctxTokens) * UInt64(max(nLayer, 1)) * UInt64(max(nEmbd, 1)) * 4 + } + sessionContextLock.unlock() + statsLock.lock() + let tps = lastTokensPerSecond + statsLock.unlock() + return LlamaPerformanceSnapshot( + modelSizeBytes: size, + parameterCount: params, + contextTokens: ctxTokens, + kvCacheBytes: kvBytes, + processFootprintBytes: Self.processFootprintBytes(), + availableSystemBytes: Self.availableMemoryBytes(), + systemCPUPercent: 0, // CPU 采集暂不启用(processCPUPercent 已注释) + tokensPerSecond: tps, + gpuLayers: gpuLayers + ) + #else + return nil + #endif + } + private func createModelParams() -> llama_model_params { var params = llama_model_default_params() params.n_gpu_layers = gpuLayers @@ -1317,6 +1747,10 @@ import Foundation if let temp = options.temperature { base.temperature = Float(temp) } + // 模型原生 context 自动接管默认(用户显式设置仍优先) + if let native = nativeContextSize, native > 0 { + base.contextSize = native + } return ResolvedGenerationOptions( base: base, @@ -1337,6 +1771,10 @@ import Foundation base.topP = 0.95 base.repeatPenalty = 1.1 base.repeatLastN = 64 + // 模型原生 context 自动接管默认(用户显式设置仍优先) + if let native = nativeContextSize, native > 0 { + base.contextSize = native + } return ResolvedGenerationOptions( base: base, @@ -1453,6 +1891,32 @@ import Foundation schema: GenerationSchema, maxTokens: Int, options: ResolvedGenerationOptions + ) async throws -> (String, LanguageModelSession.Usage) { + // 推理串行门(async 安全):结构化生成与文本生成共享同一模型/tokenCache, + // 经 actor 串行执行,避免并发访问 tokenCache / 竞争模型状态。 + await structuredGate.acquire() + do { + let result = try await self.performStructuredGeneration( + context: context, + prompt: prompt, + schema: schema, + maxTokens: maxTokens, + options: options + ) + await structuredGate.release() + return result + } catch { + await structuredGate.release() + throw error + } + } + + private func performStructuredGeneration( + context: OpaquePointer, + prompt: String, + schema: GenerationSchema, + maxTokens: Int, + options: ResolvedGenerationOptions ) async throws -> (String, LanguageModelSession.Usage) { guard let vocab = llama_model_get_vocab(model!) else { throw LlamaLanguageModelError.contextInitializationFailed @@ -1513,7 +1977,8 @@ import Foundation position: initialPosition, maximumTokens: maxTokens, endTokens: [], - tokenToTextFn: { [self] token in self.tokenToText(vocab: vocab, token: llama_token(token)) } + tokenToTextFn: { [self] token in self.tokenToText(vocab: vocab, token: llama_token(token)) }, + memoryCheckFn: { [self] in try self.checkRuntimeMemoryPressure() } ) var generator = try ConstrainedJSONGenerator(backend: backend, schema: schema, tokenCache: tokenCache) let json = try await generator.generate() @@ -1535,6 +2000,10 @@ import Foundation let tokenToTextFn: (Int) -> String? let tokensExcludedFromRepetitionPenalty: Set let endTokens: Set + /// 运行时内存护栏(结构化生成路径每 64 token 调用一次) + let memoryCheckFn: (() throws -> Void)? + /// decode 调用计数(用于按频率触发内存检查) + private var decodeCount = 0 var position: Int32 var remainingTokens: Int @@ -1550,7 +2019,8 @@ import Foundation position: Int32, maximumTokens: Int, endTokens: Set? = nil, - tokenToTextFn: @escaping (Int) -> String? + tokenToTextFn: @escaping (Int) -> String?, + memoryCheckFn: (() throws -> Void)? = nil ) { self.context = context self.vocab = vocab @@ -1561,6 +2031,7 @@ import Foundation self.remainingTokens = maximumTokens self.totalTokenBudget = maximumTokens self.eosToken = Int(llama_vocab_eos(vocab)) + self.memoryCheckFn = memoryCheckFn if let endTokens { self.endTokens = endTokens @@ -1625,6 +2096,12 @@ import Foundation } mutating func decode(_ token: Int) async throws { + decodeCount += 1 + // 结构化生成路径的运行时内存护栏(每 64 token 检查一次) + if decodeCount % 64 == 0, let memoryCheckFn { + try memoryCheckFn() + } + let llamaToken = llama_token(token) batch.pointee.n_tokens = 1 @@ -1693,6 +2170,19 @@ import Foundation options: ResolvedGenerationOptions, onToken: (String) -> Bool ) throws -> LanguageModelSession.Usage { + // 推理串行门:同一实例同时只允许一个生成,避免共享 context/batch 竞争 + inferenceGate.wait() + defer { inferenceGate.signal() } + + let generationStart = Date() + var generatedTokenCount = 0 + defer { + statsLock.lock() + let elapsed = Date().timeIntervalSince(generationStart) + lastTokensPerSecond = elapsed > 0 ? Double(generatedTokenCount) / elapsed : 0 + statsLock.unlock() + } + guard let model = self.model else { throw LlamaLanguageModelError.modelLoadFailed } @@ -1758,7 +2248,12 @@ import Foundation ) var decodedTokens = promptTokens - for _ in 0 ..< maxTokens { + for tokenIndex in 0 ..< maxTokens { + // 运行时内存护栏:KV cache 随上下文增长,内存动态变化, + // 每 64 token 检查一次系统可用内存,吃紧即中断(而非等 abort 干死机) + if tokenIndex % 64 == 0 { + try checkRuntimeMemoryPressure() + } if Task.isCancelled { break } @@ -1774,6 +2269,11 @@ import Foundation // Convert token to text and yield it if let tokenText = tokenToText(vocab: vocab, token: nextToken) { + // EOG / 会话结构标签兜底(GGUF 缺 eos_token_id 时 + // llama_vocab_is_eog 失效,按文本拦截) + if isTerminalTokenText(tokenText) { + break + } guard onToken(tokenText) else { break } @@ -1796,6 +2296,7 @@ import Foundation break } decodedTokens.append(nextToken) + generatedTokenCount += 1 } recordCachedTokens(decodedTokens, context: context) @@ -1820,6 +2321,10 @@ import Foundation options: ResolvedGenerationOptions, onToken: (String) -> Bool ) throws -> LanguageModelSession.Usage { + // 推理串行门:多模态生成与文本/JSON 生成共享同一模型实例,必须串行 + inferenceGate.wait() + defer { inferenceGate.signal() } + guard let mtmdContext, let model = self.model, let vocab = llama_model_get_vocab(model) else { @@ -1941,6 +2446,11 @@ import Foundation } if let tokenText = tokenToText(vocab: vocab, token: nextToken) { + // EOG / 会话结构标签兜底(GGUF 缺 eos_token_id 时 + // llama_vocab_is_eog 失效,按文本拦截) + if isTerminalTokenText(tokenText) { + break + } guard onToken(tokenText) else { break } @@ -2022,7 +2532,7 @@ import Foundation // Encoder models ingest the full prompt in a single llama_encode call. guard !hasEncoder || (startIndex == 0 && promptTokens.count <= batchSize) else { - throw LlamaLanguageModelError.insufficientMemory + throw LlamaLanguageModelError.insufficientMemory(requiredBytes: 0, availableBytes: 0) } if hasEncoder { @@ -2397,6 +2907,28 @@ import Foundation return Array(UnsafeBufferPointer(start: tokens, count: Int(tokenCount))) } + /// 生成过程中的终止 token 文本(模板 EOG / 会话结构标签)。 + /// 部分 GGUF 未声明 `eos_token_id`,`llama_vocab_is_eog` 识别不到 + /// `<|im_end|>` 之类的模板 EOG,导致其被解码成 ◆ 混入正文。 + /// 此处按文本兜底识别:命中即结束生成且不 yield。 + private static let terminalTokenTexts: Set = [ + "<|im_end|>", "<|im_start|>", "<|endoftext|>", "<|eot_id|>", + "<|end|>", "", "", "<|tool_call|>", "<|tool_response|>", + ] + + private func isTerminalTokenText(_ text: String) -> Bool { + let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) + return Self.terminalTokenTexts.contains(trimmed) + } + + /// 输出 token 上限:用户显式设置时取较小者;未设置时按模型 + /// 原生上下文自动推导(预留 prompt 余量),不再写死固定值。 + private static func resolveMaxTokens(requested: Int?, contextSize: UInt32) -> Int { + let cap = max(256, Int(contextSize) - 512) + guard let requested else { return cap } + return min(requested, cap) + } + private func tokenToText(vocab: OpaquePointer, token: llama_token) -> String? { // First attempt with a reasonable buffer var cap: Int32 = 64 @@ -2409,7 +2941,7 @@ import Foundation buf, cap, 0, - false + true ) if written < 0 { @@ -2423,7 +2955,7 @@ import Foundation buf, cap, 0, - false + true ) } @@ -2450,7 +2982,7 @@ import Foundation case encodingFailed case decodingFailed case invalidModelPath - case insufficientMemory + case insufficientMemory(requiredBytes: UInt64, availableBytes: UInt64) case promptExceedsContextWindow case unsupportedFeature case encoderOnlyModel @@ -2469,7 +3001,13 @@ import Foundation return "Failed to decode response" case .invalidModelPath: return "Invalid model file path" - case .insufficientMemory: + case let .insufficientMemory(requiredBytes, availableBytes): + if requiredBytes > 0 { + let gb = 1_073_741_824.0 + let req = String(format: "%.1f", Double(requiredBytes) / gb) + let avail = String(format: "%.1f", Double(availableBytes) / gb) + return "Insufficient memory to load model (needs ~\(req) GB, \(avail) GB available). Close other apps or use a smaller model." + } return "Insufficient memory for operation" case .promptExceedsContextWindow: return "Prompt is longer than the model's context window" diff --git a/Sources/AnyLanguageModel/Models/LlamaToolCallFormat.swift b/Sources/AnyLanguageModel/Models/LlamaToolCallFormat.swift index a1441b9b..c9593843 100644 --- a/Sources/AnyLanguageModel/Models/LlamaToolCallFormat.swift +++ b/Sources/AnyLanguageModel/Models/LlamaToolCallFormat.swift @@ -57,6 +57,51 @@ enum LlamaToolCallFormat: Sendable, Equatable { /// The marker that closes a Gemma 4 thought channel. static let gemmaChannelCloseMarker = "" + /// 常见思考段标签对(open/close)。不同模型的思考标记不同,统一在此剥离, + /// 避免"一个模型一套特判": + /// - Qwen3: `...` + /// - Qwen2.x 思考微调 / 旧版: `<|startofthink|>...` + /// - Llama 4: `<|begin_of_thought|>...` + /// - 其他常见格式: `...`、`...` + static let thoughtSpanPairs: [(open: String, close: String)] = [ + ("", ""), + ("<|startofthink|>", ""), + ("<|begin_of_thought|>", ""), + ("", ""), + ("", ""), + ] + + /// 从生成文本中剥离所有思考段(通用标签对)。 + /// 未闭合的思考段从开始处截断到文本末尾(同 Gemma channel 的处理)。 + /// - Parameter pairs: 模型动态检测出的标签对;nil 使用静态全集。 + static func stripThoughtSpans( + from text: String, + pairs: [(open: String, close: String)]? = nil + ) -> String { + var result = text + let spanPairs = pairs ?? thoughtSpanPairs + for pair in spanPairs { + result = stripSpans(open: pair.open, close: pair.close, from: result) + } + return result + } + + private static func stripSpans(open: String, close: String, from text: String) -> String { + var result = "" + var remainder = Substring(text) + while let openRange = remainder.range(of: open) { + result += remainder[..`, but the model volunteers thought channels /// anyway; the canonical template ships a `strip_thinking` macro for the @@ -101,12 +146,16 @@ enum LlamaToolCallFormat: Sendable, Equatable { func streamingVisibleText( in raw: String, withholdToolCalls: Bool, - holdPartialMarkers: Bool = true + holdPartialMarkers: Bool = true, + thoughtPairs: [(open: String, close: String)]? = nil ) -> String { var text = raw + let spanPairs = thoughtPairs ?? Self.thoughtSpanPairs if self == .gemma { text = Self.stripGemmaThoughtChannels(from: text) } + // 通用思考段剥离(动态检测标签对优先,兜底静态全集) + text = Self.stripThoughtSpans(from: text, pairs: spanPairs) if withholdToolCalls, let range = text.range(of: callStartMarker) { text = String(text[.. [String: JSON ] } +/// Ollama 没有 tool-choice 设置,模型可能选择不调用工具。 +/// tools 非空时在 system 中注入引导消息,保证实时数据类提问会调用工具。 +private func injectingToolGuidance(into messages: [OllamaMessage]) -> [OllamaMessage] { + let guidance = + "当用户需要实时或外部信息(当前时间、IP地址、地理位置、天气、汇率、域名信息、网页内容等)时," + + "你必须调用提供的工具获取真实结果,不得凭记忆或猜测回答。" + guard let first = messages.first, first.role == .system else { + return [OllamaMessage(role: .system, content: guidance)] + messages + } + var result = messages + result.insert(OllamaMessage(role: .system, content: guidance), at: 1) + return result +} + private func encodeChatParams(_ params: [String: JSONValue]) throws -> Data { let encoder = JSONEncoder() // Ollama reuses prompt prefixes only when their serialized bytes match. diff --git a/Sources/AnyLanguageModel/Transcript.swift b/Sources/AnyLanguageModel/Transcript.swift index 83d12b20..f194d7ba 100644 --- a/Sources/AnyLanguageModel/Transcript.swift +++ b/Sources/AnyLanguageModel/Transcript.swift @@ -840,6 +840,8 @@ extension Transcript.Reasoning { enum Origin: String, Sendable, Equatable, CaseIterable { case anthropic case ollama + /// Llama (llama.cpp / GGUF) models with embedded thinking spans. + case llama /// OpenAI's Chat Completions API, or another host that serves it. case chatCompletions = "openai.chatCompletions" } diff --git a/Tests/AnyLanguageModelTests/ThoughtParserTests.swift b/Tests/AnyLanguageModelTests/ThoughtParserTests.swift new file mode 100644 index 00000000..4dc502bf --- /dev/null +++ b/Tests/AnyLanguageModelTests/ThoughtParserTests.swift @@ -0,0 +1,90 @@ +import XCTest +#if Llama + @testable import AnyLanguageModel + + final class ThoughtParserTests: XCTestCase { + private func parseAll(_ chunks: [String]) -> String { + var parser = ThoughtParser() + for chunk in chunks { + parser.consume(chunk) + } + return parser.thinking + } + + func testThinkSingleChunk() { + let thinking = parseAll(["先想:9.9 更大,所以答案是 9.9"]) + XCTAssertEqual(thinking, "9.9 更大") + } + + func testStartOfThinkPair() { + let thinking = parseAll(["旧版:<|startofthink|>内部推理,结束"]) + XCTAssertEqual(thinking, "内部推理") + } + + func testBeginOfThoughtPair() { + let thinking = parseAll(["Llama4:<|begin_of_thought|>推理 A,结束"]) + XCTAssertEqual(thinking, "推理 A") + } + + func testReasoningPair() { + let thinking = parseAll(["步骤1答案"]) + XCTAssertEqual(thinking, "步骤1") + } + + func testThinkingPair() { + let thinking = parseAll(["思考中正文"]) + XCTAssertEqual(thinking, "思考中") + } + + func testOpenTagAcrossChunks() { + let thinking = parseAll(["跨块思考", "正文"]) + XCTAssertEqual(thinking, "跨块思考") + } + + func testCloseTagAcrossChunks() { + let thinking = parseAll(["思考", "内容正文"]) + XCTAssertEqual(thinking, "思考内容") + } + + func testLongThinkingStreamsWhileGenerating() { + // 长思考段:超过尾部保留窗口的内容在生成中即流出,末尾靠 finalize 补齐 + var parser = ThoughtParser() + parser.consume("这是比较长的思考内容超过保留窗口") + XCTAssertTrue(parser.thinking.contains("这是比较长的思考内"), "thinking=[\(parser.thinking)]") + parser.consume("继续思考") + parser.finalize() + XCTAssertEqual(parser.thinking, "这是比较长的思考内容超过保留窗口继续思考") + } + + func testUnclosedFlushedOnFinalize() { + // 生成结束:未闭合思考段完整 flush + var parser = ThoughtParser() + parser.consume("完整") + parser.consume("思考内容") + parser.finalize() + XCTAssertEqual(parser.thinking, "完整思考内容") + } + + func testPlainTextNoThinking() { + let thinking = parseAll(["没有思考标签的", "普通回复"]) + XCTAssertEqual(thinking, "") + } + + func testBodyTextIgnored() { + let parser = ThoughtParser() + var p = parser + p.consume("这是正文") + XCTAssertEqual(p.thinking, "") + } + + func testMultiplePairsSequential() { + let thinking = parseAll([ + "第一段", + "中间正文", + "<|startofthink|>第二段", + "结尾", + ]) + XCTAssertEqual(thinking, "第一段第二段") + } + } +#endif diff --git a/Tests/AnyLanguageModelTests/ThoughtSpanDetectionTests.swift b/Tests/AnyLanguageModelTests/ThoughtSpanDetectionTests.swift new file mode 100644 index 00000000..e65b5034 --- /dev/null +++ b/Tests/AnyLanguageModelTests/ThoughtSpanDetectionTests.swift @@ -0,0 +1,64 @@ +import XCTest +#if Llama + @testable import AnyLanguageModel + + final class ThoughtSpanDetectionTests: XCTestCase { + func testQwen3NewTemplate() { + let template = """ + {% if not is_tool_call %}{{'\n'}}{% endif %}<|im_start|>assistant + """ + let pairs = ThoughtParser.thoughtSpanPairs(fromTemplate: template) + XCTAssertEqual(pairs?.count, 1) + XCTAssertEqual(pairs?.first?.open, "") + XCTAssertEqual(pairs?.first?.close, "") + } + + func testQwen2OldStyleTemplate() { + let template = """ + <|im_start|>system ... <|startofthink|>{{ {content} }}<|endofthink|> + """ + let pairs = ThoughtParser.thoughtSpanPairs(fromTemplate: template) + XCTAssertEqual(pairs?.count, 1) + XCTAssertEqual(pairs?.first?.open, "<|startofthink|>") + XCTAssertEqual(pairs?.first?.close, "") + } + + func testLlama4Template() { + let template = """ + <|begin_of_thought|>{{{ content }}}<|end_of_thought|> + """ + let pairs = ThoughtParser.thoughtSpanPairs(fromTemplate: template) + XCTAssertEqual(pairs?.count, 1) + XCTAssertEqual(pairs?.first?.open, "<|begin_of_thought|>") + XCTAssertEqual(pairs?.first?.close, "") + } + + func testGenericPairDiscovery() { + // 模板含 ...,未在已知集合中显式出现于模板外 + let template = "{% if reasoning %}{{ content }}{% endif %}" + let pairs = ThoughtParser.thoughtSpanPairs(fromTemplate: template) + XCTAssertEqual(pairs?.count, 1) + XCTAssertEqual(pairs?.first?.open, "") + XCTAssertEqual(pairs?.first?.close, "") + } + + func testTemplateWithoutThinking() { + let template = "<|im_start|>user {{ content }}<|im_end|>" + let pairs = ThoughtParser.thoughtSpanPairs(fromTemplate: template) + XCTAssertNil(pairs) + } + + func testMultipleKnownPairs() { + let template = """ + {{ content }}{% if x %} + """ + let pairs = ThoughtParser.thoughtSpanPairs(fromTemplate: template) + XCTAssertEqual(pairs?.count, 2) + } + + func testDefaultPairsFallback() { + // 默认兜底全集应覆盖所有已知思考标签 + XCTAssertEqual(ThoughtParser.defaultPairs.count, 5) + } + } +#endif