Update to VectorDB

6 files changed, 106 insertions, 47 deletions

diff --git a/.gitignore b/.gitignore
index faeeddf..1194e85 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,6 +3,7 @@ models/
 prompt
 context
 npc
+game
 
 # Other files
 .DS_Store
diff --git a/Makefile b/Makefile
index 0d27bd6..807a256 100644
--- a/Makefile
+++ b/Makefile
@@ -15,18 +15,21 @@ PROMPT_HEADERS := $(PROMPT_TXT:.txt=.h)
 
 help: .help
 
-build/npc: build/prompts npc.c vectordb.c models.h # Build npc binary for testing
-	$(CC) $(CFLAGS) npc.c vectordb.c -o npc $(LDFLAGS)
-
-build/context: context.c vectordb.c models.h # Build context binary for testing
-	$(CC) $(CFLAGS) context.c vectordb.c -o context $(LDFLAGS)
-
 build/llama.cpp: .assure # Build llama.cpp libraries
 	mkdir $(LLAMA_DIR)/build && \
 		cd $(LLAMA_DIR)/build && \
 		cmake ../ -DBUILD_SHARED_LIBS=OFF && \
 		make -j8
 
+build/context: context.c vectordb.c models.h # Build context binary for testing
+	$(CC) $(CFLAGS) context.c vectordb.c -o context $(LDFLAGS)
+
+build/npc: build/prompts npc.c vectordb.c models.h # Build npc binary for testing
+	$(CC) $(CFLAGS) npc.c vectordb.c -o npc $(LDFLAGS)
+
+build/game: build/prompts game.c vectordb.c models.h # Build npc binary for testing
+	$(CC) $(CFLAGS) game.c vectordb.c -o game $(LDFLAGS)
+
 build/prompts: $(PROMPT_HEADERS) # Generate C style header
 
 run/fetch-models: .assure # Fetch GGUF models
diff --git a/context.c b/context.c
index 633b4dc..e7aa0cf 100644
--- a/context.c
+++ b/context.c
@@ -155,8 +155,9 @@ int main(int argc, char **argv) {
 		vdb_add_document(&db, line);
 	}
 
-	if (vdb_save(&db, out_file) > 0) {
-		log_message(stderr, LOG_ERROR, "Something went wrong saving file %s", out_file);
+	VectorDBErrorCode vdb_rc = vdb_save(&db, out_file);
+	if (vdb_rc != VDB_SUCCESS) {
+		log_message(stderr, LOG_ERROR, "Something went wrong saving file %s: %s", out_file, vdb_error(vdb_rc));
 		fclose(context_fp);
 		return 1;
 	}
diff --git a/npc.c b/npc.c
index 9129864..335eee7 100644
--- a/npc.c
+++ b/npc.c
@@ -340,9 +340,9 @@ int main(int argc, char **argv) {
 
 	VectorDB db = {};
 	vdb_init(&db, embed_ctx);
-	int vdb_rc = vdb_load(&db, context_file);
-	if (vdb_rc != 0) {
-		log_message(stderr, LOG_ERROR, "Failed to load vector database %s (err %d)", context_file, vdb_rc);
+	VectorDBErrorCode vdb_rc = vdb_load(&db, context_file);
+	if (vdb_rc != VDB_SUCCESS) {
+		log_message(stderr, LOG_ERROR, "Failed to load vector database %s: %s", context_file, vdb_error(vdb_rc));
 		llama_free(embed_ctx);
 		llama_model_free(model);
 		llama_backend_free();
diff --git a/vectordb.c b/vectordb.c
index b6fae64..3812ecb 100644
--- a/vectordb.c
+++ b/vectordb.c
@@ -5,19 +5,9 @@
 
 #include "llama.h"
 #include "vectordb.h"
-#include "nonstd.h"
-
-#define VDB_MAGIC 0x31424456u /* "VDB1" */
-#define VDB_VERSION 1u
-
-typedef struct {
-	uint32_t magic;
-	uint32_t version;
-	uint32_t embed_size;
-	uint32_t max_text;
-	uint32_t count;
-} VdbFileHeader;
 
+// Returns cosine similarity in range [-1, 1] (approx).
+// https://en.wikipedia.org/wiki/Cosine_similarity
 static float cosine_similarity(float *a, float *b, int n) {
 	float dot = 0, norm_a = 0, norm_b = 0;
 	for (int i = 0; i < n; i++) {
@@ -29,10 +19,10 @@ static float cosine_similarity(float *a, float *b, int n) {
 }
 
 static void embed_text(struct llama_context *ctx, const char *text, float *out) {
-	llama_token tokens[512];
+	llama_token tokens[VDB_TOKENS];
 	const struct llama_model *model = llama_get_model(ctx);
 	const struct llama_vocab *vocab = llama_model_get_vocab(model);
-	int n_tokens = llama_tokenize(vocab, text, strlen(text), tokens, 512, true, true);
+	int n_tokens = llama_tokenize(vocab, text, strlen(text), tokens, VDB_TOKENS, true, true);
 	if (n_tokens < 0) {
 		return;
 	}
@@ -56,7 +46,7 @@ void vdb_free(VectorDB *db) {
 
 void vdb_add_document(VectorDB *db, const char *text) {
 	if (db->count >= VDB_MAX_DOCS) {
-		log_message(stdout, LOG_INFO, "Vector database full");
+		printf("Vector database full\n");
 		return;
 	}
 
@@ -64,7 +54,7 @@ void vdb_add_document(VectorDB *db, const char *text) {
 	strncpy(doc->text, text, VDB_MAX_TEXT - 1);
 	doc->text[VDB_MAX_TEXT - 1] = 0;
 
-	log_message(stdout, LOG_INFO, "Embedding doc %d...", db->count);
+	printf("Embedding doc %d...\n", db->count);
 	embed_text(db->embed_ctx, text, doc->embedding);
 }
 
@@ -96,10 +86,10 @@ void vdb_search(VectorDB *db, float *query, int top_k, int *results) {
 	}
 }
 
-int vdb_save(const VectorDB *db, const char *path) {
+VectorDBErrorCode vdb_save(const VectorDB *db, const char *path) {
 	FILE *fp = fopen(path, "wb");
 	if (!fp) {
-		return 1;
+		return VDB_OPEN_ERR;
 	}
 
 	VdbFileHeader header = {
@@ -112,50 +102,54 @@ int vdb_save(const VectorDB *db, const char *path) {
 
 	if (fwrite(&header, sizeof(header), 1, fp) != 1) {
 		fclose(fp);
-		return 2;
+		return VDB_HEADER_WRITE_ERR;
 	}
 
 	if (db->count > 0) {
 		size_t wrote = fwrite(db->docs, sizeof(VectorDoc), (size_t)db->count, fp);
 		if (wrote != (size_t)db->count) {
 			fclose(fp);
-			return 3;
+			return VDB_DOC_WRITE_ERR;
 		}
 	}
 
 	if (fclose(fp) != 0) {
-		return 4;
+		return VDB_CLOSE_ERR;
 	}
 
-	return 0;
+	return VDB_SUCCESS;
 }
 
-int vdb_load(VectorDB *db, const char *path) {
+VectorDBErrorCode vdb_load(VectorDB *db, const char *path) {
 	struct llama_context *ctx = db->embed_ctx;
 	FILE *fp = fopen(path, "rb");
 	if (!fp) {
-		return -1;
+		int open_err = errno;
+		fprintf(stderr, "vdb_load: open failed: %s\n", strerror(open_err));
+		return VDB_OPEN_ERR;
 	}
 
 	VdbFileHeader header = {0};
 	if (fread(&header, sizeof(header), 1, fp) != 1) {
+		int read_err = errno;
+		fprintf(stderr, "vdb_load: header read failed: %s\n", strerror(read_err));
 		fclose(fp);
-		return -2;
+		return VDB_HEADER_READ_ERR;
 	}
 
 	if (header.magic != VDB_MAGIC || header.version != VDB_VERSION) {
 		fclose(fp);
-		return -3;
+		return VDB_MAGIC_MISMATCH_ERR;
 	}
 
 	if (header.embed_size != VDB_EMBED_SIZE || header.max_text != VDB_MAX_TEXT) {
 		fclose(fp);
-		return -4;
+		return VDB_EMBED_MISMATCH_ERR;
 	}
 
 	if (header.count > VDB_MAX_DOCS) {
 		fclose(fp);
-		return -5;
+		return VDB_COUNT_TOO_LARGE_ERR;
 	}
 
 	memset(db, 0, sizeof(VectorDB));
@@ -165,14 +159,45 @@ int vdb_load(VectorDB *db, const char *path) {
 	if (db->count > 0) {
 		size_t read = fread(db->docs, sizeof(VectorDoc), (size_t)db->count, fp);
 		if (read != (size_t)db->count) {
+			int read_err = errno;
+			fprintf(stderr, "vdb_load: doc read failed: %s\n", strerror(read_err));
 			fclose(fp);
-			return -6;
+			return VDB_DOC_READ_ERR;
 		}
 	}
 
 	if (fclose(fp) != 0) {
-		return -7;
+		int close_err = errno;
+		fprintf(stderr, "vdb_load: close failed: %s\n", strerror(close_err));
+		return VDB_CLOSE_ERR;
 	}
 
-	return 0;
+	return VDB_SUCCESS;
+}
+
+const char* vdb_error(VectorDBErrorCode err) {
+	switch (err) {
+		case VDB_SUCCESS:
+			return "Success.";
+		case VDB_OPEN_ERR:
+			return "Failed to open file.";
+		case VDB_CLOSE_ERR:
+			return "Failed to close file.";
+		case VDB_HEADER_WRITE_ERR:
+			return "Failed to write header.";
+		case VDB_HEADER_READ_ERR:
+			return "Failed to read header.";
+		case VDB_MAGIC_MISMATCH_ERR:
+			return "Header magic/version mismatch.";
+		case VDB_EMBED_MISMATCH_ERR:
+			return "Header embed/max_text mismatch.";
+		case VDB_COUNT_TOO_LARGE_ERR:
+			return "Header count too large.";
+		case VDB_DOC_WRITE_ERR:
+			return "Failed to write documents.";
+		case VDB_DOC_READ_ERR:
+			return "Failed to read documents.";
+		default:
+			return "Unknown error.";
+	}
 }
diff --git a/vectordb.h b/vectordb.h
index db1d33d..5c41a0a 100644
--- a/vectordb.h
+++ b/vectordb.h
@@ -3,9 +3,15 @@
 
 #include "llama.h"
 
-#define VDB_MAX_DOCS 1000
-#define VDB_EMBED_SIZE 768
-#define VDB_MAX_TEXT 1024
+#include <errno.h>
+
+#define VDB_MAX_DOCS    1000
+#define VDB_EMBED_SIZE  768
+#define VDB_MAX_TEXT    1024
+
+#define VDB_MAGIC       0x31424456u /* "VDB1" */
+#define VDB_VERSION     1u
+#define VDB_TOKENS      512
 
 typedef struct {
 	float embedding[VDB_EMBED_SIZE];
@@ -18,6 +24,27 @@ typedef struct {
 	struct llama_context *embed_ctx;
 } VectorDB;
 
+typedef struct {
+	uint32_t magic;
+	uint32_t version;
+	uint32_t embed_size;
+	uint32_t max_text;
+	uint32_t count;
+} VdbFileHeader;
+
+typedef enum {
+	VDB_SUCCESS                = 0,
+	VDB_OPEN_ERR               = 9001,
+	VDB_CLOSE_ERR              = 9002,
+	VDB_HEADER_WRITE_ERR       = 9003,
+	VDB_HEADER_READ_ERR        = 9004,
+	VDB_MAGIC_MISMATCH_ERR     = 9005,
+	VDB_EMBED_MISMATCH_ERR     = 9006,
+	VDB_COUNT_TOO_LARGE_ERR    = 9007,
+	VDB_DOC_WRITE_ERR          = 9008,
+	VDB_DOC_READ_ERR           = 9009,
+} VectorDBErrorCode;
+
 void vdb_init(VectorDB *db, struct llama_context *embed_ctx);
 void vdb_free(VectorDB *db);
 
@@ -26,7 +53,9 @@ void vdb_add_document(VectorDB *db, const char *text);
 void vdb_embed_query(VectorDB *db, const char *text, float *out_embedding);
 void vdb_search(VectorDB *db, float *query_embedding, int top_k, int *results);
 
-int vdb_save(const VectorDB *db, const char *path);
-int vdb_load(VectorDB *db, const char *path);
+VectorDBErrorCode vdb_save(const VectorDB *db, const char *path);
+VectorDBErrorCode vdb_load(VectorDB *db, const char *path);
+
+const char* vdb_error(VectorDBErrorCode err);
 
 #endif