rig/client/
completion.rs

1use crate::agent::AgentBuilder;
2use crate::client::FinalCompletionResponse;
3
4#[allow(deprecated)]
5use crate::completion::CompletionModelDyn;
6use crate::completion::{
7    CompletionError, CompletionModel, CompletionRequest, CompletionResponse, GetTokenUsage,
8};
9use crate::extractor::ExtractorBuilder;
10use crate::streaming::StreamingCompletionResponse;
11use crate::wasm_compat::WasmCompatSend;
12use schemars::JsonSchema;
13use serde::{Deserialize, Serialize};
14use std::future::Future;
15use std::sync::Arc;
16
17/// A provider client with completion capabilities.
18/// Clone is required for conversions between client types.
19pub trait CompletionClient {
20    /// The type of CompletionModel used by the client.
21    type CompletionModel: CompletionModel<Client = Self>;
22
23    /// Create a completion model with the given model.
24    ///
25    /// # Example with OpenAI
26    /// ```
27    /// use rig::prelude::*;
28    /// use rig::providers::openai::{Client, self};
29    ///
30    /// // Initialize the OpenAI client
31    /// let openai = Client::new("your-open-ai-api-key");
32    ///
33    /// let gpt4 = openai.completion_model(openai::GPT4);
34    /// ```
35    fn completion_model(&self, model: impl Into<String>) -> Self::CompletionModel {
36        Self::CompletionModel::make(self, model)
37    }
38
39    /// Create an agent builder with the given completion model.
40    ///
41    /// # Example with OpenAI
42    /// ```
43    /// use rig::prelude::*;
44    /// use rig::providers::openai::{Client, self};
45    ///
46    /// // Initialize the OpenAI client
47    /// let openai = Client::new("your-open-ai-api-key");
48    ///
49    /// let agent = openai.agent(openai::GPT_4)
50    ///    .preamble("You are comedian AI with a mission to make people laugh.")
51    ///    .temperature(0.0)
52    ///    .build();
53    /// ```
54    fn agent(&self, model: impl Into<String>) -> AgentBuilder<Self::CompletionModel> {
55        AgentBuilder::new(self.completion_model(model))
56    }
57
58    /// Create an extractor builder with the given completion model.
59    fn extractor<T>(&self, model: impl Into<String>) -> ExtractorBuilder<Self::CompletionModel, T>
60    where
61        T: JsonSchema + for<'a> Deserialize<'a> + Serialize + Send + Sync,
62    {
63        ExtractorBuilder::new(self.completion_model(model))
64    }
65}
66
67#[allow(deprecated)]
68#[deprecated(
69    since = "0.25.0",
70    note = "`DynClientBuilder` and related features have been deprecated and will be removed in a future release."
71)]
72/// Wraps a CompletionModel in a dyn-compatible way for AgentBuilder.
73#[derive(Clone)]
74pub struct CompletionModelHandle<'a>(Arc<dyn CompletionModelDyn + 'a>);
75
76#[allow(deprecated)]
77impl<'a> CompletionModelHandle<'a> {
78    pub fn new(handle: Arc<dyn CompletionModelDyn + 'a>) -> Self {
79        Self(handle)
80    }
81}
82
83#[allow(deprecated)]
84impl CompletionModel for CompletionModelHandle<'_> {
85    type Response = ();
86    type StreamingResponse = FinalCompletionResponse;
87    type Client = ();
88
89    /// **PANICS**: We are deprecating DynClientBuilder and related functionality, in the meantime
90    /// there may be some invalid methods which panic when called, such as this one
91    fn make(_: &Self::Client, _: impl Into<String>) -> Self {
92        panic!("Cannot create a completion model handle from a client")
93    }
94
95    fn completion(
96        &self,
97        request: CompletionRequest,
98    ) -> impl Future<Output = Result<CompletionResponse<Self::Response>, CompletionError>> + WasmCompatSend
99    {
100        self.0.completion(request)
101    }
102
103    fn stream(
104        &self,
105        request: CompletionRequest,
106    ) -> impl Future<
107        Output = Result<StreamingCompletionResponse<Self::StreamingResponse>, CompletionError>,
108    > + WasmCompatSend {
109        self.0.stream(request)
110    }
111}
112
113#[allow(deprecated)]
114#[deprecated(
115    since = "0.25.0",
116    note = "`DynClientBuilder` and related features have been deprecated and will be removed in a future release. In this case, use `CompletionClient` instead."
117)]
118pub trait CompletionClientDyn {
119    /// Create a completion model with the given name.
120    fn completion_model<'a>(&self, model: &str) -> Box<dyn CompletionModelDyn + 'a>;
121
122    /// Create an agent builder with the given completion model.
123    fn agent<'a>(&self, model: &str) -> AgentBuilder<CompletionModelHandle<'a>>;
124}
125
126#[allow(deprecated)]
127impl<T, M, R> CompletionClientDyn for T
128where
129    T: CompletionClient<CompletionModel = M>,
130    M: CompletionModel<StreamingResponse = R> + 'static,
131    R: Clone + Unpin + GetTokenUsage + 'static,
132{
133    fn completion_model<'a>(&self, model: &str) -> Box<dyn CompletionModelDyn + 'a> {
134        Box::new(self.completion_model(model))
135    }
136
137    fn agent<'a>(&self, model: &str) -> AgentBuilder<CompletionModelHandle<'a>> {
138        AgentBuilder::new(CompletionModelHandle(Arc::new(
139            self.completion_model(model),
140        )))
141    }
142}
rig/client/completion.rs

rig/client/
completion.rs
