tenant.openagents/omega

1mod models;
2
3use anyhow::{Result, anyhow};
4use aws_sdk_bedrockruntime as bedrock;
5pub use aws_sdk_bedrockruntime as bedrock_client;
6pub use aws_sdk_bedrockruntime::types::{
7    AnyToolChoice as BedrockAnyToolChoice, AutoToolChoice as BedrockAutoToolChoice,
8    ContentBlock as BedrockInnerContent, Tool as BedrockTool, ToolChoice as BedrockToolChoice,
9    ToolConfiguration as BedrockToolConfig, ToolInputSchema as BedrockToolInputSchema,
10    ToolSpecification as BedrockToolSpec,
11};
12use aws_sdk_bedrockruntime::types::{GuardrailStreamConfiguration, InferenceConfiguration};
13pub use aws_smithy_types::Blob as BedrockBlob;
14use aws_smithy_types::{Document, Number as AwsNumber};
15pub use bedrock::operation::converse_stream::ConverseStreamInput as BedrockStreamingRequest;
16pub use bedrock::types::{
17    ContentBlock as BedrockRequestContent, ConversationRole as BedrockRole,
18    ConverseOutput as BedrockResponse, ConverseStreamOutput as BedrockStreamingResponse,
19    ImageBlock as BedrockImageBlock, ImageFormat as BedrockImageFormat,
20    ImageSource as BedrockImageSource, Message as BedrockMessage,
21    ReasoningContentBlock as BedrockThinkingBlock, ReasoningTextBlock as BedrockThinkingTextBlock,
22    ResponseStream as BedrockResponseStream, SystemContentBlock as BedrockSystemContentBlock,
23    ToolResultBlock as BedrockToolResultBlock,
24    ToolResultContentBlock as BedrockToolResultContentBlock,
25    ToolResultStatus as BedrockToolResultStatus, ToolUseBlock as BedrockToolUseBlock,
26};
27use futures::stream::{self, BoxStream};
28use serde::{Deserialize, Serialize};
29use serde_json::{Number, Value};
30use std::collections::HashMap;
31use thiserror::Error;
32
33pub use crate::models::*;
34
35pub async fn stream_completion(
36    client: bedrock::Client,
37    request: Request,
38    extra_headers: http_client::CustomHeaders,
39) -> Result<BoxStream<'static, Result<BedrockStreamingResponse, anyhow::Error>>, BedrockError> {
40    let mut response = bedrock::Client::converse_stream(&client)
41        .model_id(request.model.clone())
42        .set_messages(request.messages.into());
43
44    let mut additional_fields: HashMap<String, Document> = HashMap::new();
45
46    if let Some(thinking) = &request.thinking {
47        additional_fields.extend(thinking_request_fields(thinking));
48    }
49
50    if !additional_fields.is_empty() {
51        response = response.additional_model_request_fields(Document::Object(additional_fields));
52    }
53
54    if request.tools.as_ref().is_some_and(|t| !t.tools.is_empty()) {
55        response = response.set_tool_config(request.tools);
56    }
57
58    let inference_config = InferenceConfiguration::builder()
59        .max_tokens(request.max_tokens as i32)
60        .set_temperature(request.temperature)
61        .set_top_p(request.top_p)
62        .build();
63
64    response = response.inference_config(inference_config);
65
66    for system_block in request.system {
67        response = response.system(system_block);
68    }
69
70    if let Some(guardrail_id) = &request.guardrail_identifier {
71        let version = request.guardrail_version.as_deref().unwrap_or("DRAFT");
72
73        response = response.guardrail_config(
74            GuardrailStreamConfiguration::builder()
75                .guardrail_identifier(guardrail_id)
76                .guardrail_version(version)
77                .build(),
78        );
79    }
80
81    let output = response
82        .customize()
83        .mutate_request(move |http_request| {
84            let headers = http_request.headers_mut();
85            for (name, value) in extra_headers.iter() {
86                headers.insert(
87                    name.as_str().to_owned(),
88                    value.to_str().unwrap_or("").to_owned(),
89                );
90            }
91        })
92        .send()
93        .await
94        .map_err(|err| match err {
95            bedrock::error::SdkError::ServiceError(ctx) => {
96                use bedrock::operation::converse_stream::ConverseStreamError;
97                let err = ctx.into_err();
98                match &err {
99                    ConverseStreamError::ValidationException(e) => BedrockError::Validation(
100                        e.message().unwrap_or("validation error").to_string(),
101                    ),
102                    ConverseStreamError::ThrottlingException(_) => BedrockError::RateLimited,
103                    ConverseStreamError::ServiceUnavailableException(_)
104                    | ConverseStreamError::ModelNotReadyException(_) => {
105                        BedrockError::ServiceUnavailable
106                    }
107                    ConverseStreamError::AccessDeniedException(e) => BedrockError::AccessDenied(
108                        e.message().unwrap_or("access denied").to_string(),
109                    ),
110                    ConverseStreamError::InternalServerException(e) => {
111                        BedrockError::InternalServer(
112                            e.message().unwrap_or("internal server error").to_string(),
113                        )
114                    }
115                    _ => BedrockError::Other(err.into()),
116                }
117            }
118            other => BedrockError::Other(other.into()),
119        });
120
121    let stream = Box::pin(stream::unfold(
122        output?.stream,
123        move |mut stream| async move {
124            match stream.recv().await {
125                Ok(Some(output)) => Some((Ok(output), stream)),
126                Ok(None) => None,
127                Err(err) => Some((
128                    Err(anyhow!(
129                        "{}",
130                        aws_sdk_bedrockruntime::error::DisplayErrorContext(err)
131                    )),
132                    stream,
133                )),
134            }
135        },
136    ));
137
138    Ok(stream)
139}
140
141pub fn aws_document_to_value(document: &Document) -> Value {
142    match document {
143        Document::Null => Value::Null,
144        Document::Bool(value) => Value::Bool(*value),
145        Document::Number(value) => match *value {
146            AwsNumber::PosInt(value) => Value::Number(Number::from(value)),
147            AwsNumber::NegInt(value) => Value::Number(Number::from(value)),
148            AwsNumber::Float(value) => Value::Number(Number::from_f64(value).unwrap()),
149        },
150        Document::String(value) => Value::String(value.clone()),
151        Document::Array(array) => Value::Array(array.iter().map(aws_document_to_value).collect()),
152        Document::Object(map) => Value::Object(
153            map.iter()
154                .map(|(key, value)| (key.clone(), aws_document_to_value(value)))
155                .collect(),
156        ),
157    }
158}
159
160pub fn value_to_aws_document(value: &Value) -> Document {
161    match value {
162        Value::Null => Document::Null,
163        Value::Bool(value) => Document::Bool(*value),
164        Value::Number(value) => {
165            if let Some(value) = value.as_u64() {
166                Document::Number(AwsNumber::PosInt(value))
167            } else if let Some(value) = value.as_i64() {
168                Document::Number(AwsNumber::NegInt(value))
169            } else if let Some(value) = value.as_f64() {
170                Document::Number(AwsNumber::Float(value))
171            } else {
172                Document::Null
173            }
174        }
175        Value::String(value) => Document::String(value.clone()),
176        Value::Array(array) => Document::Array(array.iter().map(value_to_aws_document).collect()),
177        Value::Object(map) => Document::Object(
178            map.iter()
179                .map(|(key, value)| (key.clone(), value_to_aws_document(value)))
180                .collect(),
181        ),
182    }
183}
184
185#[derive(Debug, Serialize, Deserialize)]
186pub enum Thinking {
187    Enabled {
188        budget_tokens: Option<u64>,
189    },
190    Adaptive {
191        effort: BedrockAdaptiveThinkingEffort,
192    },
193    /// Explicitly turns thinking off. Required by Claude Opus 5, where
194    /// adaptive thinking runs by default when the `thinking` field is
195    /// omitted; only accepted at effort `high` or below.
196    ///
197    /// <https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-opus-5.html>
198    Disabled,
199}
200
201/// Converts the request's thinking configuration into the
202/// `additionalModelRequestFields` entries understood by Anthropic models on
203/// the Converse API.
204fn thinking_request_fields(thinking: &Thinking) -> HashMap<String, Document> {
205    let mut fields = HashMap::new();
206    match thinking {
207        Thinking::Enabled {
208            budget_tokens: Some(budget_tokens),
209        } => {
210            fields.insert(
211                "thinking".to_string(),
212                Document::from(HashMap::from([
213                    ("type".to_string(), Document::String("enabled".to_string())),
214                    (
215                        "budget_tokens".to_string(),
216                        Document::Number(AwsNumber::PosInt(*budget_tokens)),
217                    ),
218                ])),
219            );
220        }
221        Thinking::Enabled {
222            budget_tokens: None,
223        } => {}
224        Thinking::Adaptive { effort: _ } => {
225            fields.insert(
226                "thinking".to_string(),
227                Document::from(HashMap::from([
228                    ("type".to_string(), Document::String("adaptive".to_string())),
229                    (
230                        "display".to_string(),
231                        Document::String("summarized".to_string()),
232                    ),
233                ])),
234            );
235        }
236        Thinking::Disabled => {
237            // On Claude Opus 5 omitting the `thinking` field means adaptive
238            // thinking runs by default, so turning it off requires this
239            // explicit opt-out. No effort is attached: `disabled` combined
240            // with effort `xhigh`/`max` is rejected with a 400.
241            fields.insert(
242                "thinking".to_string(),
243                Document::from(HashMap::from([(
244                    "type".to_string(),
245                    Document::String("disabled".to_string()),
246                )])),
247            );
248        }
249    }
250    fields
251}
252
253#[derive(Debug)]
254pub struct Request {
255    pub model: String,
256    pub max_tokens: u64,
257    pub messages: Vec<BedrockMessage>,
258    pub tools: Option<BedrockToolConfig>,
259    pub thinking: Option<Thinking>,
260    /// System content blocks in prefix order. Typically `[Text(...)]` or, when
261    /// the model supports prompt caching, `[Text(...), CachePoint(...)]` so the
262    /// system prompt anchors its own cache prefix independent of tools and
263    /// messages.
264    pub system: Vec<BedrockSystemContentBlock>,
265    pub metadata: Option<Metadata>,
266    pub stop_sequences: Vec<String>,
267    pub temperature: Option<f32>,
268    pub top_k: Option<u32>,
269    pub top_p: Option<f32>,
270    pub guardrail_identifier: Option<String>,
271    pub guardrail_version: Option<String>,
272}
273
274#[derive(Debug, Serialize, Deserialize)]
275pub struct Metadata {
276    pub user_id: Option<String>,
277}
278
279#[derive(Error, Debug)]
280pub enum BedrockError {
281    #[error("{0}")]
282    Validation(String),
283    #[error("rate limited")]
284    RateLimited,
285    #[error("service unavailable")]
286    ServiceUnavailable,
287    #[error("{0}")]
288    AccessDenied(String),
289    #[error("{0}")]
290    InternalServer(String),
291    #[error(transparent)]
292    Other(#[from] anyhow::Error),
293}
294
295#[cfg(test)]
296mod tests {
297    use super::*;
298
299    fn string_field<'a>(document: &'a Document, key: &str) -> Option<&'a str> {
300        match document {
301            Document::Object(map) => match map.get(key) {
302                Some(Document::String(value)) => Some(value.as_str()),
303                _ => None,
304            },
305            _ => None,
306        }
307    }
308
309    #[test]
310    fn test_disabled_thinking_serializes_opt_out_without_effort() {
311        let fields = thinking_request_fields(&Thinking::Disabled);
312
313        let thinking = fields.get("thinking").expect("thinking field");
314        assert_eq!(string_field(thinking, "type"), Some("disabled"));
315        // `disabled` combined with an effort of `xhigh`/`max` is a 400, so no
316        // output_config may accompany the opt-out.
317        assert!(!fields.contains_key("output_config"));
318    }
319
320    #[test]
321    fn test_enabled_thinking_serializes_budget_tokens() {
322        let fields = thinking_request_fields(&Thinking::Enabled {
323            budget_tokens: Some(4_096),
324        });
325
326        let thinking = fields.get("thinking").expect("thinking field");
327        assert_eq!(string_field(thinking, "type"), Some("enabled"));
328        match thinking {
329            Document::Object(map) => assert_eq!(
330                map.get("budget_tokens"),
331                Some(&Document::Number(AwsNumber::PosInt(4_096)))
332            ),
333            _ => panic!("thinking field should be an object"),
334        }
335
336        let fields = thinking_request_fields(&Thinking::Enabled {
337            budget_tokens: None,
338        });
339        assert!(fields.is_empty());
340    }
341}
342