syncable_cli/bedrock/
completion.rs1use super::{
4 client::Client,
5 types::{
6 assistant_content::AwsConverseOutput, completion_request::AwsCompletionRequest,
7 converse_output::InternalConverseOutput, errors::AwsSdkConverseError,
8 },
9};
10
11use rig::completion::{self, CompletionError, CompletionRequest};
12use rig::streaming::StreamingCompletionResponse;
13
14pub const AI21_JAMBA_1_5_LARGE: &str = "ai21.jamba-1-5-large-v1:0";
16pub const AI21_JAMBA_1_5_MINI: &str = "ai21.jamba-1-5-mini-v1:0";
18pub const AMAZON_NOVA_CANVAS: &str = "amazon.nova-canvas-v1:0";
20pub const AMAZON_NOVA_LITE: &str = "amazon.nova-lite-v1:0";
22pub const AMAZON_NOVA_MICRO: &str = "amazon.nova-micro-v1:0";
24pub const AMAZON_NOVA_PREMIER: &str = "amazon.nova-premier-v1:0";
26pub const AMAZON_NOVA_PRO: &str = "amazon.nova-pro-v1:0";
28pub const AMAZON_NOVA_REEL_V1_0: &str = "amazon.nova-reel-v1:0";
30pub const AMAZON_NOVA_REEL_V1_1: &str = "amazon.nova-reel-v1:1";
32pub const AMAZON_NOVA_SONIC: &str = "amazon.nova-sonic-v1:0";
34pub const AMAZON_RERANK_1_0: &str = "amazon.rerank-v1:0";
36pub const AMAZON_TITAN_EMBEDDINGS_G1_TEXT: &str = "amazon.titan-embed-text-v1";
38pub const AMAZON_TITAN_IMAGE_GENERATOR_G1_V2: &str = "amazon.titan-image-generator-v2:0";
40pub const AMAZON_TITAN_IMAGE_GENERATOR_G1: &str = "amazon.titan-image-generator-v1";
42pub const AMAZON_TITAN_MULTIMODAL_EMBEDDINGS_G1: &str = "amazon.titan-embed-image-v1";
44pub const AMAZON_TITAN_TEXT_EMBEDDINGS_V2: &str = "amazon.titan-embed-text-v2:0";
46pub const AMAZON_TITAN_TEXT_EXPRESS_V1: &str = "amazon.titan-text-express-v1";
48pub const AMAZON_TITAN_TEXT_LITE_V1: &str = "amazon.titan-text-lite-v1";
50pub const AMAZON_TITAN_TEXT_PREMIER_V1_0: &str = "amazon.titan-text-premier-v1:0";
52pub const ANTHROPIC_CLAUDE_3_HAIKU: &str = "anthropic.claude-3-haiku-20240307-v1:0";
54pub const ANTHROPIC_CLAUDE_3_OPUS: &str = "anthropic.claude-3-opus-20240229-v1:0";
56pub const ANTHROPIC_CLAUDE_3_SONNET: &str = "anthropic.claude-3-sonnet-20240229-v1:0";
58pub const ANTHROPIC_CLAUDE_3_5_HAIKU: &str = "anthropic.claude-3-5-haiku-20241022-v1:0";
60pub const ANTHROPIC_CLAUDE_3_5_SONNET_V2: &str = "anthropic.claude-3-5-sonnet-20241022-v2:0";
62pub const ANTHROPIC_CLAUDE_3_5_SONNET: &str = "anthropic.claude-3-5-sonnet-20240620-v1:0";
64pub const ANTHROPIC_CLAUDE_3_7_SONNET: &str = "anthropic.claude-3-7-sonnet-20250219-v1:0";
66pub const ANTHROPIC_CLAUDE_OPUS_4: &str = "anthropic.claude-opus-4-20250514-v1:0";
68pub const ANTHROPIC_CLAUDE_SONNET_4: &str = "anthropic.claude-sonnet-4-20250514-v1:0";
70pub const COHERE_COMMAND_LIGHT_TEXT: &str = "cohere.command-light-text-v14";
72pub const COHERE_COMMAND_R_PLUS: &str = "cohere.command-r-plus-v1:0";
74pub const COHERE_COMMAND_R: &str = "cohere.command-r-v1:0";
76pub const COHERE_COMMAND: &str = "cohere.command-text-v14";
78pub const COHERE_EMBED_ENGLISH: &str = "cohere.embed-english-v3";
80pub const COHERE_EMBED_MULTILINGUAL: &str = "cohere.embed-multilingual-v3";
82pub const COHERE_RERANK_V3_5: &str = "cohere.rerank-v3-5:0";
84pub const DEEPSEEK_R1: &str = "deepseek.r1-v1:0";
86pub const LUMA_RAY_V2_0: &str = "luma.ray-v2:0";
88pub const LLAMA_3_8B_INSTRUCT: &str = "meta.llama3-8b-instruct-v1:0";
90pub const LLAMA_3_70B_INSTRUCT: &str = "meta.llama3-70b-instruct-v1:0";
92pub const LLAMA_3_1_8B_INSTRUCT: &str = "meta.llama3-1-8b-instruct-v1:0";
94pub const LLAMA_3_1_70B_INSTRUCT: &str = "meta.llama3-1-70b-instruct-v1:0";
96pub const LLAMA_3_1_405B_INSTRUCT: &str = "meta.llama3-1-405b-instruct-v1:0";
98pub const LLAMA_3_2_1B_INSTRUCT: &str = "meta.llama3-2-1b-instruct-v1:0";
100pub const LLAMA_3_2_3B_INSTRUCT: &str = "meta.llama3-2-3b-instruct-v1:0";
102pub const LLAMA_3_2_11B_INSTRUCT: &str = "meta.llama3-2-11b-instruct-v1:0";
104pub const LLAMA_3_2_90B_INSTRUCT: &str = "meta.llama3-2-90b-instruct-v1:0";
106pub const META_LLAMA_3_3_70B_INSTRUCT: &str = "meta.llama3-3-70b-instruct-v1:0";
108pub const META_LLAMA_4_MAVERICK_17B_INSTRUCT: &str = "meta.llama4-maverick-17b-instruct-v1:0";
110pub const META_LLAMA_4_SCOUT_17B_INSTRUCT: &str = "meta.llama4-scout-17b-instruct-v1:0";
112pub const MISTRAL_7B_INSTRUCT: &str = "mistral.mistral-7b-instruct-v0:2";
114pub const MISTRAL_LARGE_24_02: &str = "mistral.mistral-large-2402-v1:0";
116pub const MISTRAL_LARGE_24_07: &str = "mistral.mistral-large-2407-v1:0";
118pub const MISTRAL_SMALL_24_02: &str = "mistral.mistral-small-2402-v1:0";
120pub const MISTRAL_MIXTRAL_8X7B_INSTRUCT_V0: &str = "mistral.mixtral-8x7b-instruct-v0:1";
122pub const MISTRAL_PIXTRAL_LARGE_2502: &str = "mistral.pixtral-large-2502-v1:0";
124pub const STABILITY_SD3_5_LARGE: &str = "stability.sd3-5-large-v1:0";
126pub const STABILITY_STABLE_IMAGE_CORE_1_0: &str = "stability.stable-image-core-v1:1";
128pub const STABILITY_STABLE_IMAGE_ULTRA_1_0: &str = "stability.stable-image-ultra-v1:1";
130pub const TWELVELABS_MARENGO_EMBED_V2_7: &str = "twelvelabs.marengo-embed-2-7-v1:0";
132pub const TWELVELABS_PEGASUS_V1_2: &str = "twelvelabs.pegasus-1-2-v1:0";
134pub const WRITER_PALMYRA_X4: &str = "writer.palmyra-x4-v1:0";
136pub const WRITER_PALMYRA_X5: &str = "writer.palmyra-x5-v1:0";
138pub const AI21_JAMBA_INSTRUCT: &str = "ai21.jamba-instruct-v1:0";
140pub const ANTHROPIC_CLAUDE_2_1: &str = "anthropic.claude-v2:1";
142pub const ANTHROPIC_CLAUDE_2: &str = "anthropic.claude-v2";
144pub const ANTHROPIC_CLAUDE_INSTANT: &str = "anthropic.claude-instant-v1";
146pub const ANTHROPIC_CLAUDE_INSTANT_V1_2: &str = "anthropic.claude-instant-v1:2";
148pub const ANTHROPIC_CLAUDE: &str = "anthropic.claude-v2:0";
150pub const STABILITY_SD3_LARGE_1_0: &str = "stability.sd3-large-v1:0";
152pub const STABILITY_SDXL_1_0: &str = "stability.stable-diffusion-xl-v1";
154pub const STABILITY_STABLE_IMAGE_CORE_1_0_V1_0: &str = "stability.stable-image-core-v1:0";
156pub const STABILITY_STABLE_IMAGE_ULTRA_1_0_V1_0: &str = "stability.stable-image-ultra-v1:0";
158
159#[derive(Clone)]
160pub struct CompletionModel {
161 pub(crate) client: Client,
162 pub model: String,
163}
164
165impl CompletionModel {
166 pub fn new(client: Client, model: impl Into<String>) -> Self {
167 Self {
168 client,
169 model: model.into(),
170 }
171 }
172}
173
174impl completion::CompletionModel for CompletionModel {
175 type Response = AwsConverseOutput;
176 type StreamingResponse = super::streaming::BedrockStreamingResponse;
177
178 type Client = Client;
179
180 fn make(client: &Self::Client, model: impl Into<String>) -> Self {
181 Self::new(client.clone(), model)
182 }
183
184 async fn completion(
185 &self,
186 completion_request: completion::CompletionRequest,
187 ) -> Result<completion::CompletionResponse<AwsConverseOutput>, CompletionError> {
188 let request = AwsCompletionRequest(completion_request);
189
190 let mut converse_builder = self
191 .client
192 .get_inner()
193 .await
194 .converse()
195 .model_id(self.model.as_str());
196
197 let tool_config = request.tools_config()?;
198 let messages = request.messages()?;
199 converse_builder = converse_builder
200 .set_additional_model_request_fields(request.additional_params())
201 .set_inference_config(request.inference_config())
202 .set_tool_config(tool_config)
203 .set_system(request.system_prompt())
204 .set_messages(Some(messages));
205
206 let response = converse_builder
207 .send()
208 .await
209 .map_err(|sdk_error| Into::<CompletionError>::into(AwsSdkConverseError(sdk_error)))?;
210
211 let response: InternalConverseOutput = response
212 .try_into()
213 .map_err(|x| CompletionError::ProviderError(format!("Type conversion error: {x}")))?;
214
215 AwsConverseOutput(response).try_into()
216 }
217
218 async fn stream(
219 &self,
220 request: CompletionRequest,
221 ) -> Result<StreamingCompletionResponse<Self::StreamingResponse>, CompletionError> {
222 CompletionModel::stream(self, request).await
223 }
224}