Skip to content

Commit

Permalink
refactor: agentic workflows now use commands
Browse files Browse the repository at this point in the history
  • Loading branch information
@olimorris
olimorris committed Apr 4, 2024
1 parent 5a4c388 commit c13d40a
Show file tree
Hide file tree
Showing 2 changed files with 95 additions and 71 deletions.
48 changes: 46 additions & 2 deletions lua/codecompanion/agent.lua
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
local config = require("codecompanion.config")
local log = require("codecompanion.utils.log")

---@class CodeCompanion.Agent
local Agent = {}
Expand All @@ -15,6 +16,8 @@ end

---@param prompts table
function Agent:workflow(prompts)
log:trace("Initiating workflow")

local starting_prompts = {}
local workflow_prompts = {}

Expand All @@ -41,12 +44,53 @@ function Agent:workflow(prompts)
::continue::
end

return require("codecompanion.strategies.chat").new({
local function send_agent_prompt(chat)
log:trace("Sending agent prompt to chat buffer")

if #workflow_prompts == 0 then
return
end

local prompt = workflow_prompts[1]
chat:add_message(prompt)

if prompt.auto_submit then
chat:submit()
end

return table.remove(workflow_prompts, 1)
end

local chat = require("codecompanion.strategies.chat").new({
type = "chat",
messages = starting_prompts,
workflow = workflow_prompts,
show_buffer = true,
})

if not chat then
return
end

local group = vim.api.nvim_create_augroup("CodeCompanionAgent", {
clear = false,
})

vim.api.nvim_create_autocmd("User", {
desc = "Listen for CodeCompanion agent messages",
group = group,
pattern = "CodeCompanionChat",
callback = function(request)
if request.buf ~= chat.bufnr or request.data.status ~= "finished" then
return
end

send_agent_prompt(chat)

if #workflow_prompts == 0 then
vim.api.nvim_del_augroup_by_id(group)
end
end,
})
end

return Agent
118 changes: 49 additions & 69 deletions lua/codecompanion/strategies/chat.lua
View file
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,47 @@ local function render_messages(bufnr, adapter, settings, messages, context)
vim.bo[bufnr].modifiable = modifiable
end

local last_role = ""
---@param bufnr integer
---@param data table
---@return nil
local function render_new_messages(bufnr, data)
local total_lines = api.nvim_buf_line_count(bufnr)
local current_line = api.nvim_win_get_cursor(0)[1]
local cursor_moved = current_line == total_lines

local lines = {}
if data.role and data.role ~= last_role then
last_role = data.role
table.insert(lines, "")
table.insert(lines, "")
table.insert(lines, string.format("# %s", data.role))
table.insert(lines, "")
end

if data.content then
for _, text in ipairs(vim.split(data.content, "\n", { plain = true, trimempty = false })) do
table.insert(lines, text)
end

local modifiable = vim.bo[bufnr].modifiable
vim.bo[bufnr].modifiable = true

local last_line = api.nvim_buf_get_lines(bufnr, total_lines - 1, total_lines, false)[1]
local last_col = last_line and #last_line or 0
api.nvim_buf_set_text(bufnr, total_lines - 1, last_col, total_lines - 1, last_col, lines)

vim.bo[bufnr].modified = false
vim.bo[bufnr].modifiable = modifiable

if cursor_moved and ui.buf_is_active(bufnr) then
ui.buf_scroll_to_end(bufnr)
elseif not ui.buf_is_active(bufnr) then
ui.buf_scroll_to_end(bufnr)
end
end
end

local display_tokens = function(bufnr)
if config.options.display.chat.show_token_count then
require("codecompanion.utils.tokens").display_tokens(bufnr)
Expand Down Expand Up @@ -465,68 +506,6 @@ function Chat:submit()
vim.bo[self.bufnr].modifiable = true
end

local role = ""
local function render_new_messages(data)
local total_lines = api.nvim_buf_line_count(self.bufnr)
local current_line = api.nvim_win_get_cursor(0)[1]
local cursor_moved = current_line == total_lines

local lines = {}
if data.role and data.role ~= role then
role = data.role
table.insert(lines, "")
table.insert(lines, "")
table.insert(lines, string.format("# %s", data.role))
table.insert(lines, "")
end

if data.content then
for _, text in ipairs(vim.split(data.content, "\n", { plain = true, trimempty = false })) do
table.insert(lines, text)
end

local modifiable = vim.bo[self.bufnr].modifiable
vim.bo[self.bufnr].modifiable = true

local last_line = api.nvim_buf_get_lines(self.bufnr, total_lines - 1, total_lines, false)[1]
local last_col = last_line and #last_line or 0
api.nvim_buf_set_text(self.bufnr, total_lines - 1, last_col, total_lines - 1, last_col, lines)

vim.bo[self.bufnr].modified = false
vim.bo[self.bufnr].modifiable = modifiable

if cursor_moved and ui.buf_is_active(self.bufnr) then
ui.buf_scroll_to_end(self.bufnr)
elseif not ui.buf_is_active(self.bufnr) then
ui.buf_scroll_to_end(self.bufnr)
end
end
end

local function apply_workflow()
local i = 0
for _, prompt in ipairs(self.workflow) do
i = i + 1
if not prompt.status then
log:trace("Workflow: %s", prompt)
prompt.status = "done"
render_new_messages({ role = prompt.role, content = prompt.content })

if prompt.auto_submit then
log:trace("Auto-submitting")
self:submit()
end

return
end

-- At the end of the prompts, ensure the chat buffer is ready for the user
if i == #self.workflow then
return render_new_messages({ role = "user", content = "" })
end
end
end

local current_message = messages[#messages]

if current_message and current_message.role == "user" and current_message.content == "" then
Expand All @@ -543,20 +522,17 @@ function Chat:submit()
end

if done then
if self.workflow then
apply_workflow()
else
render_new_messages({ role = "user", content = "" })
end
render_new_messages(self.bufnr, { role = "user", content = "" })
display_tokens(self.bufnr)
return finalize()
finalize()
return api.nvim_exec_autocmds("User", { pattern = "CodeCompanionChat", data = { status = "finished" } })
end

if data then
local result = self.adapter.callbacks.chat_output(data)

if result and result.status == "success" then
render_new_messages(result.output)
render_new_messages(self.bufnr, result.output)
elseif result and result.status == "error" then
vim.api.nvim_exec_autocmds(
"User",
Expand Down Expand Up @@ -612,6 +588,10 @@ function Chat:on_cursor_moved()
end
end

function Chat:add_message(message)
render_new_messages(self.bufnr, { role = message.role, content = message.content })
end

function Chat:complete(request, callback)
local items = {}
local cursor = api.nvim_win_get_cursor(0)
Expand Down

0 comments on commit c13d40a

Please sign in to comment.