headless_lms_chatbot/azure_chatbot/azure/
transport.rs

1//! Opening the HTTP stream to Azure and handing it on as lines.
2
3use std::pin::Pin;
4
5use bytes::Bytes;
6use futures::StreamExt;
7use futures::stream::BoxStream;
8use headless_lms_base::config::ApplicationConfiguration;
9use tokio::io::AsyncBufReadExt;
10use tokio_stream::wrappers::LinesStream;
11use tokio_util::io::StreamReader;
12use tracing::trace;
13
14use super::protocol::LLMRequest;
15use crate::chatbot_error::ChatbotResult;
16use crate::llm_utils::make_streaming_llm_request;
17use crate::prelude::*;
18
19/// The lines of an Azure response body, as the stream parsers read them.
20pub(crate) type ResponseLinesStream<'a> =
21    Pin<Box<LinesStream<StreamReader<BoxStream<'a, Result<Bytes, std::io::Error>>, Bytes>>>>;
22pub(crate) enum ResponseStreamType<'a> {
23    ToolCall(ResponseLinesStream<'a>),
24    TextResponse(ResponseLinesStream<'a>),
25}
26
27/// Wraps a byte stream as the lines the SSE parsers read, shared by production and tests so a
28/// change to how lines are framed cannot drift between them.
29pub(crate) fn lines_from_byte_stream(
30    stream: BoxStream<'_, Result<Bytes, std::io::Error>>,
31) -> ResponseLinesStream<'_> {
32    Box::pin(LinesStream::new(StreamReader::new(stream).lines()))
33}
34
35/// Makes a request to Azure and returns the resulting stream.
36pub(crate) async fn make_request_and_create_stream<'a>(
37    chat_request: &LLMRequest,
38    app_config: &ApplicationConfiguration,
39) -> ChatbotResult<ResponseLinesStream<'a>> {
40    let response = make_streaming_llm_request(chat_request, app_config).await?;
41
42    trace!("Receiving chat response with {:?}", response.version());
43
44    // Replaces the client-wide read timeout, which reqwest arms once per request rather than per
45    // chunk, so it cannot tell a stalled stream from a slow but healthy one.
46    let stream = tokio_stream::StreamExt::timeout(response.bytes_stream(), STREAM_IDLE_TIMEOUT)
47        .map(|chunk| match chunk {
48            Ok(bytes) => bytes.map_err(std::io::Error::other),
49            Err(_) => Err(std::io::Error::new(
50                std::io::ErrorKind::TimedOut,
51                format!(
52                    "The LLM stream sent nothing for {} seconds",
53                    STREAM_IDLE_TIMEOUT.as_secs()
54                ),
55            )),
56        })
57        .boxed();
58
59    Ok(lines_from_byte_stream(stream))
60}
headless_lms_chatbot/azure_chatbot/azure/transport.rs

headless_lms_chatbot/azure_chatbot/azure/
transport.rs
