ModelRouter/internal/lua/adapters/ollama.lua

local adapter = {}

adapter.name = "ollama"
adapter.version = "2.0.0"
adapter.endpoint = "/api/chat"
adapter.headers = {}

-- Ollama API 格式：{ model, messages, stream, options:{temperature,num_predict} }
-- 工具调用必须一并翻译：Ollama 的 assistant 消息用 tool_calls（arguments 是对象
-- 而非 JSON 字符串），工具结果用 tool 角色 + tool_name。此前这里只复制了
-- role/content，助手那轮的调用和 tool 消息的归属全部丢失，模型看到无来源的
-- 工具结果就只能反复重发同一个调用。
function adapter.transform_request(raw_body)
    local ok, req = pcall(json.decode, raw_body)
    if not ok then return raw_body end

    local ollama_req = {
        model = req.model or "llama3",
        stream = req.stream or false,
        options = {
            temperature = req.temperature or 0.7,
            num_predict = req.max_tokens or 2048
        }
    }

    -- 转换 messages 格式（Ollama messages 支持 images base64 数组）
    if req.messages then
        local msgs = {}
        local call_names = {}   -- tool_call_id -> 函数名
        for _, m in ipairs(req.messages) do
            local text, images
            if type(m.content) == "string" then
                text, images = m.content, nil
            else
                text = ""
                images = {}
                for _, p in ipairs(m.content or {}) do
                    if type(p) == "table" then
                        if p.type == "text" then
                            text = text .. (p.text or "")
                        elseif p.type == "image_url" and type(p.image_url) == "table" and p.image_url.url then
                            local b64 = string.match(p.image_url.url, "^data:[^,]+;base64,(.+)$")
                            if b64 then table.insert(images, b64) end
                        end
                    end
                end
                if #images == 0 then images = nil end
            end
            local msg = { role = m.role, content = text }
            if images then msg.images = images end

            -- 助手轮的工具调用：arguments 转成对象
            if type(m.tool_calls) == "table" and #m.tool_calls > 0 then
                local tcs = {}
                for _, tc in ipairs(m.tool_calls) do
                    if type(tc) == "table" then
                        local fn = tc["function"] or {}
                        local args = fn.arguments
                        if type(args) == "string" and args ~= "" then
                            local aok, decoded = pcall(json.decode, args)
                            args = aok and decoded or {}
                        elseif type(args) ~= "table" then
                            args = {}
                        end
                        if tc.id ~= nil then call_names[tc.id] = fn.name or "" end
                        table.insert(tcs, {
                            ["function"] = { name = fn.name or "", arguments = args }
                        })
                    end
                end
                if #tcs > 0 then msg.tool_calls = tcs end
            end

            -- 工具结果：Ollama 用 tool_name 标识归属（不认 tool_call_id）
            if m.role == "tool" then
                local name = call_names[m.tool_call_id] or m.name or ""
                if m.tool_call_id ~= nil then
                    msg.tool_call_id = m.tool_call_id
                end
                if name ~= "" then msg.tool_name = name end
            end

            table.insert(msgs, msg)
        end
        ollama_req.messages = msgs
    end

    -- tools 透传（Ollama 的 shape 与 OpenAI 一致）
    if type(req.tools) == "table" and #req.tools > 0 then
        ollama_req.tools = req.tools
    end

    return json.encode(ollama_req)
end

function adapter.transform_response(raw_body)
    local ok, resp = pcall(json.decode, raw_body)
    if not ok then return raw_body end

    local p = resp.prompt_eval_count or 0
    local c = resp.eval_count or 0
    local unified = {
        content = "",
        finish_reason = resp.done_reason or "",
        -- key must be token_usage to match Go's UnifiedResponse json tag
        token_usage = { prompt = p, completion = c, total = p + c }
    }

    if resp.message then
        unified.content = resp.message.content or ""
        -- Non-streaming tool calls were previously dropped: the field was
        -- initialized to an empty table and never filled, while
        -- transform_stream_chunk handled them. A non-streaming agent turn thus
        -- looked like a plain answer and the tool loop stopped.
        if type(resp.message.tool_calls) == "table" and #resp.message.tool_calls > 0 then
            local tcs = {}
            for _, tc in ipairs(resp.message.tool_calls) do
                local fn = tc["function"] or {}
                -- Ollama sends arguments as an object already
                local args = fn.arguments
                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(tcs, {
                    id = tc.id or ("call_" .. #tcs),
                    type = tc.type or "function",
                    name = fn.name or "",
                    arguments = args
                })
            end
            unified.tool_calls = tcs
            -- Ollama reports done_reason "stop" alongside tool calls
            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

    -- Ollama's terminal chunk (done=true) carries token counts but may omit
    -- message; pass them through so the gateway emits real usage.
    local uses = nil
    if chunk.done then
        local p = chunk.prompt_eval_count or 0
        local c = chunk.eval_count or 0
        if p > 0 or c > 0 then
            uses = { prompt = p, completion = c, total = p + c }
        end
    end

    -- Ollama done_reason -> OpenAI finish_reason ("length" 透传，其余归一 stop)
    local finish = nil
    if chunk.done then
        if chunk.done_reason == "length" then
            finish = "length"
        else
            finish = "stop"
        end
    end

    if not chunk.message then
        if uses ~= nil then
            return json.encode({ content = "", done = true, finish_reason = finish, usage = uses })
        end
        return json.encode({ content = "", done = true, finish_reason = finish })
    end

    local unified = {
        content = chunk.message.content or "",
        done = chunk.done or false,
        finish_reason = finish
    }
    if uses ~= nil then
        unified.usage = uses
    end
    if chunk.message.reasoning_content then
        unified.reasoning_content = chunk.message.reasoning_content
    end
    if chunk.message.tool_calls then
        local tools = {}
        for _, tc in ipairs(chunk.message.tool_calls) do
            table.insert(tools, {
                index = #tools,
                id = tc.id or ("call_" .. #tools),
                type = "function",
                ["function"] = {
                    name = tc["function"] and tc["function"].name or "",
                    arguments = tc["function"] and (tc["function"].arguments or "{}") or "{}"
                }
            })
        end
        unified.tool_calls = tools
    end
    return json.encode(unified)
end

-- 错误收敛：Ollama 常见 {error:"..."} 字符串（新版本也有对象形态）
function adapter.transform_error(status, body)
    local ok, resp = pcall(json.decode, body)
    if not ok or type(resp) ~= "table" then return nil end
    if type(resp.error) == "string" then return resp.error end
    if type(resp.error) == "table" and type(resp.error.message) == "string" then
        return resp.error.message
    end
    return nil
end

return adapter