quvyta-cli 0.1.0-alpha.5

//! Where requests go and how they are carried there.

use std::io::{self, BufRead, BufReader};
use std::sync::mpsc;
use std::thread;
use std::time::{Duration, Instant};

use serde_json::Value;

use super::reply::{self, Delta, Reply};
use super::{Key, ProviderError, error};

/// How long connecting may take. A server that does not answer within this is not there.
const CONNECT: Duration = Duration::from_secs(30);

/// How long the server may think before it starts answering. Generous rather than tight: a large
/// conversation on a busy server takes a while before the first byte, and a short limit would
/// call a working server broken. Finite, because a reply that waits forever is a stuck screen.
const FIRST_BYTE: Duration = Duration::from_secs(300);

/// The longest one reply may take to arrive, pieces and pauses together. Far longer than any
/// reply takes; it only bounds a reading thread left behind by a server that never closes.
const BODY_LIMIT: Duration = Duration::from_secs(60 * 60);

/// The version of the Messages API this module speaks.
const API_VERSION: &str = "2023-06-01";

/// A provider speaking the Anthropic Messages shape.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Endpoint {
    /// The address the API lives under, such as `https://api.example.com/anthropic`; requests go
    /// to `{address}/v1/messages`.
    pub address: String,
    /// The model asked for.
    pub model: String,
    /// The header the key goes in: `x-api-key` for most, another name for some servers. For
    /// `authorization` the key goes in as a bearer token.
    pub key_header: String,
    /// The key, or `None` for a relay on this machine that adds the key itself: QCode's, inside a
    /// profile's container, which the key never enters. Without one no key header is sent at all.
    pub key: Option<Key>,
    /// How long the server may send nothing in the middle of a reply before it is given up on.
    pub silence: Duration,
}

impl Endpoint {
    /// How long a server may go quiet in the middle of a reply: generous, since a model can think
    /// a while between two pieces, and finite, since a reply that waits for ever is a stuck screen.
    pub const SILENCE: Duration = Duration::from_secs(120);

    /// Where a message request goes. The key is never part of it, so it is safe to show.
    #[must_use]
    pub fn url(&self) -> String {
        format!("{}/v1/messages", self.address.trim().trim_end_matches('/'))
    }

    /// Sends `body` (a Messages request without `model` and `stream`, which are added here) and
    /// reads the streamed reply, handing each piece to `on_delta` as it arrives. `cancelled` is
    /// asked between events.
    ///
    /// # Errors
    ///
    /// When the server cannot be reached, refuses, or answers with something that is not a
    /// Messages reply. No error carries the key.
    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))
            // A reply is given up on after a silence (see `Watched`); this bounds how long the
            // thread that was reading it can still wait on a server that never closes.
            .timeout_recv_body(Some(BODY_LIMIT))
            // A refusal is an answer, and its body says what the server objected to; it would be
            // lost if the status alone became an error.
            .http_status_as_error(false)
            .build();
        let agent = ureq::Agent::new_with_config(config);
        let header = self.key_header.trim();
        let key = self.key.as_ref().map_or("", Key::expose);
        let secret =
            if header.eq_ignore_ascii_case("authorization") { format!("Bearer {key}") } else { key.to_owned() };
        // What the transport says can name the address; it never names a header value, so it
        // cannot name the key.
        let unreachable =
            |error: ureq::Error| ProviderError::Unreachable { url: url.clone(), reason: error.to_string() };
        let mut request =
            agent.post(&url).header("anthropic-version", API_VERSION).header("accept", "text/event-stream");
        if self.key.is_some() {
            request = request.header(header, &secret);
        }
        let mut answer =
            request.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();
            let refused = ProviderError::Refused { url, status, said: error::complaint(&said) };
            return Err(refused.without(key));
        }
        let lines = Watched::start(BufReader::new(answer.into_body().into_reader()), self.silence, cancelled);
        reply::read(lines, on_delta, cancelled).map_err(|error| error.without(key))
    }
}

/// How often a reply that waits for the server asks whether it was stopped.
const POLL: Duration = Duration::from_millis(100);

/// The lines of a reply, read on a thread of their own so that waiting for the next one can be
/// given up: when the person stops the reply, or when the server has sent nothing for the
/// endpoint's silence. The reading thread ends at the next line after it is left, or when the
/// connection ends; [`BODY_LIMIT`] bounds how long that can take.
struct Watched<'a> {
    lines: mpsc::Receiver<io::Result<String>>,
    silence: Duration,
    cancelled: &'a dyn Fn() -> bool,
}

impl<'a> Watched<'a> {
    fn start(reader: impl BufRead + Send + 'static, silence: Duration, cancelled: &'a dyn Fn() -> bool) -> Self {
        // A few lines ahead is enough to keep the connection moving; the channel does not grow
        // without end when the screen is slower than the server.
        let (sender, lines) = mpsc::sync_channel(64);
        thread::spawn(move || {
            for line in reader.lines() {
                if sender.send(line).is_err() {
                    return;
                }
            }
        });
        Self { lines, silence, cancelled }
    }
}

impl Iterator for Watched<'_> {
    type Item = io::Result<String>;

    fn next(&mut self) -> Option<Self::Item> {
        let waiting = Instant::now();
        loop {
            if (self.cancelled)() {
                return None;
            }
            let left = self.silence.saturating_sub(waiting.elapsed());
            if left.is_zero() {
                let said =
                    format!("the provider sent nothing for {} seconds and was given up on", self.silence.as_secs());
                return Some(Err(io::Error::new(io::ErrorKind::TimedOut, said)));
            }
            match self.lines.recv_timeout(left.min(POLL)) {
                Ok(line) => return Some(line),
                Err(mpsc::RecvTimeoutError::Timeout) => {}
                Err(mpsc::RecvTimeoutError::Disconnected) => return None,
            }
        }
    }
}