Skip to main content

txtai/
api.rs

1use enum_as_inner::EnumAsInner;
2use reqwest::{multipart, RequestBuilder, Response};
3use serde::Deserialize;
4use serde_json::Value;
5use std::env;
6use std::error::Error;
7
8// Generic return types
9pub type APIResponse = Result<Response, Box<dyn Error>>;
10pub type IndexResults = Result<Vec<IndexResult>, Box<dyn Error>>;
11pub type IndexResultsBatch = Result<Vec<Vec<IndexResult>>, Box<dyn Error>>;
12pub type Strings = Result<String, Box<dyn Error>>;
13pub type StringsBatch = Result<Vec<String>, Box<dyn Error>>;
14pub type Texts = Result<Text, Box<dyn Error>>;
15pub type TextsBatch = Result<Vec<Text>, Box<dyn Error>>;
16
17/// Base API definition
18pub struct API {
19    url: String,
20    token: String
21}
22
23/// Base API implementation
24impl API {
25    /// Creates an API instance.
26    ///
27    pub fn new() -> API {
28        API { 
29            url: env::var("TXTAI_API_URL").unwrap_or(String::from("")),
30            token: env::var("TXTAI_API_TOKEN").unwrap_or(String::from(""))
31        }
32    }
33
34    /// Creates an API instance.
35    /// 
36    /// # Arguments
37    /// * `url` - API url
38    pub fn with_url(url: &str) -> API {
39        API { 
40            url: url.to_string(),
41            token: env::var("TXTAI_API_TOKEN").unwrap_or(String::from(""))
42        }
43    }
44
45    /// Creates an API instance.
46    /// 
47    /// # Arguments
48    /// * `url` - API url
49    /// * `token` - API token
50    pub fn with_url_token(url: &str, token: &str) -> API {
51        API { 
52            url: url.to_string(),
53            token: token.to_string()
54        }
55    }
56
57    /// Executes a GET request. Returns Response.
58    /// 
59    /// # Arguments
60    /// * `method` - API method
61    /// * `params` - List of (name, value) tuples to use as query parameters
62    pub async fn get(&self, method: &str, params: &[(&str, &str)]) -> APIResponse {
63        // Generate url
64        let url = format!("{url}/{method}", url=self.url, method=method);
65
66        // Create client
67        let client = reqwest::Client::new();
68        let mut request = client.get(&url);
69
70        // Set headers
71        request = self.headers(request);
72
73        // Execute API call
74        Ok(request.query(&params).send().await?)
75    }
76
77    /// Executes a POST request. Returns Response.
78    /// 
79    /// # Arguments
80    /// * `method` - API method
81    /// * `json` - Value object with post body 
82    pub async fn post(&self, method: &str, json: &Value) -> APIResponse {
83        // Generate url
84        let url = format!("{url}/{method}", url=self.url, method=method);
85
86        // Create client
87        let client = reqwest::Client::new();
88        let mut request = client.post(&url);
89
90        // Set headers
91        request = self.headers(request);
92
93        // Execute API call
94        Ok(request.json(&json).send().await?)
95    }
96
97    /// Sets headers on a request.
98    ///
99    /// # Arguments
100    //  * `request` - RequestBuilder
101    pub fn headers(&self, request: RequestBuilder) -> RequestBuilder {
102        // Authorization header
103        if self.token != "" {
104            return request.header("Authorization", format!("Bearer {token}", token=self.token))
105        }
106        return request
107    }
108
109    /// Executes a multipart POST request. Returns Response.
110    ///
111    /// # Arguments
112    /// * `method` - API method
113    /// * `form` - Multipart form data
114    pub async fn post_multipart(&self, method: &str, form: multipart::Form) -> APIResponse {
115        // Generate url
116        let url = format!("{url}/{method}", url=self.url, method=method);
117
118        // Create client
119        let client = reqwest::Client::new();
120        let mut request = client.post(&url);
121
122        // Set headers
123        request = self.headers(request);
124
125        // Execute API call
126        Ok(request.multipart(form).send().await?)
127    }
128}
129
130// Index result
131#[derive(Debug, Deserialize)]
132pub struct IndexResult {
133    pub id: usize,
134    pub score: f32
135}
136
137// Text result that handles String and Vector of Strings
138#[derive(Debug, Deserialize, EnumAsInner)]
139#[serde(untagged)]
140pub enum Text {
141    String(String),
142    List(Vec<String>)
143}