use std::borrow::Cow;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use async_trait::async_trait;
use chrono::{DateTime, Datelike, NaiveDate, NaiveDateTime, TimeZone, Utc};
use cognee_embedding::EmbeddingEngine;
use cognee_graph::{GraphDBTrait, NodeData};
use cognee_llm::{GenerationOptions, Llm, LlmExt, Message};
use cognee_vector::VectorDB;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use cognee_session::SessionContext;
use crate::graph_retrieval::{
DEFAULT_TRIPLET_DISTANCE_PENALTY, GraphRetrievalConfig, RankedGraphEdge,
brute_force_triplet_search,
};
use crate::retrievers::SearchRetriever;
use crate::types::{
SearchContext, SearchError, SearchItem, SearchOutput, SearchParams, SearchType,
};
use crate::utils::{build_messages_with_history, render_graph_user_prompt, resolve_system_prompt};
const DEFAULT_TOP_K: usize = 10;
const DEFAULT_WIDE_SEARCH_TOP_K: usize = 100;
const TEMPORAL_DATA_TYPE: &str = "Event";
const TEMPORAL_FIELD_NAME: &str = "name";
const DEFAULT_TEMPORAL_INTERVAL_PROMPT: &str = "You are tasked with identifying relevant time periods where the answer to a given query should be searched.\nCurrent date is: `{time_now}`. Determine relevant period(s) and return structured intervals.\n\nExtraction rules:\n\n1. Query without specific timestamp: use the time period with starts_at set to None and ends_at set to now.\n2. Explicit time intervals: If the query specifies a range (e.g., from 2010 to 2020, between January and March 2023), extract both start and end dates. Always assign the earlier date to starts_at and the later date to ends_at.\n3. Single timestamp: If the query refers to one specific moment (e.g., in 2015, on March 5, 2022), set starts_at and ends_at to that same timestamp.\n4. Open-ended time references: For phrases such as \"before X\" or \"after X\", represent the unspecified side as None. For example: before 2009 → starts_at: None, ends_at: 2009; after 2009 → starts_at: 2009, ends_at: None.\n5. Current-time references (\"now\", \"current\", \"today\"): If the query explicitly refers to the present, set both starts_at and ends_at to now (the ingestion timestamp).\n6. \"Who is\" and \"Who was\" questions: These imply a general identity or biographical inquiry without a specific temporal scope. Set both starts_at and ends_at to None.\n7. Ordering rule: Always ensure the earlier date is assigned to starts_at and the later date to ends_at.\n8. No temporal information: If no valid or inferable time reference is found, set both starts_at and ends_at to None.";
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
struct QueryInterval {
starts_at: Option<String>,
ends_at: Option<String>,
}
#[derive(Debug, Clone)]
struct ParsedInterval {
start: Option<DateTime<Utc>>,
end: Option<DateTime<Utc>>,
}
impl QueryInterval {
fn parse(self) -> ParsedInterval {
ParsedInterval {
start: self
.starts_at
.as_deref()
.and_then(|value| parse_bound(value, false)),
end: self
.ends_at
.as_deref()
.and_then(|value| parse_bound(value, true)),
}
}
}
pub struct TemporalRetriever {
vector_db: Arc<dyn VectorDB>,
embedding_engine: Arc<dyn EmbeddingEngine>,
graph_db: Arc<dyn GraphDBTrait>,
llm: Arc<dyn Llm>,
top_k: usize,
wide_search_top_k: usize,
triplet_distance_penalty: f32,
feedback_influence: f32,
temporal_interval_prompt: Option<String>,
system_prompt: Option<String>,
system_prompt_path: Option<String>,
user_prompt_template: Option<String>,
generation_options: Option<GenerationOptions>,
}
impl TemporalRetriever {
#[allow(clippy::too_many_arguments)]
pub fn new(
vector_db: Arc<dyn VectorDB>,
embedding_engine: Arc<dyn EmbeddingEngine>,
graph_db: Arc<dyn GraphDBTrait>,
llm: Arc<dyn Llm>,
top_k: Option<usize>,
wide_search_top_k: Option<usize>,
triplet_distance_penalty: Option<f32>,
temporal_interval_prompt: Option<String>,
system_prompt: Option<String>,
system_prompt_path: Option<String>,
user_prompt_template: Option<String>,
generation_options: Option<GenerationOptions>,
) -> Self {
Self {
vector_db,
embedding_engine,
graph_db,
llm,
top_k: top_k.unwrap_or(DEFAULT_TOP_K),
wide_search_top_k: wide_search_top_k.unwrap_or(DEFAULT_WIDE_SEARCH_TOP_K),
triplet_distance_penalty: triplet_distance_penalty
.unwrap_or(DEFAULT_TRIPLET_DISTANCE_PENALTY),
feedback_influence: 0.0,
temporal_interval_prompt,
system_prompt,
system_prompt_path,
user_prompt_template,
generation_options,
}
}
async fn extract_interval(&self, query: &str) -> Result<Option<ParsedInterval>, SearchError> {
let now = chrono::Local::now().format("%d-%m-%Y").to_string();
let prompt_template = self
.temporal_interval_prompt
.as_deref()
.unwrap_or(DEFAULT_TEMPORAL_INTERVAL_PROMPT);
let system_prompt = prompt_template.replace("{time_now}", &now);
let interval = match self
.llm
.create_structured_output_with_messages::<QueryInterval>(
vec![
Message::system(system_prompt),
Message::user(query.to_string()),
],
self.generation_options.clone(),
)
.await
{
Ok(interval) => interval,
Err(_) => return Ok(None),
};
let parsed = interval.parse();
if parsed.start.is_none() && parsed.end.is_none() {
return Ok(None);
}
Ok(Some(parsed))
}
fn get_graph_retrieval_config(&self, params: &SearchParams) -> GraphRetrievalConfig {
GraphRetrievalConfig {
top_k: params.top_k_or(self.top_k),
wide_search_top_k: params.wide_search_top_k_or(self.wide_search_top_k),
triplet_distance_penalty: params
.triplet_distance_penalty_or(self.triplet_distance_penalty),
feedback_influence: params.feedback_influence_or(self.feedback_influence),
node_type: params.node_type.clone(),
node_name: params.node_name.clone(),
node_name_filter_operator: params
.node_name_filter_operator
.as_deref()
.unwrap_or("OR")
.to_string(),
}
}
async fn get_ranked_graph_edges(
&self,
query: &str,
params: &SearchParams,
) -> Result<Vec<RankedGraphEdge>, SearchError> {
brute_force_triplet_search(
query,
self.vector_db.as_ref(),
self.embedding_engine.as_ref(),
self.graph_db.as_ref(),
&self.get_graph_retrieval_config(params),
)
.await
}
fn ranked_edges_to_context(ranked_edges: Vec<RankedGraphEdge>) -> SearchContext {
ranked_edges
.into_iter()
.map(|edge| SearchItem {
id: None,
score: Some(edge.score),
payload: json!({
"source_id": edge.source_id,
"target_id": edge.target_id,
"relationship": edge.relationship_name,
"source_name": edge.source_name,
"target_name": edge.target_name,
"source_text": edge.source_text,
"target_text": edge.target_text,
"source_description": edge.source_description,
"target_description": edge.target_description,
}),
})
.collect()
}
async fn get_fallback_context(
&self,
query: &str,
params: &SearchParams,
) -> Result<SearchContext, SearchError> {
let ranked_edges = self.get_ranked_graph_edges(query, params).await?;
Ok(Self::ranked_edges_to_context(ranked_edges))
}
async fn rank_temporal_events(
&self,
query: &str,
event_ids: &HashSet<String>,
ranked_edges: &[RankedGraphEdge],
) -> Result<Vec<(String, f32)>, SearchError> {
let mut scores = HashMap::<String, f32>::new();
for edge in ranked_edges {
if event_ids.contains(&edge.source_id) {
let score = scores.entry(edge.source_id.clone()).or_insert(edge.score);
*score = score.max(edge.score);
}
if event_ids.contains(&edge.target_id) {
let score = scores.entry(edge.target_id.clone()).or_insert(edge.score);
*score = score.max(edge.score);
}
}
if self
.vector_db
.has_collection(TEMPORAL_DATA_TYPE, TEMPORAL_FIELD_NAME)
.await?
{
let query_embeddings = self.embedding_engine.embed(&[query]).await?;
let query_vector = query_embeddings.into_iter().next().ok_or_else(|| {
SearchError::InvalidInput("embedding engine returned no vectors".to_string())
})?;
let semantic_results = self
.vector_db
.search_similar(
TEMPORAL_DATA_TYPE,
TEMPORAL_FIELD_NAME,
&query_vector,
self.wide_search_top_k.max(self.top_k),
)
.await?;
for result in semantic_results {
let event_id = result.id.to_string();
if !event_ids.contains(&event_id) {
continue;
}
let score = scores.entry(event_id).or_insert(result.score);
*score = score.max(result.score);
}
}
let mut ranked = event_ids
.iter()
.map(|event_id| {
(
event_id.clone(),
scores.get(event_id).copied().unwrap_or(0.0),
)
})
.collect::<Vec<_>>();
ranked.sort_by(|left, right| {
right
.1
.partial_cmp(&left.1)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| left.0.cmp(&right.0))
});
Ok(ranked)
}
fn temporal_context_to_text(context: &SearchContext) -> String {
context
.iter()
.map(|item| {
if item.payload.get("event_id").is_some() {
let name = item
.payload
.get("event_name")
.and_then(Value::as_str)
.unwrap_or("Unnamed event");
let description = item
.payload
.get("event_description")
.and_then(Value::as_str)
.unwrap_or("No description");
let time = item
.payload
.get("event_time")
.and_then(Value::as_str)
.unwrap_or("unknown time");
return format!("{name} ({time}): {description}");
}
let source = item
.payload
.get("source_name")
.and_then(Value::as_str)
.or_else(|| item.payload.get("source_id").and_then(Value::as_str))
.unwrap_or("unknown_source");
let target = item
.payload
.get("target_name")
.and_then(Value::as_str)
.or_else(|| item.payload.get("target_id").and_then(Value::as_str))
.unwrap_or("unknown_target");
let relationship = item
.payload
.get("relationship")
.and_then(Value::as_str)
.or_else(|| {
item.payload
.get("relationship_name")
.and_then(Value::as_str)
})
.unwrap_or("related_to");
format!("{source} -[{relationship}]-> {target}")
})
.collect::<Vec<_>>()
.join("\n")
}
}
#[async_trait]
impl SearchRetriever for TemporalRetriever {
fn search_type(&self) -> SearchType {
SearchType::Temporal
}
async fn get_context(
&self,
query: &str,
params: &SearchParams,
) -> Result<SearchContext, SearchError> {
if self.graph_db.is_empty().await? {
return Ok(vec![]);
}
let Some(interval) = self.extract_interval(query).await? else {
return self.get_fallback_context(query, params).await;
};
// Fix 1: Use typed query to find Timestamp nodes instead of full graph scan.
let (candidate_timestamps, _) = self
.graph_db
.get_filtered_graph_data(&HashMap::from([(
Cow::Borrowed("type"),
vec![json!("Timestamp")],
)]))
.await?;
let interval_from_ms = interval.start.map(|dt| dt.timestamp_millis());
let interval_to_ms = interval.end.map(|dt| dt.timestamp_millis());
let matching_ts_ids: Vec<String> = candidate_timestamps
.into_iter()
.filter_map(|(id, props)| {
let time_at = props.get("time_at")?.as_i64()?;
is_within_interval_ms(time_at, interval_from_ms, interval_to_ms).then_some(id)
})
.collect();
// Fix 2: Collect Event nodes reachable within 1-2 hops from matching Timestamps.
let mut event_node_ids = HashSet::new();
for ts_id in &matching_ts_ids {
for node_props in self.graph_db.get_neighbors(ts_id).await? {
let node_type = node_props.get("type").and_then(|v| v.as_str());
match node_type {
Some("Event") => {
if let Some(id) = node_props.get("id").and_then(|v| v.as_str()) {
event_node_ids.insert(id.to_string());
}
}
Some("Interval") => {
// Hop through Interval node to reach Event nodes (hop 2).
if let Some(interval_id) = node_props.get("id").and_then(|v| v.as_str()) {
for inner_props in self.graph_db.get_neighbors(interval_id).await? {
if inner_props.get("type").and_then(|v| v.as_str()) == Some("Event")
&& let Some(id) = inner_props.get("id").and_then(|v| v.as_str())
{
event_node_ids.insert(id.to_string());
}
}
}
}
_ => {}
}
}
}
if event_node_ids.is_empty() {
return self.get_fallback_context(query, params).await;
}
let ranked_edges = self.get_ranked_graph_edges(query, params).await?;
let ranked_events = self
.rank_temporal_events(query, &event_node_ids, &ranked_edges)
.await?;
// Fetch Event nodes by ID for building the context payload.
let event_id_list: Vec<String> = ranked_events
.iter()
.take(params.top_k_or(self.top_k))
.map(|(id, _)| id.clone())
.collect();
let event_nodes = self.graph_db.get_nodes(&event_id_list).await?;
let nodes_by_id: HashMap<String, NodeData> =
event_id_list.into_iter().zip(event_nodes).collect();
let mut temporal_context = Vec::new();
for (event_id, score) in ranked_events.into_iter().take(params.top_k_or(self.top_k)) {
let Some(event_node) = nodes_by_id.get(&event_id) else {
continue;
};
temporal_context.push(SearchItem {
id: None,
score: Some(score),
payload: json!({
"event_id": event_id,
"event_name": extract_node_name(event_node),
"event_description": extract_node_description(event_node),
}),
});
}
if temporal_context.is_empty() {
return Ok(Self::ranked_edges_to_context(ranked_edges));
}
Ok(temporal_context)
}
async fn get_completion(
&self,
query: &str,
context: Option<SearchContext>,
session: &SessionContext,
params: &SearchParams,
) -> Result<SearchOutput, SearchError> {
let completion_context = match context {
Some(existing_context) => existing_context,
None => self.get_context(query, params).await?,
};
let system_prompt = resolve_system_prompt(
params
.system_prompt
.as_deref()
.or(self.system_prompt.as_deref()),
params
.system_prompt_path
.as_deref()
.or(self.system_prompt_path.as_deref()),
)?;
let user_prompt = render_graph_user_prompt(
self.user_prompt_template.as_deref(),
query,
&Self::temporal_context_to_text(&completion_context),
);
let messages = build_messages_with_history(system_prompt, user_prompt, session);
if let Some(schema) = ¶ms.response_schema {
let structured_value = self
.llm
.create_structured_output_with_messages_raw(
messages,
schema,
self.generation_options.clone(),
)
.await
.map_err(|e| SearchError::LlmError(e.to_string()))?;
Ok(SearchOutput::Structured(structured_value))
} else {
let completion = self
.llm
.generate(messages, self.generation_options.clone())
.await?;
Ok(SearchOutput::Text(completion.content))
}
}
}
fn extract_node_name(node_data: &NodeData) -> String {
node_data
.get("name")
.and_then(Value::as_str)
.or_else(|| node_data.get("title").and_then(Value::as_str))
.unwrap_or("Unnamed event")
.to_string()
}
fn extract_node_description(node_data: &NodeData) -> String {
node_data
.get("description")
.and_then(Value::as_str)
.or_else(|| node_data.get("text").and_then(Value::as_str))
.unwrap_or("")
.to_string()
}
// Fix 3: millisecond-based interval check for Timestamp nodes.
fn is_within_interval_ms(time_at_ms: i64, from_ms: Option<i64>, to_ms: Option<i64>) -> bool {
from_ms.is_none_or(|from| time_at_ms >= from) && to_ms.is_none_or(|to| time_at_ms <= to)
}
fn parse_bound(input: &str, is_end: bool) -> Option<DateTime<Utc>> {
let trimmed = input.trim();
if trimmed.is_empty() {
return None;
}
if let Ok(timestamp) = DateTime::parse_from_rfc3339(trimmed) {
return Some(timestamp.with_timezone(&Utc));
}
if let Ok(naive_dt) = NaiveDateTime::parse_from_str(trimmed, "%Y-%m-%d %H:%M:%S") {
return Some(Utc.from_utc_datetime(&naive_dt));
}
if let Ok(date) = NaiveDate::parse_from_str(trimmed, "%Y-%m-%d") {
return to_datetime(date, is_end);
}
if trimmed.len() == 7 {
let month_candidate = format!("{trimmed}-01");
if let Ok(date) = NaiveDate::parse_from_str(&month_candidate, "%Y-%m-%d") {
return if is_end {
let (next_year, next_month) = if date.month() == 12 {
(date.year() + 1, 1)
} else {
(date.year(), date.month() + 1)
};
let next_month_start = NaiveDate::from_ymd_opt(next_year, next_month, 1)?;
let month_end = next_month_start.pred_opt()?;
to_datetime(month_end, true)
} else {
to_datetime(date, false)
};
}
}
if trimmed.len() == 4
&& trimmed.chars().all(|character| character.is_ascii_digit())
&& let Ok(year) = trimmed.parse::<i32>()
{
let date = if is_end {
NaiveDate::from_ymd_opt(year, 12, 31)?
} else {
NaiveDate::from_ymd_opt(year, 1, 1)?
};
return to_datetime(date, is_end);
}
None
}
fn to_datetime(date: NaiveDate, is_end: bool) -> Option<DateTime<Utc>> {
let naive_dt = if is_end {
date.and_hms_opt(23, 59, 59)?
} else {
date.and_hms_opt(0, 0, 0)?
};
Some(Utc.from_utc_datetime(&naive_dt))
}
#[cfg(test)]
#[allow(
clippy::unwrap_used,
clippy::expect_used,
reason = "test code — panics are acceptable failures"
)]
mod tests {
use std::borrow::Cow;
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use cognee_embedding::EmbeddingResult;
use cognee_embedding::engine::EmbeddingEngine;
use cognee_graph::{EdgeData, GraphDBResult, GraphDBTrait, GraphNode, NodeData};
use cognee_llm::{
GenerationOptions, GenerationResponse, Llm, LlmError, LlmResult, Message, TokenUsage,
};
use cognee_vector::{SearchResult, VectorDB, VectorDBResult, VectorPoint};
use chrono::{TimeZone, Utc};
use serde_json::{Value, json};
use uuid::Uuid;
use cognee_session::SessionContext;
use super::{QueryInterval, TemporalRetriever};
use crate::graph_retrieval::RankedGraphEdge;
use crate::retrievers::SearchRetriever;
use crate::types::{SearchItem, SearchOutput, SearchParams};
struct TestEmbeddingEngine;
#[async_trait]
impl EmbeddingEngine for TestEmbeddingEngine {
async fn embed(&self, _texts: &[&str]) -> EmbeddingResult<Vec<Vec<f32>>> {
Ok(vec![vec![0.3, 0.7]])
}
fn dimension(&self) -> usize {
2
}
fn batch_size(&self) -> usize {
16
}
fn max_sequence_length(&self) -> usize {
512
}
}
struct TestVectorDb {
collections: HashMap<String, Vec<SearchResult>>,
}
impl TestVectorDb {
fn key(data_type: &str, field_name: &str) -> String {
format!("{data_type}_{field_name}")
}
}
#[async_trait]
impl VectorDB for TestVectorDb {
async fn create_collection(
&self,
_data_type: &str,
_field_name: &str,
_dimension: usize,
) -> VectorDBResult<()> {
Ok(())
}
async fn has_collection(&self, data_type: &str, field_name: &str) -> VectorDBResult<bool> {
Ok(self
.collections
.contains_key(&Self::key(data_type, field_name)))
}
async fn index_points(
&self,
_data_type: &str,
_field_name: &str,
_points: &[VectorPoint],
) -> VectorDBResult<()> {
Ok(())
}
async fn search_similar(
&self,
data_type: &str,
field_name: &str,
_query_vector: &[f32],
top_k: usize,
) -> VectorDBResult<Vec<SearchResult>> {
Ok(self
.collections
.get(&Self::key(data_type, field_name))
.cloned()
.unwrap_or_default()
.into_iter()
.take(top_k)
.collect())
}
async fn delete_collection(
&self,
_data_type: &str,
_field_name: &str,
) -> VectorDBResult<()> {
Ok(())
}
async fn delete_points(
&self,
_data_type: &str,
_field_name: &str,
_point_ids: &[Uuid],
) -> VectorDBResult<()> {
Ok(())
}
async fn collection_size(
&self,
data_type: &str,
field_name: &str,
) -> VectorDBResult<usize> {
Ok(self
.collections
.get(&Self::key(data_type, field_name))
.map(|items| items.len())
.unwrap_or_default())
}
}
struct TestGraphDb {
nodes: Vec<GraphNode>,
edges: Vec<EdgeData>,
/// Maps node_id -> list of neighbor NodeData returned by get_neighbors.
neighbors: HashMap<String, Vec<NodeData>>,
}
#[async_trait]
impl GraphDBTrait for TestGraphDb {
async fn initialize(&self) -> GraphDBResult<()> {
Ok(())
}
async fn is_empty(&self) -> GraphDBResult<bool> {
Ok(self.nodes.is_empty())
}
async fn query(
&self,
_query: &str,
_params: Option<HashMap<Cow<'static, str>, Value>>,
) -> GraphDBResult<Vec<Vec<Value>>> {
Ok(vec![])
}
async fn delete_graph(&self) -> GraphDBResult<()> {
Ok(())
}
async fn has_node(&self, _node_id: &str) -> GraphDBResult<bool> {
Ok(false)
}
async fn add_node_raw(&self, _node: serde_json::Value) -> GraphDBResult<()> {
Ok(())
}
async fn add_nodes_raw(&self, _nodes: Vec<serde_json::Value>) -> GraphDBResult<()> {
Ok(())
}
async fn delete_node(&self, _node_id: &str) -> GraphDBResult<()> {
Ok(())
}
async fn delete_nodes(&self, _node_ids: &[String]) -> GraphDBResult<()> {
Ok(())
}
async fn get_node(&self, _node_id: &str) -> GraphDBResult<Option<NodeData>> {
Ok(None)
}
async fn get_nodes(&self, node_ids: &[String]) -> GraphDBResult<Vec<NodeData>> {
let nodes_map: HashMap<&str, &NodeData> = self
.nodes
.iter()
.map(|(id, data)| (id.as_str(), data))
.collect();
Ok(node_ids
.iter()
.filter_map(|id| nodes_map.get(id.as_str()).map(|d| (*d).clone()))
.collect())
}
async fn has_edge(
&self,
_source_id: &str,
_target_id: &str,
_relationship_name: &str,
) -> GraphDBResult<bool> {
Ok(false)
}
async fn has_edges(&self, _edges: &[EdgeData]) -> GraphDBResult<Vec<EdgeData>> {
Ok(vec![])
}
async fn add_edge(
&self,
_source_id: &str,
_target_id: &str,
_relationship_name: &str,
_properties: Option<HashMap<Cow<'static, str>, Value>>,
) -> GraphDBResult<()> {
Ok(())
}
async fn add_edges(&self, _edges: &[EdgeData]) -> GraphDBResult<()> {
Ok(())
}
async fn get_edges(&self, _node_id: &str) -> GraphDBResult<Vec<EdgeData>> {
Ok(vec![])
}
async fn get_neighbors(&self, node_id: &str) -> GraphDBResult<Vec<NodeData>> {
Ok(self.neighbors.get(node_id).cloned().unwrap_or_default())
}
async fn get_connections(
&self,
_node_id: &str,
) -> GraphDBResult<Vec<(NodeData, HashMap<Cow<'static, str>, Value>, NodeData)>> {
Ok(vec![])
}
async fn get_graph_data(&self) -> GraphDBResult<(Vec<GraphNode>, Vec<EdgeData>)> {
Ok((self.nodes.clone(), self.edges.clone()))
}
async fn get_graph_metrics(
&self,
_include_optional: bool,
) -> GraphDBResult<HashMap<Cow<'static, str>, Value>> {
Ok(HashMap::new())
}
async fn get_filtered_graph_data(
&self,
_attribute_filters: &HashMap<Cow<'static, str>, Vec<Value>>,
) -> GraphDBResult<(Vec<GraphNode>, Vec<EdgeData>)> {
Ok((self.nodes.clone(), self.edges.clone()))
}
async fn get_nodeset_subgraph(
&self,
_node_type: &str,
_node_names: &[String],
_node_name_filter_operator: &str,
) -> GraphDBResult<(Vec<GraphNode>, Vec<EdgeData>)> {
Ok((self.nodes.clone(), self.edges.clone()))
}
}
#[derive(Default)]
struct TestLlm {
completion_response: String,
interval_response: Option<QueryInterval>,
fail_structured_output: bool,
last_messages: Mutex<Vec<Message>>,
/// When set, `create_structured_output_with_messages_raw` returns this
/// value instead of serializing `interval_response`. Used by tests that
/// exercise the response_schema path in `get_completion`.
structured_completion_response: Mutex<Option<Value>>,
/// Messages captured by the most recent `create_structured_output_with_messages_raw` call.
last_structured_messages: Mutex<Vec<Message>>,
}
#[async_trait]
impl Llm for TestLlm {
async fn generate(
&self,
messages: Vec<Message>,
_options: Option<GenerationOptions>,
) -> LlmResult<GenerationResponse> {
self.last_messages.lock().unwrap().clone_from(&messages);
Ok(GenerationResponse {
content: self.completion_response.clone(),
model: "test-model".to_string(),
usage: Some(TokenUsage {
prompt_tokens: 1,
completion_tokens: 1,
total_tokens: 2,
}),
finish_reason: Some("stop".to_string()),
})
}
async fn create_structured_output_with_messages_raw(
&self,
messages: Vec<Message>,
_json_schema: &serde_json::Value,
_options: Option<GenerationOptions>,
) -> LlmResult<serde_json::Value> {
self.last_structured_messages
.lock()
.unwrap()
.clone_from(&messages);
if self.fail_structured_output {
return Err(LlmError::ConfigError("forced failure".to_string()));
}
// If a custom structured completion response is set, return it.
if let Some(value) = self.structured_completion_response.lock().unwrap().clone() {
return Ok(value);
}
let response = self
.interval_response
.clone()
.ok_or_else(|| LlmError::ConfigError("missing interval response".to_string()))?;
serde_json::to_value(response).map_err(|error| LlmError::ConfigError(error.to_string()))
}
fn model(&self) -> &str {
"test-model"
}
}
fn event_node_data(id: &str, name: &str) -> NodeData {
HashMap::from([
(Cow::Borrowed("id"), json!(id)),
(Cow::Borrowed("name"), json!(name)),
(Cow::Borrowed("type"), json!("Event")),
(
Cow::Borrowed("description"),
json!(format!("Description for {name}")),
),
])
}
fn timestamp_node(id: &str, time_at_ms: i64) -> GraphNode {
(
id.to_string(),
HashMap::from([
(Cow::Borrowed("id"), json!(id)),
(Cow::Borrowed("type"), json!("Timestamp")),
(Cow::Borrowed("time_at"), json!(time_at_ms)),
]),
)
}
fn event_graph_node(id: &str, name: &str) -> GraphNode {
(id.to_string(), event_node_data(id, name))
}
#[tokio::test]
async fn returns_temporal_event_context_when_interval_matches() {
// 2024-03-15 00:00:00 UTC in milliseconds
let launch_event_ms: i64 = 1710460800000;
// 2020-01-10 00:00:00 UTC in milliseconds
let old_event_ms: i64 = 1578614400000;
let ts_in_2024 = "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa";
let ts_in_2020 = "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb";
let event_launch = "11111111-1111-1111-1111-111111111111";
let event_old = "22222222-2222-2222-2222-222222222222";
let vector_db = Arc::new(TestVectorDb {
collections: HashMap::from([
(
TestVectorDb::key("Entity", "name"),
vec![SearchResult {
id: uuid::Uuid::new_v4(),
score: 0.8,
metadata: HashMap::from([(String::from("type"), json!("entity"))]),
}],
),
(
TestVectorDb::key("Event", "name"),
vec![SearchResult {
id: uuid::Uuid::parse_str(event_launch).unwrap(),
score: 0.95,
metadata: HashMap::new(),
}],
),
]),
});
let embedding_engine = Arc::new(TestEmbeddingEngine);
let graph_db = Arc::new(TestGraphDb {
nodes: vec![
timestamp_node(ts_in_2024, launch_event_ms),
timestamp_node(ts_in_2020, old_event_ms),
event_graph_node(event_launch, "Launch event"),
event_graph_node(event_old, "Old event"),
],
edges: vec![
(
event_launch.to_string(),
ts_in_2024.to_string(),
"at".to_string(),
HashMap::new(),
),
(
event_old.to_string(),
ts_in_2020.to_string(),
"at".to_string(),
HashMap::new(),
),
],
neighbors: HashMap::from([
// The 2024 Timestamp node has the Launch event as a neighbor.
(
ts_in_2024.to_string(),
vec![event_node_data(event_launch, "Launch event")],
),
// The 2020 Timestamp node has the Old event as a neighbor.
(
ts_in_2020.to_string(),
vec![event_node_data(event_old, "Old event")],
),
]),
});
let llm = Arc::new(TestLlm {
completion_response: "temporal answer".to_string(),
interval_response: Some(QueryInterval {
starts_at: Some("2024-01-01".to_string()),
ends_at: Some("2024-12-31".to_string()),
}),
fail_structured_output: false,
last_messages: Mutex::new(vec![]),
structured_completion_response: Mutex::new(None),
last_structured_messages: Mutex::new(vec![]),
});
let retriever = TemporalRetriever::new(
vector_db,
embedding_engine,
graph_db,
llm,
Some(5),
Some(10),
Some(0.0),
None,
None,
None,
None,
None,
);
let context = retriever
.get_context("What happened in 2024?", &SearchParams::default())
.await
.unwrap();
assert_eq!(context.len(), 1);
assert_eq!(
context[0].payload.get("event_name").and_then(Value::as_str),
Some("Launch event")
);
}
// ── parse_bound tests ──────────────────────────────────────────────
#[test]
fn parse_bound_datetime_space_format() {
use chrono::{TimeZone, Utc};
let result = super::parse_bound("2024-01-15 10:30:00", false);
assert_eq!(
result,
Some(Utc.with_ymd_and_hms(2024, 1, 15, 10, 30, 0).unwrap())
);
}
#[test]
fn parse_bound_rfc3339_format() {
use chrono::{TimeZone, Utc};
let result = super::parse_bound("2024-01-15T10:30:00Z", false);
assert_eq!(
result,
Some(Utc.with_ymd_and_hms(2024, 1, 15, 10, 30, 0).unwrap())
);
}
#[test]
fn parse_bound_date_only_start() {
use chrono::{TimeZone, Utc};
let result = super::parse_bound("2024-03-15", false);
assert_eq!(
result,
Some(Utc.with_ymd_and_hms(2024, 3, 15, 0, 0, 0).unwrap())
);
}
#[test]
fn parse_bound_date_only_end() {
use chrono::{TimeZone, Utc};
let result = super::parse_bound("2024-03-15", true);
assert_eq!(
result,
Some(Utc.with_ymd_and_hms(2024, 3, 15, 23, 59, 59).unwrap())
);
}
#[test]
fn parse_bound_month_start() {
use chrono::{TimeZone, Utc};
let result = super::parse_bound("2024-03", false);
assert_eq!(
result,
Some(Utc.with_ymd_and_hms(2024, 3, 1, 0, 0, 0).unwrap())
);
}
#[test]
fn parse_bound_month_end() {
use chrono::{TimeZone, Utc};
let result = super::parse_bound("2024-03", true);
assert_eq!(
result,
Some(Utc.with_ymd_and_hms(2024, 3, 31, 23, 59, 59).unwrap())
);
}
#[test]
fn parse_bound_month_end_leap_year() {
use chrono::{TimeZone, Utc};
let result = super::parse_bound("2024-02", true);
assert_eq!(
result,
Some(Utc.with_ymd_and_hms(2024, 2, 29, 23, 59, 59).unwrap())
);
}
#[test]
fn parse_bound_month_end_non_leap_year() {
use chrono::{TimeZone, Utc};
let result = super::parse_bound("2023-02", true);
assert_eq!(
result,
Some(Utc.with_ymd_and_hms(2023, 2, 28, 23, 59, 59).unwrap())
);
}
#[test]
fn parse_bound_month_end_december_wrap() {
use chrono::{TimeZone, Utc};
let result = super::parse_bound("2024-12", true);
assert_eq!(
result,
Some(Utc.with_ymd_and_hms(2024, 12, 31, 23, 59, 59).unwrap())
);
}
#[test]
fn parse_bound_year_start() {
use chrono::{TimeZone, Utc};
let result = super::parse_bound("2024", false);
assert_eq!(
result,
Some(Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap())
);
}
#[test]
fn parse_bound_year_end() {
use chrono::{TimeZone, Utc};
let result = super::parse_bound("2024", true);
assert_eq!(
result,
Some(Utc.with_ymd_and_hms(2024, 12, 31, 23, 59, 59).unwrap())
);
}
#[test]
fn parse_bound_empty_and_whitespace_returns_none() {
assert_eq!(super::parse_bound("", false), None);
assert_eq!(super::parse_bound(" ", false), None);
}
#[test]
fn parse_bound_invalid_input_returns_none() {
assert_eq!(super::parse_bound("not-a-date", false), None);
assert_eq!(super::parse_bound("abc", false), None);
}
// ── is_within_interval_ms tests ───────────────────────────────────
#[test]
fn is_within_interval_ms_basic_cases() {
use super::is_within_interval_ms;
// In range
assert!(is_within_interval_ms(500, Some(100), Some(1000)));
// At lower boundary (inclusive)
assert!(is_within_interval_ms(100, Some(100), Some(1000)));
// At upper boundary (inclusive)
assert!(is_within_interval_ms(1000, Some(100), Some(1000)));
// Below range
assert!(!is_within_interval_ms(50, Some(100), Some(1000)));
// Above range
assert!(!is_within_interval_ms(1500, Some(100), Some(1000)));
}
#[test]
fn is_within_interval_ms_open_ended_bounds() {
use super::is_within_interval_ms;
// No lower bound (open start)
assert!(is_within_interval_ms(50, None, Some(1000)));
assert!(!is_within_interval_ms(1500, None, Some(1000)));
// No upper bound (open end)
assert!(is_within_interval_ms(1500, Some(100), None));
assert!(!is_within_interval_ms(50, Some(100), None));
// Both bounds None (everything matches)
assert!(is_within_interval_ms(0, None, None));
assert!(is_within_interval_ms(i64::MAX, None, None));
assert!(is_within_interval_ms(i64::MIN, None, None));
}
// ── QueryInterval::parse tests ────────────────────────────────────
#[test]
fn query_interval_parse_both_bounds() {
use chrono::{TimeZone, Utc};
let qi = QueryInterval {
starts_at: Some("2024-01-01".to_string()),
ends_at: Some("2024-12-31".to_string()),
};
let parsed = qi.parse();
assert_eq!(
parsed.start,
Some(Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap())
);
assert_eq!(
parsed.end,
Some(Utc.with_ymd_and_hms(2024, 12, 31, 23, 59, 59).unwrap())
);
}
#[test]
fn query_interval_parse_none_bounds() {
let qi = QueryInterval {
starts_at: None,
ends_at: None,
};
let parsed = qi.parse();
assert!(parsed.start.is_none());
assert!(parsed.end.is_none());
}
#[test]
fn query_interval_parse_partial_bounds() {
use chrono::{TimeZone, Utc};
// Only starts_at
let qi = QueryInterval {
starts_at: Some("2024-06".to_string()),
ends_at: None,
};
let parsed = qi.parse();
assert_eq!(
parsed.start,
Some(Utc.with_ymd_and_hms(2024, 6, 1, 0, 0, 0).unwrap())
);
assert!(parsed.end.is_none());
// Only ends_at
let qi = QueryInterval {
starts_at: None,
ends_at: Some("2024".to_string()),
};
let parsed = qi.parse();
assert!(parsed.start.is_none());
assert_eq!(
parsed.end,
Some(Utc.with_ymd_and_hms(2024, 12, 31, 23, 59, 59).unwrap())
);
}
#[tokio::test]
async fn falls_back_to_graph_context_when_interval_extraction_fails() {
let vector_db = Arc::new(TestVectorDb {
collections: HashMap::from([(
TestVectorDb::key("Entity", "name"),
vec![SearchResult {
id: uuid::Uuid::parse_str("33333333-3333-3333-3333-333333333333").unwrap(),
score: 0.9,
metadata: HashMap::new(),
}],
)]),
});
let embedding_engine = Arc::new(TestEmbeddingEngine);
let graph_db = Arc::new(TestGraphDb {
nodes: vec![
(
"33333333-3333-3333-3333-333333333333".to_string(),
HashMap::from([
(
Cow::Borrowed("id"),
json!("33333333-3333-3333-3333-333333333333"),
),
(Cow::Borrowed("name"), json!("Entity A")),
]),
),
(
"44444444-4444-4444-4444-444444444444".to_string(),
HashMap::from([
(
Cow::Borrowed("id"),
json!("44444444-4444-4444-4444-444444444444"),
),
(Cow::Borrowed("name"), json!("Entity B")),
]),
),
],
edges: vec![(
"33333333-3333-3333-3333-333333333333".to_string(),
"44444444-4444-4444-4444-444444444444".to_string(),
"connected_to".to_string(),
HashMap::new(),
)],
neighbors: HashMap::new(),
});
let llm = Arc::new(TestLlm {
completion_response: "fallback answer".to_string(),
interval_response: None,
fail_structured_output: true,
last_messages: Mutex::new(vec![]),
structured_completion_response: Mutex::new(None),
last_structured_messages: Mutex::new(vec![]),
});
let retriever = TemporalRetriever::new(
vector_db,
embedding_engine,
graph_db,
llm,
Some(3),
Some(10),
Some(0.0),
None,
None,
None,
None,
None,
);
let context = retriever
.get_context("What happened?", &SearchParams::default())
.await
.unwrap();
assert_eq!(context.len(), 1);
assert_eq!(
context[0]
.payload
.get("relationship")
.and_then(Value::as_str),
Some("connected_to")
);
}
fn build_retriever_with_llm(llm: TestLlm) -> TemporalRetriever {
TemporalRetriever::new(
Arc::new(TestVectorDb {
collections: HashMap::new(),
}),
Arc::new(TestEmbeddingEngine),
Arc::new(TestGraphDb {
nodes: vec![],
edges: vec![],
neighbors: HashMap::new(),
}),
Arc::new(llm),
Some(5),
Some(10),
Some(0.0),
None,
None,
None,
None,
None,
)
}
#[tokio::test]
async fn extract_interval_returns_parsed_interval_from_llm() {
let llm = TestLlm {
completion_response: String::new(),
interval_response: Some(QueryInterval {
starts_at: Some("2024-01-01".into()),
ends_at: Some("2024-12-31".into()),
}),
fail_structured_output: false,
last_messages: Mutex::new(vec![]),
..TestLlm::default()
};
let retriever = build_retriever_with_llm(llm);
let result = retriever
.extract_interval("What happened in 2024?")
.await
.unwrap();
let parsed = result.expect("should return Some(ParsedInterval)");
assert_eq!(
parsed.start,
Some(Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap())
);
assert_eq!(
parsed.end,
Some(Utc.with_ymd_and_hms(2024, 12, 31, 23, 59, 59).unwrap())
);
}
#[tokio::test]
async fn extract_interval_returns_none_when_llm_returns_none_none() {
let llm = TestLlm {
completion_response: String::new(),
interval_response: Some(QueryInterval {
starts_at: None,
ends_at: None,
}),
fail_structured_output: false,
last_messages: Mutex::new(vec![]),
..TestLlm::default()
};
let retriever = build_retriever_with_llm(llm);
let result = retriever
.extract_interval("Who is Einstein?")
.await
.unwrap();
assert!(
result.is_none(),
"both fields None means no interval detected"
);
}
#[tokio::test]
async fn extract_interval_returns_none_when_llm_fails() {
let llm = TestLlm {
completion_response: String::new(),
interval_response: None,
fail_structured_output: true,
last_messages: Mutex::new(vec![]),
..TestLlm::default()
};
let retriever = build_retriever_with_llm(llm);
let result = retriever.extract_interval("What happened?").await.unwrap();
assert!(result.is_none(), "error should be swallowed gracefully");
}
#[tokio::test]
async fn extract_interval_with_only_starts_at() {
let llm = TestLlm {
completion_response: String::new(),
interval_response: Some(QueryInterval {
starts_at: Some("2024-01-01".into()),
ends_at: None,
}),
fail_structured_output: false,
last_messages: Mutex::new(vec![]),
..TestLlm::default()
};
let retriever = build_retriever_with_llm(llm);
let result = retriever
.extract_interval("What happened after 2024?")
.await
.unwrap();
let parsed = result.expect("should return Some(ParsedInterval)");
assert_eq!(
parsed.start,
Some(Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap())
);
assert_eq!(parsed.end, None);
}
#[tokio::test]
async fn extract_interval_with_only_ends_at() {
let llm = TestLlm {
completion_response: String::new(),
interval_response: Some(QueryInterval {
starts_at: None,
ends_at: Some("2024-12-31".into()),
}),
fail_structured_output: false,
last_messages: Mutex::new(vec![]),
..TestLlm::default()
};
let retriever = build_retriever_with_llm(llm);
let result = retriever
.extract_interval("What happened before 2025?")
.await
.unwrap();
let parsed = result.expect("should return Some(ParsedInterval)");
assert_eq!(parsed.start, None);
assert_eq!(
parsed.end,
Some(Utc.with_ymd_and_hms(2024, 12, 31, 23, 59, 59).unwrap())
);
}
// ---- Phase 3: get_context edge case tests ----
fn build_retriever(
vector_db: Arc<dyn VectorDB>,
graph_db: Arc<dyn GraphDBTrait>,
llm: Arc<dyn Llm>,
) -> TemporalRetriever {
TemporalRetriever::new(
vector_db,
Arc::new(TestEmbeddingEngine),
graph_db,
llm,
Some(10),
Some(100),
Some(0.0),
None,
None,
None,
None,
None,
)
}
#[tokio::test]
async fn get_context_with_time_from_only() {
// Timestamps: 2020-06-15, 2024-01-15, 2024-07-15
let ts_2020 = "aa000000-0000-0000-0000-000000000001";
let ts_2024_jan = "aa000000-0000-0000-0000-000000000002";
let ts_2024_jul = "aa000000-0000-0000-0000-000000000003";
let ev_old = "bb000000-0000-0000-0000-000000000001";
let ev_jan = "bb000000-0000-0000-0000-000000000002";
let ev_jul = "bb000000-0000-0000-0000-000000000003";
let graph_db = Arc::new(TestGraphDb {
nodes: vec![
timestamp_node(ts_2020, 1592179200000),
timestamp_node(ts_2024_jan, 1705276800000),
timestamp_node(ts_2024_jul, 1721001600000),
event_graph_node(ev_old, "Old event"),
event_graph_node(ev_jan, "Jan event"),
event_graph_node(ev_jul, "Jul event"),
],
edges: vec![
(
ev_old.to_string(),
ts_2020.to_string(),
"at".to_string(),
HashMap::new(),
),
(
ev_jan.to_string(),
ts_2024_jan.to_string(),
"at".to_string(),
HashMap::new(),
),
(
ev_jul.to_string(),
ts_2024_jul.to_string(),
"at".to_string(),
HashMap::new(),
),
],
neighbors: HashMap::from([
(
ts_2020.to_string(),
vec![event_node_data(ev_old, "Old event")],
),
(
ts_2024_jan.to_string(),
vec![event_node_data(ev_jan, "Jan event")],
),
(
ts_2024_jul.to_string(),
vec![event_node_data(ev_jul, "Jul event")],
),
]),
});
let vector_db = Arc::new(TestVectorDb {
collections: HashMap::from([(
TestVectorDb::key("Event", "name"),
vec![
SearchResult {
id: uuid::Uuid::parse_str(ev_jan).unwrap(),
score: 0.9,
metadata: HashMap::new(),
},
SearchResult {
id: uuid::Uuid::parse_str(ev_jul).unwrap(),
score: 0.85,
metadata: HashMap::new(),
},
],
)]),
});
let llm = Arc::new(TestLlm {
completion_response: String::new(),
interval_response: Some(QueryInterval {
starts_at: Some("2024-01-01".to_string()),
ends_at: None,
}),
fail_structured_output: false,
last_messages: Mutex::new(vec![]),
..TestLlm::default()
});
let retriever = build_retriever(vector_db, graph_db, llm);
let context = retriever
.get_context("What happened since 2024?", &SearchParams::default())
.await
.unwrap();
assert_eq!(context.len(), 2, "should have 2 items (both 2024 events)");
let event_names: Vec<&str> = context
.iter()
.filter_map(|item| item.payload.get("event_name").and_then(Value::as_str))
.collect();
assert!(
event_names.contains(&"Jan event"),
"Jan event should be present"
);
assert!(
event_names.contains(&"Jul event"),
"Jul event should be present"
);
assert!(
!event_names.contains(&"Old event"),
"Old event should NOT be present"
);
}
#[tokio::test]
async fn get_context_with_time_to_only() {
let ts_2020 = "aa000000-0000-0000-0000-000000000001";
let ts_2024_jan = "aa000000-0000-0000-0000-000000000002";
let ts_2024_jul = "aa000000-0000-0000-0000-000000000003";
let ev_old = "bb000000-0000-0000-0000-000000000001";
let ev_jan = "bb000000-0000-0000-0000-000000000002";
let ev_jul = "bb000000-0000-0000-0000-000000000003";
let graph_db = Arc::new(TestGraphDb {
nodes: vec![
timestamp_node(ts_2020, 1592179200000),
timestamp_node(ts_2024_jan, 1705276800000),
timestamp_node(ts_2024_jul, 1721001600000),
event_graph_node(ev_old, "Old event"),
event_graph_node(ev_jan, "Jan event"),
event_graph_node(ev_jul, "Jul event"),
],
edges: vec![
(
ev_old.to_string(),
ts_2020.to_string(),
"at".to_string(),
HashMap::new(),
),
(
ev_jan.to_string(),
ts_2024_jan.to_string(),
"at".to_string(),
HashMap::new(),
),
(
ev_jul.to_string(),
ts_2024_jul.to_string(),
"at".to_string(),
HashMap::new(),
),
],
neighbors: HashMap::from([
(
ts_2020.to_string(),
vec![event_node_data(ev_old, "Old event")],
),
(
ts_2024_jan.to_string(),
vec![event_node_data(ev_jan, "Jan event")],
),
(
ts_2024_jul.to_string(),
vec![event_node_data(ev_jul, "Jul event")],
),
]),
});
let vector_db = Arc::new(TestVectorDb {
collections: HashMap::from([(
TestVectorDb::key("Event", "name"),
vec![SearchResult {
id: uuid::Uuid::parse_str(ev_old).unwrap(),
score: 0.88,
metadata: HashMap::new(),
}],
)]),
});
let llm = Arc::new(TestLlm {
completion_response: String::new(),
interval_response: Some(QueryInterval {
starts_at: None,
ends_at: Some("2021-12-31".to_string()),
}),
fail_structured_output: false,
last_messages: Mutex::new(vec![]),
..TestLlm::default()
});
let retriever = build_retriever(vector_db, graph_db, llm);
let context = retriever
.get_context("What happened before 2022?", &SearchParams::default())
.await
.unwrap();
assert_eq!(context.len(), 1, "should have 1 item (only the 2020 event)");
assert_eq!(
context[0].payload.get("event_name").and_then(Value::as_str),
Some("Old event")
);
}
#[tokio::test]
async fn get_context_falls_back_when_no_events_in_range() {
let ts_2020 = "aa000000-0000-0000-0000-000000000010";
let ts_2021 = "aa000000-0000-0000-0000-000000000011";
let ev_2020 = "bb000000-0000-0000-0000-000000000010";
let ev_2021 = "bb000000-0000-0000-0000-000000000011";
let entity_a = "cc000000-0000-0000-0000-000000000001";
let entity_b = "cc000000-0000-0000-0000-000000000002";
let graph_db = Arc::new(TestGraphDb {
nodes: vec![
timestamp_node(ts_2020, 1577836800000), // 2020-01-01
timestamp_node(ts_2021, 1609459200000), // 2021-01-01
event_graph_node(ev_2020, "Event 2020"),
event_graph_node(ev_2021, "Event 2021"),
(
entity_a.to_string(),
HashMap::from([
(Cow::Borrowed("id"), json!(entity_a)),
(Cow::Borrowed("name"), json!("Entity A")),
(Cow::Borrowed("type"), json!("Entity")),
]),
),
(
entity_b.to_string(),
HashMap::from([
(Cow::Borrowed("id"), json!(entity_b)),
(Cow::Borrowed("name"), json!("Entity B")),
(Cow::Borrowed("type"), json!("Entity")),
]),
),
],
edges: vec![
(
ev_2020.to_string(),
ts_2020.to_string(),
"at".to_string(),
HashMap::new(),
),
(
ev_2021.to_string(),
ts_2021.to_string(),
"at".to_string(),
HashMap::new(),
),
(
entity_a.to_string(),
entity_b.to_string(),
"connected_to".to_string(),
HashMap::new(),
),
],
neighbors: HashMap::from([
(
ts_2020.to_string(),
vec![event_node_data(ev_2020, "Event 2020")],
),
(
ts_2021.to_string(),
vec![event_node_data(ev_2021, "Event 2021")],
),
]),
});
let vector_db = Arc::new(TestVectorDb {
collections: HashMap::from([(
TestVectorDb::key("Entity", "name"),
vec![SearchResult {
id: uuid::Uuid::parse_str(entity_a).unwrap(),
score: 0.9,
metadata: HashMap::new(),
}],
)]),
});
let llm = Arc::new(TestLlm {
completion_response: String::new(),
interval_response: Some(QueryInterval {
starts_at: Some("2030".to_string()),
ends_at: Some("2031".to_string()),
}),
fail_structured_output: false,
last_messages: Mutex::new(vec![]),
..TestLlm::default()
});
let retriever = build_retriever(vector_db, graph_db, llm);
let context = retriever
.get_context("What happened in 2030?", &SearchParams::default())
.await
.unwrap();
// Falls back to graph triplet search; context should have "relationship" in payload
assert!(
!context.is_empty(),
"fallback should produce at least one result"
);
assert!(
context
.iter()
.any(|item| item.payload.get("relationship").is_some()),
"fallback context items should have 'relationship' in payload"
);
assert!(
context
.iter()
.all(|item| item.payload.get("event_id").is_none()),
"fallback context should NOT have 'event_id' (those are temporal items)"
);
}
#[tokio::test]
async fn get_context_on_empty_graph() {
let graph_db = Arc::new(TestGraphDb {
nodes: vec![],
edges: vec![],
neighbors: HashMap::new(),
});
let vector_db = Arc::new(TestVectorDb {
collections: HashMap::new(),
});
// LLM won't be called since the graph is empty
let llm = Arc::new(TestLlm {
completion_response: String::new(),
interval_response: None,
fail_structured_output: false,
last_messages: Mutex::new(vec![]),
..TestLlm::default()
});
let retriever = build_retriever(vector_db, graph_db, llm);
let context = retriever
.get_context("Anything?", &SearchParams::default())
.await
.unwrap();
assert!(
context.is_empty(),
"empty graph should return empty context"
);
}
#[tokio::test]
async fn get_context_respects_top_k() {
let mut nodes = Vec::new();
let mut edges = Vec::new();
let mut neighbors = HashMap::new();
let mut vector_results = Vec::new();
for i in 1..=5 {
let ts_id = format!("aa000000-0000-0000-0000-0000000000{i:02}");
let ev_id = format!("bb000000-0000-0000-0000-0000000000{i:02}");
let ev_name = format!("Event {i}");
// All in 2024: Jan through May
let time_ms = 1704067200000_i64 + (i as i64 - 1) * 30 * 86400 * 1000;
nodes.push(timestamp_node(&ts_id, time_ms));
nodes.push(event_graph_node(&ev_id, &ev_name));
edges.push((
ev_id.clone(),
ts_id.clone(),
"at".to_string(),
HashMap::new(),
));
neighbors.insert(ts_id, vec![event_node_data(&ev_id, &ev_name)]);
vector_results.push(SearchResult {
id: uuid::Uuid::parse_str(&ev_id).unwrap(),
score: 0.9 - (i as f32 * 0.01),
metadata: HashMap::new(),
});
}
let graph_db = Arc::new(TestGraphDb {
nodes,
edges,
neighbors,
});
let vector_db = Arc::new(TestVectorDb {
collections: HashMap::from([(TestVectorDb::key("Event", "name"), vector_results)]),
});
let llm = Arc::new(TestLlm {
completion_response: String::new(),
interval_response: Some(QueryInterval {
starts_at: Some("2024-01-01".to_string()),
ends_at: Some("2024-12-31".to_string()),
}),
fail_structured_output: false,
last_messages: Mutex::new(vec![]),
..TestLlm::default()
});
let retriever = build_retriever(vector_db, graph_db, llm);
let params = SearchParams {
top_k: Some(2),
..Default::default()
};
let context = retriever
.get_context("What happened in 2024?", ¶ms)
.await
.unwrap();
assert_eq!(context.len(), 2, "top_k=2 should limit results to 2 items");
}
#[tokio::test]
async fn get_context_2hop_interval_traversal() {
let ts1 = "aa000000-0000-0000-0000-000000000020";
let ts2 = "aa000000-0000-0000-0000-000000000021";
let interval_id = "ii000000-0000-0000-0000-000000000001";
let event_id = "bb000000-0000-0000-0000-000000000020";
let interval_node_data: NodeData = HashMap::from([
(Cow::Borrowed("id"), json!(interval_id)),
(Cow::Borrowed("type"), json!("Interval")),
(Cow::Borrowed("name"), json!("Feb-Mar 2024")),
]);
let graph_db = Arc::new(TestGraphDb {
nodes: vec![
timestamp_node(ts1, 1706745600000), // 2024-02-01
timestamp_node(ts2, 1709251200000), // 2024-03-01
(interval_id.to_string(), interval_node_data.clone()),
event_graph_node(event_id, "Team Meeting"),
],
edges: vec![(
event_id.to_string(),
interval_id.to_string(),
"during".to_string(),
HashMap::new(),
)],
neighbors: HashMap::from([
// Timestamp T1 -> Interval (1st hop)
(ts1.to_string(), vec![interval_node_data.clone()]),
// Timestamp T2 -> Interval (1st hop)
(ts2.to_string(), vec![interval_node_data]),
// Interval -> Event (2nd hop)
(
interval_id.to_string(),
vec![event_node_data(event_id, "Team Meeting")],
),
]),
});
let vector_db = Arc::new(TestVectorDb {
collections: HashMap::from([(
TestVectorDb::key("Event", "name"),
vec![SearchResult {
id: uuid::Uuid::parse_str(event_id).unwrap(),
score: 0.92,
metadata: HashMap::new(),
}],
)]),
});
let llm = Arc::new(TestLlm {
completion_response: String::new(),
interval_response: Some(QueryInterval {
starts_at: Some("2024-02".to_string()),
ends_at: Some("2024-03".to_string()),
}),
fail_structured_output: false,
last_messages: Mutex::new(vec![]),
..TestLlm::default()
});
let retriever = build_retriever(vector_db, graph_db, llm);
let context = retriever
.get_context(
"What meetings happened in Feb-Mar 2024?",
&SearchParams::default(),
)
.await
.unwrap();
assert_eq!(
context.len(),
1,
"should find 1 event via 2-hop traversal (Timestamp -> Interval -> Event)"
);
assert_eq!(
context[0].payload.get("event_name").and_then(Value::as_str),
Some("Team Meeting")
);
}
// -----------------------------------------------------------------------
// Phase 4 — get_completion unit tests
// -----------------------------------------------------------------------
fn default_session() -> SessionContext {
SessionContext {
session_id: None,
history: vec![],
formatted_history: String::new(),
graph_context: None,
}
}
fn make_event_context() -> Vec<SearchItem> {
vec![
SearchItem {
id: None,
score: Some(0.9),
payload: json!({
"event_id": "evt-1",
"event_name": "Product Launch",
"event_description": "Launched the new product",
"event_time": "2024-03-15",
}),
},
SearchItem {
id: None,
score: Some(0.7),
payload: json!({
"event_id": "evt-2",
"event_name": "Quarterly Review",
"event_description": "Reviewed Q1 results",
"event_time": "2024-04-01",
}),
},
]
}
fn simple_retriever(llm: Arc<TestLlm>) -> TemporalRetriever {
let vector_db = Arc::new(TestVectorDb {
collections: HashMap::new(),
});
let embedding_engine = Arc::new(TestEmbeddingEngine);
let graph_db = Arc::new(TestGraphDb {
nodes: vec![],
edges: vec![],
neighbors: HashMap::new(),
});
TemporalRetriever::new(
vector_db,
embedding_engine,
graph_db,
llm,
Some(5),
Some(10),
Some(0.0),
None,
None,
None,
None,
None,
)
}
#[tokio::test]
async fn get_completion_generates_text_from_context() {
let llm = Arc::new(TestLlm {
completion_response: "The product was launched in March 2024.".to_string(),
last_messages: Mutex::new(vec![]),
last_structured_messages: Mutex::new(vec![]),
..TestLlm::default()
});
let retriever = simple_retriever(llm);
let context = make_event_context();
let session = default_session();
let output = retriever
.get_completion(
"What happened in 2024?",
Some(context),
&session,
&SearchParams::default(),
)
.await
.unwrap();
match output {
SearchOutput::Text(text) => {
assert_eq!(text, "The product was launched in March 2024.");
}
other => panic!("Expected SearchOutput::Text, got {other:?}"),
}
}
#[tokio::test]
async fn get_completion_with_provided_context_passes_to_llm() {
let llm = Arc::new(TestLlm {
completion_response: "completion result".to_string(),
last_messages: Mutex::new(vec![]),
last_structured_messages: Mutex::new(vec![]),
..TestLlm::default()
});
let retriever = simple_retriever(Arc::clone(&llm));
let context = make_event_context();
let session = default_session();
retriever
.get_completion(
"What happened in 2024?",
Some(context),
&session,
&SearchParams::default(),
)
.await
.unwrap();
let messages = llm.last_messages.lock().unwrap();
assert_eq!(messages.len(), 2, "Expected system + user messages");
// The user prompt should contain the temporal context text.
let user_msg = &messages[1].content;
assert!(
user_msg.contains("Product Launch"),
"User prompt should contain event name from context"
);
assert!(
user_msg.contains("Quarterly Review"),
"User prompt should contain second event name from context"
);
}
#[tokio::test]
async fn get_completion_without_context_calls_get_context() {
// Setup a graph with temporal data so get_context can produce context.
let launch_event_ms: i64 = 1710460800000; // 2024-03-15 UTC
let ts_id = "ts-aaa";
let event_id = "ev-111";
let vector_db = Arc::new(TestVectorDb {
collections: HashMap::from([(
TestVectorDb::key("Entity", "name"),
vec![SearchResult {
id: Uuid::new_v4(),
score: 0.8,
metadata: HashMap::new(),
}],
)]),
});
let embedding_engine = Arc::new(TestEmbeddingEngine);
let graph_db = Arc::new(TestGraphDb {
nodes: vec![
timestamp_node(ts_id, launch_event_ms),
event_graph_node(event_id, "Launch"),
],
edges: vec![],
neighbors: HashMap::from([(
ts_id.to_string(),
vec![event_node_data(event_id, "Launch")],
)]),
});
let llm = Arc::new(TestLlm {
completion_response: "answer from internal context".to_string(),
interval_response: Some(QueryInterval {
starts_at: Some("2024-01-01".to_string()),
ends_at: Some("2024-12-31".to_string()),
}),
last_messages: Mutex::new(vec![]),
last_structured_messages: Mutex::new(vec![]),
..TestLlm::default()
});
let retriever = TemporalRetriever::new(
vector_db,
embedding_engine,
graph_db,
llm.clone(),
Some(5),
Some(10),
Some(0.0),
None,
None,
None,
None,
None,
);
let session = default_session();
let output = retriever
.get_completion(
"What happened in 2024?",
None,
&session,
&SearchParams::default(),
)
.await
.unwrap();
match output {
SearchOutput::Text(text) => {
assert_eq!(text, "answer from internal context");
}
other => panic!("Expected SearchOutput::Text, got {other:?}"),
}
// Verify that the LLM's generate was called with messages containing context.
let messages = llm.last_messages.lock().unwrap();
assert!(!messages.is_empty(), "LLM generate should have been called");
let user_msg = &messages[1].content;
assert!(
user_msg.contains("Launch"),
"User prompt should reference the event from internal context"
);
}
#[tokio::test]
async fn get_completion_with_response_schema() {
let structured_value = json!({
"answer": "The product launched in 2024",
"confidence": 0.95
});
let llm = Arc::new(TestLlm {
completion_response: "should not be used".to_string(),
structured_completion_response: Mutex::new(Some(structured_value.clone())),
last_messages: Mutex::new(vec![]),
last_structured_messages: Mutex::new(vec![]),
..TestLlm::default()
});
let retriever = simple_retriever(llm);
let context = make_event_context();
let session = default_session();
let params = SearchParams {
response_schema: Some(json!({
"type": "object",
"properties": {
"answer": { "type": "string" },
"confidence": { "type": "number" }
}
})),
..SearchParams::default()
};
let output = retriever
.get_completion("What happened in 2024?", Some(context), &session, ¶ms)
.await
.unwrap();
match output {
SearchOutput::Structured(value) => {
assert_eq!(value, structured_value);
}
other => panic!("Expected SearchOutput::Structured, got {other:?}"),
}
}
#[tokio::test]
async fn get_completion_includes_session_history() {
let llm = Arc::new(TestLlm {
completion_response: "history-aware answer".to_string(),
last_messages: Mutex::new(vec![]),
last_structured_messages: Mutex::new(vec![]),
..TestLlm::default()
});
let retriever = simple_retriever(Arc::clone(&llm));
let context = make_event_context();
let session = SessionContext {
session_id: Some("sess-1".to_string()),
history: vec![
Message::user("Previous question?".to_string()),
Message::assistant("Previous answer.".to_string()),
],
formatted_history: "Q: Previous question?\nA: Previous answer.".to_string(),
graph_context: None,
};
retriever
.get_completion(
"Follow-up question?",
Some(context),
&session,
&SearchParams::default(),
)
.await
.unwrap();
let messages = llm.last_messages.lock().unwrap();
assert_eq!(messages.len(), 2, "Expected system + user messages");
// The system prompt should contain session history (prepended via TASK:).
let system_msg = &messages[0].content;
assert!(
system_msg.contains("Previous question?"),
"System prompt should include session history"
);
assert!(
system_msg.contains("Previous answer."),
"System prompt should include session history answer"
);
}
// -----------------------------------------------------------------------
// Phase 5 — rank_temporal_events unit tests
// -----------------------------------------------------------------------
fn make_ranked_edge(source_id: &str, target_id: &str, score: f32) -> RankedGraphEdge {
RankedGraphEdge {
source_id: source_id.to_string(),
target_id: target_id.to_string(),
relationship_name: "related_to".to_string(),
score,
source_name: format!("Source-{source_id}"),
target_name: format!("Target-{target_id}"),
dataset_id: None,
source_text: None,
target_text: None,
source_description: None,
target_description: None,
}
}
fn ranking_retriever() -> TemporalRetriever {
let vector_db = Arc::new(TestVectorDb {
collections: HashMap::from([(
TestVectorDb::key("Event", "name"),
vec![
SearchResult {
id: Uuid::parse_str("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa").unwrap(),
score: 0.9,
metadata: HashMap::new(),
},
SearchResult {
id: Uuid::parse_str("bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb").unwrap(),
score: 0.5,
metadata: HashMap::new(),
},
SearchResult {
id: Uuid::parse_str("cccccccc-cccc-cccc-cccc-cccccccccccc").unwrap(),
score: 0.3,
metadata: HashMap::new(),
},
],
)]),
});
let embedding_engine = Arc::new(TestEmbeddingEngine);
let graph_db = Arc::new(TestGraphDb {
nodes: vec![],
edges: vec![],
neighbors: HashMap::new(),
});
let llm = Arc::new(TestLlm::default());
TemporalRetriever::new(
vector_db,
embedding_engine,
graph_db,
llm,
Some(5),
Some(10),
Some(0.0),
None,
None,
None,
None,
None,
)
}
#[tokio::test]
async fn rank_sorts_by_combined_score() {
let retriever = ranking_retriever();
let event_ids: HashSet<String> = [
"aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
"bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb",
"cccccccc-cccc-cccc-cccc-cccccccccccc",
]
.iter()
.map(|s| s.to_string())
.collect();
let ranked_edges = vec![
make_ranked_edge("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", "other-node", 0.8),
make_ranked_edge("bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb", "other-node", 0.4),
make_ranked_edge("cccccccc-cccc-cccc-cccc-cccccccccccc", "other-node", 0.2),
];
let ranked = retriever
.rank_temporal_events("test query", &event_ids, &ranked_edges)
.await
.unwrap();
assert_eq!(ranked.len(), 3);
assert_eq!(ranked[0].0, "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa");
assert_eq!(ranked[1].0, "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb");
assert_eq!(ranked[2].0, "cccccccc-cccc-cccc-cccc-cccccccccccc");
// Verify descending order.
assert!(ranked[0].1 >= ranked[1].1);
assert!(ranked[1].1 >= ranked[2].1);
}
#[tokio::test]
async fn rank_events_not_in_vector_get_default_score() {
let retriever = ranking_retriever();
let event_ids: HashSet<String> = [
"aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
"dddddddd-dddd-dddd-dddd-dddddddddddd", // Not in vector DB
]
.iter()
.map(|s| s.to_string())
.collect();
let ranked_edges = vec![make_ranked_edge(
"aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
"other-node",
0.8,
)];
let ranked = retriever
.rank_temporal_events("test query", &event_ids, &ranked_edges)
.await
.unwrap();
assert_eq!(ranked.len(), 2);
let unknown_event = ranked
.iter()
.find(|(id, _)| id == "dddddddd-dddd-dddd-dddd-dddddddddddd")
.unwrap();
assert!(
unknown_event.1.abs() < f32::EPSILON,
"Unknown event should have score 0.0, got {}",
unknown_event.1
);
let known_event = ranked
.iter()
.find(|(id, _)| id == "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa")
.unwrap();
assert!(
known_event.1 > unknown_event.1,
"Known event should have higher score"
);
}
#[tokio::test]
async fn rank_empty_vector_results() {
let vector_db = Arc::new(TestVectorDb {
collections: HashMap::new(),
});
let embedding_engine = Arc::new(TestEmbeddingEngine);
let graph_db = Arc::new(TestGraphDb {
nodes: vec![],
edges: vec![],
neighbors: HashMap::new(),
});
let llm = Arc::new(TestLlm::default());
let retriever = TemporalRetriever::new(
vector_db,
embedding_engine,
graph_db,
llm,
Some(5),
Some(10),
Some(0.0),
None,
None,
None,
None,
None,
);
let event_ids: HashSet<String> = ["ev-1", "ev-2"].iter().map(|s| s.to_string()).collect();
let ranked_edges = vec![
make_ranked_edge("ev-1", "other", 0.6),
make_ranked_edge("other", "ev-2", 0.3),
];
let ranked = retriever
.rank_temporal_events("query", &event_ids, &ranked_edges)
.await
.unwrap();
assert_eq!(ranked.len(), 2);
assert_eq!(ranked[0].0, "ev-1");
assert!((ranked[0].1 - 0.6).abs() < f32::EPSILON);
assert_eq!(ranked[1].0, "ev-2");
assert!((ranked[1].1 - 0.3).abs() < f32::EPSILON);
}
#[tokio::test]
async fn rank_empty_event_ids() {
let retriever = ranking_retriever();
let event_ids: HashSet<String> = HashSet::new();
let ranked_edges = vec![make_ranked_edge("some-node", "other-node", 0.5)];
let ranked = retriever
.rank_temporal_events("query", &event_ids, &ranked_edges)
.await
.unwrap();
assert!(
ranked.is_empty(),
"Empty event_ids should yield empty result"
);
}
#[tokio::test]
async fn rank_mismatched_vector_ids() {
let vector_db = Arc::new(TestVectorDb {
collections: HashMap::from([(
TestVectorDb::key("Event", "name"),
vec![SearchResult {
id: Uuid::parse_str("ffffffff-ffff-ffff-ffff-ffffffffffff").unwrap(),
score: 0.99,
metadata: HashMap::new(),
}],
)]),
});
let embedding_engine = Arc::new(TestEmbeddingEngine);
let graph_db = Arc::new(TestGraphDb {
nodes: vec![],
edges: vec![],
neighbors: HashMap::new(),
});
let llm = Arc::new(TestLlm::default());
let retriever = TemporalRetriever::new(
vector_db,
embedding_engine,
graph_db,
llm,
Some(5),
Some(10),
Some(0.0),
None,
None,
None,
None,
None,
);
let event_ids: HashSet<String> =
["ev-abc", "ev-def"].iter().map(|s| s.to_string()).collect();
let ranked_edges = vec![make_ranked_edge("ev-abc", "something", 0.4)];
let ranked = retriever
.rank_temporal_events("query", &event_ids, &ranked_edges)
.await
.unwrap();
assert_eq!(ranked.len(), 2);
let ev_abc = ranked.iter().find(|(id, _)| id == "ev-abc").unwrap();
assert!((ev_abc.1 - 0.4).abs() < f32::EPSILON);
let ev_def = ranked.iter().find(|(id, _)| id == "ev-def").unwrap();
assert!(ev_def.1.abs() < f32::EPSILON);
}
// -----------------------------------------------------------------------
// Phase 6 — temporal_context_to_text unit tests
// -----------------------------------------------------------------------
#[test]
fn context_to_text_formats_event_items() {
let context = vec![
SearchItem {
id: None,
score: Some(0.9),
payload: json!({
"event_id": "evt-1",
"event_name": "Product Launch",
"event_description": "Launched the new product",
"event_time": "2024-03-15",
}),
},
SearchItem {
id: None,
score: Some(0.7),
payload: json!({
"event_id": "evt-2",
"event_name": "Quarterly Review",
"event_description": "Reviewed Q1 results",
"event_time": "2024-04-01",
}),
},
];
let text = TemporalRetriever::temporal_context_to_text(&context);
let lines: Vec<&str> = text.lines().collect();
assert_eq!(lines.len(), 2);
assert_eq!(
lines[0],
"Product Launch (2024-03-15): Launched the new product"
);
assert_eq!(
lines[1],
"Quarterly Review (2024-04-01): Reviewed Q1 results"
);
}
#[test]
fn context_to_text_formats_triplet_items() {
let context = vec![
SearchItem {
id: None,
score: Some(0.8),
payload: json!({
"source_name": "Alice",
"target_name": "Bob",
"relationship": "knows",
}),
},
SearchItem {
id: None,
score: Some(0.6),
payload: json!({
"source_name": "Company X",
"target_name": "Product Y",
"relationship": "produces",
}),
},
];
let text = TemporalRetriever::temporal_context_to_text(&context);
let lines: Vec<&str> = text.lines().collect();
assert_eq!(lines.len(), 2);
assert_eq!(lines[0], "Alice -[knows]-> Bob");
assert_eq!(lines[1], "Company X -[produces]-> Product Y");
}
#[test]
fn context_to_text_empty_context() {
let context: Vec<SearchItem> = vec![];
let text = TemporalRetriever::temporal_context_to_text(&context);
assert_eq!(text, "");
}
#[test]
fn context_to_text_missing_fields_use_defaults() {
let context = vec![SearchItem {
id: None,
score: Some(0.5),
payload: json!({
"event_id": "evt-bare",
}),
}];
let text = TemporalRetriever::temporal_context_to_text(&context);
assert_eq!(text, "Unnamed event (unknown time): No description");
}
#[test]
fn context_to_text_mixed_items() {
let context = vec![
SearchItem {
id: None,
score: Some(0.9),
payload: json!({
"event_id": "evt-1",
"event_name": "Conference",
"event_description": "Annual tech conference",
"event_time": "2024-06-15",
}),
},
SearchItem {
id: None,
score: Some(0.7),
payload: json!({
"source_name": "Speaker",
"target_name": "Conference",
"relationship": "presents_at",
}),
},
];
let text = TemporalRetriever::temporal_context_to_text(&context);
let lines: Vec<&str> = text.lines().collect();
assert_eq!(lines.len(), 2);
assert_eq!(lines[0], "Conference (2024-06-15): Annual tech conference");
assert_eq!(lines[1], "Speaker -[presents_at]-> Conference");
}
}