swiftide_core/
indexing_traits.rs

1//! Traits in Swiftide allow for easy extendability
2//!
3//! All steps defined in the indexing pipeline and the generic transformers can also take a
4//! trait. To bring your own transformers, models and loaders, all you need to do is implement the
5//! trait and it should work out of the box.
6use crate::Embeddings;
7use crate::node::{Chunk, Node};
8use crate::{
9    SparseEmbeddings, indexing_defaults::IndexingDefaults, indexing_stream::IndexingStream,
10};
11use std::fmt::Debug;
12use std::sync::Arc;
13
14use crate::chat_completion::errors::LanguageModelError;
15use crate::prompt::Prompt;
16use anyhow::Result;
17use async_trait::async_trait;
18
19pub use dyn_clone::DynClone;
20/// All traits are easily mockable under tests
21#[cfg(feature = "test-utils")]
22#[doc(hidden)]
23use mockall::{mock, predicate::str};
24use schemars::{JsonSchema, schema_for};
25use serde::de::DeserializeOwned;
26
27#[async_trait]
28/// Transforms single nodes into single nodes
29pub trait Transformer: Send + Sync + DynClone {
30    type Input: Chunk;
31    type Output: Chunk;
32
33    async fn transform_node(&self, node: Node<Self::Input>) -> Result<Node<Self::Output>>;
34
35    /// Overrides the default concurrency of the pipeline
36    fn concurrency(&self) -> Option<usize> {
37        None
38    }
39
40    fn name(&self) -> &'static str {
41        let name = std::any::type_name::<Self>();
42        name.split("::").last().unwrap_or(name)
43    }
44}
45
46dyn_clone::clone_trait_object!(<I, O> Transformer<Input = I, Output = O>);
47
48#[cfg(feature = "test-utils")]
49mock! {
50    #[derive(Debug)]
51    pub Transformer {}
52
53    #[async_trait]
54    impl Transformer for Transformer {
55        type Input = String;
56        type Output = String;
57
58        async fn transform_node(&self, node: Node<String>) -> Result<Node<String>>;
59        fn concurrency(&self) -> Option<usize>;
60        fn name(&self) -> &'static str;
61    }
62
63    impl Clone for Transformer {
64        fn clone(&self) -> Self;
65    }
66}
67
68#[async_trait]
69impl<I: Chunk, O: Chunk> Transformer for Box<dyn Transformer<Input = I, Output = O>> {
70    type Input = I;
71    type Output = O;
72
73    async fn transform_node(&self, node: Node<Self::Input>) -> Result<Node<Self::Output>> {
74        self.as_ref().transform_node(node).await
75    }
76    fn concurrency(&self) -> Option<usize> {
77        self.as_ref().concurrency()
78    }
79    fn name(&self) -> &'static str {
80        self.as_ref().name()
81    }
82}
83
84#[async_trait]
85impl<I: Chunk, O: Chunk> Transformer for Arc<dyn Transformer<Input = I, Output = O>> {
86    type Input = I;
87    type Output = O;
88
89    async fn transform_node(&self, node: Node<Self::Input>) -> Result<Node<Self::Output>> {
90        self.as_ref().transform_node(node).await
91    }
92    fn concurrency(&self) -> Option<usize> {
93        self.as_ref().concurrency()
94    }
95    fn name(&self) -> &'static str {
96        self.as_ref().name()
97    }
98}
99
100#[async_trait]
101impl<I: Chunk, O: Chunk> Transformer for &dyn Transformer<Input = I, Output = O> {
102    type Input = I;
103    type Output = O;
104
105    async fn transform_node(&self, node: Node<Self::Input>) -> Result<Node<Self::Output>> {
106        (*self).transform_node(node).await
107    }
108    fn concurrency(&self) -> Option<usize> {
109        (*self).concurrency()
110    }
111}
112
113#[async_trait]
114/// Use a closure as a transformer
115// TODO: Find a way to make this work with full generics
116impl<F> Transformer for F
117where
118    F: Fn(Node<String>) -> Result<Node<String>> + Send + Sync + Clone,
119{
120    type Input = String;
121    type Output = String;
122
123    async fn transform_node(&self, node: Node<Self::Input>) -> Result<Node<Self::Output>> {
124        self(node)
125    }
126}
127
128#[async_trait]
129/// Transforms batched single nodes into streams of nodes
130pub trait BatchableTransformer: Send + Sync + DynClone {
131    type Input: Chunk;
132    type Output: Chunk;
133
134    /// Transforms a batch of nodes into a stream of nodes
135    async fn batch_transform(&self, nodes: Vec<Node<Self::Input>>) -> IndexingStream<Self::Output>;
136
137    /// Overrides the default concurrency of the pipeline
138    fn concurrency(&self) -> Option<usize> {
139        None
140    }
141
142    fn name(&self) -> &'static str {
143        let name = std::any::type_name::<Self>();
144        name.split("::").last().unwrap_or(name)
145    }
146
147    /// Overrides the default batch size of the pipeline
148    fn batch_size(&self) -> Option<usize> {
149        None
150    }
151}
152
153dyn_clone::clone_trait_object!(<I, O> BatchableTransformer<Input = I, Output = O>);
154
155#[cfg(feature = "test-utils")]
156mock! {
157    #[derive(Debug)]
158    pub BatchableTransformer {}
159
160    #[async_trait]
161    impl BatchableTransformer for BatchableTransformer {
162        type Input = String;
163        type Output = String;
164
165        async fn batch_transform(&self, nodes: Vec<Node<String>>) -> IndexingStream<String>;
166        fn name(&self) -> &'static str;
167        fn batch_size(&self) -> Option<usize>;
168        fn concurrency(&self) -> Option<usize>;
169    }
170
171    impl Clone for BatchableTransformer {
172        fn clone(&self) -> Self;
173    }
174}
175#[async_trait]
176/// Use a closure as a batchable transformer
177impl<F> BatchableTransformer for F
178where
179    F: Fn(Vec<Node<String>>) -> IndexingStream<String> + Send + Sync + Clone,
180{
181    type Input = String;
182    type Output = String;
183
184    async fn batch_transform(&self, nodes: Vec<Node<String>>) -> IndexingStream<String> {
185        self(nodes)
186    }
187}
188
189#[async_trait]
190impl<I: Chunk, O: Chunk> BatchableTransformer
191    for Box<dyn BatchableTransformer<Input = I, Output = O>>
192{
193    type Input = I;
194    type Output = O;
195
196    async fn batch_transform(&self, nodes: Vec<Node<Self::Input>>) -> IndexingStream<Self::Output> {
197        self.as_ref().batch_transform(nodes).await
198    }
199    fn concurrency(&self) -> Option<usize> {
200        self.as_ref().concurrency()
201    }
202    fn name(&self) -> &'static str {
203        self.as_ref().name()
204    }
205}
206
207#[async_trait]
208impl<I: Chunk, O: Chunk> BatchableTransformer
209    for Arc<dyn BatchableTransformer<Input = I, Output = O>>
210{
211    type Input = I;
212    type Output = O;
213
214    async fn batch_transform(&self, nodes: Vec<Node<Self::Input>>) -> IndexingStream<Self::Output> {
215        self.as_ref().batch_transform(nodes).await
216    }
217    fn concurrency(&self) -> Option<usize> {
218        self.as_ref().concurrency()
219    }
220    fn name(&self) -> &'static str {
221        self.as_ref().name()
222    }
223}
224
225#[async_trait]
226impl<I: Chunk, O: Chunk> BatchableTransformer for &dyn BatchableTransformer<Input = I, Output = O> {
227    type Input = I;
228    type Output = O;
229
230    async fn batch_transform(&self, nodes: Vec<Node<Self::Input>>) -> IndexingStream<Self::Output> {
231        (*self).batch_transform(nodes).await
232    }
233    fn concurrency(&self) -> Option<usize> {
234        (*self).concurrency()
235    }
236}
237
238/// Starting point of a stream
239pub trait Loader: DynClone + Send + Sync {
240    type Output: Chunk;
241
242    fn into_stream(self) -> IndexingStream<Self::Output>;
243
244    /// Intended for use with Box<dyn Loader>
245    ///
246    /// Only needed if you use trait objects (Box<dyn Loader>)
247    ///
248    /// # Example
249    ///
250    /// ```ignore
251    /// fn into_stream_boxed(self: Box<Self>) -> IndexingStream {
252    ///    self.into_stream()
253    ///  }
254    /// ```
255    fn into_stream_boxed(self: Box<Self>) -> IndexingStream<Self::Output> {
256        unimplemented!(
257            "Please implement into_stream_boxed for your loader, it needs to be implemented on the concrete type"
258        )
259    }
260
261    fn name(&self) -> &'static str {
262        let name = std::any::type_name::<Self>();
263        name.split("::").last().unwrap_or(name)
264    }
265}
266
267dyn_clone::clone_trait_object!(<O> Loader<Output = O>);
268
269#[cfg(feature = "test-utils")]
270mock! {
271    #[derive(Debug)]
272    pub Loader {}
273
274    #[async_trait]
275    impl Loader for Loader {
276        type Output = String;
277
278        fn into_stream(self) -> IndexingStream<String>;
279        fn into_stream_boxed(self: Box<Self>) -> IndexingStream<String>;
280        fn name(&self) -> &'static str;
281    }
282
283    impl Clone for Loader {
284        fn clone(&self) -> Self;
285    }
286}
287
288impl<O: Chunk> Loader for Box<dyn Loader<Output = O>> {
289    type Output = O;
290
291    fn into_stream(self) -> IndexingStream<Self::Output> {
292        Loader::into_stream_boxed(self)
293    }
294
295    fn into_stream_boxed(self: Box<Self>) -> IndexingStream<Self::Output> {
296        Loader::into_stream(*self)
297    }
298    fn name(&self) -> &'static str {
299        self.as_ref().name()
300    }
301}
302
303impl<O: Chunk> Loader for &dyn Loader<Output = O> {
304    type Output = O;
305
306    fn into_stream(self) -> IndexingStream<Self::Output> {
307        Loader::into_stream_boxed(Box::new(self))
308    }
309
310    fn into_stream_boxed(self: Box<Self>) -> IndexingStream<Self::Output> {
311        Loader::into_stream(*self)
312    }
313}
314
315#[async_trait]
316/// Turns one node into many nodes
317pub trait ChunkerTransformer: Send + Sync + DynClone {
318    type Input: Chunk;
319    type Output: Chunk;
320
321    async fn transform_node(&self, node: Node<Self::Input>) -> IndexingStream<Self::Output>;
322
323    /// Overrides the default concurrency of the pipeline
324    fn concurrency(&self) -> Option<usize> {
325        None
326    }
327
328    fn name(&self) -> &'static str {
329        let name = std::any::type_name::<Self>();
330        name.split("::").last().unwrap_or(name)
331    }
332}
333
334dyn_clone::clone_trait_object!(<I, O> ChunkerTransformer<Input = I, Output = O>);
335
336#[cfg(feature = "test-utils")]
337mock! {
338    #[derive(Debug)]
339    pub ChunkerTransformer {}
340
341    #[async_trait]
342    impl ChunkerTransformer for ChunkerTransformer {
343        type Input = String;
344        type Output = String;
345
346    async fn transform_node(&self, node: Node<String>) -> IndexingStream<String>;
347        fn name(&self) -> &'static str;
348        fn concurrency(&self) -> Option<usize>;
349    }
350
351    impl Clone for ChunkerTransformer {
352        fn clone(&self) -> Self;
353    }
354}
355#[async_trait]
356impl<I: Chunk, O: Chunk> ChunkerTransformer for Box<dyn ChunkerTransformer<Input = I, Output = O>> {
357    type Input = I;
358    type Output = O;
359
360    async fn transform_node(&self, node: Node<I>) -> IndexingStream<O> {
361        self.as_ref().transform_node(node).await
362    }
363    fn concurrency(&self) -> Option<usize> {
364        self.as_ref().concurrency()
365    }
366    fn name(&self) -> &'static str {
367        self.as_ref().name()
368    }
369}
370
371#[async_trait]
372impl<I: Chunk, O: Chunk> ChunkerTransformer for Arc<dyn ChunkerTransformer<Input = I, Output = O>> {
373    type Input = I;
374    type Output = O;
375
376    async fn transform_node(&self, node: Node<I>) -> IndexingStream<O> {
377        self.as_ref().transform_node(node).await
378    }
379    fn concurrency(&self) -> Option<usize> {
380        self.as_ref().concurrency()
381    }
382    fn name(&self) -> &'static str {
383        self.as_ref().name()
384    }
385}
386
387#[async_trait]
388impl<I: Chunk, O: Chunk> ChunkerTransformer for &dyn ChunkerTransformer<Input = I, Output = O> {
389    type Input = I;
390    type Output = O;
391
392    async fn transform_node(&self, node: Node<I>) -> IndexingStream<O> {
393        (*self).transform_node(node).await
394    }
395    fn concurrency(&self) -> Option<usize> {
396        (*self).concurrency()
397    }
398}
399
400#[async_trait]
401impl<F> ChunkerTransformer for F
402where
403    F: Fn(Node<String>) -> IndexingStream<String> + Send + Sync + Clone,
404{
405    async fn transform_node(&self, node: Node<String>) -> IndexingStream<String> {
406        self(node)
407    }
408
409    type Input = String;
410
411    type Output = String;
412}
413
414#[async_trait]
415/// Caches nodes, typically by their path and hash
416/// Recommended to namespace on the storage
417///
418/// For now just bool return value for easy filter
419pub trait NodeCache: Send + Sync + Debug + DynClone {
420    type Input: Chunk;
421
422    async fn get(&self, node: &Node<Self::Input>) -> bool;
423    async fn set(&self, node: &Node<Self::Input>);
424
425    /// Marks the node identified by `id` as cached. Used by the pipeline to cache nodes only
426    /// after they made it through the whole pipeline, so a node that failed midway is not
427    /// skipped on the next run.
428    async fn set_by_id(&self, id: uuid::Uuid);
429
430    /// Optionally provide a method to clear the cache
431    async fn clear(&self) -> Result<()> {
432        unimplemented!("Clear not implemented")
433    }
434
435    fn name(&self) -> &'static str {
436        let name = std::any::type_name::<Self>();
437        name.split("::").last().unwrap_or(name)
438    }
439}
440
441dyn_clone::clone_trait_object!(<T> NodeCache<Input = T>);
442
443#[cfg(feature = "test-utils")]
444mock! {
445    #[derive(Debug)]
446    pub NodeCache {}
447
448    #[async_trait]
449    impl NodeCache for NodeCache {
450        type Input = String;
451        async fn get(&self, node: &Node<String>) -> bool;
452        async fn set(&self, node: &Node<String>);
453        async fn set_by_id(&self, id: uuid::Uuid);
454        async fn clear(&self) -> Result<()>;
455        fn name(&self) -> &'static str;
456
457    }
458
459    impl Clone for NodeCache {
460        fn clone(&self) -> Self;
461    }
462}
463
464#[async_trait]
465impl<T: Chunk> NodeCache for Box<dyn NodeCache<Input = T>> {
466    type Input = T;
467
468    async fn get(&self, node: &Node<T>) -> bool {
469        self.as_ref().get(node).await
470    }
471    async fn set(&self, node: &Node<T>) {
472        self.as_ref().set(node).await;
473    }
474    async fn set_by_id(&self, id: uuid::Uuid) {
475        self.as_ref().set_by_id(id).await;
476    }
477    async fn clear(&self) -> Result<()> {
478        self.as_ref().clear().await
479    }
480    fn name(&self) -> &'static str {
481        self.as_ref().name()
482    }
483}
484
485#[async_trait]
486impl<T: Chunk> NodeCache for Arc<dyn NodeCache<Input = T>> {
487    type Input = T;
488    async fn get(&self, node: &Node<T>) -> bool {
489        self.as_ref().get(node).await
490    }
491    async fn set(&self, node: &Node<T>) {
492        self.as_ref().set(node).await;
493    }
494    async fn set_by_id(&self, id: uuid::Uuid) {
495        self.as_ref().set_by_id(id).await;
496    }
497    async fn clear(&self) -> Result<()> {
498        self.as_ref().clear().await
499    }
500    fn name(&self) -> &'static str {
501        self.as_ref().name()
502    }
503}
504
505#[async_trait]
506impl<T: Chunk> NodeCache for &dyn NodeCache<Input = T> {
507    type Input = T;
508    async fn get(&self, node: &Node<T>) -> bool {
509        (*self).get(node).await
510    }
511    async fn set(&self, node: &Node<T>) {
512        (*self).set(node).await;
513    }
514    async fn set_by_id(&self, id: uuid::Uuid) {
515        (*self).set_by_id(id).await;
516    }
517    async fn clear(&self) -> Result<()> {
518        (*self).clear().await
519    }
520}
521
522#[async_trait]
523/// Embeds a list of strings and returns its embeddings.
524/// Assumes the strings will be moved.
525pub trait EmbeddingModel: Send + Sync + Debug + DynClone {
526    async fn embed(&self, input: Vec<String>) -> Result<Embeddings, LanguageModelError>;
527
528    fn name(&self) -> &'static str {
529        let name = std::any::type_name::<Self>();
530        name.split("::").last().unwrap_or(name)
531    }
532}
533
534dyn_clone::clone_trait_object!(EmbeddingModel);
535
536#[cfg(feature = "test-utils")]
537mock! {
538    #[derive(Debug)]
539    pub EmbeddingModel {}
540
541    #[async_trait]
542    impl EmbeddingModel for EmbeddingModel {
543        async fn embed(&self, input: Vec<String>) -> Result<Embeddings, LanguageModelError>;
544        fn name(&self) -> &'static str;
545    }
546
547    impl Clone for EmbeddingModel {
548        fn clone(&self) -> Self;
549    }
550}
551
552#[async_trait]
553impl EmbeddingModel for Box<dyn EmbeddingModel> {
554    async fn embed(&self, input: Vec<String>) -> Result<Embeddings, LanguageModelError> {
555        self.as_ref().embed(input).await
556    }
557
558    fn name(&self) -> &'static str {
559        self.as_ref().name()
560    }
561}
562
563#[async_trait]
564impl EmbeddingModel for Arc<dyn EmbeddingModel> {
565    async fn embed(&self, input: Vec<String>) -> Result<Embeddings, LanguageModelError> {
566        self.as_ref().embed(input).await
567    }
568
569    fn name(&self) -> &'static str {
570        self.as_ref().name()
571    }
572}
573
574#[async_trait]
575impl EmbeddingModel for &dyn EmbeddingModel {
576    async fn embed(&self, input: Vec<String>) -> Result<Embeddings, LanguageModelError> {
577        (*self).embed(input).await
578    }
579}
580
581#[async_trait]
582/// Embeds a list of strings and returns its embeddings.
583/// Assumes the strings will be moved.
584pub trait SparseEmbeddingModel: Send + Sync + Debug + DynClone {
585    async fn sparse_embed(
586        &self,
587        input: Vec<String>,
588    ) -> Result<SparseEmbeddings, LanguageModelError>;
589
590    fn name(&self) -> &'static str {
591        let name = std::any::type_name::<Self>();
592        name.split("::").last().unwrap_or(name)
593    }
594}
595
596dyn_clone::clone_trait_object!(SparseEmbeddingModel);
597
598#[cfg(feature = "test-utils")]
599mock! {
600    #[derive(Debug)]
601    pub SparseEmbeddingModel {}
602
603    #[async_trait]
604    impl SparseEmbeddingModel for SparseEmbeddingModel {
605        async fn sparse_embed(&self, input: Vec<String>) -> Result<SparseEmbeddings, LanguageModelError>;
606        fn name(&self) -> &'static str;
607    }
608
609    impl Clone for SparseEmbeddingModel {
610        fn clone(&self) -> Self;
611    }
612}
613
614#[async_trait]
615impl SparseEmbeddingModel for Box<dyn SparseEmbeddingModel> {
616    async fn sparse_embed(
617        &self,
618        input: Vec<String>,
619    ) -> Result<SparseEmbeddings, LanguageModelError> {
620        self.as_ref().sparse_embed(input).await
621    }
622
623    fn name(&self) -> &'static str {
624        self.as_ref().name()
625    }
626}
627
628#[async_trait]
629impl SparseEmbeddingModel for Arc<dyn SparseEmbeddingModel> {
630    async fn sparse_embed(
631        &self,
632        input: Vec<String>,
633    ) -> Result<SparseEmbeddings, LanguageModelError> {
634        self.as_ref().sparse_embed(input).await
635    }
636
637    fn name(&self) -> &'static str {
638        self.as_ref().name()
639    }
640}
641
642#[async_trait]
643impl SparseEmbeddingModel for &dyn SparseEmbeddingModel {
644    async fn sparse_embed(
645        &self,
646        input: Vec<String>,
647    ) -> Result<SparseEmbeddings, LanguageModelError> {
648        (*self).sparse_embed(input).await
649    }
650}
651
652#[async_trait]
653/// Given a string prompt, queries an LLM
654pub trait SimplePrompt: Debug + Send + Sync + DynClone {
655    // Takes a simple prompt, prompts the llm and returns the response
656    async fn prompt(&self, prompt: Prompt) -> Result<String, LanguageModelError>;
657
658    fn name(&self) -> &'static str {
659        let name = std::any::type_name::<Self>();
660        name.split("::").last().unwrap_or(name)
661    }
662}
663
664dyn_clone::clone_trait_object!(SimplePrompt);
665
666#[cfg(feature = "test-utils")]
667mock! {
668    #[derive(Debug)]
669    pub SimplePrompt {}
670
671    #[async_trait]
672    impl SimplePrompt for SimplePrompt {
673        async fn prompt(&self, prompt: Prompt) -> Result<String, LanguageModelError>;
674        fn name(&self) -> &'static str;
675    }
676
677    impl Clone for SimplePrompt {
678        fn clone(&self) -> Self;
679    }
680}
681
682#[async_trait]
683impl SimplePrompt for Box<dyn SimplePrompt> {
684    async fn prompt(&self, prompt: Prompt) -> Result<String, LanguageModelError> {
685        self.as_ref().prompt(prompt).await
686    }
687
688    fn name(&self) -> &'static str {
689        self.as_ref().name()
690    }
691}
692
693#[async_trait]
694impl SimplePrompt for Arc<dyn SimplePrompt> {
695    async fn prompt(&self, prompt: Prompt) -> Result<String, LanguageModelError> {
696        self.as_ref().prompt(prompt).await
697    }
698
699    fn name(&self) -> &'static str {
700        self.as_ref().name()
701    }
702}
703
704#[async_trait]
705impl SimplePrompt for &dyn SimplePrompt {
706    async fn prompt(&self, prompt: Prompt) -> Result<String, LanguageModelError> {
707        (*self).prompt(prompt).await
708    }
709}
710
711#[async_trait]
712/// Persists nodes
713pub trait Persist: Debug + Send + Sync + DynClone {
714    type Input: Chunk;
715    type Output: Chunk;
716
717    async fn setup(&self) -> Result<()>;
718    async fn store(&self, node: Node<Self::Input>) -> Result<Node<Self::Output>>;
719    async fn batch_store(&self, nodes: Vec<Node<Self::Input>>) -> IndexingStream<Self::Output>;
720    fn batch_size(&self) -> Option<usize> {
721        None
722    }
723
724    fn name(&self) -> &'static str {
725        let name = std::any::type_name::<Self>();
726        name.split("::").last().unwrap_or(name)
727    }
728}
729
730dyn_clone::clone_trait_object!(<I, O> Persist<Input = I, Output = O>);
731
732#[cfg(feature = "test-utils")]
733mock! {
734    #[derive(Debug)]
735    pub Persist {}
736
737    #[async_trait]
738    impl Persist for Persist {
739        type Input = String;
740        type Output = String;
741
742        async fn setup(&self) -> Result<()>;
743        async fn store(&self, node: Node<String>) -> Result<Node<String>>;
744        async fn batch_store(&self, nodes: Vec<Node<String>>) -> IndexingStream<String>;
745        fn batch_size(&self) -> Option<usize>;
746
747        fn name(&self) -> &'static str;
748    }
749
750    impl Clone for Persist {
751        fn clone(&self) -> Self;
752    }
753}
754
755#[async_trait]
756impl<I: Chunk, O: Chunk> Persist for Box<dyn Persist<Input = I, Output = O>> {
757    type Input = I;
758    type Output = O;
759
760    async fn setup(&self) -> Result<()> {
761        self.as_ref().setup().await
762    }
763    async fn store(&self, node: Node<I>) -> Result<Node<O>> {
764        self.as_ref().store(node).await
765    }
766    async fn batch_store(&self, nodes: Vec<Node<I>>) -> IndexingStream<O> {
767        self.as_ref().batch_store(nodes).await
768    }
769    fn batch_size(&self) -> Option<usize> {
770        self.as_ref().batch_size()
771    }
772    fn name(&self) -> &'static str {
773        self.as_ref().name()
774    }
775}
776
777#[async_trait]
778impl<I: Chunk, O: Chunk> Persist for Arc<dyn Persist<Input = I, Output = O>> {
779    type Input = I;
780    type Output = O;
781
782    async fn setup(&self) -> Result<()> {
783        self.as_ref().setup().await
784    }
785    async fn store(&self, node: Node<I>) -> Result<Node<O>> {
786        self.as_ref().store(node).await
787    }
788    async fn batch_store(&self, nodes: Vec<Node<I>>) -> IndexingStream<O> {
789        self.as_ref().batch_store(nodes).await
790    }
791    fn batch_size(&self) -> Option<usize> {
792        self.as_ref().batch_size()
793    }
794    fn name(&self) -> &'static str {
795        self.as_ref().name()
796    }
797}
798
799#[async_trait]
800impl<I: Chunk, O: Chunk> Persist for &dyn Persist<Input = I, Output = O> {
801    type Input = I;
802    type Output = O;
803
804    async fn setup(&self) -> Result<()> {
805        (*self).setup().await
806    }
807    async fn store(&self, node: Node<I>) -> Result<Node<O>> {
808        (*self).store(node).await
809    }
810    async fn batch_store(&self, nodes: Vec<Node<I>>) -> IndexingStream<O> {
811        (*self).batch_store(nodes).await
812    }
813    fn batch_size(&self) -> Option<usize> {
814        (*self).batch_size()
815    }
816}
817
818/// Allows for passing defaults from the pipeline to the transformer
819/// Required for batch transformers as at least a marker, implementation is not required
820pub trait WithIndexingDefaults {
821    fn with_indexing_defaults(&mut self, _indexing_defaults: IndexingDefaults) {}
822}
823
824/// Allows for passing defaults from the pipeline to the batch transformer
825/// Required for batch transformers as at least a marker, implementation is not required
826pub trait WithBatchIndexingDefaults {
827    fn with_indexing_defaults(&mut self, _indexing_defaults: IndexingDefaults) {}
828}
829
830impl<I, O> WithIndexingDefaults for dyn Transformer<Input = I, Output = O> {}
831impl<I, O> WithIndexingDefaults for Box<dyn Transformer<Input = I, Output = O>> {
832    fn with_indexing_defaults(&mut self, indexing_defaults: IndexingDefaults) {
833        self.as_mut().with_indexing_defaults(indexing_defaults);
834    }
835}
836impl<I, O> WithBatchIndexingDefaults for dyn BatchableTransformer<Input = I, Output = O> {}
837impl<I, O> WithBatchIndexingDefaults for Box<dyn BatchableTransformer<Input = I, Output = O>> {
838    fn with_indexing_defaults(&mut self, indexing_defaults: IndexingDefaults) {
839        self.as_mut().with_indexing_defaults(indexing_defaults);
840    }
841}
842
843impl<F> WithIndexingDefaults for F where F: Fn(Node<String>) -> Result<Node<String>> {}
844impl<F> WithBatchIndexingDefaults for F where F: Fn(Vec<Node<String>>) -> IndexingStream<String> {}
845
846#[cfg(feature = "test-utils")]
847impl WithIndexingDefaults for MockTransformer {}
848// //
849#[cfg(feature = "test-utils")]
850impl WithBatchIndexingDefaults for MockBatchableTransformer {}
851
852#[async_trait]
853/// Given a string prompt, queries an LLM to return structured data
854pub trait StructuredPrompt: Debug + Send + Sync + DynClone {
855    async fn structured_prompt<T: DeserializeOwned + JsonSchema>(
856        &self,
857        prompt: Prompt,
858    ) -> Result<T, LanguageModelError>;
859
860    fn name(&self) -> &'static str {
861        let name = std::any::type_name::<Self>();
862        name.split("::").last().unwrap_or(name)
863    }
864}
865
866/// Helper trait object to call structured prompt with dynamic dispatch
867///
868/// Internally Swiftide only implements this trait, as implementing `DynStructuredPrompt` gives
869/// `StructuredPrompt` for free
870#[async_trait]
871pub trait DynStructuredPrompt: Debug + Send + Sync + DynClone {
872    async fn structured_prompt_dyn(
873        &self,
874        prompt: Prompt,
875        schema: schemars::Schema,
876    ) -> Result<serde_json::Value, LanguageModelError>;
877
878    fn name(&self) -> &'static str {
879        let name = std::any::type_name::<Self>();
880        name.split("::").last().unwrap_or(name)
881    }
882}
883
884dyn_clone::clone_trait_object!(DynStructuredPrompt);
885
886#[async_trait]
887impl<C> StructuredPrompt for C
888where
889    C: DynStructuredPrompt + Debug + Send + Sync + DynClone,
890{
891    async fn structured_prompt<T: DeserializeOwned + JsonSchema>(
892        &self,
893        prompt: Prompt,
894    ) -> Result<T, LanguageModelError> {
895        // Call with T = serde_json::Value
896        let schema = schema_for!(T);
897        let val = self.structured_prompt_dyn(prompt, schema).await?;
898
899        let parsed = serde_json::from_value(val).map_err(LanguageModelError::permanent)?;
900
901        Ok(parsed)
902    }
903}
swiftide_core/indexing_traits.rs

swiftide_core/
indexing_traits.rs
