feat: implement wasi_client for aws operations. Use reqwest based wasi async client

Author: iambenkay

Cargo.lock

@@ -16,6 +16,7 @@ aws-types = { version = "1.3.4", default-features = false }
 aws-smithy-wasm = { version = "0.1.4", default-features = false }
 aws-sdk-bedrockruntime = { version = "1.56.0", default-features = false }
 aws-smithy-types = { version = "1.3.1" }
+aws-smithy-runtime-api = "1.8.3"
 
 wasi-async-runtime = "0.1.2"
 
@@ -26,11 +27,15 @@ golem-llm = { workspace = true }
 
 golem-rust = { workspace = true }
 log = { workspace = true }
-reqwest = { workspace = true }
+reqwest = { git = "https://github.com/golemcloud/reqwest", branch = "update-july-2025", features = [
+    "json",
+    "async",
+] }
 serde = { workspace = true }
 serde_json = { workspace = true }
 wit-bindgen-rt = { workspace = true }
 base64 = { workspace = true }
+bytes = "1.10.1"
 
 [lib]
 crate-type = ["cdylib"]

llm/bedrock/Cargo.toml

@@ -7,10 +7,43 @@ pub fn get_async_runtime() -> AsyncRuntime {
 pub struct AsyncRuntime;
 
 impl AsyncRuntime {
-    pub fn block_on<F>(self, f: F) -> F::Output
+    pub fn block_on<F, Fut>(self, f: F) -> Fut::Output
     where
-        F: Future,
+        F: FnOnce(wasi_async_runtime::Reactor) -> Fut,
+        Fut: Future,
     {
-        wasi_async_runtime::block_on(|_| f)
+        wasi_async_runtime::block_on(f)
+    }
+}
+
+#[derive(Clone)]
+pub struct UnsafeFuture<Fut> {
+    inner: Fut,
+}
+
+impl<F> UnsafeFuture<F>
+where
+    F: Future,
+{
+    pub fn new(inner: F) -> Self {
+        Self { inner }
+    }
+}
+
+unsafe impl<F> Send for UnsafeFuture<F> where F: Future {}
+unsafe impl<F> Sync for UnsafeFuture<F> where F: Future {}
+
+impl<F> Future for UnsafeFuture<F>
+where
+    F: Future,
+{
+    type Output = F::Output;
+
+    fn poll(
+        mut self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Self::Output> {
+        let pinned_future = unsafe { self.as_mut().map_unchecked_mut(|this| &mut this.inner) };
+        pinned_future.poll(cx)
     }
 }

llm/bedrock/src/async_utils.rs

@@ -1,7 +1,8 @@
 use crate::{
-    async_utils,
+    async_utils::UnsafeFuture,
     conversions::{self, from_converse_sdk_error, from_converse_stream_sdk_error, BedrockInput},
     stream::BedrockChatStream,
+    wasi_client::WasiClient,
 };
 use aws_config::BehaviorVersion;
 use aws_sdk_bedrockruntime::{
@@ -12,7 +13,6 @@ use aws_sdk_bedrockruntime::{
         converse_stream::builders::ConverseStreamFluentBuilder,
     },
 };
-use aws_smithy_wasm::wasi::WasiHttpClientBuilder;
 use aws_types::region;
 use golem_llm::{
     config::{get_config_key, get_config_key_or_none},
@@ -27,69 +27,59 @@ pub struct Bedrock {
 }
 
 impl Bedrock {
-    pub fn new() -> Result<Self, llm::Error> {
+    pub async fn new(reactor: wasi_async_runtime::Reactor) -> Result<Self, llm::Error> {
         let environment = BedrockEnvironment::load_from_env()?;
 
-        let wasi_http = WasiHttpClientBuilder::new().build();
-
-        let runtime = async_utils::get_async_runtime();
-
-        runtime.block_on(async {
-            let sdk_config = aws_config::defaults(BehaviorVersion::latest())
-                .region(environment.aws_region())
-                .http_client(wasi_http)
-                .credentials_provider(environment.aws_credentials())
-                .sleep_impl(WasiSleep)
-                .load()
-                .await;
-            let client = bedrock::Client::new(&sdk_config);
-            Ok(Self { client })
-        })
+        let sdk_config = aws_config::defaults(BehaviorVersion::latest())
+            .region(environment.aws_region())
+            .http_client(WasiClient::new(reactor.clone()))
+            .credentials_provider(environment.aws_credentials())
+            .sleep_impl(WasiSleep::new(reactor))
+            .load()
+            .await;
+        let client = bedrock::Client::new(&sdk_config);
+        Ok(Self { client })
     }
 
-    pub fn converse(
+    pub async fn converse(
         &self,
         messages: Vec<llm::Message>,
         config: llm::Config,
         tool_results: Option<Vec<(llm::ToolCall, llm::ToolResult)>>,
     ) -> llm::ChatEvent {
         let bedrock_input = BedrockInput::from(messages, config, tool_results);
 
-        let runtime = async_utils::get_async_runtime();
-
         match bedrock_input {
             Err(err) => llm::ChatEvent::Error(err),
             Ok(input) => {
                 trace!("Sending request to AWS Bedrock: {input:?}");
-                runtime.block_on(async {
-                    let model_id = input.model_id.clone();
-                    let response = self
-                        .init_converse(input)
-                        .send()
-                        .await
-                        .map_err(|e| from_converse_sdk_error(model_id, e));
-
-                    match response {
-                        Err(err) => llm::ChatEvent::Error(err),
-                        Ok(response) => {
-                            let event = match response.stop_reason() {
-                                bedrock::types::StopReason::ToolUse => {
-                                    conversions::converse_output_to_tool_calls(response)
-                                        .map(llm::ChatEvent::ToolRequest)
-                                }
-                                _ => conversions::converse_output_to_complete_response(response)
-                                    .map(llm::ChatEvent::Message),
-                            };
-
-                            event.unwrap_or_else(llm::ChatEvent::Error)
-                        }
+                let model_id = input.model_id.clone();
+                let response = self
+                    .init_converse(input)
+                    .send()
+                    .await
+                    .map_err(|e| from_converse_sdk_error(model_id, e));
+
+                match response {
+                    Err(err) => llm::ChatEvent::Error(err),
+                    Ok(response) => {
+                        let event = match response.stop_reason() {
+                            bedrock::types::StopReason::ToolUse => {
+                                conversions::converse_output_to_tool_calls(response)
+                                    .map(llm::ChatEvent::ToolRequest)
+                            }
+                            _ => conversions::converse_output_to_complete_response(response)
+                                .map(llm::ChatEvent::Message),
+                        };
+
+                        event.unwrap_or_else(llm::ChatEvent::Error)
                     }
-                })
+                }
             }
         }
     }
 
-    pub fn converse_stream(
+    pub async fn converse_stream(
         &self,
         messages: Vec<llm::Message>,
         config: llm::Config,
@@ -99,22 +89,19 @@ impl Bedrock {
         match bedrock_input {
             Err(err) => BedrockChatStream::failed(err),
             Ok(input) => {
-                let runtime = async_utils::get_async_runtime();
                 trace!("Sending request to AWS Bedrock: {input:?}");
-                runtime.block_on(async {
-                    let model_id = input.model_id.clone();
-                    let response = self
-                        .init_converse_stream(input)
-                        .send()
-                        .await
-                        .map_err(|e| from_converse_stream_sdk_error(model_id, e));
-
-                    trace!("Creating AWS Bedrock event stream");
-                    match response {
-                        Ok(response) => BedrockChatStream::new(response.stream),
-                        Err(error) => BedrockChatStream::failed(error),
-                    }
-                })
+                let model_id = input.model_id.clone();
+                let response = self
+                    .init_converse_stream(input)
+                    .send()
+                    .await
+                    .map_err(|e| from_converse_stream_sdk_error(model_id, e));
+
+                trace!("Creating AWS Bedrock event stream");
+                match response {
+                    Ok(response) => BedrockChatStream::new(response.stream),
+                    Err(error) => BedrockChatStream::failed(error),
+                }
             }
         }
     }
@@ -146,15 +133,15 @@ impl Bedrock {
 }
 
 #[derive(Debug)]
-struct BedrockEnvironment {
+pub struct BedrockEnvironment {
     access_key_id: String,
     region: String,
     secret_access_key: String,
     session_token: Option<String>,
 }
 
 impl BedrockEnvironment {
-    fn load_from_env() -> Result<Self, llm::Error> {
+    pub fn load_from_env() -> Result<Self, llm::Error> {
         Ok(Self {
             access_key_id: get_config_key("AWS_ACCESS_KEY_ID")?,
             region: get_config_key("AWS_REGION")?,
@@ -179,12 +166,32 @@ impl BedrockEnvironment {
 }
 
 #[derive(Debug, Clone)]
-struct WasiSleep;
+struct WasiSleep {
+    reactor: wasi_async_runtime::Reactor,
+}
+
+impl WasiSleep {
+    fn new(reactor: wasi_async_runtime::Reactor) -> Self {
+        Self { reactor }
+    }
+}
+
 impl AsyncSleep for WasiSleep {
     fn sleep(&self, duration: std::time::Duration) -> Sleep {
-        Sleep::new(Box::pin(async move {
+        let reactor = self.reactor.clone();
+
+        let fut = async move {
             let nanos = duration.as_nanos() as u64;
-            monotonic_clock::subscribe_duration(nanos).block();
-        }))
+            let pollable = monotonic_clock::subscribe_duration(nanos);
+
+            reactor
+                .clone()
+                .wait_for(unsafe { std::mem::transmute(pollable) })
+                .await;
+        };
+        Sleep::new(Box::pin(UnsafeFuture::new(fut)))
     }
 }
+
+unsafe impl Send for WasiSleep {}
+unsafe impl Sync for WasiSleep {}

llm/bedrock/src/client.rs

@@ -14,6 +14,8 @@ use aws_sdk_bedrockruntime::{
 };
 use golem_llm::golem::llm::llm;
 
+use crate::async_utils::get_async_runtime;
+
 #[derive(Debug)]
 pub struct BedrockInput {
     pub model_id: String,
@@ -225,34 +227,39 @@ fn get_image_content_block_from_url(url: &str) -> Result<bedrock::types::Content
 }
 
 fn get_bytes_from_url(url: &str) -> Result<Vec<u8>, llm::Error> {
-    let client = reqwest::Client::builder()
-        .build()
-        .expect("Failed to initialize HTTP client");
-
-    let response = client.get(url).send().map_err(|err| {
-        custom_error(
-            llm::ErrorCode::InvalidRequest,
-            format!("Could not read image bytes from url: {url}, cause: {err}"),
-        )
-    })?;
-    if !response.status().is_success() {
-        return Err(custom_error(
-            llm::ErrorCode::InvalidRequest,
-            format!(
-                "Could not read image bytes from url: {url}, cause: request failed with status: {}",
-                response.status()
-            ),
-        ));
-    }
+    let runtime = get_async_runtime();
 
-    let bytes = response.bytes().map_err(|err| {
-        custom_error(
-            llm::ErrorCode::InvalidRequest,
-            format!("Could not read image bytes from url: {url}, cause: {err}"),
-        )
-    })?;
+    runtime.block_on(|reactor| async {
+        let client = reqwest::Client::builder(reactor)
+            .build()
+            .expect("Failed to initialize HTTP client");
+    
+        let response = client.get(url).send().await.map_err(|err| {
+            custom_error(
+                llm::ErrorCode::InvalidRequest,
+                format!("Could not read image bytes from url: {url}, cause: {err}"),
+            )
+        })?;
+        if !response.status().is_success() {
+            return Err(custom_error(
+                llm::ErrorCode::InvalidRequest,
+                format!(
+                    "Could not read image bytes from url: {url}, cause: request failed with status: {}",
+                    response.status()
+                ),
+            ));
+        }
+    
+        let bytes = response.bytes().await.map_err(|err| {
+            custom_error(
+                llm::ErrorCode::InvalidRequest,
+                format!("Could not read image bytes from url: {url}, cause: {err}"),
+            )
+        })?;
+    
+        Ok(bytes.to_vec())
 
-    Ok(bytes.to_vec())
+    })
 }
 
 fn str_to_bedrock_mime_type(mime_type: &str) -> Result<ImageFormat, llm::Error> {