swiftide_core/
query.rs

1//! A query is the main object going through a query pipeline
2//!
3//! It acts as a statemachine, with the following transitions:
4//!
5//! `states::Pending`: No documents have been retrieved
6//! `states::Retrieved`: Documents have been retrieved
7//! `states::Answered`: The query has been answered
8use derive_builder::Builder;
9
10use crate::{Embedding, SparseEmbedding, document::Document, util::debug_long_utf8};
11
12/// A query is the main object going through a query pipeline
13///
14/// It acts as a statemachine, with the following transitions:
15///
16/// `states::Pending`: No documents have been retrieved
17/// `states::Retrieved`: Documents have been retrieved
18/// `states::Answered`: The query has been answered
19#[derive(Clone, Default, Builder, PartialEq)]
20#[builder(setter(into))]
21pub struct Query<STATE: QueryState> {
22    original: String,
23    #[builder(default = "self.original.clone().unwrap_or_default()")]
24    current: String,
25    #[builder(default = STATE::default())]
26    state: STATE,
27    #[builder(default)]
28    transformation_history: Vec<TransformationEvent>,
29
30    // TODO: How would this work when doing a rollup query?
31    #[builder(default)]
32    pub embedding: Option<Embedding>,
33
34    #[builder(default)]
35    pub sparse_embedding: Option<SparseEmbedding>,
36
37    /// Documents the query will operate on
38    ///
39    /// A query can retrieve multiple times, accumulating documents
40    #[builder(default)]
41    pub documents: Vec<Document>,
42}
43
44impl<STATE: std::fmt::Debug + QueryState> std::fmt::Debug for Query<STATE> {
45    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46        f.debug_struct("Query")
47            .field(
48                "original",
49                &debug_long_utf8(&self.original, 100).lines().take(1),
50            )
51            .field(
52                "current",
53                &debug_long_utf8(&self.current, 100).lines().take(1),
54            )
55            .field("state", &self.state)
56            .field("num_transformations", &self.transformation_history.len())
57            .field("embedding", &self.embedding.is_some())
58            .field("num_documents", &self.documents.len())
59            .finish()
60    }
61}
62
63impl<STATE: Clone + QueryState> Query<STATE> {
64    pub fn builder() -> QueryBuilder<STATE> {
65        QueryBuilder::default().clone()
66    }
67
68    /// Return the query it started with
69    pub fn original(&self) -> &str {
70        &self.original
71    }
72
73    /// Return the current query (or after retrieval!)
74    pub fn current(&self) -> &str {
75        &self.current
76    }
77
78    fn transition_to<NEWSTATE: QueryState>(self, new_state: NEWSTATE) -> Query<NEWSTATE> {
79        Query {
80            state: new_state,
81            original: self.original,
82            current: self.current,
83            transformation_history: self.transformation_history,
84            embedding: self.embedding,
85            sparse_embedding: self.sparse_embedding,
86            documents: self.documents,
87        }
88    }
89
90    #[allow(dead_code)]
91    pub fn history(&self) -> &Vec<TransformationEvent> {
92        &self.transformation_history
93    }
94
95    /// Returns the current documents that will be used as context for answer generation
96    pub fn documents(&self) -> &[Document] {
97        &self.documents
98    }
99
100    /// Returns the current documents as mutable
101    pub fn documents_mut(&mut self) -> &mut Vec<Document> {
102        &mut self.documents
103    }
104}
105
106impl<STATE: Clone + CanRetrieve> Query<STATE> {
107    /// Add retrieved documents and transition to `states::Retrieved`
108    pub fn retrieved_documents(mut self, documents: Vec<Document>) -> Query<states::Retrieved> {
109        self.documents.extend(documents.clone());
110        self.transformation_history
111            .push(TransformationEvent::Retrieved {
112                before: self.current.clone(),
113                after: self.current.clone(),
114                documents,
115            });
116
117        let state = states::Retrieved;
118
119        self.transition_to(state)
120    }
121}
122
123impl Query<states::Pending> {
124    pub fn new(query: impl Into<String>) -> Self {
125        Self {
126            original: query.into(),
127            ..Default::default()
128        }
129    }
130
131    /// Transforms the current query
132    pub fn transformed_query(&mut self, new_query: impl Into<String>) {
133        let new_query = new_query.into();
134
135        self.transformation_history
136            .push(TransformationEvent::Transformed {
137                before: self.current.clone(),
138                after: new_query.clone(),
139            });
140
141        self.current = new_query;
142    }
143}
144
145impl Query<states::Retrieved> {
146    pub fn new() -> Self {
147        Self::default()
148    }
149
150    /// Transforms the current response
151    pub fn transformed_response(&mut self, new_response: impl Into<String>) {
152        let new_response = new_response.into();
153
154        self.transformation_history
155            .push(TransformationEvent::Transformed {
156                before: self.current.clone(),
157                after: new_response.clone(),
158            });
159
160        self.current = new_response;
161    }
162
163    /// Transition the query to `states::Answered`
164    #[must_use]
165    pub fn answered(mut self, answer: impl Into<String>) -> Query<states::Answered> {
166        self.current = answer.into();
167        let state = states::Answered;
168        self.transition_to(state)
169    }
170}
171
172impl Query<states::Answered> {
173    pub fn new() -> Self {
174        Self::default()
175    }
176
177    /// Returns the answer of the query
178    pub fn answer(&self) -> &str {
179        &self.current
180    }
181}
182
183/// Marker trait for query states
184pub trait QueryState: Send + Sync + Default {}
185/// Marker trait for query states that can still retrieve
186pub trait CanRetrieve: QueryState {}
187
188/// States of a query
189pub mod states {
190    use super::{CanRetrieve, QueryState};
191
192    #[derive(Debug, Default, Clone, PartialEq)]
193    /// The query is pending and has not been used
194    pub struct Pending;
195
196    #[derive(Debug, Default, Clone, PartialEq)]
197    /// Documents have been retrieved
198    pub struct Retrieved;
199
200    #[derive(Debug, Default, Clone, PartialEq)]
201    /// The query has been answered
202    pub struct Answered;
203
204    impl QueryState for Pending {}
205    impl QueryState for Retrieved {}
206    impl QueryState for Answered {}
207
208    impl CanRetrieve for Pending {}
209    impl CanRetrieve for Retrieved {}
210}
211
212impl<T: AsRef<str>> From<T> for Query<states::Pending> {
213    fn from(original: T) -> Self {
214        Self {
215            original: original.as_ref().to_string(),
216            current: original.as_ref().to_string(),
217            state: states::Pending,
218            ..Default::default()
219        }
220    }
221}
222
223#[derive(Clone, PartialEq)]
224/// Records changes to a query
225pub enum TransformationEvent {
226    Transformed {
227        before: String,
228        after: String,
229    },
230    Retrieved {
231        before: String,
232        after: String,
233        documents: Vec<Document>,
234    },
235}
236
237impl TransformationEvent {
238    /// Returns true if the event is a retrieval
239    pub fn is_retrieval(&self) -> bool {
240        matches!(self, TransformationEvent::Retrieved { .. })
241    }
242
243    /// Returns true if the event is a transformation
244    pub fn is_transformation(&self) -> bool {
245        matches!(self, TransformationEvent::Transformed { .. })
246    }
247
248    /// Returns the query before the transformation/retrieval
249    pub fn before(&self) -> &str {
250        match self {
251            TransformationEvent::Transformed { before, .. }
252            | TransformationEvent::Retrieved { before, .. } => before,
253        }
254    }
255
256    /// Returns the query after the transformation/retrieval
257    pub fn after(&self) -> &str {
258        match self {
259            TransformationEvent::Transformed { after, .. }
260            | TransformationEvent::Retrieved { after, .. } => after,
261        }
262    }
263
264    /// Returns the documents retrieved, if any
265    pub fn documents(&self) -> Option<&[Document]> {
266        match self {
267            TransformationEvent::Retrieved { documents, .. } => Some(documents),
268            TransformationEvent::Transformed { .. } => None,
269        }
270    }
271}
272
273impl std::fmt::Debug for TransformationEvent {
274    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
275        match self {
276            TransformationEvent::Transformed { before, after } => {
277                write!(
278                    f,
279                    "Transformed: {} -> {}",
280                    debug_long_utf8(before, 100),
281                    debug_long_utf8(after, 100)
282                )
283            }
284            TransformationEvent::Retrieved {
285                before,
286                after,
287                documents,
288            } => {
289                write!(
290                    f,
291                    "Retrieved: {} -> {}\nDocuments: {:?}",
292                    debug_long_utf8(before, 100),
293                    debug_long_utf8(after, 100),
294                    documents.len()
295                )
296            }
297        }
298    }
299}
300
301#[cfg(test)]
302mod tests {
303    use super::*;
304
305    #[test]
306    fn test_query_initial_state() {
307        let query = Query::<states::Pending>::from("test query");
308        assert_eq!(query.original(), "test query");
309        assert_eq!(query.current(), "test query");
310        assert_eq!(query.history().len(), 0);
311    }
312
313    #[test]
314    fn test_query_transformed_query() {
315        let mut query = Query::<states::Pending>::from("test query");
316        query.transformed_query("new query");
317        assert_eq!(query.current(), "new query");
318        assert_eq!(query.history().len(), 1);
319        if let TransformationEvent::Transformed { before, after } = &query.history()[0] {
320            assert_eq!(before, "test query");
321            assert_eq!(after, "new query");
322        } else {
323            panic!("Unexpected event in history");
324        }
325    }
326
327    #[test]
328    fn test_query_retrieved_documents() {
329        let query = Query::<states::Pending>::from("test query");
330        let documents: Vec<Document> = vec!["doc1".into(), "doc2".into()];
331        let query = query.retrieved_documents(documents.clone());
332        assert_eq!(query.documents(), &documents);
333        assert_eq!(query.history().len(), 1);
334        if let TransformationEvent::Retrieved {
335            before,
336            after,
337            documents: retrieved_docs,
338        } = &query.history()[0]
339        {
340            assert_eq!(before, "test query");
341            assert_eq!(after, "test query");
342            assert_eq!(retrieved_docs, &documents);
343        } else {
344            panic!("Unexpected event in history");
345        }
346    }
347
348    #[test]
349    fn test_query_transformed_response() {
350        let query = Query::<states::Pending>::from("test query");
351        let documents = vec!["doc1".into(), "doc2".into()];
352        let mut query = query.retrieved_documents(documents.clone());
353        query.transformed_response("new response");
354
355        assert_eq!(query.current(), "new response");
356        assert_eq!(query.history().len(), 2);
357        assert_eq!(query.documents(), &documents);
358        assert_eq!(query.original, "test query");
359        if let TransformationEvent::Transformed { before, after } = &query.history()[1] {
360            assert_eq!(before, "test query");
361            assert_eq!(after, "new response");
362        } else {
363            panic!("Unexpected event in history");
364        }
365    }
366
367    #[test]
368    fn test_query_answered() {
369        let query = Query::<states::Pending>::from("test query");
370        let documents = vec!["doc1".into(), "doc2".into()];
371        let query = query.retrieved_documents(documents);
372        let query = query.answered("the answer");
373
374        assert_eq!(query.answer(), "the answer");
375    }
376}
swiftide_core/query.rs

swiftide_core/
query.rs
