1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
type t
type space = L2 | Cosine | IP
type dtype = F32 | F16 | I8
type vector = (float, Bigarray.float32_elt, Bigarray.c_layout) Bigarray.Array1.t
type search_result = {
key : string;
value : string;
distance : float;
metadata : string option;
}
type key_result = {
key : string;
distance : float;
}
type stats = {
dim : int;
space : space;
dtype : dtype;
connectivity : int;
expansion_add : int;
expansion_search : int;
count : int64;
capacity : int64;
fill_ratio : float;
vec_cf_count : int64;
has_metadata : bool;
sidecar_enabled : bool;
}
type item = {
key : string;
value : string;
vector : vector;
metadata : string option;
}
exception Vec_error of int * string
let () = Callback.register_exception "snkv_vec_error" (Vec_error (0, ""))
external open_raw : string option -> int -> int -> int -> int -> int -> int -> string option -> t = "ocaml_kvstore_vec_open_bytecode" "ocaml_kvstore_vec_open"
external close_raw : t -> unit = "ocaml_kvstore_vec_close"
external put_raw : t -> string -> string -> vector -> int64 -> string option -> unit = "ocaml_kvstore_vec_put_bytecode" "ocaml_kvstore_vec_put"
external put_batch_raw : t -> item array -> int64 -> unit = "ocaml_kvstore_vec_put_batch"
external kv_put_raw : t -> string -> string -> unit = "ocaml_kvstore_vec_kv_put"
external get_raw : t -> string -> string = "ocaml_kvstore_vec_get"
external get_vector_raw : t -> string -> vector = "ocaml_kvstore_vec_get_vector"
external get_metadata_raw : t -> string -> string option = "ocaml_kvstore_vec_get_metadata"
external contains_raw : t -> string -> bool = "ocaml_kvstore_vec_contains"
external count_raw : t -> int64 = "ocaml_kvstore_vec_count"
external delete_raw : t -> string -> unit = "ocaml_kvstore_vec_delete"
external search_raw : t -> vector -> int -> bool -> int -> float -> search_result array = "ocaml_kvstore_vec_search_bytecode" "ocaml_kvstore_vec_search"
external search_keys_raw : t -> vector -> int -> key_result array = "ocaml_kvstore_vec_search_keys"
external stats_raw : t -> int * int * int * int * int * int * int64 * int64 * float * int64 * bool * bool = "ocaml_kvstore_vec_stats"
external purge_expired_raw : t -> int = "ocaml_kvstore_vec_purge_expired"
external drop_index_raw : t -> unit = "ocaml_kvstore_vec_drop_index"
let protect f =
match f () with
| v -> Ok v
| exception Vec_error (rc, _msg) -> Error (Kv.error_code_of_int rc)
let space_to_int = function L2 -> 0 | Cosine -> 1 | IP -> 2
let dtype_to_int = function F32 -> 0 | F16 -> 1 | I8 -> 2
let space_of_int = function 1 -> Cosine | 2 -> IP | _ -> L2
let dtype_of_int = function 1 -> F16 | 2 -> I8 | _ -> F32
let open_ ?filename ?(space = L2) ?(connectivity = 0) ?(expansion_add = 0) ?(expansion_search = 0) ?(dtype = F32) ?password dim =
protect (fun () -> open_raw filename dim (space_to_int space) connectivity expansion_add expansion_search (dtype_to_int dtype) password)
let open_exn ?filename ?space ?connectivity ?expansion_add ?expansion_search ?dtype ?password dim =
match open_ ?filename ?space ?connectivity ?expansion_add ?expansion_search ?dtype ?password dim with
| Ok t -> t
| Error e -> failwith (Kv.error_to_string e)
let close t = close_raw t
let put ?(expire_ms = 0L) ?metadata t ~key ~value ~vector =
protect (fun () -> put_raw t key value vector expire_ms metadata)
let put_batch t ?(expire_ms = 0L) items =
if Array.length items = 0 then Ok ()
else protect (fun () -> put_batch_raw t items expire_ms)
let kv_put t ~key ~value = protect (fun () -> kv_put_raw t key value)
let get t ~key = protect (fun () -> get_raw t key)
let get_vector t ~key = protect (fun () -> get_vector_raw t key)
let get_metadata t ~key =
match protect (fun () -> get_metadata_raw t key) with
| Error Kv.NotFound -> Ok None
| Ok v -> Ok v
| Error e -> Error e
let contains t ~key = protect (fun () -> contains_raw t key)
let count t = protect (fun () -> count_raw t)
let delete t ~key = protect (fun () -> delete_raw t key)
let search ?(rerank = false) ?(oversample = 0) ?(max_distance = 0.0) t ~query ~top_k =
match protect (fun () -> search_raw t query top_k rerank oversample max_distance) with
| Ok a -> Ok (Array.to_list a)
| Error e -> Error e
let search_keys t ~query ~top_k =
match protect (fun () -> search_keys_raw t query top_k) with
| Ok a -> Ok (Array.to_list a)
| Error e -> Error e
let stats t =
match protect (fun () -> stats_raw t) with
| Error e -> Error e
| Ok (dim, space, dtype, connectivity, expansion_add, expansion_search, count, capacity, fill_ratio, vec_cf_count, has_metadata, sidecar_enabled) ->
Ok { dim; space = space_of_int space; dtype = dtype_of_int dtype; connectivity; expansion_add; expansion_search; count; capacity; fill_ratio; vec_cf_count; has_metadata; sidecar_enabled }
let purge_expired t = protect (fun () -> purge_expired_raw t)
let drop_index t = protect (fun () -> drop_index_raw t)
let of_float_array a = Bigarray.Array1.of_array Bigarray.float32 Bigarray.c_layout a