97 lines
3.1 KiB
Rust
97 lines
3.1 KiB
Rust
|
use futures::{Stream, StreamExt};
|
||
|
|
use rig::agent::{MultiTurnStreamItem, StreamingError};
|
||
|
|
use rig::streaming::StreamedAssistantContent;
|
||
|
|
use std::io::Write;
|
||
|
|
|
||
|
|
/// Drains a multi-turn prompt stream, writing each text chunk to `writer` as
|
||
|
|
/// it arrives and returning the full accumulated text. Non-text items (tool
|
||
|
|
/// calls, reasoning, completion-call metadata, ...) are ignored — this is
|
||
|
|
/// only concerned with the assistant's prose.
|
||
|
|
///
|
||
|
|
/// `writer` is taken by reference rather than locked internally so the
|
||
|
|
/// caller controls the lock's lifetime: locking once around a whole report
|
||
|
|
/// (as `write_report` does) avoids re-acquiring it on every chunk, the way
|
||
|
|
/// `print!` would.
|
||
|
|
pub(crate) async fn write_text_stream<R>(
|
||
|
|
mut stream: impl Stream<Item = Result<MultiTurnStreamItem<R>, StreamingError>> + Unpin,
|
||
|
|
writer: &mut impl Write,
|
||
|
|
) -> anyhow::Result<String>
|
||
|
|
where
|
||
|
|
R: Clone,
|
||
|
|
{
|
||
|
|
let mut text = String::new();
|
||
|
|
|
||
|
|
while let Some(chunk) = stream.next().await {
|
||
|
|
if let MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Text(chunk)) =
|
||
|
|
chunk?
|
||
|
|
{
|
||
|
|
write!(writer, "{}", chunk.text)?;
|
||
|
|
writer.flush()?;
|
||
|
|
text.push_str(&chunk.text);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
Ok(text)
|
||
|
|
}
|
||
|
|
|
||
|
|
#[cfg(test)]
|
||
|
|
mod tests {
|
||
|
|
use super::*;
|
||
|
|
use futures::stream;
|
||
|
|
use rig::completion::CompletionError;
|
||
|
|
use rig::message::Text;
|
||
|
|
|
||
|
|
#[derive(Clone)]
|
||
|
|
struct DummyResponse;
|
||
|
|
|
||
|
|
fn text_item(text: &str) -> Result<MultiTurnStreamItem<DummyResponse>, StreamingError> {
|
||
|
|
Ok(MultiTurnStreamItem::StreamAssistantItem(
|
||
|
|
StreamedAssistantContent::Text(Text::new(text)),
|
||
|
|
))
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn writes_and_accumulates_text_chunks() {
|
||
|
|
let items = vec![text_item("Hello, "), text_item("world!")];
|
||
|
|
let mut written = Vec::new();
|
||
|
|
|
||
|
|
let accumulated = write_text_stream(stream::iter(items), &mut written)
|
||
|
|
.await
|
||
|
|
.unwrap();
|
||
|
|
|
||
|
|
assert_eq!(accumulated, "Hello, world!");
|
||
|
|
assert_eq!(String::from_utf8(written).unwrap(), "Hello, world!");
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn ignores_non_text_items() {
|
||
|
|
let final_item = Ok(MultiTurnStreamItem::final_response(
|
||
|
|
rig::OneOrMany::one(rig::message::AssistantContent::text("ignored")),
|
||
|
|
rig::completion::Usage::new(),
|
||
|
|
));
|
||
|
|
let items = vec![text_item("kept"), final_item];
|
||
|
|
let mut written = Vec::new();
|
||
|
|
|
||
|
|
let accumulated = write_text_stream(stream::iter(items), &mut written)
|
||
|
|
.await
|
||
|
|
.unwrap();
|
||
|
|
|
||
|
|
assert_eq!(accumulated, "kept");
|
||
|
|
assert_eq!(String::from_utf8(written).unwrap(), "kept");
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn propagates_stream_errors() {
|
||
|
|
let error = StreamingError::Completion(CompletionError::RequestError(Box::new(
|
||
|
|
std::io::Error::other("boom"),
|
||
|
|
)));
|
||
|
|
let items = vec![text_item("kept"), Err(error)];
|
||
|
|
let mut written = Vec::new();
|
||
|
|
|
||
|
|
let result = write_text_stream(stream::iter(items), &mut written).await;
|
||
|
|
|
||
|
|
assert!(result.is_err());
|
||
|
|
assert_eq!(String::from_utf8(written).unwrap(), "kept");
|
||
|
|
}
|
||
|
|
}
|