use anyhow::{Context, Result, bail};
use oxibrain::Brain;
use oxibrain_core::retrieval::{Query, QueryMode};
use oxibrain_ports::{FakeClock, Timestamp};
use oxibrain_store::project::{DeclObject, Declaration, EntityRef};
use serde::Deserialize;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tempfile::TempDir;
#[derive(Debug, Deserialize)]
struct Manifest {
#[allow(dead_code)]
version: String,
#[allow(dead_code)]
categories: Vec<Category>,
episodes: Vec<EpisodeEntry>,
questions: Vec<QuestionEntry>,
}
#[derive(Debug, Deserialize)]
struct Category {
#[allow(dead_code)]
name: String,
#[allow(dead_code)]
description: String,
}
#[derive(Debug, Deserialize)]
struct EpisodeEntry {
#[allow(dead_code)]
id: String,
#[allow(dead_code)]
shape: String,
#[allow(dead_code)]
lang: String,
file: String,
}
#[derive(Debug, Deserialize)]
struct QuestionEntry {
#[allow(dead_code)]
id: String,
#[allow(dead_code)]
category: String,
file: String,
}
#[derive(Debug, Deserialize)]
struct EpisodeFile {
#[allow(dead_code)]
shape: String,
#[allow(dead_code)]
lang: String,
#[allow(dead_code)]
occurred_at: String,
#[allow(dead_code)]
content: String,
entities: Vec<EpisodeEntity>,
statements: Vec<EpisodeStatement>,
}
#[derive(Debug, Deserialize)]
struct EpisodeEntity {
surface: String,
#[serde(rename = "type")]
ty: String,
}
#[derive(Debug, Deserialize)]
struct EpisodeStatement {
predicate: String,
subject_surface: String,
object_surface: Option<String>,
object_literal_type: Option<String>,
object_literal_value: Option<String>,
valid_from: String,
#[serde(default)]
valid_to: Option<String>,
}
#[derive(Debug, Deserialize)]
struct QuestionFile {
#[allow(dead_code)]
id: String,
#[allow(dead_code)]
category: String,
question: String,
#[serde(default)]
as_of: Option<String>,
answer: String,
#[allow(dead_code)]
supporting_episodes: Vec<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Arm {
B,
C,
}
impl Arm {
fn label(self) -> &'static str {
match self {
Arm::B => "b (lexical)",
Arm::C => "c (hybrid)",
}
}
fn query_mode(self) -> QueryMode {
match self {
Arm::B => QueryMode::Lexical,
Arm::C => QueryMode::Hybrid,
}
}
}
#[derive(Debug, Clone)]
struct ArmResult {
passed: bool,
approx_tokens: usize,
}
pub async fn run(suite: &str, corpus_dir: &Path) -> Result<()> {
if suite != "gate" {
bail!("gate runner only handles the 'gate' suite (got '{suite}')");
}
let manifest_path = corpus_dir.join("manifest.toml");
let manifest_text = std::fs::read_to_string(&manifest_path)
.with_context(|| format!("read {}", manifest_path.display()))?;
let manifest: Manifest = toml::from_str(&manifest_text).context("parse manifest.toml")?;
let dir = TempDir::new().context("create temp dir")?;
let clock_ts = Timestamp::from_millis(1_700_000_000_000); let clock = Arc::new(FakeClock::new(clock_ts));
let brain = Brain::with_clock(oxibrain::BrainConfig::at(dir.path()), clock)
.await
.context("open brain")?;
let space = "gate";
let space_id = brain.ensure_space(space).await.context("ensure space")?;
let mut ingested = 0usize;
let mut entity_surfaces: Vec<String> = Vec::new();
for ep_entry in &manifest.episodes {
let ep_path = corpus_dir.join(&ep_entry.file);
let ep_text = std::fs::read_to_string(&ep_path)
.with_context(|| format!("read {}", ep_path.display()))?;
let ep: EpisodeFile =
toml::from_str(&ep_text).with_context(|| format!("parse {}", ep_path.display()))?;
for e in &ep.entities {
entity_surfaces.push(e.surface.clone());
}
for st in &ep.statements {
let decl = statement_to_declaration(st, &ep.entities)?;
brain
.declare(&space_id, decl)
.await
.with_context(|| format!("declare from {}", ep_entry.id))?;
ingested += 1;
}
}
entity_surfaces.sort();
entity_surfaces.dedup();
brain
.rebuild_indexes(&space_id)
.await
.context("rebuild indexes")?;
let mut results: Vec<QuestionResult> = Vec::new();
for q_entry in &manifest.questions {
let q_path = corpus_dir.join(&q_entry.file);
let q_text = std::fs::read_to_string(&q_path)
.with_context(|| format!("read {}", q_path.display()))?;
let q: QuestionFile =
toml::from_str(&q_text).with_context(|| format!("parse {}", q_path.display()))?;
let keyword = extract_keywords(&q.question, &entity_surfaces);
let as_of = q.as_of.as_deref().map(parse_iso_to_timestamp).transpose()?;
let arm_b = run_arm(&brain, &space_id, &keyword, Arm::B, as_of, &q.answer).await?;
let arm_c = run_arm(&brain, &space_id, &keyword, Arm::C, as_of, &q.answer).await?;
let recall = brain
.assemble_context(&space_id, &keyword, 3000)
.await
.context("assemble_context")?;
let recall_tokens = recall.total_tokens;
let brief_tokens = brief_token_cost(&brain, &space_id, &keyword, &entity_surfaces).await?;
results.push(QuestionResult {
id: q.id,
category: q.category,
answer: q.answer,
arm_b,
arm_c,
recall_tokens,
brief_tokens,
});
}
print_report(&results, ingested);
Ok(())
}
#[derive(Debug, Clone)]
struct QuestionResult {
id: String,
category: String,
answer: String,
arm_b: ArmResult,
arm_c: ArmResult,
recall_tokens: usize,
brief_tokens: usize,
}
fn statement_to_declaration(
st: &EpisodeStatement,
entities: &[EpisodeEntity],
) -> Result<Declaration> {
let valid_from = parse_iso_to_timestamp(&st.valid_from)?.millis();
let valid_to = match &st.valid_to {
Some(s) => parse_iso_to_timestamp(s)?.millis(),
None => oxibrain_ports::TIME_MAX.millis(),
};
let subject = EntityRef {
surface: st.subject_surface.clone(),
ty: entity_type(entities, &st.subject_surface)
.unwrap_or("Person")
.to_string(),
};
let object = if let Some(surface) = &st.object_surface {
let ty = entity_type(entities, surface).ok_or_else(|| {
anyhow::anyhow!("object surface '{surface}' not declared in episode entities")
})?;
DeclObject::Entity {
surface: surface.clone(),
ty: ty.to_string(),
}
} else if let (Some(lt), Some(val)) = (&st.object_literal_type, &st.object_literal_value) {
DeclObject::Literal {
literal_type: lt.clone(),
value: val.clone(),
}
} else {
bail!(
"statement '{}' has neither object_surface nor object_literal_*",
st.subject_surface
);
};
Ok(Declaration::AddStatement {
subject,
predicate: st.predicate.clone(),
object,
polarity: "affirm".into(),
valid_from,
valid_to,
})
}
fn extract_keywords(question: &str, entity_surfaces: &[String]) -> String {
let lower = question.to_lowercase();
let mut found: Vec<&str> = Vec::new();
for surface in entity_surfaces {
if lower.contains(&surface.to_lowercase()) {
found.push(surface.as_str());
}
}
if found.is_empty() {
question.to_string()
} else {
found.join(" ")
}
}
async fn brief_token_cost(
brain: &Brain,
space_id: &str,
keyword: &str,
entity_surfaces: &[String],
) -> Result<usize> {
let mut total = 0usize;
let mut seen = std::collections::HashSet::new();
for surface in entity_surfaces {
if !keyword.to_lowercase().contains(&surface.to_lowercase()) {
continue;
}
if !seen.insert(surface.clone()) {
continue;
}
if let Ok(Some(id)) = brain
.resolve_entity_id(
space_id,
&entity_type_for(brain, space_id, surface).await,
surface,
)
.await
{
let page = brain
.brief(space_id, &id)
.await
.context("brief cost page")?;
total += page.len() / 4;
}
}
Ok(total)
}
async fn entity_type_for(brain: &Brain, space_id: &str, surface: &str) -> String {
for ty in ["Person", "Organization", "Project", "Place", "Concept"] {
if brain
.resolve_entity_id(space_id, ty, surface)
.await
.ok()
.flatten()
.is_some()
{
return ty.to_string();
}
}
"Concept".to_string()
}
fn entity_type<'a>(entities: &'a [EpisodeEntity], surface: &str) -> Option<&'a str> {
entities
.iter()
.find(|e| e.surface == surface)
.map(|e| e.ty.as_str())
}
fn parse_iso_to_timestamp(s: &str) -> Result<Timestamp> {
let parts: Vec<&str> = s.split('-').collect();
if parts.len() != 3 {
bail!("not a YYYY-MM-DD date: {s}");
}
let y: i64 = parts[0]
.parse()
.with_context(|| format!("year parse: {s}"))?;
let m: i64 = parts[1]
.parse()
.with_context(|| format!("month parse: {s}"))?;
let d: i64 = parts[2]
.parse()
.with_context(|| format!("day parse: {s}"))?;
let days = days_from_civil(y, m, d);
Ok(Timestamp(days * 86_400_000 + 12 * 3_600_000)) }
fn days_from_civil(y: i64, m: i64, d: i64) -> i64 {
let y = if m <= 2 { y - 1 } else { y };
let era = if y >= 0 { y } else { y - 399 } / 400;
let yoe = y - era * 400; let m = if m > 2 { m - 3 } else { m + 9 }; let doy = (153 * m + 2) / 5 + d - 1; let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy; era * 146_097 + doe - 719_468
}
async fn run_arm(
brain: &Brain,
space_id: &str,
text: &str,
arm: Arm,
as_of: Option<Timestamp>,
answer: &str,
) -> Result<ArmResult> {
let query = Query {
text: text.to_string(),
mode: arm.query_mode(),
space: space_id.to_string(),
as_of,
limit: 15,
min_confidence: 0.0,
};
let result = brain.query(query).await.context(arm.label())?;
let mut ids: Vec<String> = Vec::new();
for item in result.items.iter().take(15) {
if let oxibrain_core::rank::TargetId::Statement { id } = &item.target {
ids.push(id.clone());
}
}
if arm == Arm::C {
let mut entity_ids: Vec<String> = Vec::new();
for item in result.items.iter().take(15) {
if let oxibrain_core::rank::TargetId::Entity { id } = &item.target {
entity_ids.push(id.clone());
}
}
if !entity_ids.is_empty() {
let entity_stmts = brain
.statements_for_entities(space_id, &entity_ids)
.await
.unwrap_or_default();
ids.extend(entity_stmts);
}
}
let rendered_lines = brain.render_statements(space_id, &ids).await?;
let rendered = rendered_lines.join("\n");
let passed = !rendered.is_empty()
&& answer_matches(answer.to_lowercase().as_str(), &rendered.to_lowercase());
let approx_tokens = rendered.len() / 4;
Ok(ArmResult {
passed,
approx_tokens,
})
}
fn answer_matches(needle: &str, haystack_lower: &str) -> bool {
let words: Vec<&str> = needle.split_whitespace().collect();
if words.is_empty() {
return false;
}
words.iter().all(|w| haystack_lower.contains(w))
}
fn print_report(results: &[QuestionResult], ingested: usize) {
println!("═══ oxibrain gate (golden-only) ═══");
println!("Episodes ingested: {ingested} declarations");
println!("Questions: {}", results.len());
println!();
for r in results {
let mark_b = if r.arm_b.passed { "✓" } else { "✗" };
let mark_c = if r.arm_c.passed { "✓" } else { "✗" };
println!(
" {} [{}] b={} ({:>3} tok) c={} ({:>3} tok) answer: {}",
r.id,
r.category,
mark_b,
r.arm_b.approx_tokens,
mark_c,
r.arm_c.approx_tokens,
truncate(&r.answer, 40),
);
}
println!();
let categories: std::collections::BTreeSet<&str> =
results.iter().map(|r| r.category.as_str()).collect();
println!("Per-category accuracy (c − b):");
for cat in categories {
let in_cat: Vec<&QuestionResult> = results.iter().filter(|r| r.category == cat).collect();
let n = in_cat.len() as f64;
let b = in_cat.iter().filter(|r| r.arm_b.passed).count() as f64;
let c = in_cat.iter().filter(|r| r.arm_c.passed).count() as f64;
println!(
" {:<22} b={:.0}/{:.0} c={:.0}/{:.0} delta(c−b) = {:+.0}",
cat,
b,
n,
c,
n,
(c - b),
);
}
println!();
let b_tokens: usize = results.iter().map(|r| r.arm_b.approx_tokens).sum();
let c_tokens: usize = results.iter().map(|r| r.arm_c.approx_tokens).sum();
let n = results.len().max(1);
println!(
"Tokens/query (approx): b={:.0} c={:.0} delta(c−b) = {:+} tok/q",
b_tokens as f64 / n as f64,
c_tokens as f64 / n as f64,
(c_tokens as isize - b_tokens as isize) / n as isize,
);
let recall_tokens: usize = results.iter().map(|r| r.recall_tokens).sum();
let brief_tokens: usize = results.iter().map(|r| r.brief_tokens).sum();
println!(
"Tokens/answer: M8 recall-only={:.0} M9 brief/navigate={:.0} delta = {:+} tok/q",
recall_tokens as f64 / n as f64,
brief_tokens as f64 / n as f64,
(brief_tokens as isize - recall_tokens as isize) / n as isize,
);
}
fn truncate(s: &str, n: usize) -> String {
if s.chars().count() <= n {
s.to_string()
} else {
let mut out: String = s.chars().take(n).collect();
out.push('…');
out
}
}
pub async fn run_with_dir(suite: &str, corpus_dir: Option<PathBuf>) -> Result<()> {
let dir = match corpus_dir {
Some(d) => d,
None => default_corpus_dir()?,
};
run(suite, &dir).await
}
fn default_corpus_dir() -> Result<PathBuf> {
let mut here = std::env::current_dir().context("cwd")?;
loop {
let candidate = here.join("eval").join("golden");
if candidate.join("manifest.toml").is_file() {
return Ok(candidate);
}
if !here.pop() {
bail!("could not locate eval/golden/manifest.toml from CWD");
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn days_from_civil_known_dates() {
assert_eq!(days_from_civil(1970, 1, 1), 0);
assert_eq!(days_from_civil(2023, 11, 14), 19675);
assert_eq!(days_from_civil(2023, 6, 1), 19509);
}
#[test]
fn answer_matches_substring() {
assert!(answer_matches("acme corp", "works for acme corp in 2023"));
assert!(answer_matches(
"alice smith",
"alice smith works on project x"
));
assert!(!answer_matches("missing", "alice works on project x"));
}
}