rig_bedrock/
completion.rs

1//! All supported models <https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html>
2
3use crate::{
4    client::Client,
5    types::{
6        assistant_content::AwsConverseOutput, completion_request::AwsCompletionRequest,
7        converse_output::InternalConverseOutput, errors::AwsSdkConverseError,
8    },
9};
10
11use aws_sdk_bedrockruntime::types as aws_bedrock;
12use rig_core::completion::{self, CompletionError, CompletionRequest};
13use rig_core::streaming::StreamingCompletionResponse;
14use rig_core::telemetry::ProviderResponseExt;
15use rig_core::telemetry::{CompletionOperation, CompletionSpanBuilder, SpanCombinator};
16use tracing::Instrument;
17
18// Model identifiers below were verified against `bedrock:ListFoundationModels`
19// in us-east-1, us-west-2, eu-central-1 and ap-northeast-1, and each was
20// invoked with `Converse` in us-east-1. Bedrock answers a retired identifier
21// with `ResourceNotFoundException` ("This model version has reached the end of
22// its life") and a profile-only identifier invoked bare with
23// `ValidationException` ("Invocation of model ID … with on-demand throughput
24// isn't supported"), so an identifier that fails either check is not shipped.
25//
26// A `us.` prefix marks a *cross-region inference profile*, which is the only
27// invocable form for the models that carry one. The prefix names a region
28// family: callers in Europe or Asia-Pacific substitute `eu.` or `apac.` (a
29// `global.` profile also exists for some Anthropic models). See
30// <https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles-support.html>.
31
32/// `amazon.nova-lite-v1:0`
33pub const AMAZON_NOVA_LITE: &str = "amazon.nova-lite-v1:0";
34/// `amazon.nova-micro-v1:0`
35pub const AMAZON_NOVA_MICRO: &str = "amazon.nova-micro-v1:0";
36/// `amazon.nova-pro-v1:0`
37pub const AMAZON_NOVA_PRO: &str = "amazon.nova-pro-v1:0";
38/// `amazon.nova-canvas-v1:0` image generation model
39pub const AMAZON_NOVA_CANVAS: &str = "amazon.nova-canvas-v1:0";
40/// `amazon.nova-reel-v1:0` video generation model
41pub const AMAZON_NOVA_REEL_V1_0: &str = "amazon.nova-reel-v1:0";
42/// `amazon.nova-reel-v1:1` video generation model
43pub const AMAZON_NOVA_REEL_V1_1: &str = "amazon.nova-reel-v1:1";
44/// `amazon.nova-sonic-v1:0` speech model
45pub const AMAZON_NOVA_SONIC: &str = "amazon.nova-sonic-v1:0";
46/// `amazon.rerank-v1:0` rerank model
47pub const AMAZON_RERANK_1_0: &str = "amazon.rerank-v1:0";
48/// `amazon.titan-embed-text-v1` embedding model
49pub const AMAZON_TITAN_EMBEDDINGS_G1_TEXT: &str = "amazon.titan-embed-text-v1";
50/// `amazon.titan-embed-image-v1` multimodal embedding model
51pub const AMAZON_TITAN_MULTIMODAL_EMBEDDINGS_G1: &str = "amazon.titan-embed-image-v1";
52/// `amazon.titan-embed-text-v2:0` embedding model
53pub const AMAZON_TITAN_TEXT_EMBEDDINGS_V2: &str = "amazon.titan-embed-text-v2:0";
54
55/// `us.anthropic.claude-haiku-4-5-20251001-v1:0` (cross-region profile)
56pub const ANTHROPIC_CLAUDE_HAIKU_4_5: &str = "us.anthropic.claude-haiku-4-5-20251001-v1:0";
57/// `us.anthropic.claude-sonnet-4-5-20250929-v1:0` (cross-region profile)
58pub const ANTHROPIC_CLAUDE_SONNET_4_5: &str = "us.anthropic.claude-sonnet-4-5-20250929-v1:0";
59/// `us.anthropic.claude-opus-4-5-20251101-v1:0` (cross-region profile)
60pub const ANTHROPIC_CLAUDE_OPUS_4_5: &str = "us.anthropic.claude-opus-4-5-20251101-v1:0";
61/// `us.anthropic.claude-sonnet-4-6` (cross-region profile)
62pub const ANTHROPIC_CLAUDE_SONNET_4_6: &str = "us.anthropic.claude-sonnet-4-6";
63/// `us.anthropic.claude-sonnet-5` (cross-region profile)
64pub const ANTHROPIC_CLAUDE_SONNET_5: &str = "us.anthropic.claude-sonnet-5";
65/// `us.anthropic.claude-opus-5` (cross-region profile)
66pub const ANTHROPIC_CLAUDE_OPUS_5: &str = "us.anthropic.claude-opus-5";
67
68/// `cohere.embed-english-v3` embedding model
69pub const COHERE_EMBED_ENGLISH: &str = "cohere.embed-english-v3";
70/// `cohere.embed-multilingual-v3` embedding model
71pub const COHERE_EMBED_MULTILINGUAL: &str = "cohere.embed-multilingual-v3";
72/// `cohere.rerank-v3-5:0` rerank model
73pub const COHERE_RERANK_V3_5: &str = "cohere.rerank-v3-5:0";
74
75/// `us.deepseek.r1-v1:0` (cross-region profile)
76pub const DEEPSEEK_R1: &str = "us.deepseek.r1-v1:0";
77
78/// `luma.ray-v2:0` video generation model
79pub const LUMA_RAY_V2_0: &str = "luma.ray-v2:0";
80
81/// `meta.llama3-8b-instruct-v1:0`
82pub const LLAMA_3_8B_INSTRUCT: &str = "meta.llama3-8b-instruct-v1:0";
83/// `meta.llama3-70b-instruct-v1:0`
84pub const LLAMA_3_70B_INSTRUCT: &str = "meta.llama3-70b-instruct-v1:0";
85/// `meta.llama3-1-8b-instruct-v1:0`
86pub const LLAMA_3_1_8B_INSTRUCT: &str = "meta.llama3-1-8b-instruct-v1:0";
87/// `meta.llama3-1-70b-instruct-v1:0`
88pub const LLAMA_3_1_70B_INSTRUCT: &str = "meta.llama3-1-70b-instruct-v1:0";
89/// `us.meta.llama3-3-70b-instruct-v1:0` (cross-region profile)
90pub const META_LLAMA_3_3_70B_INSTRUCT: &str = "us.meta.llama3-3-70b-instruct-v1:0";
91/// `us.meta.llama4-maverick-17b-instruct-v1:0` (cross-region profile)
92pub const META_LLAMA_4_MAVERICK_17B_INSTRUCT: &str = "us.meta.llama4-maverick-17b-instruct-v1:0";
93/// `us.meta.llama4-scout-17b-instruct-v1:0` (cross-region profile)
94pub const META_LLAMA_4_SCOUT_17B_INSTRUCT: &str = "us.meta.llama4-scout-17b-instruct-v1:0";
95
96/// `mistral.mistral-7b-instruct-v0:2`
97pub const MISTRAL_7B_INSTRUCT: &str = "mistral.mistral-7b-instruct-v0:2";
98/// `mistral.mistral-large-2402-v1:0`
99pub const MISTRAL_LARGE_24_02: &str = "mistral.mistral-large-2402-v1:0";
100/// `mistral.mistral-small-2402-v1:0`
101pub const MISTRAL_SMALL_24_02: &str = "mistral.mistral-small-2402-v1:0";
102/// `mistral.mixtral-8x7b-instruct-v0:1`
103pub const MISTRAL_MIXTRAL_8X7B_INSTRUCT_V0: &str = "mistral.mixtral-8x7b-instruct-v0:1";
104/// `us.mistral.pixtral-large-2502-v1:0` (cross-region profile)
105pub const MISTRAL_PIXTRAL_LARGE_2502: &str = "us.mistral.pixtral-large-2502-v1:0";
106
107/// `stability.sd3-5-large-v1:0` image generation model
108pub const STABILITY_SD3_5_LARGE: &str = "stability.sd3-5-large-v1:0";
109/// `stability.stable-image-core-v1:1` image generation model
110pub const STABILITY_STABLE_IMAGE_CORE_1_0: &str = "stability.stable-image-core-v1:1";
111/// `stability.stable-image-ultra-v1:1` image generation model
112pub const STABILITY_STABLE_IMAGE_ULTRA_1_0: &str = "stability.stable-image-ultra-v1:1";
113
114/// `twelvelabs.pegasus-1-2-v1:0` video-understanding model
115pub const TWELVELABS_PEGASUS_V1_2: &str = "twelvelabs.pegasus-1-2-v1:0";
116
117/// `us.writer.palmyra-x4-v1:0` (cross-region profile)
118pub const WRITER_PALMYRA_X4: &str = "us.writer.palmyra-x4-v1:0";
119/// `us.writer.palmyra-x5-v1:0` (cross-region profile)
120pub const WRITER_PALMYRA_X5: &str = "us.writer.palmyra-x5-v1:0";
121
122#[derive(Clone)]
123pub struct CompletionModel {
124    pub(crate) client: Client,
125    pub model: String,
126    /// When enabled, cache checkpoints are inserted into Converse API requests
127    /// to take advantage of [Bedrock prompt caching](https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html).
128    /// A checkpoint is placed after the system prompt and after the last message
129    /// in the chat history. Disabled by default.
130    pub prompt_caching: bool,
131    /// Guardrail applied to every Converse request from this model, if any.
132    /// Set through [`CompletionModel::with_guardrail`].
133    pub guardrail: Option<aws_bedrock::GuardrailConfiguration>,
134}
135
136impl CompletionModel {
137    pub fn new(client: Client, model: impl Into<String>) -> Self {
138        Self {
139            client,
140            model: model.into(),
141            prompt_caching: false,
142            guardrail: None,
143        }
144    }
145
146    /// Enable Bedrock prompt caching for this model.
147    ///
148    /// When enabled, `CachePoint` blocks are inserted after the serialized
149    /// `system` content and after the final `messages` entry in each Converse
150    /// API request. This allows Bedrock to cache and reuse repeated prompt
151    /// prefixes, reducing both latency and input token costs.
152    ///
153    /// This currently covers the `system` and `messages` request fields only.
154    /// Some Bedrock models also support caching `tools`, but that is not wired
155    /// up here yet.
156    ///
157    /// Cacheability and token thresholds are model-specific. If the cached
158    /// prefix is too short or the model does not support caching for a given
159    /// field, Bedrock ignores the checkpoint. See the Bedrock prompt caching
160    /// support table for current limits and field support.
161    pub fn with_prompt_caching(mut self) -> Self {
162        self.prompt_caching = true;
163        self
164    }
165
166    /// Apply a [Bedrock guardrail](https://docs.aws.amazon.com/bedrock/latest/userguide/guardrails.html)
167    /// to every Converse request this model issues.
168    ///
169    /// `identifier` is the guardrail id or ARN and `version` its version (or
170    /// `DRAFT`). When `trace` is enabled, Bedrock returns its assessment on
171    /// [`InternalConverseOutput::trace`](crate::types::converse_output::InternalConverseOutput::trace) —
172    /// the only place it explains *why* a turn came back with
173    /// [`StopReason::GuardrailIntervened`](crate::types::converse_output::StopReason::GuardrailIntervened),
174    /// which the normalized response reports as a content filter and nothing
175    /// more. Reach it through
176    /// [`raw_completion`](CompletionModel::raw_completion).
177    pub fn with_guardrail(
178        mut self,
179        identifier: impl Into<String>,
180        version: impl Into<String>,
181        trace: aws_bedrock::GuardrailTrace,
182    ) -> Self {
183        self.guardrail = Some(
184            aws_bedrock::GuardrailConfiguration::builder()
185                .guardrail_identifier(identifier)
186                .guardrail_version(version)
187                .trace(trace)
188                .build(),
189        );
190        self
191    }
192}
193
194pub(crate) fn resolve_request_model(
195    default_model: &str,
196    completion_request: &CompletionRequest,
197) -> String {
198    completion_request
199        .model
200        .clone()
201        .unwrap_or_else(|| default_model.to_string())
202}
203
204impl CompletionModel {
205    /// Execute a completion and return Bedrock's own Converse output.
206    ///
207    /// This is the escape hatch for fields rig does not normalize;
208    /// [`completion::CompletionModel::completion`] calls it and maps the
209    /// result, so there is exactly one request either way.
210    pub async fn raw_completion(
211        &self,
212        completion_request: completion::CompletionRequest,
213    ) -> Result<AwsConverseOutput, CompletionError> {
214        let request_model = resolve_request_model(&self.model, &completion_request);
215
216        let span =
217            CompletionSpanBuilder::new("aws_bedrock", &request_model, CompletionOperation::Chat)
218                .system_instructions(
219                    completion_request.preamble.as_deref(),
220                    completion_request.record_telemetry_content,
221                )
222                .build();
223
224        let request = AwsCompletionRequest {
225            inner: completion_request,
226            prompt_caching: self.prompt_caching,
227        };
228
229        let mut converse_builder = self
230            .client
231            .get_inner()
232            .await
233            .converse()
234            .model_id(request_model.clone());
235
236        let tool_config = request.tools_config()?;
237        let messages = request.messages()?;
238        let output_config = request.output_config()?;
239        converse_builder = converse_builder
240            .set_additional_model_request_fields(request.additional_params())
241            .set_inference_config(request.inference_config())
242            .set_tool_config(tool_config)
243            .set_system(request.system_prompt()?)
244            .set_messages(Some(messages))
245            .set_output_config(output_config)
246            .set_guardrail_config(self.guardrail.clone());
247
248        async move {
249            let response = converse_builder.send().await.map_err(|sdk_error| {
250                Into::<CompletionError>::into(AwsSdkConverseError(sdk_error))
251            })?;
252
253            let response: InternalConverseOutput = response.try_into().map_err(|x| {
254                CompletionError::ProviderError(format!("Type conversion error: {x}"))
255            })?;
256
257            let aws_output = AwsConverseOutput(response);
258
259            let span = tracing::Span::current();
260            span.record_response_metadata(&aws_output);
261            span.record_token_usage(&aws_output.get_usage().unwrap_or_default());
262
263            Ok(aws_output)
264        }
265        .instrument(span)
266        .await
267    }
268}
269
270impl completion::CompletionModel for CompletionModel {
271    async fn completion(
272        &self,
273        completion_request: completion::CompletionRequest,
274    ) -> Result<completion::CompletionResponse, CompletionError> {
275        // Capture before `try_into` consumes the raw value.
276        let raw = self.raw_completion(completion_request).await?;
277        let captured = serde_json::to_value(&raw)?;
278        let response: completion::CompletionResponse = raw.try_into()?;
279        Ok(response.with_raw(captured))
280    }
281
282    async fn stream(
283        &self,
284        request: CompletionRequest,
285    ) -> Result<StreamingCompletionResponse, CompletionError> {
286        CompletionModel::stream(self, request).await
287    }
288}
rig_bedrock/completion.rs

rig_bedrock/
completion.rs
