ModelRouter/internal/lua/adapters/gemini.lua

local adapter = {}

adapter.name = "gemini"
adapter.version = "2.0.0"
adapter.endpoint = "/v1/models"
adapter.headers = {}

-- Gemini API: POST /v1/models/{model}:generateContent
-- Auth: API key in query param ?key=XXX or Authorization: Bearer XXX
function adapter.transform_request(raw_body)
    local ok, req = pcall(json.decode, raw_body)
    if not ok then return raw_body end

    -- 将 OpenAI 风格 content（字符串或 [{type:*}] 数组）拆成 Gemini parts
    local function to_parts(content)
        if type(content) == "string" then
            return { { text = content } }
        end
        local parts = {}
        for _, p in ipairs(content or {}) do
            if p.type == "text" then
                table.insert(parts, { text = p.text })
            elseif p.type == "image_url" and type(p.image_url) == "table" and p.image_url.url then
                local mt, b64 = string.match(p.image_url.url, "^data:([^,]+);base64,(.+)$")
                if b64 then
                    table.insert(parts, { inline_data = { mime_type = mt or "image/png", data = b64 } })
                end
            end
        end
        return parts
    end

    local contents = {}
    for _, m in ipairs(req.messages or {}) do
        table.insert(contents, {
            role = (m.role == "assistant") and "model" or m.role,
            parts = to_parts(m.content)
        })
    end

    local gemini_req = {
        contents = contents,
        generationConfig = {
            temperature = req.temperature or 0.7,
            maxOutputTokens = req.max_tokens or 4096,
        }
    }

    if req.stream then
        gemini_req.stream = true
    end

    return json.encode(gemini_req)
end

-- Gemini 的 endpoint 动态拼接：/v1/models/{model}:generateContent
function adapter.transform_response(raw_body)
    local ok, resp = pcall(json.decode, raw_body)
    if not ok then return raw_body end

    local unified = {
        content = "",
        finish_reason = "",
        token_usage = { prompt = 0, completion = 0, total = 0 }
    }

    if resp.usageMetadata then
        unified.token_usage.prompt = resp.usageMetadata.promptTokenCount or 0
        unified.token_usage.completion = resp.usageMetadata.candidatesTokenCount or 0
        unified.token_usage.total = resp.usageMetadata.totalTokenCount or 0
        -- Gemini reports context-cache reads as cachedContentTokenCount;
        -- normalize into OpenAI-standard prompt_tokens_details.cached_tokens
        -- so clients and the audit trail see the hit count. Emitted even when
        -- 0 so a reported miss stays distinguishable from "not reported".
        if resp.usageMetadata.cachedContentTokenCount ~= nil then
            unified.token_usage.prompt_tokens_details = {
                cached_tokens = resp.usageMetadata.cachedContentTokenCount
            }
        end
    end

    if resp.candidates and #resp.candidates > 0 then
        local cand = resp.candidates[1]
        local tools = {}
        if cand.content and cand.content.parts then
            for _, part in ipairs(cand.content.parts) do
                if part.text then
                    unified.content = unified.content .. part.text
                elseif part.functionCall then
                    -- Non-streaming tool calls used to be dropped here while
                    -- transform_stream_chunk handled them, so a non-streaming
                    -- agent turn looked like a plain text answer and the tool
                    -- loop died. Gemini's args are already an object.
                    local args = part.functionCall.args
                    if type(args) == "string" then
                        local aok, decoded = pcall(json.decode, args)
                        args = aok and decoded or {}
                    elseif type(args) ~= "table" then
                        args = {}
                    end
                    table.insert(tools, {
                        id = part.functionCall.id or ("call_" .. #tools),
                        type = "function",
                        name = part.functionCall.name or "",
                        arguments = args
                    })
                end
            end
        end
        if cand.finishReason then
            unified.finish_reason = cand.finishReason
        end
        if #tools > 0 then
            unified.tool_calls = tools
            -- Gemini reports finishReason "STOP" even when it emitted a
            -- functionCall; clients keyed on finish_reason would treat that as
            -- a completed answer and never run the tool.
            unified.finish_reason = "tool_calls"
        end
    end

    return json.encode(unified)
end

function adapter.transform_stream_chunk(raw_chunk)
    local ok, chunk = pcall(json.decode, raw_chunk)
    if not ok then return "" end

    -- Gemini attaches usageMetadata to the final chunk (alongside or after
    -- candidates). Keys map to Go's TokenUsage json tags (prompt/...).
    local uses = nil
    if type(chunk.usageMetadata) == "table" then
        local p = chunk.usageMetadata.promptTokenCount or 0
        local c = chunk.usageMetadata.candidatesTokenCount or 0
        local t = chunk.usageMetadata.totalTokenCount or 0
        if p > 0 or c > 0 or t > 0 then
            uses = { prompt = p, completion = c, total = t }
            -- Emit details whenever the field is present, even at 0, so a
            -- reported cache miss stays distinguishable from "not reported".
            if chunk.usageMetadata.cachedContentTokenCount ~= nil then
                uses.prompt_tokens_details = {
                    cached_tokens = chunk.usageMetadata.cachedContentTokenCount
                }
            end
        end
    end

    if not chunk.candidates or #chunk.candidates == 0 then
        if uses ~= nil then
            return json.encode({ usage = uses, done = false })
        end
        return ""
    end
    local cand = chunk.candidates[1]
    -- Gemini finishReason -> OpenAI finish_reason
    local finish = nil
    if cand.finishReason ~= nil then
        if cand.finishReason == "MAX_TOKENS" then
            finish = "length"
        elseif cand.finishReason == "SAFETY" or cand.finishReason == "RECITATION"
            or cand.finishReason == "BLOCKLIST" then
            finish = "content_filter"
        else
            finish = "stop"
        end
    end
    local unified = { content = "", done = (finish ~= nil), finish_reason = finish }
    local reasoning = ""
    local tools = {}
    if cand.content and cand.content.parts then
        for _, part in ipairs(cand.content.parts) do
            if part.text then
                unified.content = (unified.content or "") .. part.text
            elseif part.reasoning_content then
                reasoning = reasoning .. part.reasoning_content
            elseif part.functionCall then
                table.insert(tools, {
                    index = #tools,
                    id = part.functionCall.id or ("call_" .. #tools),
                    type = "function",
                    ["function"] = {
                        name = part.functionCall.name or "",
                        arguments = part.functionCall.args or "{}"
                    }
                })
            end
        end
    end
    if reasoning ~= "" then unified.reasoning_content = reasoning end
    if #tools > 0 then unified.tool_calls = tools end
    if uses ~= nil then
        unified.usage = uses
    end
    return json.encode(unified)
end

-- 错误收敛：Gemini REST 信封 {error:{code, message, status}}
function adapter.transform_error(status, body)
    local ok, resp = pcall(json.decode, body)
    if not ok or type(resp) ~= "table" then return nil end
    local e = resp.error
    if type(e) == "table" then
        if type(e.message) == "string" then
            if type(e.status) == "string" then
                return e.status .. ": " .. e.message
            end
            return e.message
        end
    end
    return nil
end

return adapter