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