Skip to content

Commit 5288763

Browse files
committed
feat: use reader for passing body to avoid expensive clone
1 parent fccd5ba commit 5288763

File tree

1 file changed

+27
-8
lines changed

1 file changed

+27
-8
lines changed

llm/bedrock/src/wasi_client.rs

Lines changed: 27 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use std::{str::FromStr, sync::Arc};
1+
use std::{io::Read, str::FromStr, sync::Arc};
22

33
use aws_sdk_bedrockruntime::{config, error::ConnectorError};
44
use aws_smithy_runtime_api::{
@@ -82,13 +82,11 @@ impl WasiConnector {
8282
);
8383
}
8484

85-
let mut request = self.0.request(method, url).headers(header_map);
86-
87-
if let Some(bytes) = parts.body.bytes() {
88-
request = request.body(bytes.to_owned());
89-
}
90-
91-
request
85+
let body_reader = BodyReader::new(parts.body);
86+
self.0
87+
.request(method, url)
88+
.headers(header_map)
89+
.body(reqwest::Body::new(body_reader))
9290
.send()
9391
.await
9492
.map_err(|e| ConnectorError::other(e.into(), None))
@@ -138,3 +136,24 @@ impl HttpConnector for SharedWasiConnector {
138136
HttpConnectorFuture::new(UnsafeFuture::new(future))
139137
}
140138
}
139+
140+
struct BodyReader {
141+
body: SdkBody,
142+
}
143+
144+
impl BodyReader {
145+
fn new(body: SdkBody) -> Self {
146+
Self { body }
147+
}
148+
}
149+
150+
impl Read for BodyReader {
151+
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
152+
let bytes = self.body.bytes();
153+
154+
match bytes {
155+
Some(mut bytes) => bytes.read(buf),
156+
None => Ok(0),
157+
}
158+
}
159+
}

0 commit comments

Comments
 (0)