refactor: message content (#1424)

This commit is contained in:
yetone
2025-02-28 14:43:39 +08:00
committed by GitHub
parent ebadba7420
commit ae8497faf1
10 changed files with 187 additions and 125 deletions

View File

@@ -16,74 +16,7 @@ M.role_map = {
assistant = "assistant",
}
function M.parse_messages(opts)
---@type AvanteBedrockClaudeMessage[]
local messages = {}
for _, message in ipairs(opts.messages) do
table.insert(messages, {
role = M.role_map[message.role],
content = {
{
type = "text",
text = message.content,
},
},
})
end
if opts.tool_histories then
for _, tool_history in ipairs(opts.tool_histories) do
if tool_history.tool_use then
local msg = {
role = "assistant",
content = {},
}
if tool_history.tool_use.thinking_contents then
for _, thinking_content in ipairs(tool_history.tool_use.thinking_contents) do
msg.content[#msg.content + 1] = {
type = "thinking",
thinking = thinking_content.content,
signature = thinking_content.signature,
}
end
end
if tool_history.tool_use.response_contents then
for _, response_content in ipairs(tool_history.tool_use.response_contents) do
msg.content[#msg.content + 1] = {
type = "text",
text = response_content,
}
end
end
msg.content[#msg.content + 1] = {
type = "tool_use",
id = tool_history.tool_use.id,
name = tool_history.tool_use.name,
input = vim.json.decode(tool_history.tool_use.input_json),
}
messages[#messages + 1] = msg
end
if tool_history.tool_result then
messages[#messages + 1] = {
role = "user",
content = {
{
type = "tool_result",
tool_use_id = tool_history.tool_result.tool_use_id,
content = tool_history.tool_result.content,
is_error = tool_history.tool_result.is_error,
},
},
}
end
end
end
return messages
end
M.parse_messages = Claude.parse_messages
M.parse_response = Claude.parse_response
---@param prompt_opts AvantePromptOptions

View File

@@ -43,7 +43,10 @@ function M.parse_messages(opts)
---@type {idx: integer, length: integer}[]
local messages_with_length = {}
for idx, message in ipairs(opts.messages) do
table.insert(messages_with_length, { idx = idx, length = Utils.tokens.calculate_tokens(message.content) })
table.insert(
messages_with_length,
{ idx = idx, length = Utils.tokens.calculate_message_content_tokens(message.content) }
)
end
table.sort(messages_with_length, function(a, b) return a.length > b.length end)
@@ -55,15 +58,46 @@ function M.parse_messages(opts)
end
for idx, message in ipairs(opts.messages) do
local content = {}
if type(message.content) == "string" then
table.insert(content, {
type = "text",
text = message.content,
cache_control = top_three[idx] and { type = "ephemeral" } or nil,
})
else
local message_content = message.content
---@cast message_content AvanteLLMMessageContentItem[]
for _, item in ipairs(message_content) do
if type(item) == "string" then
table.insert(content, {
type = "text",
text = item,
cache_control = top_three[idx] and { type = "ephemeral" } or nil,
})
elseif item.type == "text" then
table.insert(content, {
type = "text",
text = item.text,
cache_control = top_three[idx] and { type = "ephemeral" } or nil,
})
elseif item.type == "thinking" then
table.insert(content, {
type = "thinking",
thinking = item.thinking,
signature = item.signature,
})
elseif item.type == "redacted_thinking" then
table.insert(content, {
type = "redacted_thinking",
data = item.data,
})
end
end
end
table.insert(messages, {
role = M.role_map[message.role],
content = {
{
type = "text",
text = message.content,
cache_control = top_three[idx] and { type = "ephemeral" } or nil,
},
},
content = content,
})
end
@@ -89,12 +123,20 @@ function M.parse_messages(opts)
role = "assistant",
content = {},
}
if tool_history.tool_use.thinking_contents then
for _, thinking_content in ipairs(tool_history.tool_use.thinking_contents) do
if tool_history.tool_use.thinking_blocks then
for _, thinking_block in ipairs(tool_history.tool_use.thinking_blocks) do
msg.content[#msg.content + 1] = {
type = "thinking",
thinking = thinking_content.content,
signature = thinking_content.signature,
thinking = thinking_block.thinking,
signature = thinking_block.signature,
}
end
end
if tool_history.tool_use.redacted_thinking_blocks then
for _, thinking_block in ipairs(tool_history.tool_use.redacted_thinking_blocks) do
msg.content[#msg.content + 1] = {
type = "redacted_thinking",
data = thinking_block.data,
}
end
end
@@ -135,6 +177,8 @@ function M.parse_messages(opts)
end
function M.parse_response(ctx, data_stream, event_state, opts)
if ctx.resp_filename == nil then ctx.resp_filename = vim.fn.tempname() .. ".txt" end
vim.fn.writefile({ data_stream }, ctx.resp_filename, "a")
if event_state == nil then
if data_stream:match('"message_start"') then
event_state = "message_start"
@@ -208,24 +252,33 @@ function M.parse_response(ctx, data_stream, event_state, opts)
:filter(function(content_block_) return content_block_.stoppped and content_block_.type == "text" end)
:map(function(content_block_) return content_block_.text end)
:totable()
local thinking_contents = vim
local thinking_blocks = vim
.iter(ctx.content_blocks)
:filter(function(content_block_) return content_block_.stoppped and content_block_.type == "thinking" end)
:map(
function(content_block_)
return { content = content_block_.thinking, signature = content_block_.signature }
return { thinking = content_block_.thinking, signature = content_block_.signature }
end
)
:totable()
local redacted_thinking_blocks = vim
.iter(ctx.content_blocks)
:filter(
function(content_block_) return content_block_.stoppped and content_block_.type == "redacted_thinking" end
)
:map(function(content_block_) return { data = content_block_.data } end)
:totable()
return {
name = content_block.name,
id = content_block.id,
input_json = content_block.input_json,
response_contents = response_contents,
thinking_contents = thinking_contents,
thinking_blocks = thinking_blocks,
redacted_thinking_blocks = redacted_thinking_blocks,
}
end)
:totable()
Utils.debug("resp filename", ctx.resp_filename)
opts.on_stop({
reason = "tool_use",
usage = jsn.usage,

View File

@@ -33,9 +33,15 @@ function M.parse_messages(opts)
end
end
prev_role = role
table.insert(contents, { role = M.role_map[role] or role, parts = {
{ text = message.content },
} })
local parts = {}
for _, item in ipairs(message.content) do
if type(item) == "string" then
table.insert(parts, { text = item })
elseif item.type == "text" then
table.insert(parts, { text = item.text })
end
end
table.insert(contents, { role = M.role_map[role] or role, parts = parts })
end)
if Clipboard.support_paste_image() and opts.image_paths then

View File

@@ -71,18 +71,30 @@ function M.parse_messages(opts)
-- NOTE: Handle the case where the selected model is the `o1` model
-- "o1" models are "smart" enough to understand user prompt as a system prompt in this context
if M.is_o_series_model(base.model) then
table.insert(messages, { role = "user", content = opts.system_prompt })
table.insert(messages, { role = "user", content = { { type = "text", text = opts.system_prompt } } })
else
table.insert(messages, { role = "system", content = opts.system_prompt })
end
vim
.iter(opts.messages)
:each(function(msg) table.insert(messages, { role = M.role_map[msg.role], content = msg.content }) end)
vim.iter(opts.messages):each(function(msg)
if type(msg.content) == "string" then
table.insert(messages, { role = M.role_map[msg.role], content = msg.content })
else
local content = {}
for _, item in ipairs(msg.content) do
if type(item) == "string" then
table.insert(content, { type = "text", text = item })
elseif item.type == "text" then
table.insert(content, { type = "text", text = item.text })
end
end
table.insert(messages, { role = M.role_map[msg.role], content = content })
end
end)
if Config.behaviour.support_paste_from_clipboard and opts.image_paths and #opts.image_paths > 0 then
local message_content = messages[#messages].content
if type(message_content) ~= "table" then message_content = { type = "text", text = message_content } end
if type(message_content) == "string" then message_content = { { type = "text", text = message_content } } end
for _, image_path in ipairs(opts.image_paths) do
table.insert(message_content, {
type = "image_url",
@@ -103,7 +115,7 @@ function M.parse_messages(opts)
if role == M.role_map["user"] then
table.insert(final_messages, { role = M.role_map["assistant"], content = "Ok, I understand." })
else
table.insert(final_messages, { role = M.role_map["user"], content = "Ok" })
table.insert(final_messages, { role = M.role_map["user"], content = { { type = "text", text = "Ok" } } })
end
end
prev_role = role