feat: tools (#1180)

* feat: tools

* feat: claude use tools

* feat: openai use tools

lua/avante/providers/bedrock/claude.lua

@@ -6,6 +6,8 @@
 ---@field role "user" | "assistant"
 ---@field content [AvanteBedrockClaudeTextMessage][]
 
+local Claude = require("avante.providers.claude")
+
 ---@class AvanteBedrockModelHandler
 local M = {}
 
@@ -33,25 +35,7 @@ M.parse_messages = function(opts)
   return messages
 end
 
-M.parse_response = function(ctx, data_stream, event_state, opts)
-  if event_state == nil then
-    if data_stream:match('"content_block_delta"') then
-      event_state = "content_block_delta"
-    elseif data_stream:match('"message_stop"') then
-      event_state = "message_stop"
-    end
-  end
-  if event_state == "content_block_delta" then
-    local ok, json = pcall(vim.json.decode, data_stream)
-    if not ok then return end
-    opts.on_chunk(json.delta.text)
-  elseif event_state == "message_stop" then
-    opts.on_complete(nil)
-    return
-  elseif event_state == "error" then
-    opts.on_complete(vim.json.decode(data_stream))
-  end
-end
+M.parse_response = Claude.parse_response
 
 ---@param prompt_opts AvantePromptOptions
 ---@param body_opts table
@@ -60,7 +44,6 @@ M.build_bedrock_payload = function(prompt_opts, body_opts)
   local system_prompt = prompt_opts.system_prompt or ""
   local messages = M.parse_messages(prompt_opts)
   local max_tokens = body_opts.max_tokens or 2000
-  local temperature = body_opts.temperature or 0.7
   local payload = {
     anthropic_version = "bedrock-2023-05-31",
     max_tokens = max_tokens,

lua/avante/providers/claude.lua

@@ -17,6 +17,44 @@ local P = require("avante.providers")
 ---@field role "user" | "assistant"
 ---@field content [AvanteClaudeTextMessage | AvanteClaudeImageMessage][]
 
+---@class AvanteClaudeTool
+---@field name string
+---@field description string
+---@field input_schema AvanteClaudeToolInputSchema
+
+---@class AvanteClaudeToolInputSchema
+---@field type "object"
+---@field properties table<string, AvanteClaudeToolInputSchemaProperty>
+---@field required string[]
+
+---@class AvanteClaudeToolInputSchemaProperty
+---@field type "string" | "number" | "boolean"
+---@field description string
+---@field enum? string[]
+
+---@param tool AvanteLLMTool
+---@return AvanteClaudeTool
+local function transform_tool(tool)
+  local input_schema_properties = {}
+  local required = {}
+  for _, field in ipairs(tool.param.fields) do
+    input_schema_properties[field.name] = {
+      type = field.type,
+      description = field.description,
+    }
+    if not field.optional then table.insert(required, field.name) end
+  end
+  return {
+    name = tool.name,
+    description = tool.description,
+    input_schema = {
+      type = "object",
+      properties = input_schema_properties,
+      required = required,
+    },
+  }
+end
+
 ---@class AvanteProviderFunctor
 local M = {}
 
@@ -74,26 +112,101 @@ M.parse_messages = function(opts)
     messages[#messages].content = message_content
   end
 
+  if opts.tool_use then
+    local msg = {
+      role = "assistant",
+      content = {},
+    }
+    if opts.response_content then
+      msg.content[#msg.content + 1] = {
+        type = "text",
+        text = opts.response_content,
+      }
+    end
+    msg.content[#msg.content + 1] = {
+      type = "tool_use",
+      id = opts.tool_use.id,
+      name = opts.tool_use.name,
+      input = vim.json.decode(opts.tool_use.input_json),
+    }
+    messages[#messages + 1] = msg
+  end
+
+  if opts.tool_result then
+    messages[#messages + 1] = {
+      role = "user",
+      content = {
+        {
+          type = "tool_result",
+          tool_use_id = opts.tool_result.tool_use_id,
+          content = opts.tool_result.content,
+          is_error = opts.tool_result.is_error,
+        },
+      },
+    }
+  end
+
   return messages
 end
 
 M.parse_response = function(ctx, data_stream, event_state, opts)
   if event_state == nil then
-    if data_stream:match('"content_block_delta"') then
-      event_state = "content_block_delta"
+    if data_stream:match('"message_start"') then
+      event_state = "message_start"
+    elseif data_stream:match('"message_delta"') then
+      event_state = "message_delta"
     elseif data_stream:match('"message_stop"') then
       event_state = "message_stop"
+    elseif data_stream:match('"content_block_start"') then
+      event_state = "content_block_start"
+    elseif data_stream:match('"content_block_delta"') then
+      event_state = "content_block_delta"
+    elseif data_stream:match('"content_block_stop"') then
+      event_state = "content_block_stop"
     end
   end
-  if event_state == "content_block_delta" then
-    local ok, json = pcall(vim.json.decode, data_stream)
+  if event_state == "message_start" then
+    local ok, jsn = pcall(vim.json.decode, data_stream)
     if not ok then return end
-    opts.on_chunk(json.delta.text)
-  elseif event_state == "message_stop" then
-    opts.on_complete(nil)
+    opts.on_start(jsn.message.usage)
+  elseif event_state == "content_block_start" then
+    local ok, jsn = pcall(vim.json.decode, data_stream)
+    if not ok then return end
+    if jsn.content_block.type == "tool_use" then
+      ctx.tool_use = {
+        name = jsn.content_block.name,
+        id = jsn.content_block.id,
+        input_json = "",
+      }
+    elseif jsn.content_block.type == "text" then
+      ctx.response_content = ""
+    end
+  elseif event_state == "content_block_delta" then
+    local ok, jsn = pcall(vim.json.decode, data_stream)
+    if not ok then return end
+    if ctx.tool_use and jsn.delta.type == "input_json_delta" then
+      ctx.tool_use.input_json = ctx.tool_use.input_json .. jsn.delta.partial_json
+      return
+    elseif ctx.response_content and jsn.delta.type == "text_delta" then
+      ctx.response_content = ctx.response_content .. jsn.delta.text
+    end
+    opts.on_chunk(jsn.delta.text)
+  elseif event_state == "message_delta" then
+    local ok, jsn = pcall(vim.json.decode, data_stream)
+    if not ok then return end
+    if jsn.delta.stop_reason == "end_turn" then
+      opts.on_stop({ reason = "complete", usage = jsn.usage })
+    elseif jsn.delta.stop_reason == "tool_use" then
+      opts.on_stop({
+        reason = "tool_use",
+        usage = jsn.usage,
+        tool_use = ctx.tool_use,
+        response_content = ctx.response_content,
+      })
+    end
     return
   elseif event_state == "error" then
-    opts.on_complete(vim.json.decode(data_stream))
+    opts.on_stop({ reason = "error", error = vim.json.decode(data_stream) })
   end
 end
 
@@ -113,6 +226,13 @@ M.parse_curl_args = function(provider, prompt_opts)
 
   local messages = M.parse_messages(prompt_opts)
 
+  local tools = {}
+  if prompt_opts.tools then
+    for _, tool in ipairs(prompt_opts.tools) do
+      table.insert(tools, transform_tool(tool))
+    end
+  end
+
   return {
     url = Utils.url_join(base.endpoint, "/v1/messages"),
     proxy = base.proxy,
@@ -128,6 +248,7 @@ M.parse_curl_args = function(provider, prompt_opts)
         },
       },
       messages = messages,
+      tools = tools,
       stream = true,
     }, body_opts),
   }

lua/avante/providers/cohere.lua

@@ -62,7 +62,7 @@ M.parse_stream_data = function(data, opts)
   local json = vim.json.decode(data)
   if json.type ~= nil then
     if json.type == "message-end" and json.delta.finish_reason == "COMPLETE" then
-      opts.on_complete(nil)
+      opts.on_stop({ reason = "complete" })
       return
     end
     if json.type == "content-delta" then opts.on_chunk(json.delta.message.content.text) end

lua/avante/providers/gemini.lua

@@ -66,17 +66,17 @@ end
 
 M.parse_response = function(ctx, data_stream, _, opts)
   local ok, json = pcall(vim.json.decode, data_stream)
-  if not ok then opts.on_complete(json) end
+  if not ok then opts.on_stop({ reason = "error", error = json }) end
   if json.candidates then
     if #json.candidates > 0 then
       if json.candidates[1].finishReason and json.candidates[1].finishReason == "STOP" then
         opts.on_chunk(json.candidates[1].content.parts[1].text)
-        opts.on_complete(nil)
+        opts.on_stop({ reason = "complete" })
       else
         opts.on_chunk(json.candidates[1].content.parts[1].text)
       end
     else
-      opts.on_complete(nil)
+      opts.on_stop({ reason = "complete" })
     end
   end
 end

lua/avante/providers/init.lua

@@ -11,17 +11,28 @@ local DressingConfig = {
 local DressingState = { winid = nil, input_winid = nil, input_bufnr = nil }
 
 ---@class AvanteHandlerOptions: table<[string], string>
----@field on_chunk AvanteChunkParser
----@field on_complete AvanteCompleteParser
+---@field on_start AvanteLLMStartCallback
+---@field on_chunk AvanteLLMChunkCallback
+---@field on_stop AvanteLLMStopCallback
 ---
 ---@class AvanteLLMMessage
 ---@field role "user" | "assistant"
 ---@field content string
 ---
+---@class AvanteLLMToolResult
+---@field tool_name string
+---@field tool_use_id string
+---@field content string
+---@field is_error? boolean
+---
 ---@class AvantePromptOptions: table<[string], string>
 ---@field system_prompt string
 ---@field messages AvanteLLMMessage[]
 ---@field image_paths? string[]
+---@field tools? AvanteLLMTool[]
+---@field tool_result? AvanteLLMToolResult
+---@field tool_use? AvanteLLMToolUse
+---@field response_content? string
 ---
 ---@class AvanteGeminiMessage
 ---@field role "user"
@@ -35,8 +46,9 @@ local DressingState = { winid = nil, input_winid = nil, input_bufnr = nil }
 ---@alias AvanteCurlArgsParser fun(opts: AvanteProvider | AvanteProviderFunctor | AvanteBedrockProviderFunctor, code_opts: AvantePromptOptions): AvanteCurlOutput
 ---
 ---@class ResponseParser
----@field on_chunk fun(chunk: string): any
----@field on_complete fun(err: string|nil): any
+---@field on_start AvanteLLMStartCallback
+---@field on_chunk AvanteLLMChunkCallback
+---@field on_stop AvanteLLMStopCallback
 ---@alias AvanteResponseParser fun(ctx: any, data_stream: string, event_state: string, opts: ResponseParser): nil
 ---
 ---@class AvanteDefaultBaseProvider: table<string, any>
@@ -54,9 +66,31 @@ local DressingState = { winid = nil, input_winid = nil, input_bufnr = nil }
 ---@field temperature? number
 ---@field max_tokens? number
 ---
+---@class AvanteLLMUsage
+---@field input_tokens number
+---@field cache_creation_input_tokens number
+---@field cache_read_input_tokens number
+---@field output_tokens number
+---
+---@class AvanteLLMToolUse
+---@field name string
+---@field id string
+---@field input_json string
+---
+---@class AvanteLLMStartCallbackOptions
+---@field usage? AvanteLLMUsage
+---
+---@class AvanteLLMStopCallbackOptions
+---@field reason "complete" | "tool_use" | "error"
+---@field error? string | table
+---@field usage? AvanteLLMUsage
+---@field tool_use? AvanteLLMToolUse
+---@field response_content? string
+---
 ---@alias AvanteStreamParser fun(line: string, handler_opts: AvanteHandlerOptions): nil
----@alias AvanteChunkParser fun(chunk: string): any
----@alias AvanteCompleteParser fun(err: string|nil): nil
+---@alias AvanteLLMStartCallback fun(opts: AvanteLLMStartCallbackOptions): nil
+---@alias AvanteLLMChunkCallback fun(chunk: string): any
+---@alias AvanteLLMStopCallback fun(opts: AvanteLLMStopCallbackOptions): nil
 ---@alias AvanteLLMConfigHandler fun(opts: AvanteSupportedProvider): AvanteDefaultBaseProvider, table<string, any>
 ---
 ---@class AvanteProvider: AvanteSupportedProvider

lua/avante/providers/openai.lua

@@ -24,12 +24,72 @@ local P = require("avante.providers")
 ---@field index integer
 ---@field logprobs integer
 ---
+---@class OpenAIMessageToolCallFunction
+---@field name string
+---@field arguments string
+---
+---@class OpenAIMessageToolCall
+---@field id string
+---@field type "function"
+---@field function OpenAIMessageToolCallFunction
+---
 ---@class OpenAIMessage
 ---@field role? "user" | "system" | "assistant"
 ---@field content? string
 ---@field reasoning_content? string
 ---@field reasoning? string
+---@field tool_calls? OpenAIMessageToolCall[]
 ---
+---@class AvanteOpenAITool
+---@field type "function"
+---@field function AvanteOpenAIToolFunction
+---
+---@class AvanteOpenAIToolFunction
+---@field name string
+---@field description string
+---@field parameters AvanteOpenAIToolFunctionParameters
+---@field strict boolean
+---
+---@class AvanteOpenAIToolFunctionParameters
+---@field type string
+---@field properties table<string, AvanteOpenAIToolFunctionParameterProperty>
+---@field required string[]
+---@field additionalProperties boolean
+---
+---@class AvanteOpenAIToolFunctionParameterProperty
+---@field type string
+---@field description string
+
+---@param tool AvanteLLMTool
+---@return AvanteOpenAITool
+local function transform_tool(tool)
+  local input_schema_properties = {}
+  local required = {}
+  for _, field in ipairs(tool.param.fields) do
+    input_schema_properties[field.name] = {
+      type = field.type,
+      description = field.description,
+    }
+    if not field.optional then table.insert(required, field.name) end
+  end
+  local res = {
+    type = "function",
+    ["function"] = {
+      name = tool.name,
+      description = tool.description,
+    },
+  }
+  if vim.tbl_count(input_schema_properties) > 0 then
+    res["function"].parameters = {
+      type = "object",
+      properties = input_schema_properties,
+      required = required,
+      additionalProperties = false,
+    }
+  end
+  return res
+end
+
 ---@class AvanteProviderFunctor
 local M = {}
 
@@ -107,12 +167,34 @@ M.parse_messages = function(opts)
     table.insert(final_messages, { role = M.role_map[role] or role, content = message.content })
   end)
 
+  if opts.tool_result then
+    table.insert(final_messages, {
+      role = M.role_map["assistant"],
+      tool_calls = {
+        {
+          id = opts.tool_use.id,
+          type = "function",
+          ["function"] = {
+            name = opts.tool_use.name,
+            arguments = opts.tool_use.input_json,
+          },
+        },
+      },
+    })
+    local result_content = opts.tool_result.content or ""
+    table.insert(final_messages, {
+      role = "tool",
+      tool_call_id = opts.tool_result.tool_use_id,
+      content = opts.tool_result.is_error and "Error: " .. result_content or result_content,
+    })
+  end
+
   return final_messages
 end
 
 M.parse_response = function(ctx, data_stream, _, opts)
   if data_stream:match('"%[DONE%]":') then
-    opts.on_complete(nil)
+    opts.on_stop({ reason = "complete" })
     return
   end
   if data_stream:match('"delta":') then
@@ -121,7 +203,14 @@ M.parse_response = function(ctx, data_stream, _, opts)
     if jsn.choices and jsn.choices[1] then
       local choice = jsn.choices[1]
       if choice.finish_reason == "stop" or choice.finish_reason == "eos_token" then
-        opts.on_complete(nil)
+        opts.on_stop({ reason = "complete" })
+      elseif choice.finish_reason == "tool_calls" then
+        opts.on_stop({
+          reason = "tool_use",
+          usage = jsn.usage,
+          tool_use = ctx.tool_use,
+          response_content = ctx.response_content,
+        })
       elseif choice.delta.reasoning_content and choice.delta.reasoning_content ~= vim.NIL then
         if ctx.returned_think_start_tag == nil or not ctx.returned_think_start_tag then
           ctx.returned_think_start_tag = true
@@ -136,6 +225,17 @@ M.parse_response = function(ctx, data_stream, _, opts)
         end
         ctx.last_think_content = choice.delta.reasoning
         opts.on_chunk(choice.delta.reasoning)
+      elseif choice.delta.tool_calls then
+        local tool_call = choice.delta.tool_calls[1]
+        if not ctx.tool_use then
+          ctx.tool_use = {
+            name = tool_call["function"].name,
+            id = tool_call.id,
+            input_json = "",
+          }
+        else
+          ctx.tool_use.input_json = ctx.tool_use.input_json .. tool_call["function"].arguments
+        end
       elseif choice.delta.content then
         if
           ctx.returned_think_start_tag ~= nil and (ctx.returned_think_end_tag == nil or not ctx.returned_think_end_tag)
@@ -164,7 +264,7 @@ M.parse_response_without_stream = function(data, _, opts)
     local choice = json.choices[1]
     if choice.message and choice.message.content then
       opts.on_chunk(choice.message.content)
-      vim.schedule(function() opts.on_complete(nil) end)
+      vim.schedule(function() opts.on_stop({ reason = "complete" }) end)
     end
   end
 end
@@ -198,6 +298,13 @@ M.parse_curl_args = function(provider, code_opts)
     body_opts.temperature = 1
   end
 
+  local tools = {}
+  if code_opts.tools then
+    for _, tool in ipairs(code_opts.tools) do
+      table.insert(tools, transform_tool(tool))
+    end
+  end
+
   Utils.debug("endpoint", base.endpoint)
   Utils.debug("model", base.model)
 
@@ -210,6 +317,7 @@ M.parse_curl_args = function(provider, code_opts)
       model = base.model,
       messages = M.parse_messages(code_opts),
       stream = stream,
+      tools = tools,
     }, body_opts),
   }
 end