1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
//! The `SearchModel` trait — the provider-facing interface for web search.
//!
//! Designed per [RFC-0007](../../rfc/0007-search-model-trait.md). Covers the
//! common `query → results[]` pattern shared by search providers (Tavily,
//! Serper, Exa, etc.). Provider-specific fields are passed via
//! `provider_options`; the core trait only models the shared structure.
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use ts_rs::TS;
use crate::error::AiMuxError;
use crate::shared::{SharedHeaders, SharedProviderMetadata, SharedProviderOptions};
use crate::types::Warning;
use crate::{AbortSignal, retry, timeout};
/// Options passed to [`SearchModel::do_search`].
#[derive(Debug, Clone, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct SearchCallOptions {
/// The search query string.
pub query: String,
/// Maximum number of results to return.
pub max_results: Option<u32>,
/// Whether to include raw page content in results.
///
/// Provider support varies; providers that cannot honor this should
/// issue a warning rather than erroring.
pub include_raw_content: Option<bool>,
/// Optional time range filter (e.g. `"day"`, `"week"`, `"month"`,
/// `"year"`). Provider support varies.
pub time_range: Option<String>,
/// Optional list of domains to include in results.
pub include_domains: Option<Vec<String>>,
/// Optional list of domains to exclude from results.
pub exclude_domains: Option<Vec<String>>,
/// Abort signal for cancelling the operation.
#[serde(skip)]
#[ts(skip)]
pub abort_signal: Option<AbortSignal>,
/// Per-call retry override. `None` uses the model default.
pub max_retries: Option<u32>,
/// Per-call operation timeout.
pub timeout: Option<crate::options::TimeoutConfiguration>,
/// Additional provider-specific options, keyed by provider name.
pub provider_options: Option<SharedProviderOptions>,
/// Additional HTTP headers to send with the request.
pub headers: Option<SharedHeaders>,
}
impl SearchCallOptions {
/// Create options with a query and all else unset.
pub fn new(query: impl Into<String>) -> Self {
Self {
query: query.into(),
max_results: None,
include_raw_content: None,
time_range: None,
include_domains: None,
exclude_domains: None,
abort_signal: None,
max_retries: None,
timeout: None,
provider_options: None,
headers: None,
}
}
}
/// A single search result item.
#[derive(Debug, Clone, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct SearchResultItem {
/// The title of the result (e.g. page title).
pub title: Option<String>,
/// The URL of the result.
pub url: Option<String>,
/// A snippet/summary of the result content.
pub content: Option<String>,
/// Raw page content (when `include_raw_content` is requested and
/// supported by the provider).
pub raw_content: Option<String>,
/// A relevance score (0.0–1.0) if the provider returns one.
pub score: Option<f64>,
/// Provider-specific metadata for this result.
pub provider_metadata: Option<SharedProviderMetadata>,
}
/// The result of [`SearchModel::do_search`].
#[derive(Debug, Clone, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct SearchResult {
/// Ordered list of search results.
pub results: Vec<SearchResultItem>,
/// An optional direct answer / summary (some providers return
/// an AI-generated answer alongside results).
pub answer: Option<String>,
/// Additional provider-specific metadata.
pub provider_metadata: Option<SharedProviderMetadata>,
/// Warnings for the call.
pub warnings: Vec<Warning>,
/// Optional response information for debugging.
pub response: Option<SearchResponse>,
}
/// Optional response information for a search call.
#[derive(Debug, Clone, Default, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct SearchResponse {
/// Response headers.
pub headers: Option<SharedHeaders>,
/// The response body (opaque JSON).
pub body: Option<serde_json::Value>,
}
/// The unified search model trait (provider-facing).
///
/// Designed per [RFC-0007](../../rfc/0007-search-model-trait.md). Providers
/// implement `do_search`; users never call it directly.
#[async_trait]
pub trait SearchModel: Send + Sync {
/// Specification version (always `"v4"`).
fn specification_version(&self) -> &'static str {
"v4"
}
/// Provider name, e.g. `"tavily"`.
fn provider(&self) -> &str;
/// Provider-specific model ID (some providers use fixed IDs like
/// `"tavily-search"`; others accept endpoint-specific names).
fn model_id(&self) -> &str;
fn retry_config(&self) -> crate::retry::RetryConfig {
crate::retry::RetryConfig::default()
}
/// Execute a search query and return results.
///
/// Naming: the `do_` prefix prevents accidental direct usage by users.
async fn do_search(&self, options: &SearchCallOptions) -> Result<SearchResult, AiMuxError>;
}
/// User-facing search with Core-owned retry and timeout.
///
/// # Errors
///
/// Returns the provider failure, retry exhaustion, timeout, or caller abort.
pub async fn search(
model: &dyn SearchModel,
options: SearchCallOptions,
) -> Result<SearchResult, AiMuxError> {
let timeout = timeout::OperationTimeout::new(options.timeout.unwrap_or_default())?;
let abort_signal = options.abort_signal.clone();
let retries = retry::prepare_retries(
options.max_retries,
model.retry_config(),
abort_signal.clone(),
);
timeout::run(
retries.retry(|| model.do_search(&options)),
abort_signal.as_ref(),
timeout,
)
.await
}