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 citations: bool,
112 pub references: bool,
113 pub recommend: bool,
114}
115
116#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
118#[serde(rename_all = "snake_case")]
119#[non_exhaustive]
120pub enum Filter {
121 Years,
122 WorkType,
123 OpenAccess,
124}
125
126#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
128pub struct YearRange {
129 pub from: Option<i16>,
130 pub to: Option<i16>,
131}
132
133#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
135#[serde(transparent)]
136pub struct Cursor(pub String);
137
138#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
139#[non_exhaustive]
140pub struct PageReq {
141 pub limit: u32,
144 pub cursor: Option<Cursor>,
146}
147
148impl PageReq {
149 pub const DEFAULT_LIMIT: u32 = 25;
150
151 pub fn first(limit: u32) -> Self {
152 PageReq {
153 limit,
154 cursor: None,
155 }
156 }
157}
158
159impl Default for PageReq {
160 fn default() -> Self {
161 PageReq::first(Self::DEFAULT_LIMIT)
162 }
163}
164
165#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
167#[serde(rename_all = "snake_case")]
168#[non_exhaustive]
169pub enum SearchMode {
170 #[default]
172 Keywords,
173 Citation,
177}
178
179#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
180#[non_exhaustive]
181pub struct SearchQuery {
182 pub text: String,
183 pub mode: SearchMode,
184 pub years: Option<YearRange>,
185 pub work_type: Option<WorkType>,
186 pub open_access: Option<bool>,
187 pub fields: FieldSet,
188 pub page: PageReq,
189 pub extra: BTreeMap<String, String>,
191}
192
193impl SearchQuery {
194 pub fn new(text: impl Into<String>) -> Self {
195 SearchQuery {
196 text: text.into(),
197 fields: FieldSet::DEFAULT,
198 ..SearchQuery::default()
199 }
200 }
201
202 pub fn citation(reference: impl Into<String>) -> Self {
204 SearchQuery {
205 mode: SearchMode::Citation,
206 ..SearchQuery::new(reference)
207 }
208 }
209
210 pub fn filters(&self) -> Vec<Filter> {
212 [
213 self.years.map(|_| Filter::Years),
214 self.work_type.as_ref().map(|_| Filter::WorkType),
215 self.open_access.map(|_| Filter::OpenAccess),
216 ]
217 .into_iter()
218 .flatten()
219 .collect()
220 }
221}
222
223#[derive(Debug, Clone, PartialEq)]
225pub struct Lookup {
226 pub id: WorkId,
228 pub record: Option<SourceRecord>,
230}
231
232#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
234#[non_exhaustive]
235pub struct Page {
236 pub hits: Vec<Hit>,
237 pub total: Option<u64>,
239 pub next: Option<Cursor>,
241}
242
243impl Page {
244 pub fn new(hits: Vec<Hit>, total: Option<u64>, next: Option<Cursor>) -> Self {
245 Page { hits, total, next }
246 }
247}
248
249#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
251pub struct Hit {
252 pub record: SourceRecord,
253 pub score: Option<f64>,
256}
257
258impl Hit {
259 pub fn new(record: SourceRecord, score: Option<f64>) -> Self {
260 Hit { record, score }
261 }
262}
263
264#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
267#[non_exhaustive]
268pub struct Usage {
269 pub cost: Option<f64>,
271 pub remaining: Option<f64>,
273 pub resets_at: Option<Timestamp>,
275}
276
277impl Usage {
278 pub fn then(self, later: Usage) -> Usage {
281 Usage {
282 cost: match (self.cost, later.cost) {
283 (Some(a), Some(b)) => Some(a + b),
284 (a, b) => a.or(b),
285 },
286 remaining: later.remaining.or(self.remaining),
287 resets_at: later.resets_at.or(self.resets_at),
288 }
289 }
290}
291
292#[derive(Debug, Clone, PartialEq)]
294pub struct Response<T> {
295 pub value: T,
296 pub usage: Option<Usage>,
297}
298
299impl<T> Response<T> {
300 pub fn new(value: T) -> Self {
301 Response { value, usage: None }
302 }
303
304 pub fn with_usage(value: T, usage: Usage) -> Self {
305 Response {
306 value,
307 usage: Some(usage),
308 }
309 }
310}
311
312#[async_trait]
318pub trait Provider: Send + Sync {
319 fn id(&self) -> ProviderId;
320
321 fn capabilities(&self) -> &Capabilities;
322
323 async fn lookup(&self, ids: &[WorkId], fields: FieldSet) -> Result<Response<Vec<Lookup>>> {
326 let _ = (ids, fields);
327 Err(Error::Unsupported("lookup"))
328 }
329
330 async fn search(&self, query: &SearchQuery) -> Result<Response<Page>> {
331 let _ = query;
332 Err(Error::Unsupported("search"))
333 }
334
335 async fn citations(
337 &self,
338 id: &WorkId,
339 page: &PageReq,
340 fields: FieldSet,
341 ) -> Result<Response<Page>> {
342 let _ = (id, page, fields);
343 Err(Error::Unsupported("citations"))
344 }
345
346 async fn references(
348 &self,
349 id: &WorkId,
350 page: &PageReq,
351 fields: FieldSet,
352 ) -> Result<Response<Page>> {
353 let _ = (id, page, fields);
354 Err(Error::Unsupported("references"))
355 }
356
357 async fn recommend(
359 &self,
360 positive: &[WorkId],
361 negative: &[WorkId],
362 page: &PageReq,
363 fields: FieldSet,
364 ) -> Result<Response<Page>> {
365 let _ = (positive, negative, page, fields);
366 Err(Error::Unsupported("recommend"))
367 }
368}
369
370#[cfg(test)]
371mod tests {
372 use super::*;
373
374 struct Nothing(Capabilities);
375
376 #[async_trait]
377 impl Provider for Nothing {
378 fn id(&self) -> ProviderId {
379 ProviderId::from_static("nothing")
380 }
381
382 fn capabilities(&self) -> &Capabilities {
383 &self.0
384 }
385 }
386
387 #[test]
388 fn usage_adds_costs_and_keeps_latest_budget() {
389 let first = Usage {
390 cost: Some(0.001),
391 remaining: Some(0.9),
392 resets_at: None,
393 };
394 let second = Usage {
395 cost: Some(0.0001),
396 remaining: None,
397 resets_at: None,
398 };
399 let total = first.then(second);
400 assert!((total.cost.unwrap() - 0.0011).abs() < 1e-12);
401 assert_eq!(total.remaining, Some(0.9));
402 }
403
404 #[test]
405 fn provider_ids_are_validated() {
406 assert!(ProviderId::parse("openalex").is_some());
407 assert!(ProviderId::parse("s2-graph").is_some());
408 assert!(ProviderId::parse("OpenAlex").is_none());
409 assert!(ProviderId::parse("").is_none());
410 assert!(serde_json::from_str::<ProviderId>(r#""bad id""#).is_err());
411 }
412
413 #[test]
414 fn query_reports_the_filters_it_uses() {
415 let mut q = SearchQuery::new("transformers");
416 assert!(q.filters().is_empty());
417 q.years = Some(YearRange {
418 from: Some(2017),
419 to: None,
420 });
421 q.open_access = Some(true);
422 assert_eq!(q.filters(), [Filter::Years, Filter::OpenAccess]);
423 }
424
425 #[test]
426 fn unimplemented_operations_are_unsupported_and_dyn_compatible() {
427 let p: Box<dyn Provider> = Box::new(Nothing(Capabilities::default()));
428 let query = SearchQuery::new("x");
429 let result = block_on_ready(p.search(&query));
430 assert!(matches!(result, Err(Error::Unsupported("search"))));
431 }
432
433 fn block_on_ready<F: std::future::Future>(fut: F) -> F::Output {
435 use std::task::{Context, Poll, Waker};
436 let mut fut = std::pin::pin!(fut);
437 match fut.as_mut().poll(&mut Context::from_waker(Waker::noop())) {
438 Poll::Ready(out) => out,
439 Poll::Pending => panic!("future was not immediately ready"),
440 }
441 }
442}