rig/client/
completion.rs

1use crate::agent::AgentBuilder;
2use crate::client::{AsCompletion, ProviderClient};
3use crate::completion::{
4    CompletionError, CompletionModel, CompletionModelDyn, CompletionRequest, CompletionResponse,
5};
6use crate::extractor::ExtractorBuilder;
7use crate::streaming::StreamingCompletionResponse;
8use schemars::JsonSchema;
9use serde::{Deserialize, Serialize};
10use std::future::Future;
11use std::sync::Arc;
12
13/// A provider client with completion capabilities.
14/// Clone is required for conversions between client types.
15pub trait CompletionClient: ProviderClient + Clone {
16    /// The type of CompletionModel used by the client.
17    type CompletionModel: CompletionModel;
18
19    /// Create a completion model with the given name.
20    ///
21    /// # Example with OpenAI
22    /// ```
23    /// use rig::prelude::*;
24    /// use rig::providers::openai::{Client, self};
25    ///
26    /// // Initialize the OpenAI client
27    /// let openai = Client::new("your-open-ai-api-key");
28    ///
29    /// let gpt4 = openai.completion_model(openai::GPT_4);
30    /// ```
31    fn completion_model(&self, model: &str) -> Self::CompletionModel;
32
33    /// Create an agent builder with the given completion model.
34    ///
35    /// # Example with OpenAI
36    /// ```
37    /// use rig::prelude::*;
38    /// use rig::providers::openai::{Client, self};
39    ///
40    /// // Initialize the OpenAI client
41    /// let openai = Client::new("your-open-ai-api-key");
42    ///
43    /// let agent = openai.agent(openai::GPT_4)
44    ///    .preamble("You are comedian AI with a mission to make people laugh.")
45    ///    .temperature(0.0)
46    ///    .build();
47    /// ```
48    fn agent(&self, model: &str) -> AgentBuilder<Self::CompletionModel> {
49        AgentBuilder::new(self.completion_model(model))
50    }
51
52    /// Create an extractor builder with the given completion model.
53    fn extractor<T: JsonSchema + for<'a> Deserialize<'a> + Serialize + Send + Sync>(
54        &self,
55        model: &str,
56    ) -> ExtractorBuilder<T, Self::CompletionModel> {
57        ExtractorBuilder::new(self.completion_model(model))
58    }
59}
60
61/// Wraps a CompletionModel in a dyn-compatible way for AgentBuilder.
62#[derive(Clone)]
63pub struct CompletionModelHandle<'a> {
64    pub inner: Arc<dyn CompletionModelDyn + 'a>,
65}
66
67impl CompletionModel for CompletionModelHandle<'_> {
68    type Response = ();
69    type StreamingResponse = ();
70
71    fn completion(
72        &self,
73        request: CompletionRequest,
74    ) -> impl Future<Output = Result<CompletionResponse<Self::Response>, CompletionError>> + Send
75    {
76        self.inner.completion(request)
77    }
78
79    fn stream(
80        &self,
81        request: CompletionRequest,
82    ) -> impl Future<
83        Output = Result<StreamingCompletionResponse<Self::StreamingResponse>, CompletionError>,
84    > + Send {
85        self.inner.stream(request)
86    }
87}
88
89pub trait CompletionClientDyn: ProviderClient {
90    /// Create a completion model with the given name.
91    fn completion_model<'a>(&self, model: &str) -> Box<dyn CompletionModelDyn + 'a>;
92
93    /// Create an agent builder with the given completion model.
94    fn agent<'a>(&self, model: &str) -> AgentBuilder<CompletionModelHandle<'a>>;
95}
96
97impl<
98        T: CompletionClient<CompletionModel = M>,
99        M: CompletionModel<StreamingResponse = R> + 'static,
100        R: Clone + Unpin + 'static,
101    > CompletionClientDyn for T
102{
103    fn completion_model<'a>(&self, model: &str) -> Box<dyn CompletionModelDyn + 'a> {
104        Box::new(self.completion_model(model))
105    }
106
107    fn agent<'a>(&self, model: &str) -> AgentBuilder<CompletionModelHandle<'a>> {
108        AgentBuilder::new(CompletionModelHandle {
109            inner: Arc::new(self.completion_model(model)),
110        })
111    }
112}
113
114impl<T: CompletionClientDyn + Clone + 'static> AsCompletion for T {
115    fn as_completion(&self) -> Option<Box<dyn CompletionClientDyn>> {
116        Some(Box::new(self.clone()))
117    }
118}
rig/client/completion.rs

rig/client/
completion.rs
