69 lines
2 KiB
Rust
69 lines
2 KiB
Rust
//! Query processing and document retrieval.
|
|
|
|
use snafu::{ResultExt, Snafu};
|
|
use tokio::task::JoinError;
|
|
|
|
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("Embedding task failed to execute."))]
|
|
EmbedJoin { source: JoinError },
|
|
#[snafu(display("Failed to retrieve similar documents."))]
|
|
Query {
|
|
source: storage::queries::QueryError,
|
|
},
|
|
#[snafu(display("Failed to rerank documents."))]
|
|
Rerank { source: text_encoder::RerankError },
|
|
#[snafu(display("Reranking task failed to execute."))]
|
|
RerankJoin { source: JoinError },
|
|
}
|
|
|
|
/// 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 embedder = embedder.clone();
|
|
let encoding = encodings[0].clone();
|
|
let embeddings = tokio::task::spawn_blocking(move || embedder.embed(encoding))
|
|
.await
|
|
.context(EmbedJoinSnafu)?
|
|
.context(EmbedSnafu)?;
|
|
|
|
let documents = db
|
|
.query(embeddings, (limit * 10) as i32)
|
|
.await
|
|
.context(QuerySnafu)?;
|
|
|
|
let reranker = reranker.clone();
|
|
let tokenizer = tokenizer.clone();
|
|
let query = query.to_string();
|
|
let reranked_docs = tokio::task::spawn_blocking(move || {
|
|
reranker.rerank(&query, documents, &tokenizer, limit)
|
|
})
|
|
.await
|
|
.context(RerankJoinSnafu)?
|
|
.context(RerankSnafu)?;
|
|
|
|
Ok(reranked_docs)
|
|
}
|