mirror of
https://github.com/fabro-sh/fabro.git
synced 2026-09-15 23:32:46 +00:00
Take environment helpers, compaction accounting, and search providers from pebble
Steps 4, 6, and 7 of .ai/plans/pebble-absorbs-embedder-concerns.md, pinning pebble fc907a1 with its `search-providers` feature. The sandbox's `Environment` adapter uses pebble's `environment::support` for the glob grammar check, the tree-order sort of a listing, and the capture accounting, in place of its own copies; the adapter itself stays. The stage reads the compactions a prompt performed from the report, as a breakdown of the usage it already billed. `web_search.rs` keeps the vault-backed credentials and the Brave-over-Venice preference, and hands pebble's `Brave` or `Venice` provider a fabro HTTP client; the providers and their tests are pebble's now. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Bryan Helmkamp
parent
d4e483925c
commit
6592505d0e
6 changed files with 87 additions and 441 deletions
26
Cargo.lock
generated
26
Cargo.lock
generated
|
|
@ -2063,7 +2063,7 @@ dependencies = [
|
|||
"libc",
|
||||
"option-ext",
|
||||
"redox_users",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -2190,7 +2190,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4345,7 +4345,7 @@ dependencies = [
|
|||
"js-sys",
|
||||
"log",
|
||||
"wasm-bindgen",
|
||||
"windows-core 0.61.2",
|
||||
"windows-core 0.62.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -5291,7 +5291,7 @@ version = "0.50.3"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5"
|
||||
dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -5856,7 +5856,7 @@ dependencies = [
|
|||
[[package]]
|
||||
name = "pebble-agent"
|
||||
version = "0.1.0"
|
||||
source = "git+https://github.com/lithoscomputer/pebble?rev=49da13798c753eb312b765596840bdc200b819d7#49da13798c753eb312b765596840bdc200b819d7"
|
||||
source = "git+https://github.com/lithoscomputer/pebble?rev=fc907a1b7df9bba8eff3bb0a79ab2d089b31b990#fc907a1b7df9bba8eff3bb0a79ab2d089b31b990"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"futures-util",
|
||||
|
|
@ -5873,7 +5873,7 @@ dependencies = [
|
|||
[[package]]
|
||||
name = "pebble-coding-agent"
|
||||
version = "0.1.0"
|
||||
source = "git+https://github.com/lithoscomputer/pebble?rev=49da13798c753eb312b765596840bdc200b819d7#49da13798c753eb312b765596840bdc200b819d7"
|
||||
source = "git+https://github.com/lithoscomputer/pebble?rev=fc907a1b7df9bba8eff3bb0a79ab2d089b31b990#fc907a1b7df9bba8eff3bb0a79ab2d089b31b990"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"futures-util",
|
||||
|
|
@ -6243,7 +6243,7 @@ dependencies = [
|
|||
"once_cell",
|
||||
"socket2",
|
||||
"tracing",
|
||||
"windows-sys 0.60.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -6720,7 +6720,7 @@ dependencies = [
|
|||
"errno 0.3.14",
|
||||
"libc",
|
||||
"linux-raw-sys",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -6779,7 +6779,7 @@ dependencies = [
|
|||
"security-framework",
|
||||
"security-framework-sys",
|
||||
"webpki-root-certs",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -7423,7 +7423,7 @@ version = "1.4.8"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b"
|
||||
dependencies = [
|
||||
"errno 0.2.8",
|
||||
"errno 0.3.14",
|
||||
"libc",
|
||||
]
|
||||
|
||||
|
|
@ -7988,7 +7988,7 @@ dependencies = [
|
|||
"getrandom 0.4.1",
|
||||
"once_cell",
|
||||
"rustix",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -8023,7 +8023,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "230a1b821ccbd75b185820a1f1ff7b14d21da1e442e22c0863ea5f08771a8874"
|
||||
dependencies = [
|
||||
"rustix",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -9090,7 +9090,7 @@ version = "0.1.11"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
|
||||
dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
|
|||
|
|
@ -112,8 +112,8 @@ futures-util = "0.3"
|
|||
# the merge commit once it lands. Pebble pins the same lithos-llm rev as
|
||||
# fabro, and its lockfile policy is that every shared crate resolves to the
|
||||
# version lithos-llm locks.
|
||||
pebble-agent = { git = "https://github.com/lithoscomputer/pebble", rev = "49da13798c753eb312b765596840bdc200b819d7" }
|
||||
pebble-coding-agent = { git = "https://github.com/lithoscomputer/pebble", rev = "49da13798c753eb312b765596840bdc200b819d7", features = ["mcp"] }
|
||||
pebble-agent = { git = "https://github.com/lithoscomputer/pebble", rev = "fc907a1b7df9bba8eff3bb0a79ab2d089b31b990" }
|
||||
pebble-coding-agent = { git = "https://github.com/lithoscomputer/pebble", rev = "fc907a1b7df9bba8eff3bb0a79ab2d089b31b990", features = ["mcp", "search-providers"] }
|
||||
sandbox-driver = { git = "https://github.com/lithoscomputer/sandbox-driver", rev = "23062b6ad62ff4665cbbcb7dce037ec9c4c34318" }
|
||||
sandbox-driver-protocol = { git = "https://github.com/lithoscomputer/sandbox-driver", rev = "23062b6ad62ff4665cbbcb7dce037ec9c4c34318" }
|
||||
sandbox-driver-host = { git = "https://github.com/lithoscomputer/sandbox-driver", rev = "23062b6ad62ff4665cbbcb7dce037ec9c4c34318" }
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@ use fabro_types::settings::run::ResolvedMcpEntry;
|
|||
use fabro_util::exit::{self, ErrorExt, ExitClass};
|
||||
use fabro_util::home::Home;
|
||||
use fabro_util::terminal::Styles;
|
||||
use fabro_workflow::web_search::{SearchBackend, SearchSecrets};
|
||||
use fabro_workflow::web_search::{self, SearchSecrets};
|
||||
use lithos_llm::catalog::ProviderId;
|
||||
use pebble_agent::{ToolCallRequest, ToolSystemError};
|
||||
use pebble_coding_agent::environment::Environment;
|
||||
|
|
@ -636,8 +636,8 @@ async fn run_session(
|
|||
if let Some(routes) = sandbox.port_routes() {
|
||||
builder = builder.port_routes(routes);
|
||||
}
|
||||
if let Some(search) = SearchBackend::from_secrets(&cli_search_secrets()) {
|
||||
builder = builder.search_provider(Arc::new(search));
|
||||
if let Some(search) = web_search::search_provider(&cli_search_secrets()) {
|
||||
builder = builder.search_provider(search);
|
||||
}
|
||||
let mut agent = builder
|
||||
.build()
|
||||
|
|
|
|||
|
|
@ -16,12 +16,11 @@ use std::sync::Arc;
|
|||
|
||||
use async_trait::async_trait;
|
||||
use fabro_types::CommandOutputStream;
|
||||
use fabro_util::workspace_glob::WorkspaceGlob;
|
||||
use pebble_coding_agent::environment::support::{capture_stats, tree_order, validate_glob};
|
||||
use pebble_coding_agent::environment::{
|
||||
DirEntry, EnvResult, Environment, EnvironmentError, EnvironmentErrorKind, ExecOutcome,
|
||||
ExecOutputSink, ExecOutputStream, ExecRequest, ExecResult, GrepOptions,
|
||||
};
|
||||
use pebble_coding_agent::tools::OutputCaptureStats as PebbleCaptureStats;
|
||||
use sandbox_driver::FileKind;
|
||||
|
||||
use crate::driver_sandbox::RunSandbox;
|
||||
|
|
@ -136,9 +135,8 @@ impl Environment for RunSandbox {
|
|||
.collect();
|
||||
// The driver lists in flat lexicographic order of the whole relative
|
||||
// path, where `foo-bar` sorts between `foo` and `foo/x`. Pebble lists
|
||||
// in tree order: by name within each directory. Comparing the paths
|
||||
// segment by segment is that order.
|
||||
entries.sort_by(|left, right| left.name.split('/').cmp(right.name.split('/')));
|
||||
// in tree order, and says how.
|
||||
tree_order(&mut entries);
|
||||
Ok(entries)
|
||||
}
|
||||
|
||||
|
|
@ -162,22 +160,11 @@ impl Environment for RunSandbox {
|
|||
}
|
||||
|
||||
async fn glob(&self, pattern: &str, path: Option<&str>) -> EnvResult<Vec<String>> {
|
||||
// Validated here rather than by the run sandbox's own glob so the
|
||||
// reason reaches the model in pebble's words, and so the patterns
|
||||
// pebble rejects (a trailing `/`, a `/` or wildcard inside `[...]`)
|
||||
// are rejected even though fabro's glob would accept them.
|
||||
if let Err(reason) = validate_pebble_glob(pattern) {
|
||||
return Err(EnvironmentError::new(
|
||||
EnvironmentErrorKind::InvalidInput,
|
||||
format!("Invalid glob pattern {pattern:?}: {reason}"),
|
||||
));
|
||||
}
|
||||
if let Err(error) = WorkspaceGlob::try_new(pattern) {
|
||||
return Err(EnvironmentError::new(
|
||||
EnvironmentErrorKind::InvalidInput,
|
||||
format!("Invalid glob pattern {pattern:?}: {error}"),
|
||||
));
|
||||
}
|
||||
// Validated by pebble's own grammar before the driver sees the
|
||||
// pattern, so the reason reaches the model in pebble's words and the
|
||||
// patterns pebble rejects are rejected even where fabro's glob would
|
||||
// accept them.
|
||||
validate_glob(pattern)?;
|
||||
Self::glob(self, pattern, path)
|
||||
.await
|
||||
.map_err(|error| environment_error("Failed to match files", error))
|
||||
|
|
@ -226,8 +213,14 @@ impl Environment for RunSandbox {
|
|||
duration_ms: streaming.result.duration_ms,
|
||||
},
|
||||
streams_separated: streaming.streams_separated,
|
||||
stdout_capture: capture_stats(streaming.stdout_capture),
|
||||
stderr_capture: capture_stats(streaming.stderr_capture),
|
||||
stdout_capture: capture_stats(
|
||||
streaming.stdout_capture.observed_bytes,
|
||||
output_bytes_cap,
|
||||
),
|
||||
stderr_capture: capture_stats(
|
||||
streaming.stderr_capture.observed_bytes,
|
||||
output_bytes_cap,
|
||||
),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -242,39 +235,6 @@ impl RunSandbox {
|
|||
|
||||
/// Pebble's glob grammar, beyond what fabro's glob already rejects.
|
||||
///
|
||||
/// A pattern names files, so one that ends with `/` is a mistake rather than a
|
||||
/// directory; and `/`, `*`, or `?` inside a character class never mean what a
|
||||
/// model meant by them. The messages are pebble's own, so a model corrects
|
||||
/// itself the same way wherever pebble runs.
|
||||
fn validate_pebble_glob(pattern: &str) -> Result<(), &'static str> {
|
||||
let trimmed = pattern.strip_prefix("./").unwrap_or(pattern);
|
||||
if trimmed.is_empty() {
|
||||
return Err("pattern cannot be empty");
|
||||
}
|
||||
if trimmed.ends_with('/') {
|
||||
return Err(
|
||||
"pattern ends with \"/\"; glob matches files, drop the trailing slash or add a \
|
||||
filename pattern",
|
||||
);
|
||||
}
|
||||
let mut in_class = false;
|
||||
for character in trimmed.chars() {
|
||||
match (in_class, character) {
|
||||
(false, '[') => in_class = true,
|
||||
(true, ']') => in_class = false,
|
||||
(true, '/') => return Err("a \"/\" cannot appear inside a character class"),
|
||||
(true, '*' | '?') => {
|
||||
return Err("wildcards are not valid inside a character class");
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
if in_class {
|
||||
return Err("pattern has an unclosed character class");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// A path with its redundant separators and `.` segments removed, for
|
||||
/// deciding whether two spellings name the same file.
|
||||
fn normalize(path: &str) -> String {
|
||||
|
|
@ -314,14 +274,6 @@ fn adapt_output_sink(sink: ExecOutputSink) -> CommandOutputCallback {
|
|||
})
|
||||
}
|
||||
|
||||
fn capture_stats(stats: sandbox::OutputCaptureStats) -> PebbleCaptureStats {
|
||||
PebbleCaptureStats {
|
||||
observed_bytes: stats.observed_bytes,
|
||||
retained_bytes: stats.retained_bytes,
|
||||
omitted_bytes: stats.omitted_bytes,
|
||||
}
|
||||
}
|
||||
|
||||
/// A sandbox failure as pebble classifies it, keeping the driver cause.
|
||||
fn environment_error(message: &str, error: crate::Error) -> EnvironmentError {
|
||||
let kind = if error.is_not_found() {
|
||||
|
|
@ -402,16 +354,6 @@ mod tests {
|
|||
drop(directory);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pebbles_glob_grammar_is_enforced_before_the_driver() {
|
||||
for pattern in ["nested/", "[a/]", "[a*]", "[ab"] {
|
||||
assert!(validate_pebble_glob(pattern).is_err(), "{pattern}");
|
||||
}
|
||||
for pattern in ["**/*.txt", "?.txt", "[ab].txt", "./src/**"] {
|
||||
assert!(validate_pebble_glob(pattern).is_ok(), "{pattern}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_path_spelled_two_ways_is_one_path() {
|
||||
assert_eq!(normalize("/work//a/./b.txt"), "/work/a/b.txt");
|
||||
|
|
|
|||
|
|
@ -66,7 +66,7 @@ use crate::model_fallback::{ModelFallbackNotice, ModelFallbackPolicy};
|
|||
use crate::outcome::billed_model_usage_from_llm;
|
||||
use crate::services::FabroRunToolServices;
|
||||
use crate::steering_hub::{ActiveControlHandle, SteeringHub, SteeringItem};
|
||||
use crate::web_search::{SearchBackend, SearchSecrets};
|
||||
use crate::web_search::{self, SearchSecrets};
|
||||
|
||||
/// The share of the model's context window at which an agent stage compacts
|
||||
/// its conversation. Fabro's own agent loop used this value; pebble's default
|
||||
|
|
@ -403,6 +403,18 @@ impl LiveAgent {
|
|||
self.tool_duration = self.tool_duration.saturating_add(report.timing.tool);
|
||||
self.files_touched
|
||||
.extend(report.files_touched.iter().cloned());
|
||||
for compaction in &report.compactions {
|
||||
// The summary call's usage is already in `report.usage`; this is
|
||||
// the breakdown, for anyone asking why a stage cost what it did.
|
||||
tracing::debug!(
|
||||
reason = ?compaction.reason,
|
||||
original_turns = compaction.original_turn_count,
|
||||
preserved_turns = compaction.preserved_turn_count,
|
||||
usage = ?compaction.usage,
|
||||
cost_usd_micros = ?compaction.cost_usd_micros,
|
||||
"agent stage compacted its conversation"
|
||||
);
|
||||
}
|
||||
if report.last_file_touched.is_some() {
|
||||
self.last_file_touched.clone_from(&report.last_file_touched);
|
||||
}
|
||||
|
|
@ -668,8 +680,8 @@ impl PebbleBackend {
|
|||
if let Some(human_input) = bindings.human_input {
|
||||
builder = builder.human_input(Arc::clone(human_input));
|
||||
}
|
||||
if let Some(search) = SearchBackend::from_secrets(&self.search_secrets) {
|
||||
builder = builder.search_provider(Arc::new(search));
|
||||
if let Some(search) = web_search::search_provider(&self.search_secrets) {
|
||||
builder = builder.search_provider(search);
|
||||
}
|
||||
if provider.profile_kind == AgentProfileKind::Claude5 {
|
||||
builder = builder.web_fetch_summarizer(route.selector());
|
||||
|
|
|
|||
|
|
@ -1,23 +1,14 @@
|
|||
//! Built-in `web_search` backends for workflow agents.
|
||||
//! The built-in `web_search` backend for workflow agents.
|
||||
//!
|
||||
//! Agents always call the same tool. Brave is preferred when its credential
|
||||
//! is present; otherwise Venice is used when its credential is present. With
|
||||
//! neither, the agent gets no search tool.
|
||||
//! Agents always call the same tool. Pebble ships the Brave and Venice
|
||||
//! providers; fabro supplies the credential, read from the vault, and the
|
||||
//! preference: Brave when its credential is present, else Venice when its is.
|
||||
//! With neither, the agent gets no search tool.
|
||||
|
||||
use std::fmt::Write as _;
|
||||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use pebble_coding_agent::extensions::{
|
||||
SearchError, SearchErrorKind, SearchProvider, SearchRequest, SearchResult,
|
||||
};
|
||||
|
||||
const BRAVE_SEARCH_URL: &str = "https://api.search.brave.com/res/v1/web/search";
|
||||
const VENICE_SEARCH_URL: &str = "https://api.venice.ai/api/v1/augment/search";
|
||||
const VENICE_QUERY_MAX_CHARS: usize = 400;
|
||||
const VENICE_REQUEST_TIMEOUT: Duration = Duration::from_mins(1);
|
||||
const MAX_RESULTS: u32 = 20;
|
||||
use pebble_coding_agent::extensions::SearchProvider;
|
||||
use pebble_coding_agent::search::providers::{Brave, Venice};
|
||||
|
||||
/// Credentials for the built-in search backends, read from the vault.
|
||||
#[derive(Clone, Default)]
|
||||
|
|
@ -38,342 +29,43 @@ impl std::fmt::Debug for SearchSecrets {
|
|||
}
|
||||
}
|
||||
|
||||
/// One of the search services a workflow agent can call.
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum SearchBackend {
|
||||
Brave {
|
||||
api_key: String,
|
||||
search_url: String,
|
||||
},
|
||||
Venice {
|
||||
api_key: String,
|
||||
search_url: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl SearchBackend {
|
||||
#[must_use]
|
||||
pub fn from_secrets(secrets: &SearchSecrets) -> Option<Self> {
|
||||
match (
|
||||
secrets.brave_search_api_key.as_ref(),
|
||||
secrets.venice_api_key.as_ref(),
|
||||
) {
|
||||
(Some(api_key), _) => Some(Self::brave(api_key.clone())),
|
||||
(None, Some(api_key)) => Some(Self::venice(api_key.clone())),
|
||||
(None, None) => None,
|
||||
}
|
||||
/// The provider an agent's `web_search` calls, when a credential names one.
|
||||
///
|
||||
/// The HTTP client is fabro's, so its proxy and TLS policy apply.
|
||||
#[must_use]
|
||||
pub fn search_provider(secrets: &SearchSecrets) -> Option<Arc<dyn SearchProvider>> {
|
||||
let client = fabro_http::HttpClientBuilder::new().build().ok()?;
|
||||
match (
|
||||
secrets.brave_search_api_key.as_ref(),
|
||||
secrets.venice_api_key.as_ref(),
|
||||
) {
|
||||
(Some(api_key), _) => Some(Arc::new(Brave::new(api_key.clone(), client))),
|
||||
(None, Some(api_key)) => Some(Arc::new(Venice::new(api_key.clone(), client))),
|
||||
(None, None) => None,
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn brave(api_key: String) -> Self {
|
||||
Self::Brave {
|
||||
api_key,
|
||||
search_url: BRAVE_SEARCH_URL.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn venice(api_key: String) -> Self {
|
||||
Self::Venice {
|
||||
api_key,
|
||||
search_url: VENICE_SEARCH_URL.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn with_search_url(mut self, url: &str) -> Self {
|
||||
match &mut self {
|
||||
Self::Brave { search_url, .. } | Self::Venice { search_url, .. } => {
|
||||
*search_url = url.to_string();
|
||||
}
|
||||
}
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SearchProvider for SearchBackend {
|
||||
async fn search(&self, request: SearchRequest) -> Result<Vec<SearchResult>, SearchError> {
|
||||
let max_results = request.max_results.clamp(1, MAX_RESULTS);
|
||||
match self {
|
||||
Self::Brave {
|
||||
api_key,
|
||||
search_url,
|
||||
} => search_brave(api_key, search_url, &request.query, max_results).await,
|
||||
Self::Venice {
|
||||
api_key,
|
||||
search_url,
|
||||
} => {
|
||||
if request.query.chars().count() > VENICE_QUERY_MAX_CHARS {
|
||||
return Err(SearchError::new(
|
||||
SearchErrorKind::InvalidRequest,
|
||||
format!(
|
||||
"query exceeds Venice Search maximum of {VENICE_QUERY_MAX_CHARS} \
|
||||
characters"
|
||||
),
|
||||
));
|
||||
}
|
||||
search_venice(api_key, search_url, &request.query, max_results).await
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn search_http_client() -> fabro_http::HttpClient {
|
||||
static CLIENT: OnceLock<fabro_http::HttpClient> = OnceLock::new();
|
||||
CLIENT
|
||||
.get_or_init(|| {
|
||||
#[cfg(test)]
|
||||
{
|
||||
fabro_http::test_http_client().expect("Search HTTP client should build")
|
||||
}
|
||||
#[cfg(not(test))]
|
||||
{
|
||||
fabro_http::http_client().expect("Search HTTP client should build")
|
||||
}
|
||||
})
|
||||
.clone()
|
||||
}
|
||||
|
||||
fn request_failed(error: impl std::fmt::Display) -> SearchError {
|
||||
SearchError::new(
|
||||
SearchErrorKind::Execution,
|
||||
format!("HTTP request failed: {error}"),
|
||||
)
|
||||
}
|
||||
|
||||
fn parse_failed(error: impl std::fmt::Display) -> SearchError {
|
||||
SearchError::new(
|
||||
SearchErrorKind::Execution,
|
||||
format!("Failed to parse response: {error}"),
|
||||
)
|
||||
}
|
||||
|
||||
async fn search_brave(
|
||||
api_key: &str,
|
||||
search_url: &str,
|
||||
query: &str,
|
||||
max_results: u32,
|
||||
) -> Result<Vec<SearchResult>, SearchError> {
|
||||
let resp = search_http_client()
|
||||
.get(search_url)
|
||||
.header("X-Subscription-Token", api_key)
|
||||
.header("Accept", "application/json")
|
||||
.query(&[("q", query), ("count", &max_results.to_string())])
|
||||
.send()
|
||||
.await
|
||||
.map_err(request_failed)?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
return Err(SearchError::new(
|
||||
SearchErrorKind::Execution,
|
||||
format!("Brave Search API returned status {}", resp.status()),
|
||||
));
|
||||
}
|
||||
|
||||
let body: serde_json::Value = resp.json().await.map_err(parse_failed)?;
|
||||
Ok(brave_results(&body))
|
||||
}
|
||||
|
||||
async fn search_venice(
|
||||
api_key: &str,
|
||||
search_url: &str,
|
||||
query: &str,
|
||||
max_results: u32,
|
||||
) -> Result<Vec<SearchResult>, SearchError> {
|
||||
let resp = search_http_client()
|
||||
.post(search_url)
|
||||
.timeout(VENICE_REQUEST_TIMEOUT)
|
||||
.bearer_auth(api_key)
|
||||
.header("Accept", "application/json")
|
||||
.json(&serde_json::json!({
|
||||
"query": query,
|
||||
"limit": max_results,
|
||||
"search_provider": "brave",
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.map_err(request_failed)?;
|
||||
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
return Err(SearchError::new(
|
||||
SearchErrorKind::Execution,
|
||||
venice_status_error(status.as_u16(), &resp),
|
||||
));
|
||||
}
|
||||
|
||||
let body: serde_json::Value = resp.json().await.map_err(parse_failed)?;
|
||||
Ok(venice_results(&body))
|
||||
}
|
||||
|
||||
fn venice_status_error(status: u16, resp: &fabro_http::Response) -> String {
|
||||
let mut message = format!("Venice Search API returned status {status}");
|
||||
if status == 402 {
|
||||
if let Some(balance) = header_str(resp, "x-venice-balance-usd") {
|
||||
let _ = write!(message, " (balance USD {balance})");
|
||||
} else if let Some(balance) = header_str(resp, "x-venice-balance-diem") {
|
||||
let _ = write!(message, " (balance DIEM {balance})");
|
||||
}
|
||||
}
|
||||
message
|
||||
}
|
||||
|
||||
fn header_str(resp: &fabro_http::Response, name: &str) -> Option<String> {
|
||||
resp.headers()
|
||||
.get(name)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_owned)
|
||||
}
|
||||
|
||||
fn brave_results(body: &serde_json::Value) -> Vec<SearchResult> {
|
||||
body.get("web")
|
||||
.and_then(|web| web.get("results"))
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.map(|results| {
|
||||
results
|
||||
.iter()
|
||||
.map(|result| {
|
||||
SearchResult::new(
|
||||
json_str(result, "title"),
|
||||
json_str(result, "url"),
|
||||
json_str(result, "description"),
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn venice_results(body: &serde_json::Value) -> Vec<SearchResult> {
|
||||
body.get("results")
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.map(|results| {
|
||||
results
|
||||
.iter()
|
||||
.map(|result| {
|
||||
let hit = SearchResult::new(
|
||||
json_str(result, "title"),
|
||||
json_str(result, "url"),
|
||||
json_str(result, "content"),
|
||||
);
|
||||
match optional_json_str(result, "date") {
|
||||
Some(date) => hit.with_published_at(date),
|
||||
None => hit,
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn json_str(value: &serde_json::Value, key: &str) -> String {
|
||||
optional_json_str(value, key).unwrap_or_else(|| match key {
|
||||
"title" => "(no title)".to_string(),
|
||||
"url" => "(no url)".to_string(),
|
||||
_ => String::new(),
|
||||
})
|
||||
}
|
||||
|
||||
fn optional_json_str(value: &serde_json::Value, key: &str) -> Option<String> {
|
||||
value
|
||||
.get(key)
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.filter(|text| !text.is_empty())
|
||||
.map(str::to_owned)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use httpmock::Method::{GET, POST};
|
||||
use httpmock::MockServer;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn secrets(brave: Option<&str>, venice: Option<&str>) -> SearchSecrets {
|
||||
SearchSecrets {
|
||||
brave_search_api_key: brave.map(str::to_string),
|
||||
venice_api_key: venice.map(str::to_string),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn brave_is_preferred_and_venice_is_the_fallback() {
|
||||
assert!(matches!(
|
||||
SearchBackend::from_secrets(&secrets(Some("b"), Some("v"))),
|
||||
Some(SearchBackend::Brave { .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
SearchBackend::from_secrets(&secrets(None, Some("v"))),
|
||||
Some(SearchBackend::Venice { .. })
|
||||
));
|
||||
assert!(SearchBackend::from_secrets(&secrets(None, None)).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn brave_results_are_returned_in_order() {
|
||||
let server = MockServer::start_async().await;
|
||||
let mock = server
|
||||
.mock_async(|when, then| {
|
||||
when.method(GET)
|
||||
.path("/search")
|
||||
.header("X-Subscription-Token", "brave-key")
|
||||
.query_param("q", "fabro")
|
||||
.query_param("count", "2");
|
||||
then.status(200).json_body(serde_json::json!({
|
||||
"web": {"results": [
|
||||
{"title": "One", "url": "https://one", "description": "first"},
|
||||
{"title": "Two", "url": "https://two", "description": "second"}
|
||||
]}
|
||||
}));
|
||||
})
|
||||
.await;
|
||||
let backend = SearchBackend::brave("brave-key".to_string())
|
||||
.with_search_url(&format!("{}/search", server.base_url()));
|
||||
|
||||
let results = backend
|
||||
.search(SearchRequest::new("fabro", 2))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
mock.assert_async().await;
|
||||
assert_eq!(results.len(), 2);
|
||||
assert_eq!(results[0].title, "One");
|
||||
assert_eq!(results[1].snippet, "second");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn venice_results_carry_dates_and_reject_long_queries() {
|
||||
let server = MockServer::start_async().await;
|
||||
let mock = server
|
||||
.mock_async(|when, then| {
|
||||
when.method(POST)
|
||||
.path("/augment")
|
||||
.header("Authorization", "Bearer venice-key");
|
||||
then.status(200).json_body(serde_json::json!({
|
||||
"results": [
|
||||
{"title": "One", "url": "https://one", "content": "first", "date": "2026-01-01"}
|
||||
]
|
||||
}));
|
||||
})
|
||||
.await;
|
||||
let backend = SearchBackend::venice("venice-key".to_string())
|
||||
.with_search_url(&format!("{}/augment", server.base_url()));
|
||||
|
||||
let results = backend
|
||||
.search(SearchRequest::new("fabro", 5))
|
||||
.await
|
||||
.unwrap();
|
||||
mock.assert_async().await;
|
||||
assert_eq!(results[0].published_at.as_deref(), Some("2026-01-01"));
|
||||
|
||||
let error = backend
|
||||
.search(SearchRequest::new(
|
||||
"x".repeat(VENICE_QUERY_MAX_CHARS + 1),
|
||||
5,
|
||||
))
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(error.kind(), SearchErrorKind::InvalidRequest);
|
||||
fn brave_is_preferred_and_neither_means_no_tool() {
|
||||
let both = SearchSecrets {
|
||||
brave_search_api_key: Some("b".into()),
|
||||
venice_api_key: Some("v".into()),
|
||||
};
|
||||
assert!(search_provider(&both).is_some());
|
||||
let venice_only = SearchSecrets {
|
||||
brave_search_api_key: None,
|
||||
venice_api_key: Some("v".into()),
|
||||
};
|
||||
assert!(search_provider(&venice_only).is_some());
|
||||
assert!(search_provider(&SearchSecrets::default()).is_none());
|
||||
assert_eq!(
|
||||
format!("{both:?}"),
|
||||
"SearchSecrets { brave_search_configured: true, venice_configured: true }",
|
||||
"a debug rendering never carries a key"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue