use std::collections::HashMap;
use std::sync::Arc;
use async_trait::async_trait;
use bytes::Bytes;
use chrono::TimeZone;
use futures::stream::{self, BoxStream, StreamExt};
use lunaris::{Extractor, Lunaris};
use lunaris_core::storage::keyword::{KeywordHit, KeywordPort};
use lunaris_core::storage::types::{
CypherQuery, Filter, GraphResult, Lsn, QueueMsg, Row, VectorHit, WriteOp,
};
use lunaris_core::{
BiTemporal, Embedder, Episode, Hlc, HlcClock, LunarisError, StorageCapabilities, StorageError,
StoragePort,
};
use lunaris_extract::{ChunkInput, RawExtraction, RawExtractionBatch};
use parking_lot::Mutex;
use ulid::Ulid;
const DIM: usize = 768;
struct StubEmbedder;
#[async_trait]
impl Embedder for StubEmbedder {
fn dim(&self) -> usize {
DIM
}
async fn embed_batch(&self, inputs: &[&str]) -> Result<Vec<Vec<f32>>, LunarisError> {
Ok(inputs.iter().map(|_| vec![0.1f32; DIM]).collect())
}
}
#[derive(Default)]
struct OpRecordingStorage {
rows: Mutex<HashMap<Vec<u8>, Vec<u8>>>,
batches: Mutex<Vec<Vec<WriteOp>>>,
}
#[async_trait]
impl StoragePort for OpRecordingStorage {
async fn atomic_write(
&self,
_scope: &lunaris_core::Scope,
ops: &[WriteOp],
) -> Result<Lsn, StorageError> {
for op in ops {
if let WriteOp::KvPut { key, value } = op {
self.rows.lock().insert(key.clone(), value.clone());
}
}
self.batches.lock().push(ops.to_vec());
Ok(Lsn { wall_ms: 1, counter: 1 })
}
async fn vector_search(
&self,
_scope: &lunaris_core::Scope,
_index: &str,
_query: &[f32],
_k: usize,
_filter: Option<&Filter>,
_as_of: Option<Hlc>,
_rerank: bool,
) -> Result<Vec<VectorHit>, StorageError> {
Ok(Vec::new())
}
async fn graph_traverse(
&self,
_scope: &lunaris_core::Scope,
_q: &CypherQuery,
_as_of: Option<Hlc>,
) -> Result<GraphResult, StorageError> {
Ok(GraphResult::default())
}
async fn scan_range(
&self,
_scope: &lunaris_core::Scope,
_prefix: &[u8],
_as_of: Option<Hlc>,
) -> Result<BoxStream<'_, Result<(Bytes, Bytes), StorageError>>, StorageError> {
Ok(stream::iter(Vec::<Result<(Bytes, Bytes), StorageError>>::new()).boxed())
}
async fn read_as_of(
&self,
_scope: &lunaris_core::Scope,
key: &[u8],
_as_of: Hlc,
) -> Result<Option<Row<Bytes>>, StorageError> {
Ok(self.rows.lock().get(key).cloned().map(|v| Row {
key: key.to_vec(),
value: Bytes::from(v),
bt: BiTemporal::at(Hlc::ZERO, Hlc::ZERO),
}))
}
async fn publish(
&self,
_s: &lunaris_core::Scope,
_t: &str,
_p: u16,
_payload: Bytes,
) -> Result<u64, StorageError> {
Ok(0)
}
async fn subscribe(
&self,
_s: &lunaris_core::Scope,
_g: &str,
_t: &str,
_p: u16,
) -> Result<BoxStream<'static, Result<QueueMsg, StorageError>>, StorageError> {
Err(StorageError::NotSupported("OpRecordingStorage::subscribe"))
}
fn capabilities(&self) -> StorageCapabilities {
StorageCapabilities {
bi_temporal_native: false,
graph_native: false,
rerank_native: false,
queue_native: false,
max_vector_dim: DIM as u32,
native_rrf: false,
max_scopes_recommended: 0,
cypher_dialect: lunaris_core::CypherDialect::Legacy,
graph_decay_native: false,
graph_navigate_native: false,
}
}
}
#[async_trait]
impl KeywordPort for OpRecordingStorage {
async fn keyword_search(
&self,
_scope: &lunaris_core::Scope,
_index: &str,
_query: &str,
_k: usize,
_filter: Option<&Filter>,
_as_of: Option<Hlc>,
) -> Result<Vec<KeywordHit>, StorageError> {
Ok(Vec::new())
}
}
#[derive(Default)]
struct CapturingExtractor {
seen: Mutex<Vec<Vec<ChunkInput>>>,
}
#[async_trait]
impl Extractor for CapturingExtractor {
async fn extract(
&self,
_episode_id: Ulid,
chunks: &[ChunkInput],
) -> Result<RawExtractionBatch, LunarisError> {
self.seen.lock().push(chunks.to_vec());
Ok(RawExtractionBatch {
by_chunk: chunks
.iter()
.map(|c| RawExtraction {
source_chunk_id: c.chunk_id,
entities: vec![],
relations: vec![],
facts: vec![],
})
.collect(),
})
}
fn applies(&self) -> bool {
true
}
}
fn make_handle(
rec: Arc<OpRecordingStorage>,
extractor: Arc<CapturingExtractor>,
clock: Arc<HlcClock>,
) -> Lunaris {
let handle = Lunaris::with_parts_keyword(
rec.clone() as Arc<dyn StoragePort>,
rec as Arc<dyn KeywordPort>,
Arc::new(StubEmbedder) as Arc<dyn Embedder>,
clock,
);
handle.graph_pipeline().enable();
handle.graph_pipeline().set_extractor(extractor);
handle
}
#[tokio::test]
async fn dated_episode_threads_reference_time_to_every_chunk_input() {
let rec = Arc::new(OpRecordingStorage::default());
let extractor = Arc::new(CapturingExtractor::default());
let clock = HlcClock::new(0);
let handle = make_handle(rec, extractor.clone(), clock.clone());
let para = "The launch retrospective covered the relay outage in depth. \
Multiple teams contributed root-cause notes and follow-ups.\n\n";
let content = para.repeat(40);
let mut ep = Episode::new(lunaris_core::Scope::dev(), "dated.md", content, &clock);
ep.t_ref = Some(chrono::Utc.with_ymd_and_hms(2023, 5, 30, 23, 40, 0).unwrap());
handle.ingest(ep).await.expect("graph-ON ingest must succeed");
let seen = extractor.seen.lock();
assert!(!seen.is_empty(), "extractor must have been invoked");
let inputs: Vec<&ChunkInput> = seen.iter().flatten().collect();
assert!(!inputs.is_empty());
for (i, c) in inputs.iter().enumerate() {
assert_eq!(
c.reference_time_iso.as_deref(),
Some("2023-05-30"),
"chunk input {i} must carry the episode's date-only reference time"
);
}
}
#[tokio::test]
async fn undated_episode_leaves_reference_time_none() {
let rec = Arc::new(OpRecordingStorage::default());
let extractor = Arc::new(CapturingExtractor::default());
let clock = HlcClock::new(0);
let handle = make_handle(rec, extractor.clone(), clock.clone());
let ep = Episode::new(lunaris_core::Scope::dev(), "undated.md", "Some short note.", &clock);
assert!(ep.t_ref.is_none(), "Episode::new must not invent a t_ref");
handle.ingest(ep).await.expect("ingest must succeed");
let seen = extractor.seen.lock();
let inputs: Vec<&ChunkInput> = seen.iter().flatten().collect();
assert!(!inputs.is_empty());
for c in inputs {
assert_eq!(c.reference_time_iso, None);
}
}
#[tokio::test]
async fn working_memory_write_dated_stamps_t_ref() {
let rec = Arc::new(OpRecordingStorage::default());
let clock = HlcClock::new(0);
let handle = Lunaris::with_parts_keyword(
rec.clone() as Arc<dyn StoragePort>,
rec.clone() as Arc<dyn KeywordPort>,
Arc::new(StubEmbedder) as Arc<dyn Embedder>,
clock,
);
let wm =
lunaris::WorkingMemory::new(Arc::new(handle), lunaris_core::Scope::dev(), "helios:fs/t/");
let t = chrono::Utc.with_ymd_and_hms(2023, 6, 1, 9, 0, 0).unwrap();
wm.write_dated("s1.md", serde_json::Value::String("doc body".into()), t)
.await
.expect("write_dated must succeed");
let episodes: Vec<Episode> = rec
.rows
.lock()
.values()
.filter_map(|v| serde_json::from_slice::<Episode>(v).ok())
.filter(|e| e.source.contains("s1.md"))
.collect();
assert_eq!(episodes.len(), 1, "exactly one persisted Episode for the write");
assert_eq!(episodes[0].t_ref, Some(t), "write_dated must stamp t_ref on the Episode");
}