diff --git a/lua/claude_handler.lua b/lua/claude_handler.lua index 8531a1c..4f144ad 100644 --- a/lua/claude_handler.lua +++ b/lua/claude_handler.lua @@ -49,6 +49,10 @@ function _M.handle() -- Translate Claude request to OpenAI format local oai_request = claude_fmt.request_to_openai(claude_body) + + -- Thinking mode handling (operates on OpenAI request format) + oai_request, model = e2ee.handle_thinking(oai_request) + local oai_body = cjson.encode(oai_request) if not is_streaming then diff --git a/lua/e2ee_discovery.lua b/lua/e2ee_discovery.lua index 8ee8d71..832fb03 100644 --- a/lua/e2ee_discovery.lua +++ b/lua/e2ee_discovery.lua @@ -99,11 +99,18 @@ function _M.resolve_chute_id(model, api_key) return model end + -- Strip the :THINKING suffix before map lookup. v1/models returns + -- base model IDs only; :THINKING is a proxy-side flag and must not + -- leak into the lookup key. Other suffixes (e.g. LoRA names) are + -- left intact so they surface as proper errors upstream rather than + -- silently resolving to the base chute. + local base_model = model:match("^(.-):THINKING$") or model + local now = ngx.now() -- Check cache if model_map and now < model_map_expires then - local entry = model_map[model] + local entry = model_map[base_model] if entry then return check_confidential(model, entry) end @@ -114,7 +121,7 @@ function _M.resolve_chute_id(model, api_key) if not map then -- If we have a stale cache, try it if model_map then - local entry = model_map[model] + local entry = model_map[base_model] if entry then return check_confidential(model, entry) end @@ -125,7 +132,7 @@ function _M.resolve_chute_id(model, api_key) model_map = map model_map_expires = now + MODEL_MAP_TTL - local entry = map[model] + local entry = map[base_model] return check_confidential(model, entry) end diff --git a/lua/e2ee_handler.lua b/lua/e2ee_handler.lua index 71938e5..b605c9f 100644 --- a/lua/e2ee_handler.lua +++ b/lua/e2ee_handler.lua @@ -17,6 +17,31 @@ local API_BASE = "https://api.chutes.ai" local _M = {} +-- Per-model thinking defaults. Mirrors chutes-api's invocation router +-- (api/invocation/router.py). The /v1/models endpoint exposes a +-- "reasoning" capability flag in supported_features but not the +-- per-model default value, so these stay duplicated until the API +-- grows a dedicated attribute. +local THINKING_DEFAULT_ON_PREFIXES = { + "deepseek-ai/DeepSeek-V3.2-Speciale", + "zai-org/GLM-4.7", + "moonshotai/Kimi-K2.5", +} +local THINKING_DEFAULT_ON_EXACT = { + ["deepseek-ai/DeepSeek-V3.2-Speciale"] = true, + ["deepseek-ai/DeepSeek-V3.2-Speciale-TEE"] = true, + ["zai-org/GLM-4.7"] = true, + ["zai-org/GLM-4.7-TEE"] = true, + ["moonshotai/Kimi-K2.5"] = true, + ["moonshotai/Kimi-K2.5-TEE"] = true, +} +local MIMO_PREFIX = "XiaomiMiMo/MiMo-V2-Flash" +local MIMO_EXACT_TEE = "XiaomiMiMo/MiMo-V2-Flash-TEE" + +local function starts_with(s, prefix) + return s:sub(1, #prefix) == prefix +end + --- Extract API key from Authorization header or x-api-key function _M.get_api_key() local headers = ngx.req.get_headers() @@ -40,6 +65,68 @@ function _M.get_api_key() return key end +--- Handle thinking mode for OpenAI format requests +-- Takes an OpenAI request table, applies thinking logic, and returns the modified table +-- Also strips :THINKING suffix from model name and returns the model name +function _M.handle_thinking(oai_request) + if not oai_request or not oai_request.model then + return oai_request, nil + end + + local model = oai_request.model + local enable_thinking = false + local think_header = ngx.req.get_headers()["X-Enable-Thinking"] + if think_header and think_header:lower() == "true" then + enable_thinking = true + end + if model:match(":THINKING$") then + model = model:sub(1, -#":THINKING" - 1) + oai_request.model = model + enable_thinking = true + end + if enable_thinking then + if not oai_request.chat_template_kwargs then + oai_request.chat_template_kwargs = {} + end + oai_request.chat_template_kwargs.thinking = true + oai_request.chat_template_kwargs.enable_thinking = true + end + + local kwargs = oai_request.chat_template_kwargs + if kwargs then + -- Normalize the two spellings so downstream sees both. + if kwargs.thinking ~= nil and kwargs.enable_thinking == nil then + kwargs.enable_thinking = kwargs.thinking + end + if kwargs.enable_thinking ~= nil and kwargs.thinking == nil then + kwargs.thinking = kwargs.enable_thinking + end + + if kwargs.thinking == nil then + for _, prefix in ipairs(THINKING_DEFAULT_ON_PREFIXES) do + if starts_with(model, prefix) then + kwargs.thinking = true + kwargs.enable_thinking = true + break + end + end + end + + if kwargs.thinking == nil and starts_with(model, MIMO_PREFIX) then + kwargs.thinking = false + kwargs.enable_thinking = false + end + else + if THINKING_DEFAULT_ON_EXACT[model] then + oai_request.chat_template_kwargs = { thinking = true, enable_thinking = true } + elseif model == MIMO_EXACT_TEE then + oai_request.chat_template_kwargs = { thinking = false, enable_thinking = false } + end + end + + return oai_request, model +end + --- Send error response function _M.send_error(status, message) ngx.status = status @@ -345,6 +432,12 @@ function _M.handle() return _M.send_error(400, "missing 'model' field") end + -- Thinking mode handling + payload, model = _M.handle_thinking(payload) + + -- Re-encode body with any chat_template_kwargs modifications + body = cjson.encode(payload) + local is_streaming = (payload.stream == true) local original_path = ngx.var.uri diff --git a/lua/responses_handler.lua b/lua/responses_handler.lua index 10bf166..02d3a71 100644 --- a/lua/responses_handler.lua +++ b/lua/responses_handler.lua @@ -49,6 +49,10 @@ function _M.handle() -- Translate Responses request to OpenAI format local oai_request = resp_fmt.request_to_openai(resp_body) + + -- Thinking mode handling (operates on OpenAI request format) + oai_request, model = e2ee.handle_thinking(oai_request) + local oai_body = cjson.encode(oai_request) if not is_streaming then diff --git a/tests/test_thinking.lua b/tests/test_thinking.lua new file mode 100644 index 0000000..04d6dc8 --- /dev/null +++ b/tests/test_thinking.lua @@ -0,0 +1,271 @@ +-- +-- Unit tests for e2ee_handler.handle_thinking +-- +-- Runs the :THINKING suffix / X-Enable-Thinking header logic in isolation +-- by stubbing the openresty globals (ngx, cjson.safe, resty.http) and the +-- sibling lua modules that e2ee_handler requires but doesn't exercise here. +-- +-- Usage: +-- luajit tests/test_thinking.lua +-- +-- Exits non-zero on any assertion failure. +-- + +-- Make lua/ importable regardless of cwd. +local script_path = arg and arg[0] or "" +local script_dir = script_path:match("(.*/)") or "./" +package.path = script_dir .. "../lua/?.lua;" .. package.path + +-- Stub ngx. The only surface handle_thinking touches is ngx.req.get_headers(). +local stub_headers = {} +_G.ngx = { + req = { + get_headers = function() return stub_headers end, + }, + log = function() end, + INFO = 0, WARN = 0, ERR = 0, + now = function() return 0 end, +} + +-- Stub modules that e2ee_handler requires but handle_thinking doesn't use. +package.loaded["cjson.safe"] = { encode = function(x) return x end, decode = function(x) return x end } +package.loaded["resty.http"] = { new = function() return {} end } +package.loaded["e2ee_crypto"] = {} +package.loaded["e2ee_discovery"] = {} + +local e2ee = require("e2ee_handler") + +local failures = 0 +local total = 0 + +local function set_header(name, value) + stub_headers = {} + if value ~= nil then stub_headers[name] = value end +end + +local function assert_eq(label, actual, expected) + total = total + 1 + local ok + if type(expected) == "table" and type(actual) == "table" then + ok = true + for k, v in pairs(expected) do + if actual[k] ~= v then ok = false; break end + end + if ok then + for k, v in pairs(actual) do + if expected[k] ~= v then ok = false; break end + end + end + else + ok = (actual == expected) + end + if not ok then + failures = failures + 1 + io.stderr:write(string.format( + "FAIL %s\n expected: %s\n actual: %s\n", + label, tostring(expected), tostring(actual) + )) + else + io.stdout:write("ok " .. label .. "\n") + end +end + +-- 1. Missing model → returns request, nil; no mutation. +set_header() +do + local req = { stream = true } + local out, model = e2ee.handle_thinking(req) + assert_eq("nil model returns nil", model, nil) + assert_eq("nil model leaves request untouched", out.chat_template_kwargs, nil) +end + +-- 2. Plain model, no suffix, no header, no kwargs → unchanged. +set_header() +do + local req = { model = "deepseek-ai/DeepSeek-V3.1-TEE" } + local out, model = e2ee.handle_thinking(req) + assert_eq("plain model name preserved", model, "deepseek-ai/DeepSeek-V3.1-TEE") + assert_eq("plain model payload.model preserved", out.model, "deepseek-ai/DeepSeek-V3.1-TEE") + assert_eq("plain model leaves kwargs absent", out.chat_template_kwargs, nil) +end + +-- 3. :THINKING suffix → strip, set both keys. +set_header() +do + local req = { model = "zai-org/GLM-5.1-TEE:THINKING" } + local out, model = e2ee.handle_thinking(req) + assert_eq(":THINKING strips suffix (return)", model, "zai-org/GLM-5.1-TEE") + assert_eq(":THINKING strips suffix (payload)", out.model, "zai-org/GLM-5.1-TEE") + assert_eq(":THINKING sets thinking=true", out.chat_template_kwargs.thinking, true) + assert_eq(":THINKING sets enable_thinking=true", out.chat_template_kwargs.enable_thinking, true) +end + +-- 4. X-Enable-Thinking: true → sets both keys, no suffix. +set_header("X-Enable-Thinking", "true") +do + local req = { model = "zai-org/GLM-5.1-TEE" } + local out = e2ee.handle_thinking(req) + assert_eq("header=true sets thinking", out.chat_template_kwargs.thinking, true) + assert_eq("header=true sets enable_thinking", out.chat_template_kwargs.enable_thinking, true) + assert_eq("header=true keeps model intact", out.model, "zai-org/GLM-5.1-TEE") +end + +-- 5. X-Enable-Thinking: TRUE (uppercase) → case-insensitive. +set_header("X-Enable-Thinking", "TRUE") +do + local req = { model = "zai-org/GLM-5.1-TEE" } + local out = e2ee.handle_thinking(req) + assert_eq("header TRUE case-insensitive", out.chat_template_kwargs.thinking, true) +end + +-- 6. X-Enable-Thinking: false → does not set kwargs. +set_header("X-Enable-Thinking", "false") +do + local req = { model = "deepseek-ai/DeepSeek-V3.1-TEE" } + local out = e2ee.handle_thinking(req) + assert_eq("header=false leaves kwargs absent", out.chat_template_kwargs, nil) +end + +-- 7. Only `thinking` key → `enable_thinking` propagated. +set_header() +do + local req = { + model = "deepseek-ai/DeepSeek-V3.1-TEE", + chat_template_kwargs = { thinking = true }, + } + local out = e2ee.handle_thinking(req) + assert_eq("thinking=true propagates to enable_thinking", out.chat_template_kwargs.enable_thinking, true) +end + +-- 8. Only `enable_thinking` key → `thinking` propagated. +set_header() +do + local req = { + model = "deepseek-ai/DeepSeek-V3.1-TEE", + chat_template_kwargs = { enable_thinking = false }, + } + local out = e2ee.handle_thinking(req) + assert_eq("enable_thinking=false propagates to thinking", out.chat_template_kwargs.thinking, false) +end + +-- 9. Per-model prefix default-on (GLM-4.7) when kwargs present but no thinking key. +set_header() +do + local req = { + model = "zai-org/GLM-4.7", + chat_template_kwargs = { some_other = "x" }, + } + local out = e2ee.handle_thinking(req) + assert_eq("GLM-4.7 prefix defaults thinking=true", out.chat_template_kwargs.thinking, true) + assert_eq("GLM-4.7 prefix defaults enable_thinking=true", out.chat_template_kwargs.enable_thinking, true) + assert_eq("GLM-4.7 prefix preserves other keys", out.chat_template_kwargs.some_other, "x") +end + +-- 10. Per-model exact default-on (GLM-4.7-TEE) when kwargs absent. +set_header() +do + local req = { model = "zai-org/GLM-4.7-TEE" } + local out = e2ee.handle_thinking(req) + assert_eq("GLM-4.7-TEE exact defaults thinking=true", out.chat_template_kwargs.thinking, true) + assert_eq("GLM-4.7-TEE exact defaults enable_thinking=true", out.chat_template_kwargs.enable_thinking, true) +end + +-- 11. Per-model exact default-on for DeepSeek-V3.2-Speciale-TEE and Kimi-K2.5-TEE. +set_header() +do + local req = { model = "deepseek-ai/DeepSeek-V3.2-Speciale-TEE" } + local out = e2ee.handle_thinking(req) + assert_eq("DeepSeek-V3.2-Speciale-TEE exact thinking=true", out.chat_template_kwargs.thinking, true) +end +do + local req = { model = "moonshotai/Kimi-K2.5-TEE" } + local out = e2ee.handle_thinking(req) + assert_eq("Kimi-K2.5-TEE exact thinking=true", out.chat_template_kwargs.thinking, true) +end + +-- 12. MiMo-V2-Flash prefix default-off when kwargs present. +set_header() +do + local req = { + model = "XiaomiMiMo/MiMo-V2-Flash-Pro", + chat_template_kwargs = { some_other = "x" }, + } + local out = e2ee.handle_thinking(req) + assert_eq("MiMo-V2-Flash prefix defaults thinking=false", out.chat_template_kwargs.thinking, false) + assert_eq("MiMo-V2-Flash prefix defaults enable_thinking=false", out.chat_template_kwargs.enable_thinking, false) +end + +-- 13. MiMo-V2-Flash-TEE exact default-off when kwargs absent. +set_header() +do + local req = { model = "XiaomiMiMo/MiMo-V2-Flash-TEE" } + local out = e2ee.handle_thinking(req) + assert_eq("MiMo-V2-Flash-TEE exact thinking=false", out.chat_template_kwargs.thinking, false) + assert_eq("MiMo-V2-Flash-TEE exact enable_thinking=false", out.chat_template_kwargs.enable_thinking, false) +end + +-- 14. :THINKING on MiMo → user override wins over default-off. +set_header() +do + local req = { model = "XiaomiMiMo/MiMo-V2-Flash-TEE:THINKING" } + local out, model = e2ee.handle_thinking(req) + assert_eq(":THINKING on MiMo strips suffix", model, "XiaomiMiMo/MiMo-V2-Flash-TEE") + assert_eq(":THINKING on MiMo sets thinking=true", out.chat_template_kwargs.thinking, true) +end + +-- 15. Non-reasoning model + no suffix + no header → no kwargs injected. +set_header() +do + local req = { model = "deepseek-ai/DeepSeek-V3.1-TEE" } + local out = e2ee.handle_thinking(req) + assert_eq("non-reasoning model leaves kwargs absent", out.chat_template_kwargs, nil) +end + +-- 16. :THINKING preserves unrelated kwargs. +set_header() +do + local req = { + model = "zai-org/GLM-5.1-TEE:THINKING", + chat_template_kwargs = { tools_in_user_message = false }, + } + local out = e2ee.handle_thinking(req) + assert_eq(":THINKING preserves unrelated kwargs", + out.chat_template_kwargs.tools_in_user_message, false) + assert_eq(":THINKING still sets thinking=true", + out.chat_template_kwargs.thinking, true) +end + +-- 17. :THINKING appears mid-string (not suffix) → not stripped. +set_header() +do + local req = { model = "some:THINKING:other" } + local out, model = e2ee.handle_thinking(req) + assert_eq("non-suffix :THINKING not stripped", model, "some:THINKING:other") + assert_eq("non-suffix :THINKING leaves kwargs absent", out.chat_template_kwargs, nil) +end + +-- 18. Explicit thinking=false overrides per-model default-on. +set_header() +do + local req = { + model = "zai-org/GLM-4.7", + chat_template_kwargs = { thinking = false }, + } + local out = e2ee.handle_thinking(req) + assert_eq("explicit thinking=false overrides GLM-4.7 default", out.chat_template_kwargs.thinking, false) + assert_eq("explicit thinking=false propagates enable_thinking=false", out.chat_template_kwargs.enable_thinking, false) +end + +-- 19. DeepSeek-V3.2-Speciale prefix (bare, no TEE) with kwargs present. +set_header() +do + local req = { + model = "deepseek-ai/DeepSeek-V3.2-Speciale", + chat_template_kwargs = {}, + } + local out = e2ee.handle_thinking(req) + assert_eq("DeepSeek-V3.2-Speciale prefix default thinking=true", out.chat_template_kwargs.thinking, true) +end + +io.stdout:write(string.format("\n%d/%d passed\n", total - failures, total)) +os.exit(failures == 0 and 0 or 1)