Skip to main content

link_r/
query.rs

1// Author: support@lituus.io
2// Copyright: lituus-io, all rights reserved
3//
4//! Query inputs and search results.
5//!
6//! [`PreparedQuery`] is the embedder-agnostic search input the index consumes: a
7//! dense vector plus normalized BM25 terms plus a [`Filter`]. The text-level
8//! `Query` (which turns text into a vector via an [`Embedder`](crate::embed)) is
9//! assembled by the facade. [`Hit`]/[`Hits`] borrow the index, so results are
10//! allocation-free to read.
11
12use crate::index::fuse::DEFAULT_RRF_K;
13use crate::index::meta::{DocMeta, Filter};
14use crate::resource::{DocId, ResourceKind};
15use compact_str::CompactString;
16
17/// How dense and sparse result lists are combined.
18#[derive(Clone, Copy, Debug)]
19pub enum Fusion {
20    /// Reciprocal Rank Fusion with the given constant (rank-based; scale-free).
21    Rrf {
22        /// The RRF damping constant.
23        k: u32,
24    },
25}
26
27impl Default for Fusion {
28    fn default() -> Self {
29        Self::Rrf { k: DEFAULT_RRF_K }
30    }
31}
32
33/// Ranking parameters for a hybrid search.
34#[derive(Clone, Copy, Debug)]
35pub struct RankParams {
36    /// Weight of the dense (semantic) arm.
37    pub dense_weight: f32,
38    /// Weight of the sparse (BM25) arm.
39    pub sparse_weight: f32,
40    /// How the arms are fused.
41    pub fusion: Fusion,
42    /// Per-arm candidate pool size as a multiple of `limit` (oversampling before
43    /// fusion). Clamped to a sane floor internally.
44    pub candidate_multiplier: usize,
45    /// One-hop knowledge-graph boost applied after fusion (0 = off, the default):
46    /// each candidate gains `graph_boost × mean(fused score of its in/out neighbors
47    /// within the candidate pool)`, rewarding well-connected hub/authority pages.
48    pub graph_boost: f32,
49}
50
51impl Default for RankParams {
52    fn default() -> Self {
53        Self {
54            dense_weight: 1.0,
55            sparse_weight: 1.0,
56            fusion: Fusion::default(),
57            candidate_multiplier: 10,
58            graph_boost: 0.0,
59        }
60    }
61}
62
63/// An embedder-agnostic search input.
64#[derive(Debug)]
65pub struct PreparedQuery<'a> {
66    /// The query embedding (same dimension and metric as the index).
67    pub vector: &'a [f32],
68    /// Normalized query terms for BM25.
69    pub terms: &'a [CompactString],
70    /// Structured prefilter.
71    pub filter: &'a Filter,
72    /// Maximum hits to return.
73    pub limit: usize,
74    /// Ranking parameters.
75    pub rank: RankParams,
76}
77
78/// A single search result, borrowing the index it came from.
79#[derive(Clone, Copy, Debug)]
80pub struct Hit<'a> {
81    /// The resolvable remote URL (the value of the index).
82    pub url: &'a str,
83    /// The fused relevance score (higher is better).
84    pub score: f32,
85    /// Display title, if any.
86    pub title: Option<&'a str>,
87    /// A short display snippet.
88    pub snippet: &'a str,
89    /// Content kind.
90    pub kind: ResourceKind,
91    doc: &'a DocMeta,
92}
93
94impl<'a> Hit<'a> {
95    /// The document's structured tags.
96    pub fn tags(&self) -> impl Iterator<Item = &'a str> {
97        self.doc.tags.iter().map(CompactString::as_str)
98    }
99}
100
101/// The ranked results of a search, borrowing the index.
102#[derive(Debug)]
103pub struct Hits<'a> {
104    docs: &'a [DocMeta],
105    ranked: Vec<(DocId, f32)>,
106}
107
108impl<'a> Hits<'a> {
109    /// Construct from the index's metadata and a ranked id list.
110    #[must_use]
111    pub(crate) fn new(docs: &'a [DocMeta], ranked: Vec<(DocId, f32)>) -> Self {
112        Self { docs, ranked }
113    }
114
115    /// Number of hits.
116    #[must_use]
117    pub fn len(&self) -> usize {
118        self.ranked.len()
119    }
120
121    /// Whether there are no hits.
122    #[must_use]
123    pub fn is_empty(&self) -> bool {
124        self.ranked.is_empty()
125    }
126
127    /// Iterate hits in descending relevance.
128    pub fn iter(&self) -> impl Iterator<Item = Hit<'a>> + '_ {
129        let docs = self.docs;
130        self.ranked.iter().map(move |&(id, score)| {
131            let doc = &docs[id as usize];
132            Hit {
133                url: &doc.url,
134                score,
135                title: doc.title.as_deref(),
136                snippet: &doc.snippet,
137                kind: doc.kind,
138                doc,
139            }
140        })
141    }
142}
143
144impl<'a> IntoIterator for &Hits<'a> {
145    type Item = Hit<'a>;
146    type IntoIter = std::vec::IntoIter<Hit<'a>>;
147
148    fn into_iter(self) -> Self::IntoIter {
149        self.iter().collect::<Vec<_>>().into_iter()
150    }
151}