1#ifndef VECTORDB_H
2#define VECTORDB_H
3
4#include "llama.h"
5
6#define VDB_MAX_DOCS 1000
7#define VDB_EMBED_SIZE 768
8#define VDB_MAX_TEXT 1024
9
10#define VDB_MAGIC 0x31424456u /* "VDB1" */
11#define VDB_VERSION 1u
12#define VDB_TOKENS 512
13
14typedef struct {
15 float embedding[VDB_EMBED_SIZE];
16 char text[VDB_MAX_TEXT];
17} VectorDoc;
18
19typedef struct {
20 VectorDoc docs[VDB_MAX_DOCS];
21 int count;
22 struct llama_context *embed_ctx;
23} VectorDB;
24
25typedef struct {
26 uint32_t magic;
27 uint32_t version;
28 uint32_t embed_size;
29 uint32_t max_text;
30 uint32_t count;
31} VdbFileHeader;
32
33typedef enum {
34 VDB_SUCCESS = 0,
35 VDB_OPEN_ERR = 9001,
36 VDB_CLOSE_ERR = 9002,
37 VDB_HEADER_WRITE_ERR = 9003,
38 VDB_HEADER_READ_ERR = 9004,
39 VDB_MAGIC_MISMATCH_ERR = 9005,
40 VDB_EMBED_MISMATCH_ERR = 9006,
41 VDB_COUNT_TOO_LARGE_ERR = 9007,
42 VDB_DOC_WRITE_ERR = 9008,
43 VDB_DOC_READ_ERR = 9009,
44} VectorDBErrorCode;
45
46void vdb_init(VectorDB *db, struct llama_context *embed_ctx);
47void vdb_free(VectorDB *db);
48
49void vdb_add_document(VectorDB *db, const char *text);
50
51void vdb_embed_query(VectorDB *db, const char *text, float *out_embedding);
52void vdb_search(VectorDB *db, float *query_embedding, int top_k, int *results);
53
54VectorDBErrorCode vdb_save(const VectorDB *db, const char *path);
55VectorDBErrorCode vdb_load(VectorDB *db, const char *path);
56
57const char* vdb_error(VectorDBErrorCode err);
58
59#endif