f2c10783db
deep_research is the only project this repo is meant to showcase, so the
Cargo workspace wrapping it and an unrelated side project no longer earns
its keep:
- swear_cleanup moved to a new standalone local repo (~/dev/swear_cleanup,
not pushed anywhere) via `git subtree split`, with its pre-workspace-
split history (when it lived at src/swear_cleanup/ in a single shared
crate) spliced onto its post-split history rather than starting from a
single flattened snapshot. FINDINGS.md, which was sitting at this repo's
root but was actually swear_cleanup's own build log, went with it.
- deep_research/{src,Cargo.toml,README.md,docs} moved to the repo root;
the [workspace] table collapsed into a plain [package] manifest with
dependency versions inlined from the old [workspace.dependencies].
- Cargo.toml keeps an explicit empty [workspace] table (not just omitted)
so that checking this repo out as a nested git worktree — this
project's own normal workflow — can't accidentally inherit a stale
ancestor directory's workspace manifest, which is exactly what broke
the build while testing this change from a worktree.
- .forgejo/workflows/deep_research-ci.yml -> ci.yml, dropping the now-
meaningless -p deep_research scoping and path filters (redundant when
it's the only thing in the repo).
- README.md and docs/case-study.md updated for the flattened commands
(cargo run/test with no -p flag); their relative links to each other
and to src/ were already correct since both moved together.
Verified: cargo build/test/clippy/fmt all clean from the new repo root.
108 lines
3.6 KiB
Rust
108 lines
3.6 KiB
Rust
use crate::progress::Spinner;
|
|
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.
|
|
///
|
|
/// `spinner` stays up until the stream actually produces its first item,
|
|
/// covering the gap between the prompt being sent and generation starting
|
|
/// (otherwise the terminal would go blank for however long that takes)
|
|
/// rather than being dropped by the caller before this is even called.
|
|
pub(crate) async fn write_text_stream<R>(
|
|
mut stream: impl Stream<Item = Result<MultiTurnStreamItem<R>, StreamingError>> + Unpin,
|
|
writer: &mut impl Write,
|
|
mut spinner: Spinner,
|
|
) -> anyhow::Result<String>
|
|
where
|
|
R: Clone,
|
|
{
|
|
let mut text = String::new();
|
|
|
|
while let Some(chunk) = stream.next().await {
|
|
spinner.stop();
|
|
|
|
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, Spinner::start(false, ""))
|
|
.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, Spinner::start(false, ""))
|
|
.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, Spinner::start(false, "")).await;
|
|
|
|
assert!(result.is_err());
|
|
assert_eq!(String::from_utf8(written).unwrap(), "kept");
|
|
}
|
|
}
|