use crate::error::LaserError;
use crate::laser::Laser;
use crate::types::ConversationId;
impl Laser {
pub fn graph(&self, name: impl Into<String>) -> GraphHandle<'_> {
GraphHandle {
laser: self,
name: self.resource_name(&name.into()),
start: None,
hops: Vec::new(),
node_filter: None,
edge_filter: None,
return_: laser_wire::graph::GraphReturn::Nodes,
limit: laser_wire::limits::DEFAULT_RECALL_LIMIT,
as_of: None,
conversation: None,
stream: None,
session: None,
producer: None,
source: None,
}
}
}
pub struct GraphHandle<'a> {
laser: &'a Laser,
name: String,
start: Option<laser_wire::graph::GraphStart>,
hops: Vec<laser_wire::graph::Hop>,
node_filter: Option<laser_wire::query::Filter>,
edge_filter: Option<laser_wire::query::Filter>,
return_: laser_wire::graph::GraphReturn,
limit: usize,
as_of: Option<u64>,
conversation: Option<String>,
stream: Option<String>,
session: Option<laser_wire::agent::SessionRef>,
producer: Option<laser_wire::graph::ProducerInfo>,
source: Option<laser_wire::graph::SourceRef>,
}
impl GraphHandle<'_> {
#[must_use]
pub fn in_session(mut self, session: laser_wire::agent::SessionRef) -> Self {
self.session = Some(session);
self
}
#[must_use]
pub fn produced_by(mut self, producer: laser_wire::graph::ProducerInfo) -> Self {
self.producer = Some(producer);
self
}
#[must_use]
pub fn sourced_from(mut self, source: laser_wire::graph::SourceRef) -> Self {
self.source = Some(source);
self
}
#[must_use]
pub fn conversation(mut self, conversation: ConversationId) -> Self {
self.conversation = Some(conversation.to_string());
self.stream = self.laser.resource_stream().map(str::to_owned);
self
}
#[must_use]
pub fn start_ids(mut self, ids: Vec<laser_wire::graph::NodeId>) -> Self {
self.start = Some(laser_wire::graph::GraphStart::Ids(ids));
self
}
#[must_use]
pub fn start_match(mut self, filter: laser_wire::query::Filter) -> Self {
self.start = Some(laser_wire::graph::GraphStart::Match(filter));
self
}
#[must_use]
pub fn start_nearest(mut self, embedding: Vec<f32>, k: usize) -> Self {
self.start = Some(laser_wire::graph::GraphStart::Nearest { embedding, k });
self
}
#[must_use]
pub fn out(mut self, edge_type: impl Into<String>) -> Self {
self.hops.push(laser_wire::graph::Hop {
edge_type: Some(edge_type.into()),
dir: laser_wire::graph::EdgeDir::Out,
max: 1,
});
self
}
#[must_use]
pub fn incoming(mut self, edge_type: impl Into<String>) -> Self {
self.hops.push(laser_wire::graph::Hop {
edge_type: Some(edge_type.into()),
dir: laser_wire::graph::EdgeDir::In,
max: 1,
});
self
}
#[must_use]
pub fn both(mut self, edge_type: impl Into<String>) -> Self {
self.hops.push(laser_wire::graph::Hop {
edge_type: Some(edge_type.into()),
dir: laser_wire::graph::EdgeDir::Both,
max: 1,
});
self
}
#[must_use]
pub fn return_edges(mut self) -> Self {
self.return_ = laser_wire::graph::GraphReturn::Edges;
self
}
#[must_use]
pub fn return_triplets(mut self) -> Self {
self.return_ = laser_wire::graph::GraphReturn::Triplets;
self
}
#[must_use]
pub fn return_paths(mut self) -> Self {
self.return_ = laser_wire::graph::GraphReturn::Paths;
self
}
#[must_use]
pub fn limit(mut self, limit: usize) -> Self {
self.limit = limit;
self
}
#[must_use]
pub fn as_of(mut self, micros: u64) -> Self {
self.as_of = Some(micros);
self
}
pub async fn fetch(self) -> Result<laser_wire::graph::GraphResult, LaserError> {
use laser_wire::graph::{GraphQuery, GraphStart};
self.require_graph()?;
let query = GraphQuery {
v: laser_wire::codes::GRAPH_OP_VERSION,
graph: self.name,
start: self.start.unwrap_or(GraphStart::Ids(Vec::new())),
traverse: self.hops,
node_filter: self.node_filter,
edge_filter: self.edge_filter,
return_: self.return_,
limit: self.limit,
fork: None,
consistency: laser_wire::query::Consistency::Eventual,
as_of: self.as_of,
conversation: self.conversation,
stream: self.stream,
};
let payload = laser_wire::framing::encode_named(&query)
.map_err(|error| LaserError::Codec(format!("encode graph query: {error}")))?;
let payload = self
.laser
.send_raw_with_response(laser_wire::codes::AGDX_GRAPH_QUERY_CODE, payload)
.await?;
decode_graph_reply(&payload)
}
pub async fn neighbors(
self,
node: laser_wire::graph::NodeId,
dir: laser_wire::graph::EdgeDir,
edge_type: Option<String>,
depth: u32,
) -> Result<laser_wire::graph::GraphResult, LaserError> {
use laser_wire::graph::GraphNeighbors;
self.require_graph()?;
let request = GraphNeighbors {
v: laser_wire::codes::GRAPH_OP_VERSION,
graph: self.name,
node,
dir,
edge_type,
depth,
limit: self.limit,
as_of: self.as_of,
conversation: self.conversation,
stream: self.stream,
};
let payload = laser_wire::framing::encode_named(&request)
.map_err(|error| LaserError::Codec(format!("encode graph neighbors: {error}")))?;
let payload = self
.laser
.send_raw_with_response(laser_wire::codes::AGDX_GRAPH_NEIGHBORS_CODE, payload)
.await?;
decode_graph_reply(&payload)
}
pub async fn link(
self,
from: impl Into<String>,
relation: impl Into<String>,
to: impl Into<String>,
) -> Result<(), LaserError> {
use laser_wire::graph::GraphEdge;
let from = entity_node(&from.into());
let to = entity_node(&to.into());
let edge = GraphEdge::relate(&from, relation, &to);
self.upsert(vec![from, to], vec![edge]).await
}
pub async fn relink(
self,
from: impl Into<String>,
relation: impl Into<String>,
to: impl Into<String>,
) -> Result<usize, LaserError> {
use laser_wire::graph::EdgeDir;
let laser = self.laser;
let graph = self.name.clone();
let from = from.into();
let relation = relation.into();
let to = to.into();
let from_node = entity_node(&from);
let to_node = entity_node(&to);
let live = laser
.graph(&graph)
.neighbors(from_node.id, EdgeDir::Out, Some(relation.clone()), 1)
.await?;
let now = now_micros();
let superseded: Vec<laser_wire::graph::GraphEdge> = live
.edges
.into_iter()
.filter(|edge| edge.to != to_node.id && edge.valid_to.is_none())
.map(|mut edge| {
edge.valid_to = Some(now);
edge
})
.collect();
let closed = superseded.len();
if !superseded.is_empty() {
laser.graph(&graph).upsert(Vec::new(), superseded).await?;
}
self.link(from, relation, to).await?;
Ok(closed)
}
pub async fn unlink(
self,
from: impl Into<String>,
relation: impl Into<String>,
to: impl Into<String>,
) -> Result<(), LaserError> {
use laser_wire::graph::GraphEdge;
let from = entity_node(&from.into());
let to = entity_node(&to.into());
let mut edge = GraphEdge::relate(&from, relation, &to);
edge.valid_to = Some(now_micros());
self.upsert(Vec::new(), vec![edge]).await
}
pub async fn upsert(
self,
nodes: Vec<laser_wire::graph::GraphNode>,
edges: Vec<laser_wire::graph::GraphEdge>,
) -> Result<(), LaserError> {
use laser_wire::graph::GraphUpsert;
use laser_wire::validate::Validate;
self.require_graph()?;
let (mut nodes, mut edges) = (nodes, edges);
for node in &mut nodes {
if node.producer.is_none() {
node.producer.clone_from(&self.producer);
}
if node.source.is_none() {
node.source.clone_from(&self.source);
}
}
for edge in &mut edges {
if edge.producer.is_none() {
edge.producer.clone_from(&self.producer);
}
if edge.source.is_none() {
edge.source.clone_from(&self.source);
}
}
let request = GraphUpsert {
v: laser_wire::codes::GRAPH_OP_VERSION,
graph: self.name,
session: self.session,
nodes,
edges,
};
request.validate()?;
let payload = laser_wire::framing::encode_named(&request)
.map_err(|error| LaserError::Codec(format!("encode graph upsert: {error}")))?;
let payload = self
.laser
.send_raw_with_response(laser_wire::codes::AGDX_GRAPH_UPSERT_CODE, payload)
.await?;
decode_graph_reply(&payload).map(|_| ())
}
fn require_graph(&self) -> Result<(), LaserError> {
laser_wire::graph::validate_graph_name(&self.name)?;
if self.laser.current_capabilities().graph {
Ok(())
} else {
Err(LaserError::unsupported(
"graph",
"graph traversal is not served by this deployment",
))
}
}
}
fn decode_graph_reply(payload: &[u8]) -> Result<laser_wire::graph::GraphResult, LaserError> {
use laser_wire::graph::GraphReply;
match crate::error::decode_managed_reply::<GraphReply>(payload)? {
GraphReply::Ok(result) => Ok(result),
GraphReply::Err(error) => Err(error.into()),
_ => Err(LaserError::Protocol(
"graph: unknown reply variant".to_owned(),
)),
}
}
fn entity_node(id: &str) -> laser_wire::graph::GraphNode {
let (label, value) = id.split_once(':').unwrap_or(("entity", id));
laser_wire::graph::GraphNode::entity(label, value)
}
fn now_micros() -> u64 {
use std::time::{SystemTime, UNIX_EPOCH};
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| u64::try_from(duration.as_micros()).unwrap_or(u64::MAX))
.unwrap_or(0)
}
#[cfg(test)]
mod tests {
use crate::laser::{Laser, ResourceNaming};
use crate::types::ConversationId;
fn laser() -> Laser {
Laser::from_client(crate::iggy::prelude::IggyClient::default()).with_default_stream("acme")
}
#[test]
fn given_a_default_stream_when_opening_a_graph_then_should_scope_the_name_and_the_lens() {
let laser = laser();
let handle = laser
.graph("knowledge")
.conversation(ConversationId::derive("lens"));
assert_eq!(handle.name, "stream:acme/knowledge");
assert_eq!(handle.stream.as_deref(), Some("acme"));
}
#[test]
fn given_bare_naming_when_opening_a_graph_then_should_send_the_name_and_lens_unchanged() {
let laser = laser().with_resource_naming(ResourceNaming::Bare);
let handle = laser
.graph("knowledge")
.conversation(ConversationId::derive("lens"));
assert_eq!(handle.name, "knowledge");
assert_eq!(handle.stream, None);
}
}