tidaldb/tidal-server/tests/vector_search.rs

// Integration-test exemption (same posture as the other tidal-server tests).
#![allow(clippy::unwrap_used, clippy::cast_possible_truncation)]
//! End-to-end coverage for the m12p1 `POST /vector_search` recall probe.
//!
//! Drives the real standalone handler in-process via `tower::ServiceExt::oneshot`
//! (no TCP bind): seed items + embeddings, then POST a query vector and assert
//! the raw k-NN result — closest-first ordering, `k` honored — plus the boundary
//! 400s (empty vector, dimension mismatch). This pins the surface the
//! `tidal-stress --verify-recall` harness measures against.

use std::sync::Arc;

use axum::{
    body::Body,
    http::{Method, Request, StatusCode},
};
use tidal_server::{router::build_router, state::ServerState};
use tidaldb::TidalDb;
use tower::ServiceExt;

/// Dimensionality of the default schema's `content_vector` slot.
const DIM: usize = 128;

fn make_app() -> axum::Router {
    let (schema, profiles) = tidal_server::config::load_schema(None).unwrap();
    let db = TidalDb::builder()
        .ephemeral()
        .with_schema(schema)
        .with_profiles(profiles)
        .open()
        .unwrap();
    let state = Arc::new(ServerState::new(db));
    build_router(
        state,
        Arc::new(tidal_server::cluster::security::ClusterCreds::unauthenticated()),
    )
}

/// A 128-dim one-hot-ish vector: component `axis` set to `mag`, rest 0 — except
/// we nudge a second axis a hair so no vector is exactly zero-norm.
fn axis_vector(axis: usize, mag: f32) -> Vec<f32> {
    let mut v = vec![0.0_f32; DIM];
    v[axis] = mag;
    v[(axis + 1) % DIM] = 0.01;
    v
}

async fn post_json(app: &axum::Router, uri: &str, body: serde_json::Value) -> StatusCode {
    app.clone()
        .oneshot(
            Request::builder()
                .method(Method::POST)
                .uri(uri)
                .header("Content-Type", "application/json")
                .body(Body::from(serde_json::to_vec(&body).unwrap()))
                .unwrap(),
        )
        .await
        .unwrap()
        .status()
}

async fn post_json_full(
    app: &axum::Router,
    uri: &str,
    body: serde_json::Value,
) -> (StatusCode, serde_json::Value) {
    let response = app
        .clone()
        .oneshot(
            Request::builder()
                .method(Method::POST)
                .uri(uri)
                .header("Content-Type", "application/json")
                .body(Body::from(serde_json::to_vec(&body).unwrap()))
                .unwrap(),
        )
        .await
        .unwrap();
    let status = response.status();
    let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
        .await
        .unwrap();
    let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap_or(serde_json::Value::Null);
    (status, json)
}

async fn seed(app: &axum::Router) {
    // Items on distinct axes: item 1 ~ axis 0, item 3 ~ axis 0 (close to item 1),
    // item 2 ~ axis 64 (far). A query on axis 0 must rank 1 and 3 above 2.
    for (id, v) in [
        (1u64, axis_vector(0, 1.0)),
        (2u64, axis_vector(64, 1.0)),
        (3u64, axis_vector(0, 0.8)),
    ] {
        let s = post_json(
            app,
            "/items",
            serde_json::json!({ "entity_id": id, "metadata": {} }),
        )
        .await;
        assert_eq!(s, StatusCode::CREATED, "item {id}");
        let s = post_json(
            app,
            "/embeddings",
            serde_json::json!({ "entity_id": id, "values": v }),
        )
        .await;
        assert_eq!(s, StatusCode::NO_CONTENT, "embedding {id}");
    }
}

#[tokio::test]
async fn vector_search_returns_nearest_closest_first() {
    let app = make_app();
    seed(&app).await;

    let (status, body) = post_json_full(
        &app,
        "/vector_search",
        serde_json::json!({ "vector": axis_vector(0, 1.0), "k": 3 }),
    )
    .await;
    assert_eq!(status, StatusCode::OK, "body: {body}");

    let items = body["items"].as_array().expect("items array");
    assert_eq!(items.len(), 3, "k=3 nearest");

    // Closest-first: an axis-0 query ranks the two axis-0 items (1, 3) above the
    // far axis-64 item (2).
    let ids: Vec<u64> = items
        .iter()
        .map(|it| it["entity_id"].as_u64().unwrap())
        .collect();
    assert_eq!(ids[0], 1, "the exact-axis item is nearest");
    assert!(
        ids[..2].contains(&3),
        "the near-axis item ranks above the far one; got {ids:?}"
    );
    assert_eq!(ids[2], 2, "the far axis-64 item is last");

    // Distances are present and ascending.
    let dists: Vec<f64> = items
        .iter()
        .map(|it| it["distance"].as_f64().unwrap())
        .collect();
    for w in dists.windows(2) {
        assert!(w[0] <= w[1], "distances must ascend: {dists:?}");
    }
}

#[tokio::test]
async fn vector_search_k_defaults_to_ten_and_clamps_to_corpus() {
    let app = make_app();
    seed(&app).await;
    // k omitted → defaults to 10, but only 3 items exist, so 3 come back.
    let (status, body) = post_json_full(
        &app,
        "/vector_search",
        serde_json::json!({ "vector": axis_vector(0, 1.0) }),
    )
    .await;
    assert_eq!(status, StatusCode::OK);
    assert_eq!(body["items"].as_array().unwrap().len(), 3);
}

#[tokio::test]
async fn vector_search_empty_vector_is_400() {
    let app = make_app();
    seed(&app).await;
    let v: Vec<f32> = vec![];
    let status = post_json(&app, "/vector_search", serde_json::json!({ "vector": v })).await;
    assert_eq!(status, StatusCode::BAD_REQUEST);
}

#[tokio::test]
async fn vector_search_dimension_mismatch_is_400() {
    let app = make_app();
    seed(&app).await;
    // 4-dim query against a 128-dim slot → a client error, not a 500.
    let status = post_json(
        &app,
        "/vector_search",
        serde_json::json!({ "vector": vec![0.1f32, 0.2, 0.3, 0.4] }),
    )
    .await;
    assert_eq!(status, StatusCode::BAD_REQUEST);
}