Assistants stream their answers; the time limit is for silence 5436f0dd9429

/// The longest a streamed answer may take in all, however steadily it

/// comes: the limit a call is given is for silence (oxygenhub#248).

const STREAM_LIMIT: Duration = Duration::from_secs(15 * 60);

/// line of what the endpoint said, or why no answer came. The answer is

/// asked for as a stream, and `wait` is the longest the provider may be

/// silent, before it starts or between parts, so a long answer that keeps

/// coming is not cut off (oxygenhub#248); a provider that answers whole

/// instead is read whole within `wait`.

    ask_streaming(app, a, system, user, wait, |_| {}).await

}

/// [`ask`], handing each part of the answer to `part` as it arrives.

pub async fn ask_streaming(

    app: &Shared,

    a: &Assistant,

    system: &str,

    user: &str,

    wait: Duration,

    mut part: impl FnMut(&str),

) -> Result<String, String> {

    let client = crate::net::client(&a.endpoint, &app.config.outbound.allow, STREAM_LIMIT.max(wait)).await?;

                "stream": true,

                "stream": true,

    let quiet = || format!("no answer within {} seconds", wait.as_secs());

    let r = tokio::time::timeout(wait, req.json(&body).send()).await.map_err(|_| quiet())?.map_err(|e| {

            quiet()

    let streamed =

        r.headers().get(reqwest::header::CONTENT_TYPE).and_then(|v| v.to_str().ok()).is_some_and(|t| t.starts_with("text/event-stream"));

    if status.is_success() && streamed {

        let (answer, stop) = read_stream(&a.api, r, wait, &mut part).await.map_err(|e| scrub(&e, &token))?;

        return finish(answer, stop);

    }

    let text: String =

        tokio::time::timeout(wait, r.text()).await.map_err(|_| quiet())?.unwrap_or_default().chars().take(MAX_ANSWER).collect();

    if !answer.is_empty() {

        part(&answer);

    }

}

/// Read a streamed answer (server-sent events), handing out its text as it

/// comes; at most `quiet` between parts.

async fn read_stream(api: &str, mut r: reqwest::Response, quiet: Duration, part: &mut impl FnMut(&str)) -> Result<(String, Stop), String> {

    let mut events = Events::new(api);

    loop {

        let chunk = match tokio::time::timeout(quiet, r.chunk()).await {

            Err(_) if events.text.is_empty() => return Err(format!("no answer within {} seconds", quiet.as_secs())),

            // It went quiet part of the way through: what came is kept.

            Err(_) => return Ok((events.text, Stop::Quiet)),

            Ok(Err(e)) if events.text.is_empty() => return Err(e.to_string()),

            Ok(Err(_)) => return Ok((events.text, Stop::Quiet)),

            Ok(Ok(None)) => return Ok((events.text, events.stop.unwrap_or(Stop::Done))),

            Ok(Ok(Some(c))) => c,

        };

        let before = events.text.len();

        events.feed(&chunk)?;

        if events.text.len() > before {

            part(&events.text[before..]);

        }

        if events.done || events.text.len() >= MAX_ANSWER {

            return Ok((events.text, events.stop.unwrap_or(Stop::Done)));

        }

    }

}

/// Server-sent events of a streamed answer, as Anthropic's messages API

/// and the chat-completions APIs send them: text deltas, why it stopped,

/// and errors. Fed bytes as they come; a line is read once it is whole.

struct Events {

    anthropic: bool,

    buf: Vec<u8>,

    text: String,

    stop: Option<Stop>,

    done: bool,

}

impl Events {

    fn new(api: &str) -> Events {

        Events { anthropic: api == "anthropic-messages", buf: Vec::new(), text: String::new(), stop: None, done: false }

    }

    fn feed(&mut self, bytes: &[u8]) -> Result<(), String> {

        self.buf.extend_from_slice(bytes);

        while let Some(end) = self.buf.iter().position(|b| *b == b'\n') {

            let line: Vec<u8> = self.buf.drain(..=end).collect();

            let line = String::from_utf8_lossy(&line);

            let Some(data) = line.trim_end_matches(['\r', '\n']).strip_prefix("data:") else {

                continue;

            };

            let data = data.trim_start();

            if data == "[DONE]" {

                self.done = true;

                continue;

            }

            let Ok(v) = serde_json::from_str::<serde_json::Value>(data) else {

                continue;

            };

            self.event(&v)?;

        }

        Ok(())

    }

    fn event(&mut self, v: &serde_json::Value) -> Result<(), String> {

        if self.anthropic {

            match v["type"].as_str() {

                Some("content_block_delta") if v["delta"]["type"] == "text_delta" => {

                    self.text.push_str(v["delta"]["text"].as_str().unwrap_or(""));

                }

                Some("message_delta") => {

                    self.stop = match v["delta"]["stop_reason"].as_str() {

                        Some("max_tokens") => Some(Stop::Length),

                        Some("refusal") => Some(Stop::Declined(v["delta"]["stop_details"]["category"].as_str().map(str::to_string))),

                        Some(_) => Some(Stop::Done),

                        None => self.stop.take(),

                    };

                }

                Some("message_stop") => self.done = true,

                Some("error") => {

                    return Err(v["error"]["message"].as_str().unwrap_or("the provider reported an error").to_string());

                }

                _ => {}

            }

        } else {

            if let Some(e) = v["error"]["message"].as_str() {

                return Err(e.to_string());

            }

            let choice = &v["choices"][0];

            if let Some(t) = choice["delta"]["content"].as_str() {

                self.text.push_str(t);

            }

            match choice["finish_reason"].as_str() {

                Some("length") => self.stop = Some(Stop::Length),

                Some("content_filter") => self.stop = Some(Stop::Declined(None)),

                Some(_) => self.stop = Some(Stop::Done),

                None => {}

            }

        }

        Ok(())

    }

    /// The provider went silent part of the way through a stream.

    Quiet,

        Stop::Quiet if answer.trim().is_empty() => Err("it went quiet before writing an answer".into()),

        Stop::Quiet => Ok(format!("{}\n\n*(Cut short: the assistant went quiet part of the way through.)*", answer.trim_end())),

    /// Anthropic's events, split anywhere across chunks, give the text,

    /// leave thinking out, and say why it stopped (oxygenhub#248).

    #[test]

    fn anthropic_events_give_the_text_whole() {

        let sse = "event: message_start\ndata: {\"type\":\"message_start\"}\n\n\

                   data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"hmm\"}}\n\n\

                   data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"Adds a \"}}\n\n\

                   data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"line — café.\"}}\n\n\

                   data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}}\n\n\

                   data: {\"type\":\"message_stop\"}\n\n";

        for size in [1, 3, 7, 64, sse.len()] {

            let mut e = Events::new("anthropic-messages");

            for chunk in sse.as_bytes().chunks(size) {

                e.feed(chunk).unwrap();

            }

            assert_eq!((e.text.as_str(), e.stop.as_ref(), e.done), ("Adds a line — café.", Some(&Stop::Done), true), "chunks of {size}");

        }

        let mut cut = Events::new("anthropic-messages");

        cut.feed(b"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"max_tokens\"}}\n").unwrap();

        assert_eq!(cut.stop, Some(Stop::Length));

        let mut bad = Events::new("anthropic-messages");

        let e = bad.feed(b"data: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"Overloaded\"}}\n").unwrap_err();

        assert_eq!(e, "Overloaded");

    }

    #[test]

    fn chat_completion_events_give_the_text() {

        let sse = "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\n\

                   data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\

                   data: {\"choices\":[{\"delta\":{\"content\":\", world\"},\"finish_reason\":\"length\"}]}\n\n\

                   data: [DONE]\n\n";

        let mut e = Events::new("chat-completions");

        for chunk in sse.as_bytes().chunks(5) {

            e.feed(chunk).unwrap();

        }

        assert_eq!((e.text.as_str(), e.stop.as_ref(), e.done), ("Hello, world", Some(&Stop::Length), true));

    }

    /// A stream that takes longer than the silence limit in all, but never

    /// pauses that long, is read whole; one that stops talking is cut off

    /// with what it said kept.

    #[tokio::test]

    async fn a_slow_steady_stream_is_read_and_a_silent_one_cut() {

        async fn serve(parts: usize, gap_ms: u64, then_silent: bool) -> String {

            let router = axum::Router::new().route(

                "/v1/messages",

                axum::routing::post(move || async move {

                    let (tx, rx) = tokio::sync::mpsc::channel::<Result<String, std::io::Error>>(4);

                    tokio::spawn(async move {

                        for i in 0..parts {

                            tokio::time::sleep(Duration::from_millis(gap_ms)).await;

                            let ev = format!(

                                "data: {{\"type\":\"content_block_delta\",\"delta\":{{\"type\":\"text_delta\",\"text\":\"{i} \"}}}}\n\n"

                            );

                            if tx.send(Ok(ev)).await.is_err() {

                                return;

                            }

                        }

                        if then_silent {

                            tokio::time::sleep(Duration::from_secs(5)).await;

                        }

                        let _ = tx.send(Ok("data: {\"type\":\"message_stop\"}\n\n".into())).await;

                    });

                    let body = axum::body::Body::from_stream(tokio_stream::wrappers::ReceiverStream::new(rx));

                    ([(axum::http::header::CONTENT_TYPE, "text/event-stream")], body)

                }),

            );

            let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();

            let port = listener.local_addr().unwrap().port();

            tokio::spawn(async move { axum::serve(listener, router).await.unwrap() });

            format!("http://127.0.0.1:{port}/v1/messages")

        }

        let quiet = Duration::from_millis(400);

        // Six parts 200 ms apart: 1.2 s in all, three times the limit.

        let url = serve(6, 200, false).await;

        let r = reqwest::Client::new().post(&url).send().await.unwrap();

        let mut parts = Vec::new();

        let (text, stop) = read_stream("anthropic-messages", r, quiet, &mut |p: &str| parts.push(p.to_string())).await.unwrap();

        assert_eq!((text.as_str(), stop), ("0 1 2 3 4 5 ", Stop::Done));

        assert_eq!(parts.len(), 6, "each part handed out as it came");

        // Two parts, then silence: cut off, the two kept.

        let url = serve(2, 50, true).await;

        let r = reqwest::Client::new().post(&url).send().await.unwrap();

        let (text, stop) = read_stream("anthropic-messages", r, quiet, &mut |_: &str| {}).await.unwrap();

        assert_eq!((text.as_str(), stop), ("0 1 ", Stop::Quiet));

        assert!(finish(text, Stop::Quiet).unwrap().contains("went quiet"));

    }