Skip to main content

aspero_core/
provider.rs

1//! The `Provider` trait: the contract every literature source implements.
2//!
3//! Rules every implementation follows (the engine relies on them):
4//!
5//! 1. **No merging, no caching, no retries.** A provider makes the request,
6//!    maps the response, and returns. Retries, backoff, rate limiting, caching
7//!    and merging are the engine's job, so they are consistent and shared.
8//! 2. **Declared capabilities are honoured.** Every [`Filter`] listed in
9//!    [`Capabilities::search_filters`] is applied when present in a query. The
10//!    engine reports filters a provider doesn't declare as ignored, so nothing
11//!    is dropped silently. Unknown keys in [`SearchQuery::extra`] are
12//!    rejected with [`Error::BadRequest`].
13//! 3. **Batch lookups answer every ID.** [`Provider::lookup`] returns one
14//!    [`Lookup`] per input ID, in input order, with `record: None` for IDs the
15//!    provider doesn't know. Whole-call failures are `Err`.
16//! 4. **Errors are mapped** into [`crate::Error`], including non-provider
17//!    bodies (e.g. an API gateway's 429 page).
18//! 5. **Usage is reported** when the provider exposes it (credits spent,
19//!    budget remaining), so agents can pace themselves.
20//!
21//! Each provider also owns its HTTP mapping, credentials and default
22//! [`ResiliencePolicy`](crate::ResiliencePolicy).
23
24use std::borrow::Cow;
25use std::collections::BTreeMap;
26use std::fmt;
27
28use async_trait::async_trait;
29use jiff::Timestamp;
30use serde::{Deserialize, Serialize};
31
32use crate::error::{Error, Result};
33use crate::id::{IdKind, WorkId};
34use crate::record::{FieldSet, SourceRecord};
35use crate::work::WorkType;
36
37/// Stable short name of a provider, e.g. `openalex`. Lowercase ASCII letters,
38/// digits and `-`.
39#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize)]
40pub struct ProviderId(Cow<'static, str>);
41
42impl ProviderId {
43    /// For built-in providers.
44    ///
45    /// # Panics
46    /// If `id` is not a valid provider name (checked at compile time when used
47    /// in a `const`).
48    pub const fn from_static(id: &'static str) -> Self {
49        assert!(Self::is_valid(id), "invalid provider id");
50        ProviderId(Cow::Borrowed(id))
51    }
52
53    pub fn parse(id: &str) -> Option<Self> {
54        Self::is_valid(id).then(|| ProviderId(Cow::Owned(id.to_owned())))
55    }
56
57    pub fn as_str(&self) -> &str {
58        &self.0
59    }
60
61    const fn is_valid(id: &str) -> bool {
62        let b = id.as_bytes();
63        if b.is_empty() {
64            return false;
65        }
66        let mut i = 0;
67        while i < b.len() {
68            if !(b[i].is_ascii_lowercase() || b[i].is_ascii_digit() || b[i] == b'-') {
69                return false;
70            }
71            i += 1;
72        }
73        true
74    }
75}
76
77impl fmt::Display for ProviderId {
78    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
79        f.write_str(&self.0)
80    }
81}
82
83impl<'de> Deserialize<'de> for ProviderId {
84    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
85        let s = String::deserialize(d)?;
86        ProviderId::parse(&s)
87            .ok_or_else(|| serde::de::Error::custom(format!("invalid provider id {s:?}")))
88    }
89}
90
91/// What a provider can do. The engine routes on this before calling.
92#[derive(Debug, Clone, Default, PartialEq, Eq)]
93#[non_exhaustive]
94pub struct Capabilities {
95    /// Identifier kinds [`Provider::lookup`] accepts.
96    pub lookup_kinds: Vec<IdKind>,
97    /// Most IDs per `lookup` call; the engine splits larger batches. 0 is
98    /// treated as 1.
99    pub max_batch: usize,
100    /// Fields this provider can supply at all.
101    pub fields: FieldSet,
102    pub search: bool,
103    /// Supports [`SearchMode::Citation`]: a matcher built for reference
104    /// strings, not just keyword search.
105    pub citation_search: bool,
106    /// Reports editorial notices in detail and in both directions (a work's
107    /// corrections, a correction's work), so the engine asks it about
108    /// records other providers supplied.
109    pub notices: bool,
110    pub search_filters: Vec<Filter>,
111    /// Identifier kinds whose "not found" here is final: no other provider
112    /// can know them (doi.org for DOIs), so lookups stop asking.
113    pub final_misses: Vec<IdKind>,
114    pub citations: bool,
115    pub references: bool,
116    pub recommend: bool,
117}
118
119/// Common search filters (see [`SearchQuery`]).
120#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
121#[serde(rename_all = "snake_case")]
122#[non_exhaustive]
123pub enum Filter {
124    Years,
125    WorkType,
126    OpenAccess,
127    /// Volume and first page ([`SearchQuery::place`]).
128    Place,
129}
130
131/// Where a work appeared, to search by: volume and first page.
132#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
133#[non_exhaustive]
134pub struct Place {
135    pub volume: String,
136    pub first_page: String,
137}
138
139impl Place {
140    pub fn new(volume: impl Into<String>, first_page: impl Into<String>) -> Self {
141        Place {
142            volume: volume.into(),
143            first_page: first_page.into(),
144        }
145    }
146}
147
148/// Inclusive publication-year range; either end may be open.
149#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
150pub struct YearRange {
151    pub from: Option<i16>,
152    pub to: Option<i16>,
153}
154
155/// An opaque, provider-specific continuation token.
156#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
157#[serde(transparent)]
158pub struct Cursor(pub String);
159
160#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
161#[non_exhaustive]
162pub struct PageReq {
163    /// Requested page size. Providers cap it at their own maximum and may
164    /// return fewer, never more.
165    pub limit: u32,
166    /// `None` for the first page.
167    pub cursor: Option<Cursor>,
168}
169
170impl PageReq {
171    pub const DEFAULT_LIMIT: u32 = 25;
172
173    pub fn first(limit: u32) -> Self {
174        PageReq {
175            limit,
176            cursor: None,
177        }
178    }
179}
180
181impl Default for PageReq {
182    fn default() -> Self {
183        PageReq::first(Self::DEFAULT_LIMIT)
184    }
185}
186
187/// How [`SearchQuery::text`] should be interpreted.
188#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
189#[serde(rename_all = "snake_case")]
190#[non_exhaustive]
191pub enum SearchMode {
192    /// Free-text keywords.
193    #[default]
194    Keywords,
195    /// A whole, possibly messy reference string ("LeCun Y, Bengio Y. Deep
196    /// learning. Nature. 2015"), to be matched to the work it cites. Only
197    /// sent to providers declaring [`Capabilities::citation_search`].
198    Citation,
199}
200
201#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
202#[non_exhaustive]
203pub struct SearchQuery {
204    pub text: String,
205    pub mode: SearchMode,
206    pub years: Option<YearRange>,
207    pub work_type: Option<WorkType>,
208    pub open_access: Option<bool>,
209    /// Volume and first page, for citations that give no title.
210    pub place: Option<Place>,
211    pub fields: FieldSet,
212    pub page: PageReq,
213    /// Provider-specific parameters, passed through by name.
214    pub extra: BTreeMap<String, String>,
215}
216
217impl SearchQuery {
218    pub fn new(text: impl Into<String>) -> Self {
219        SearchQuery {
220            text: text.into(),
221            fields: FieldSet::DEFAULT,
222            ..SearchQuery::default()
223        }
224    }
225
226    /// A [`SearchMode::Citation`] query for a reference string.
227    pub fn citation(reference: impl Into<String>) -> Self {
228        SearchQuery {
229            mode: SearchMode::Citation,
230            ..SearchQuery::new(reference)
231        }
232    }
233
234    /// The common filters this query actually uses.
235    pub fn filters(&self) -> Vec<Filter> {
236        [
237            self.years.map(|_| Filter::Years),
238            self.work_type.as_ref().map(|_| Filter::WorkType),
239            self.open_access.map(|_| Filter::OpenAccess),
240            self.place.as_ref().map(|_| Filter::Place),
241        ]
242        .into_iter()
243        .flatten()
244        .collect()
245    }
246}
247
248/// The answer for one requested ID in a batch lookup.
249#[derive(Debug, Clone, PartialEq)]
250pub struct Lookup {
251    /// The ID as requested.
252    pub id: WorkId,
253    /// `None`: the provider doesn't know this ID.
254    pub record: Option<SourceRecord>,
255}
256
257/// A page of works: search hits in rank order, or one side of the citation graph.
258#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
259#[non_exhaustive]
260pub struct Page {
261    pub hits: Vec<Hit>,
262    /// Total matches, if the provider reports it.
263    pub total: Option<u64>,
264    /// `None` on the last page.
265    pub next: Option<Cursor>,
266}
267
268impl Page {
269    pub fn new(hits: Vec<Hit>, total: Option<u64>, next: Option<Cursor>) -> Self {
270        Page { hits, total, next }
271    }
272}
273
274/// One work in a [`Page`].
275#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
276pub struct Hit {
277    pub record: SourceRecord,
278    /// The provider's relevance score, if it gives one. Only comparable
279    /// within one provider's results for one query.
280    pub score: Option<f64>,
281}
282
283impl Hit {
284    pub fn new(record: SourceRecord, score: Option<f64>) -> Self {
285        Hit { record, score }
286    }
287}
288
289/// Usage accounting reported by the provider, in its own units (OpenAlex:
290/// USD). Any part may be unknown.
291#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
292#[non_exhaustive]
293pub struct Usage {
294    /// Spent by this call.
295    pub cost: Option<f64>,
296    /// Left in the current budget period.
297    pub remaining: Option<f64>,
298    /// When the budget resets.
299    pub resets_at: Option<Timestamp>,
300}
301
302impl Usage {
303    /// Usage of this call followed by `later`: costs add up; the remaining
304    /// budget and reset time are the most recent known.
305    pub fn then(self, later: Usage) -> Usage {
306        Usage {
307            cost: match (self.cost, later.cost) {
308                (Some(a), Some(b)) => Some(a + b),
309                (a, b) => a.or(b),
310            },
311            remaining: later.remaining.or(self.remaining),
312            resets_at: later.resets_at.or(self.resets_at),
313        }
314    }
315}
316
317/// A provider result plus the usage it cost.
318#[derive(Debug, Clone, PartialEq)]
319pub struct Response<T> {
320    pub value: T,
321    pub usage: Option<Usage>,
322}
323
324impl<T> Response<T> {
325    pub fn new(value: T) -> Self {
326        Response { value, usage: None }
327    }
328
329    pub fn with_usage(value: T, usage: Usage) -> Self {
330        Response {
331            value,
332            usage: Some(usage),
333        }
334    }
335}
336
337/// A source of scholarly metadata. See the [module docs](self) for the rules
338/// implementations follow.
339///
340/// Operations default to [`Error::Unsupported`]; a provider overrides those
341/// it declares in [`Provider::capabilities`].
342#[async_trait]
343pub trait Provider: Send + Sync {
344    fn id(&self) -> ProviderId;
345
346    fn capabilities(&self) -> &Capabilities;
347
348    /// Looks up works by ID. One [`Lookup`] per input, in input order.
349    /// Callers pass at most `max_batch` IDs, all of a kind in `lookup_kinds`.
350    async fn lookup(&self, ids: &[WorkId], fields: FieldSet) -> Result<Response<Vec<Lookup>>> {
351        let _ = (ids, fields);
352        Err(Error::Unsupported("lookup"))
353    }
354
355    async fn search(&self, query: &SearchQuery) -> Result<Response<Page>> {
356        let _ = query;
357        Err(Error::Unsupported("search"))
358    }
359
360    /// Works that cite `id`.
361    async fn citations(
362        &self,
363        id: &WorkId,
364        page: &PageReq,
365        fields: FieldSet,
366    ) -> Result<Response<Page>> {
367        let _ = (id, page, fields);
368        Err(Error::Unsupported("citations"))
369    }
370
371    /// Works that `id` cites.
372    async fn references(
373        &self,
374        id: &WorkId,
375        page: &PageReq,
376        fields: FieldSet,
377    ) -> Result<Response<Page>> {
378        let _ = (id, page, fields);
379        Err(Error::Unsupported("references"))
380    }
381
382    /// Works similar to `positive` and unlike `negative`.
383    async fn recommend(
384        &self,
385        positive: &[WorkId],
386        negative: &[WorkId],
387        page: &PageReq,
388        fields: FieldSet,
389    ) -> Result<Response<Page>> {
390        let _ = (positive, negative, page, fields);
391        Err(Error::Unsupported("recommend"))
392    }
393}
394
395#[cfg(test)]
396mod tests {
397    use super::*;
398
399    struct Nothing(Capabilities);
400
401    #[async_trait]
402    impl Provider for Nothing {
403        fn id(&self) -> ProviderId {
404            ProviderId::from_static("nothing")
405        }
406
407        fn capabilities(&self) -> &Capabilities {
408            &self.0
409        }
410    }
411
412    #[test]
413    fn usage_adds_costs_and_keeps_latest_budget() {
414        let first = Usage {
415            cost: Some(0.001),
416            remaining: Some(0.9),
417            resets_at: None,
418        };
419        let second = Usage {
420            cost: Some(0.0001),
421            remaining: None,
422            resets_at: None,
423        };
424        let total = first.then(second);
425        assert!((total.cost.unwrap() - 0.0011).abs() < 1e-12);
426        assert_eq!(total.remaining, Some(0.9));
427    }
428
429    #[test]
430    fn provider_ids_are_validated() {
431        assert!(ProviderId::parse("openalex").is_some());
432        assert!(ProviderId::parse("s2-graph").is_some());
433        assert!(ProviderId::parse("OpenAlex").is_none());
434        assert!(ProviderId::parse("").is_none());
435        assert!(serde_json::from_str::<ProviderId>(r#""bad id""#).is_err());
436    }
437
438    #[test]
439    fn query_reports_the_filters_it_uses() {
440        let mut q = SearchQuery::new("transformers");
441        assert!(q.filters().is_empty());
442        q.years = Some(YearRange {
443            from: Some(2017),
444            to: None,
445        });
446        q.open_access = Some(true);
447        assert_eq!(q.filters(), [Filter::Years, Filter::OpenAccess]);
448    }
449
450    #[test]
451    fn unimplemented_operations_are_unsupported_and_dyn_compatible() {
452        let p: Box<dyn Provider> = Box::new(Nothing(Capabilities::default()));
453        let query = SearchQuery::new("x");
454        let result = block_on_ready(p.search(&query));
455        assert!(matches!(result, Err(Error::Unsupported("search"))));
456    }
457
458    /// Polls a future that is immediately ready (no runtime in `core`).
459    fn block_on_ready<F: std::future::Future>(fut: F) -> F::Output {
460        use std::task::{Context, Poll, Waker};
461        let mut fut = std::pin::pin!(fut);
462        match fut.as_mut().poll(&mut Context::from_waker(Waker::noop())) {
463            Poll::Ready(out) => out,
464            Poll::Pending => panic!("future was not immediately ready"),
465        }
466    }
467}