chatbar.lua

local M = {}
local U = require('vigpt.utils')

-- These variables are managed here
M.chatbar_buffer = nil
M.chatbar_window = nil
M.active_buffer = nil

function M.toggle_chatbar(selected_llm)
  if M.chatbar_window and vim.api.nvim_win_is_valid(M.chatbar_window) then
    vim.api.nvim_win_close(M.chatbar_window, true)
    M.chatbar_window = nil
    print("ChatBar closed.")
  else
    M.active_buffer = vim.api.nvim_get_current_buf()

    if not (M.chatbar_buffer and vim.api.nvim_buf_is_valid(M.chatbar_buffer)) then
      M.chatbar_buffer = vim.api.nvim_create_buf(false, true)
      vim.api.nvim_buf_set_name(M.chatbar_buffer, "ChatBar")
      vim.bo[M.chatbar_buffer].buftype = "nofile"
      vim.bo[M.chatbar_buffer].bufhidden = "hide"
      vim.bo[M.chatbar_buffer].swapfile = false
      vim.bo[M.chatbar_buffer].filetype = "markdown"
    end

    vim.cmd("vsplit")
    vim.cmd("wincmd L")
    vim.cmd("vertical resize 71")
    M.chatbar_window = vim.api.nvim_get_current_win()
    vim.api.nvim_win_set_buf(M.chatbar_window, M.chatbar_buffer)

    vim.wo.number = false
    vim.wo.relativenumber = false
    vim.wo.wrap = true
    vim.wo.signcolumn = "no"
    vim.wo.cursorline = false
    vim.wo[M.chatbar_window].statusline = "  Model: " .. selected_llm

    vim.api.nvim_buf_set_keymap(M.chatbar_buffer, "n", "gb", "<nop>", { noremap = true, silent = true })
    vim.api.nvim_buf_set_keymap(M.chatbar_buffer, "n", "gB", "<nop>", { noremap = true, silent = true })
    vim.api.nvim_buf_set_keymap(M.chatbar_buffer, "n", "j", "gj", { noremap = true, silent = true })
    vim.api.nvim_buf_set_keymap(M.chatbar_buffer, "n", "k", "gk", { noremap = true, silent = true })
    vim.api.nvim_buf_set_keymap(M.chatbar_buffer, "n", "<Esc>", "", {
      noremap = true,
      silent = true,
      callback = function()
        if M.chatbar_window and vim.api.nvim_win_is_valid(M.chatbar_window) then
          vim.api.nvim_win_close(M.chatbar_window, true)
          M.chatbar_window = nil
          print("ChatBar closed.")
        end
      end
    })

    -- <C-CR> will be mapped by main module to M.send_to_llm

    local current_line_count = vim.api.nvim_buf_line_count(M.chatbar_buffer)
    if current_line_count == 1 and vim.api.nvim_buf_get_lines(M.chatbar_buffer, 0, -1, false)[1] == "" then
      vim.api.nvim_buf_set_lines(M.chatbar_buffer, 0, -1, false, { "USER »", "" })
      vim.api.nvim_win_set_cursor(M.chatbar_window, {2, 0})
    end

    print("ChatBar Opened (Markdown highlighting)")
  end
end

function M.clear_chatbar()
  if M.chatbar_buffer and vim.api.nvim_buf_is_valid(M.chatbar_buffer) then
    vim.api.nvim_buf_set_lines(M.chatbar_buffer, 0, -1, false, {})
    print("ChatBar cleared.")
    vim.api.nvim_buf_set_lines(M.chatbar_buffer, 0, -1, false, { "USER »", "" })
    vim.api.nvim_win_set_cursor(M.chatbar_window, {2, 0})
  else
    print("ChatBar is not open!")
  end
end

return M

init.lua

local M = {}

local chatbar = require('vigpt.chatbar')
local system_msg = require('vigpt.system_message')
local llm_select = require('vigpt.llm_select')
local request = require('vigpt.request')

-- Expose some variables so we can dynamically call them
M.selected_llm = llm_select.selected_llm
M.system_message = system_msg.system_message

function M.toggle_chatbar()
  chatbar.toggle_chatbar(llm_select.selected_llm)
end

function M.clear_chatbar()
  chatbar.clear_chatbar()
end

function M.edit_system_message_floating()
  system_msg.edit_system_message_floating()
end

function M.select_llm_floating()
  llm_select.select_llm_floating(chatbar.chatbar_window)
end

function M.send_to_llm()
  request.send_to_llm(llm_select.selected_llm, system_msg.system_message, chatbar.chatbar_buffer, chatbar.chatbar_window, chatbar.active_buffer)
end

return M

llm_select.lua

local M = {}
local llm_options = { 
  "llama 3.1 8b", 
  "o1-mini", 
  "gpt-4o-mini",
  "claude-3-5-sonnet-latest",
  "claude-3-5-haiku-latest"
}

M.selected_llm = "llama 3.1 8b"
M.original_llm = M.selected_llm
M.llm_win = nil
M.llm_buf = nil

-- Update the chatbar statusline if needed (we'll call back into main M)
local function update_chatbar_statusline(chatbar_window, llm)
  if chatbar_window and vim.api.nvim_win_is_valid(chatbar_window) then
    vim.wo[chatbar_window].statusline = "  Model: " .. llm
  end
end

function M.select_llm_floating(chatbar_window)
  if M.llm_win and vim.api.nvim_win_is_valid(M.llm_win) then
    -- Confirm current highlighted model
    local cursor = vim.api.nvim_win_get_cursor(M.llm_win)
    local line = cursor[1]
    if llm_options[line] then
      M.selected_llm = llm_options[line]
      print("Selected LLM:", M.selected_llm)
    else
      M.selected_llm = M.original_llm
      print("No model selected, reverting to:", M.selected_llm)
    end
    update_chatbar_statusline(chatbar_window, M.selected_llm)
    vim.api.nvim_win_close(M.llm_win, true)
    M.llm_win = nil
    M.llm_buf = nil
    return
  end

  M.original_llm = M.selected_llm
  M.llm_buf = vim.api.nvim_create_buf(false, true)
  vim.api.nvim_buf_set_option(M.llm_buf, 'bufhidden', 'wipe')
  vim.api.nvim_buf_set_option(M.llm_buf, 'filetype', 'llm_menu')

  vim.api.nvim_buf_set_lines(M.llm_buf, 0, -1, false, llm_options)
  local width = 30
  local height = #llm_options
  local opts = {
    relative = 'editor',
    width = width,
    height = height,
    col = (vim.o.columns - width)/2,
    row = (vim.o.lines - height)/2,
    style = 'minimal',
    border = 'rounded'
  }

  M.llm_win = vim.api.nvim_open_win(M.llm_buf, true, opts)
  vim.wo[M.llm_win].cursorline = true
  vim.wo[M.llm_win].wrap = false
  vim.wo[M.llm_win].winhighlight = "Normal:Normal,FloatBorder:FloatBorder"

  local function confirm_selection_and_close()
    local cursor = vim.api.nvim_win_get_cursor(M.llm_win)
    local line = cursor[1]
    local chosen = llm_options[line]
    if chosen then
      M.selected_llm = chosen
      print("Selected LLM:", chosen)
      update_chatbar_statusline(chatbar_window, M.selected_llm)
    else
      M.selected_llm = M.original_llm
      print("No valid model selected, reverting to:", M.selected_llm)
    end
    vim.api.nvim_win_close(M.llm_win, true)
    M.llm_win = nil
    M.llm_buf = nil
  end

  local function discard_changes()
    M.selected_llm = M.original_llm
    print("LLM selection discarded, reverting to:", M.selected_llm)
    vim.api.nvim_win_close(M.llm_win, true)
    M.llm_win = nil
    M.llm_buf = nil
  end

  vim.api.nvim_buf_set_keymap(M.llm_buf, 'n', '<CR>', '', { noremap = true, silent = true, callback = confirm_selection_and_close })
  vim.api.nvim_buf_set_keymap(M.llm_buf, 'n', '<Esc>', '', { noremap = true, silent = true, callback = discard_changes })
end

return M

messages.lua

local U = require('vigpt.utils')

local M = {}

-- Parse the ChatBar into structured messages
function M.parse_chatbar_messages(chatbar_buffer)
  if not (chatbar_buffer and vim.api.nvim_buf_is_valid(chatbar_buffer)) then
    return {}
  end

  local lines = vim.api.nvim_buf_get_lines(chatbar_buffer, 0, -1, false)
  local messages = {}
  local current_role = nil
  local current_content = {}

  local function flush_message()
    if current_role and #current_content > 0 then
      local msg_content = table.concat(current_content, "\n"):match("^%s*(.-)%s*$") or ""
      if msg_content ~= "" then
        table.insert(messages, { role = current_role, content = msg_content })
      end
      current_content = {}
    end
  end

  for _, line in ipairs(lines) do
    if line:match("^USER »") then
      flush_message()
      current_role = "user"
    elseif line:match("^LLM »") then
      flush_message()
      current_role = "assistant"
    else
      if current_role then
        table.insert(current_content, line)
      end
    end
  end
  flush_message()

  return messages
end

-- Replace @buffer and @all references
function M.replace_references_in_message(msg, active_buffer)
  local content = msg.content
  local modified_lines = {}

  for line in content:gmatch("([^\n]*)\n?") do
    if line:find("@buffer") then
      local formatted = U.fetch_and_format_buffer(active_buffer)
      if formatted then
        local before, after = line:match("(.-)@buffer(.*)")
        if before and before ~= "" then table.insert(modified_lines, before) end
        vim.list_extend(modified_lines, formatted)
        if after and after ~= "" then table.insert(modified_lines, after) end
      else
        table.insert(modified_lines, line)
      end
    elseif line:find("@all") then
      local before, after = line:match("(.-)@all(.*)")
      if before and before ~= "" then table.insert(modified_lines, before) end
      for _, buffer in ipairs(vim.api.nvim_list_bufs()) do
        if vim.api.nvim_buf_is_loaded(buffer) and vim.api.nvim_buf_get_option(buffer, "buftype") == "" then
          local formatted = U.fetch_and_format_buffer(buffer)
          if formatted then
            vim.list_extend(modified_lines, formatted)
          end
        end
      end
      if after and after ~= "" then table.insert(modified_lines, after) end
    else
      table.insert(modified_lines, line)
    end
  end

  msg.content = table.concat(modified_lines, "\n")
end

return M

models.lua

local M = {}

-- Common parse_response for OpenAI streaming
local function parse_openai_response_line(line, append_to_chatbar_fn)
  local sse_data = line:match("^data:%s*(.*)")
  if sse_data and sse_data ~= "[DONE]" then
    local ok, json_data = pcall(vim.fn.json_decode, sse_data)
    if ok and json_data and json_data.choices and json_data.choices[1] then
      local delta = json_data.choices[1].delta or {}
      if delta.content then
        append_to_chatbar_fn(delta.content)
      end
    end
  elseif sse_data == "[DONE]" then
    -- End of stream
  else
    -- If not SSE, could append raw line for debugging
    append_to_chatbar_fn(line)
  end
end

M.registry = {

  -- Llama 3.1 8B
  ["llama 3.1 8b"] = {
    -- For local model, you might have a different endpoint and parsing logic
    api_type = "local",
    endpoint = "http://localhost:1234/v1/chat/completions",
    headers_fn = function()
      return { ["Content-Type"] = "application/json" }
    end,
    transform_request_fn = function(system_message, messages)
      return {
        model = "lmstudio-community/Meta-Llama-3.1-8B-Instruct-GGUF",
        messages = messages,
        -- temperature = 0.7,
        max_tokens = -1,
        stream = true
      }
    end,
    parse_response_fn = function(line, append_fn)
      -- SSE parsing logic for local model here
      local sse_data = line:match("^data:%s*(.*)")
      if sse_data then
        local ok, json_data = pcall(vim.fn.json_decode, sse_data)
        if ok and json_data and json_data.choices and json_data.choices[1] then
          local delta = json_data.choices[1].delta or {}
          if delta.content then
            append_fn(delta.content)
          end
        end
      else
        append_fn(line)
      end
    end
  },

  -- ChatGPT o1-mini
  ["o1-mini"] = {
    api_type = "openai",
    endpoint = "https://api.openai.com/v1/chat/completions",
    headers_fn = function()
      return {
        ["Content-Type"] = "application/json",
        ["Authorization"] = "Bearer " .. vim.fn.getenv("NVIM_OPENAI") or ""
      }
    end,
    transform_request_fn = function(system_message, messages)
      return {
        model = "o1-mini",
        messages = messages,
        -- temperature = 1.0,
        stream = true,
      }
    end,
    parse_response_fn = parse_openai_response_line
  },

  -- ChatGPT 4o-mini
  ["gpt-4o-mini"] = {
    api_type = "openai",
    endpoint = "https://api.openai.com/v1/chat/completions",
    headers_fn = function()
      return {
        ["Content-Type"] = "application/json",
        ["Authorization"] = "Bearer " .. vim.fn.getenv("NVIM_OPENAI") or ""
      }
    end,
    transform_request_fn = function(system_message, messages)
      return {
        model = "gpt-4o-mini",
        messages = messages,
        -- temperature = 1.0,
        stream = true,
      }
    end,
    parse_response_fn = parse_openai_response_line
  },

  -- Claude 3.5 Sonnet
  ["claude-3-5-sonnet-latest"] = {
    api_type = "anthropic",
    endpoint = "https://api.anthropic.com/v1/messages",
    headers_fn = function()
      local api_key = vim.fn.getenv("NVIM_ANTHROPIC")
      return {
        ["Content-Type"] = "application/json",
        ["x-api-key"] = api_key,
        ["anthropic-version"] = "2023-06-01",
        ["accept"] = "text/event-stream"
      }
    end,
    transform_request_fn = function(system_message, messages)
      local transformed_messages = {}
      for _, msg in ipairs(messages) do
        if msg.role ~= "system" then
          table.insert(transformed_messages, {
            role = msg.role == "assistant" and "assistant" or "user",
            content = msg.content
          })
        end
      end

      -- Remove max_tokens, it's not needed
      local request = {
        model = "claude-3-5-sonnet-latest",
        stream = true,
        system = system_message,
        messages = transformed_messages,
        max_tokens = 4096
      }
      return request
    end,
    parse_response_fn = function(line, append_to_chatbar_fn)
      if line:match("^event:") then
        return
      end

      local sse_data = line:match("^data:%s*(.*)")
      if sse_data and sse_data ~= "[DONE]" then
        local ok, json_data = pcall(vim.fn.json_decode, sse_data)
        if ok and json_data and json_data.delta and json_data.delta.text then
          append_to_chatbar_fn(json_data.delta.text)
        end
      end
    end
  },

  -- Claude 3.5 Haiku
  ["claude-3-5-haiku-latest"] = {
    api_type = "anthropic",
    endpoint = "https://api.anthropic.com/v1/messages",
    headers_fn = function()
      return {
        ["Content-Type"] = "application/json",
        ["x-api-key"] = vim.fn.getenv("NVIM_ANTHROPIC"),
        ["anthropic-version"] = "2023-06-01",  -- Updated API version
        ["Accept"] = "text/event-stream"
      }
    end,
    transform_request_fn = function(system_message, messages)
      -- Transform messages to Claude's format
      local transformed_messages = {}
      for _, msg in ipairs(messages) do
        if msg.role ~= "system" then
          table.insert(transformed_messages, {
            role = msg.role == "assistant" and "assistant" or "user",
            content = msg.content
          })
        end
      end

      return {
        model = "claude-3-5-haiku-latest",  -- Updated model name
        max_tokens = 1024,
        stream = true,
        system = system_message,
        messages = transformed_messages,
        max_tokens = 4096
      }
    end,
    parse_response_fn = function(line, append_to_chatbar_fn)
      if line:match("^event:") then
        return
      end

      local sse_data = line:match("^data:%s*(.*)")
      if sse_data and sse_data ~= "[DONE]" then
        local ok, json_data = pcall(vim.fn.json_decode, sse_data)
        if ok and json_data and json_data.delta and json_data.delta.text then
          append_to_chatbar_fn(json_data.delta.text)
        end
      end
    end
  },

}

return M

system_message.lua

local M = {}

M.system_message = [[
You are an expert at JavaScript in the browser as well as Deno (the server-side JavaScript runtime). Help me write JavaScript code snippets for various purposes, following these strict style rules:

1. JavaScript only: Never use TypeScript.
2. No semicolons: Never put semicolons at the ends of lines. Only use them in for loops (e.g. for (var i = 0; i < 10; i++) {}) or where strictly necessary (e.g. IIFE).
3. No arrow functions: Avoid arrow functions. Use them only in rare cases where you need a short one-line function with an implied return.
4. var not const: Default to using var for declarations and let for block scoping in loops. Never use const.
5. Space in function declarations: Put a space between the function name and the opening parenthesis in a function declaration, and between the function keyword and the opening parenthesis in an anonymous function. For example:

function myFunc (x) {
  // ...
}
var someFunc = function (y) {
  // ...
}

6. Concise code: Write the fewest lines possible. Omit extra whitespace or verbose syntax.
7. Use async/await: Prefer async/await over Promise chaining methods (.then(), .catch()).
8. Comment functions: Above each function, write a one-line comment describing: The input and return value for pure functions & The side effects for any functions that produce side effects

Follow Kyle Simpson's Functional Lite JavaScript principles. For code that is purely functional, describe its inputs and returns. For code that causes side effects (like reading or writing files), note the side effects.
]]

M.original_system_message = M.system_message
M.system_win = nil
M.system_buf = nil

function M.edit_system_message_floating()
  if M.system_win and vim.api.nvim_win_is_valid(M.system_win) then
    -- Confirm changes
    local new_lines = vim.api.nvim_buf_get_lines(M.system_buf, 0, -1, false)
    if #new_lines > 0 then
      M.system_message = table.concat(new_lines, "\n")
      print("System message updated:", M.system_message)
    end
    vim.api.nvim_win_close(M.system_win, true)
    M.system_win = nil
    M.system_buf = nil
    return
  end

  M.original_system_message = M.system_message
  M.system_buf = vim.api.nvim_create_buf(false, true)
  vim.api.nvim_buf_set_option(M.system_buf, 'bufhidden', 'wipe')
  vim.api.nvim_buf_set_option(M.system_buf, 'filetype', 'system_msg')

  local lines = vim.split(M.system_message, "\n", { plain = true })
  vim.api.nvim_buf_set_lines(M.system_buf, 0, -1, false, lines)

  local width = 96
  local height = 48
  local opts = {
    relative = 'editor',
    width = width,
    height = height,
    col = (vim.o.columns - width)/2,
    row = (vim.o.lines - height)/2,
    style = 'minimal',
    border = 'rounded'
  }

  M.system_win = vim.api.nvim_open_win(M.system_buf, true, opts)
  vim.wo[M.system_win].cursorline = true
  vim.wo[M.system_win].wrap = true
  vim.wo[M.system_win].linebreak = true
  vim.wo[M.system_win].winhighlight = "Normal:Normal,FloatBorder:FloatBorder"
  vim.api.nvim_buf_set_keymap(
    M.system_buf,
    'n',
    'j',
    'gj',
    { noremap = true, silent = true }
  )

  vim.api.nvim_buf_set_keymap(
    M.system_buf,
    'n',
    'k',
    'gk',
    { noremap = true, silent = true }
  )


  local function confirm_system_message()
    local new_lines = vim.api.nvim_buf_get_lines(M.system_buf, 0, -1, false)
    if #new_lines > 0 then
      M.system_message = table.concat(new_lines, "\n")
      print("System message updated:", M.system_message)
    end
    vim.api.nvim_win_close(M.system_win, true)
    M.system_win = nil
    M.system_buf = nil
  end

  local function discard_changes()
    M.system_message = M.original_system_message
    vim.api.nvim_win_close(M.system_win, true)
    M.system_win = nil
    M.system_buf = nil
  end

  vim.api.nvim_buf_set_keymap(M.system_buf, 'n', '<CR>', '', { noremap = true, silent = true, callback = confirm_system_message })
  vim.api.nvim_buf_set_keymap(M.system_buf, 'n', '<Esc>', '', { noremap = true, silent = true, callback = discard_changes })
end

return M

utils.lua

local U = {}

function U.append_to_chatbar(chatbar_buffer, lines)
  if chatbar_buffer and vim.api.nvim_buf_is_valid(chatbar_buffer) then
    local current_lines = vim.api.nvim_buf_line_count(chatbar_buffer)
    vim.api.nvim_buf_set_lines(chatbar_buffer, current_lines, current_lines, false, lines)
  end
end

function U.fetch_and_format_buffer(buffer)
  if not vim.api.nvim_buf_is_valid(buffer) then return nil end
  local filename = vim.api.nvim_buf_get_name(buffer)
  filename = filename ~= "" and vim.fn.fnamemodify(filename, ":t") or "[No Name]"
  local content = vim.api.nvim_buf_get_lines(buffer, 0, -1, false)

  local formatted = {
    "", -- Blank line before filename
    "_" .. filename .. "_", -- Underscores around the filename
    "", -- Blank line after filename
    "```", -- Opening backticks for content
  }
  vim.list_extend(formatted, content)
  table.insert(formatted, "```") -- Closing backticks
  table.insert(formatted, "") -- Final blank line after content
  return formatted
end

function U.count_buf_lines(bufnr)
  if vim.api.nvim_buf_is_valid(bufnr) then
    return vim.api.nvim_buf_line_count(bufnr)
  end
  return 0
end

return U