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}