Skip to main content

rig_core/
operation.rs

1//! The operations a [`Wire`](crate::wire::Wire) can perform.
2//!
3//! Each operation declares its request, event, end and response types and
4//! the fold one reply goes through. [`Completion`] streams events; most
5//! others answer with one whole document, folded by [`Whole`]. Listings
6//! follow their cursors above the driver: each page is one call.
7//!
8//! ```
9//! use rig_core::operation::{Completion, Embedding};
10//! use rig_core::wire::Operation;
11//!
12//! fn streams<Op: Operation>() {}
13//! streams::<Completion>();
14//! streams::<Embedding>();
15//! ```
16
17use crate::error::ProviderError;
18use crate::wire::{Fold, Operation, Reply};
19
20mod cached_content;
21pub(crate) mod completion;
22mod listing;
23mod modality;
24mod verify;
25
26pub use cached_content::{CachedContentFold, ContextCache};
27pub use completion::{
28    CallFragment, CallPart, Completion, Finish, IfMalformed, ReasoningPart, Seal, TextPart, Turn,
29};
30pub use listing::{ModelListing, ModelPage};
31#[cfg(feature = "audio")]
32pub use modality::AudioGeneration;
33#[cfg(feature = "image")]
34pub use modality::ImageGeneration;
35pub use modality::{Embedding, ImageEmbedding, Rerank, RerankRequest, Transcription};
36pub use verify::{Verify, VerifyDecoder};
37
38/// The fold of an operation whose reply is one whole answer: it has no
39/// events, and the provider's end is the response. The stamp writes what
40/// the driver learned about the reply onto it.
41///
42/// ```
43/// use std::convert::Infallible;
44/// use rig_core::operation::Whole;
45/// use rig_core::wire::{Call, Free, Operation};
46///
47/// struct Echo;
48///
49/// impl Operation for Echo {
50///     type Request = String;
51///     type Event = Infallible;
52///     type End = String;
53///     type Response = String;
54///     type Fold = Whole<Self>;
55///     type Emit = Free;
56///
57///     fn fold(_request: &String, _call: &mut Call<'_>) -> Whole<Self> {
58///         Whole::<Self>::stamping(|response, reply| response.push_str(&reply.provider))
59///     }
60/// }
61/// ```
62pub struct Whole<Op: Operation> {
63    stamp: fn(&mut Op::Response, &Reply),
64}
65
66impl<Op: Operation> Whole<Op> {
67    /// A fold that stamps nothing.
68    pub fn new() -> Self {
69        Self::stamping(|_, _| {})
70    }
71
72    /// A fold that writes what the driver learned about the reply onto the
73    /// response with `stamp`.
74    pub fn stamping(stamp: fn(&mut Op::Response, &Reply)) -> Self {
75        Self { stamp }
76    }
77}
78
79impl<Op: Operation> Default for Whole<Op> {
80    fn default() -> Self {
81        Self::new()
82    }
83}
84
85impl<Op> Fold<Op> for Whole<Op>
86where
87    Op: Operation<Event = std::convert::Infallible>,
88    Op::End: Into<Op::Response>,
89{
90    fn absorb(&mut self, event: &std::convert::Infallible) -> Result<(), ProviderError> {
91        match *event {}
92    }
93
94    fn finish(self, end: Op::End, reply: Reply) -> Result<Op::Response, ProviderError> {
95        let mut response = end.into();
96        (self.stamp)(&mut response, &reply);
97        Ok(response)
98    }
99}