little-librarian/src/query.rs

//! Query processing and document retrieval.

use snafu::{ResultExt, Snafu};

use crate::{
    storage::{
        self,
        queries::{DocumentMatch, Queries},
    },
    text_encoder::{self, TextEncoder},
    tokenize::{self, Tokenizer},
};

/// Errors that occur during query processing.
#[derive(Debug, Snafu)]
pub enum AskError {
    #[snafu(display("Failed to encode query."))]
    Encode { source: tokenize::EncodeError },
    #[snafu(display("Failed to embed query."))]
    Embed { source: text_encoder::EmbedError },
    #[snafu(display("Failed to retrieve similar documents."))]
    Query {
        source: storage::queries::QueryError,
    },
    #[snafu(display("Failed to rerank documents."))]
    Rerank { source: text_encoder::RerankError },
}

/// Process a user query and return ranked document matches.
pub async fn ask(
    query: &str,
    db: &Queries,
    tokenizer: &Tokenizer,
    embedder: &TextEncoder,
    reranker: &TextEncoder,
    chunk_size: usize,
    limit: usize,
) -> Result<Vec<DocumentMatch>, AskError> {
    let encodings = tokenizer.encode(query, chunk_size).context(EncodeSnafu)?;
    let embeddings = embedder.embed(encodings[0].clone()).context(EmbedSnafu)?;
    let documents = db
        .query(embeddings, (limit * 10) as i32)
        .await
        .context(QuerySnafu)?;

    let reranked_docs = reranker
        .rerank(query, documents, tokenizer, limit)
        .context(RerankSnafu)?;

    Ok(reranked_docs)
}