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
    end

    if resp.candidates and #resp.candidates > 0 then
        local cand = resp.candidates[1]
        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
                end
            end
        end
        if cand.finishReason then
            unified.finish_reason = cand.finishReason
        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 }
        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

return adapter