feat: supports fetching models during configuration initialization (#1161)

1 files changed, 26 insertions, 1 deletions

diff --git a/src/utils/request.rs b/src/utils/request.rs
index efcd102..838110d 100644
--- a/src/utils/request.rs
+++ b/src/utils/request.rs
@@ -28,7 +28,7 @@ const USER_AGENT: &str = "curl/8.6.0";
 
 lazy_static::lazy_static! {
     static ref CLIENT: Result<reqwest::Client> = {
-        let builder = reqwest::ClientBuilder::new().timeout(Duration::from_secs(30));
+        let builder = reqwest::ClientBuilder::new().timeout(Duration::from_secs(16));
         let client = builder.build()?;
         Ok(client)
     };
@@ -158,6 +158,31 @@ pub async fn fetch_with_loaders(
     Ok(result)
 }
 
+pub async fn fetch_models(api_base: &str, api_key: Option<&str>) -> Result<Vec<String>> {
+    let client = match *CLIENT {
+        Ok(ref client) => client,
+        Err(ref err) => bail!("{err}"),
+    };
+    let mut request_builder = client.get(format!("{}/models", api_base.trim_end_matches('/')));
+    if let Some(api_key) = api_key {
+        request_builder = request_builder.bearer_auth(api_key);
+    }
+    let res_body: Value = request_builder.send().await?.json().await?;
+    let result: Vec<String> = res_body
+        .get("data")
+        .and_then(|v| v.as_array())
+        .map(|v| {
+            v.iter()
+                .filter_map(|v| v.get("id").and_then(|v| v.as_str().map(|v| v.to_string())))
+                .collect()
+        })
+        .unwrap_or_default();
+    if result.is_empty() {
+        bail!("No models")
+    }
+    Ok(result)
+}
+
 #[derive(Debug, Clone, Default)]
 pub struct CrawlOptions {
     extract: Option<String>,