Skip to main content

llm/catalog/
bedrock.rs

1use std::borrow::Cow;
2use std::str::FromStr;
3
4use crate::catalog::transport::ModelTransport;
5use crate::catalog::{BedrockFoundationModel, ModelPricing};
6use crate::{ReasoningDisabledSupport, ReasoningEffort};
7
8#[derive(Debug, Clone, PartialEq, Eq, Hash)]
9pub enum BedrockModel {
10    Foundation(BedrockFoundationModel),
11    Profile(String),
12}
13
14impl BedrockModel {
15    pub fn model_id(&self) -> Cow<'static, str> {
16        match self {
17            Self::Foundation(m) => Cow::Borrowed(m.model_id()),
18            Self::Profile(s) => Cow::Owned(s.clone()),
19        }
20    }
21
22    pub fn display_name(&self) -> Cow<'static, str> {
23        match self {
24            Self::Foundation(m) => Cow::Borrowed(m.display_name()),
25            Self::Profile(s) => Cow::Owned(format!("Bedrock {s}")),
26        }
27    }
28
29    pub fn context_window(&self) -> Option<u32> {
30        match self {
31            Self::Foundation(m) => Some(m.context_window()),
32            Self::Profile(_) => None,
33        }
34    }
35
36    pub fn reasoning_levels(&self) -> &'static [ReasoningEffort] {
37        match self {
38            Self::Foundation(m) => m.reasoning_levels(),
39            Self::Profile(_) => &[],
40        }
41    }
42
43    pub fn reasoning_disabled_support(&self) -> crate::reasoning::ReasoningDisabledSupport {
44        match self {
45            Self::Foundation(model) => model.reasoning_disabled_support(),
46            Self::Profile(_) => ReasoningDisabledSupport::Unsupported,
47        }
48    }
49
50    pub fn supports_reasoning(&self) -> bool {
51        self.reasoning_levels().iter().any(|effort| effort.is_enabled())
52    }
53
54    pub fn supports_prompt_caching(&self) -> bool {
55        match self {
56            Self::Foundation(m) => m.supports_prompt_caching(),
57            Self::Profile(_) => false,
58        }
59    }
60
61    pub fn supports_image(&self) -> bool {
62        match self {
63            Self::Foundation(m) => m.supports_image(),
64            Self::Profile(_) => false,
65        }
66    }
67
68    pub fn supports_audio(&self) -> bool {
69        match self {
70            Self::Foundation(m) => m.supports_audio(),
71            Self::Profile(_) => false,
72        }
73    }
74
75    pub fn pricing(&self) -> Option<ModelPricing> {
76        match self {
77            Self::Foundation(m) => m.pricing(),
78            Self::Profile(_) => None,
79        }
80    }
81
82    pub fn transport(&self) -> Option<ModelTransport> {
83        match self {
84            Self::Foundation(m) => m.transport(),
85            Self::Profile(_) => None,
86        }
87    }
88}
89
90impl FromStr for BedrockModel {
91    type Err = String;
92
93    fn from_str(s: &str) -> Result<Self, Self::Err> {
94        match s.parse::<BedrockFoundationModel>() {
95            Ok(m) => Ok(Self::Foundation(m)),
96            Err(_) if is_bedrock_inference_profile_arn(s) => Err(
97                "Bedrock inference profile ARNs must be configured as providers.bedrock.inferenceProfileArn; keep model as bedrock:<model-id>".to_string(),
98            ),
99            Err(_) => Ok(Self::Profile(s.to_string())),
100        }
101    }
102}
103
104fn is_bedrock_inference_profile_arn(s: &str) -> bool {
105    let Some(rest) = s.strip_prefix("arn:") else {
106        return false;
107    };
108    let parts: Vec<&str> = rest.split(':').collect();
109    matches!(
110        parts.as_slice(),
111        [partition, "bedrock", _, _, resource, ..]
112            if partition.starts_with("aws")
113                && (resource.starts_with("inference-profile/")
114                    || resource.starts_with("application-inference-profile/"))
115    )
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121
122    #[test]
123    fn foundation_model_parses() {
124        let model: BedrockModel = "anthropic.claude-sonnet-4-5-20250929-v1:0".parse().unwrap();
125        assert!(matches!(model, BedrockModel::Foundation(_)));
126    }
127
128    #[test]
129    fn unknown_profile_id_falls_through_to_profile_variant() {
130        let model: BedrockModel = "us.anthropic.claude-future-model-v99:0".parse().unwrap();
131        assert!(matches!(model, BedrockModel::Profile(_)));
132        assert_eq!(model.context_window(), None);
133    }
134
135    #[test]
136    fn inference_profile_arn_is_rejected() {
137        let error =
138            "arn:aws:bedrock:us-west-2:000000000000:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0"
139                .parse::<BedrockModel>()
140                .unwrap_err();
141        assert!(error.contains("providers.bedrock.inferenceProfileArn"));
142    }
143
144    #[test]
145    fn application_inference_profile_arn_is_rejected() {
146        let error = "arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/000000000000"
147            .parse::<BedrockModel>()
148            .unwrap_err();
149        assert!(error.contains("providers.bedrock.inferenceProfileArn"));
150    }
151
152    #[test]
153    fn gov_cloud_arn_is_rejected() {
154        let error = "arn:aws-us-gov:bedrock:us-gov-west-1:000000000000:application-inference-profile/000000000000"
155            .parse::<BedrockModel>()
156            .unwrap_err();
157        assert!(error.contains("providers.bedrock.inferenceProfileArn"));
158    }
159
160    #[test]
161    fn non_bedrock_arn_falls_through_to_profile() {
162        let model: BedrockModel = "arn:aws:s3:us-west-2:000000000000:bucket/foo".parse().unwrap();
163        assert!(matches!(model, BedrockModel::Profile(_)));
164    }
165}