From 2a6fd554998c64733e2d97aecf653f6e48e0f8b4 Mon Sep 17 00:00:00 2001 From: Mitja Felicijan Date: Fri, 13 Feb 2026 18:07:45 +0100 Subject: Store context documents to Vector Database --- vectordb.c | 118 ++++++++++++++++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 102 insertions(+), 16 deletions(-) (limited to 'vectordb.c') diff --git a/vectordb.c b/vectordb.c index 5e45cc4..b6fae64 100644 --- a/vectordb.c +++ b/vectordb.c @@ -1,33 +1,38 @@ #include #include #include +#include #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; static float cosine_similarity(float *a, float *b, int n) { - float dot = 0, normA = 0, normB = 0; + float dot = 0, norm_a = 0, norm_b = 0; for (int i = 0; i < n; i++) { dot += a[i] * b[i]; - normA += a[i] * a[i]; - normB += b[i] * b[i]; + norm_a += a[i] * a[i]; + norm_b += b[i] * b[i]; } - return dot / (sqrtf(normA) * sqrtf(normB) + 1e-8f); + return dot / (sqrtf(norm_a) * sqrtf(norm_b) + 1e-8f); } static void embed_text(struct llama_context *ctx, const char *text, float *out) { llama_token tokens[512]; 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, 512, true, true); if (n_tokens < 0) { return; } @@ -46,12 +51,12 @@ void vdb_init(VectorDB *db, struct llama_context *embed_ctx) { } void vdb_free(VectorDB *db) { - (void)db; // nothing yet (future persistence etc.) + (void)db; } void vdb_add_document(VectorDB *db, const char *text) { if (db->count >= VDB_MAX_DOCS) { - printf("VectorDB full!\n"); + log_message(stdout, LOG_INFO, "Vector database full"); return; } @@ -59,7 +64,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; - printf("Embedding doc %d...\n", db->count); + log_message(stdout, LOG_INFO, "Embedding doc %d...", db->count); embed_text(db->embed_ctx, text, doc->embedding); } @@ -90,3 +95,84 @@ void vdb_search(VectorDB *db, float *query, int top_k, int *results) { } } } + +int vdb_save(const VectorDB *db, const char *path) { + FILE *fp = fopen(path, "wb"); + if (!fp) { + return 1; + } + + VdbFileHeader header = { + .magic = VDB_MAGIC, + .version = VDB_VERSION, + .embed_size = VDB_EMBED_SIZE, + .max_text = VDB_MAX_TEXT, + .count = (uint32_t)db->count, + }; + + if (fwrite(&header, sizeof(header), 1, fp) != 1) { + fclose(fp); + return 2; + } + + 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; + } + } + + if (fclose(fp) != 0) { + return 4; + } + + return 0; +} + +int vdb_load(VectorDB *db, const char *path) { + struct llama_context *ctx = db->embed_ctx; + FILE *fp = fopen(path, "rb"); + if (!fp) { + return -1; + } + + VdbFileHeader header = {0}; + if (fread(&header, sizeof(header), 1, fp) != 1) { + fclose(fp); + return -2; + } + + if (header.magic != VDB_MAGIC || header.version != VDB_VERSION) { + fclose(fp); + return -3; + } + + if (header.embed_size != VDB_EMBED_SIZE || header.max_text != VDB_MAX_TEXT) { + fclose(fp); + return -4; + } + + if (header.count > VDB_MAX_DOCS) { + fclose(fp); + return -5; + } + + memset(db, 0, sizeof(VectorDB)); + db->embed_ctx = ctx; + db->count = (int)header.count; + + if (db->count > 0) { + size_t read = fread(db->docs, sizeof(VectorDoc), (size_t)db->count, fp); + if (read != (size_t)db->count) { + fclose(fp); + return -6; + } + } + + if (fclose(fp) != 0) { + return -7; + } + + return 0; +} -- cgit v1.2.3