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
20pub(crate) mod completion;
21mod listing;
22mod modality;
23mod verify;
24
25pub use completion::{Block, CallFragment, Completion, Finish, Turn, merge};
26pub use listing::{ModelListing, ModelPage};
27#[cfg(feature = "audio")]
28pub use modality::AudioGeneration;
29#[cfg(feature = "image")]
30pub use modality::ImageGeneration;
31pub use modality::{Embedding, ImageEmbedding, Rerank, RerankRequest, Transcription};
32pub use verify::{Verify, VerifyDecoder};
33
34/// The fold of an operation whose reply is one whole answer: it has no
35/// events, and the provider's end is the response. The stamp writes what
36/// the driver learned about the reply onto it.
37///
38/// ```
39/// use std::convert::Infallible;
40/// use rig_core::operation::Whole;
41/// use rig_core::wire::{Call, Free, Operation};
42///
43/// struct Echo;
44///
45/// impl Operation for Echo {
46///     type Request = String;
47///     type Event = Infallible;
48///     type End = String;
49///     type Response = String;
50///     type Fold = Whole<Self>;
51///     type Emit = Free;
52///
53///     fn fold(_request: &String, _call: &mut Call<'_>) -> Whole<Self> {
54///         Whole::<Self>::stamping(|response, reply| response.push_str(&reply.provider))
55///     }
56/// }
57/// ```
58pub struct Whole<Op: Operation> {
59    stamp: fn(&mut Op::Response, &Reply),
60}
61
62impl<Op: Operation> Whole<Op> {
63    /// A fold that stamps nothing.
64    pub fn new() -> Self {
65        Self::stamping(|_, _| {})
66    }
67
68    /// A fold that writes what the driver learned about the reply onto the
69    /// response with `stamp`.
70    pub fn stamping(stamp: fn(&mut Op::Response, &Reply)) -> Self {
71        Self { stamp }
72    }
73}
74
75impl<Op: Operation> Default for Whole<Op> {
76    fn default() -> Self {
77        Self::new()
78    }
79}
80
81impl<Op> Fold<Op> for Whole<Op>
82where
83    Op: Operation<Event = std::convert::Infallible>,
84    Op::End: Into<Op::Response>,
85{
86    fn absorb(&mut self, event: &std::convert::Infallible) -> Result<(), ProviderError> {
87        match *event {}
88    }
89
90    fn finish(self, end: Op::End, reply: Reply) -> Result<Op::Response, ProviderError> {
91        let mut response = end.into();
92        (self.stamp)(&mut response, &reply);
93        Ok(response)
94    }
95}