rig/client/completion.rs
1use crate::agent::AgentBuilder;
2use crate::completion::CompletionModel;
3use crate::extractor::ExtractorBuilder;
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6
7/// A provider client with completion capabilities.
8/// Clone is required for conversions between client types.
9pub trait CompletionClient {
10 /// The type of CompletionModel used by the client.
11 type CompletionModel: CompletionModel<Client = Self>;
12
13 /// Create a completion model with the given model.
14 ///
15 /// # Example with OpenAI
16 /// ```
17 /// use rig::prelude::*;
18 /// use rig::providers::openai::{Client, self};
19 ///
20 /// // Initialize the OpenAI client
21 /// let openai = Client::new("your-open-ai-api-key");
22 ///
23 /// let gpt4 = openai.completion_model(openai::GPT4);
24 /// ```
25 fn completion_model(&self, model: impl Into<String>) -> Self::CompletionModel {
26 Self::CompletionModel::make(self, model)
27 }
28
29 /// Create an agent builder with the given completion model.
30 ///
31 /// # Example with OpenAI
32 /// ```
33 /// use rig::prelude::*;
34 /// use rig::providers::openai::{Client, self};
35 ///
36 /// // Initialize the OpenAI client
37 /// let openai = Client::new("your-open-ai-api-key");
38 ///
39 /// let agent = openai.agent(openai::GPT_4)
40 /// .preamble("You are comedian AI with a mission to make people laugh.")
41 /// .temperature(0.0)
42 /// .build();
43 /// ```
44 fn agent(&self, model: impl Into<String>) -> AgentBuilder<Self::CompletionModel> {
45 AgentBuilder::new(self.completion_model(model))
46 }
47
48 /// Create an extractor builder with the given completion model.
49 fn extractor<T>(&self, model: impl Into<String>) -> ExtractorBuilder<Self::CompletionModel, T>
50 where
51 T: JsonSchema + for<'a> Deserialize<'a> + Serialize + Send + Sync,
52 {
53 ExtractorBuilder::new(self.completion_model(model))
54 }
55}