local edit = require 'bugabinga.pivi.edit' local IMAGES = { png = 'image/png', jpg = 'image/jpeg', jpeg = 'image/jpeg', gif = 'image/gif', webp = 'image/webp', } --- What a written prompt turns out to mean. --- Pi reads a leading bang as a shell command rather than something said, and --- doubles it to keep the output out of the conversation. --- @param text string --- @return table action kind, and what that kind needs local parse = function( text ) vim.validate( 'text', text, 'string' ) local written = vim.trim( text ) local quiet = written:match '^!!%s*(.*)$' if quiet then return { kind = 'bash', command = vim.trim( quiet ), exclude = true, } end local loud = written:match '^!%s*(.*)$' if loud then return { kind = 'bash', command = vim.trim( loud ), exclude = false, } end return { kind = 'prompt', message = written, } end local image_type = function( path ) local suffix = path:match '%.([%a%d]+)$' if not suffix then return nil end return IMAGES[suffix:lower()] end local absolute = function( path, cwd ) local expanded = vim.fn.expand( path ) if expanded:match '^/' or expanded:match '^%a:' then return vim.fs.normalize( expanded ) end return vim.fs.normalize( ( cwd or vim.uv.cwd() ) .. '/' .. expanded ) end --- Reads a file for the model. --- A loaded buffer wins over what is on disk, because that is what the user is --- looking at and what every other pivi read already answers with. local contents = function( path ) local text = edit.read( path ) return text end local encode_image = function( path ) local read = io.open( path, 'rb' ) if not read then return nil end local bytes = read:read '*a' read:close() if not bytes or bytes == '' then return nil end local ok, encoded = pcall( vim.base64.encode, bytes ) if not ok then return nil end return encoded end --- Replaces every file reference with the file itself. --- The shape matches what Pi builds for the same reference elsewhere, so a --- model sees one convention rather than two. --- @param text string --- @param options table? cwd --- @return string message --- @return table images --- @return string[] missing references that could not be read local expand = function( text, options ) vim.validate( 'text', text, 'string' ) options = options or {} local images = {} local missing = {} local attached = {} local message = text:gsub( '@([%w%._%-~/\\:]+)', function( reference ) local path = absolute( reference, options.cwd ) local mime = image_type( path ) if mime then local data = encode_image( path ) if not data then table.insert( missing, reference ) return '@' .. reference end table.insert( images, { type = 'image', data = data, mimeType = mime, } ) return ( '' ):format( path ) end local body = contents( path ) if not body then table.insert( missing, reference ) return '@' .. reference end if attached[path] then return ( '' ):format( path ) end attached[path] = true return ( '\n%s\n' ):format( path, body ) end ) return message, images, missing end --- Every reference a written prompt names. local references = function( text ) local found = {} for reference in ( text or '' ):gmatch '@([%w%._%-~/\\:]+)' do table.insert( found, reference ) end return found end return { image_types = IMAGES, parse = parse, expand = expand, references = references, absolute = absolute, }