use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use anyhow::{Result, anyhow};
use reqwest::Client;
use serde_json::Value;
use crate::api::http_proxy::resolve_http_proxy_url_for_target;
use crate::types::{
AssistantMessage, AssistantMessageEvent, Model, OnPayloadCallback, OnResponseCallback, ProviderEnv,
ProviderResponse, StopReason, StreamOptions,
};
use crate::utils::error_body::{error_body_from_response, format_provider_error, normalize_provider_error};
use crate::utils::event_stream::AssistantMessageEventStream;
use crate::utils::headers::{has_header, headers_to_record, merge_provider_headers};
pub fn build_http_client(timeout_ms: Option<u64>) -> Result<Client> {
build_http_client_for_target(timeout_ms, None, None)
}
pub fn build_http_client_for_target(
timeout_ms: Option<u64>,
target_url: Option<&str>,
env: Option<&ProviderEnv>,
) -> Result<Client> {
let mut builder = Client::builder();
if let Some(ms) = timeout_ms {
builder = builder.timeout(std::time::Duration::from_millis(ms));
}
if let Some(target_url) = target_url
&& let Some(proxy_url) = resolve_http_proxy_url_for_target(target_url, env)?
{
let proxy = reqwest::Proxy::all(proxy_url.as_str())?;
builder = builder.proxy(proxy);
}
Ok(builder.build()?)
}
pub fn get_client_api_key(provider: &str, api_key: Option<&str>, headers: &HashMap<String, String>) -> Result<String> {
if let Some(key) = api_key {
return Ok(key.to_string());
}
if has_header(headers, "authorization") || has_header(headers, "cf-aig-authorization") {
return Ok("unused".to_string());
}
Err(anyhow!("No API key for provider: {provider}"))
}
pub async fn apply_on_payload(callback: Option<&OnPayloadCallback>, payload: Value, model: &Model) -> Value {
if let Some(cb) = callback {
let m = model.clone();
let original = payload.clone();
if let Some(next) = cb(payload, m).await {
return next;
}
return original;
}
payload
}
pub async fn apply_on_response(callback: Option<&OnResponseCallback>, response: ProviderResponse, model: &Model) {
if let Some(cb) = callback {
let m = model.clone();
cb(response, m).await;
}
}
pub fn merge_model_headers(model: &Model, options: Option<&StreamOptions>) -> HashMap<String, String> {
let base = model.headers.clone().unwrap_or_default();
merge_provider_headers(&base, options.and_then(|o| o.headers.as_ref()))
}
pub const REQUEST_ABORTED: &str = "Request aborted";
pub fn is_request_aborted(token: &Option<tokio_util::sync::CancellationToken>) -> bool {
token.as_ref().is_some_and(|t| t.is_cancelled())
}
pub fn request_aborted_error() -> anyhow::Error {
anyhow!(REQUEST_ABORTED)
}
pub fn is_abort_error(error: &anyhow::Error) -> bool {
error.to_string() == REQUEST_ABORTED
}
pub async fn send_with_abort(
token: &Option<tokio_util::sync::CancellationToken>,
request: reqwest::RequestBuilder,
) -> Result<reqwest::Response> {
if is_request_aborted(token) {
return Err(request_aborted_error());
}
match token {
Some(token) => {
let token = token.clone();
tokio::select! {
result = request.send() => result.map_err(Into::into),
_ = token.cancelled() => Err(request_aborted_error()),
}
}
None => request.send().await.map_err(Into::into),
}
}
pub fn finish_stream_error(
stream: &AssistantMessageEventStream,
output: &mut AssistantMessage,
error: anyhow::Error,
aborted: bool,
) {
output.stop_reason = if aborted {
StopReason::Aborted
} else {
StopReason::Error
};
output.error_message = Some(format_provider_error(&normalize_provider_error(&error), None));
stream.push(AssistantMessageEvent::Error {
reason: output.stop_reason,
error: output.clone(),
});
stream.end();
}
pub async fn check_response_ok(response: reqwest::Response) -> Result<reqwest::Response> {
if response.status().is_success() {
return Ok(response);
}
let status = response.status();
let body = error_body_from_response(response).await;
Err(anyhow!("{status}: {body}"))
}
pub type StreamTask = Pin<Box<dyn Future<Output = ()> + Send>>;
pub fn spawn_stream_task(fut: impl Future<Output = ()> + Send + 'static) -> StreamTask {
Box::pin(async move {
tokio::spawn(fut);
})
}
pub fn wrap_on_payload<F>(f: F) -> OnPayloadCallback
where
F: Fn(Value, Model) -> Pin<Box<dyn Future<Output = Option<Value>> + Send>> + Send + Sync + 'static,
{
Arc::new(f)
}
pub fn wrap_on_response<F>(f: F) -> OnResponseCallback
where
F: Fn(ProviderResponse, Model) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync + 'static,
{
Arc::new(f)
}
pub async fn invoke_on_response_from_reqwest(
callback: Option<&OnResponseCallback>,
response: &reqwest::Response,
model: &Model,
) {
let provider_response = ProviderResponse {
status: response.status().as_u16(),
headers: headers_to_record(response.headers()),
};
apply_on_response(callback, provider_response, model).await;
}