Add Bedrock integration tests (0xPlaygrounds#1707)

* Add Bedrock integration tests

* tests

* tests

Author: gold-silver-copper

tests/integrations.rs

@@ -6,6 +6,9 @@
     clippy::unreachable
 )]
 
+#[cfg(feature = "bedrock")]
+#[path = "integrations/bedrock/mod.rs"]
+mod bedrock;
 #[cfg(feature = "lancedb")]
 #[path = "integrations/lancedb/mod.rs"]
 mod lancedb;

tests/integrations/bedrock/adaptive_thinking.rs

@@ -0,0 +1,95 @@
+//! Live Bedrock Anthropic adaptive-thinking regression tests.
+
+use futures::StreamExt;
+use rig::agent::AgentBuilder;
+use rig::client::CompletionClient;
+use rig::completion::{CompletionModel as _, Prompt};
+use rig::streaming::StreamedAssistantContent;
+use serde_json::json;
+
+use super::{
+    anthropic_adaptive_model, anthropic_signature_only_model, client,
+    support::{ALPHA_SIGNAL_OUTPUT, AlphaSignal, assert_contains_all_case_insensitive},
+};
+
+fn adaptive_thinking_params() -> serde_json::Value {
+    json!({
+        "thinking": {
+            "type": "adaptive"
+        }
+    })
+}
+
+#[tokio::test]
+#[ignore = "requires AWS credentials and Bedrock Anthropic adaptive-thinking model access"]
+async fn adaptive_thinking_prompt_caching_tool_roundtrip_regression() {
+    let model = client()
+        .completion_model(anthropic_adaptive_model())
+        .with_prompt_caching();
+    let agent = AgentBuilder::new(model)
+        .preamble(
+            "You must call tools when the user asks for their result. \
+             After a tool result is available, answer with the exact result.",
+        )
+        .max_tokens(2048)
+        .additional_params(adaptive_thinking_params())
+        .tool(AlphaSignal)
+        .build();
+
+    let response = agent
+        .prompt("Call `lookup_harbor_label` exactly once, then answer with the exact tool output.")
+        .await
+        .expect("adaptive-thinking prompt-caching tool roundtrip should succeed");
+
+    assert_contains_all_case_insensitive(&response, &[ALPHA_SIGNAL_OUTPUT]);
+}
+
+#[tokio::test]
+#[ignore = "requires AWS credentials and Bedrock Anthropic adaptive-thinking model access"]
+async fn streaming_emits_signature_only_adaptive_reasoning_regression() {
+    let model = client().completion_model(anthropic_signature_only_model());
+    let request = model
+        .completion_request("What is 2 + 2? Answer with only the number.")
+        .max_tokens(2048)
+        .additional_params(adaptive_thinking_params())
+        .build();
+    let mut stream = model
+        .stream(request)
+        .await
+        .expect("adaptive-thinking Bedrock stream should start");
+
+    let mut reasoning_chunks = 0;
+    let mut signature_chunks = 0;
+    let mut signature_only_chunks = 0;
+    let mut got_final = false;
+
+    while let Some(item) = stream.next().await {
+        match item.expect("adaptive-thinking Bedrock stream item should succeed") {
+            StreamedAssistantContent::Reasoning(reasoning) => {
+                reasoning_chunks += 1;
+                if reasoning.first_signature().is_some() {
+                    signature_chunks += 1;
+                    if reasoning.display_text().is_empty() {
+                        signature_only_chunks += 1;
+                    }
+                }
+            }
+            StreamedAssistantContent::Final(_) => got_final = true,
+            _ => {}
+        }
+    }
+
+    assert!(got_final, "stream should emit a final response");
+    assert!(
+        reasoning_chunks > 0,
+        "expected at least one adaptive-thinking reasoning chunk"
+    );
+    assert!(
+        signature_chunks > 0,
+        "expected adaptive-thinking reasoning to include a Bedrock signature"
+    );
+    assert!(
+        signature_only_chunks > 0,
+        "expected at least one signature-only reasoning chunk"
+    );
+}

tests/integrations/bedrock/agent.rs

@@ -0,0 +1,86 @@
+//! AWS Bedrock completion smoke tests inspired by the OpenAI and Anthropic provider tests.
+
+use rig::agent::AgentBuilder;
+use rig::client::CompletionClient;
+use rig::completion::Prompt;
+
+use super::{
+    BEDROCK_COMPLETION_MODEL, client,
+    support::{
+        Adder, BASIC_PREAMBLE, BASIC_PROMPT, CONTEXT_DOCS, CONTEXT_PROMPT,
+        STREAMING_TOOLS_PREAMBLE, STREAMING_TOOLS_PROMPT, Subtract,
+        assert_contains_any_case_insensitive, assert_mentions_expected_number,
+        assert_nonempty_response,
+    },
+};
+
+#[tokio::test]
+#[ignore = "requires AWS credentials and Bedrock model access"]
+async fn completion_smoke() {
+    let agent = client()
+        .agent(BEDROCK_COMPLETION_MODEL)
+        .preamble(BASIC_PREAMBLE)
+        .build();
+
+    let response = agent
+        .prompt(BASIC_PROMPT)
+        .await
+        .expect("completion should succeed");
+
+    assert_nonempty_response(&response);
+}
+
+#[tokio::test]
+#[ignore = "requires AWS credentials and Bedrock model access"]
+async fn completion_with_context_smoke() {
+    let agent = client()
+        .agent(BEDROCK_COMPLETION_MODEL)
+        .preamble("Answer the user using only the supplied context.")
+        .context(CONTEXT_DOCS[0])
+        .context(CONTEXT_DOCS[1])
+        .context(CONTEXT_DOCS[2])
+        .build();
+
+    let response = agent
+        .prompt(CONTEXT_PROMPT)
+        .await
+        .expect("context completion should succeed");
+
+    assert_contains_any_case_insensitive(&response, &["ancient tool", "farm"]);
+}
+
+#[tokio::test]
+#[ignore = "requires AWS credentials and Bedrock model access"]
+async fn tool_roundtrip_smoke() {
+    let agent = client()
+        .agent(BEDROCK_COMPLETION_MODEL)
+        .preamble(STREAMING_TOOLS_PREAMBLE)
+        .max_tokens(1024)
+        .tool(Adder)
+        .tool(Subtract)
+        .build();
+
+    let response = agent
+        .prompt(STREAMING_TOOLS_PROMPT)
+        .await
+        .expect("tool prompt should succeed");
+
+    assert_mentions_expected_number(&response, -3);
+}
+
+#[tokio::test]
+#[ignore = "requires AWS credentials and Bedrock model access"]
+async fn prompt_caching_completion_smoke() {
+    let bedrock_client = client();
+    let model = bedrock_client
+        .completion_model(BEDROCK_COMPLETION_MODEL)
+        .with_prompt_caching();
+    let agent = AgentBuilder::new(model).preamble(BASIC_PREAMBLE).build();
+
+    let response = agent
+        .prompt(BASIC_PROMPT)
+        .await
+        .expect("prompt-caching completion should succeed");
+
+    assert_nonempty_response(&response);
+}

tests/integrations/bedrock/documents.rs

@@ -0,0 +1,74 @@
+//! AWS Bedrock document prompt smoke tests inspired by Anthropic document tests.
+
+use rig::OneOrMany;
+use rig::client::CompletionClient;
+use rig::completion::Prompt;
+use rig::message::{Document, DocumentMediaType, DocumentSourceKind, Message, UserContent};
+
+use super::{
+    BEDROCK_COMPLETION_MODEL, client,
+    support::{assert_contains_any_case_insensitive, assert_nonempty_response},
+};
+
+fn rust_document() -> String {
+    r#"
+The Rust Programming Language
+
+Rust is a systems programming language focused on three goals: safety, speed,
+and concurrency. It accomplishes these goals without a garbage collector.
+
+Key Features:
+- Zero-cost abstractions
+- Move semantics
+- Guaranteed memory safety
+- Threads without data races
+"#
+    .trim()
+    .to_string()
+}
+
+#[tokio::test]
+#[ignore = "requires AWS credentials and Bedrock model access"]
+async fn plaintext_document_prompt() {
+    let agent = client()
+        .agent(BEDROCK_COMPLETION_MODEL)
+        .preamble("Summarize the provided document.")
+        .temperature(0.5)
+        .build();
+
+    let document = Document {
+        data: DocumentSourceKind::String(rust_document()),
+        media_type: Some(DocumentMediaType::TXT),
+        additional_params: None,
+    };
+    let response = agent
+        .prompt(document)
+        .await
+        .expect("document prompt should succeed");
+
+    assert_nonempty_response(&response);
+    assert_contains_any_case_insensitive(&response, &["safety", "speed", "concurrency"]);
+}
+
+#[tokio::test]
+#[ignore = "requires AWS credentials and Bedrock model access"]
+async fn plaintext_document_with_instruction() {
+    let agent = client()
+        .agent(BEDROCK_COMPLETION_MODEL)
+        .preamble("Answer from the provided document.")
+        .temperature(0.5)
+        .build();
+
+    let response = agent
+        .prompt(Message::User {
+            content: OneOrMany::many(vec![
+                UserContent::document(rust_document(), Some(DocumentMediaType::TXT)),
+                UserContent::text("List the three main goals of Rust mentioned in this document."),
+            ])
+            .expect("content should be non-empty"),
+        })
+        .await
+        .expect("instruction prompt should succeed");
+
+    assert_contains_any_case_insensitive(&response, &["safety", "speed", "concurrency"]);
+}

tests/integrations/bedrock/embeddings.rs

@@ -0,0 +1,27 @@
+//! AWS Bedrock embedding smoke test inspired by provider embedding coverage.
+
+use rig::client::EmbeddingsClient;
+use rig::embeddings::EmbeddingModel as _;
+
+use super::{
+    BEDROCK_EMBEDDING_MODEL, client,
+    support::{EMBEDDING_INPUTS, assert_embeddings_nonempty_and_consistent},
+};
+
+#[tokio::test]
+#[ignore = "requires AWS credentials and Bedrock embedding model access"]
+async fn embeddings_smoke() {
+    let model = client().embedding_model_with_ndims(BEDROCK_EMBEDDING_MODEL, 256);
+    let embeddings = model
+        .embed_texts(EMBEDDING_INPUTS.into_iter().map(str::to_string))
+        .await
+        .expect("embedding request should succeed");
+
+    assert_embeddings_nonempty_and_consistent(&embeddings, EMBEDDING_INPUTS.len());
+    assert!(
+        embeddings
+            .iter()
+            .all(|embedding| embedding.vec.len() == 256),
+        "Titan text embeddings v2 should return the requested 256 dimensions"
+    );
+}

tests/integrations/bedrock/extractor.rs

@@ -0,0 +1,62 @@
+//! AWS Bedrock extractor smoke tests inspired by the provider extractor tests.
+
+use rig::client::CompletionClient;
+use rig::message::Message;
+
+use super::{
+    BEDROCK_COMPLETION_MODEL, client,
+    support::{EXTRACTOR_TEXT, SmokePerson, assert_nonempty_response},
+};
+
+fn assert_smoke_person(person: &SmokePerson) {
+    let first_name = person
+        .first_name
+        .as_deref()
+        .expect("first_name should be present");
+    let last_name = person
+        .last_name
+        .as_deref()
+        .expect("last_name should be present");
+    let job = person.job.as_deref().expect("job should be present");
+
+    assert_nonempty_response(first_name);
+    assert_nonempty_response(last_name);
+    assert_nonempty_response(job);
+}
+
+#[tokio::test]
+#[ignore = "requires AWS credentials and Bedrock model access"]
+async fn extractor_smoke() {
+    let extractor = client()
+        .extractor::<SmokePerson>(BEDROCK_COMPLETION_MODEL)
+        .build();
+
+    let response = extractor
+        .extract_with_usage(EXTRACTOR_TEXT)
+        .await
+        .expect("extractor request should succeed");
+
+    assert_smoke_person(&response.data);
+    assert!(response.usage.total_tokens > 0, "usage should be populated");
+}
+
+#[tokio::test]
+#[ignore = "requires AWS credentials and Bedrock model access"]
+async fn extractor_with_chat_history_smoke() {
+    let extractor = client()
+        .extractor::<SmokePerson>(BEDROCK_COMPLETION_MODEL)
+        .build();
+
+    let response = extractor
+        .extract_with_chat_history_with_usage(
+            "The text is about Ada Lovelace, a mathematician.",
+            vec![Message::user(
+                "Extract the person's name and job from the next message.",
+            )],
+        )
+        .await
+        .expect("extractor request with chat history should succeed");
+
+    assert_smoke_person(&response.data);
+    assert!(response.usage.total_tokens > 0, "usage should be populated");
+}

tests/integrations/bedrock/image_generation.rs

@@ -0,0 +1,25 @@
+//! AWS Bedrock image generation smoke test inspired by OpenAI image generation tests.
+
+use rig::image_generation::ImageGenerationModel;
+use rig::prelude::ImageGenerationClient;
+
+use super::{
+    BEDROCK_IMAGE_MODEL, client,
+    support::{IMAGE_PROMPT, assert_nonempty_bytes},
+};
+
+#[tokio::test]
+#[ignore = "requires AWS credentials and Bedrock image generation model access"]
+async fn image_generation_smoke() {
+    let model = client().image_generation_model(BEDROCK_IMAGE_MODEL);
+    let response = model
+        .image_generation_request()
+        .prompt(IMAGE_PROMPT)
+        .width(512)
+        .height(512)
+        .send()
+        .await
+        .expect("image generation request should succeed");
+
+    assert_nonempty_bytes(&response.image);
+}