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}