Skip to content
Enno edited this page Sep 7, 2024 · 2 revisions

Here's an example for a custom source to read a YAML role file ~/.config/roles.yaml used in aichat versions up to 0.20 :

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

In aichat versions newer than 0.20, these roles are split in markdown files, and accordingly the custom role source becomes

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

function! VimAIGetRoles() abort
  let folder = g:vim_ai_roles_md_folder
  if !isdirectory(folder)
    throw 'Cannot read folder ' . folder
  endif

  let roles = {}
  let files = split(globpath(folder, '*.md'), '\n')

  for file in files
    let current_role = fnamemodify(file, ':t:r')
    let roles[current_role] = {}

    let lines = readfile(file)
    let in_attributes = 0
    for line in lines
      if line == '---'
        let in_attributes += 1
        if in_attributes == 2
          let roles[current_role]['prompt'] = []
        endif
      elseif in_attributes == 1
        let parts = split(line, ':', 2)
        if len(parts) == 2
          let roles[current_role][trim(parts[0])] = trim(parts[1])
        endif
      elseif in_attributes == 2
        call add(roles[current_role]['prompt'], line)
      endif
    endfor
    if has_key(roles[current_role], 'prompt')
      let roles[current_role]['prompt'] = join(roles[current_role]['prompt'], "\n").."\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