use async_trait::async_trait;
use thiserror::Error;
use tokio_util::sync::CancellationToken;
use toolkit_security::AccessScope;
use crate::models::{
ComponentReadiness, DeleteOutcome, DeleteRequest, Direction, EdgeKey, EdgeRef, EdgeView,
EmbeddingSpaceId, EngineCapabilities, GraphRevision, GtsTypeId, HopBudget, IngestOutcome,
IngestRequest, ItemError, LabelAssignment, LabelFilter, LabelId, LabelRecord, LabelSpec,
NodeId, NodeKey, NodeRow, NodeView, Page, ProjectionRequest, ReadSnapshot, RegisteredType,
RemainingBudget, RevisionOutcome, SearchRequest, SearchResponse, SourceNamespaceOwner,
StoreCapabilities, Subject, TenantId, TopologyPage, TopologyRequest, TruncationReason,
TypeIdSet, TypeQuery, TypeRecord, TypeRegistration, TypeRegistrationOptions,
};
pub struct StoreCtx<'a> {
pub tenant: TenantId,
pub scope: &'a AccessScope,
pub subject: Subject,
pub snapshot: Option<&'a ReadSnapshot>,
pub budget: RemainingBudget,
pub cancel: CancellationToken,
}
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum GraphStoreError {
#[error("scope unservable: {reason}")]
ScopeUnservable { reason: String },
#[error("unsupported: {what}")]
Unsupported { what: &'static str },
#[error("{} item(s) failed validation", items.len())]
Validation { items: Vec<ItemError> },
#[error("conflict: {reason}")]
Conflict { reason: String },
#[error("serialization failure")]
Serialization,
#[error("stale generation: recorded {recorded}, offered {offered}")]
StaleGeneration { recorded: i64, offered: i64 },
#[error("idempotency key reused with a different request")]
IdempotencyMismatch,
#[error("idempotency receipt expired")]
IdempotencyExpired,
#[error("not found")]
NotFound,
#[error("limit exceeded: {what}")]
LimitExceeded { what: String },
#[error("source namespace `{namespace}` is owned by another producer")]
SourceNamespaceForbidden { namespace: String },
#[error("invalid query: {what}")]
InvalidQuery { what: String },
#[error("store corrupt: {reason}")]
Corrupt { reason: String },
#[error("store unavailable: {reason}")]
Unavailable { reason: String },
#[error("deadline exceeded")]
Deadline,
#[error("cancelled")]
Cancelled,
#[error("internal store error: {0}")]
Internal(String),
}
#[async_trait]
pub trait GraphStoreV1: Send + Sync + 'static {
fn capabilities(&self) -> StoreCapabilities;
async fn register_types_with(
&self,
ctx: &StoreCtx<'_>,
batch: Vec<TypeRegistration>,
options: TypeRegistrationOptions,
) -> Result<Vec<RegisteredType>, GraphStoreError>;
async fn register_types(
&self,
ctx: &StoreCtx<'_>,
batch: Vec<TypeRegistration>,
) -> Result<Vec<TypeRecord>, GraphStoreError> {
let registered = self
.register_types_with(ctx, batch, TypeRegistrationOptions::default())
.await?;
Ok(registered.into_iter().map(|item| item.record).collect())
}
async fn get_type(
&self,
ctx: &StoreCtx<'_>,
id: &GtsTypeId,
) -> Result<TypeRecord, GraphStoreError>;
async fn list_types(
&self,
ctx: &StoreCtx<'_>,
query: TypeQuery,
) -> Result<Page<TypeRecord>, GraphStoreError>;
async fn probe_readiness(&self) -> Vec<ComponentReadiness>;
async fn list_source_namespaces(
&self,
ctx: &StoreCtx<'_>,
) -> Result<Vec<SourceNamespaceOwner>, GraphStoreError>;
async fn transfer_source_namespace(
&self,
ctx: &StoreCtx<'_>,
namespace: &str,
owner_principal: &str,
) -> Result<SourceNamespaceOwner, GraphStoreError>;
async fn resolve_type_set(
&self,
ctx: &StoreCtx<'_>,
patterns: &[String],
) -> Result<TypeIdSet, GraphStoreError>;
async fn ingest(
&self,
ctx: &StoreCtx<'_>,
req: IngestRequest,
embedding: EmbeddingPlan,
) -> Result<IngestOutcome, GraphStoreError>;
async fn soft_delete(
&self,
ctx: &StoreCtx<'_>,
req: DeleteRequest,
) -> Result<DeleteOutcome, GraphStoreError>;
async fn upsert_label(
&self,
_ctx: &StoreCtx<'_>,
_label: LabelSpec,
) -> Result<LabelRecord, GraphStoreError> {
Err(GraphStoreError::Unsupported { what: "labels" })
}
async fn delete_label(
&self,
_ctx: &StoreCtx<'_>,
_id: LabelId,
) -> Result<RevisionOutcome, GraphStoreError> {
Err(GraphStoreError::Unsupported { what: "labels" })
}
async fn list_labels(&self, _ctx: &StoreCtx<'_>) -> Result<Vec<LabelRecord>, GraphStoreError> {
Err(GraphStoreError::Unsupported { what: "labels" })
}
async fn assign_labels(
&self,
_ctx: &StoreCtx<'_>,
_req: LabelAssignment,
) -> Result<RevisionOutcome, GraphStoreError> {
Err(GraphStoreError::Unsupported { what: "labels" })
}
async fn begin_read(&self, ctx: &StoreCtx<'_>) -> Result<ReadSnapshot, GraphStoreError>;
async fn end_read(&self, snapshot: ReadSnapshot) -> Result<(), GraphStoreError>;
async fn revision(&self, ctx: &StoreCtx<'_>) -> Result<GraphRevision, GraphStoreError>;
async fn get_node(
&self,
ctx: &StoreCtx<'_>,
key: &NodeKey,
adjacency_limit: u32,
) -> Result<NodeView, GraphStoreError>;
async fn hydrate_nodes(
&self,
ctx: &StoreCtx<'_>,
ids: &[NodeId],
) -> Result<Vec<NodeView>, GraphStoreError>;
async fn node_types(
&self,
_ctx: &StoreCtx<'_>,
_ids: &[NodeId],
) -> Result<Vec<(NodeId, GtsTypeId)>, GraphStoreError> {
Err(GraphStoreError::Unsupported { what: "node_types" })
}
async fn get_edge(
&self,
ctx: &StoreCtx<'_>,
key: &EdgeKey,
) -> Result<EdgeView, GraphStoreError>;
async fn search(
&self,
ctx: &StoreCtx<'_>,
req: SearchRequest,
vector: Option<VectorArm>,
) -> Result<SearchResponse, GraphStoreError>;
async fn project_table(
&self,
ctx: &StoreCtx<'_>,
req: ProjectionRequest,
) -> Result<toolkit_odata::Page<NodeRow>, GraphStoreError>;
async fn load_topology(
&self,
ctx: &StoreCtx<'_>,
req: TopologyRequest,
) -> Result<TopologyPage, GraphStoreError>;
async fn resolve_node_ids(
&self,
ctx: &StoreCtx<'_>,
keys: &[NodeKey],
) -> Result<Vec<(NodeKey, NodeId)>, GraphStoreError>;
async fn embedding_state(
&self,
ctx: &StoreCtx<'_>,
keys: &[NodeKey],
) -> Result<Vec<Option<EmbeddingState>>, GraphStoreError>;
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct EmbeddingState {
pub input_hash: Option<String>,
pub vector_epoch: Option<i64>,
}
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum GraphEngineError {
#[error("scope not enforceable: {reason}")]
ScopeNotEnforceable { reason: String },
#[error("unsupported: {what}")]
Unsupported { what: &'static str },
#[error("engine unavailable: {reason}")]
Unavailable { reason: String },
#[error("deadline exceeded")]
Deadline,
#[error("cancelled")]
Cancelled,
#[error("internal engine error: {0}")]
Internal(String),
}
pub struct ExpandRequest {
pub frontier: Vec<NodeId>,
pub direction: Direction,
pub edge_types: Option<TypeIdSet>,
pub labels: Option<LabelFilter>,
pub budget: HopBudget,
pub with_degrees: bool,
}
pub struct ExpandResponse {
pub reached: Vec<NodeId>,
pub degrees: Vec<u32>,
pub edges: Vec<EdgeRef>,
pub truncated: Option<TruncationReason>,
pub served_by: HopBackend,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HopBackend {
Pattern,
TwoQuery,
}
pub struct EngineCursor {
pub revision: GraphRevision,
}
pub struct ShortestPathRequest {
pub from: NodeId,
pub to: NodeId,
pub max_depth: u8,
}
pub struct PathResponse {
pub nodes: Vec<NodeId>,
pub edges: Vec<EdgeRef>,
}
pub struct PatternRequest {
pub pattern: String,
}
pub struct PatternResponse {
pub rows: Vec<Vec<NodeId>>,
}
#[async_trait]
pub trait GraphEngineV1: Send + Sync + 'static {
fn capabilities(&self) -> EngineCapabilities;
async fn cursor(&self, ctx: &StoreCtx<'_>) -> Result<EngineCursor, GraphEngineError>;
async fn expand(
&self,
ctx: &StoreCtx<'_>,
req: ExpandRequest,
) -> Result<ExpandResponse, GraphEngineError>;
async fn shortest_path(
&self,
ctx: &StoreCtx<'_>,
req: ShortestPathRequest,
) -> Result<PathResponse, GraphEngineError>;
async fn match_pattern(
&self,
ctx: &StoreCtx<'_>,
req: PatternRequest,
) -> Result<PatternResponse, GraphEngineError>;
}
#[derive(Clone, Debug, PartialEq)]
pub struct NodeEmbedding {
pub vector: Option<Vec<f32>>,
pub input_hash: String,
}
impl NodeEmbedding {
#[must_use]
pub fn computed(vector: Vec<f32>, input_hash: String) -> Self {
Self {
vector: Some(vector),
input_hash,
}
}
#[must_use]
pub fn skipped(input_hash: String) -> Self {
Self {
vector: None,
input_hash,
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct VectorArm {
pub query_vector: Vec<f32>,
pub epoch: i64,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct EmbeddingPlan {
pub epoch: Option<i64>,
pub nodes: Vec<NodeEmbedding>,
}
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum EmbeddingProviderError {
#[error("provider unavailable: {reason}")]
Unavailable { reason: String },
#[error("embedding space mismatch")]
SpaceMismatch,
#[error("deadline exceeded")]
Deadline,
#[error("cancelled")]
Cancelled,
#[error("internal provider error: {0}")]
Internal(String),
}
pub struct EmbedRequest {
pub inputs: Vec<String>,
pub budget: RemainingBudget,
pub cancel: CancellationToken,
}
pub struct EmbedResponse {
pub vectors: Vec<Vec<f32>>,
pub space: EmbeddingSpaceId,
}
#[async_trait]
pub trait EmbeddingProviderV1: Send + Sync + 'static {
fn embedding_space(&self) -> &EmbeddingSpaceId;
fn dimension(&self) -> u32;
async fn embed(&self, req: EmbedRequest) -> Result<EmbedResponse, EmbeddingProviderError>;
async fn health(&self) -> Result<(), EmbeddingProviderError>;
}