use std::io::BufReader;
use std::time::Duration;
use serde_json::Value;
use super::reply::{self, Delta, Reply};
use super::{Key, ProviderError, error};
const CONNECT: Duration = Duration::from_secs(30);
const FIRST_BYTE: Duration = Duration::from_secs(300);
const API_VERSION: &str = "2023-06-01";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Endpoint {
pub address: String,
pub model: String,
pub key_header: String,
pub key: Key,
}
impl Endpoint {
#[must_use]
pub fn url(&self) -> String {
format!("{}/v1/messages", self.address.trim().trim_end_matches('/'))
}
pub fn stream(
&self,
body: &Value,
on_delta: &mut dyn FnMut(Delta),
cancelled: &dyn Fn() -> bool,
) -> Result<Reply, ProviderError> {
let url = self.url();
let mut body = body.clone();
if let Some(object) = body.as_object_mut() {
object.insert("model".to_owned(), Value::String(self.model.clone()));
object.insert("stream".to_owned(), Value::Bool(true));
}
let config = ureq::Agent::config_builder()
.timeout_connect(Some(CONNECT))
.timeout_recv_response(Some(FIRST_BYTE))
.http_status_as_error(false)
.build();
let agent = ureq::Agent::new_with_config(config);
let header = self.key_header.trim();
let secret = if header.eq_ignore_ascii_case("authorization") {
format!("Bearer {}", self.key.expose())
} else {
self.key.expose().to_owned()
};
let unreachable =
|error: ureq::Error| ProviderError::Unreachable { url: url.clone(), reason: error.to_string() };
let mut answer = agent
.post(&url)
.header("anthropic-version", API_VERSION)
.header("accept", "text/event-stream")
.header(header, &secret)
.content_type("application/json")
.send(body.to_string().as_str())
.map_err(unreachable)?;
let status = answer.status().as_u16();
if !(200..300).contains(&status) {
let said = answer.body_mut().read_to_string().unwrap_or_default();
return Err(ProviderError::Refused { url, status, said: error::complaint(&said) });
}
let lines = BufReader::new(answer.into_body().into_reader());
reply::read(lines, on_delta, cancelled)
}
}