1use 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#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize)]
40pub struct ProviderId(Cow<'static, str>);
41
42impl ProviderId {
43 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#[derive(Debug, Clone, Default, PartialEq, Eq)]
93#[non_exhaustive]
94pub struct Capabilities {
95 pub lookup_kinds: Vec<IdKind>,
97 pub max_batch: usize,
100 pub fields: FieldSet,
102 pub search: bool,
103 pub citation_search: bool,
106 pub notices: bool,
110 pub search_filters: Vec<Filter>,
111 pub final_misses: Vec<IdKind>,
114 pub citations: bool,
115 pub references: bool,
116 pub recommend: bool,
117}
118
119#[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 Place,
129}
130
131#[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#[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#[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 pub limit: u32,
166 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#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
189#[serde(rename_all = "snake_case")]
190#[non_exhaustive]
191pub enum SearchMode {
192 #[default]
194 Keywords,
195 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 pub place: Option<Place>,
211 pub fields: FieldSet,
212 pub page: PageReq,
213 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 pub fn citation(reference: impl Into<String>) -> Self {
228 SearchQuery {
229 mode: SearchMode::Citation,
230 ..SearchQuery::new(reference)
231 }
232 }
233
234 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#[derive(Debug, Clone, PartialEq)]
250pub struct Lookup {
251 pub id: WorkId,
253 pub record: Option<SourceRecord>,
255}
256
257#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
259#[non_exhaustive]
260pub struct Page {
261 pub hits: Vec<Hit>,
262 pub total: Option<u64>,
264 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#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
276pub struct Hit {
277 pub record: SourceRecord,
278 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#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
292#[non_exhaustive]
293pub struct Usage {
294 pub cost: Option<f64>,
296 pub remaining: Option<f64>,
298 pub resets_at: Option<Timestamp>,
300}
301
302impl Usage {
303 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#[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#[async_trait]
343pub trait Provider: Send + Sync {
344 fn id(&self) -> ProviderId;
345
346 fn capabilities(&self) -> &Capabilities;
347
348 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 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 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 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 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}