use std::collections::BTreeSet;
use std::time::{Duration, Instant};
use time::OffsetDateTime;
use uuid::Uuid;
pub type TenantId = Uuid;
pub type NodeId = i64;
pub type EdgeId = i64;
pub type NodeKey = String;
pub type EdgeKey = String;
pub type GtsTypeId = String;
pub type LabelId = i32;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct UnknownVariant {
pub expected: &'static str,
pub found: String,
}
impl std::fmt::Display for UnknownVariant {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"`{}` is not a known {} in this version",
self.found, self.expected
)
}
}
impl std::error::Error for UnknownVariant {}
macro_rules! closed_enum {
($name:ident, $label:literal { $($variant:ident => $spelling:literal),+ $(,)? }) => {
impl $name {
#[must_use]
pub fn as_str(&self) -> &'static str {
match self {
$(Self::$variant => $spelling,)+
}
}
}
impl std::str::FromStr for $name {
type Err = UnknownVariant;
fn from_str(value: &str) -> Result<Self, Self::Err> {
match value {
$($spelling => Ok(Self::$variant),)+
other => Err(UnknownVariant {
expected: $label,
found: other.to_owned(),
}),
}
}
}
impl TryFrom<&str> for $name {
type Error = UnknownVariant;
fn try_from(value: &str) -> Result<Self, Self::Error> {
value.parse()
}
}
impl std::fmt::Display for $name {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
};
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TypeKind {
Node,
Edge,
Attribute,
}
closed_enum!(TypeKind, "type kind" {
Node => "node",
Edge => "edge",
Attribute => "attribute",
});
#[derive(Clone, Debug, PartialEq)]
pub struct TypeRegistration {
pub type_id: GtsTypeId,
pub schema: serde_json::Value,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct EffectiveTraits {
pub family: Option<String>,
pub scope_managed: bool,
pub emit_events: bool,
pub index: Vec<String>,
pub full_text_search: Vec<String>,
pub vector_search: Vec<String>,
pub src_types: Vec<String>,
pub dst_types: Vec<String>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ReadinessState {
Healthy,
Degraded,
Unhealthy,
NotImplemented,
}
closed_enum!(ReadinessState, "readiness state" {
Healthy => "healthy",
Degraded => "degraded",
Unhealthy => "unhealthy",
NotImplemented => "not_implemented",
});
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ComponentReadiness {
pub component: String,
pub state: ReadinessState,
pub problem: Option<String>,
pub blocked: Option<String>,
pub recovery: Option<String>,
}
impl ComponentReadiness {
#[must_use]
pub fn healthy(component: &str) -> Self {
Self {
component: component.to_owned(),
state: ReadinessState::Healthy,
problem: None,
blocked: None,
recovery: None,
}
}
#[must_use]
pub fn new(
component: &str,
state: ReadinessState,
problem: &str,
blocked: &str,
recovery: &str,
) -> Self {
Self {
component: component.to_owned(),
state,
problem: Some(problem.to_owned()),
blocked: Some(blocked.to_owned()),
recovery: Some(recovery.to_owned()),
}
}
#[must_use]
pub fn fatal(&self) -> bool {
self.state == ReadinessState::Unhealthy && self.component != EMBEDDING_SPACE
}
}
pub const DATABASE: &str = "database_and_migrations";
pub const SQLPGQ: &str = "server_major_and_sqlpgq";
pub const EMBEDDING_PROVIDER: &str = "embedding_provider";
pub const EMBEDDING_SPACE: &str = "embedding_space_identity";
pub const GRAPH_ENGINE: &str = "graph_engine_plugin";
pub const AUTHZ: &str = "authz_resolver";
pub const TYPES_REGISTRY: &str = "types_registry";
pub const DYNAMIC_INDEXES: &str = "dynamic_indexes";
pub const TENANT_RECONCILIATION: &str = "tenant_reconciliation";
pub const METRIC_ANNOTATION: &str = "metric_annotation_source";
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Readiness {
pub ready: bool,
pub components: Vec<ComponentReadiness>,
}
impl Readiness {
#[must_use]
pub fn of(components: Vec<ComponentReadiness>) -> Self {
Self {
ready: !components.iter().any(ComponentReadiness::fatal),
components,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SourceNamespaceOwner {
pub namespace: String,
pub owner_principal: String,
pub claimed_at: OffsetDateTime,
pub previous_owner: Option<String>,
pub transferred_at: Option<OffsetDateTime>,
pub transferred_by: Option<Subject>,
}
#[derive(Clone, Debug, PartialEq)]
pub struct TypeRecord {
pub type_id: GtsTypeId,
pub type_uuid: Uuid,
pub kind: TypeKind,
pub is_abstract: bool,
pub schema: serde_json::Value,
pub effective_traits: EffectiveTraits,
pub created_at: OffsetDateTime,
pub revision: i32,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct TypeQuery {
pub kind: Option<TypeKind>,
pub pattern: Option<String>,
pub top: Option<u32>,
pub cursor: Option<String>,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct TypeIdSet(pub BTreeSet<GtsTypeId>);
impl TypeIdSet {
#[must_use]
pub fn intersect(&self, other: &Self) -> Self {
Self(self.0.intersection(&other.0).cloned().collect())
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
#[must_use]
pub fn contains(&self, type_id: &str) -> bool {
self.0.contains(type_id)
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum OnExisting {
#[default]
Reject,
Update,
}
#[derive(Clone, Debug, PartialEq)]
pub enum MigrationStep {
Rename { from: String, to: String },
Default {
path: String,
value: serde_json::Value,
},
Drop { path: String },
}
impl MigrationStep {
#[must_use]
pub fn paths(&self) -> Vec<&str> {
match self {
Self::Rename { from, to } => vec![from.as_str(), to.as_str()],
Self::Default { path, .. } | Self::Drop { path } => vec![path.as_str()],
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct MigrationSpec {
pub type_id: GtsTypeId,
pub steps: Vec<MigrationStep>,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct TypeRegistrationOptions {
pub on_existing: OnExisting,
pub revalidate: bool,
pub dry_run: bool,
pub migrations: Vec<MigrationSpec>,
}
impl TypeRegistrationOptions {
#[must_use]
pub fn migration_for(&self, type_id: &str) -> Option<&MigrationSpec> {
self.migrations.iter().find(|m| m.type_id == type_id)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TypeChangeState {
New,
Unchanged,
Compatible,
Incompatible,
Undecidable,
}
closed_enum!(TypeChangeState, "type change state" {
New => "new",
Unchanged => "unchanged",
Compatible => "compatible",
Incompatible => "incompatible",
Undecidable => "undecidable",
});
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SchemaDiagnostic {
pub location: String,
pub finding: String,
pub message: String,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TraitChange {
pub trait_name: String,
pub added: Vec<String>,
pub removed: Vec<String>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TypeChange {
pub type_id: GtsTypeId,
pub state: TypeChangeState,
pub backward: String,
pub forward: String,
pub diagnostics: Vec<SchemaDiagnostic>,
pub traits_changed: Vec<TraitChange>,
pub rows: Option<u64>,
pub rows_rewritten: Option<u64>,
pub levels_not_evolvable_in_place: Vec<String>,
pub migration_required: bool,
pub admissible: bool,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AdmissionBasis {
SchemaProved,
DataBacked { rows_validated: u64 },
Migrated {
rows_scanned: u64,
rows_rewritten: u64,
},
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TypeOutcome {
Created,
Unchanged,
Updated,
}
closed_enum!(TypeOutcome, "type outcome" {
Created => "created",
Unchanged => "unchanged",
Updated => "updated",
});
#[derive(Clone, Debug, PartialEq)]
pub struct RegisteredType {
pub record: TypeRecord,
pub outcome: TypeOutcome,
pub basis: Option<AdmissionBasis>,
pub change: Option<TypeChange>,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
pub struct GraphRevision {
pub source_epoch: i64,
pub revision: i64,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ReadSnapshot {
pub id: Uuid,
pub revision: GraphRevision,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Subject {
pub subject_id: Uuid,
pub subject_type: Option<GtsTypeId>,
}
impl Subject {
#[must_use]
pub fn principal(&self) -> String {
self.subject_id.to_string()
}
#[must_use]
pub fn from_security_context(ctx: &toolkit_security::SecurityContext) -> Self {
Self {
subject_id: ctx.subject_id(),
subject_type: ctx.subject_type().map(ToOwned::to_owned),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ElementEnvelope {
pub tenant_id: Uuid,
pub key: String,
pub created_at: OffsetDateTime,
pub created_by: Subject,
pub updated_at: OffsetDateTime,
pub updated_by: Subject,
pub deleted_at: Option<OffsetDateTime>,
pub deleted_by: Option<Subject>,
pub graph_revision: GraphRevision,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct NodeSpec {
pub node_key: NodeKey,
pub type_id: GtsTypeId,
pub name: Option<String>,
pub payload: Option<serde_json::Value>,
pub expected_version: Option<i64>,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct EdgeSpec {
pub type_id: GtsTypeId,
pub src_node_key: NodeKey,
pub dst_node_key: NodeKey,
pub discriminator: Option<String>,
pub payload: Option<serde_json::Value>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ReplaceScope {
pub attribute: String,
pub value: String,
pub generation: i64,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct IngestOptions {
pub create_phantoms: Option<bool>,
pub report_per_item: bool,
pub embed: Option<bool>,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct IngestRequest {
pub nodes: Vec<NodeSpec>,
pub edges: Vec<EdgeSpec>,
pub options: IngestOptions,
pub replace_scope: Option<ReplaceScope>,
pub idempotency_key: Option<String>,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct IngestCounts {
pub nodes_inserted: u64,
pub nodes_updated: u64,
pub nodes_unchanged: u64,
pub edges_inserted: u64,
pub edges_updated: u64,
pub edges_unchanged: u64,
pub phantoms_created: u64,
pub phantoms_materialized: u64,
pub scope_removed_nodes: u64,
pub scope_removed_edges: u64,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ItemFamily {
Node,
Edge,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ItemOutcome {
Inserted,
Updated,
Unchanged,
Materialized,
}
closed_enum!(ItemOutcome, "item outcome" {
Inserted => "inserted",
Updated => "updated",
Unchanged => "unchanged",
Materialized => "materialized",
});
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ItemError {
pub index: usize,
pub family: ItemFamily,
pub gts_type: Option<GtsTypeId>,
pub pointer: Option<String>,
pub message: String,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct IngestOutcome {
pub revision: GraphRevision,
pub replayed: bool,
pub counts: IngestCounts,
pub per_item_nodes: Option<Vec<ItemOutcome>>,
pub per_item_edges: Option<Vec<ItemOutcome>>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum DeleteRequest {
Node(NodeKey),
Edge(EdgeKey),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct DeleteOutcome {
pub revision: GraphRevision,
pub tombstoned_nodes: u64,
pub tombstoned_edges: u64,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AdjacencySide {
Outgoing,
Incoming,
}
closed_enum!(AdjacencySide, "adjacency side" {
Outgoing => "outgoing",
Incoming => "incoming",
});
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AdjacencyEntry {
pub edge_key: EdgeKey,
pub edge_type_id: GtsTypeId,
pub side: AdjacencySide,
pub neighbor_key: NodeKey,
pub neighbor_type_id: GtsTypeId,
}
#[derive(Clone, Debug, PartialEq)]
pub struct NodeView {
pub node_key: NodeKey,
pub type_id: GtsTypeId,
pub name: Option<String>,
pub payload: Option<serde_json::Value>,
pub has_embedding: bool,
pub labels: Vec<String>,
pub adjacency: Vec<AdjacencyEntry>,
pub adjacency_truncated: bool,
pub envelope: ElementEnvelope,
}
#[derive(Clone, Debug, PartialEq)]
pub struct EdgeView {
pub edge_key: EdgeKey,
pub edge_type_id: GtsTypeId,
pub src: NodeKey,
pub dst: NodeKey,
pub discriminator: Option<String>,
pub payload: Option<serde_json::Value>,
pub envelope: ElementEnvelope,
}
#[derive(Clone, Debug, PartialEq)]
pub struct NodeRow {
pub node_key: NodeKey,
pub type_id: GtsTypeId,
pub name: Option<String>,
pub payload: Option<serde_json::Value>,
pub envelope: ElementEnvelope,
}
#[derive(Clone, Debug, PartialEq)]
pub struct Page<T> {
pub items: Vec<T>,
pub next_cursor: Option<String>,
pub revision: GraphRevision,
}
#[derive(toolkit_odata_macros::ODataFilterable)]
pub struct NodeQuery {
#[odata(filter(kind = "String"))]
pub node_key: String,
#[odata(filter(kind = "String"))]
pub name: String,
#[odata(filter(kind = "DateTimeUtc"))]
pub created_at: time::OffsetDateTime,
#[odata(filter(kind = "DateTimeUtc"))]
pub updated_at: time::OffsetDateTime,
}
pub use NodeQueryFilterField as NodeFilterField;
#[derive(Clone, Debug, Default)]
pub struct ProjectionRequest {
pub type_set: Option<TypeIdSet>,
pub query: toolkit_odata::ODataQuery,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SearchArm {
Lexical,
Vector,
}
closed_enum!(SearchArm, "search arm" {
Lexical => "lexical",
Vector => "vector",
});
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SearchMode {
Lexical,
Vector,
Hybrid,
}
closed_enum!(SearchMode, "search mode" {
Lexical => "lexical",
Vector => "vector",
Hybrid => "hybrid",
});
#[derive(Clone, Debug, PartialEq)]
pub struct SearchRequest {
pub mode: SearchMode,
pub query: Option<String>,
pub arm_limit: u32,
pub limit: u32,
pub type_patterns: Vec<String>,
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct ArmHit {
pub arm: SearchArm,
pub rank: u32,
pub score: f64,
}
#[derive(Clone, Debug, PartialEq)]
pub struct SearchHit {
pub node_key: NodeKey,
pub type_id: GtsTypeId,
pub name: Option<String>,
pub score: f64,
pub arms: Vec<ArmHit>,
pub snippet: Option<String>,
}
#[derive(Clone, Debug, PartialEq)]
pub struct SearchResponse {
pub hits: Vec<SearchHit>,
pub revision: GraphRevision,
pub truncated: Option<TruncationReason>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Direction {
Outgoing,
Incoming,
Either,
}
closed_enum!(Direction, "direction" {
Outgoing => "outgoing",
Incoming => "incoming",
Either => "either",
});
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct HopBudget {
pub max_frontier: u32,
pub max_edges_scanned: u64,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TruncationReason {
FrontierCap,
EdgeScanCap,
NodeBudget,
ResponseBytes,
}
closed_enum!(TruncationReason, "truncation reason" {
FrontierCap => "frontier_cap",
EdgeScanCap => "edge_scan_cap",
NodeBudget => "node_budget",
ResponseBytes => "response_bytes",
});
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct EdgeRef {
pub edge_key: EdgeKey,
pub edge_type_id: GtsTypeId,
pub src: NodeKey,
pub dst: NodeKey,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct LabelFilter {
pub any_of: Vec<String>,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct TraverseRequest {
pub seeds: Vec<NodeKey>,
pub depth: u8,
pub edge_type_patterns: Vec<String>,
pub node_type_patterns: Vec<String>,
pub max_nodes: Option<u32>,
}
#[derive(Clone, Debug, PartialEq)]
pub struct NeighborhoodRequest {
pub root: NodeKey,
pub depth: u8,
pub node_budget: Option<u32>,
pub include_phantoms: bool,
}
#[derive(Clone, Debug, PartialEq)]
pub struct TraversalResponse {
pub nodes: Vec<NodeView>,
pub edges: Vec<EdgeRef>,
pub seeds: Vec<NodeKey>,
pub truncated: Option<TruncationReason>,
pub revision: GraphRevision,
pub consistent_snapshot: bool,
}
#[derive(Clone, Debug, PartialEq)]
pub struct LabelSpec {
pub name: String,
pub description: Option<String>,
pub style: Option<serde_json::Value>,
pub applies_to: LabelAppliesTo,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LabelAppliesTo {
Node,
Edge,
Both,
}
#[derive(Clone, Debug, PartialEq)]
pub struct LabelRecord {
pub id: LabelId,
pub spec: LabelSpec,
pub created_at: OffsetDateTime,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct LabelAssignment {
pub target: LabelTarget,
pub attach: Vec<LabelId>,
pub detach: Vec<LabelId>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum LabelTarget {
Node(NodeKey),
Edge(EdgeKey),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct RevisionOutcome {
pub revision: GraphRevision,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct TopologyRequest {
pub cursor: Option<String>,
pub page_size: Option<u32>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TopologyPage {
pub nodes: Vec<(NodeKey, GtsTypeId)>,
pub edges: Vec<EdgeRef>,
pub next_cursor: Option<String>,
pub schema_version: u32,
}
#[expect(
clippy::struct_excessive_bools,
reason = "a capability set is independent yes/no facts read by name, not a \
parameter list; collapsing them into flags would hide which \
capability a store lacks at the call site"
)]
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct StoreCapabilities {
pub scope_replace: bool,
pub snapshots: bool,
pub vector_search: bool,
pub labels: bool,
pub chunks: bool,
pub topology: bool,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct EngineCapabilities {
pub shortest_path: bool,
pub match_pattern: bool,
}
#[derive(Clone, Copy, Debug)]
pub struct RemainingBudget {
deadline: Instant,
}
impl RemainingBudget {
#[must_use]
pub fn starting_now(total: Duration) -> Self {
Self {
deadline: Instant::now() + total,
}
}
#[must_use]
pub fn remaining(&self) -> Duration {
self.deadline.saturating_duration_since(Instant::now())
}
#[must_use]
pub fn is_exhausted(&self) -> bool {
self.remaining().is_zero()
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct EmbeddingSpaceId {
pub identity_hash: String,
pub model_artifact: String,
pub tokenizer_artifact: String,
pub preprocessing: serde_json::Value,
pub pooling: serde_json::Value,
pub normalization: serde_json::Value,
pub dimension: u32,
}
impl EmbeddingSpaceId {
#[must_use]
pub fn new(
model_artifact: impl Into<String>,
tokenizer_artifact: impl Into<String>,
preprocessing: serde_json::Value,
pooling: serde_json::Value,
normalization: serde_json::Value,
dimension: u32,
) -> Self {
let model_artifact = model_artifact.into();
let tokenizer_artifact = tokenizer_artifact.into();
let identity_hash = identity_hash(
&model_artifact,
&tokenizer_artifact,
&preprocessing,
&pooling,
&normalization,
dimension,
);
Self {
identity_hash,
model_artifact,
tokenizer_artifact,
preprocessing,
pooling,
normalization,
dimension,
}
}
}
#[must_use]
pub fn canonical_json(value: &serde_json::Value) -> serde_json::Value {
match value {
serde_json::Value::Object(map) => serde_json::Value::Object(
map.iter()
.map(|(key, inner)| (key.clone(), canonical_json(inner)))
.collect::<std::collections::BTreeMap<_, _>>()
.into_iter()
.collect(),
),
serde_json::Value::Array(items) => {
serde_json::Value::Array(items.iter().map(canonical_json).collect())
}
serde_json::Value::Number(number) => serde_json::Value::Number(canonical_number(number)),
other => other.clone(),
}
}
const EXACT_INTEGER_LIMIT: f64 = 9_007_199_254_740_992.0;
fn canonical_number(number: &serde_json::Number) -> serde_json::Number {
if number.is_f64()
&& let Some(float) = number.as_f64()
&& float.fract() == 0.0
&& float.abs() < EXACT_INTEGER_LIMIT
{
#[expect(
clippy::cast_possible_truncation,
reason = "the magnitude bound above is exactly the range this cast is lossless over"
)]
return serde_json::Number::from(float as i64);
}
number.clone()
}
fn identity_hash(
model_artifact: &str,
tokenizer_artifact: &str,
preprocessing: &serde_json::Value,
pooling: &serde_json::Value,
normalization: &serde_json::Value,
dimension: u32,
) -> String {
let mut hasher = aws_lc_rs::digest::Context::new(&aws_lc_rs::digest::SHA256);
let preprocessing = canonical_json(preprocessing).to_string();
let pooling = canonical_json(pooling).to_string();
let normalization = canonical_json(normalization).to_string();
for part in [
model_artifact.as_bytes(),
tokenizer_artifact.as_bytes(),
preprocessing.as_bytes(),
pooling.as_bytes(),
normalization.as_bytes(),
&dimension.to_be_bytes(),
] {
hasher.update(&(part.len() as u64).to_be_bytes());
hasher.update(part);
}
hex::encode(hasher.finish())
}
#[cfg(test)]
mod embedding_space_tests {
use super::EmbeddingSpaceId;
fn space(pooling: &str, dimension: u32) -> EmbeddingSpaceId {
EmbeddingSpaceId::new(
"all-MiniLM-L6-v2@sha256:abc",
"bert-wordpiece@sha256:def",
serde_json::json!({ "lowercase": true }),
serde_json::json!({ "strategy": pooling }),
serde_json::json!({ "l2": true }),
dimension,
)
}
#[test]
fn the_same_identity_hashes_the_same_however_its_numbers_are_written() {
let configured = |preprocessing: &str| {
EmbeddingSpaceId::new(
"all-MiniLM-L6-v2@sha256:abc",
"bert-wordpiece@sha256:def",
serde_json::from_str(preprocessing).expect("the fixture is JSON"),
serde_json::json!({ "strategy": "mean" }),
serde_json::json!({ "l2": true }),
384,
)
};
assert_eq!(
configured(r#"{"max_length": 512}"#).identity_hash,
configured(r#"{"max_length": 512.0}"#).identity_hash
);
assert_ne!(
configured(r#"{"max_length": 512}"#).identity_hash,
configured(r#"{"max_length": 256}"#).identity_hash,
"folding spellings together must not fold values together"
);
}
#[test]
fn the_same_identity_hashes_the_same_however_the_json_is_ordered() {
let one = EmbeddingSpaceId::new(
"m",
"t",
serde_json::json!({ "a": 1, "b": 2 }),
serde_json::json!({}),
serde_json::json!({}),
384,
);
let other = EmbeddingSpaceId::new(
"m",
"t",
serde_json::json!({ "b": 2, "a": 1 }),
serde_json::json!({}),
serde_json::json!({}),
384,
);
assert_eq!(one.identity_hash, other.identity_hash);
}
#[test]
fn pooling_alone_changes_the_identity() {
assert_ne!(
space("mean", 384).identity_hash,
space("cls", 384).identity_hash
);
}
#[test]
fn dimension_alone_changes_the_identity() {
assert_ne!(
space("mean", 384).identity_hash,
space("mean", 768).identity_hash
);
}
}
#[cfg(test)]
mod closed_enum_tests {
use super::*;
macro_rules! contract_case {
($case:ident, $name:ident, $label:literal, [$($variant:expr),+ $(,)?]) => {
#[test]
fn $case() {
let mut seen: Vec<&'static str> = Vec::new();
$(
let spelling = $variant.as_str();
assert!(
!seen.contains(&spelling),
"two {} variants share the spelling `{spelling}`",
$label
);
seen.push(spelling);
assert_eq!(
spelling.parse::<$name>().expect("its own spelling decodes"),
$variant,
"{} does not round-trip through `{spelling}`",
$label
);
)+
for unknown in ["", "UNKNOWN", "healthy_", " node", "something_new"] {
assert!(!seen.contains(&unknown), "the fixture must be unknown");
let refused = unknown
.parse::<$name>()
.expect_err("an unknown value is never a known variant");
assert_eq!(refused.found, unknown);
assert_eq!(refused.expected, $label);
}
}
};
}
contract_case!(
a_type_kind,
TypeKind,
"type kind",
[TypeKind::Node, TypeKind::Edge, TypeKind::Attribute]
);
contract_case!(
a_readiness_state,
ReadinessState,
"readiness state",
[
ReadinessState::Healthy,
ReadinessState::Degraded,
ReadinessState::Unhealthy,
ReadinessState::NotImplemented,
]
);
contract_case!(
a_type_change_state,
TypeChangeState,
"type change state",
[
TypeChangeState::New,
TypeChangeState::Unchanged,
TypeChangeState::Compatible,
TypeChangeState::Incompatible,
TypeChangeState::Undecidable,
]
);
contract_case!(
a_type_outcome,
TypeOutcome,
"type outcome",
[
TypeOutcome::Created,
TypeOutcome::Unchanged,
TypeOutcome::Updated
]
);
contract_case!(
an_item_outcome,
ItemOutcome,
"item outcome",
[
ItemOutcome::Inserted,
ItemOutcome::Updated,
ItemOutcome::Unchanged,
ItemOutcome::Materialized,
]
);
contract_case!(
an_adjacency_side,
AdjacencySide,
"adjacency side",
[AdjacencySide::Outgoing, AdjacencySide::Incoming]
);
contract_case!(
a_search_arm,
SearchArm,
"search arm",
[SearchArm::Lexical, SearchArm::Vector]
);
contract_case!(
a_search_mode,
SearchMode,
"search mode",
[SearchMode::Lexical, SearchMode::Vector, SearchMode::Hybrid]
);
contract_case!(
a_direction,
Direction,
"direction",
[Direction::Outgoing, Direction::Incoming, Direction::Either]
);
contract_case!(
a_truncation_reason,
TruncationReason,
"truncation reason",
[
TruncationReason::FrontierCap,
TruncationReason::EdgeScanCap,
TruncationReason::NodeBudget,
TruncationReason::ResponseBytes,
]
);
}