refactor(history): change HistoryMessage:new() to accept role and item

Change the constructor to accept role and a single content item instead of raw AvanteLLMMessage instance. This will help when we will start changing message.content from being a string or a list of tables to just a string or a single table.
2025-07-18 15:08:19 -07:00
parent 2d9f8fd252
commit aa606b6147
9 changed files with 171 additions and 252 deletions
--- a/lua/avante/providers/claude.lua
+++ b/lua/avante/providers/claude.lua
@@ -161,6 +161,24 @@ function M:parse_response(ctx, data_stream, event_state, opts)
    end
  end
  if ctx.content_blocks == nil then ctx.content_blocks = {} end
+
+  ---@param content AvanteLLMMessageContentItem
+  ---@param uuid? string
+  ---@return avante.HistoryMessage
+  local function new_assistant_message(content, uuid)
+    assert(
+      event_state == "content_block_start"
+        or event_state == "content_block_delta"
+        or event_state == "content_block_stop",
+      "called with unexpected event_state: " .. event_state
+    )
+    return HistoryMessage:new("assistant", content, {
+      state = event_state == "content_block_stop" and "generated" or "generating",
+      turn_id = ctx.turn_id,
+      uuid = uuid,
+    })
+  end
+
  if event_state == "message_start" then
    local ok, jsn = pcall(vim.json.decode, data_stream)
    if not ok then return end
@@ -172,55 +190,33 @@ function M:parse_response(ctx, data_stream, event_state, opts)
    content_block.stoppped = false
    ctx.content_blocks[jsn.index + 1] = content_block
    if content_block.type == "text" then
-      local msg = HistoryMessage:new({
-        role = "assistant",
-        content = content_block.text,
-      }, {
-        state = "generating",
-        turn_id = ctx.turn_id,
-      })
+      local msg = new_assistant_message(content_block.text)
      content_block.uuid = msg.uuid
      if opts.on_messages_add then opts.on_messages_add({ msg }) end
-    end
-    if content_block.type == "thinking" then
+    elseif content_block.type == "thinking" then
      if opts.on_chunk then opts.on_chunk("<think>\n") end
      if opts.on_messages_add then
-        local msg = HistoryMessage:new({
-          role = "assistant",
-          content = {
-            {
-              type = "thinking",
-              thinking = content_block.thinking,
-              signature = content_block.signature,
-            },
-          },
-        }, {
-          state = "generating",
-          turn_id = ctx.turn_id,
+        local msg = new_assistant_message({
+          type = "thinking",
+          thinking = content_block.thinking,
+          signature = content_block.signature,
        })
        content_block.uuid = msg.uuid
        opts.on_messages_add({ msg })
      end
-    end
-    if content_block.type == "tool_use" and opts.on_messages_add then
-      local incomplete_json = JsonParser.parse(content_block.input_json)
-      local msg = HistoryMessage:new({
-        role = "assistant",
-        content = {
-          {
-            type = "tool_use",
-            name = content_block.name,
-            id = content_block.id,
-            input = incomplete_json or {},
-          },
-        },
-      }, {
-        state = "generating",
-        turn_id = ctx.turn_id,
-      })
-      content_block.uuid = msg.uuid
-      opts.on_messages_add({ msg })
-      -- opts.on_stop({ reason = "tool_use", streaming_tool_use = true })
+    elseif content_block.type == "tool_use" then
+      if opts.on_messages_add then
+        local incomplete_json = JsonParser.parse(content_block.input_json)
+        local msg = new_assistant_message({
+          type = "tool_use",
+          name = content_block.name,
+          id = content_block.id,
+          input = incomplete_json or {},
+        })
+        content_block.uuid = msg.uuid
+        opts.on_messages_add({ msg })
+        -- opts.on_stop({ reason = "tool_use", streaming_tool_use = true })
+      end
    end
  elseif event_state == "content_block_delta" then
    local ok, jsn = pcall(vim.json.decode, data_stream)
@@ -233,33 +229,21 @@ function M:parse_response(ctx, data_stream, event_state, opts)
    elseif jsn.delta.type == "thinking_delta" then
      content_block.thinking = content_block.thinking .. jsn.delta.thinking
      if opts.on_chunk then opts.on_chunk(jsn.delta.thinking) end
-      local msg = HistoryMessage:new({
-        role = "assistant",
-        content = {
-          {
-            type = "thinking",
-            thinking = content_block.thinking,
-            signature = content_block.signature,
-          },
-        },
-      }, {
-        state = "generating",
-        uuid = content_block.uuid,
-        turn_id = ctx.turn_id,
-      })
-      if opts.on_messages_add then opts.on_messages_add({ msg }) end
+      if opts.on_messages_add then
+        local msg = new_assistant_message({
+          type = "thinking",
+          thinking = content_block.thinking,
+          signature = content_block.signature,
+        }, content_block.uuid)
+        opts.on_messages_add({ msg })
+      end
    elseif jsn.delta.type == "text_delta" then
      content_block.text = content_block.text .. jsn.delta.text
      if opts.on_chunk then opts.on_chunk(jsn.delta.text) end
-      local msg = HistoryMessage:new({
-        role = "assistant",
-        content = content_block.text,
-      }, {
-        state = "generating",
-        uuid = content_block.uuid,
-        turn_id = ctx.turn_id,
-      })
-      if opts.on_messages_add then opts.on_messages_add({ msg }) end
+      if opts.on_messages_add then
+        local msg = new_assistant_message(content_block.text, content_block.uuid)
+        opts.on_messages_add({ msg })
+      end
    elseif jsn.delta.type == "signature_delta" then
      if ctx.content_blocks[jsn.index + 1].signature == nil then ctx.content_blocks[jsn.index + 1].signature = "" end
      ctx.content_blocks[jsn.index + 1].signature = ctx.content_blocks[jsn.index + 1].signature .. jsn.delta.signature
@@ -270,36 +254,11 @@ function M:parse_response(ctx, data_stream, event_state, opts)
    local content_block = ctx.content_blocks[jsn.index + 1]
    content_block.stoppped = true
    if content_block.type == "text" then
-      local msg = HistoryMessage:new({
-        role = "assistant",
-        content = content_block.text,
-      }, {
-        state = "generated",
-        uuid = content_block.uuid,
-        turn_id = ctx.turn_id,
-      })
-      if opts.on_messages_add then opts.on_messages_add({ msg }) end
-    end
-    if content_block.type == "tool_use" then
-      local complete_json = vim.json.decode(content_block.input_json)
-      local msg = HistoryMessage:new({
-        role = "assistant",
-        content = {
-          {
-            type = "tool_use",
-            name = content_block.name,
-            id = content_block.id,
-            input = complete_json or {},
-          },
-        },
-      }, {
-        state = "generated",
-        uuid = content_block.uuid,
-        turn_id = ctx.turn_id,
-      })
-      if opts.on_messages_add then opts.on_messages_add({ msg }) end
-    end
-    if content_block.type == "thinking" then
+      if opts.on_messages_add then
+        local msg = new_assistant_message(content_block.text, content_block.uuid)
+        opts.on_messages_add({ msg })
+      end
+    elseif content_block.type == "thinking" then
      if opts.on_chunk then
        if content_block.thinking and content_block.thinking ~= vim.NIL and content_block.thinking:sub(-1) ~= "\n" then
          opts.on_chunk("\n</think>\n\n")
@@ -307,21 +266,25 @@ function M:parse_response(ctx, data_stream, event_state, opts)
          opts.on_chunk("</think>\n\n")
        end
      end
-      local msg = HistoryMessage:new({
-        role = "assistant",
-        content = {
-          {
-            type = "thinking",
-            thinking = content_block.thinking,
-            signature = content_block.signature,
-          },
-        },
-      }, {
-        state = "generated",
-        uuid = content_block.uuid,
-        turn_id = ctx.turn_id,
-      })
-      if opts.on_messages_add then opts.on_messages_add({ msg }) end
+      if opts.on_messages_add then
+        local msg = new_assistant_message({
+          type = "thinking",
+          thinking = content_block.thinking,
+          signature = content_block.signature,
+        }, content_block.uuid)
+        opts.on_messages_add({ msg })
+      end
+    elseif content_block.type == "tool_use" then
+      if opts.on_messages_add then
+        local complete_json = vim.json.decode(content_block.input_json)
+        local msg = new_assistant_message({
+          type = "tool_use",
+          name = content_block.name,
+          id = content_block.id,
+          input = complete_json or {},
+        }, content_block.uuid)
+        opts.on_messages_add({ msg })
+      end
    end
  elseif event_state == "message_delta" then
    local ok, jsn = pcall(vim.json.decode, data_stream)
--- a/lua/avante/providers/ollama.lua
+++ b/lua/avante/providers/ollama.lua
@@ -135,27 +135,24 @@ function M:is_disable_stream() return false end
 ---@param tool_calls avante.OllamaToolCall[]
 ---@param opts AvanteLLMStreamOptions
 function M:add_tool_use_messages(tool_calls, opts)
-  local msgs = {}
-  for _, tool_call in ipairs(tool_calls) do
-    local id = Utils.uuid()
-    local func = tool_call["function"]
-    local msg = HistoryMessage:new({
-      role = "assistant",
-      content = {
-        {
-          type = "tool_use",
-          name = func.name,
-          id = id,
-          input = func.arguments,
-        },
-      },
-    }, {
-      state = "generated",
-      uuid = id,
-    })
-    table.insert(msgs, msg)
+  if opts.on_messages_add then
+    local msgs = {}
+    for _, tool_call in ipairs(tool_calls) do
+      local id = Utils.uuid()
+      local func = tool_call["function"]
+      local msg = HistoryMessage:new("assistant", {
+        type = "tool_use",
+        name = func.name,
+        id = id,
+        input = func.arguments,
+      }, {
+        state = "generated",
+        uuid = id,
+      })
+      table.insert(msgs, msg)
+    end
+    opts.on_messages_add(msgs)
  end
-  if opts.on_messages_add then opts.on_messages_add(msgs) end
 end

 function M:parse_stream_data(ctx, data, opts)
--- a/lua/avante/providers/openai.lua
+++ b/lua/avante/providers/openai.lua
@@ -231,10 +231,7 @@ function M:add_text_message(ctx, text, state, opts)
  local content =
    ctx.content:gsub("<tool_code>", ""):gsub("</tool_code>", ""):gsub("<tool_call>", ""):gsub("</tool_call>", "")
  ctx.content = content
-  local msg = HistoryMessage:new({
-    role = "assistant",
-    content = ctx.content,
-  }, {
+  local msg = HistoryMessage:new("assistant", ctx.content, {
    state = state,
    uuid = ctx.content_uuid,
    original_content = ctx.content,
@@ -299,16 +296,11 @@ function M:add_text_message(ctx, text, state, opts)
        has_tool_use = true
        local msg_uuid = ctx.content_uuid .. "-" .. idx
        local tool_use_id = msg_uuid
-        local msg_ = HistoryMessage:new({
-          role = "assistant",
-          content = {
-            {
-              type = "tool_use",
-              name = item.tool_name,
-              id = tool_use_id,
-              input = input,
-            },
-          },
+        local msg_ = HistoryMessage:new("assistant", {
+          type = "tool_use",
+          name = item.tool_name,
+          id = tool_use_id,
+          input = input,
        }, {
          state = state,
          uuid = msg_uuid,
@@ -342,15 +334,10 @@ end
 function M:add_thinking_message(ctx, text, state, opts)
  if ctx.reasonging_content == nil then ctx.reasonging_content = "" end
  ctx.reasonging_content = ctx.reasonging_content .. text
-  local msg = HistoryMessage:new({
-    role = "assistant",
-    content = {
-      {
-        type = "thinking",
-        thinking = ctx.reasonging_content,
-        signature = "",
-      },
-    },
+  local msg = HistoryMessage:new("assistant", {
+    type = "thinking",
+    thinking = ctx.reasonging_content,
+    signature = "",
  }, {
    state = state,
    uuid = ctx.reasonging_content_uuid,
@@ -362,16 +349,11 @@ end

 function M:add_tool_use_message(ctx, tool_use, state, opts)
  local jsn = JsonParser.parse(tool_use.input_json)
-  local msg = HistoryMessage:new({
-    role = "assistant",
-    content = {
-      {
-        type = "tool_use",
-        name = tool_use.name,
-        id = tool_use.id,
-        input = jsn or {},
-      },
-    },
+  local msg = HistoryMessage:new("assistant", {
+    type = "tool_use",
+    name = tool_use.name,
+    id = tool_use.id,
+    input = jsn or {},
  }, {
    state = state,
    uuid = tool_use.uuid,