diff --git a/README.md b/README.md index 9f694828..86a086fd 100644 --- a/README.md +++ b/README.md @@ -687,6 +687,7 @@ In addition to the standard OpenAI fields (`temperature`, `top_p`, `max_tokens`, |---|---|---| | `kv_bits` | `int` (4 or 8) | Enable **MLX-native quantized KV cache** for this request. Uses `QuantizedKVCache` (standard group quantization) instead of `KVCacheSimple`. Separate from `--turbo-kv`. Reduces KV memory ~2–4× at mild quality cost. | | `enable_thinking` | `bool` | Force-enable or disable chain-of-thought thinking blocks for Gemma-4 / Qwen3. | +| `chat_template_kwargs` | `object` | Passed to the chat template as sent, any JSON value: `{"thinking": true}` for DeepSeek V3.1, `{"reasoning_effort": "high"}` for gpt-oss. An `enable_thinking` inside it is overridden by the top-level field above. | | `kv_group_size` | `int` | Group size for `kv_bits` quantization (default: `64`). | | `top_k` | `int` | Per-request top-k sampling override (0 = disabled). | | `min_p` | `float` | Per-request min-p sampling threshold (0 = disabled). | diff --git a/Sources/MLXInferenceCore/HFModelSearch.swift b/Sources/MLXInferenceCore/HFModelSearch.swift index ba96de55..fb818edc 100644 --- a/Sources/MLXInferenceCore/HFModelSearch.swift +++ b/Sources/MLXInferenceCore/HFModelSearch.swift @@ -160,6 +160,12 @@ public final class HFModelSearchService: ObservableObject { private var currentSort = HFSortOption.trending private var currentSizeFilter = HFSizeFilter.all private var debounceTask: Task? = nil + /// Counts searches. A page fetch remembers the count it started under and publishes + /// nothing once a newer search has started: cancelling the debounce task does not stop + /// a fetch that is already awaiting the network, and a `loadMore()` fetch is not in + /// that task at all. Before this, a superseded fetch wrote "Search failed: cancelled" + /// and `isSearching = false` over the search that replaced it. + private var generation = 0 private init() {} @@ -168,6 +174,7 @@ public final class HFModelSearchService: ObservableObject { /// Debounced search — safe to call on every keystroke. public func search(query: String, sort: HFSortOption = .trending, sizeFilter: HFSizeFilter = .all) { debounceTask?.cancel() + generation += 1 debounceTask = Task { // 300ms debounce try? await Task.sleep(nanoseconds: 300_000_000) @@ -225,6 +232,10 @@ public final class HFModelSearchService: ObservableObject { private func fetchPage() async { print("HFSearch: fetchPage started. Query: '\(currentQuery)' Sort: \(currentSort.rawValue)") + let generation = self.generation + // True once a newer search has started (or this fetch's task was cancelled with + // it): the published state belongs to that search, so this one stops here. + func superseded() -> Bool { Task.isCancelled || generation != self.generation } isSearching = true errorMessage = nil @@ -233,7 +244,9 @@ public final class HFModelSearchService: ObservableObject { // and hit the model detail endpoint. This is needed because the HF // search API filters on pipeline_tag which many MLX uploads don't match. if isDirectRepoQuery(currentQuery) { - if let result = await fetchDirectRepo(currentQuery) { + let result = await fetchDirectRepo(currentQuery) + if superseded() { return } + if let result { results = [result] } else { // Also try with prefix variants (e.g. user typed partial name without quant suffix) @@ -283,6 +296,7 @@ public final class HFModelSearchService: ObservableObject { do { let (data, response) = try await URLSession.shared.data(from: urlToFetch) + if superseded() { return } guard let http = response as? HTTPURLResponse, http.statusCode == 200 else { errorMessage = "HuggingFace API unavailable" break @@ -335,6 +349,7 @@ public final class HFModelSearchService: ObservableObject { if let size = size { page[index].usedStorage = size } } } + if superseded() { return } localResults.append(contentsOf: page) } @@ -345,17 +360,22 @@ public final class HFModelSearchService: ObservableObject { hasMore = true } } catch is CancellationError { - break + return + } catch let error as URLError where error.code == .cancelled { + // The request was cancelled with this fetch's task: a newer search's doing. + return } catch { + if superseded() { return } errorMessage = "Search failed: \(error.localizedDescription)" break } } + if superseded() { return } if !localResults.isEmpty { results.append(contentsOf: localResults) } - + isSearching = false print("HFSearch: fetchPage finished") } diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index a4f818ff..ec93a94c 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -2516,6 +2516,18 @@ func handleChatCompletion( ) } + // ── chat_template_kwargs is free-form JSON, but enable_thinking stays a boolean ── + // It decides the precedence below, so a string or a number there is a client error, + // answered the way the whole field used to be answered before it took other types. + if let raw = chatReq.chatTemplateKwargs?["enable_thinking"], !(raw.value is Bool) { + let errBody = "{\"error\":{\"message\":\"Invalid type for 'chat_template_kwargs.enable_thinking': expected a boolean.\",\"type\":\"invalid_request_error\",\"param\":\"chat_template_kwargs.enable_thinking\"}}" + return Response( + status: .badRequest, + headers: jsonHeaders(), + body: .init(byteBuffer: ByteBuffer(string: errBody)) + ) + } + var generateParams = GenerateParameters( maxTokens: tokenLimit, maxKVSize: config.ctxSize, @@ -2583,7 +2595,7 @@ func handleChatCompletion( var enableThinking: Bool if let explicitTopLevel = chatReq.enableThinking { enableThinking = explicitTopLevel - } else if let kwargs = chatReq.chatTemplateKwargs, let perRequest = kwargs["enable_thinking"] { + } else if let perRequest = chatReq.chatTemplateKwargs?["enable_thinking"]?.value as? Bool { enableThinking = perRequest // per-request override wins } else { enableThinking = config.thinking // fall back to server --thinking flag @@ -2617,7 +2629,8 @@ func handleChatCompletion( // The Jinja template evaluates `not enable_thinking | default(false)`. If we pass nil instead of // true, it evaluates to false and still breaks. We MUST explicitly pass the boolean. - let templateContext: [String: any Sendable] = ["enable_thinking": enableThinking] + let templateContext = chatTemplateContext( + kwargs: chatReq.chatTemplateKwargs, enableThinking: enableThinking) let userInput = UserInput(chat: chatMessages, tools: toolSpecs, additionalContext: templateContext) print("[Server Debug] Created UserInput with \(userInput.images.count) images and \(userInput.audios.count) audio inputs.") let lmInput = try await container.prepare(input: userInput) @@ -4419,6 +4432,18 @@ func checkStopSequences(_ text: String, stopSequences: [String], lookback: Int? return (String(text[text.startIndex.. [String: any Sendable] { + var context = AnyCodable.toSendable(kwargs) ?? [:] + context["enable_thinking"] = enableThinking + return context +} + /// The sampler seed for a request `seed`. OpenAI's `seed` is a plain integer, so negative /// values do arrive, and `UInt64(seed)` trapped on them — one such request took the whole /// server down (#219). Taking the bit pattern keeps every distinct Int a distinct seed. @@ -4920,8 +4945,11 @@ struct ChatCompletionRequest: Decodable { let seed: Int? let streamOptions: StreamOptions? let responseFormat: ResponseFormat? - /// Per-request Jinja template kwargs (e.g. {"enable_thinking": false} for Qwen3/Qwen3.5) - let chatTemplateKwargs: [String: Bool]? + /// Per-request Jinja template kwargs, passed to the chat template as sent: any JSON + /// value, since templates read more than booleans (`{"enable_thinking": false}` for + /// Qwen3, `{"thinking": true}` for DeepSeek V3.1, `{"reasoning_effort": "high"}` for + /// gpt-oss). `enable_thinking` keeps its own precedence rules in `handleChatCompletion`. + let chatTemplateKwargs: [String: AnyCodable]? /// Per-request reasoning-token cap; overrides `--thinking-budget`. let thinkingBudget: Int? /// Top-level thinking override emitted by Aegis-AI gateway diff --git a/docs/api-response-formats.md b/docs/api-response-formats.md index 14932c80..cd08740d 100644 --- a/docs/api-response-formats.md +++ b/docs/api-response-formats.md @@ -185,6 +185,7 @@ for await (const chunk of stream) { | `response_format` | object | none | `{"type": "json_object"}` for JSON mode | | `tools` | array | none | Tool/function definitions | | `enable_thinking` | bool | false | Enable `` reasoning tokens | +| `chat_template_kwargs` | object | none | Passed to the chat template as sent (e.g. `{"thinking": true}`, `{"reasoning_effort": "high"}`) | --- diff --git a/tests/SwiftBuddyTests/HFModelSearchTests.swift b/tests/SwiftBuddyTests/HFModelSearchTests.swift index ca9186d1..c2577db6 100644 --- a/tests/SwiftBuddyTests/HFModelSearchTests.swift +++ b/tests/SwiftBuddyTests/HFModelSearchTests.swift @@ -50,6 +50,32 @@ final class HFModelSearchTests: XCTestCase { XCTAssertNil(service.errorMessage) } + /// A search that replaces one still awaiting the network must not inherit its state. + /// The replaced fetch's request is cancelled, which it used to report as + /// "Search failed: cancelled" over the new search, and it also cleared `isSearching` + /// while the new search was still running. That is the flake CI hit in + /// testStrictMLXFilterDisabled whenever an earlier test left a fetch in flight. + @MainActor + func testSupersededSearchPublishesNothing() async throws { + let service = HFModelSearchService.shared + service.errorMessage = nil + service.results = [] + service.strictMLX = true + + service.search(query: "gpt", sort: .trending) + // Past the 300 ms debounce, so the first fetch is on the wire when it is replaced. + try? await Task.sleep(nanoseconds: 450_000_000) + service.search(query: "mistral", sort: .trending) + + for _ in 0..<40 { + try? await Task.sleep(nanoseconds: 100_000_000) + if !service.isSearching && service.results.count > 0 { break } + } + + try XCTSkipIf(service.isSearching || service.results.count == 0, "Skipping due to transient HF rate limiting on GitHub Actions IP bounds") + XCTAssertNil(service.errorMessage, "the replaced search's cancellation leaked into the new search") + } + // Feature 3: Empty query trending @MainActor func testFeature3_EmptyQueryTrending() async throws { diff --git a/tests/SwiftLMTests/ChatRequestParsingTests.swift b/tests/SwiftLMTests/ChatRequestParsingTests.swift index cc7b1c06..b7eed6c5 100644 --- a/tests/SwiftLMTests/ChatRequestParsingTests.swift +++ b/tests/SwiftLMTests/ChatRequestParsingTests.swift @@ -322,4 +322,62 @@ final class ChatRequestParsingTests: XCTestCase { XCTAssertEqual(req.maxCompletionTokens, 64) XCTAssertNil(req.maxTokens) } + + // MARK: - chat_template_kwargs is passed to the template as sent (#259) + + /// The field decoded as `[String: Bool]`, so a string or a number anywhere in it failed + /// the whole request, and only `enable_thinking` was read. Templates take other names + /// and types: DeepSeek-V3.1 `thinking`, gpt-oss `reasoning_effort` and `model_identity`. + func testChatTemplateKwargsDecodeAnyJSONValue() throws { + let req = try decode(""" + { + "messages": [{"role": "user", "content": "hi"}], + "chat_template_kwargs": { + "enable_thinking": false, + "thinking": true, + "reasoning_effort": "high", + "budget": 3, + "nested": {"a": [1, "b", null]} + } + } + """) + let kwargs = try XCTUnwrap(req.chatTemplateKwargs) + XCTAssertEqual(kwargs["enable_thinking"]?.value as? Bool, false) + XCTAssertEqual(kwargs["thinking"]?.value as? Bool, true) + XCTAssertEqual(kwargs["reasoning_effort"]?.value as? String, "high") + XCTAssertEqual(kwargs["budget"]?.value as? Int, 3) + let nested = try XCTUnwrap(kwargs["nested"]?.value as? [String: Any]) + let list = try XCTUnwrap(nested["a"] as? [Any]) + XCTAssertEqual(list.count, 3) + XCTAssertEqual(list[0] as? Int, 1) + XCTAssertEqual(list[1] as? String, "b") + XCTAssertTrue(list[2] is NSNull) + } + + func testChatTemplateKwargsAbsentDecodesToNil() throws { + let req = try decode(#"{"messages":[{"role":"user","content":"hi"}]}"#) + XCTAssertNil(req.chatTemplateKwargs) + } + + /// Every kwarg reaches the template context; `enable_thinking` carries the verdict of + /// the precedence rules in `handleChatCompletion`, not the raw kwarg. + func testChatTemplateContextForwardsEveryKwargAndSetsEnableThinking() throws { + let req = try decode(""" + { + "messages": [{"role": "user", "content": "hi"}], + "chat_template_kwargs": {"enable_thinking": true, "thinking": true, "reasoning_effort": "low"} + } + """) + let context = chatTemplateContext(kwargs: req.chatTemplateKwargs, enableThinking: false) + XCTAssertEqual(context["enable_thinking"] as? Bool, false) + XCTAssertEqual(context["thinking"] as? Bool, true) + XCTAssertEqual(context["reasoning_effort"] as? String, "low") + XCTAssertEqual(context.count, 3) + } + + func testChatTemplateContextWithoutKwargsHasOnlyEnableThinking() { + let context = chatTemplateContext(kwargs: nil, enableThinking: true) + XCTAssertEqual(context["enable_thinking"] as? Bool, true) + XCTAssertEqual(context.count, 1) + } } diff --git a/tests/test-contract.sh b/tests/test-contract.sh index f8023a2c..1e69c978 100755 --- a/tests/test-contract.sh +++ b/tests/test-contract.sh @@ -553,6 +553,31 @@ PY early_exit_then_follow_up "stop sequence" "stop-seq" "the" "" early_exit_then_follow_up "client disconnect" "disconnect" "" "leave" +# ── 17. chat_template_kwargs is passed to the template as sent ────────────── +# #259: the field decoded as [String: Bool], so a string or a number anywhere in it was a +# 400, and every key but enable_thinking was dropped before the template saw it. Templates +# read other names (DeepSeek-V3.1 `thinking`, gpt-oss `reasoning_effort`). The CI model's +# template ignores names it does not know, so this checks the wire contract only: any JSON +# value is accepted, and a non-boolean enable_thinking is still refused by name. +log "Test 17: chat_template_kwargs takes any JSON value" +KW=$(chat '{"model":"x","messages":[{"role":"user","content":"Say hello in one short sentence."}],"max_tokens":16,"temperature":0,"chat_template_kwargs":{"reasoning_effort":"low","budget":3,"thinking":true,"nested":{"a":[1,"b"]}}}' || true) +KW_CONTENT=$(echo "$KW" | python3 -c 'import json,sys; print(json.load(sys.stdin)["choices"][0]["message"].get("content") or "")' 2>/dev/null || echo "__no_body__") +if [ "$KW_CONTENT" = "__no_body__" ] || [ -z "$KW_CONTENT" ]; then + fail "string, number and nested kwargs were rejected: $(echo "$KW" | head -c 160)" +else + pass "string, number, boolean and nested kwargs are accepted" +fi + +log "Test 17b: a non-boolean enable_thinking is a 400 that names the param" +CODE=$(curl -s -o /tmp/SwiftLM-contract-kwargs.json -w '%{http_code}' "$URL/v1/chat/completions" \ + -H 'Content-Type: application/json' -d '{"model":"x","messages":[{"role":"user","content":"hi"}],"max_tokens":4,"chat_template_kwargs":{"enable_thinking":"yes"}}') +ERR_TYPE=$(python3 -c 'import json; e=json.load(open("/tmp/SwiftLM-contract-kwargs.json"))["error"]; print(e.get("type"), e.get("param"))' 2>/dev/null || echo "unparseable") +if [ "$CODE" = "400" ] && [ "$ERR_TYPE" = "invalid_request_error chat_template_kwargs.enable_thinking" ]; then + pass "400 invalid_request_error naming param 'chat_template_kwargs.enable_thinking'" +else + fail "got HTTP $CODE with error '$ERR_TYPE' (expected 400 invalid_request_error chat_template_kwargs.enable_thinking)" +fi + log "═══════════════════════════════════════" log "Results: $PASS passed, $FAIL failed, $SKIP skipped, $TOTAL total" log "═══════════════════════════════════════"