ferrin_core/prompt/
download.rs

1//! URL downloads for file parts the model cannot fetch itself.
2
3use std::fmt;
4use std::sync::Arc;
5
6use bytes::Bytes;
7use ferrin_provider_util::SharedTransport;
8use ferrin_provider_util::UrlPolicy;
9use ferrin_provider_util::secure_url::DownloadError;
10use ferrin_provider_util::secure_url::DownloadErrorKind;
11use ferrin_spec::BoxFuture;
12use ferrin_spec::MediaType;
13use tokio::task::JoinSet;
14use tokio_util::sync::CancellationToken;
15use url::Url;
16
17use crate::error::Error;
18use crate::limits::DEFAULT_MAX_PARALLEL_DOWNLOADS;
19
20/// One URL to download.
21#[derive(Debug, Clone, PartialEq, Eq)]
22pub struct DownloadRequest {
23    /// The URL.
24    pub url: Url,
25    /// Whether the model can fetch the URL itself (matched against
26    /// `supported_urls`).
27    pub is_url_supported_by_model: bool,
28}
29
30/// A downloaded file.
31#[derive(Debug, Clone, PartialEq, Eq)]
32pub struct DownloadedFile {
33    /// The bytes.
34    pub data: Bytes,
35    /// Media type from the response headers, if any.
36    pub media_type: Option<MediaType>,
37}
38
39/// Downloads files referenced by URL.
40///
41/// Implementations return one entry per request in the same order; `None`
42/// keeps the URL in the prompt for the provider to fetch.
43pub trait DownloadFn: Send + Sync {
44    /// Downloads `requests`.
45    fn download(
46        &self,
47        requests: Vec<DownloadRequest>,
48        cancellation: CancellationToken,
49    ) -> BoxFuture<'_, Result<Vec<Option<DownloadedFile>>, Error>>;
50}
51
52/// The default downloader: fetches only URLs the model does not support,
53/// through the secure URL policy (HTTPS, no private networks, size limit).
54#[derive(Clone)]
55pub struct DefaultDownloader {
56    transport: SharedTransport,
57    policy: Arc<UrlPolicy>,
58    max_parallel: usize,
59}
60
61impl DefaultDownloader {
62    /// Creates a downloader on `transport` with the default policy.
63    #[must_use]
64    pub fn new(transport: SharedTransport) -> Self {
65        Self {
66            transport,
67            policy: Arc::new(UrlPolicy::default()),
68            max_parallel: DEFAULT_MAX_PARALLEL_DOWNLOADS,
69        }
70    }
71
72    /// Creates a downloader on the default HTTP transport.
73    ///
74    /// # Errors
75    ///
76    /// Returns [`Error::Other`] when the transport cannot be built.
77    pub fn try_default() -> Result<Self, Error> {
78        let transport = ferrin_provider_util::default_transport().map_err(Error::other)?;
79        Ok(Self::new(transport))
80    }
81
82    /// Replaces the URL policy.
83    #[must_use]
84    pub fn with_policy(mut self, policy: UrlPolicy) -> Self {
85        self.policy = Arc::new(policy);
86        self
87    }
88
89    /// Sets the maximum number of concurrent downloads (at least 1).
90    #[must_use]
91    pub fn with_max_parallel(mut self, max_parallel: usize) -> Self {
92        self.max_parallel = max_parallel.max(1);
93        self
94    }
95
96    async fn fetch_one(
97        transport: SharedTransport,
98        policy: Arc<UrlPolicy>,
99        url: Url,
100        cancellation: CancellationToken,
101    ) -> Result<DownloadedFile, Error> {
102        let downloaded =
103            ferrin_provider_util::secure_url::fetch(transport.as_ref(), url, &policy, cancellation)
104                .await
105                .map_err(download_error)?;
106        Ok(DownloadedFile {
107            data: downloaded.data,
108            media_type: downloaded.media_type,
109        })
110    }
111}
112
113impl fmt::Debug for DefaultDownloader {
114    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
115        f.debug_struct("DefaultDownloader")
116            .field("policy", &self.policy)
117            .field("max_parallel", &self.max_parallel)
118            .finish_non_exhaustive()
119    }
120}
121
122impl DownloadFn for DefaultDownloader {
123    fn download(
124        &self,
125        requests: Vec<DownloadRequest>,
126        cancellation: CancellationToken,
127    ) -> BoxFuture<'_, Result<Vec<Option<DownloadedFile>>, Error>> {
128        Box::pin(async move {
129            let mut results: Vec<Option<DownloadedFile>> = vec![None; requests.len()];
130            let mut pending = requests
131                .into_iter()
132                .enumerate()
133                .filter(|(_, request)| !request.is_url_supported_by_model);
134            let mut tasks: JoinSet<(usize, Result<DownloadedFile, Error>)> = JoinSet::new();
135            let mut spawn_next = |tasks: &mut JoinSet<(usize, Result<DownloadedFile, Error>)>| {
136                if let Some((index, request)) = pending.next() {
137                    let transport = Arc::clone(&self.transport);
138                    let policy = Arc::clone(&self.policy);
139                    let cancellation = cancellation.clone();
140                    tasks.spawn(async move {
141                        let result =
142                            Self::fetch_one(transport, policy, request.url, cancellation).await;
143                        (index, result)
144                    });
145                    true
146                } else {
147                    false
148                }
149            };
150            for _ in 0..self.max_parallel {
151                if !spawn_next(&mut tasks) {
152                    break;
153                }
154            }
155            while let Some(joined) = tasks.join_next().await {
156                let (index, result) = joined
157                    .map_err(|error| Error::message(format!("download task failed: {error}")))?;
158                match result {
159                    Ok(file) => results[index] = Some(file),
160                    Err(error) => {
161                        tasks.abort_all();
162                        return Err(error);
163                    }
164                }
165                spawn_next(&mut tasks);
166            }
167            Ok(results)
168        })
169    }
170}
171
172fn download_error(error: DownloadError) -> Error {
173    if error.is_cancelled() {
174        return Error::Cancelled;
175    }
176    let url = error.url().clone();
177    let status_code = match error.kind() {
178        DownloadErrorKind::Status { status, .. } => Some(*status),
179        _ => None,
180    };
181    Error::download(url, status_code, Some(Box::new(error)))
182}
ferrin_core/prompt/download.rs

ferrin_core/prompt/
download.rs
