Compare commits
No commits in common. "53b4f572df9f8a9768a3fe2e7ef16001925e08d5" and "8f4af514a365018b8bc7e7acf386f33b197c2949" have entirely different histories.
53b4f572df
...
8f4af514a3
7 changed files with 72 additions and 82 deletions
49
Cargo.lock
generated
49
Cargo.lock
generated
|
|
@ -5495,9 +5495,9 @@ checksum = "47b34b781b31e5d73e9fbc8689c70551fd1ade9a19e3e28cfec8580a79290cc4"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "rig"
|
name = "rig"
|
||||||
version = "0.42.0"
|
version = "0.41.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "b989f4060e5f3c1ea4b8b6999f6759949d820704286cc6976b7a688cc0133029"
|
checksum = "2ce03971e6115d30ef53fb3244d06718a4e62bbde82c103065600c09459b989a"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"rig-agent",
|
"rig-agent",
|
||||||
"rig-core",
|
"rig-core",
|
||||||
|
|
@ -5510,9 +5510,9 @@ dependencies = [
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "rig-agent"
|
name = "rig-agent"
|
||||||
version = "0.42.0"
|
version = "0.41.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "b66477c0d0a786358d078e7e3541a73089dd9009100d6e2b3fe3b0d730e8550f"
|
checksum = "2b0796bbf47d7b76670401aac975bc619cf7fba3482b22dfe14992edaa9c2e04"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"async-stream",
|
"async-stream",
|
||||||
"fastrand",
|
"fastrand",
|
||||||
|
|
@ -5532,9 +5532,9 @@ dependencies = [
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "rig-core"
|
name = "rig-core"
|
||||||
version = "0.42.0"
|
version = "0.41.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "432d83e0facf16749f91fe729cbffca84437e8062d2f4e92f4f12e903693922d"
|
checksum = "35f5520515ae8f6851adcbc6fde9eea8e96f657418c062e16c82cd81cce44e8e"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"as-any",
|
"as-any",
|
||||||
"async-stream",
|
"async-stream",
|
||||||
|
|
@ -5556,7 +5556,6 @@ dependencies = [
|
||||||
"schemars 1.2.2",
|
"schemars 1.2.2",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
"sha2",
|
|
||||||
"thiserror",
|
"thiserror",
|
||||||
"tokio",
|
"tokio",
|
||||||
"tokio-tungstenite",
|
"tokio-tungstenite",
|
||||||
|
|
@ -5567,9 +5566,9 @@ dependencies = [
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "rig-derive"
|
name = "rig-derive"
|
||||||
version = "0.42.0"
|
version = "0.41.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "de0a33f1bac45f16e50146c248bcbbfaa44518c7252d274e972c7f4ad71aaba7"
|
checksum = "eb868fcebdf3ba425e3afad2e4926bb6d9e1188a856843b00bcee2e15c07424f"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"convert_case",
|
"convert_case",
|
||||||
"proc-macro-crate",
|
"proc-macro-crate",
|
||||||
|
|
@ -5580,9 +5579,9 @@ dependencies = [
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "rig-fastembed"
|
name = "rig-fastembed"
|
||||||
version = "0.42.0"
|
version = "0.41.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "f1dff99f9bea13108bb14df5431b6fc0738beb87f1fe5ebfcf26ae9fe54a188d"
|
checksum = "69ba93356242940918fac76a4b1cfe78b964f757e5cb65bb3feebd597635cfb0"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"fastembed",
|
"fastembed",
|
||||||
"rig-core",
|
"rig-core",
|
||||||
|
|
@ -5594,9 +5593,9 @@ dependencies = [
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "rig-helixdb"
|
name = "rig-helixdb"
|
||||||
version = "0.42.0"
|
version = "0.41.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "b5fa2b188a1e9e65d29c34738c97465f0822c57540d6d45fd4f04dad906c966a"
|
checksum = "d1d40ae8fa89af33793626b898ec527e246183ef8b4b5095afee5f6a65cbbb1a"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"reqwest 0.13.4",
|
"reqwest 0.13.4",
|
||||||
"rig-core",
|
"rig-core",
|
||||||
|
|
@ -5607,11 +5606,12 @@ dependencies = [
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "rig-lancedb"
|
name = "rig-lancedb"
|
||||||
version = "0.42.0"
|
version = "0.41.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "c5fe3555202c5bd4e30948414c2e234fddd9260d4520046733fcf0d58718f3a1"
|
checksum = "f703b9e081ee4d77e60953598a5071a5cff4b8b702a2f06d251310950b0dae7c"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"arrow-array",
|
"arrow-array",
|
||||||
|
"deranged",
|
||||||
"futures",
|
"futures",
|
||||||
"lancedb",
|
"lancedb",
|
||||||
"rig-core",
|
"rig-core",
|
||||||
|
|
@ -5621,9 +5621,9 @@ dependencies = [
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "rig-milvus"
|
name = "rig-milvus"
|
||||||
version = "0.42.0"
|
version = "0.41.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "af9c7b7ff7ed9c16da983a21f1e81f425c9805355427a92f4d2d97ba1fe2aa12"
|
checksum = "76662eb7f75de4e54b5f71e616430dc7f75d96d32f109d9435ad426b511eb565"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"reqwest 0.13.4",
|
"reqwest 0.13.4",
|
||||||
"rig-core",
|
"rig-core",
|
||||||
|
|
@ -6709,9 +6709,9 @@ dependencies = [
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "tokio-tungstenite"
|
name = "tokio-tungstenite"
|
||||||
version = "0.29.0"
|
version = "0.28.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "8f72a05e828585856dacd553fba484c242c46e391fb0e58917c942ee9202915c"
|
checksum = "d25a406cddcc431a75d3d9afc6a7c0f7428d4891dd973e4d54c56b46127bf857"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"futures-util",
|
"futures-util",
|
||||||
"log",
|
"log",
|
||||||
|
|
@ -6898,9 +6898,9 @@ checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "tungstenite"
|
name = "tungstenite"
|
||||||
version = "0.29.0"
|
version = "0.28.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "6c01152af293afb9c7c2a57e4b559c5620b421f6d133261c60dd2d0cdb38e6b8"
|
checksum = "8628dcc84e5a09eb3d8423d6cb682965dea9133204e8fb3efee74c2a0c259442"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"bytes",
|
"bytes",
|
||||||
"data-encoding",
|
"data-encoding",
|
||||||
|
|
@ -6912,6 +6912,7 @@ dependencies = [
|
||||||
"rustls-pki-types",
|
"rustls-pki-types",
|
||||||
"sha1",
|
"sha1",
|
||||||
"thiserror",
|
"thiserror",
|
||||||
|
"utf-8",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
|
|
@ -7032,6 +7033,12 @@ dependencies = [
|
||||||
"serde",
|
"serde",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "utf-8"
|
||||||
|
version = "0.7.6"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "utf8-ranges"
|
name = "utf8-ranges"
|
||||||
version = "1.0.5"
|
version = "1.0.5"
|
||||||
|
|
|
||||||
|
|
@ -16,7 +16,7 @@ clap = { version = "4", features = ["derive"] }
|
||||||
futures = "0.3"
|
futures = "0.3"
|
||||||
indicatif = "0.18.6"
|
indicatif = "0.18.6"
|
||||||
reqwest = { version = "0.13.4", features = ["query", "json"] }
|
reqwest = { version = "0.13.4", features = ["query", "json"] }
|
||||||
rig = "0.42.0"
|
rig = "0.41.0"
|
||||||
schemars = "1"
|
schemars = "1"
|
||||||
scraper = "0.27"
|
scraper = "0.27"
|
||||||
serde = { version = "1.0.229", features = ["derive"] }
|
serde = { version = "1.0.229", features = ["derive"] }
|
||||||
|
|
|
||||||
|
|
@ -149,11 +149,11 @@ mod tests {
|
||||||
fn assistant_text_ignores_tool_calls() {
|
fn assistant_text_ignores_tool_calls() {
|
||||||
let message = Message::Assistant {
|
let message = Message::Assistant {
|
||||||
id: None,
|
id: None,
|
||||||
content: vec![AssistantContent::tool_call(
|
content: rig::OneOrMany::one(AssistantContent::tool_call(
|
||||||
"call-1",
|
"call-1",
|
||||||
"search_web",
|
"search_web",
|
||||||
serde_json::json!({ "query": "test"}),
|
serde_json::json!({ "query": "test" }),
|
||||||
)],
|
)),
|
||||||
};
|
};
|
||||||
assert!(extract_assistant_text(&message).is_empty());
|
assert!(extract_assistant_text(&message).is_empty());
|
||||||
}
|
}
|
||||||
|
|
@ -166,13 +166,13 @@ mod tests {
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn tool_result_text_extracts_and_truncates() {
|
fn tool_result_text_extracts_and_truncates() {
|
||||||
let short = Message::tool_result("call-1", "call-1-name", "short result");
|
let short = Message::tool_result("call-1", "short result");
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
tool_result_text(&short, 100),
|
tool_result_text(&short, 100),
|
||||||
vec!["short result".to_string()]
|
vec!["short result".to_string()]
|
||||||
);
|
);
|
||||||
|
|
||||||
let long = Message::tool_result("call-2", "call-2-result", "0123456789");
|
let long = Message::tool_result("call-2", "0123456789");
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
tool_result_text(&long, 5),
|
tool_result_text(&long, 5),
|
||||||
vec!["01234 ...[truncated]".to_string()]
|
vec!["01234 ...[truncated]".to_string()]
|
||||||
|
|
@ -191,11 +191,11 @@ mod tests {
|
||||||
Message::system("preamble"),
|
Message::system("preamble"),
|
||||||
Message::Assistant {
|
Message::Assistant {
|
||||||
id: None,
|
id: None,
|
||||||
content: vec![AssistantContent::tool_call(
|
content: rig::OneOrMany::one(AssistantContent::tool_call(
|
||||||
"call-1",
|
"call-1",
|
||||||
"search_web",
|
"search_web",
|
||||||
serde_json::json!({ "query": "test" }),
|
serde_json::json!({ "query": "test" }),
|
||||||
)],
|
)),
|
||||||
},
|
},
|
||||||
];
|
];
|
||||||
|
|
||||||
|
|
@ -209,9 +209,9 @@ mod tests {
|
||||||
fn partial_findings_from_history_collects_assistant_text_and_tool_results_in_order() {
|
fn partial_findings_from_history_collects_assistant_text_and_tool_results_in_order() {
|
||||||
let history = vec![
|
let history = vec![
|
||||||
Message::assistant("Checking sources..."),
|
Message::assistant("Checking sources..."),
|
||||||
Message::tool_result("call-1", "call-1-name", "Result A [1]"),
|
Message::tool_result("call-1", "Result A [1]"),
|
||||||
Message::assistant("Cross-checking..."),
|
Message::assistant("Cross-checking..."),
|
||||||
Message::tool_result("call-2", "call-2-name", "Result B [2]"),
|
Message::tool_result("call-2", "Result B [2]"),
|
||||||
];
|
];
|
||||||
|
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
|
|
@ -233,11 +233,11 @@ mod tests {
|
||||||
fn transcript_lines_renders_tool_calls_with_their_arguments() {
|
fn transcript_lines_renders_tool_calls_with_their_arguments() {
|
||||||
let message = Message::Assistant {
|
let message = Message::Assistant {
|
||||||
id: None,
|
id: None,
|
||||||
content: vec![AssistantContent::tool_call(
|
content: rig::OneOrMany::one(AssistantContent::tool_call(
|
||||||
"call-1",
|
"call-1",
|
||||||
"search_web",
|
"search_web",
|
||||||
serde_json::json!({ "query": "uruguay senior engineers" }),
|
serde_json::json!({ "query": "uruguay senior engineers" }),
|
||||||
)],
|
)),
|
||||||
};
|
};
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
transcript_lines(&message),
|
transcript_lines(&message),
|
||||||
|
|
@ -247,7 +247,7 @@ mod tests {
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn transcript_lines_prefixes_and_truncates_tool_results() {
|
fn transcript_lines_prefixes_and_truncates_tool_results() {
|
||||||
let short = Message::tool_result("call-1", "call-1-name", "found via search [1]");
|
let short = Message::tool_result("call-1", "found via search [1]");
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
transcript_lines(&short),
|
transcript_lines(&short),
|
||||||
vec!["Result: found via search [1]".to_string()]
|
vec!["Result: found via search [1]".to_string()]
|
||||||
|
|
@ -264,13 +264,13 @@ mod tests {
|
||||||
let history = vec![
|
let history = vec![
|
||||||
Message::Assistant {
|
Message::Assistant {
|
||||||
id: None,
|
id: None,
|
||||||
content: vec![AssistantContent::tool_call(
|
content: rig::OneOrMany::one(AssistantContent::tool_call(
|
||||||
"call-1",
|
"call-1",
|
||||||
"search_web",
|
"search_web",
|
||||||
serde_json::json!({ "query": "test" }),
|
serde_json::json!({ "query": "test" }),
|
||||||
)],
|
)),
|
||||||
},
|
},
|
||||||
Message::tool_result("call-1", "call-1-name", "1. Example\n https://example.com\n snippet"),
|
Message::tool_result("call-1", "1. Example\n https://example.com\n snippet"),
|
||||||
Message::assistant("that source looks solid"),
|
Message::assistant("that source looks solid"),
|
||||||
];
|
];
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -43,17 +43,20 @@ pub(crate) async fn research(
|
||||||
findings = gathered.findings;
|
findings = gathered.findings;
|
||||||
incomplete = gathered.incomplete;
|
incomplete = gathered.incomplete;
|
||||||
|
|
||||||
// Prevent repeat iterations where researcher hits same dead end.
|
// The researcher ran out of turns mid-investigation rather than
|
||||||
|
// concluding on its own — another round would just repeat the same
|
||||||
|
// dead end, so stop and write up whatever was gathered.
|
||||||
if incomplete {
|
if incomplete {
|
||||||
tracing::info!(round, "researcher exhausted its turn budget");
|
tracing::info!(round, "researcher exhausted its turn budget");
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
|
|
||||||
let review = review::review_findings(&client, topic, &findings, show_progress).await?;
|
let review = review::review_findings(&client, topic, &findings, show_progress).await?;
|
||||||
|
let approved = review.approved;
|
||||||
|
|
||||||
tracing::info!(round, review.approved, "review verdict");
|
tracing::info!(round, approved, "review verdict");
|
||||||
|
|
||||||
if review.approved || round == MAX_RESEARCH_ROUNDS {
|
if approved || round == MAX_RESEARCH_ROUNDS {
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,4 @@
|
||||||
use rig::client::AgentClientExt;
|
use rig::client::AgentClientExt;
|
||||||
use rig::completion::TypedPrompt;
|
|
||||||
use rig::providers::ollama;
|
use rig::providers::ollama;
|
||||||
use rig::schemars::JsonSchema;
|
use rig::schemars::JsonSchema;
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
@ -26,18 +25,10 @@ pub(crate) struct Review {
|
||||||
pub(crate) gaps: String,
|
pub(crate) gaps: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Retries for a review that fails to come back as valid structured output
|
/// Uses rig's typed extractor — a forced tool call into a `submit(Review)`
|
||||||
/// (e.g. a transient network error), on top of the initial attempt.
|
/// schema — rather than parsing free-text output, so the verdict and its
|
||||||
const REVIEW_RETRIES: usize = 2;
|
/// two feedback fields always come back structured instead of relying on
|
||||||
|
/// scanning prose for a trailing yes/no.
|
||||||
/// Uses rig's typed-prompt API — `Native` structured output constraining the
|
|
||||||
/// model's own reply to the `Review` schema — rather than a forced tool call
|
|
||||||
/// or parsing free-text output. A forced tool call needs `tool_choice`
|
|
||||||
/// support, which Ollama doesn't have: the reviewer's small local model would
|
|
||||||
/// just answer in prose and never call the tool, so extraction only ever
|
|
||||||
/// exhausted its retries and errored out. `Native` mode instead uses Ollama's
|
|
||||||
/// own `format` JSON-schema constraint, which is honored regardless of
|
|
||||||
/// tool-calling ability.
|
|
||||||
#[tracing::instrument(skip(client, findings), fields(gen_ai.agent.name = "reviewer"))]
|
#[tracing::instrument(skip(client, findings), fields(gen_ai.agent.name = "reviewer"))]
|
||||||
pub(crate) async fn review_findings(
|
pub(crate) async fn review_findings(
|
||||||
client: &ollama::Client,
|
client: &ollama::Client,
|
||||||
|
|
@ -46,7 +37,7 @@ pub(crate) async fn review_findings(
|
||||||
show_progress: bool,
|
show_progress: bool,
|
||||||
) -> anyhow::Result<Review> {
|
) -> anyhow::Result<Review> {
|
||||||
let reviewer = client
|
let reviewer = client
|
||||||
.agent(REVIEWER_MODEL)
|
.extractor::<Review>(REVIEWER_MODEL)
|
||||||
.preamble(
|
.preamble(
|
||||||
"You are a skeptical fact-checker reviewing another researcher's notes before they \
|
"You are a skeptical fact-checker reviewing another researcher's notes before they \
|
||||||
get turned into a report. Approve only if every conclusion in the findings is \
|
get turned into a report. Approve only if every conclusion in the findings is \
|
||||||
|
|
@ -56,36 +47,20 @@ pub(crate) async fn review_findings(
|
||||||
Always separate the solid, well-supported findings from the gaps so a follow-up \
|
Always separate the solid, well-supported findings from the gaps so a follow-up \
|
||||||
research pass knows what to keep and what to dig into further.",
|
research pass knows what to keep and what to dig into further.",
|
||||||
)
|
)
|
||||||
|
.retries(2)
|
||||||
.build();
|
.build();
|
||||||
|
|
||||||
let spinner = Spinner::start(
|
let spinner = Spinner::start(
|
||||||
show_progress,
|
show_progress,
|
||||||
format!("{REVIEW_EMOJI} Reviewing findings..."),
|
format!("{REVIEW_EMOJI} Reviewing findings..."),
|
||||||
);
|
);
|
||||||
|
let review = reviewer
|
||||||
let prompt = format!("Topic: {topic}\n\nResearch findings to review:\n{findings}");
|
.extract(format!(
|
||||||
|
"Topic: {topic}\n\nResearch findings to review:\n{findings}"
|
||||||
let mut last_error = None;
|
))
|
||||||
let mut extracted_review = None;
|
.await?;
|
||||||
|
|
||||||
for attempt in 0..=REVIEW_RETRIES {
|
|
||||||
match reviewer.prompt_typed(prompt.clone()).await {
|
|
||||||
Ok(r) => {
|
|
||||||
extracted_review = Some(r);
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
Err(e) => {
|
|
||||||
tracing::warn!("Attempt {attempt} to extract JSON failed: {e:?}.");
|
|
||||||
last_error = Some(e);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
drop(spinner);
|
drop(spinner);
|
||||||
|
|
||||||
let review: Review = extracted_review
|
|
||||||
.ok_or_else(|| last_error.expect("loop always sets last_error on failure"))?;
|
|
||||||
|
|
||||||
tracing::info!(approved = review.approved, gaps = %review.gaps, "review complete");
|
tracing::info!(approved = review.approved, gaps = %review.gaps, "review complete");
|
||||||
|
|
||||||
Ok(review)
|
Ok(review)
|
||||||
|
|
|
||||||
|
|
@ -18,11 +18,14 @@ use std::io::Write;
|
||||||
/// covering the gap between the prompt being sent and generation starting
|
/// covering the gap between the prompt being sent and generation starting
|
||||||
/// (otherwise the terminal would go blank for however long that takes)
|
/// (otherwise the terminal would go blank for however long that takes)
|
||||||
/// rather than being dropped by the caller before this is even called.
|
/// rather than being dropped by the caller before this is even called.
|
||||||
pub(crate) async fn write_text_stream(
|
pub(crate) async fn write_text_stream<R>(
|
||||||
mut stream: impl Stream<Item = Result<MultiTurnStreamItem, StreamingError>> + Unpin,
|
mut stream: impl Stream<Item = Result<MultiTurnStreamItem<R>, StreamingError>> + Unpin,
|
||||||
writer: &mut impl Write,
|
writer: &mut impl Write,
|
||||||
mut spinner: Spinner,
|
mut spinner: Spinner,
|
||||||
) -> anyhow::Result<String> {
|
) -> anyhow::Result<String>
|
||||||
|
where
|
||||||
|
R: Clone,
|
||||||
|
{
|
||||||
let mut text = String::new();
|
let mut text = String::new();
|
||||||
|
|
||||||
while let Some(chunk) = stream.next().await {
|
while let Some(chunk) = stream.next().await {
|
||||||
|
|
@ -47,7 +50,10 @@ mod tests {
|
||||||
use rig::completion::CompletionError;
|
use rig::completion::CompletionError;
|
||||||
use rig::message::Text;
|
use rig::message::Text;
|
||||||
|
|
||||||
fn text_item(text: &str) -> Result<MultiTurnStreamItem, StreamingError> {
|
#[derive(Clone)]
|
||||||
|
struct DummyResponse;
|
||||||
|
|
||||||
|
fn text_item(text: &str) -> Result<MultiTurnStreamItem<DummyResponse>, StreamingError> {
|
||||||
Ok(MultiTurnStreamItem::StreamAssistantItem(
|
Ok(MultiTurnStreamItem::StreamAssistantItem(
|
||||||
StreamedAssistantContent::Text(Text::new(text)),
|
StreamedAssistantContent::Text(Text::new(text)),
|
||||||
))
|
))
|
||||||
|
|
@ -70,7 +76,7 @@ mod tests {
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn ignores_non_text_items() {
|
async fn ignores_non_text_items() {
|
||||||
let final_item = Ok(MultiTurnStreamItem::final_response(
|
let final_item = Ok(MultiTurnStreamItem::final_response(
|
||||||
vec![rig::message::AssistantContent::text("ignored")],
|
rig::OneOrMany::one(rig::message::AssistantContent::text("ignored")),
|
||||||
rig::completion::Usage::new(),
|
rig::completion::Usage::new(),
|
||||||
));
|
));
|
||||||
let items = vec![text_item("kept"), final_item];
|
let items = vec![text_item("kept"), final_item];
|
||||||
|
|
|
||||||
|
|
@ -44,7 +44,6 @@ pub(crate) async fn search_web(
|
||||||
) -> Result<String, ToolExecutionError> {
|
) -> Result<String, ToolExecutionError> {
|
||||||
progress::set_activity(format!("{SEARCH_EMOJI} Searching: {query}"));
|
progress::set_activity(format!("{SEARCH_EMOJI} Searching: {query}"));
|
||||||
|
|
||||||
// TODO re-use same client object?
|
|
||||||
let response = reqwest::Client::new()
|
let response = reqwest::Client::new()
|
||||||
.get(format!("{}/search", searxng_base_url()))
|
.get(format!("{}/search", searxng_base_url()))
|
||||||
.query(&[("q", query.as_str()), ("format", "json")])
|
.query(&[("q", query.as_str()), ("format", "json")])
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue