palace-rs 0.20.0

//! Semantic search over the palace.
//!
//! Embeds a query, does cosine similarity over all drawer embeddings,
//! returns verbatim text with similarity scores. Port of searcher.py.

use crate::embedder::embed_query;
use crate::query_intent::classify;
use crate::ranker::{hybrid_search, HybridResult};
use crate::store::{preference_search_filtered, source_context, DrawerFilter, SearchResult};
use anyhow::Result;
use rusqlite::Connection;

/// Print search results to stdout in the same format as the Python version.
pub fn search_and_print(
    conn: &Connection,
    query: &str,
    wing: Option<&str>,
    room: Option<&str>,
    n_results: usize,
    rerank: bool,
) -> Result<()> {
    let rewritten = crate::query_rewriter::rewrite(query);
    let sanitized_query = crate::query_sanitizer::sanitize_query(&rewritten);
    let effective_query = if sanitized_query.is_empty() {
        query
    } else {
        &sanitized_query
    };
    let intent = classify(effective_query);
    let embedding = embed_query(effective_query)?;
    let filter = DrawerFilter {
        wing: wing.map(String::from),
        room: room.map(String::from),
    };
    let mut results = hybrid_search(conn, effective_query, Some(&embedding), &filter, n_results)?;
    if crate::reranker::should_rerank(rerank) {
        crate::reranker::rerank(effective_query, &mut results);
    }

    println!("\n{}", "=".repeat(60));
    println!("  Results for: \"{query}\"");
    println!("  Intent: {}", intent.as_str());
    if let Some(w) = wing {
        println!("  Wing: {w}");
    }
    if let Some(r) = room {
        println!("  Room: {r}");
    }
    println!("{}\n", "=".repeat(60));

    if results.is_empty() {
        println!("  No results found for: \"{query}\"");
        return Ok(());
    }

    for (i, result) in results.iter().enumerate() {
        println!(
            "  [{}] {} / {}",
            i + 1,
            result.drawer.wing,
            result.drawer.room
        );
        println!("      Source: {}", result.drawer.source_file);
        println!("      Match:  {:.3}", result.combined);
        println!("      cosine={:.3} bm25={:.3}", result.cosine, result.bm25);
        if let Some(score) = result.rerank_score {
            println!("      rerank={score:.3}");
        }
        println!();
        for line in result.drawer.text.trim().lines() {
            println!("      {line}");
        }
        println!();
        println!("  {}", "-".repeat(56));
    }
    println!();
    Ok(())
}

/// Programmatic search returning structured results (used by MCP).
pub fn search_memories(
    conn: &Connection,
    query: &str,
    wing: Option<&str>,
    room: Option<&str>,
    n_results: usize,
) -> serde_json::Value {
    search_memories_with(conn, query, wing, room, n_results, SearchOptions::default())
}

/// Programmatic search with optional local reranking.
pub fn search_memories_with_options(
    conn: &Connection,
    query: &str,
    wing: Option<&str>,
    room: Option<&str>,
    n_results: usize,
    rerank: bool,
) -> serde_json::Value {
    let options = SearchOptions {
        rerank,
        ..SearchOptions::default()
    };
    search_memories_with(conn, query, wing, room, n_results, options)
}

/// What a programmatic search does beyond ranking.
#[derive(Debug, Clone, Copy)]
pub struct SearchOptions {
    /// Rerank the top hybrid candidates with the local interaction reranker.
    pub rerank: bool,
    /// Attach each hit's neighbouring chunks from the same source.
    pub context: bool,
}

impl Default for SearchOptions {
    fn default() -> Self {
        Self {
            rerank: false,
            context: true,
        }
    }
}

/// Programmatic search shaped by `options`.
pub fn search_memories_with(
    conn: &Connection,
    query: &str,
    wing: Option<&str>,
    room: Option<&str>,
    n_results: usize,
    options: SearchOptions,
) -> serde_json::Value {
    let rewritten = crate::query_rewriter::rewrite(query);
    let sanitized_query = crate::query_sanitizer::sanitize_query(&rewritten);
    let effective_query = if sanitized_query.is_empty() {
        query
    } else {
        &sanitized_query
    };
    let intent = classify(effective_query);
    let embedding = match embed_query(effective_query) {
        Ok(e) => e,
        Err(e) => {
            return serde_json::json!({
                "error": format!("Embedding error: {e}"),
                "hint": "Run: palace mine <dir>"
            })
        }
    };

    let filter = DrawerFilter {
        wing: wing.map(String::from),
        room: room.map(String::from),
    };

    let mut results =
        match hybrid_search(conn, effective_query, Some(&embedding), &filter, n_results) {
            Ok(r) => r,
            Err(e) => {
                return serde_json::json!({
                    "error": format!("Search error: {e}"),
                })
            }
        };

    // Merge dedicated preference recall pass — surfaces preference drawers even
    // when BM25 has no keyword overlap (the known R@1 weakness).
    merge_preference_results(conn, &embedding, &filter, &mut results, n_results);
    let rerank_enabled = crate::reranker::should_rerank(options.rerank);
    if rerank_enabled {
        crate::reranker::rerank(effective_query, &mut results);
    }
    let query_id = query_id(
        effective_query,
        &results
            .iter()
            .map(|result| result.drawer.id.as_str())
            .collect::<Vec<_>>(),
    );

    let hit_ids: std::collections::HashSet<&str> = results
        .iter()
        .map(|result| result.drawer.id.as_str())
        .collect();
    let hits: Vec<serde_json::Value> = results
        .iter()
        .map(|r| {
            let mut hit = serde_json::json!({
                "id": r.drawer.id,
                "text": r.drawer.text,
                "wing": r.drawer.wing,
                "room": r.drawer.room,
                "source_file": r.drawer.source_file,
                "created_at": r.drawer.created_at,
                "filed_at": r.drawer.filed_at,
                "similarity": r.combined,
                "combined": r.combined,
                "cosine": r.cosine,
                "bm25": r.bm25,
                "coding_boost": r.coding_boost,
                "preference_match": r.preference_match,
                "date_boost": r.date_boost,
                "source_boost": r.source_boost,
                "mined_discount": r.mined_discount,
                "rerank_score": r.rerank_score,
                "intent": intent.as_str(),
            });
            if options.context {
                hit["source_context"] = neighbouring_chunks(conn, &r.drawer.id, &hit_ids).into();
            }
            hit
        })
        .collect();

    serde_json::json!({
        "query": query,
        "query_id": query_id,
        "sanitized_query": effective_query,
        "intent": intent.as_str(),
        "rerank_enabled": rerank_enabled,
        "filters": {
            "wing": wing,
            "room": room,
        },
        "results": hits,
    })
}

/// The chunks right before and after a hit in its source, as `id`,
/// `chunk_index` and `text`. The hit itself and chunks that are hits of the
/// same search are left out: the response already carries them.
fn neighbouring_chunks(
    conn: &Connection,
    drawer_id: &str,
    hit_ids: &std::collections::HashSet<&str>,
) -> Vec<serde_json::Value> {
    let Some(drawer) = crate::store::get_drawer(conn, drawer_id).ok().flatten() else {
        return Vec::new();
    };
    if drawer.source_file.is_empty() {
        return Vec::new();
    }
    source_context(
        conn,
        &drawer.wing,
        &drawer.room,
        &drawer.source_file,
        drawer.chunk_index,
        1,
    )
    .unwrap_or_default()
    .into_iter()
    .filter(|neighbour| !hit_ids.contains(neighbour.id.as_str()))
    .map(|neighbour| {
        serde_json::json!({
            "id": neighbour.id,
            "chunk_index": neighbour.chunk_index,
            "text": neighbour.content,
        })
    })
    .collect()
}

/// Merge preference-tagged drawers into hybrid results, deduplicating by ID.
///
/// Preference drawers are weighted at 0.4 of their cosine similarity so they
/// supplement (rather than displace) strong non-preference matches. Only the top
/// half of preference candidates (cosine ≥ 0.25) are considered.
fn merge_preference_results(
    conn: &Connection,
    query_vec: &[f32],
    filter: &DrawerFilter,
    results: &mut Vec<HybridResult>,
    n_results: usize,
) {
    let pref_candidates = match preference_search_filtered(conn, query_vec, filter, 10) {
        Ok(v) => v,
        Err(_) => return,
    };

    let existing_ids: std::collections::HashSet<String> =
        results.iter().map(|r| r.drawer.id.clone()).collect();

    for pref in pref_candidates {
        if pref.similarity < 0.25 {
            break; // already sorted descending — skip low-relevance preferences
        }
        if existing_ids.contains(&pref.id) {
            continue;
        }
        let combined = pref.similarity * 0.4;
        results.push(HybridResult {
            cosine: pref.similarity,
            bm25: 0.0,
            coding_boost: 0.0,
            preference_match: pref.similarity,
            date_boost: 0.0,
            source_boost: 0.0,
            mined_discount: 0.0,
            rerank_score: None,
            combined,
            drawer: SearchResult {
                similarity: (combined * 1000.0).round() / 1000.0,
                ..pref
            },
        });
    }

    results.sort_by(|a, b| {
        b.combined
            .partial_cmp(&a.combined)
            .unwrap_or(std::cmp::Ordering::Equal)
    });
    results.truncate(n_results);
}

fn query_id(query: &str, top_drawer_ids: &[&str]) -> String {
    let bucket = chrono::Utc::now().timestamp() / 300;
    let joined_ids = top_drawer_ids.join(",");
    let hash = blake3::hash(format!("{query}|{joined_ids}|{bucket}").as_bytes());
    format!("query_{}", &hash.to_hex()[..16])
}