tftio-kb 2.5.4

//! Port of Haskell `test/EmbeddingSpec.hs`.
//!
//! Stub `EmbeddingClient` implementations only — no network traffic.
//! Tests exercise the v9 design where embedding generation happens at
//! the handler layer (here: the test driver) and the precomputed vector
//! is passed to storage as `Option<Vec<f32>>`.

use std::sync::{
    Arc, Mutex,
    atomic::{AtomicUsize, Ordering},
};

use async_trait::async_trait;
use kb::embedding::{EmbeddingClient, EmbeddingError, decode_embedding, encode_embedding};
use kb::storage::{self, init_db, insert_node_with, open_db, search_hybrid, update_node_with};
use tftio_org::ast::*;

type TestResult = Result<(), Box<dyn std::error::Error>>;

// ── Stub clients (async) ──────────────────────────────────────────────

struct ConstClient(Vec<f32>);

#[async_trait]
impl EmbeddingClient for ConstClient {
    async fn embed(&self, _input: &str) -> Result<Vec<f32>, EmbeddingError> {
        Ok(self.0.clone())
    }
}

struct FailingClient;

#[async_trait]
impl EmbeddingClient for FailingClient {
    async fn embed(&self, _input: &str) -> Result<Vec<f32>, EmbeddingError> {
        Err(EmbeddingError::EmptyResponse)
    }
}

/// Per-node deterministic vector, so different nodes have different
/// distances and hybrid ranking is observable.
struct DeterministicClient(Mutex<Vec<(String, Vec<f32>)>>);

#[async_trait]
impl EmbeddingClient for DeterministicClient {
    async fn embed(&self, input: &str) -> Result<Vec<f32>, EmbeddingError> {
        // Scope the lock guard so it is released before the fallback return,
        // and recover it from poisoning rather than unwrapping.
        let matched = {
            let map = self
                .0
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            map.iter()
                .find_map(|(key, vec)| input.contains(key).then(|| vec.clone()))
        };
        Ok(matched.unwrap_or_else(|| vec![0.0_f32, 0.0, 0.0]))
    }
}

struct CallCounter {
    inner: ConstClient,
    count: Arc<AtomicUsize>,
}

#[async_trait]
impl EmbeddingClient for CallCounter {
    async fn embed(&self, input: &str) -> Result<Vec<f32>, EmbeddingError> {
        self.count.fetch_add(1, Ordering::SeqCst);
        self.inner.embed(input).await
    }
}

fn fresh() -> Result<(tempfile::TempDir, rusqlite::Connection), Box<dyn std::error::Error>> {
    let dir = tempfile::tempdir()?;
    let path = dir.path().join("emb.db");
    let path_str = path.to_str().ok_or("temp db path is not valid UTF-8")?;
    let conn = open_db(path_str)?;
    init_db(&conn)?;
    Ok((dir, conn))
}

fn simple_doc(title: &str, body: &str) -> Document {
    Document {
        blocks: vec![
            Block::Heading {
                level: 1,
                title: Title(title.into()),
                tags: vec![],
                children: vec![],
            },
            Block::Paragraph {
                inlines: vec![Inline::Plain(body.into())],
            },
        ],
    }
}

/// Driver helper: compute the embedding for a doc via the supplied
/// async client (the role the handler layer plays in production), then
/// pass the precomputed vector into storage. Mirrors what the axum and
/// MCP handlers do.
async fn compute_payload<C: EmbeddingClient>(client: &C, doc: &Document) -> Option<Vec<f32>> {
    let title = match doc.blocks.first() {
        Some(Block::Heading { title, .. }) => title.0.clone(),
        _ => "(untitled)".into(),
    };
    let body = match doc.blocks.get(1) {
        Some(Block::Paragraph { inlines }) => inlines
            .iter()
            .filter_map(|i| match i {
                Inline::Plain(s) => Some(s.as_str()),
                _ => None,
            })
            .collect::<Vec<_>>()
            .join(" "),
        _ => String::new(),
    };
    let payload = format!("{title}\n{body}");
    client.embed(&payload).await.ok()
}

// ── Codec ─────────────────────────────────────────────────────────────

#[test]
fn codec_round_trip() -> TestResult {
    let v: Vec<f32> = vec![0.0, 1.0, -1.0, 0.5, f32::MIN, f32::MAX];
    let bytes = encode_embedding(&v);
    assert_eq!(bytes.len(), 4 * v.len());
    assert_eq!(decode_embedding(&bytes)?, v);
    Ok(())
}

#[test]
fn codec_encode_of_empty_is_empty() -> TestResult {
    let v: Vec<f32> = vec![];
    assert!(encode_embedding(&v).is_empty());
    assert_eq!(decode_embedding(&encode_embedding(&v))?, v);
    Ok(())
}

#[test]
fn codec_decode_rejects_non_multiple_of_four() -> TestResult {
    let bad: Vec<u8> = vec![0; 7];
    let Err(err) = decode_embedding(&bad) else {
        return Err("decode_embedding should reject a non-multiple-of-4 length".into());
    };
    assert!(err.to_string().contains("multiple of 4"), "got: {err}");
    Ok(())
}

// ── Write-path: handler computes, storage writes ──────────────────────

fn count_embeddings(conn: &rusqlite::Connection) -> Result<i64, rusqlite::Error> {
    conn.query_row("SELECT count(*) FROM embeddings", [], |r| r.get(0))
}

fn embedding_for(conn: &rusqlite::Connection, node_id: &str) -> Option<Vec<u8>> {
    conn.query_row(
        "SELECT embedding FROM embeddings WHERE node_id = ?1",
        [node_id],
        |r| r.get::<_, Vec<u8>>(0),
    )
    .ok()
}

#[tokio::test]
async fn insert_node_writes_embedding_when_client_configured() -> TestResult {
    let (_d, c) = fresh()?;
    let client = ConstClient(vec![0.1, 0.2, 0.3]);
    let doc = simple_doc("T", "body");
    let v = compute_payload(&client, &doc).await;
    insert_node_with(&c, "n1", &doc, v, Some("stub"))?;
    assert_eq!(count_embeddings(&c)?, 1);
    let blob = embedding_for(&c, "n1").ok_or("no embedding row for n1")?;
    assert_eq!(decode_embedding(&blob)?, vec![0.1, 0.2, 0.3]);
    Ok(())
}

#[tokio::test]
async fn update_node_re_embeds_via_insert_or_replace() -> TestResult {
    let (_d, c) = fresh()?;
    let c1 = ConstClient(vec![0.5, 0.5]);
    let doc1 = simple_doc("T", "alpha");
    let v1 = compute_payload(&c1, &doc1).await;
    insert_node_with(&c, "n1", &doc1, v1, Some("stub"))?;
    assert_eq!(count_embeddings(&c)?, 1);

    let c2 = ConstClient(vec![0.9, 0.1]);
    let doc2 = simple_doc("T", "beta");
    let v2 = compute_payload(&c2, &doc2).await;
    update_node_with(&c, "n1", &doc2, v2, Some("stub"))?;
    assert_eq!(count_embeddings(&c)?, 1);
    let blob = embedding_for(&c, "n1").ok_or("no embedding row for n1")?;
    assert_eq!(decode_embedding(&blob)?, vec![0.9, 0.1]);
    Ok(())
}

#[tokio::test]
async fn client_failure_does_not_block_node_write() -> TestResult {
    let (_d, c) = fresh()?;
    let client = FailingClient;
    let doc = simple_doc("T", "body");
    // Handler-layer behaviour: failed embed -> proceed with None.
    let v = compute_payload(&client, &doc).await;
    assert!(v.is_none());
    insert_node_with(&c, "n1", &doc, v, Some("stub"))?;
    assert!(storage::get_node(&c, "n1")?.is_some());
    assert_eq!(count_embeddings(&c)?, 0);
    Ok(())
}

#[test]
fn no_client_means_no_embeddings_rows() -> TestResult {
    let (_d, c) = fresh()?;
    insert_node_with(&c, "n1", &simple_doc("T", "body"), None, None)?;
    assert_eq!(count_embeddings(&c)?, 0);
    Ok(())
}

// ── Hybrid search ─────────────────────────────────────────────────────

#[test]
fn search_hybrid_falls_back_to_fts_when_no_client() -> TestResult {
    let (_d, c) = fresh()?;
    insert_node_with(&c, "a", &simple_doc("Apple", "first"), None, None)?;
    insert_node_with(&c, "b", &simple_doc("Banana", "second"), None, None)?;
    let hits = search_hybrid(&c, "Apple", None)?;
    assert_eq!(hits, vec![NodeId("a".into())]);
    Ok(())
}

#[tokio::test]
async fn search_hybrid_blends_fts_and_vector_via_rrf() -> TestResult {
    let (_d, c) = fresh()?;

    let det = DeterministicClient(Mutex::new(vec![
        ("Apple".to_string(), vec![1.0_f32, 0.0, 0.0]),
        ("Banana".to_string(), vec![0.9_f32, 0.0, 0.0]),
        ("shared".to_string(), vec![1.0_f32, 0.0, 0.0]),
    ]));
    let doc_a = simple_doc("Apple", "shared body");
    let doc_b = simple_doc("Banana", "completely different");
    let va = compute_payload(&det, &doc_a).await;
    let vb = compute_payload(&det, &doc_b).await;
    insert_node_with(&c, "a", &doc_a, va, Some("stub"))?;
    insert_node_with(&c, "b", &doc_b, vb, Some("stub"))?;

    let q = "shared Banana";
    let qv = det.embed(q).await?;
    let hits = search_hybrid(&c, q, Some((qv, "stub")))?;
    let ids: Vec<String> = hits.into_iter().map(|n| n.0).collect();
    assert!(
        ids.contains(&"a".to_string()),
        "RRF result missing a: {ids:?}"
    );
    assert!(
        ids.contains(&"b".to_string()),
        "RRF result missing b: {ids:?}"
    );
    Ok(())
}

#[tokio::test]
async fn search_hybrid_handles_client_failure_gracefully() -> TestResult {
    let (_d, c) = fresh()?;
    let working = ConstClient(vec![1.0]);
    let doc = simple_doc("Apple", "first");
    let v = compute_payload(&working, &doc).await;
    insert_node_with(&c, "a", &doc, v, Some("stub"))?;
    // Handler-layer behaviour with a failing client: query embedding =
    // None; result equals the FTS ranking.
    let failing = FailingClient;
    let qe = failing.embed("Apple").await.ok().map(|v| (v, "stub"));
    assert!(qe.is_none());
    let hits = search_hybrid(&c, "Apple", qe)?;
    assert_eq!(hits, vec![NodeId("a".into())]);
    Ok(())
}

#[tokio::test]
async fn embed_called_once_on_insert_one_on_update() -> TestResult {
    let (_d, c) = fresh()?;
    let count = Arc::new(AtomicUsize::new(0));
    let client = CallCounter {
        inner: ConstClient(vec![1.0_f32, 0.0]),
        count: Arc::clone(&count),
    };
    let doc1 = simple_doc("T", "body");
    let v1 = compute_payload(&client, &doc1).await;
    insert_node_with(&c, "n1", &doc1, v1, Some("stub"))?;
    let doc2 = simple_doc("T", "body2");
    let v2 = compute_payload(&client, &doc2).await;
    update_node_with(&c, "n1", &doc2, v2, Some("stub"))?;
    assert_eq!(count.load(Ordering::SeqCst), 2);
    Ok(())
}