From 5dfd0c04bdcfa683fda6fbfe323d1573bc30348d Mon Sep 17 00:00:00 2001
From: HugoCasa <hugo@casademont.ch>
Date: Thu, 23 Nov 2023 15:03:12 +0100
Subject: [PATCH] feat: query embeddings from s3 (#2683)

* feat: query embeddings from s3

* fix: better error handling + logging
---
 backend/windmill-api/src/embeddings.rs | 68 +++++++++++++++++++-------
 1 file changed, 51 insertions(+), 17 deletions(-)

diff --git a/backend/windmill-api/src/embeddings.rs b/backend/windmill-api/src/embeddings.rs
index 49b148822d..47628a6797 100644
--- a/backend/windmill-api/src/embeddings.rs
+++ b/backend/windmill-api/src/embeddings.rs
@@ -1,6 +1,6 @@
 use std::{collections::HashMap, path::PathBuf, sync::Arc};
 
-use anyhow::{self, Error, Result};
+use anyhow::{anyhow, Error, Result};
 use axum::{
     extract::{Path, Query},
     routing::get,
@@ -228,14 +228,31 @@ impl EmbeddingsDb {
         self.db
             .create_collection("resource_types".to_string(), 384, Distance::Cosine)?;
 
-        let response = http_get_from_hub(
-            &HTTP_CLIENT,
-            "https://hub.windmill.dev/scripts/embeddings",
-            false,
-            None,
-            pg_db,
-        )
-        .await?;
+        let response = HTTP_CLIENT
+            .get("https://bucket.windmillhub.com/embeddings/scripts_embeddings.json")
+            .send()
+            .await;
+        let response =
+            if response.is_err() || response.as_ref().unwrap().error_for_status_ref().is_err() {
+                tracing::warn!("Failed to get scripts embeddings from bucket, trying hub...");
+                http_get_from_hub(
+                    &HTTP_CLIENT,
+                    "https://hub.windmill.dev/scripts/embeddings",
+                    false,
+                    None,
+                    pg_db,
+                )
+                .await?
+            } else {
+                response.unwrap()
+            };
+        if response.error_for_status_ref().is_err() {
+            return Err(anyhow!(
+                "Failed to get scripts embeddings from hub with error code: {}",
+                response.status()
+            ));
+        }
+
         let hub_scripts = response.json::<Vec<HubScript>>().await?;
 
         for script in &hub_scripts {
@@ -257,14 +274,31 @@ impl EmbeddingsDb {
             self.db.insert_into_collection("scripts", embedding)?;
         }
 
-        let response = http_get_from_hub(
-            &HTTP_CLIENT,
-            "https://hub.windmill.dev/resource_types/embeddings",
-            false,
-            None,
-            pg_db,
-        )
-        .await?;
+        let response = HTTP_CLIENT
+            .get("https://bucket.windmillhub.com/embeddings/resource_types_embeddings.json")
+            .send()
+            .await;
+        let response = if response.is_err()
+            || response.as_ref().unwrap().error_for_status_ref().is_err()
+        {
+            tracing::warn!("Failed to get resource types embeddings from bucket, trying hub...");
+            http_get_from_hub(
+                &HTTP_CLIENT,
+                "https://hub.windmill.dev/resource_types/embeddings",
+                false,
+                None,
+                pg_db,
+            )
+            .await?
+        } else {
+            response.unwrap()
+        };
+        if response.error_for_status_ref().is_err() {
+            return Err(anyhow!(
+                "Failed to get resource types embeddings from hub with error code: {}",
+                response.status()
+            ));
+        }
         let hub_resource_types = response.json::<Vec<HubResourceType>>().await?;
 
         let resource_types: Vec<ResourceType> =