Advanced Features: Vector Search, Temporal Queries, and GraphRAG

# Create nodes representing documents, each with a 4-dimensional embedding
doc1 <- client$create_node(
  labels     = c("Document"),
  properties = list(title = "Graph Databases", topic = "databases"),
  embedding  = c(0.9, 0.1, 0.2, 0.05)
)

doc2 <- client$create_node(
  labels     = c("Document"),
  properties = list(title = "Vector Search", topic = "search"),
  embedding  = c(0.1, 0.85, 0.3, 0.1)
)

doc3 <- client$create_node(
  labels     = c("Document"),
  properties = list(title = "Neural Networks", topic = "ML"),
  embedding  = c(0.15, 0.8, 0.9, 0.7)
)

doc4 <- client$create_node(
  labels     = c("Document"),
  properties = list(title = "Knowledge Graphs", topic = "databases"),
  embedding  = c(0.85, 0.2, 0.25, 0.1)
)

results <- client$vector_search(
  query_vector = c(0.88, 0.15, 0.2, 0.08),
  k            = 3L
)

for (r in results) {
  cat(sprintf("Node %d -- distance: %.4f\n", r$node_id, r$distance))
}
#> Node 1 -- distance: 0.0048
#> Node 4 -- distance: 0.0129
#> Node 2 -- distance: 0.5877

# First, connect the documents with edges
client$create_edge(doc1, doc2, "RELATED_TO")
client$create_edge(doc1, doc3, "RELATED_TO")
client$create_edge(doc1, doc4, "RELATED_TO")

# Rank doc1's neighbors by similarity to a "search-oriented" concept
ranked <- client$semantic_neighbors(
  node_id   = doc1,
  concept   = c(0.1, 0.9, 0.3, 0.1),
  direction = "outgoing",
  k         = 5L
)

for (r in ranked) {
  cat(sprintf("Neighbor %d -- distance: %.4f\n", r$node_id, r$distance))
}

# Walk the graph following the "ML" concept
path <- client$semantic_walk(
  start    = doc1,
  concept  = c(0.1, 0.8, 0.9, 0.7),
  max_hops = 4L
)

cat("Walk path:", paste(path, collapse = " -> "), "\n")

# Create a social network with temporal edges
alice <- client$create_node(c("Person"), list(name = "Alice"))
bob   <- client$create_node(c("Person"), list(name = "Bob"))
carol <- client$create_node(c("Person"), list(name = "Carol"))
dave  <- client$create_node(c("Person"), list(name = "Dave"))

# Helper: convert date string to epoch milliseconds
to_epoch_ms <- function(date_str) {
  as.numeric(as.POSIXct(date_str, tz = "UTC")) * 1000
}

# Alice knew Bob from 2020 to 2022
client$create_edge(
  source     = alice,
  target     = bob,
  edge_type  = "FRIENDS",
  valid_from = to_epoch_ms("2020-01-01"),
  valid_to   = to_epoch_ms("2022-01-01")
)

# Alice has known Carol since 2021 (no end date -- still active)
client$create_edge(
  source     = alice,
  target     = carol,
  edge_type  = "FRIENDS",
  valid_from = to_epoch_ms("2021-06-01")
)

# Alice has known Dave since 2023
client$create_edge(
  source     = alice,
  target     = dave,
  edge_type  = "FRIENDS",
  valid_from = to_epoch_ms("2023-01-01")
)

# Extract a 2-hop subgraph around Alice, linearized as structured text
sg <- client$extract_subgraph(
  center    = alice,
  hops      = 2L,
  max_nodes = 20L,
  format    = "structured"
)

cat("Nodes:", sg$node_count, "\n")
cat("Edges:", sg$edge_count, "\n")
cat("\n", sg$text, "\n")

answer <- client$graph_rag(
  question  = "Who are Alice's current friends and what do they work on?",
  anchor    = alice,
  hops      = 2L,
  max_nodes = 30L,
  format    = "prose"
)

cat(answer$answer, "\n")

answer <- client$graph_rag(
  question            = "What research topics are related to graph databases?",
  question_embedding  = c(0.9, 0.1, 0.2, 0.05),
  hops                = 3L,
  max_nodes           = 50L,
  format              = "structured"
)

cat(answer$answer, "\n")

# PageRank
scores <- client$run_pagerank(damping = 0.85, max_iterations = 100L)

# Degree and betweenness centrality
deg <- client$run_degree_centrality(direction = "both")
btw <- client$run_betweenness_centrality()

# Louvain community detection
louvain <- client$run_louvain()
cat("Communities found:", louvain$num_communities, "\n")

# Connected components (weakly connected by default; strong = TRUE for SCCs)
cc <- client$run_connected_components(strong = FALSE)
cat("Number of components:", cc$count, "\n")

# Depth-first traversal (a companion to bfs())
visited <- client$dfs(alice, max_depth = 3L)

# ...as of a point in time
visited_past <- client$dfs_at(alice, max_depth = 3L, timestamp = 1672531200000)

# All node IDs carrying a label
people <- client$find_by_label("Person")

# All edges of a given type, each as {edge_id, source, target}
knows_edges <- client$find_edge_by_type("KNOWS")

# Bulk-delete every node with a label (and its edges); returns the count
removed <- client$delete_by_label("Temporary")

# Raw subgraph (nodes + edges) around a center node, for visualization
sg <- client$get_subgraph(alice, hops = 2L, max_nodes = 100L)

# Graph-wide statistics
stats <- client$graph_stats()
cat("Nodes:", stats$total_nodes, " Edges:", stats$total_edges, "\n")

# install.packages("arrow")  # if not already installed
library(AstraeaDB)

ac <- ArrowClient$new("grpc://localhost:7689")
ac$connect()

# Execute a GQL query -- returns an Arrow Table
table <- ac$query("MATCH (p:Person) RETURN p.name, p.age")

# Convert to data.frame
df <- as.data.frame(table)

# Or use the convenience method
df <- ac$query_df("MATCH (p:Person) RETURN p.name, p.age")

ac$query_batches(
  "MATCH (n) RETURN n",
  callback = function(batch) {
    cat("Received batch with", nrow(batch), "rows\n")
    # Process each batch incrementally
  }
)

ac$disconnect()

ac <- astraea_arrow_connect("grpc://localhost:7689")
# ... work ...
ac$disconnect()

uc <- UnifiedClient$new(
  host       = "127.0.0.1",
  port       = 7687L,
  flight_uri = "grpc://localhost:7689"
)
uc$connect()

# Check which transports are active
uc$is_arrow_enabled()
#> [1] TRUE

# CRUD operations go through JSON/TCP
node_id <- uc$create_node(c("Person"), list(name = "Grace", age = 42))

# Queries go through Arrow Flight (or fall back to JSON/TCP)
df <- uc$query_df("MATCH (p:Person) RETURN p.name, p.age")

# All other operations are available as usual
uc$neighbors(node_id, direction = "outgoing")
uc$vector_search(c(0.5, 0.5, 0.5, 0.5), k = 3L)

uc$disconnect()

# AstraeaClient with authentication
client <- AstraeaClient$new(
  host       = "127.0.0.1",
  port       = 7687L,
  auth_token = "my-secret-token"
)
client$connect()

# All operations now carry the token
client$ping()
client$create_node(c("Person"), list(name = "Secured"))
client$disconnect()

client <- astraea_connect(auth_token = "my-secret-token")
# ... work ...
client$disconnect()

uc <- UnifiedClient$new(
  host       = "127.0.0.1",
  port       = 7687L,
  auth_token = "my-secret-token"
)
uc$connect()
# ... work ...
uc$disconnect()

tryCatch(
  {
    client <- astraea_connect(auth_token = "wrong-token")
    client$create_node(c("Test"), list(x = 1))
  },
  error = function(e) {
    message("Auth error: ", conditionMessage(e))
  }
)