Refactoring of context and request entry into a single thing prompt

1 files changed, 415 insertions, 0 deletions

diff --git a/lua/99/prompt.lua b/lua/99/prompt.lua
new file mode 100644
index 0000000..d4f5095
--- /dev/null
+++ b/lua/99/prompt.lua
@@ -0,0 +1,415 @@
+local BaseProvider = require("99.providers")
+local Logger = require("99.logger.logger")
+local utils = require("99.utils")
+local random_file = utils.random_file
+local copy = utils.copy
+local get_id = require("99.id")
+local Range = require("99.geo").Range
+local Time = require("99.time")
+
+local filetype_map = {
+  typescriptreact = "typescript",
+}
+
+-- luacheck: ignore
+--- @alias _99.Prompt.Data _99.Prompt.Data.Search | _99.Prompt.Data.Tutorial | _99.Prompt.Data.Visual
+--- @alias _99.Prompt.Operation "visual" | "tutorial" | "search"
+--- @alias _99.Prompt.EndingState "failed" | "success" | "cancelled"
+--- @alias _99.Prompt.State "ready" | "requesting" | _99.Prompt.EndingState
+--- @alias _99.Prompt.Cleanup fun(): nil
+
+--- @class _99.Prompt.Data.Search
+--- @field type "search"
+--- @field qfix_items _99.Search.Result[]
+
+--- @class _99.Prompt.Data.Visual
+--- @field type "visual"
+--- @field buffer number
+--- @field file_type string
+--- @field range _99.Range
+
+--- @class _99.Prompt.Data.Tutorial
+--- @field type "tutorial"
+--- @field buffer number
+--- @field window number
+--- @field xid number TODO: we should probably get rid of this.  The request pattern is not quite correct
+--- @field tutorial string[]
+
+--- @class _99.Prompt
+--- @field md_file_names string[]
+--- @field model string
+--- @field operation _99.Prompt.Operation
+--- @field state _99.Prompt.State
+--- @field full_path string
+--- @field started_at number
+--- @field data _99.Prompt.Data
+--- @field agent_context string[]
+--- @field tmp_file string
+--- @field marks table<string, _99.Mark>
+--- @field logger _99.Logger
+--- @field xid number
+--- @field clean_ups (fun(): nil)[]
+--- @field _99 _99.State
+---@diagnostic disable-next-line: undefined-doc-name
+--- @field _proc vim.SystemObj?
+local Prompt = {}
+Prompt.__index = Prompt
+
+--- @type _99.Prompt[]
+Prompt.__previous_contexts = {}
+
+--- @type table<number, _99.Prompt>
+Prompt.__context_by_id = {}
+
+--- @param context  _99.Prompt
+--- @param _99 _99.State
+local function set_defaults(context, _99)
+  local xid = get_id()
+  local full_path = vim.api.nvim_buf_get_name(0)
+
+  context.state = "ready"
+  context._99 = _99
+  context.clean_ups = {}
+  context.md_file_names = copy(_99.md_files)
+  context.model = _99.model
+  context.agent_context = {}
+  context.tmp_file = random_file(_99:tmp_dir())
+  context.logger = Logger:set_id(xid)
+  context.xid = xid
+  context.full_path = full_path
+  context.marks = {}
+  context.started_at = Time.now()
+end
+
+--- TODO: Work item for "TODO implementation"
+function Prompt.todo(_99)
+  _ = _99
+  assert(false, "not implemented")
+end
+
+function Prompt.vibe(_99, opts)
+  _ = _99
+  _ = opts
+  assert(false, "not implemented")
+end
+
+--- @param _99 _99.State
+--- @return _99.Prompt
+function Prompt.visual(_99)
+  _99:refresh_rules()
+
+  local file_type = vim.bo[0].ft
+  local buffer = vim.api.nvim_get_current_buf()
+  file_type = filetype_map[file_type] or file_type
+
+  local mds = {}
+  for _, md in ipairs(_99.md_files) do
+    table.insert(mds, md)
+  end
+
+  --- @type _99.Prompt
+  local context = setmetatable({}, Prompt)
+  set_defaults(context, _99)
+  context.operation = "visual"
+  context.data = {
+    type = "visual",
+    buffer = buffer,
+    file_type = file_type,
+    range = Range.zero(),
+  }
+  context.logger:debug("99 Request", "method", "visual")
+
+  return context
+end
+
+--- @param _99 _99.State
+--- @return _99.Prompt
+function Prompt.tutorial(_99)
+  _99:refresh_rules()
+
+  --- @type _99.Prompt
+  local context = setmetatable({}, Prompt)
+  set_defaults(context, _99)
+  context.operation = "tutorial"
+  context.data = {
+    type = "tutorial",
+    xid = context.xid, -- TODO: i want to get rid of this when i implement rehydration of the data.
+    buffer = 0,
+    window = 0,
+    tutorial = {},
+  }
+  context.logger:debug("99 Request", "method", "tutorial")
+
+  return context
+end
+
+--- @param _99 _99.State
+--- @return _99.Prompt
+function Prompt.search(_99)
+  _99:refresh_rules()
+
+  --- @type _99.Prompt
+  local context = setmetatable({}, Prompt)
+  set_defaults(context, _99)
+  context.operation = "search"
+  context.data = {
+    type = "search",
+    qfix_items = {},
+  }
+  context.logger:debug("99 Request", "method", "search")
+
+  return context
+end
+
+--- @param obs _99.Providers.Observer | nil
+function Prompt:_observer(obs)
+  return {
+    on_start = function()
+      self.state = "requesting"
+      self._99:track_request(self)
+
+      if obs then
+        obs.on_start()
+      end
+    end,
+    on_complete = function(status, res)
+      self.state = status
+      if obs then
+        obs.on_complete(status, res)
+      end
+    end,
+    on_stderr = function(line)
+      if obs then
+        obs.on_stderr(line)
+      end
+    end,
+    on_stdout = function(line)
+      if obs then
+        obs.on_stdout(line)
+      end
+    end,
+  }
+end
+
+--- @return boolean
+function Prompt:valid()
+  local t = self.data.type
+  return t == "visual" or t == "search" or t == "tutorial"
+end
+
+--- @param observer _99.Providers.Observer?
+function Prompt:start_request(observer)
+  local l = self.logger
+  l:assert(
+    self.state == "ready",
+    'state is not "ready" when attempting to start a request'
+  )
+
+  local ok = self:finalize()
+  l:assert(ok, "context failed to finalize")
+
+  --- TODO: create a prompt context class that can actually organize.
+  --- do not do this during the request context refactoring, but next
+  local prompt = table.concat(self.agent_context, "\n")
+  local obs = self:_observer(observer)
+  local provider = self._99.provider_override or BaseProvider.OpenCodeProvider
+
+  self:save_prompt(prompt)
+  l:debug("start", "prompt", prompt)
+
+  provider:make_request(prompt, self, obs)
+end
+
+function Prompt:is_cancelled()
+  return self.state == "cancelled"
+end
+
+---@diagnostic disable-next-line: undefined-doc-name
+--- @param proc vim.SystemObj?
+function Prompt:_set_process(proc)
+  self._proc = proc
+end
+
+function Prompt:cancel()
+  if self:is_cancelled() then
+    return
+  end
+
+  self.logger:debug("cancel")
+  self.state = "cancelled"
+  local proc = self._proc
+  ---@diagnostic disable-next-line: undefined-field
+  if proc and proc.pid then
+    self._proc = nil
+    pcall(function()
+      local sigterm = (vim.uv and vim.uv.constants and vim.uv.constants.SIGTERM)
+        or 15
+      ---@diagnostic disable-next-line: undefined-field
+      proc:kill(sigterm)
+    end)
+  end
+end
+
+--- @return _99.Prompt.Data.Visual
+function Prompt:visual_data()
+  assert(
+    self.data.type == "visual",
+    "you cannot get visual data if its not type visual"
+  )
+  return self.data --[[@as _99.Prompt.Data.Visual]]
+end
+
+--- @return _99.Prompt.Data.Tutorial
+function Prompt:tutorial_data()
+  assert(
+    self.data.type == "tutorial",
+    "you cannot get tutorial data if its not type tutorial"
+  )
+  return self.data --[[@as _99.Prompt.Data.Tutorial]]
+end
+
+--- @return _99.Prompt.Data.Search
+function Prompt:search_data()
+  assert(
+    self.data.type == "search",
+    "you cannot get search data if its not type search"
+  )
+  return self.data --[[@as _99.Prompt.Data.Search]]
+end
+
+function Prompt:stop()
+  self:cancel()
+  for _, cb in ipairs(self.clean_ups) do
+    cb()
+  end
+end
+
+--- @param clean_up fun(): nil
+function Prompt:add_clean_up(clean_up)
+  table.insert(self.clean_ups, clean_up)
+end
+
+--- @param md_file_name string
+--- @return self
+function Prompt:add_md_file_name(md_file_name)
+  table.insert(self.md_file_names, md_file_name)
+  return self
+end
+
+--- @param content string
+--- @return self
+function Prompt:add_prompt_content(content)
+  table.insert(self.agent_context, content)
+  return self
+end
+
+--- @param refs _99.Reference[]
+function Prompt:add_references(refs)
+  for _, ref in ipairs(refs) do
+    self.logger:debug("adding reference to context")
+    table.insert(self.agent_context, ref.content)
+  end
+end
+
+function Prompt:_read_md_files()
+  local cwd = vim.uv.cwd()
+  local dir = vim.fn.fnamemodify(self.full_path, ":h")
+
+  while dir:find(cwd, 1, true) == 1 do
+    for _, md_file_name in ipairs(self.md_file_names) do
+      local md_path = dir .. "/" .. md_file_name
+      local file = io.open(md_path, "r")
+      if file then
+        local content = file:read("*a")
+        file:close()
+        self.logger:info(
+          "Context#adding md file to the context",
+          "md_path",
+          md_path
+        )
+        table.insert(self.agent_context, content)
+      end
+    end
+
+    if dir == cwd then
+      break
+    end
+
+    dir = vim.fn.fnamemodify(dir, ":h")
+  end
+end
+
+--- @return string[]
+function Prompt:content()
+  return self.agent_context
+end
+
+--- @return boolean
+function Prompt:_ready_request_files()
+  local response_file = self.tmp_file
+  local prompt_file = self.tmp_file .. "-prompt"
+
+  local dir = vim.fs.dirname(prompt_file)
+
+  if dir and not vim.uv.fs_stat(dir) then
+    vim.fn.mkdir(dir, "p")
+  end
+
+  local files = { prompt_file, response_file }
+  for _, f in ipairs(files) do
+    local file = io.open(f, "w")
+    if file then
+      file:write("")
+      file:close()
+    else
+      self.logger:error("unable to create prompt file")
+      return false
+    end
+  end
+  return true
+end
+
+--- @param prompt string
+function Prompt:save_prompt(prompt)
+  local prompt_file = self.tmp_file .. "-prompt"
+  local file = io.open(prompt_file, "w")
+  if file then
+    file:write(prompt)
+    file:close()
+    self.logger:debug("saved prompt to file", "path", prompt_file)
+  else
+    self.logger:error("failed to save prompt", "path", prompt_file)
+  end
+end
+
+--- @return boolean, self
+function Prompt:finalize()
+  if self:_ready_request_files() == false then
+    return false, self
+  end
+  self:_read_md_files()
+
+  local ok, visual_data = pcall(self.visual_data, self)
+  if ok then
+    local f_loc =
+      self._99.prompts.get_file_location(self.full_path, visual_data.range)
+    table.insert(self.agent_context, f_loc)
+    table.insert(
+      self.agent_context,
+      self._99.prompts.get_range_text(visual_data.range)
+    )
+  end
+  table.insert(
+    self.agent_context,
+    self._99.prompts.tmp_file_location(self.tmp_file)
+  )
+  return true, self
+end
+
+function Prompt:clear_marks()
+  for _, mark in pairs(self.marks) do
+    mark:delete()
+  end
+end
+
+return Prompt