shithub: oai

ref: e6f8cc2410632b4c05c9fd438034fafc19b95798
dir: /ocomplete.c/

#include <u.h>
#include <libc.h>
#include <bio.h>
#include <json.h>
#include "oai.h"

static void
usage(void)
{
	fprint(2, "usage: %s [-d] [-k apikey] [-m model] [-u baseurl]\n", argv0);
	exits("usage");
}

char* winid;
char tag[2048];

void
readfilename(void)
{
	int fd;
	char *s;
	
	snprint(tag, sizeof tag, "/mnt/acme/%s/tag", winid);
	fd = open(tag, OREAD);
	tag[0] = 0;
	if (fd < 0) {
		return;
	}
	
	read(fd, tag, sizeof tag);
	close(fd);
	
	s = strchr(tag, ' ');
	if (s)
		*s = 0;
}

char*
charpos(char *body, int pos)
{
	Rune r;
	int n;
	int ispos = 0;
	while (ispos < pos && *body) {
		n = chartorune(&r, body);
		ispos += n;
		body += n;
	}
	return body;
}

void
main(int argc, char **argv)
{
	char buf[128];
	ORequest req;
	OResult res;
	char *key = nil;
	char *body, *selection;
	char *p1body;
	int ctlfd, addrfd;
	int addr, addr2;
	char *p1pos, *p2pos;
	int n;
	char *toks[2];
	Biobuf *bodyin, *selout;
	char *url = nil;
	
	req.prompts = nil;
	req.model = nil;
	
	ARGBEGIN{
	case 'h':
		usage();
	case 'k':
		key = EARGF(usage());
		break;
	case 'm':
		req.model = EARGF(usage());
		break;
	case 'u':
		url = EARGF(usage());
		break;
	case 'd':
		oaidebug++;
		break;
	}ARGEND;
	
	winid = getenv("winid");
	if (!winid || winid[0] == 0)
		sysfatal("not in acme");
	
	readfilename();
	
	if (!initoai(url, key))
		sysfatal("initoai: %r");
	
	/* get addr */
	snprint(buf, sizeof buf, "/mnt/acme/%s/ctl", winid);
	ctlfd = open(buf, ORDWR);
	snprint(buf, sizeof buf, "/mnt/acme/%s/addr", winid);
	addrfd = open(buf, OREAD);
	
	if (ctlfd < 0)
		sysfatal("ctl: %r");
	if (addrfd < 0)
		sysfatal("addr: %r");
	
	fprint(ctlfd, "addr=dot\n");
	if (read(addrfd, buf, sizeof buf) <= 0)
		sysfatal("read: %r");
	close(ctlfd);
	close(addrfd);
	
	n = getfields(buf, toks, 2, 1, " ");
	if (!n)
		sysfatal("invalid addr");
	addr2 = -1;
	addr = atoi(toks[0]);
	if (n >= 2)
		addr2 = atoi(toks[1]);
	
	/* make sure addr < addr2 */
	if (addr > addr2) {
		n = addr;
		addr = addr2;
		addr2 = n;
	}
	
	/* buffer body */
	snprint(buf, sizeof buf, "/mnt/acme/%s/body", winid);
	bodyin = Bopen(buf, OREAD);
	if (!bodyin)
		sysfatal("Bopen body: %r");
	body = Brdstr(bodyin, 0, 1);
	Bterm(bodyin);
	
	/* split body in p1body, selection (if applicable) and p1pos */
	p1pos = charpos(body, addr);
	p2pos = nil;
	if (addr != addr2)
		p2pos = charpos(body, addr2);
	
	p1body = mallocz(p1pos-body + 1, 1);
	memcpy(p1body, body, p1pos-body);
	
	selection = nil;
	if (p2pos) {
		selection = mallocz(p2pos-p1pos + 1, 1);
		memcpy(selection, p1pos, p2pos-p1pos);
		p1pos = p2pos;
	}
	
	/* build prompts */
	addstrprompt(&req, "system", "You are the Acme Assistant. Your task is to fill in a gap in a file or replace the selected part. User provides you with the relevant parts of the file. Do not directly respond to the user, just give the text. Do not add markdown formatting, only do plain text or the format of the file. Follow a potential user request if it's in the selection.");
	
	addstrprompt(&req, "user", "The name of the file is: %s", tag);
	addstrprompt(&req, "assistant", "Give me the first part of the file");
	addstrprompt(&req, "user", "%s", p1body);
	if (selection) {
		addstrprompt(&req, "assistant", "Give me part that's selected");
		addstrprompt(&req, "user", "%s", selection);
	}
	addstrprompt(&req, "assistant", "Give me the remaining part of the file");
	addstrprompt(&req, "user", "%s", p1pos);
	
	free(body);
	free(p1body);
	free(selection);
	
	res = makerequest(req);
	if (!res.message)
		sysfatal("invalid result");
	
	snprint(buf, sizeof buf, "/mnt/acme/%s/wrsel", winid);
	selout = Bopen(buf, OWRITE);
	Bprint(selout, "%s", res.message);
	Bterm(selout);
	
	exits(nil);
}