Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge 2ed3ba8d6b into 0e575b2809

Browse source
This commit is contained in:
Shardul Vaidya 2025-08-27 00:39:16 +08:00 • committed by GitHub
commit 2e8a45b41a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 22 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

9
crates/bedrock/src/bedrock.rs
View file

@ -3,6 +3,7 @@ mod models;
use anyhow::{Context, Error, Result, anyhow}; use anyhow::{Context, Error, Result, anyhow};
use aws_sdk_bedrockruntime as bedrock; use aws_sdk_bedrockruntime as bedrock;
pub use aws_sdk_bedrockruntime as bedrock_client; pub use aws_sdk_bedrockruntime as bedrock_client;
use aws_sdk_bedrockruntime::types::InferenceConfiguration;
pub use aws_sdk_bedrockruntime::types::{ pub use aws_sdk_bedrockruntime::types::{
AnyToolChoice as BedrockAnyToolChoice, AutoToolChoice as BedrockAutoToolChoice, AnyToolChoice as BedrockAnyToolChoice, AutoToolChoice as BedrockAutoToolChoice,
ContentBlock as BedrockInnerContent, Tool as BedrockTool, ToolChoice as BedrockToolChoice, ContentBlock as BedrockInnerContent, Tool as BedrockTool, ToolChoice as BedrockToolChoice,
@ -58,6 +59,14 @@ pub async fn stream_completion(
response = response.set_tool_config(request.tools); response = response.set_tool_config(request.tools);
} }
let inference_config = InferenceConfiguration::builder()
.max_tokens(request.max_tokens as i32)
.set_temperature(request.temperature)
.set_top_p(request.top_p)
.build();
response = response.inference_config(inference_config);
let output = response let output = response
.send() .send()
.await .await

28
crates/bedrock/src/models.rs
View file

@ -151,12 +151,12 @@ impl Model {
pub fn id(&self) -> &str { pub fn id(&self) -> &str {
match self { match self {
Model::ClaudeSonnet4 => "claude-4-sonnet", Model::ClaudeSonnet4 => "claude-sonnet-4",
Model::ClaudeSonnet4Thinking => "claude-4-sonnet-thinking", Model::ClaudeSonnet4Thinking => "claude-sonnet-4-thinking",
Model::ClaudeOpus4 => "claude-4-opus", Model::ClaudeOpus4 => "claude-opus-4",
Model::ClaudeOpus4_1 => "claude-4-opus-1", Model::ClaudeOpus4_1 => "claude-opus-4-1",
Model::ClaudeOpus4Thinking => "claude-4-opus-thinking", Model::ClaudeOpus4Thinking => "claude-opus-4-thinking",
Model::ClaudeOpus4_1Thinking => "claude-4-opus-1-thinking", Model::ClaudeOpus4_1Thinking => "claude-opus-4-1-thinking",
Model::Claude3_5SonnetV2 => "claude-3-5-sonnet-v2", Model::Claude3_5SonnetV2 => "claude-3-5-sonnet-v2",
Model::Claude3_5Sonnet => "claude-3-5-sonnet", Model::Claude3_5Sonnet => "claude-3-5-sonnet",
Model::Claude3Opus => "claude-3-opus", Model::Claude3Opus => "claude-3-opus",
@ -359,14 +359,12 @@ impl Model {
pub fn max_output_tokens(&self) -> u64 { pub fn max_output_tokens(&self) -> u64 {
match self { match self {
Self::Claude3Opus | Self::Claude3Sonnet | Self::Claude3_5Haiku => 4_096, Self::Claude3Opus | Self::Claude3Sonnet | Self::Claude3_5Haiku => 4_096,
Self::Claude3_7Sonnet Self::Claude3_7Sonnet | Self::Claude3_7SonnetThinking => 128_000,
| Self::Claude3_7SonnetThinking Self::ClaudeSonnet4 | Self::ClaudeSonnet4Thinking => 64_000,
| Self::ClaudeSonnet4 Self::ClaudeOpus4
| Self::ClaudeSonnet4Thinking | Self::ClaudeOpus4Thinking
| Self::ClaudeOpus4
| Model::ClaudeOpus4Thinking
| Self::ClaudeOpus4_1 | Self::ClaudeOpus4_1
| Model::ClaudeOpus4_1Thinking => 128_000, | Self::ClaudeOpus4_1Thinking => 32_000,
Self::Claude3_5SonnetV2 | Self::PalmyraWriterX4 | Self::PalmyraWriterX5 => 8_192, Self::Claude3_5SonnetV2 | Self::PalmyraWriterX4 | Self::PalmyraWriterX5 => 8_192,
Self::Custom { Self::Custom {
max_output_tokens, .. max_output_tokens, ..
@ -784,10 +782,10 @@ mod tests {
); );
// Test thinking models have different friendly IDs but same request IDs // Test thinking models have different friendly IDs but same request IDs
assert_eq!(Model::ClaudeSonnet4.id(), "claude-4-sonnet"); assert_eq!(Model::ClaudeSonnet4.id(), "claude-sonnet-4");
assert_eq!( assert_eq!(
Model::ClaudeSonnet4Thinking.id(), Model::ClaudeSonnet4Thinking.id(),
"claude-4-sonnet-thinking" "claude-sonnet-4-thinking"
); );
assert_eq!( assert_eq!(
Model::ClaudeSonnet4.request_id(), Model::ClaudeSonnet4.request_id(),