Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633//! query request builders and response views.//!//! builders return tuples, which `std.json` writes as arrays, so they compose//! into any request body and nest freely://!//! ```zig//! .filters = filter.@"and"(.{//! filter.eq("kind", "observation"),//! filter.notEq("status", "superseded"),//! }),//! .rank_by = rank.max(.{ rank.bm25("name", q), rank.bm25("caption", q) }),//! ```
const std = @import("std");const json = std.json;const Allocator = std.mem.Allocator;
fn Pair(comptime B: type) type { return struct { []const u8, B };}
fn Triple(comptime C: type) type { return struct { []const u8, []const u8, C };}
pub const rank = struct { /// `["attr", "ANN", [..]]` pub fn ann(attribute: []const u8, vector: []const f32) Triple([]const f32) { return .{ attribute, "ANN", vector }; }
/// exact nearest neighbour; the api requires `filters` alongside it. pub fn knn(attribute: []const u8, vector: []const f32) Triple([]const f32) { return .{ attribute, "kNN", vector }; }
/// `["attr", "BM25", "text"]` pub fn bm25(attribute: []const u8, text: []const u8) Triple([]const u8) { return .{ attribute, "BM25", text }; }
pub fn asc(attribute: []const u8) Pair([]const u8) { return .{ attribute, "asc" }; }
pub fn desc(attribute: []const u8) Pair([]const u8) { return .{ attribute, "desc" }; }
pub fn max(clauses: anytype) Pair(@TypeOf(clauses)) { return .{ "Max", clauses }; }
pub fn sum(clauses: anytype) Pair(@TypeOf(clauses)) { return .{ "Sum", clauses }; }
pub fn product(weight: f64, clause: anytype) struct { []const u8, f64, @TypeOf(clause) } { return .{ "Product", weight, clause }; }};
pub const filter = struct { /// any operator the api accepts: `clause("permissions", "ContainsAny", ids)`. pub fn clause(attribute: []const u8, comptime operator: []const u8, value: anytype) Triple(@TypeOf(value)) { return .{ attribute, operator, value }; }
pub fn eq(attribute: []const u8, value: anytype) Triple(@TypeOf(value)) { return clause(attribute, "Eq", value); } pub fn notEq(attribute: []const u8, value: anytype) Triple(@TypeOf(value)) { return clause(attribute, "NotEq", value); } pub fn in(attribute: []const u8, values: anytype) Triple(@TypeOf(values)) { return clause(attribute, "In", values); } pub fn notIn(attribute: []const u8, values: anytype) Triple(@TypeOf(values)) { return clause(attribute, "NotIn", values); } pub fn lt(attribute: []const u8, value: anytype) Triple(@TypeOf(value)) { return clause(attribute, "Lt", value); } pub fn lte(attribute: []const u8, value: anytype) Triple(@TypeOf(value)) { return clause(attribute, "Lte", value); } pub fn gt(attribute: []const u8, value: anytype) Triple(@TypeOf(value)) { return clause(attribute, "Gt", value); } pub fn gte(attribute: []const u8, value: anytype) Triple(@TypeOf(value)) { return clause(attribute, "Gte", value); } pub fn contains(attribute: []const u8, value: anytype) Triple(@TypeOf(value)) { return clause(attribute, "Contains", value); } pub fn notContains(attribute: []const u8, value: anytype) Triple(@TypeOf(value)) { return clause(attribute, "NotContains", value); } pub fn containsAny(attribute: []const u8, values: anytype) Triple(@TypeOf(values)) { return clause(attribute, "ContainsAny", values); } pub fn notContainsAny(attribute: []const u8, values: anytype) Triple(@TypeOf(values)) { return clause(attribute, "NotContainsAny", values); } pub fn glob(attribute: []const u8, pattern: []const u8) Triple([]const u8) { return clause(attribute, "Glob", pattern); } pub fn notGlob(attribute: []const u8, pattern: []const u8) Triple([]const u8) { return clause(attribute, "NotGlob", pattern); } pub fn regex(attribute: []const u8, pattern: []const u8) Triple([]const u8) { return clause(attribute, "Regex", pattern); } pub fn containsAllTokens(attribute: []const u8, text: []const u8) Triple([]const u8) { return clause(attribute, "ContainsAllTokens", text); } pub fn containsAnyToken(attribute: []const u8, text: []const u8) Triple([]const u8) { return clause(attribute, "ContainsAnyToken", text); }
pub fn @"and"(clauses: anytype) Pair(@TypeOf(clauses)) { return .{ "And", clauses }; } pub fn @"or"(clauses: anytype) Pair(@TypeOf(clauses)) { return .{ "Or", clauses }; } pub fn not(inner: anytype) Pair(@TypeOf(inner)) { return .{ "Not", inner }; }};
/// `include_attributes`: every attribute, or a named subset. `all` is the/// schema-evolution-safe choice — naming an attribute a namespace predates/// is a 400.pub const Attributes = union(enum) { every, some: []const []const u8,
/// `include_attributes: true` pub const all: Attributes = .every;
/// `include_attributes: [names...]` pub fn only(names: []const []const u8) Attributes { return .{ .some = names }; }
pub fn jsonStringify(self: Attributes, jw: *json.Stringify) !void { switch (self) { .every => try jw.write(true), .some => |names| try jw.write(names), } }};
/// a document id: string (≤64 bytes), unsigned integer, or uuid string.pub const Id = union(enum) { string: []const u8, integer: u64,
pub fn jsonStringify(self: Id, jw: *json.Stringify) !void { switch (self) { .string => |s| try jw.write(s), .integer => |i| try jw.write(i), } }
pub fn jsonParse(alloc: Allocator, source: anytype, options: json.ParseOptions) !Id { const value = try json.Value.jsonParse(alloc, source, options); return fromValue(value) orelse error.UnexpectedToken; }
pub fn jsonParseFromValue(alloc: Allocator, source: json.Value, options: json.ParseOptions) !Id { _ = alloc; _ = options; return fromValue(source) orelse error.UnexpectedToken; }
pub fn fromValue(value: json.Value) ?Id { return switch (value) { .string => |s| .{ .string = s }, .integer => |i| if (i >= 0) .{ .integer = @intCast(i) } else null, else => null, }; }
pub fn eql(a: Id, b: Id) bool { return switch (a) { .string => |s| b == .string and std.mem.eql(u8, s, b.string), .integer => |i| b == .integer and i == b.integer, }; }};
/// one result row. attribute getters return null when the attribute is absent/// or has another type, which is what a pre-schema-change namespace looks like.pub const Row = struct { object: json.ObjectMap,
pub fn from(value: json.Value) ?Row { return switch (value) { .object => |o| .{ .object = o }, else => null, }; }
pub fn get(self: Row, key: []const u8) ?json.Value { return self.object.get(key); }
pub fn id(self: Row) ?Id { return Id.fromValue(self.get("id") orelse return null); }
/// `$dist`: distance for ANN, score for BM25, fused rank for RRF. pub fn dist(self: Row) ?f64 { return floatOf(self.get("$dist") orelse return null); }
pub fn string(self: Row, key: []const u8) ?[]const u8 { return stringOf(self.get(key) orelse return null); }
pub fn int(self: Row, key: []const u8) ?i64 { return switch (self.get(key) orelse return null) { .integer => |i| i, else => null, }; }
pub fn uint(self: Row, key: []const u8) ?u64 { return uintOf(self.get(key) orelse return null); }
pub fn float(self: Row, key: []const u8) ?f64 { return floatOf(self.get(key) orelse return null); }
pub fn boolean(self: Row, key: []const u8) ?bool { return switch (self.get(key) orelse return null) { .bool => |b| b, else => null, }; }
/// borrowed string array; null when the attribute is absent or not an array of strings. pub fn strings(self: Row, alloc: Allocator, key: []const u8) !?[]const []const u8 { const items = switch (self.get(key) orelse return null) { .array => |a| a.items, else => return null, }; const out = try alloc.alloc([]const u8, items.len); errdefer alloc.free(out); for (items, out) |item, *slot| slot.* = stringOf(item) orelse { alloc.free(out); return null; }; return out; }
/// caller owns the slice. requires `vector` (or the named attribute) in `include_attributes`. pub fn vector(self: Row, alloc: Allocator, key: []const u8) !?[]f32 { const items = switch (self.get(key) orelse return null) { .array => |a| a.items, else => return null, }; const out = try alloc.alloc(f32, items.len); errdefer alloc.free(out); for (items, out) |item, *slot| slot.* = @floatCast(floatOf(item) orelse { alloc.free(out); return null; }); return out; }
/// deserialize into a caller struct; unknown attributes are ignored and /// `$dist` is reachable as `@"$dist": ?f64`. memory comes from `alloc`. pub fn as(self: Row, comptime T: type, alloc: Allocator) !T { return json.parseFromValueLeaky(T, alloc, .{ .object = self.object }, .{ .ignore_unknown_fields = true }); }};
pub const Rows = struct { items: []const json.Value, index: usize = 0,
pub fn len(self: Rows) usize { return self.items.len; }
pub fn at(self: Rows, i: usize) ?Row { if (i >= self.items.len) return null; return Row.from(self.items[i]); }
pub fn next(self: *Rows) ?Row { while (self.index < self.items.len) { const value = self.items[self.index]; self.index += 1; if (Row.from(value)) |row| return row; } return null; }};
pub const Performance = struct { server_total_ms: ?u64 = null, query_execution_ms: ?u64 = null, exhaustive_search_count: ?u64 = null, approx_namespace_size: ?u64 = null, cache_hit_ratio: ?f64 = null, cache_temperature: ?[]const u8 = null,};
pub const Billing = struct { billable_logical_bytes_queried: ?u64 = null, billable_logical_bytes_returned: ?u64 = null,};
/// a query response body, or one entry of a multi-query's `results`.pub const View = struct { value: json.Value,
pub fn get(self: View, key: []const u8) ?json.Value { return switch (self.value) { .object => |o| o.get(key), else => null, }; }
pub fn rows(self: View) Rows { return .{ .items = arrayOf(self.get("rows")) }; }
/// multi-query sub-results, in request order. pub fn results(self: View) Results { return .{ .items = arrayOf(self.get("results")) }; }
pub fn aggregations(self: View) ?json.ObjectMap { return switch (self.get("aggregations") orelse return null) { .object => |o| o, else => null, }; }
pub fn aggregationGroups(self: View) Rows { return .{ .items = arrayOf(self.get("aggregation_groups")) }; }
pub fn performance(self: View) Performance { const row = Row.from(self.get("performance") orelse return .{}) orelse return .{}; return .{ .server_total_ms = row.uint("server_total_ms"), .query_execution_ms = row.uint("query_execution_ms"), .exhaustive_search_count = row.uint("exhaustive_search_count"), .approx_namespace_size = row.uint("approx_namespace_size"), .cache_hit_ratio = row.float("cache_hit_ratio"), .cache_temperature = row.string("cache_temperature"), }; }
pub fn billing(self: View) Billing { const row = Row.from(self.get("billing") orelse return .{}) orelse return .{}; return .{ .billable_logical_bytes_queried = row.uint("billable_logical_bytes_queried"), .billable_logical_bytes_returned = row.uint("billable_logical_bytes_returned"), }; }
/// every row as `T`; see `Row.as`. pub fn rowsAs(self: View, comptime T: type, alloc: Allocator) ![]T { var it = self.rows(); const out = try alloc.alloc(T, it.len()); var n: usize = 0; while (it.next()) |row| : (n += 1) out[n] = try row.as(T, alloc); return out[0..n]; }};
pub const Results = struct { items: []const json.Value, index: usize = 0,
pub fn len(self: Results) usize { return self.items.len; }
pub fn at(self: Results, i: usize) ?View { if (i >= self.items.len) return null; return .{ .value = self.items[i] }; }
pub fn next(self: *Results) ?View { if (self.index >= self.items.len) return null; defer self.index += 1; return .{ .value = self.items[self.index] }; }};
/// an owned query response. everything borrowed from it (rows, strings)/// lives until `deinit`.pub const Result = struct { parsed: json.Parsed(json.Value), request_bytes: usize = 0, response_bytes: usize,
pub fn parse(alloc: Allocator, body: []const u8) !Result { const parsed = json.parseFromSlice(json.Value, alloc, body, .{ .allocate = .alloc_always }) catch return error.UnexpectedResponse; return .{ .parsed = parsed, .response_bytes = body.len }; }
pub fn deinit(self: Result) void { self.parsed.deinit(); }
pub fn view(self: Result) View { return .{ .value = self.parsed.value }; }
pub fn rows(self: Result) Rows { return self.view().rows(); }
pub fn results(self: Result) Results { return self.view().results(); }
pub fn aggregations(self: Result) ?json.ObjectMap { return self.view().aggregations(); }
pub fn aggregationGroups(self: Result) Rows { return self.view().aggregationGroups(); }
pub fn performance(self: Result) Performance { return self.view().performance(); }
pub fn billing(self: Result) Billing { return self.view().billing(); }
pub fn rowsAs(self: Result, comptime T: type, alloc: Allocator) ![]T { return self.view().rowsAs(T, alloc); }};
pub fn stringOf(value: json.Value) ?[]const u8 { return switch (value) { .string => |s| s, else => null, };}
pub fn floatOf(value: json.Value) ?f64 { return switch (value) { .float => |f| f, .integer => |i| @floatFromInt(i), else => null, };}
pub fn uintOf(value: json.Value) ?u64 { return switch (value) { .integer => |i| if (i >= 0) @intCast(i) else null, else => null, };}
fn arrayOf(value: ?json.Value) []const json.Value { return switch (value orelse return &.{}) { .array => |a| a.items, else => &.{}, };}
fn expectJson(expected: []const u8, value: anytype) !void { const body = try json.Stringify.valueAlloc(std.testing.allocator, value, .{}); defer std.testing.allocator.free(body); try std.testing.expectEqualStrings(expected, body);}
test "rank builders" { try expectJson("[\"vector\",\"ANN\",[0.5,0.25]]", rank.ann("vector", &.{ 0.5, 0.25 })); try expectJson("[\"text\",\"BM25\",\"puffer fish\"]", rank.bm25("text", "puffer fish")); try expectJson("[\"created_at\",\"desc\"]", rank.desc("created_at")); try expectJson("[\"id\",\"asc\"]", rank.asc("id")); try expectJson( "[\"Max\",[[\"name\",\"BM25\",\"q\"],[\"caption\",\"BM25\",\"q\"]]]", rank.max(.{ rank.bm25("name", "q"), rank.bm25("caption", "q") }), ); try expectJson( "[\"Sum\",[[\"Product\",2,[\"title\",\"BM25\",\"q\"]],[\"body\",\"BM25\",\"q\"]]]", rank.sum(.{ rank.product(2, rank.bm25("title", "q")), rank.bm25("body", "q") }), );}
test "filter builders" { try expectJson("[\"kind\",\"Eq\",\"summary\"]", filter.eq("kind", "summary")); try expectJson("[\"ts\",\"Gt\",1000]", filter.gt("ts", 1000)); try expectJson("[\"tags\",\"ContainsAny\",[\"a\",\"b\"]]", filter.containsAny("tags", .{ "a", "b" })); const ids = [_][]const u8{ "x", "y" }; try expectJson("[\"id\",\"In\",[\"x\",\"y\"]]", filter.in("id", &ids)); try expectJson( "[\"And\",[[\"kind\",\"Eq\",\"observation\"],[\"Not\",[\"status\",\"Eq\",\"superseded\"]]]]", filter.@"and"(.{ filter.eq("kind", "observation"), filter.not(filter.eq("status", "superseded")) }), ); try expectJson("[\"path\",\"Glob\",\"docs/*\"]", filter.glob("path", "docs/*")); try expectJson("[\"perms\",\"NotContainsAny\",[1,2]]", filter.clause("perms", "NotContainsAny", .{ 1, 2 }));}
test "a whole query body" { try expectJson( "{\"rank_by\":[\"vector\",\"ANN\",[1,0]],\"top_k\":3,\"filters\":[\"kind\",\"Eq\",\"summary\"],\"include_attributes\":true}", .{ .rank_by = rank.ann("vector", &.{ 1, 0 }), .top_k = 3, .filters = filter.eq("kind", "summary"), .include_attributes = Attributes.all, }, ); try expectJson( "{\"queries\":[{\"rank_by\":[\"vector\",\"ANN\",[1,0]],\"top_k\":3},{\"rank_by\":[\"content\",\"BM25\",\"q\"],\"top_k\":3}],\"rerank_by\":[\"RRF\"]}", .{ .queries = .{ .{ .rank_by = rank.ann("vector", &.{ 1, 0 }), .top_k = 3 }, .{ .rank_by = rank.bm25("content", "q"), .top_k = 3 }, }, .rerank_by = .{"RRF"}, }, ); try expectJson("{\"include_attributes\":[\"a\",\"b\"]}", .{ .include_attributes = Attributes.only(&.{ "a", "b" }) });}
test "Id stringifies bare" { try expectJson("[\"id\",\"Gt\",\"abc\"]", filter.gt("id", Id{ .string = "abc" })); try expectJson("[\"id\",\"Gt\",42]", filter.gt("id", Id{ .integer = 42 }));}
const fixture = \\{"rows":[ \\ {"id":"a1","$dist":0.12,"title":"hello","ts":1700000000,"has_publication":true,"tags":["x","y"],"vector":[0.5,0.25]}, \\ {"id":7,"$dist":1,"title":null} \\],"billing":{"billable_logical_bytes_queried":1024,"billable_logical_bytes_returned":256}, \\"performance":{"cache_hit_ratio":0.5,"cache_temperature":"warm","server_total_ms":17,"exhaustive_search_count":0}};
test "Result rows and getters" { const result = try Result.parse(std.testing.allocator, fixture); defer result.deinit();
var rows = result.rows(); try std.testing.expectEqual(@as(usize, 2), rows.len());
const a = rows.next().?; try std.testing.expect(a.id().?.eql(.{ .string = "a1" })); try std.testing.expectEqual(@as(f64, 0.12), a.dist().?); try std.testing.expectEqualStrings("hello", a.string("title").?); try std.testing.expectEqual(@as(i64, 1700000000), a.int("ts").?); try std.testing.expectEqual(true, a.boolean("has_publication").?); try std.testing.expectEqual(@as(?[]const u8, null), a.string("missing")); try std.testing.expectEqual(@as(?i64, null), a.int("title"));
const tags = (try a.strings(std.testing.allocator, "tags")).?; defer std.testing.allocator.free(tags); try std.testing.expectEqualStrings("y", tags[1]);
const vec = (try a.vector(std.testing.allocator, "vector")).?; defer std.testing.allocator.free(vec); try std.testing.expectEqualSlices(f32, &.{ 0.5, 0.25 }, vec);
const b = rows.next().?; try std.testing.expect(b.id().?.eql(.{ .integer = 7 })); try std.testing.expectEqual(@as(f64, 1), b.dist().?); try std.testing.expectEqual(@as(?[]const u8, null), b.string("title")); try std.testing.expectEqual(@as(?Row, null), rows.next());
const perf = result.performance(); try std.testing.expectEqualStrings("warm", perf.cache_temperature.?); try std.testing.expectEqual(@as(u64, 17), perf.server_total_ms.?); try std.testing.expectEqual(@as(f64, 0.5), perf.cache_hit_ratio.?);
const billing = result.billing(); try std.testing.expectEqual(@as(u64, 1024), billing.billable_logical_bytes_queried.?); try std.testing.expectEqual(@as(u64, 256), billing.billable_logical_bytes_returned.?); try std.testing.expectEqual(@as(usize, 0), result.request_bytes); try std.testing.expectEqual(fixture.len, result.response_bytes);}
test "rowsAs into a caller struct" { const result = try Result.parse(std.testing.allocator, fixture); defer result.deinit();
const Hit = struct { id: Id, @"$dist": f64, title: ?[]const u8 = null }; var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); const hits = try result.rowsAs(Hit, arena.allocator()); try std.testing.expectEqual(@as(usize, 2), hits.len); try std.testing.expectEqualStrings("hello", hits[0].title.?); try std.testing.expectEqual(@as(u64, 7), hits[1].id.integer); try std.testing.expectEqual(@as(?[]const u8, null), hits[1].title);}
test "multi-query results" { const result = try Result.parse(std.testing.allocator, \\{"results":[{"rows":[{"id":1,"$dist":0.1}]},{"rows":[{"id":2,"$dist":0.2},{"id":3,"$dist":0.3}]}]} ); defer result.deinit(); try std.testing.expectEqual(@as(usize, 0), result.rows().len()); var results = result.results(); try std.testing.expectEqual(@as(usize, 2), results.len()); _ = results.next().?; const second = results.next().?; try std.testing.expectEqual(@as(usize, 2), second.rows().len()); try std.testing.expectEqual(@as(u64, 3), second.rows().at(1).?.id().?.integer); try std.testing.expectEqual(@as(?View, null), results.next());}
test "aggregations" { const result = try Result.parse(std.testing.allocator, \\{"aggregations":{"n":42},"aggregation_groups":[{"color":"blue","n":2}]} ); defer result.deinit(); try std.testing.expectEqual(@as(u64, 42), uintOf(result.aggregations().?.get("n").?).?); try std.testing.expectEqualStrings("blue", result.aggregationGroups().at(0).?.string("color").?);}
test "malformed body" { try std.testing.expectError(error.UnexpectedResponse, Result.parse(std.testing.allocator, "<html>"));}