Update: 将子项目从 submodule 转为完整内容
- 移除 GovAI, nomifun-tauri, 算力盒子 的 submodule 引用 - 添加所有子项目的完整源代码 - 保留原始 .git 为 .git.bak 备份
This commit is contained in:
freedak
@@ -0,0 +1,185 @@
|
||||
use async_trait::async_trait;
|
||||
use reqwest::header::{CONTENT_TYPE, HeaderMap, HeaderValue};
|
||||
use serde_json::{Value, json};
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use nomi_types::llm::{LlmEvent, LlmRequest, ThinkingConfig};
|
||||
|
||||
use super::anthropic_shared;
|
||||
use crate::{LlmProvider, ProviderError};
|
||||
use nomi_config::compat::ProviderCompat;
|
||||
|
||||
pub struct AnthropicProvider {
|
||||
client: reqwest::Client,
|
||||
api_key: String,
|
||||
base_url: String,
|
||||
cache_enabled: bool,
|
||||
compat: ProviderCompat,
|
||||
}
|
||||
|
||||
impl AnthropicProvider {
|
||||
pub fn new(api_key: &str, base_url: &str, compat: ProviderCompat) -> Self {
|
||||
Self {
|
||||
client: crate::http_client(),
|
||||
api_key: api_key.to_string(),
|
||||
base_url: base_url.to_string(),
|
||||
cache_enabled: true,
|
||||
compat,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_cache(mut self, enabled: bool) -> Self {
|
||||
self.cache_enabled = enabled;
|
||||
self
|
||||
}
|
||||
|
||||
fn build_headers(&self) -> Result<HeaderMap, ProviderError> {
|
||||
let mut headers = HeaderMap::new();
|
||||
let api_key = HeaderValue::from_str(&self.api_key)
|
||||
.map_err(|e| ProviderError::Connection(format!("Invalid x-api-key header: {}", e)))?;
|
||||
headers.insert("x-api-key", api_key);
|
||||
headers.insert("anthropic-version", HeaderValue::from_static("2023-06-01"));
|
||||
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
|
||||
if self.cache_enabled {
|
||||
headers.insert(
|
||||
"anthropic-beta",
|
||||
HeaderValue::from_static("prompt-caching-2024-07-31"),
|
||||
);
|
||||
}
|
||||
Ok(headers)
|
||||
}
|
||||
|
||||
fn build_request_body(&self, request: &LlmRequest) -> Value {
|
||||
// Build system prompt with optional cache_control
|
||||
let system = if self.cache_enabled {
|
||||
json!([{
|
||||
"type": "text",
|
||||
"text": &request.system,
|
||||
"cache_control": { "type": "ephemeral" }
|
||||
}])
|
||||
} else {
|
||||
json!(&request.system)
|
||||
};
|
||||
|
||||
let mut body = json!({
|
||||
"model": request.model,
|
||||
"max_tokens": request.max_tokens,
|
||||
"system": system,
|
||||
"messages": anthropic_shared::build_messages(&request.messages, &self.compat),
|
||||
"stream": true
|
||||
});
|
||||
|
||||
if !request.tools.is_empty() {
|
||||
let mut tools = anthropic_shared::build_tools(&request.tools);
|
||||
// Mark last tool with cache_control to cache the entire tools block
|
||||
if let Some(last) = tools.last_mut().filter(|_| self.cache_enabled) {
|
||||
last["cache_control"] = json!({ "type": "ephemeral" });
|
||||
}
|
||||
body["tools"] = json!(tools);
|
||||
}
|
||||
|
||||
if let Some(ThinkingConfig::Enabled { budget_tokens }) = &request.thinking {
|
||||
body["thinking"] = json!({
|
||||
"type": "enabled",
|
||||
"budget_tokens": budget_tokens
|
||||
});
|
||||
}
|
||||
|
||||
body
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LlmProvider for AnthropicProvider {
|
||||
async fn stream(
|
||||
&self,
|
||||
request: &LlmRequest,
|
||||
) -> Result<mpsc::Receiver<LlmEvent>, ProviderError> {
|
||||
let url = format!("{}/v1/messages", self.base_url);
|
||||
let body = self.build_request_body(request);
|
||||
|
||||
tracing::debug!(target: "nomi_providers", body = %serde_json::to_string_pretty(&body).unwrap_or_default(), "outgoing request");
|
||||
|
||||
let response = crate::retry::with_initial_connect_retry(|| async {
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.headers(self.build_headers()?)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
let retry_after_ms = crate::parse_retry_after_ms(response.headers()).unwrap_or(5000);
|
||||
let body_text = response.text().await.unwrap_or_default();
|
||||
if status.as_u16() == 429 {
|
||||
return Err(ProviderError::RateLimited {
|
||||
retry_after_ms,
|
||||
message: crate::non_empty_rate_limit_message(body_text),
|
||||
});
|
||||
}
|
||||
return Err(ProviderError::Api {
|
||||
status: status.as_u16(),
|
||||
message: body_text,
|
||||
});
|
||||
}
|
||||
Ok(response)
|
||||
})
|
||||
.await?;
|
||||
|
||||
let (tx, rx) = mpsc::channel(64);
|
||||
let client = self.client.clone();
|
||||
let headers = self.build_headers()?;
|
||||
let url_clone = url.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
match anthropic_shared::process_sse_stream(response, &tx).await {
|
||||
anthropic_shared::StreamOutcome::Ok => {}
|
||||
anthropic_shared::StreamOutcome::FailedPartial(e) => {
|
||||
let _ = tx.send(LlmEvent::Error(e.to_string())).await;
|
||||
}
|
||||
anthropic_shared::StreamOutcome::FailedEmpty(e) => {
|
||||
if e.is_retryable() {
|
||||
let mut backoff = std::time::Duration::from_secs(1);
|
||||
let mut final_err = Some(e);
|
||||
for attempt in 1..=crate::retry::MAX_STREAM_RETRIES {
|
||||
backoff = crate::retry::backoff_sleep(attempt, backoff).await;
|
||||
match crate::retry::send_and_check(&client, &url_clone, &headers, &body)
|
||||
.await
|
||||
{
|
||||
Ok(resp) => {
|
||||
let outcome =
|
||||
anthropic_shared::process_sse_stream(resp, &tx).await;
|
||||
match crate::retry::evaluate_outcome(outcome, attempt) {
|
||||
Ok(None) => {
|
||||
final_err = None;
|
||||
break;
|
||||
}
|
||||
Ok(Some(e)) => {
|
||||
final_err = Some(e);
|
||||
break;
|
||||
}
|
||||
Err(_) => continue,
|
||||
}
|
||||
}
|
||||
Err(e) if attempt == crate::retry::MAX_STREAM_RETRIES => {
|
||||
final_err = Some(e);
|
||||
break;
|
||||
}
|
||||
Err(_) => continue,
|
||||
}
|
||||
}
|
||||
if let Some(err) = final_err {
|
||||
let _ = tx.send(LlmEvent::Error(err.to_string())).await;
|
||||
}
|
||||
} else {
|
||||
let _ = tx.send(LlmEvent::Error(e.to_string())).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(rx)
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,548 @@
|
||||
// AWS Bedrock provider for Claude models.
|
||||
// Uses AWS SigV4 authentication and AWS event stream binary framing.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use aws_credential_types::Credentials;
|
||||
use aws_sigv4::http_request::{
|
||||
self as sigv4_http, PayloadChecksumKind, SignableBody, SignableRequest, SignatureLocation,
|
||||
SigningSettings,
|
||||
};
|
||||
use reqwest::header::{CONTENT_TYPE, HeaderMap, HeaderValue};
|
||||
use serde_json::{Value, json};
|
||||
use std::time::SystemTime;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use base64::Engine as _;
|
||||
|
||||
use nomi_types::llm::{LlmEvent, LlmRequest, ThinkingConfig};
|
||||
use nomi_types::message::{StopReason, TokenUsage};
|
||||
|
||||
use super::anthropic_shared;
|
||||
use crate::{LlmProvider, ProviderError};
|
||||
use nomi_config::compat::{self, ProviderCompat};
|
||||
|
||||
pub struct BedrockProvider {
|
||||
client: reqwest::Client,
|
||||
region: String,
|
||||
credentials: AwsCredentials,
|
||||
cache_enabled: bool,
|
||||
compat: ProviderCompat,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum AwsCredentials {
|
||||
Explicit {
|
||||
access_key_id: String,
|
||||
secret_access_key: String,
|
||||
session_token: Option<String>,
|
||||
},
|
||||
Profile(String),
|
||||
Environment,
|
||||
}
|
||||
|
||||
impl BedrockProvider {
|
||||
pub fn new(
|
||||
region: &str,
|
||||
credentials: AwsCredentials,
|
||||
cache_enabled: bool,
|
||||
compat: ProviderCompat,
|
||||
) -> Self {
|
||||
Self {
|
||||
client: crate::http_client(),
|
||||
region: region.to_string(),
|
||||
credentials,
|
||||
cache_enabled,
|
||||
compat,
|
||||
}
|
||||
}
|
||||
|
||||
fn build_request_body(&self, request: &LlmRequest) -> Value {
|
||||
let system = if self.cache_enabled {
|
||||
json!([{
|
||||
"type": "text",
|
||||
"text": &request.system,
|
||||
"cache_control": { "type": "ephemeral" }
|
||||
}])
|
||||
} else {
|
||||
json!(&request.system)
|
||||
};
|
||||
|
||||
let mut body = json!({
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"max_tokens": request.max_tokens,
|
||||
"system": system,
|
||||
"messages": anthropic_shared::build_messages(&request.messages, &self.compat)
|
||||
});
|
||||
|
||||
if !request.tools.is_empty() {
|
||||
let mut tools = anthropic_shared::build_tools(&request.tools);
|
||||
if self.compat.sanitize_schema() {
|
||||
for tool in &mut tools {
|
||||
if let Some(schema) = tool.get("input_schema").cloned() {
|
||||
tool["input_schema"] = compat::sanitize_json_schema(&schema);
|
||||
}
|
||||
}
|
||||
}
|
||||
if self.cache_enabled
|
||||
&& let Some(last) = tools.last_mut()
|
||||
{
|
||||
last["cache_control"] = json!({ "type": "ephemeral" });
|
||||
}
|
||||
body["tools"] = json!(tools);
|
||||
}
|
||||
|
||||
if let Some(ThinkingConfig::Enabled { budget_tokens }) = &request.thinking {
|
||||
body["thinking"] = json!({
|
||||
"type": "enabled",
|
||||
"budget_tokens": budget_tokens
|
||||
});
|
||||
}
|
||||
|
||||
body
|
||||
}
|
||||
|
||||
fn build_url(&self, model: &str) -> String {
|
||||
format!(
|
||||
"https://bedrock-runtime.{}.amazonaws.com/model/{}/invoke-with-response-stream",
|
||||
self.region, model
|
||||
)
|
||||
}
|
||||
|
||||
fn resolve_credentials(&self) -> Result<Credentials, ProviderError> {
|
||||
match &self.credentials {
|
||||
AwsCredentials::Explicit {
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
session_token,
|
||||
} => Ok(Credentials::new(
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
session_token.clone(),
|
||||
None,
|
||||
"nomi",
|
||||
)),
|
||||
AwsCredentials::Profile(profile) => Self::credentials_from_sdk(Some(profile.clone())),
|
||||
AwsCredentials::Environment => Self::credentials_from_sdk(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn credentials_from_sdk(profile: Option<String>) -> Result<Credentials, ProviderError> {
|
||||
// Use a short-lived tokio runtime to resolve credentials synchronously.
|
||||
// This is called once per LLM request so the overhead is acceptable.
|
||||
let rt = tokio::runtime::Handle::try_current();
|
||||
|
||||
let resolve = async move {
|
||||
let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
|
||||
if let Some(p) = profile {
|
||||
loader = loader.profile_name(p);
|
||||
}
|
||||
let config = loader.load().await;
|
||||
let provider = config.credentials_provider().ok_or_else(|| {
|
||||
ProviderError::Connection(
|
||||
"No AWS credentials found. Set AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY, \
|
||||
AWS_PROFILE, or configure credentials in ~/.aws/credentials"
|
||||
.into(),
|
||||
)
|
||||
})?;
|
||||
|
||||
use aws_credential_types::provider::ProvideCredentials;
|
||||
let creds = provider
|
||||
.provide_credentials()
|
||||
.await
|
||||
.map_err(|e| ProviderError::Connection(format!("AWS credential error: {}", e)))?;
|
||||
|
||||
Ok(Credentials::new(
|
||||
creds.access_key_id(),
|
||||
creds.secret_access_key(),
|
||||
creds.session_token().map(|s| s.to_string()),
|
||||
creds.expiry(),
|
||||
"nomi-sdk",
|
||||
))
|
||||
};
|
||||
|
||||
match rt {
|
||||
Ok(_handle) => {
|
||||
// Already inside a tokio runtime — use spawn_blocking to avoid nested block_on
|
||||
std::thread::scope(|s| {
|
||||
s.spawn(|| {
|
||||
tokio::runtime::Runtime::new()
|
||||
.map_err(|e| {
|
||||
ProviderError::Connection(format!("Runtime error: {}", e))
|
||||
})?
|
||||
.block_on(resolve)
|
||||
})
|
||||
.join()
|
||||
.unwrap()
|
||||
})
|
||||
}
|
||||
Err(_) => {
|
||||
// No runtime — safe to create one
|
||||
tokio::runtime::Runtime::new()
|
||||
.map_err(|e| ProviderError::Connection(format!("Runtime error: {}", e)))?
|
||||
.block_on(resolve)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn sign_bedrock_request(
|
||||
region: &str,
|
||||
method: &str,
|
||||
url: &str,
|
||||
headers: &HeaderMap,
|
||||
body: &[u8],
|
||||
credentials: &Credentials,
|
||||
) -> Result<HeaderMap, ProviderError> {
|
||||
let mut signing_settings = SigningSettings::default();
|
||||
signing_settings.payload_checksum_kind = PayloadChecksumKind::XAmzSha256;
|
||||
signing_settings.signature_location = SignatureLocation::Headers;
|
||||
|
||||
let identity = credentials.clone().into();
|
||||
let signing_params = aws_sigv4::sign::v4::SigningParams::builder()
|
||||
.identity(&identity)
|
||||
.region(region)
|
||||
.name("bedrock")
|
||||
.time(SystemTime::now())
|
||||
.settings(signing_settings)
|
||||
.build()
|
||||
.map_err(|e| ProviderError::Connection(format!("SigV4 params error: {}", e)))?;
|
||||
|
||||
// Build header pairs for signing
|
||||
let header_pairs: Vec<(&str, &str)> = headers
|
||||
.iter()
|
||||
.filter_map(|(name, value)| value.to_str().ok().map(|v| (name.as_str(), v)))
|
||||
.collect();
|
||||
|
||||
let signable_request = SignableRequest::new(
|
||||
method,
|
||||
url,
|
||||
header_pairs.into_iter(),
|
||||
SignableBody::Bytes(body),
|
||||
)
|
||||
.map_err(|e| ProviderError::Connection(format!("Signable request error: {}", e)))?;
|
||||
|
||||
let (signing_instructions, _signature) =
|
||||
sigv4_http::sign(signable_request, &signing_params.into())
|
||||
.map_err(|e| ProviderError::Connection(format!("SigV4 signing error: {}", e)))?
|
||||
.into_parts();
|
||||
|
||||
let mut signed_headers = headers.clone();
|
||||
for (name, value) in signing_instructions.headers() {
|
||||
signed_headers.insert(
|
||||
reqwest::header::HeaderName::from_bytes(name.as_bytes())
|
||||
.map_err(|e| ProviderError::Connection(format!("Header name error: {}", e)))?,
|
||||
HeaderValue::from_str(value)
|
||||
.map_err(|e| ProviderError::Connection(format!("Header value error: {}", e)))?,
|
||||
);
|
||||
}
|
||||
|
||||
Ok(signed_headers)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LlmProvider for BedrockProvider {
|
||||
async fn stream(
|
||||
&self,
|
||||
request: &LlmRequest,
|
||||
) -> Result<mpsc::Receiver<LlmEvent>, ProviderError> {
|
||||
let url = self.build_url(&request.model);
|
||||
let body = self.build_request_body(request);
|
||||
|
||||
tracing::debug!(target: "nomi_providers", body = %serde_json::to_string_pretty(&body).unwrap_or_default(), "outgoing request");
|
||||
|
||||
let body_bytes = serde_json::to_vec(&body)
|
||||
.map_err(|e| ProviderError::Connection(format!("JSON serialize error: {}", e)))?;
|
||||
|
||||
let credentials = self.resolve_credentials()?;
|
||||
|
||||
// Each attempt re-signs: SigV4 signatures embed a timestamp and are only
|
||||
// valid within a short window, so a retried request needs a fresh
|
||||
// signature. `send_signed` builds headers, signs, sends, and maps a
|
||||
// non-2xx status — used for both the initial connect-retry and the
|
||||
// mid-stream retry loop (parity with the other providers, which retry via
|
||||
// crate::retry). (Phase 1 provider retry parity)
|
||||
let send_signed = |region: &str,
|
||||
client: &reqwest::Client,
|
||||
url: &str,
|
||||
body_bytes: &[u8],
|
||||
credentials: &Credentials| {
|
||||
let region = region.to_owned();
|
||||
let client = client.clone();
|
||||
let url = url.to_owned();
|
||||
let body_bytes = body_bytes.to_vec();
|
||||
let credentials = credentials.clone();
|
||||
async move {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
|
||||
let signed = Self::sign_bedrock_request(
|
||||
®ion, "POST", &url, &headers, &body_bytes, &credentials,
|
||||
)?;
|
||||
let response = client
|
||||
.post(&url)
|
||||
.headers(signed)
|
||||
.body(body_bytes.clone())
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| ProviderError::Connection(e.to_string()))?;
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
let retry_after_ms =
|
||||
crate::parse_retry_after_ms(response.headers()).unwrap_or(5000);
|
||||
let body_text = response.text().await.unwrap_or_default();
|
||||
if status.as_u16() == 429 {
|
||||
return Err(ProviderError::RateLimited {
|
||||
retry_after_ms,
|
||||
message: crate::non_empty_rate_limit_message(body_text),
|
||||
});
|
||||
}
|
||||
return Err(ProviderError::Api {
|
||||
status: status.as_u16(),
|
||||
message: format_bedrock_error(status.as_u16(), &body_text),
|
||||
});
|
||||
}
|
||||
Ok(response)
|
||||
}
|
||||
};
|
||||
|
||||
// Initial request with connect-failure retry (status/rate-limit errors
|
||||
// are surfaced immediately, same as the other providers).
|
||||
let response = crate::retry::with_initial_connect_retry(|| {
|
||||
send_signed(&self.region, &self.client, &url, &body_bytes, &credentials)
|
||||
})
|
||||
.await?;
|
||||
|
||||
let (tx, rx) = mpsc::channel(64);
|
||||
|
||||
// Owned copies for the spawned task's mid-stream retry loop (it outlives
|
||||
// `&self`, so it cannot borrow region/client/credentials).
|
||||
let region = self.region.clone();
|
||||
let client = self.client.clone();
|
||||
let url_owned = url.clone();
|
||||
|
||||
// AWS event stream uses binary framing.
|
||||
tokio::spawn(async move {
|
||||
match process_aws_event_stream(response, &tx).await {
|
||||
anthropic_shared::StreamOutcome::Ok => {}
|
||||
anthropic_shared::StreamOutcome::FailedPartial(e) => {
|
||||
// Content already emitted — replaying would duplicate it.
|
||||
let _ = tx.send(LlmEvent::Error(e.to_string())).await;
|
||||
}
|
||||
anthropic_shared::StreamOutcome::FailedEmpty(e) => {
|
||||
if e.is_retryable() {
|
||||
let mut backoff = std::time::Duration::from_secs(1);
|
||||
let mut final_err = Some(e);
|
||||
for attempt in 1..=crate::retry::MAX_STREAM_RETRIES {
|
||||
backoff = crate::retry::backoff_sleep(attempt, backoff).await;
|
||||
match send_signed(®ion, &client, &url_owned, &body_bytes, &credentials)
|
||||
.await
|
||||
{
|
||||
Ok(resp) => {
|
||||
let outcome = process_aws_event_stream(resp, &tx).await;
|
||||
match crate::retry::evaluate_outcome(outcome, attempt) {
|
||||
Ok(None) => {
|
||||
final_err = None;
|
||||
break;
|
||||
}
|
||||
Ok(Some(err)) => {
|
||||
final_err = Some(err);
|
||||
break;
|
||||
}
|
||||
Err(_) => continue,
|
||||
}
|
||||
}
|
||||
Err(err) if attempt == crate::retry::MAX_STREAM_RETRIES => {
|
||||
final_err = Some(err);
|
||||
break;
|
||||
}
|
||||
Err(_) => continue,
|
||||
}
|
||||
}
|
||||
if let Some(err) = final_err {
|
||||
let _ = tx.send(LlmEvent::Error(err.to_string())).await;
|
||||
}
|
||||
} else {
|
||||
let _ = tx.send(LlmEvent::Error(e.to_string())).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(rx)
|
||||
}
|
||||
}
|
||||
|
||||
/// Process the AWS event stream (binary framed) from Bedrock
|
||||
async fn process_aws_event_stream(
|
||||
response: reqwest::Response,
|
||||
tx: &mpsc::Sender<LlmEvent>,
|
||||
) -> anthropic_shared::StreamOutcome {
|
||||
use futures::StreamExt;
|
||||
|
||||
let mut state = anthropic_shared::StreamState::new();
|
||||
let mut buffer = Vec::new();
|
||||
let mut stream = response.bytes_stream();
|
||||
let mut emitted_content = false;
|
||||
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = match chunk {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
let err = ProviderError::Connection(e.to_string());
|
||||
return if emitted_content {
|
||||
anthropic_shared::StreamOutcome::FailedPartial(err)
|
||||
} else {
|
||||
anthropic_shared::StreamOutcome::FailedEmpty(err)
|
||||
};
|
||||
}
|
||||
};
|
||||
buffer.extend_from_slice(&chunk);
|
||||
|
||||
// Parse complete AWS event stream messages from buffer
|
||||
while let Some((event_data, consumed)) = parse_aws_event(&buffer) {
|
||||
buffer = buffer[consumed..].to_vec();
|
||||
|
||||
if let Some(payload) = event_data {
|
||||
// The payload contains an SSE-like structure with "bytes" field
|
||||
if let Ok(wrapper) = serde_json::from_slice::<Value>(&payload) {
|
||||
// Bedrock wraps the payload in {"bytes": "base64-encoded-data"}
|
||||
if let Some(b64) = wrapper["bytes"].as_str()
|
||||
&& let Ok(decoded) = base64::engine::general_purpose::STANDARD.decode(b64)
|
||||
&& let Ok(inner) = String::from_utf8(decoded)
|
||||
{
|
||||
tracing::debug!(target: "nomi_providers", chunk = %inner, "bedrock event chunk");
|
||||
// Inner payload is JSON with event type hints
|
||||
if let Ok(json_val) = serde_json::from_str::<Value>(&inner) {
|
||||
let event_type = json_val["type"].as_str().unwrap_or("");
|
||||
let events =
|
||||
anthropic_shared::parse_sse_data(event_type, &inner, &mut state);
|
||||
for event in events {
|
||||
if matches!(
|
||||
event,
|
||||
LlmEvent::TextDelta(_)
|
||||
| LlmEvent::ThinkingDelta(_)
|
||||
| LlmEvent::ThinkingSignature(_)
|
||||
| LlmEvent::ToolUse { .. }
|
||||
) {
|
||||
emitted_content = true;
|
||||
}
|
||||
if tx.send(event).await.is_err() {
|
||||
return anthropic_shared::StreamOutcome::Ok;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// If we haven't sent a Done event, send one now
|
||||
if state.input_tokens > 0 || state.output_tokens > 0 {
|
||||
let _ = tx
|
||||
.send(LlmEvent::Done {
|
||||
stop_reason: StopReason::EndTurn,
|
||||
usage: TokenUsage {
|
||||
input_tokens: state.input_tokens,
|
||||
output_tokens: state.output_tokens,
|
||||
cache_creation_tokens: state.cache_creation_tokens,
|
||||
cache_read_tokens: state.cache_read_tokens,
|
||||
},
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
anthropic_shared::StreamOutcome::Ok
|
||||
}
|
||||
|
||||
/// Parse one AWS event stream message from the buffer.
|
||||
/// Returns (Some(payload), bytes_consumed) if a complete message is found,
|
||||
/// or None if more data is needed.
|
||||
///
|
||||
/// AWS event stream binary format:
|
||||
/// - Prelude: total_len (4 bytes, big-endian) + headers_len (4 bytes) + prelude_crc (4 bytes)
|
||||
/// - Headers: variable length
|
||||
/// - Payload: variable length
|
||||
/// - Message CRC: 4 bytes
|
||||
fn parse_aws_event(buffer: &[u8]) -> Option<(Option<Vec<u8>>, usize)> {
|
||||
if buffer.len() < 12 {
|
||||
return None; // Need at least the prelude
|
||||
}
|
||||
|
||||
let total_len = u32::from_be_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize;
|
||||
let headers_len = u32::from_be_bytes([buffer[4], buffer[5], buffer[6], buffer[7]]) as usize;
|
||||
|
||||
if buffer.len() < total_len {
|
||||
return None; // Incomplete message
|
||||
}
|
||||
|
||||
// Prelude is 12 bytes (total_len + headers_len + prelude_crc)
|
||||
// Payload starts after prelude + headers
|
||||
let payload_start = 12 + headers_len;
|
||||
// Payload ends 4 bytes before total_len (message CRC)
|
||||
let payload_end = total_len - 4;
|
||||
|
||||
if payload_start <= payload_end {
|
||||
let payload = buffer[payload_start..payload_end].to_vec();
|
||||
Some((Some(payload), total_len))
|
||||
} else {
|
||||
// Empty payload (e.g., initial response event)
|
||||
Some((None, total_len))
|
||||
}
|
||||
}
|
||||
|
||||
/// Format Bedrock error responses with actionable hints
|
||||
fn format_bedrock_error(status: u16, body: &str) -> String {
|
||||
// Try to extract the AWS error type from the response
|
||||
let error_type = serde_json::from_str::<Value>(body).ok().and_then(|v| {
|
||||
v.get("__type")
|
||||
.or_else(|| v.get("type"))
|
||||
.and_then(|t| t.as_str().map(String::from))
|
||||
});
|
||||
|
||||
let hint = match status {
|
||||
403 => Some(
|
||||
"Check IAM permissions: the role/user needs bedrock:InvokeModelWithResponseStream. \
|
||||
Also verify the model is enabled in the Bedrock console for your account.",
|
||||
),
|
||||
404 => Some(
|
||||
"Model not found in this region. Verify the model ID and that it's available in \
|
||||
your configured AWS region.",
|
||||
),
|
||||
400 => {
|
||||
if body.contains("schema") || body.contains("Schema") {
|
||||
Some(
|
||||
"Request schema validation failed. If using tools, try enabling sanitize_schema=true in [providers.bedrock.compat].",
|
||||
)
|
||||
} else {
|
||||
Some("Bad request — check model parameters and message format.")
|
||||
}
|
||||
}
|
||||
503 | 529 => Some(
|
||||
"Service overloaded or throttled. You may have exceeded your provisioned throughput quota. \
|
||||
Retry after a moment or request a quota increase.",
|
||||
),
|
||||
_ => None,
|
||||
};
|
||||
|
||||
let type_info = error_type.map(|t| format!(" [{}]", t)).unwrap_or_default();
|
||||
|
||||
match hint {
|
||||
Some(h) => format!("{}{}\nHint: {}", body, type_info, h),
|
||||
None => format!("{}{}", body, type_info),
|
||||
}
|
||||
}
|
||||
|
||||
/// Build AwsCredentials from nomi-config's BedrockConfig
|
||||
pub fn credentials_from_config(bc: &nomi_config::config::BedrockConfig) -> AwsCredentials {
|
||||
if let (Some(key_id), Some(secret)) = (&bc.access_key_id, &bc.secret_access_key) {
|
||||
AwsCredentials::Explicit {
|
||||
access_key_id: key_id.clone(),
|
||||
secret_access_key: secret.clone(),
|
||||
session_token: bc.session_token.clone(),
|
||||
}
|
||||
} else if let Some(profile) = &bc.profile {
|
||||
AwsCredentials::Profile(profile.clone())
|
||||
} else {
|
||||
AwsCredentials::Environment
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,206 @@
|
||||
pub mod anthropic;
|
||||
pub mod anthropic_shared;
|
||||
pub mod bedrock;
|
||||
pub mod openai;
|
||||
pub mod retry;
|
||||
pub mod vertex;
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use nomi_config::config::{Config, ProviderType};
|
||||
use nomi_types::llm::{LlmEvent, LlmRequest};
|
||||
|
||||
/// Unified interface for LLM API providers
|
||||
#[async_trait]
|
||||
pub trait LlmProvider: Send + Sync {
|
||||
async fn stream(&self, request: &LlmRequest)
|
||||
-> Result<mpsc::Receiver<LlmEvent>, ProviderError>;
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum ProviderError {
|
||||
#[error("HTTP error: {0}")]
|
||||
Http(#[from] reqwest::Error),
|
||||
#[error("API error {status}: {message}")]
|
||||
Api { status: u16, message: String },
|
||||
#[error("SSE parse error: {0}")]
|
||||
Parse(String),
|
||||
#[error("Rate limited, retry after {retry_after_ms}ms: {message}")]
|
||||
RateLimited {
|
||||
retry_after_ms: u64,
|
||||
message: String,
|
||||
},
|
||||
#[error("Prompt too long: {0}")]
|
||||
PromptTooLong(String),
|
||||
#[error("Connection error: {0}")]
|
||||
Connection(String),
|
||||
}
|
||||
|
||||
impl ProviderError {
|
||||
pub fn is_retryable(&self) -> bool {
|
||||
match self {
|
||||
ProviderError::RateLimited { .. } | ProviderError::Connection(_) => true,
|
||||
// Transient server-side faults (500/502/503/504) from an overloaded
|
||||
// gateway are the most common spurious failure and are safe to retry
|
||||
// on the pre-response / empty-content paths. 4xx are terminal.
|
||||
ProviderError::Api { status, .. } => *status >= 500,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a `Retry-After` HTTP header into milliseconds, honouring the provider's
|
||||
/// requested backoff instead of a fixed guess. Supports the delta-seconds form
|
||||
/// (what LLM gateways send); returns `None` for an absent, non-numeric, or
|
||||
/// HTTP-date value (caller falls back to its default). Clamped to 120s so a
|
||||
/// hostile/huge value can't wedge the agent.
|
||||
pub(crate) fn parse_retry_after_ms(headers: &reqwest::header::HeaderMap) -> Option<u64> {
|
||||
let secs: u64 = headers
|
||||
.get(reqwest::header::RETRY_AFTER)?
|
||||
.to_str()
|
||||
.ok()?
|
||||
.trim()
|
||||
.parse()
|
||||
.ok()?;
|
||||
Some(secs.saturating_mul(1000).min(120_000))
|
||||
}
|
||||
|
||||
/// Connection timeout for provider HTTP clients. Bounds the TCP/TLS connect
|
||||
/// phase so an unreachable or non-responsive gateway fails fast.
|
||||
const HTTP_CONNECT_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
|
||||
/// Idle read timeout for provider HTTP clients. Applies to each read of the
|
||||
/// (streaming) response, so a gateway that accepts the request but then stalls
|
||||
/// — sending no further bytes — surfaces an error instead of hanging the turn
|
||||
/// forever. Active streaming resets this on every chunk, so it only trips on a
|
||||
/// genuine stall. The health-check probe has its own 30s wrapper; the live
|
||||
/// conversation path previously had NO timeout at all, which turned an upstream
|
||||
/// stall into a silent freeze (no output, no error).
|
||||
const HTTP_READ_TIMEOUT: Duration = Duration::from_secs(120);
|
||||
|
||||
/// Shared reqwest client for all LLM providers, configured with connection and
|
||||
/// idle-read timeouts. A stalled upstream produces a `reqwest` timeout error,
|
||||
/// which the SSE loop converts into `LlmEvent::Error` (surfaced to the user as
|
||||
/// `Nomi agent error: ...`) instead of an indefinite hang.
|
||||
pub(crate) fn http_client() -> reqwest::Client {
|
||||
let builder = reqwest::Client::builder()
|
||||
.connect_timeout(HTTP_CONNECT_TIMEOUT)
|
||||
.read_timeout(HTTP_READ_TIMEOUT);
|
||||
nomifun_net::proxy::apply_detected_proxy(builder)
|
||||
.build()
|
||||
.unwrap_or_else(|_| reqwest::Client::new())
|
||||
}
|
||||
|
||||
pub(crate) fn non_empty_rate_limit_message(body: String) -> String {
|
||||
if body.trim().is_empty() {
|
||||
"HTTP 429 Too Many Requests".to_owned()
|
||||
} else {
|
||||
body
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a provider from resolved config
|
||||
pub fn create_provider(config: &Config) -> Arc<dyn LlmProvider> {
|
||||
let compat = config.compat.clone();
|
||||
|
||||
match config.provider {
|
||||
ProviderType::Anthropic => Arc::new(
|
||||
anthropic::AnthropicProvider::new(&config.api_key, &config.base_url, compat)
|
||||
.with_cache(config.prompt_caching),
|
||||
),
|
||||
ProviderType::OpenAI => Arc::new(openai::OpenAIProvider::new(
|
||||
&config.api_key,
|
||||
&config.base_url,
|
||||
compat,
|
||||
)),
|
||||
ProviderType::Bedrock => {
|
||||
let bc = config.bedrock.clone().unwrap_or_default();
|
||||
let region = bc
|
||||
.region
|
||||
.clone()
|
||||
.or_else(|| std::env::var("AWS_REGION").ok())
|
||||
.or_else(|| std::env::var("AWS_DEFAULT_REGION").ok())
|
||||
.unwrap_or_else(|| "us-east-1".to_string());
|
||||
let credentials = bedrock::credentials_from_config(&bc);
|
||||
Arc::new(bedrock::BedrockProvider::new(
|
||||
®ion,
|
||||
credentials,
|
||||
config.prompt_caching,
|
||||
compat,
|
||||
))
|
||||
}
|
||||
ProviderType::Vertex => {
|
||||
let vc = config.vertex.clone().unwrap_or_default();
|
||||
let project_id = vc.project_id.clone().unwrap_or_default();
|
||||
let region = vc
|
||||
.region
|
||||
.clone()
|
||||
.unwrap_or_else(|| "us-central1".to_string());
|
||||
let auth = vertex::auth_from_config(&vc);
|
||||
Arc::new(vertex::VertexProvider::new(
|
||||
&project_id,
|
||||
®ion,
|
||||
auth,
|
||||
config.prompt_caching,
|
||||
compat,
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod retryable_tests {
|
||||
use super::ProviderError;
|
||||
use super::parse_retry_after_ms;
|
||||
|
||||
#[test]
|
||||
fn parse_retry_after_seconds_clamped() {
|
||||
use reqwest::header::{HeaderMap, HeaderValue, RETRY_AFTER};
|
||||
let mut h = HeaderMap::new();
|
||||
h.insert(RETRY_AFTER, HeaderValue::from_static("30"));
|
||||
assert_eq!(parse_retry_after_ms(&h), Some(30_000));
|
||||
|
||||
let mut huge = HeaderMap::new();
|
||||
huge.insert(RETRY_AFTER, HeaderValue::from_static("99999"));
|
||||
assert_eq!(parse_retry_after_ms(&huge), Some(120_000)); // clamped
|
||||
|
||||
// Absent / non-numeric (HTTP-date) -> None (caller uses its default).
|
||||
assert_eq!(parse_retry_after_ms(&HeaderMap::new()), None);
|
||||
let mut date = HeaderMap::new();
|
||||
date.insert(RETRY_AFTER, HeaderValue::from_static("Wed, 21 Oct 2025 07:28:00 GMT"));
|
||||
assert_eq!(parse_retry_after_ms(&date), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transient_5xx_is_retryable_but_4xx_is_not() {
|
||||
// Transient server-side faults (overloaded gateways) are the most common
|
||||
// spurious failure and are safe to retry on the pre-response / empty
|
||||
// paths; client errors (4xx) are terminal. (Phase 1)
|
||||
let api = |status| ProviderError::Api {
|
||||
status,
|
||||
message: "x".to_string(),
|
||||
};
|
||||
assert!(api(500).is_retryable());
|
||||
assert!(api(502).is_retryable());
|
||||
assert!(api(503).is_retryable());
|
||||
assert!(api(504).is_retryable());
|
||||
assert!(!api(400).is_retryable());
|
||||
assert!(!api(404).is_retryable());
|
||||
assert!(!api(429).is_retryable(), "429 is surfaced as RateLimited, not Api");
|
||||
|
||||
assert!(
|
||||
ProviderError::RateLimited {
|
||||
retry_after_ms: 0,
|
||||
message: "x".to_string()
|
||||
}
|
||||
.is_retryable()
|
||||
);
|
||||
assert!(ProviderError::Connection("x".to_string()).is_retryable());
|
||||
assert!(!ProviderError::PromptTooLong("x".to_string()).is_retryable());
|
||||
assert!(!ProviderError::Parse("x".to_string()).is_retryable());
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,239 @@
|
||||
use std::future::Future;
|
||||
use std::time::Duration;
|
||||
|
||||
use reqwest::header::HeaderMap;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::ProviderError;
|
||||
use super::anthropic_shared::StreamOutcome;
|
||||
|
||||
pub const MAX_STREAM_RETRIES: u32 = 2;
|
||||
pub const MAX_INITIAL_CONNECT_RETRIES: u32 = 2;
|
||||
const MAX_BACKOFF: Duration = Duration::from_secs(15);
|
||||
const INITIAL_CONNECT_BACKOFF: Duration = Duration::from_millis(300);
|
||||
const MAX_INITIAL_CONNECT_BACKOFF: Duration = Duration::from_secs(2);
|
||||
|
||||
/// Retry initial request failures that occur before an HTTP response exists.
|
||||
/// HTTP status errors and rate limits are intentionally not retried here.
|
||||
pub async fn with_initial_connect_retry<F, Fut, T>(f: F) -> Result<T, ProviderError>
|
||||
where
|
||||
F: Fn() -> Fut,
|
||||
Fut: Future<Output = Result<T, ProviderError>>,
|
||||
{
|
||||
let mut backoff = INITIAL_CONNECT_BACKOFF;
|
||||
for attempt in 0..=MAX_INITIAL_CONNECT_RETRIES {
|
||||
match f().await {
|
||||
Ok(val) => return Ok(val),
|
||||
Err(e) if is_initial_connect_error(&e) && attempt < MAX_INITIAL_CONNECT_RETRIES => {
|
||||
tracing::warn!(
|
||||
attempt = attempt + 1,
|
||||
max_retries = MAX_INITIAL_CONNECT_RETRIES,
|
||||
error = %e,
|
||||
"retrying initial provider request after connect failure"
|
||||
);
|
||||
tokio::time::sleep(backoff).await;
|
||||
backoff = (backoff * 2).min(MAX_INITIAL_CONNECT_BACKOFF);
|
||||
}
|
||||
Err(e) => return Err(e),
|
||||
}
|
||||
}
|
||||
unreachable!()
|
||||
}
|
||||
|
||||
fn is_initial_connect_error(error: &ProviderError) -> bool {
|
||||
match error {
|
||||
ProviderError::Http(err) => err.is_connect(),
|
||||
ProviderError::Connection(_) => true,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Send an HTTP request and check status, returning the response on success.
|
||||
/// Used by provider-specific retry loops to avoid duplicating request logic.
|
||||
pub async fn send_and_check(
|
||||
client: &reqwest::Client,
|
||||
url: &str,
|
||||
headers: &HeaderMap,
|
||||
body: &Value,
|
||||
) -> Result<reqwest::Response, ProviderError> {
|
||||
let response = client
|
||||
.post(url)
|
||||
.headers(headers.clone())
|
||||
.json(body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| ProviderError::Connection(e.to_string()))?;
|
||||
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
let body_text = response.text().await.unwrap_or_default();
|
||||
return Err(ProviderError::Api {
|
||||
status: status.as_u16(),
|
||||
message: body_text,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
/// Sleep with exponential backoff and log the retry attempt.
|
||||
/// Returns the next backoff duration.
|
||||
pub async fn backoff_sleep(attempt: u32, current_backoff: Duration) -> Duration {
|
||||
tracing::warn!(
|
||||
attempt,
|
||||
max = MAX_STREAM_RETRIES,
|
||||
"retrying stream after mid-stream disconnect"
|
||||
);
|
||||
tokio::time::sleep(current_backoff).await;
|
||||
(current_backoff * 2).min(MAX_BACKOFF)
|
||||
}
|
||||
|
||||
/// Evaluate a `StreamOutcome` within a retry loop. Returns:
|
||||
/// - `Ok(None)` — stream succeeded, stop retrying
|
||||
/// - `Ok(Some(err))` — non-retryable failure, caller should emit error
|
||||
/// - `Err(err)` — retryable failure, caller should continue loop
|
||||
pub fn evaluate_outcome(
|
||||
outcome: StreamOutcome,
|
||||
attempt: u32,
|
||||
) -> Result<Option<ProviderError>, ProviderError> {
|
||||
match outcome {
|
||||
StreamOutcome::Ok => Ok(None),
|
||||
StreamOutcome::FailedPartial(e) => Ok(Some(e)),
|
||||
StreamOutcome::FailedEmpty(e) => {
|
||||
if attempt == MAX_STREAM_RETRIES {
|
||||
Ok(Some(e))
|
||||
} else {
|
||||
Err(e)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU32, Ordering};
|
||||
|
||||
use super::*;
|
||||
use crate::ProviderError;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_initial_connect_retry_succeeds_after_connection_failures() {
|
||||
tokio::time::pause();
|
||||
|
||||
let counter = Arc::new(AtomicU32::new(0));
|
||||
let result = with_initial_connect_retry(|| {
|
||||
let counter = Arc::clone(&counter);
|
||||
async move {
|
||||
let attempt = counter.fetch_add(1, Ordering::SeqCst);
|
||||
if attempt < 2 {
|
||||
Err(ProviderError::Connection("connection refused".into()))
|
||||
} else {
|
||||
Ok(attempt)
|
||||
}
|
||||
}
|
||||
})
|
||||
.await;
|
||||
|
||||
assert_eq!(result.unwrap(), 2);
|
||||
assert_eq!(counter.load(Ordering::SeqCst), 3);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_initial_connect_retry_does_not_retry_rate_limit() {
|
||||
let counter = Arc::new(AtomicU32::new(0));
|
||||
let result = with_initial_connect_retry(|| {
|
||||
let counter = Arc::clone(&counter);
|
||||
async move {
|
||||
counter.fetch_add(1, Ordering::SeqCst);
|
||||
Err::<(), _>(ProviderError::RateLimited {
|
||||
retry_after_ms: 5000,
|
||||
message: "Too Many Requests".into(),
|
||||
})
|
||||
}
|
||||
})
|
||||
.await;
|
||||
|
||||
assert!(matches!(
|
||||
result.unwrap_err(),
|
||||
ProviderError::RateLimited {
|
||||
retry_after_ms: 5000,
|
||||
..
|
||||
}
|
||||
));
|
||||
assert_eq!(counter.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
// --- evaluate_outcome tests ---
|
||||
|
||||
#[test]
|
||||
fn test_evaluate_outcome_ok_stops_retry() {
|
||||
let result = evaluate_outcome(StreamOutcome::Ok, 1);
|
||||
assert!(matches!(result, Ok(None)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_evaluate_outcome_failed_partial_always_stops() {
|
||||
let err = ProviderError::Connection("disconnect".into());
|
||||
let result = evaluate_outcome(StreamOutcome::FailedPartial(err), 1);
|
||||
// FailedPartial means content was already emitted — cannot retry regardless of attempt
|
||||
let Ok(Some(e)) = result else {
|
||||
panic!("expected Ok(Some(err))")
|
||||
};
|
||||
assert!(matches!(e, ProviderError::Connection(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_evaluate_outcome_failed_partial_on_last_attempt() {
|
||||
let err = ProviderError::Connection("disconnect".into());
|
||||
let result = evaluate_outcome(StreamOutcome::FailedPartial(err), MAX_STREAM_RETRIES);
|
||||
let Ok(Some(_)) = result else {
|
||||
panic!("expected Ok(Some(err))")
|
||||
};
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_evaluate_outcome_failed_empty_retries_when_not_exhausted() {
|
||||
let err = ProviderError::Connection("disconnect".into());
|
||||
// attempt 1 < MAX_STREAM_RETRIES(2), should signal "continue retrying"
|
||||
let result = evaluate_outcome(StreamOutcome::FailedEmpty(err), 1);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_evaluate_outcome_failed_empty_stops_on_last_attempt() {
|
||||
let err = ProviderError::Connection("disconnect".into());
|
||||
// attempt == MAX_STREAM_RETRIES, should stop and return error
|
||||
let result = evaluate_outcome(StreamOutcome::FailedEmpty(err), MAX_STREAM_RETRIES);
|
||||
let Ok(Some(e)) = result else {
|
||||
panic!("expected Ok(Some(err))")
|
||||
};
|
||||
assert!(matches!(e, ProviderError::Connection(_)));
|
||||
}
|
||||
|
||||
// --- backoff_sleep tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_backoff_sleep_doubles_duration() {
|
||||
tokio::time::pause();
|
||||
|
||||
let next = backoff_sleep(1, Duration::from_secs(1)).await;
|
||||
assert_eq!(next, Duration::from_secs(2));
|
||||
|
||||
let next = backoff_sleep(2, Duration::from_secs(4)).await;
|
||||
assert_eq!(next, Duration::from_secs(8));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_backoff_sleep_caps_at_max() {
|
||||
tokio::time::pause();
|
||||
|
||||
// 10s * 2 = 20s, but MAX_BACKOFF is 15s
|
||||
let next = backoff_sleep(1, Duration::from_secs(10)).await;
|
||||
assert_eq!(next, Duration::from_secs(15));
|
||||
|
||||
// Already at max
|
||||
let next = backoff_sleep(2, Duration::from_secs(15)).await;
|
||||
assert_eq!(next, Duration::from_secs(15));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,418 @@
|
||||
// Google Vertex AI provider for Claude models.
|
||||
// Uses GCP OAuth2 authentication. Response is standard SSE (same as Anthropic).
|
||||
|
||||
use async_trait::async_trait;
|
||||
use jsonwebtoken::{Algorithm, EncodingKey, Header};
|
||||
use reqwest::header::{AUTHORIZATION, CONTENT_TYPE, HeaderMap, HeaderValue};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Value, json};
|
||||
use std::sync::Mutex;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use nomi_types::llm::{LlmEvent, LlmRequest, ThinkingConfig};
|
||||
|
||||
use super::anthropic_shared;
|
||||
use crate::{LlmProvider, ProviderError};
|
||||
use nomi_config::compat::ProviderCompat;
|
||||
|
||||
pub struct VertexProvider {
|
||||
client: reqwest::Client,
|
||||
project_id: String,
|
||||
region: String,
|
||||
auth: GcpAuth,
|
||||
cache_enabled: bool,
|
||||
compat: ProviderCompat,
|
||||
/// Cached access token
|
||||
cached_token: Mutex<Option<CachedToken>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum GcpAuth {
|
||||
ServiceAccount { key_file: String },
|
||||
ApplicationDefault,
|
||||
MetadataServer,
|
||||
}
|
||||
|
||||
struct CachedToken {
|
||||
token: String,
|
||||
expires_at: u64,
|
||||
}
|
||||
|
||||
impl VertexProvider {
|
||||
pub fn new(
|
||||
project_id: &str,
|
||||
region: &str,
|
||||
auth: GcpAuth,
|
||||
cache_enabled: bool,
|
||||
compat: ProviderCompat,
|
||||
) -> Self {
|
||||
Self {
|
||||
client: crate::http_client(),
|
||||
project_id: project_id.to_string(),
|
||||
region: region.to_string(),
|
||||
auth,
|
||||
cache_enabled,
|
||||
compat,
|
||||
cached_token: Mutex::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn build_url(&self, model: &str) -> String {
|
||||
format!(
|
||||
"https://{}-aiplatform.googleapis.com/v1/projects/{}/locations/{}/publishers/anthropic/models/{}:streamRawPredict",
|
||||
self.region, self.project_id, self.region, model
|
||||
)
|
||||
}
|
||||
|
||||
fn build_request_body(&self, request: &LlmRequest) -> Value {
|
||||
let system = if self.cache_enabled {
|
||||
json!([{
|
||||
"type": "text",
|
||||
"text": &request.system,
|
||||
"cache_control": { "type": "ephemeral" }
|
||||
}])
|
||||
} else {
|
||||
json!(&request.system)
|
||||
};
|
||||
|
||||
let mut body = json!({
|
||||
"anthropic_version": "vertex-2023-10-16",
|
||||
"max_tokens": request.max_tokens,
|
||||
"system": system,
|
||||
"messages": anthropic_shared::build_messages(&request.messages, &self.compat),
|
||||
"stream": true
|
||||
});
|
||||
|
||||
if !request.tools.is_empty() {
|
||||
let mut tools = anthropic_shared::build_tools(&request.tools);
|
||||
if let Some(last) = tools.last_mut().filter(|_| self.cache_enabled) {
|
||||
last["cache_control"] = json!({ "type": "ephemeral" });
|
||||
}
|
||||
body["tools"] = json!(tools);
|
||||
}
|
||||
|
||||
if let Some(ThinkingConfig::Enabled { budget_tokens }) = &request.thinking {
|
||||
body["thinking"] = json!({
|
||||
"type": "enabled",
|
||||
"budget_tokens": budget_tokens
|
||||
});
|
||||
}
|
||||
|
||||
body
|
||||
}
|
||||
|
||||
async fn get_access_token(&self) -> Result<String, ProviderError> {
|
||||
// Check cache first
|
||||
{
|
||||
let cached = self.cached_token.lock().map_err(|_| {
|
||||
ProviderError::Connection("Vertex token cache lock poisoned".to_string())
|
||||
})?;
|
||||
if let Some(token) = cached.as_ref() {
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs();
|
||||
if token.expires_at > now + 60 {
|
||||
return Ok(token.token.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let (token, expires_in) = match &self.auth {
|
||||
GcpAuth::ServiceAccount { key_file } => {
|
||||
self.get_service_account_token(key_file).await?
|
||||
}
|
||||
GcpAuth::ApplicationDefault => self.get_adc_token().await?,
|
||||
GcpAuth::MetadataServer => self.get_metadata_token().await?,
|
||||
};
|
||||
|
||||
// Cache the token
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs();
|
||||
let mut cached = self.cached_token.lock().map_err(|_| {
|
||||
ProviderError::Connection("Vertex token cache lock poisoned".to_string())
|
||||
})?;
|
||||
*cached = Some(CachedToken {
|
||||
token: token.clone(),
|
||||
expires_at: now + expires_in,
|
||||
});
|
||||
|
||||
Ok(token)
|
||||
}
|
||||
|
||||
async fn get_service_account_token(
|
||||
&self,
|
||||
key_file: &str,
|
||||
) -> Result<(String, u64), ProviderError> {
|
||||
let key_json = std::fs::read_to_string(key_file)
|
||||
.map_err(|e| ProviderError::Connection(format!("Failed to read key file: {}", e)))?;
|
||||
|
||||
let sa: ServiceAccountKey = serde_json::from_str(&key_json)
|
||||
.map_err(|e| ProviderError::Connection(format!("Failed to parse key file: {}", e)))?;
|
||||
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs();
|
||||
|
||||
let claims = JwtClaims {
|
||||
iss: sa.client_email.clone(),
|
||||
scope: "https://www.googleapis.com/auth/cloud-platform".to_string(),
|
||||
aud: sa.token_uri.clone(),
|
||||
iat: now,
|
||||
exp: now + 3600,
|
||||
};
|
||||
|
||||
let encoding_key = EncodingKey::from_rsa_pem(sa.private_key.as_bytes())
|
||||
.map_err(|e| ProviderError::Connection(format!("Invalid RSA key: {}", e)))?;
|
||||
|
||||
let header = Header::new(Algorithm::RS256);
|
||||
let jwt = jsonwebtoken::encode(&header, &claims, &encoding_key)
|
||||
.map_err(|e| ProviderError::Connection(format!("JWT encode error: {}", e)))?;
|
||||
|
||||
// Exchange JWT for access token
|
||||
let resp = self
|
||||
.client
|
||||
.post(&sa.token_uri)
|
||||
.form(&[
|
||||
("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer"),
|
||||
("assertion", &jwt),
|
||||
])
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| ProviderError::Connection(format!("Token exchange error: {}", e)))?;
|
||||
|
||||
let token_resp: GoogleTokenResponse = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| ProviderError::Connection(format!("Token parse error: {}", e)))?;
|
||||
|
||||
Ok((token_resp.access_token, token_resp.expires_in))
|
||||
}
|
||||
|
||||
async fn get_adc_token(&self) -> Result<(String, u64), ProviderError> {
|
||||
// Read Application Default Credentials
|
||||
let adc_path = dirs::home_dir()
|
||||
.ok_or_else(|| ProviderError::Connection("Cannot determine home dir".into()))?
|
||||
.join(".config/gcloud/application_default_credentials.json");
|
||||
|
||||
let adc_json = std::fs::read_to_string(&adc_path).map_err(|e| {
|
||||
ProviderError::Connection(format!(
|
||||
"Failed to read ADC at {}: {}. Run 'gcloud auth application-default login'.",
|
||||
adc_path.display(),
|
||||
e
|
||||
))
|
||||
})?;
|
||||
|
||||
let adc: AdcCredentials = serde_json::from_str(&adc_json)
|
||||
.map_err(|e| ProviderError::Connection(format!("Failed to parse ADC: {}", e)))?;
|
||||
|
||||
// Use refresh token to get access token
|
||||
let resp = self
|
||||
.client
|
||||
.post("https://oauth2.googleapis.com/token")
|
||||
.form(&[
|
||||
("client_id", adc.client_id.as_str()),
|
||||
("client_secret", adc.client_secret.as_str()),
|
||||
("refresh_token", adc.refresh_token.as_str()),
|
||||
("grant_type", "refresh_token"),
|
||||
])
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| ProviderError::Connection(format!("ADC token refresh error: {}", e)))?;
|
||||
|
||||
let token_resp: GoogleTokenResponse = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| ProviderError::Connection(format!("Token parse error: {}", e)))?;
|
||||
|
||||
Ok((token_resp.access_token, token_resp.expires_in))
|
||||
}
|
||||
|
||||
async fn get_metadata_token(&self) -> Result<(String, u64), ProviderError> {
|
||||
let resp = self
|
||||
.client
|
||||
.get("http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token")
|
||||
.header("Metadata-Flavor", "Google")
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| ProviderError::Connection(format!("Metadata server error: {}", e)))?;
|
||||
|
||||
let token_resp: GoogleTokenResponse = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| ProviderError::Connection(format!("Token parse error: {}", e)))?;
|
||||
|
||||
Ok((token_resp.access_token, token_resp.expires_in))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LlmProvider for VertexProvider {
|
||||
async fn stream(
|
||||
&self,
|
||||
request: &LlmRequest,
|
||||
) -> Result<mpsc::Receiver<LlmEvent>, ProviderError> {
|
||||
let url = self.build_url(&request.model);
|
||||
let body = self.build_request_body(request);
|
||||
|
||||
tracing::debug!(target: "nomi_providers", body = %serde_json::to_string_pretty(&body).unwrap_or_default(), "outgoing request");
|
||||
|
||||
let access_token = self.get_access_token().await?;
|
||||
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
|
||||
headers.insert(
|
||||
AUTHORIZATION,
|
||||
HeaderValue::from_str(&format!("Bearer {}", access_token))
|
||||
.map_err(|e| ProviderError::Connection(format!("Header error: {}", e)))?,
|
||||
);
|
||||
|
||||
let response = crate::retry::with_initial_connect_retry(|| async {
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.headers(headers.clone())
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
let retry_after_ms = crate::parse_retry_after_ms(response.headers()).unwrap_or(5000);
|
||||
let body_text = response.text().await.unwrap_or_default();
|
||||
if status.as_u16() == 429 {
|
||||
return Err(ProviderError::RateLimited {
|
||||
retry_after_ms,
|
||||
message: crate::non_empty_rate_limit_message(body_text),
|
||||
});
|
||||
}
|
||||
return Err(ProviderError::Api {
|
||||
status: status.as_u16(),
|
||||
message: body_text,
|
||||
});
|
||||
}
|
||||
Ok(response)
|
||||
})
|
||||
.await?;
|
||||
|
||||
let (tx, rx) = mpsc::channel(64);
|
||||
let client = self.client.clone();
|
||||
let url_clone = url.clone();
|
||||
let headers_clone = {
|
||||
let mut h = HeaderMap::new();
|
||||
h.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
|
||||
h.insert(
|
||||
AUTHORIZATION,
|
||||
HeaderValue::from_str(&format!("Bearer {}", access_token))
|
||||
.map_err(|e| ProviderError::Connection(format!("Header error: {}", e)))?,
|
||||
);
|
||||
h
|
||||
};
|
||||
|
||||
// Vertex uses standard SSE (same as Anthropic)
|
||||
tokio::spawn(async move {
|
||||
match anthropic_shared::process_sse_stream(response, &tx).await {
|
||||
anthropic_shared::StreamOutcome::Ok => {}
|
||||
anthropic_shared::StreamOutcome::FailedPartial(e) => {
|
||||
let _ = tx.send(LlmEvent::Error(e.to_string())).await;
|
||||
}
|
||||
anthropic_shared::StreamOutcome::FailedEmpty(e) => {
|
||||
if e.is_retryable() {
|
||||
let mut backoff = std::time::Duration::from_secs(1);
|
||||
let mut final_err = Some(e);
|
||||
for attempt in 1..=crate::retry::MAX_STREAM_RETRIES {
|
||||
backoff = crate::retry::backoff_sleep(attempt, backoff).await;
|
||||
match crate::retry::send_and_check(
|
||||
&client,
|
||||
&url_clone,
|
||||
&headers_clone,
|
||||
&body,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(resp) => {
|
||||
let outcome =
|
||||
anthropic_shared::process_sse_stream(resp, &tx).await;
|
||||
match crate::retry::evaluate_outcome(outcome, attempt) {
|
||||
Ok(None) => {
|
||||
final_err = None;
|
||||
break;
|
||||
}
|
||||
Ok(Some(e)) => {
|
||||
final_err = Some(e);
|
||||
break;
|
||||
}
|
||||
Err(_) => continue,
|
||||
}
|
||||
}
|
||||
Err(e) if attempt == crate::retry::MAX_STREAM_RETRIES => {
|
||||
final_err = Some(e);
|
||||
break;
|
||||
}
|
||||
Err(_) => continue,
|
||||
}
|
||||
}
|
||||
if let Some(err) = final_err {
|
||||
let _ = tx.send(LlmEvent::Error(err.to_string())).await;
|
||||
}
|
||||
} else {
|
||||
let _ = tx.send(LlmEvent::Error(e.to_string())).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(rx)
|
||||
}
|
||||
}
|
||||
|
||||
// --- Internal types ---
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ServiceAccountKey {
|
||||
client_email: String,
|
||||
private_key: String,
|
||||
token_uri: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct JwtClaims {
|
||||
iss: String,
|
||||
scope: String,
|
||||
aud: String,
|
||||
iat: u64,
|
||||
exp: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct GoogleTokenResponse {
|
||||
access_token: String,
|
||||
#[serde(default = "default_expires_in")]
|
||||
expires_in: u64,
|
||||
}
|
||||
|
||||
fn default_expires_in() -> u64 {
|
||||
3600
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct AdcCredentials {
|
||||
client_id: String,
|
||||
client_secret: String,
|
||||
refresh_token: String,
|
||||
}
|
||||
|
||||
/// Build GcpAuth from nomi-config's VertexConfig
|
||||
pub fn auth_from_config(vc: &nomi_config::config::VertexConfig) -> GcpAuth {
|
||||
if let Some(creds_file) = &vc.credentials_file {
|
||||
GcpAuth::ServiceAccount {
|
||||
key_file: creds_file.clone(),
|
||||
}
|
||||
} else {
|
||||
GcpAuth::ApplicationDefault
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user