Add daemon mode.

1 files changed, 111 insertions, 183 deletions

diff --git a/main.c b/main.c
index 6a2eb36..e89ed7a 100644
--- a/main.c
+++ b/main.c
@@ -2,229 +2,157 @@
 #include <stdbool.h>
 #include <ctype.h>
 #include <err.h>
+#include <errno.h>
+#include <signal.h>
+#include <string.h>
 #include <unistd.h>
+#include <sys/un.h>
+#include <sys/stat.h>
+#include <sys/socket.h>
 
-#include "mem.h"
-#include "llama.h"
+#include "llm.h"
 
+#define SOCK_PATH  "/var/run/lex.sock"
 #define MODEL_PATH "qwen2.5-3b-instruct-q4_k_m.gguf"
 
-#define SYSTEM_PROMPT \
-    "You are a dictionary assistant. Define the target word in the context provided.\n" \
-    "DO NOT repeat target word in definition. Choose one or two of Formality levels.\n\n" \
-    "FOLLOW the output format:\n" \
-    "<WORD> (<part of speech>) — <Formality: Conversational|Formal|Literary|Figurative|Archaic>\n\n" \
-    "DEFINITION:\n" \
-    "<terse, highly accurate definition in context>\n\n" \
-    "GENERAL TONE:\n" \
-    "<One terse sentence of describing general tone. One terse sentence of standard conversational synonyms.>\n\n" \
-    "SYNONYMS:\n" \
-    "<Three to four standard conversational synonyms.>\n\n" \
-    "EXAMPLES:\n" \
-    "1. <example 1>\n" \
-    "2. <example 2>\n\n"
-
-#define MAX_TOKENS  300
-
-static void process_request(struct llama_model *model, const char *user_prompt)
+static void usage(const char *bin)
 {
-	int n_prompt_tokens, i;
-	char *prompt;
-	int prompt_len;
-
-	llama_token *prompt_tokens;
-	llama_token new_token_id;
-
-	struct llama_batch batch;
-	struct llama_context *ctx;
-	struct llama_context_params cparams;
-	struct llama_sampler *smpl;
-	struct llama_sampler_chain_params sparams;
-
-	struct llama_chat_message messages[] = {
-		{ "system", SYSTEM_PROMPT },
-		{ "user",   user_prompt }
-	};
-
-	const char *tmpl = llama_model_chat_template(model, NULL);
-	if (!tmpl)
-		fprintf(stderr, "Warning: model has no embedded chat template\n");
-
-	prompt_len = llama_chat_apply_template(NULL, messages, 2, true, NULL, 0);
-	if (prompt_len <= 0) {
-		fprintf(stderr, "Error: failed to calculate chat template size\n");
-		return;
-	}
-
-	prompt = MALLOC((size_t)prompt_len + 1);
-	if (llama_chat_apply_template(NULL, messages, 2, true, prompt, prompt_len + 1) < 0) {
-		fprintf(stderr, "Error: failed to apply chat template\n");
-		free(prompt);
-		return;
-	}
-
-	cparams = llama_context_default_params();
-	cparams.n_ctx = 1024;    /* context size in tokens */
-	cparams.n_threads = 4;
-	cparams.n_threads_batch = 4;
-
-	ctx = llama_init_from_model(model, cparams);
-	if (!ctx) {
-		fprintf(stderr, "Error: failed to create context\n");
-		free(prompt);
-		return;
-	}
-
-	const struct llama_vocab *vocab = llama_model_get_vocab(model);
-	if (!vocab) {
-		fprintf(stderr, "Error: failed to obtain model vocabulary\n");
-		free(prompt);
-		llama_free(ctx);
-		return;
-	}
-
-	n_prompt_tokens = -llama_tokenize(vocab, 
-		prompt, prompt_len, NULL, 0, false, true);
+	errx(1, "usage: %s [-m model_path] [-d] [prompt]", bin);
+}
 
-	if (n_prompt_tokens <= 0) {
-		fprintf(stderr, "Error: tokenization sizing failed\n");
-		free(prompt);
-		llama_free(ctx);
-		return;
-	}
+static void loop(struct llama_model *model, int lsock)
+{
+	int csock, c;
+	FILE *cfp;
+	size_t len;
+	char line[4096];
+
+	for (;;) {
+		csock = accept(lsock, NULL, NULL);
+		if (csock == -1) {
+			if (errno != EINTR)
+				warn("accept()");
+			continue;
+		}
 
-	if (n_prompt_tokens + MAX_TOKENS > (int)cparams.n_ctx) {
-		fprintf(stderr, "Error: token count exceeds context size\n");
-		free(prompt);
-		llama_free(ctx);
-		return;
-	}
+		if (!(cfp = fdopen(csock, "r+"))) {
+			warn("cfp fdopen()");
+			close(csock);
+			continue;
+		}
 
-	prompt_tokens = MALLOC((size_t)n_prompt_tokens * sizeof(llama_token));
+		if (fgets(line, sizeof(line), cfp) != NULL) {
+			len = strcspn(line, "\n");
+			if (line[len] == '\0' && len == sizeof(line) - 1) {
+				warnx("request too large");
+				// drain and skip
+				while ((c = fgetc(cfp)) != '\n' && c != EOF)
+					;
+				fclose(cfp);
+				continue;
+			}
+
+			line[len] = '\0';
+			if (line[0] != '\0')
+				llm_process_request(model, line, cfp);
+		}
 
-	if (llama_tokenize(vocab, prompt, prompt_len, prompt_tokens, 
-		n_prompt_tokens, false, true) < 0) {
-		fprintf(stderr, "Error: Tokenization failed\n");
-		free(prompt);
-		free(prompt_tokens);
-		llama_free(ctx);
-		return;
+		fclose(cfp);
 	}
+}
 
-	free(prompt); /* Prompt buffer is fully tokenized and no longer needed */
+int sock_init(void)
+{
+	int sock;
+	struct sockaddr_un addr;
 
-	/* Ingest prompt tokens in one batch (parallelizes matrix 
-	 * multiplications across tokens in the batch) */
-	batch = llama_batch_get_one(prompt_tokens, n_prompt_tokens);
+	sock = socket(AF_UNIX, SOCK_STREAM, 0);
+	if (sock == -1)
+		err(1, "socket()");
 
-	if (llama_decode(ctx, batch) != 0) {
-		fprintf(stderr, "Error: Prompt evaluation failed\n");
-		free(prompt_tokens);
-		llama_free(ctx);
-		return;
-	}
+	memset(&addr, 0, sizeof(addr));
+	addr.sun_family = AF_UNIX;
+	strlcpy(addr.sun_path, SOCK_PATH, sizeof(addr.sun_path));
 
-	free(prompt_tokens);
+	unlink(SOCK_PATH); /* remove stale socket from a previous run */
+	
+	if (bind(sock, (struct sockaddr *)&addr, sizeof(addr)) == -1)
+		err(1, "bind()");
 
-	/* Generation loop */
-	sparams = llama_sampler_chain_default_params();
-	smpl = llama_sampler_chain_init(sparams);
+	if (listen(sock, 1) == -1) /* backlog = 1 */
+		err(1, "listen()");
 
-	llama_sampler_chain_add(smpl, llama_sampler_init_penalties(
-		64,      /* last_n: lookback window (64 is standard) */
-		1.1f,    /* repeat_penalty */
-		0.0f,    /* frequency_penalty */
-		0.0f     /* presence_penalty */
-	));
+	return sock;
+}
 
-	/* Pick the top token */
-	llama_sampler_chain_add(smpl, llama_sampler_init_greedy());
+int main(int argc , char *argv[])
+{
+	int daemon_mode, sock, opt;
+	const char *prog;
+	const char *prompt;
+	const char *model_path;
+	struct llama_model *model;
 
-	for (i = 0; i < MAX_TOKENS; i++) {	
-		/* Model outputs next tokens for every token in the prompt. 
-		 * We need the one after the last token in the prompt */
-		new_token_id = llama_sampler_sample(smpl, ctx, -1);
+	if (pledge("stdio rpath wpath cpath unix", NULL) == -1)
+		err(1, "pledge failed");
 
-		/* Tell the sampler chain which token was chosen */
-		llama_sampler_accept(smpl, new_token_id);
+	sock = -1;
+	daemon_mode = 0;
+	prompt = NULL;
+	prog = argv[0];
+	model_path = MODEL_PATH;
 
-		/* Check for end-of-generation (EOG) tokens:
-		 * EOS: end-of-sequence
-		 * EOT: end-of-turn */
-		if (llama_vocab_is_eog(vocab, new_token_id))
+	while ((opt = getopt(argc, argv, "dm:")) != -1) {
+		switch (opt) {
+		case 'd':
+			daemon_mode = 1;
 			break;
-
-		/* Convert numeric token id to printable text */
-		char buf[128];
-		int n = llama_token_to_piece(vocab, new_token_id, buf, 
-			sizeof(buf), 0, false);
-
-		if (n > 0) {
-			fwrite(buf, 1, (size_t)n, stdout);
-			fflush(stdout);
-		}
-
-		/* Create batch with 1 token for the next forward pass */
-		batch = llama_batch_get_one(&new_token_id, 1);
-
-		if (llama_decode(ctx, batch) != 0) {
-			fprintf(stderr, "llama_decode failed!\n");
+		case 'm':
+			model_path = optarg;
 			break;
+		default:
+			usage(prog);
 		}
 	}
 
-	printf("\n\n");
-	fflush(stdout);
-
-	//llama_perf_context_print(ctx);
+	argc -= optind;
+	argv += optind;
 
-	llama_sampler_free(smpl);
-	llama_free(ctx);
-}
-
-int main(int argc , char *argv[])
-{
-	const char *prompt;
-	const char *model_path;
-	struct llama_model *model;
-	struct llama_model_params mparams;
-
-	if (argc == 2) {
-		model_path = MODEL_PATH;
-		prompt = argv[1];
-	} else if (argc == 3) {
-		model_path = argv[1];
-		prompt = argv[2];
+	if (daemon_mode) {
+		if (argc != 0)  // -d takes no positional args
+			usage(prog);
 	} else {
-		errx(1, "usage: %s [model_path] <prompt>", argv[0]);
+		if (argc != 1) // normal mode requires one prompt
+			usage(prog);
+
+		prompt = argv[0];
 	}
 
 	if (unveil(model_path, "r") == -1)
 		err(1, "unveil %s failed", model_path);
-
+	if (daemon_mode && unveil(SOCK_PATH, "rwc") == -1)
+		err(1, "unveil %s failed", SOCK_PATH);
 	if (unveil(NULL, NULL) == -1)
 		err(1, "unveil lock failed");
 
-	if (pledge("stdio rpath", NULL) == -1)
-		err(1, "initial pledge failed");
-
-	llama_backend_init();
-
-	mparams = llama_model_default_params();
-	mparams.n_gpu_layers = 0; /* force all layers onto CPU */
-	mparams.load_mode = LLAMA_LOAD_MODE_MMAP;
+	if (daemon_mode) {
+		sock = sock_init();
+		if (pledge("stdio rpath unix", NULL) == -1)
+			err(1, "Failed to narrow pledge after sock_init()");
+	}
 
-	model = llama_model_load_from_file(model_path, mparams);
+	model = llm_init(model_path);
 	if (!model)
-		errx(1, "failed to load model from file %s", model_path);
+		errx(1, "Failed to load model");
 
-	if (pledge("stdio", NULL) == -1)
-		err(1, "secondary pledge failed");
+	if (daemon_mode)
+		loop(model, sock); /* doesn't return */
 
-	process_request(model, prompt);
+	if (pledge("stdio", NULL) == -1)
+		err(1, "Failed to narrow pledge for normal mode");
 
-	llama_model_free(model);
-	llama_backend_free();
+	llm_process_request(model, prompt, stdout);
+	llm_free(model);
 
 	return 0;
 }