Skip to content
Enno edited this page Mar 11, 2024 · 1 revision

Here's an example for a custom source to read a YAML role file ~/.config/roles.yaml:

let g:vim_ai_roles_config_function = 'VimAIGetRoles'
let g:vim_ai_roles_yaml_file = $XDG_CONFIG_HOME . '/aichat/roles.yaml'

function! VimAIGetRoles() abort
  let file = g:vim_ai_roles_yaml_file
  if !filereadable(file)
    throw 'cannot read file ' . file
  endif

  let lines = readfile(file)
  let roles = {}
  let current_role = ''
  let current_attribute = ''

  for line in lines
    if line =~ '^-\s*name:\s*'
      let current_role = matchstr(line, '^-\s*name:\s*\zs.*')->trim()
      let roles[current_role] = {}
    elseif line =~ '^\s\+\S\+:\s*[>|]'
      let current_attribute = matchstr(line, '^\s\+\zs\S\+\ze:\s*[>|]')->trim()
      let roles[current_role][current_attribute] = ''
    elseif line =~ '^\s\+\S\+:\s*[^>|]'
      let current_attribute = matchstr(line, '^\s\+\zs[^:[:space:]]\+\ze:\s*[^>|]')->trim()
      let roles[current_role][current_attribute] = matchstr(line, '^\s*' . current_attribute . ':\zs.*')->trim()
    else
      let roles[current_role][current_attribute] .= line->trim() . "\n"
    endif
  endfor

  " convert to vim-ai config dictionary form
  for role in keys(roles)
    let roles[role].options = {}
    if has_key(roles[role], 'prompt')
      let roles[role].options.initial_prompt = ">>> system\n" . roles[role].prompt
    endif
    if has_key(roles[role], 'temperature')
      let roles[role].options.temperature = roles[role].temperature
      unlet roles[role].temperature
    endif
  endfor
  return roles
endfunction
Clone this wiki locally