trace_moe/
tracemoe.rs

1use reqwest::header::{HeaderName, HeaderValue};
2use reqwest::multipart::{Form, Part};
3use serde::Deserialize;
4
5use crate::client::Client;
6use crate::error::Result;
7
8const DEFAULT_BASE: &str = "https://api.trace.moe/";
9
10/// Create a `Client` configured for the trace.moe API.
11///
12/// If an API key is provided, it will be sent via the `x-trace-key` header.
13/// For higher quotas, obtain a key from the trace.moe dashboard.
14pub fn new_client_with_key(api_key: Option<&str>) -> Result<Client> {
15    let mut client = Client::new(DEFAULT_BASE)?;
16    if let Some(key) = api_key {
17        client = client.with_default_header(
18            HeaderName::from_static("x-trace-key"),
19            HeaderValue::from_str(key)?,
20        );
21    }
22    Ok(client)
23}
24
25#[derive(Debug, Deserialize)]
26#[serde(rename_all = "camelCase")]
27/// Response returned by trace.moe search endpoints.
28pub struct SearchResponse<TAnilist = i64> {
29    pub frame_count: i64,
30    pub error: String,
31    pub result: Vec<SearchResult<TAnilist>>,
32}
33
34#[derive(Debug, Deserialize)]
35#[serde(rename_all = "camelCase")]
36/// A single search hit.
37pub struct SearchResult<TAnilist = i64> {
38    pub anilist: TAnilist,
39    pub filename: String,
40    pub episode: Option<Episode>,
41    pub duration: f64,
42    pub from: f64,
43    pub to: f64,
44    pub at: f64,
45    pub similarity: f64,
46    pub image: String,
47    pub video: String,
48}
49
50#[derive(Debug, Deserialize)]
51#[serde(untagged)]
52/// Episode info in results may be a number or text.
53pub enum Episode {
54    Number(i64),
55    Text(String),
56}
57
58#[derive(Debug, Deserialize)]
59#[serde(rename_all = "camelCase")]
60/// AniList title variants.
61pub struct AnilistInfoTitle {
62    pub native: Option<String>,
63    pub romaji: Option<String>,
64    pub english: Option<String>,
65}
66
67#[derive(Debug, Deserialize)]
68#[serde(rename_all = "camelCase")]
69/// AniList metadata when `anilist_info` is requested.
70pub struct AnilistInfo {
71    pub id: i64,
72    pub id_mal: Option<i64>,
73    pub title: AnilistInfoTitle,
74    pub synonyms: Vec<String>,
75    pub is_adult: bool,
76}
77
78#[derive(Debug, Deserialize)]
79#[serde(rename_all = "camelCase")]
80/// Response from `GET /me` describing quota and concurrency.
81pub struct MeResponse {
82    pub id: String,
83    pub priority: i64,
84    pub concurrency: i64,
85    pub quota: i64,
86    pub quota_used: i64,
87}
88
89#[derive(Default, Debug, Clone)]
90/// Search query parameters for `search` endpoints.
91pub struct SearchQuery {
92    pub url: Option<String>,
93    pub anilist_id: Option<i64>,
94    pub cut_borders: Option<bool>,
95    pub anilist_info: Option<bool>,
96}
97
98impl Client {
99    /// Search by image URL with optional parameters.
100    pub async fn tracemoe_search_by_url<TAnilist: for<'de> serde::Deserialize<'de>>(
101        &self,
102        query: &SearchQuery,
103    ) -> Result<SearchResponse<TAnilist>> {
104        let path = build_query_path("search", query);
105        self.get_json::<SearchResponse<TAnilist>>(path).await
106    }
107
108    /// Upload image bytes to search.
109    pub async fn tracemoe_search_upload<TAnilist: for<'de> serde::Deserialize<'de>>(
110        &self,
111        bytes: impl Into<Vec<u8>>,
112    ) -> Result<SearchResponse<TAnilist>> {
113        let form = Form::new().part("image", Part::bytes(bytes.into()));
114        let resp = self
115            .request(reqwest::Method::POST, "search")?
116            .multipart(form)
117            .send()
118            .await?;
119        Self::parse_json(resp).await
120    }
121
122    /// Get current account quota and concurrency information.
123    pub async fn tracemoe_me(&self) -> Result<MeResponse> {
124        self.get_json("me").await
125    }
126}
127
128/// Build a relative path with query string from a `SearchQuery`.
129fn build_query_path(base: &str, query: &SearchQuery) -> String {
130    let mut qp = url::form_urlencoded::Serializer::new(String::new());
131    if let Some(v) = &query.url { qp.append_pair("url", v); }
132    if let Some(v) = query.anilist_id { qp.append_pair("anilistID", &v.to_string()); }
133    if let Some(true) = query.cut_borders { qp.append_pair("cutBorders", ""); }
134    if let Some(true) = query.anilist_info { qp.append_pair("anilistInfo", ""); }
135    let qs = qp.finish();
136    if qs.is_empty() { base.to_string() } else { format!("{}?{}", base, qs) }
137}
trace_moe/tracemoe.rs

trace_moe/
tracemoe.rs
