Function Calling in Neovim

function_calling.lua

local Job = require("plenary.job")
local json = vim.fn.json_encode
local function add(a, b)
  print("a + b = ", a + b)
  return a + b
end

local function Multiplication(a, b)
  print("a * b = ", a * b)
  return a * b
end

local function call_llm(prompt, f)
  local api_key = os.getenv("SILICONFLOW_TOKEN")
  if not api_key then
    print("Please set SILICONFLOW_TOKEN environment variable.")
    return
  end

  local functions = {
    {
      type = "function",
      ["function"] = {
        name = "add",
        description = "Compute the sum of two numbers",
        parameters = {
          type = "object",
          properties = {
            a = {
              type = "integer",
              description = "A number",
            },
            b = {
              type = "integer",
              description = "A number",
            },
          },
          required = { "a", "b" },
        },
      },
    },
    {
      type = "function",
      ["function"] = {
        name = "Multiplication",
        description = "Compute the product of two numbers",
        parameters = {
          type = "object",
          properties = {
            b = {
              type = "integer",
              description = "A number",
            },
            a = {
              type = "integer",
              description = "A number",
            },
          },
          required = { "b", "a" },
        },
      },
    },
  }

  local body = {
    -- model = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
    -- model = "internlm/internlm2_5-7b-chat",
    -- model = "THUDM/GLM-4-9B-0414",
    model = "Qwen/Qwen2.5-7B-Instruct",
    messages = {
      { role = "user", content = prompt },
    },
    tools = functions,
    -- functions = functions,
    -- function_call = "auto",
    stream = false,
  }

  Job:new({
    command = "curl",
    args = {
      "-s",
      "-X",
      "POST",
      "https://api.siliconflow.cn/v1/chat/completions",
      "-H",
      "Content-Type: application/json",
      "-H",
      "Authorization: Bearer " .. api_key,
      "-d",
      json(body),
    },
    on_stdout = vim.schedule_wrap(function(_, chunk)
      local j = vim.json.decode(chunk)
      local msg = j.choices[1].message
      local N = vim.tbl_count(msg.tool_calls)

      for i = 1, N do
        local name = msg.tool_calls[i]["function"].name
        print("name", name)
        local id = msg.tool_calls[i].id

        local params = vim.json.decode(msg.tool_calls[i]["function"].arguments)
        local keys = vim.tbl_filter(function(item)
          return item["function"].name == name
        end, functions)[1]["function"].parameters.required

        local p = {}

        for _, k in pairs(keys) do
          table.insert(p, params[k])
        end
        local fstring = string.dump(f[name])
        local func = load(fstring)

        local res = func(unpack(p))
        table.insert(body.messages, msg)
        table.insert(body.messages, { role = "tool", content = tostring(res), tool_call_id = id })
      end
    end),
    on_exit = vim.schedule_wrap(function()
      Job:new({
        command = "curl",
        args = {
          "-s",
          "-X",
          "POST",
          "https://api.siliconflow.cn/v1/chat/completions",
          "-H",
          "Content-Type: application/json",
          "-H",
          "Authorization: Bearer " .. api_key,
          "-d",
          json(body),
        },
        on_stdout = vim.schedule_wrap(function(_, c)
          print("cb", c)
        end),
      }):start()
    end),
  }):start()
end

call_llm("2 + 3 = ?, 4 * 5 = ?", { add = add, Multiplication = Multiplication })