From d4087dbaf7208da5ae842ec5ba2e9c442bfa4f94 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Mon, 10 Feb 2025 08:01:21 +0800
Subject: feat: supports fetching models during configuration initialization
 (#1161)

---
 src/utils/request.rs | 27 ++++++++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

(limited to 'src/utils')

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>,
-- 
cgit v1.2.3