Skip to content
Enno edited this page Feb 15, 2025 · 4 revisions

Here's an example for a custom source to read a YAML role file ~/.config/roles.yaml used in aichat versions newer than 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 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 role = fnamemodify(file, ':t:r')
    let roles[role] = {}

    " read Aichat roles into Vim dictionary
    let lines = readfile(file)
    let separation_line_count = 0
    for line in lines
      if line == '---'
        let separation_line_count += 1
        if separation_line_count == 2
          let roles[role]['prompt'] = []
        endif
      elseif separation_line_count == 1
        let parts = split(line, '\v^[^:]+\zs:\s+')
        if len(parts) == 2
          let roles[role][trim(parts[0])] = trim(parts[1])
        endif
      elseif separation_line_count == 2
        call add(roles[role]['prompt'], line)
      endif
    endfor
    if has_key(roles[role], 'prompt')
      let roles[role]['prompt'] = join(roles[role]['prompt'], "\n").."\n"
    endif
  endfor

  " convert this Vim dictionary to Vim-ai config dictionary
  for role in keys(roles)
    let roles[role].options = {}
    if has_key(roles[role], 'model')
      " only keep OpenAI models
      if roles[role].model =~# '\v^openai:'
        let roles[role].options.model = roles[role].model[len('openai:'):]
      endif
      unlet roles[role].model
    endif
    if has_key(roles[role], 'temperature')
      let roles[role].options.temperature = roles[role].temperature
      unlet roles[role].temperature
    endif
    if has_key(roles[role], 'prompt')
      let roles[role].options.initial_prompt = ">>> system\n"..roles[role].prompt
      unlet roles[role].prompt
    endif
  endfor

  " special case for roles using the O1 model as it
  " - neither accepts system prompt
  " - nor support streaming
  for role in keys(roles)
    if has_key(roles[role].options, 'model') && roles[role].options.model ==# 'o1'
      let roles[role].options.streaming = '0'
      let roles[role].options.initial_prompt =
            \ substitute(roles[role].options.initial_prompt, '>>> system\n', '>>> user\n', '')
    endif
  endfor
  return roles
endfunction
Clone this wiki locally