diff --git a/lua/claude_handler.lua b/lua/claude_handler.lua index 8531a1c..9538cb4 100644 --- a/lua/claude_handler.lua +++ b/lua/claude_handler.lua @@ -54,7 +54,7 @@ function _M.handle() if not is_streaming then -- Non-streaming local decrypted, round_err = e2ee.e2ee_round_trip( - api_key, model, oai_body, false, "/v1/chat/completions" + api_key, model, oai_body, false, "/v1/chat/completions", nil, oai_request ) if not decrypted then @@ -101,7 +101,8 @@ function _M.handle() ngx.flush(true) end end - end) + end, + oai_request) if round_err then if round_err.raw then diff --git a/lua/e2ee_handler.lua b/lua/e2ee_handler.lua index 71938e5..7108ee8 100644 --- a/lua/e2ee_handler.lua +++ b/lua/e2ee_handler.lua @@ -14,9 +14,20 @@ local cjson = require("cjson.safe") local http = require("resty.http") local API_BASE = "https://api.chutes.ai" +local payload_cjson local _M = {} +local function get_payload_cjson() + if not payload_cjson then + payload_cjson = cjson.new() + if payload_cjson.decode_array_with_array_mt then + payload_cjson.decode_array_with_array_mt(true) + end + end + return payload_cjson +end + --- Extract API key from Authorization header or x-api-key function _M.get_api_key() local headers = ngx.req.get_headers() @@ -107,9 +118,70 @@ end --- E2EE round-trip: encrypt request, send, decrypt response -- For non-streaming: returns (decrypted_json_string, nil) or (nil, {status=N, message=S}) -- For streaming: calls on_chunk(line) for each decrypted SSE data line, on_chunk(nil) at end -function _M.e2ee_round_trip(api_key, model, body_json, is_streaming, e2e_path, on_chunk) +function _M.e2ee_round_trip(api_key, model, body_json, is_streaming, e2e_path, on_chunk, request_payload) local err + local thinking + local base_model = model:match("^(.-):THINKING$") + local changed = false + + if base_model then + model = base_model + thinking = true + end + + local h = ngx.var and ngx.var.http_x_enable_thinking + if type(h) == "string" then + h = h:lower() + if h == "true" then + thinking = true + elseif h == "false" then + thinking = false + end + end + + local kwargs_hint = type(request_payload) == "table" + and type(request_payload.chat_template_kwargs) == "table" + and request_payload.chat_template_kwargs or nil + local normalize_kwargs = kwargs_hint + and ((kwargs_hint.thinking ~= nil and kwargs_hint.enable_thinking == nil) + or (kwargs_hint.enable_thinking ~= nil and kwargs_hint.thinking == nil)) + + if thinking ~= nil or normalize_kwargs then + local json = get_payload_cjson() + local payload = json.decode(body_json) + if payload then + if base_model then + payload.model = model + changed = true + end + + local kwargs = type(payload.chat_template_kwargs) == "table" and payload.chat_template_kwargs or nil + if kwargs and getmetatable(kwargs) == json.array_mt then + kwargs = nil + end + if thinking ~= nil then + kwargs = kwargs or {} + kwargs.thinking = thinking + kwargs.enable_thinking = thinking + payload.chat_template_kwargs = kwargs + changed = true + elseif kwargs then + if kwargs.thinking ~= nil and kwargs.enable_thinking == nil then + kwargs.enable_thinking = kwargs.thinking + changed = true + elseif kwargs.enable_thinking ~= nil and kwargs.thinking == nil then + kwargs.thinking = kwargs.enable_thinking + changed = true + end + end + + if changed then + body_json = json.encode(payload) + end + end + end + -- Resolve model -> chute_id local chute_id chute_id, err = discovery.resolve_chute_id(model, api_key) @@ -354,7 +426,9 @@ function _M.handle() end if not is_streaming then - local decrypted, round_err = _M.e2ee_round_trip(api_key, model, body, false, original_path) + local decrypted, round_err = _M.e2ee_round_trip( + api_key, model, body, false, original_path, nil, payload + ) if not decrypted then if round_err.raw then ngx.status = round_err.status @@ -384,7 +458,8 @@ function _M.handle() ngx.flush(true) end end - end) + end, + payload) if round_err then if round_err.raw then diff --git a/lua/responses_handler.lua b/lua/responses_handler.lua index 10bf166..7b277c2 100644 --- a/lua/responses_handler.lua +++ b/lua/responses_handler.lua @@ -54,7 +54,7 @@ function _M.handle() if not is_streaming then -- Non-streaming local decrypted, round_err = e2ee.e2ee_round_trip( - api_key, model, oai_body, false, "/v1/chat/completions" + api_key, model, oai_body, false, "/v1/chat/completions", nil, oai_request ) if not decrypted then @@ -101,7 +101,8 @@ function _M.handle() ngx.flush(true) end end - end) + end, + oai_request) if round_err then if round_err.raw then