elph_ai/api/
common.rs

1use std::collections::HashMap;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::Arc;
5
6use anyhow::{Result, anyhow};
7use reqwest::Client;
8use serde_json::Value;
9
10use crate::api::http_proxy::resolve_http_proxy_url_for_target;
11use crate::types::{
12    AssistantMessage, AssistantMessageEvent, Model, OnPayloadCallback, OnResponseCallback, ProviderEnv,
13    ProviderResponse, StopReason, StreamOptions,
14};
15use crate::utils::error_body::{error_body_from_response, format_provider_error, normalize_provider_error};
16use crate::utils::event_stream::AssistantMessageEventStream;
17use crate::utils::headers::{has_header, headers_to_record, merge_provider_headers};
18
19pub fn build_http_client(timeout_ms: Option<u64>) -> Result<Client> {
20    build_http_client_for_target(timeout_ms, None, None)
21}
22
23pub fn build_http_client_for_target(
24    timeout_ms: Option<u64>,
25    target_url: Option<&str>,
26    env: Option<&ProviderEnv>,
27) -> Result<Client> {
28    let mut builder = Client::builder();
29    if let Some(ms) = timeout_ms {
30        builder = builder.timeout(std::time::Duration::from_millis(ms));
31    }
32    if let Some(target_url) = target_url
33        && let Some(proxy_url) = resolve_http_proxy_url_for_target(target_url, env)?
34    {
35        let proxy = reqwest::Proxy::all(proxy_url.as_str())?;
36        builder = builder.proxy(proxy);
37    }
38    Ok(builder.build()?)
39}
40
41pub fn get_client_api_key(provider: &str, api_key: Option<&str>, headers: &HashMap<String, String>) -> Result<String> {
42    if let Some(key) = api_key {
43        return Ok(key.to_string());
44    }
45    if has_header(headers, "authorization") || has_header(headers, "cf-aig-authorization") {
46        return Ok("unused".to_string());
47    }
48    Err(anyhow!("No API key for provider: {provider}"))
49}
50
51pub async fn apply_on_payload(callback: Option<&OnPayloadCallback>, payload: Value, model: &Model) -> Value {
52    if let Some(cb) = callback {
53        let m = model.clone();
54        let original = payload.clone();
55        if let Some(next) = cb(payload, m).await {
56            return next;
57        }
58        return original;
59    }
60    payload
61}
62
63pub async fn apply_on_response(callback: Option<&OnResponseCallback>, response: ProviderResponse, model: &Model) {
64    if let Some(cb) = callback {
65        let m = model.clone();
66        cb(response, m).await;
67    }
68}
69
70pub fn merge_model_headers(model: &Model, options: Option<&StreamOptions>) -> HashMap<String, String> {
71    let base = model.headers.clone().unwrap_or_default();
72    merge_provider_headers(&base, options.and_then(|o| o.headers.as_ref()))
73}
74
75pub const REQUEST_ABORTED: &str = "Request aborted";
76
77pub fn is_request_aborted(token: &Option<tokio_util::sync::CancellationToken>) -> bool {
78    token.as_ref().is_some_and(|t| t.is_cancelled())
79}
80
81pub fn request_aborted_error() -> anyhow::Error {
82    anyhow!(REQUEST_ABORTED)
83}
84
85pub fn is_abort_error(error: &anyhow::Error) -> bool {
86    error.to_string() == REQUEST_ABORTED
87}
88
89pub async fn send_with_abort(
90    token: &Option<tokio_util::sync::CancellationToken>,
91    request: reqwest::RequestBuilder,
92) -> Result<reqwest::Response> {
93    if is_request_aborted(token) {
94        return Err(request_aborted_error());
95    }
96    match token {
97        Some(token) => {
98            let token = token.clone();
99            tokio::select! {
100                result = request.send() => result.map_err(Into::into),
101                _ = token.cancelled() => Err(request_aborted_error()),
102            }
103        }
104        None => request.send().await.map_err(Into::into),
105    }
106}
107
108pub fn finish_stream_error(
109    stream: &AssistantMessageEventStream,
110    output: &mut AssistantMessage,
111    error: anyhow::Error,
112    aborted: bool,
113) {
114    output.stop_reason = if aborted {
115        StopReason::Aborted
116    } else {
117        StopReason::Error
118    };
119    output.error_message = Some(format_provider_error(&normalize_provider_error(&error), None));
120    stream.push(AssistantMessageEvent::Error {
121        reason: output.stop_reason,
122        error: output.clone(),
123    });
124    stream.end();
125}
126
127pub async fn check_response_ok(response: reqwest::Response) -> Result<reqwest::Response> {
128    if response.status().is_success() {
129        return Ok(response);
130    }
131    let status = response.status();
132    let body = error_body_from_response(response).await;
133    Err(anyhow!("{status}: {body}"))
134}
135
136pub type StreamTask = Pin<Box<dyn Future<Output = ()> + Send>>;
137
138pub fn spawn_stream_task(fut: impl Future<Output = ()> + Send + 'static) -> StreamTask {
139    Box::pin(async move {
140        tokio::spawn(fut);
141    })
142}
143
144pub fn wrap_on_payload<F>(f: F) -> OnPayloadCallback
145where
146    F: Fn(Value, Model) -> Pin<Box<dyn Future<Output = Option<Value>> + Send>> + Send + Sync + 'static,
147{
148    Arc::new(f)
149}
150
151pub fn wrap_on_response<F>(f: F) -> OnResponseCallback
152where
153    F: Fn(ProviderResponse, Model) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync + 'static,
154{
155    Arc::new(f)
156}
157
158pub async fn invoke_on_response_from_reqwest(
159    callback: Option<&OnResponseCallback>,
160    response: &reqwest::Response,
161    model: &Model,
162) {
163    let provider_response = ProviderResponse {
164        status: response.status().as_u16(),
165        headers: headers_to_record(response.headers()),
166    };
167    apply_on_response(callback, provider_response, model).await;
168}
elph_ai/api/common.rs

elph_ai/api/
common.rs
