Add daemon mode.

4 files changed, 330 insertions, 186 deletions

diff --git a/Makefile b/Makefile
index 38babd7..93ac6f7 100644
--- a/Makefile
+++ b/Makefile
@@ -2,7 +2,7 @@ CC            = cc
 CXX           = c++
 CFLAGS        = -march=native -O2 -std=c11 -Wall -Wextra
 CXXFLAGS      = -march=native -O2 -O2 -Wall -Wextra
-INCLUDES      = -Ideps/include
+INCLUDES      = -I. -Ideps/include
 
 LLAMA_TAG     = b10107
 LLAMA_URL     = https://github.com/ggml-org/llama.cpp/archive/refs/tags/$(LLAMA_TAG).tar.gz
@@ -18,8 +18,8 @@ STATIC_LIBS   = deps/lib/libllama.a \
 SYS_LIBS      = -static -pthread -lm -lc++ -lc++abi
 
 TARGET        = lex
-SRC           = main.c
-OBJS          = main.o
+SRC           = llm.c main.c
+OBJS          = $(SRC:.c=.o)
 
 .PHONY: all clean distclean
 
@@ -40,6 +40,9 @@ $(TARGET): deps/lib/libllama.a $(OBJS)
 .c.o:
 	$(CC) $(CFLAGS) $(INCLUDES) -c $< -o $@
 
+# Rebuild objects if the shared header changes
+$(OBJS): llm.h
+
 # Download llama source and build statically
 deps/lib/libllama.a deps/lib/libggml.a deps/lib/libggml-cpu.a:
 	@mkdir -p deps/lib deps/include
diff --git a/llm.c b/llm.c
new file mode 100644
index 0000000..dc3608f
--- /dev/null
+++ b/llm.c
@@ -0,0 +1,200 @@
+#include <err.h>
+
+#include "llm.h"
+#include "mem.h"
+
+#define MAX_TOKENS  300
+
+#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"
+
+struct llama_model *llm_init(const char *model_path)
+{
+	struct llama_model *model;
+	struct llama_model_params mparams;
+
+	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;
+
+	model = llama_model_load_from_file(model_path, mparams);
+	return model;
+}
+
+void llm_free(struct llama_model *model)
+{
+	llama_model_free(model);
+	llama_backend_free();
+}
+
+void llm_process_request(struct llama_model *model, 
+	const char *user_prompt, FILE *out)
+{
+	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);
+
+	if (n_prompt_tokens <= 0) {
+		fprintf(stderr, "Error: tokenization sizing failed\n");
+		free(prompt);
+		llama_free(ctx);
+		return;
+	}
+
+	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;
+	}
+
+	prompt_tokens = MALLOC((size_t)n_prompt_tokens * sizeof(llama_token));
+
+	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;
+	}
+
+	free(prompt); /* Prompt buffer is fully tokenized and no longer needed */
+
+	/* Ingest prompt tokens in one batch (parallelizes matrix 
+	 * multiplications across tokens in the batch) */
+	batch = llama_batch_get_one(prompt_tokens, n_prompt_tokens);
+
+	if (llama_decode(ctx, batch) != 0) {
+		fprintf(stderr, "Error: Prompt evaluation failed\n");
+		free(prompt_tokens);
+		llama_free(ctx);
+		return;
+	}
+
+	free(prompt_tokens);
+
+	/* Generation loop */
+	sparams = llama_sampler_chain_default_params();
+	smpl = llama_sampler_chain_init(sparams);
+
+	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 */
+	));
+
+	/* Pick the top token */
+	llama_sampler_chain_add(smpl, llama_sampler_init_greedy());
+
+	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);
+
+		/* Tell the sampler chain which token was chosen */
+		llama_sampler_accept(smpl, new_token_id);
+
+		/* Check for end-of-generation (EOG) tokens:
+		 * EOS: end-of-sequence
+		 * EOT: end-of-turn */
+		if (llama_vocab_is_eog(vocab, new_token_id))
+			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, out);
+			fflush(out);
+		}
+
+		/* 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");
+			break;
+		}
+	}
+
+	fprintf(out, "\n\n");
+	fflush(out);
+
+	//llama_perf_context_print(ctx);
+
+	llama_sampler_free(smpl);
+	llama_free(ctx);
+}
+
diff --git a/llm.h b/llm.h
new file mode 100644
index 0000000..5bec282
--- /dev/null
+++ b/llm.h
@@ -0,0 +1,13 @@
+#ifndef LLM_H
+#define LLM_H
+
+#include "llama.h"
+
+struct llama_model *llm_init(const char *model_path);
+
+void llm_process_request(struct llama_model *model, 
+	const char *user_prompt, FILE *out);
+
+void llm_free(struct llama_model *model);
+
+#endif /* LLM_H */
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;
 }