crawlkit_extensions/jina/
mod.rs1use std::collections::HashMap;
27use std::time::Duration;
28
29use async_trait::async_trait;
30use backon::{ExponentialBuilder, Retryable};
31use crawlkit_core::{CrawlError, HttpClient, Response};
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
35pub enum JinaFormat {
36 #[default]
38 Markdown,
39 Html,
41 Text,
43}
44
45impl JinaFormat {
46 fn accept_header(&self) -> &'static str {
47 match self {
48 Self::Markdown => "text/markdown",
49 Self::Html => "text/html",
50 Self::Text => "text/plain",
51 }
52 }
53}
54
55pub struct JinaClient {
59 name: String,
60 api_token: Option<String>,
61 client: reqwest::Client,
62 max_retries: usize,
63 format: JinaFormat,
64}
65
66impl JinaClient {
67 pub fn builder() -> JinaClientBuilder {
69 JinaClientBuilder::default()
70 }
71
72 fn from_builder(builder: JinaClientBuilder) -> Self {
74 let client = reqwest::Client::builder()
75 .timeout(builder.timeout)
76 .build()
77 .expect("failed to build reqwest client");
78
79 Self {
80 name: builder.name,
81 api_token: builder.api_token,
82 client,
83 max_retries: builder.max_retries,
84 format: builder.format,
85 }
86 }
87}
88
89impl Default for JinaClient {
90 fn default() -> Self {
91 Self::builder().build()
92 }
93}
94
95pub struct JinaClientBuilder {
97 name: String,
98 api_token: Option<String>,
99 timeout: Duration,
100 max_retries: usize,
101 format: JinaFormat,
102}
103
104impl Default for JinaClientBuilder {
105 fn default() -> Self {
106 Self {
107 name: "jina".to_string(),
108 api_token: None,
109 timeout: Duration::from_secs(60),
110 max_retries: 3,
111 format: JinaFormat::default(),
112 }
113 }
114}
115
116impl JinaClientBuilder {
117 pub fn with_name(mut self, name: impl Into<String>) -> Self {
119 self.name = name.into();
120 self
121 }
122
123 pub fn with_token(mut self, token: impl Into<String>) -> Self {
125 self.api_token = Some(token.into());
126 self
127 }
128
129 pub fn with_token_from_env(mut self, env_var: &str) -> Self {
131 self.api_token = std::env::var(env_var).ok().filter(|s| !s.is_empty());
132 self
133 }
134
135 pub fn with_timeout(mut self, timeout: Duration) -> Self {
137 self.timeout = timeout;
138 self
139 }
140
141 pub fn with_max_retries(mut self, max_retries: usize) -> Self {
143 self.max_retries = max_retries;
144 self
145 }
146
147 pub fn with_format(mut self, format: JinaFormat) -> Self {
153 self.format = format;
154 self
155 }
156
157 pub fn build(self) -> JinaClient {
159 JinaClient::from_builder(self)
160 }
161}
162
163#[async_trait]
164impl HttpClient for JinaClient {
165 async fn get(&self, url: &str, _headers: &HashMap<String, String>) -> crawlkit_core::Result<Response> {
166 let client = self.client.clone();
167 let reader_url = format!("https://r.jina.ai/{url}");
168 let token = self.api_token.clone();
169 let max_retries = self.max_retries;
170 let accept = self.format.accept_header().to_string();
171
172 let fetch = || {
173 let client = client.clone();
174 let reader_url = reader_url.clone();
175 let token = token.clone();
176 let accept = accept.clone();
177 async move {
178 let mut req = client
179 .get(&reader_url)
180 .header("Accept", &accept);
181
182 if let Some(ref t) = token {
183 req = req.header("Authorization", format!("Bearer {t}"));
184 }
185
186 let resp = req.send().await.map_err(|e| {
187 anyhow::anyhow!("Jina 请求失败: {e}")
188 })?;
189
190 let status = resp.status().as_u16();
191 let final_url = resp.url().to_string();
192 let body = resp.text().await.map_err(|e| {
193 anyhow::anyhow!("Jina 读取响应失败: {e}")
194 })?;
195
196 if status >= 400 {
197 return Err(anyhow::anyhow!("Jina 返回错误状态 {status}: {body}"));
198 }
199
200 Ok((final_url, status, body))
201 }
202 };
203
204 let (final_url, status, body) = fetch
205 .retry(&ExponentialBuilder::default().with_max_times(max_retries))
206 .await
207 .map_err(|e| CrawlError::Http(format!("Jina 请求失败(重试 {max_retries} 次): {e}")))?;
208
209 Ok(Response {
210 url: final_url,
211 status,
212 headers: Default::default(),
213 body,
214 })
215 }
216
217 async fn post(
218 &self,
219 url: &str,
220 headers: &HashMap<String, String>,
221 _body: Vec<u8>,
222 ) -> crawlkit_core::Result<Response> {
223 self.get(url, headers).await
224 }
225
226 fn name(&self) -> &str {
227 &self.name
228 }
229}
230
231#[cfg(test)]
232mod tests {
233 use super::*;
234
235 #[test]
236 fn test_jina_format_default() {
237 assert_eq!(JinaFormat::default(), JinaFormat::Markdown);
238 }
239
240 #[test]
241 fn test_jina_format_accept_header() {
242 assert_eq!(JinaFormat::Markdown.accept_header(), "text/markdown");
243 assert_eq!(JinaFormat::Html.accept_header(), "text/html");
244 assert_eq!(JinaFormat::Text.accept_header(), "text/plain");
245 }
246
247 #[test]
248 fn test_jina_builder_defaults() {
249 let client = JinaClient::builder().build();
250 assert_eq!(client.name(), "jina");
251 assert!(client.api_token.is_none());
252 }
253
254 #[test]
255 fn test_jina_builder_with_token() {
256 let client = JinaClient::builder()
257 .with_token("test-token")
258 .build();
259 assert_eq!(client.api_token.as_deref(), Some("test-token"));
260 }
261
262 #[test]
263 fn test_jina_builder_with_name() {
264 let client = JinaClient::builder()
265 .with_name("custom")
266 .build();
267 assert_eq!(client.name(), "custom");
268 }
269}