feat: keep the resource type search fresh and serve it without embeddings

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

ai_evals/adapters/frontend/mockBackend.ts

@@ -12,7 +12,6 @@ import type {
 	ListableResource,
 	ListableVariable,
 	Resource,
-	ResourceType,
 	Script
 } from '../../../frontend/src/lib/gen'
 import type {
@@ -100,7 +99,7 @@ export interface BenchmarkWorkspaceResource {
 	description?: string
 }
 
-/** A resource type the instance holds. Without a schema, a lookup by name misses it. */
+/** A resource type the instance holds. */
 export interface BenchmarkWorkspaceResourceType {
 	name: string
 	schema?: unknown
@@ -352,18 +351,8 @@ export function listBenchmarkResourceTypeNames(workspace: string): string[] | nu
 	return benchmarkResourceTypes(workspace).map((rt) => rt.name)
 }
 
-/** A resource type of a benchmark workspace, as `ResourceService.getResourceType` returns it. */
-export function getBenchmarkResourceType(workspace: string, name: string): ResourceType | null {
-	const seed = benchmarkResourceTypes(workspace).find(
-		(rt) => rt.name === name && rt.schema !== undefined
-	)
-	return seed
-		? { workspace_id: 'admins', name: seed.name, schema: seed.schema, description: seed.description }
-		: null
-}
-
-/** Stands in for the embedding search: a type matches when a word of the query appears in its
- * name or description, ranked by how many do. */
+/** Stands in for the resource type search: a type matches when a word of the query appears in
+ * its name or description, ranked by how many do. */
 export function queryBenchmarkResourceTypes(
 	workspace: string,
 	text: string,

ai_evals/adapters/frontend/vitestAdapter.test.ts

@@ -104,7 +104,6 @@ vi.mock('$lib/gen', async () => {
 		getBenchmarkVariableByPath,
 		hasBenchmarkWorkspace,
 		getBenchmarkResource,
-		getBenchmarkResourceType,
 		listBenchmarkAiProviderResources,
 		listBenchmarkPlainResources,
 		listBenchmarkResourceTypeNames,
@@ -471,17 +470,7 @@ vi.mock('$lib/gen', async () => {
 					: actual.ResourceService.queryResourceTypes(data),
 			listResourceTypeNames: async (data: { workspace: string }) =>
 				listBenchmarkResourceTypeNames(data.workspace) ??
-				actual.ResourceService.listResourceTypeNames(data),
-			getResourceType: async (data: { workspace: string; path: string }) => {
-				if (!hasBenchmarkWorkspace(data.workspace)) {
-					return actual.ResourceService.getResourceType(data)
-				}
-				const resourceType = getBenchmarkResourceType(data.workspace, data.path)
-				if (!resourceType) {
-					throw new Error(`Resource type "${data.path}" not found in benchmark workspace`)
-				}
-				return resourceType
-			}
+				actual.ResourceService.listResourceTypeNames(data)
 		}),
 		VariableService: wrapService(actual.VariableService, {
 			existsVariable: async (data: { workspace: string; path: string }) =>

backend/windmill-api-embeddings/src/lib.rs

@@ -9,16 +9,11 @@ use windmill_common::DEFAULT_HUB_BASE_URL;
 #[cfg(feature = "embedding")]
 use windmill_common::HUB_BASE_URL;
 
-use axum::Router;
-
-#[cfg(feature = "embedding")]
 use axum::{
     extract::{Path, Query},
-    Json,
+    routing::get,
+    Extension, Json, Router,
 };
-
-#[cfg(feature = "embedding")]
-use axum::routing::get;
 #[cfg(feature = "embedding")]
 use candle_core::{Device, Tensor};
 #[cfg(feature = "embedding")]
@@ -27,10 +22,7 @@ use candle_nn::VarBuilder;
 use candle_transformers::models::bert::{BertModel, Config, DTYPE};
 #[cfg(feature = "embedding")]
 use hf_hub::api::tokio::Api;
-#[cfg(feature = "embedding")]
-use serde::Deserialize;
-#[cfg(feature = "embedding")]
-use serde::Serialize;
+use serde::{Deserialize, Serialize};
 #[cfg(feature = "embedding")]
 use sqlx::{Pool, Postgres};
 #[cfg(feature = "embedding")]
@@ -44,9 +36,7 @@ use tokenizers::Tokenizer;
 use tokio::sync::RwLock;
 #[cfg(feature = "embedding")]
 use windmill_common::utils::http_get_from_hub;
-
-#[cfg(feature = "embedding")]
-use windmill_common::error::JsonResult;
+use windmill_common::{db::DB, error::JsonResult};
 
 #[cfg(feature = "embedding")]
 use windmill_store::resources::ResourceType;
@@ -60,6 +50,7 @@ lazy_static::lazy_static! {
     // full pulling interval, so a transient startup error doesn't leave the
     // embeddings DB uninitialized for a whole day.
     pub static ref HUB_EMBEDDINGS_RETRY_INTERVAL_SECS: u64 = std::env::var("HUB_EMBEDDINGS_RETRY_INTERVAL_SECS").ok().map(|x| x.parse::<u64>().ok()).flatten().unwrap_or(60);
+    pub static ref RESOURCE_TYPE_EMBEDDINGS_CHECK_INTERVAL_SECS: u64 = std::env::var("RESOURCE_TYPE_EMBEDDINGS_CHECK_INTERVAL_SECS").ok().map(|x| x.parse::<u64>().ok()).flatten().unwrap_or(60);
 }
 
 #[cfg(feature = "embedding")]
@@ -104,14 +95,12 @@ async fn query_hub_scripts(
     }
 }
 
-#[cfg(feature = "embedding")]
 #[derive(Deserialize)]
 struct ResourceTypesQuery {
     text: String,
     limit: Option<i64>,
 }
 
-#[cfg(feature = "embedding")]
 #[derive(Serialize)]
 pub struct ResourceTypeResult {
     name: String,
@@ -138,24 +127,80 @@ fn trim_to_top_score<T>(
         .take_while(|r| (top_score - score(r)) / top_score <= max_relative_drop)
         .collect()
 }
-#[cfg(feature = "embedding")]
 async fn query_resource_types(
+    Extension(db): Extension<DB>,
     Query(query): Query<ResourceTypesQuery>,
     Path(w_id): Path<String>,
 ) -> JsonResult<Vec<ResourceTypeResult>> {
-    let embeddings_db = EMBEDDINGS_DB.read().await;
-
-    if let Some(embeddings_db) = embeddings_db.as_ref() {
-        let results = embeddings_db
-            .query_resource_types(w_id, &query.text, query.limit)
-            .await?;
-
-        Ok(Json(results))
-    } else {
-        Err(windmill_common::error::Error::internal_err(
-            "Embeddings db not initialized".to_string(),
-        ))
+    #[cfg(feature = "embedding")]
+    if let Some(embeddings_db) = EMBEDDINGS_DB.read().await.as_ref() {
+        return Ok(Json(
+            embeddings_db
+                .query_resource_types(w_id, &query.text, query.limit)
+                .await?,
+        ));
     }
+    Ok(Json(
+        match_resource_types_by_text(&db, &w_id, &query.text, query.limit).await?,
+    ))
+}
+
+/// Serves the search when there is no embeddings index: a build without the `embedding`
+/// feature, or one whose index never loaded (air-gapped, `DISABLE_EMBEDDING`). The AI chat
+/// has no other way to read the instance's resource types.
+async fn match_resource_types_by_text(
+    db: &DB,
+    w_id: &str,
+    text: &str,
+    limit: Option<i64>,
+) -> windmill_common::error::Result<Vec<ResourceTypeResult>> {
+    let words: Vec<String> = text
+        .to_lowercase()
+        .split(|c: char| !c.is_alphanumeric())
+        .filter(|word| word.len() >= 3)
+        .map(str::to_string)
+        .collect();
+    let rows: Vec<(String, Option<String>, Option<serde_json::Value>)> = sqlx::query_as(
+        "SELECT name, description, schema FROM resource_type WHERE workspace_id = $1 OR workspace_id = 'admins'",
+    )
+    .bind(w_id)
+    .fetch_all(db)
+    .await?;
+
+    let mut results: Vec<ResourceTypeResult> = rows
+        .into_iter()
+        .map(|(name, description, schema)| ResourceTypeResult {
+            score: text_match_score(&words, &name, description.as_deref().unwrap_or_default()),
+            name,
+            schema,
+        })
+        .filter(|rt| rt.score > 0.0)
+        .collect();
+    results.sort_by(|a, b| {
+        b.score
+            .total_cmp(&a.score)
+            .then_with(|| a.name.cmp(&b.name))
+    });
+    results.truncate(limit.unwrap_or(10).max(0) as usize);
+    Ok(results)
+}
+
+/// A word found in the name counts twice a word found in the description: names are what
+/// a query for a service spells out ("postgres" in `postgresql`, "sheets" in `gsheets`).
+fn text_match_score(words: &[String], name: &str, description: &str) -> f32 {
+    let description = description.to_lowercase();
+    words
+        .iter()
+        .map(|word| {
+            if name.contains(word.as_str()) {
+                2.0
+            } else if description.contains(word.as_str()) {
+                1.0
+            } else {
+                0.0
+            }
+        })
+        .sum()
 }
 
 #[cfg(feature = "embedding")]
@@ -288,6 +333,11 @@ impl ModelInstance {
 pub struct EmbeddingsDb {
     db: Db,
     model_instance: Arc<ModelInstance>,
+    /// Kept from the last full fill, so the resource types can be re-indexed between fills
+    /// without downloading the hub's embeddings again.
+    hub_resource_type_embeddings: Arc<HashMap<String, Vec<f32>>>,
+    /// The `resource_types_fingerprint` the resource types collection was built from.
+    resource_types_fingerprint: Option<String>,
 }
 
 #[cfg(feature = "embedding")]
@@ -295,7 +345,12 @@ impl EmbeddingsDb {
     pub async fn new(pg_db: &Pool<Postgres>, model_instance: Arc<ModelInstance>) -> Result<Self> {
         let db = Db::new();
 
-        let mut embeddings_db = Self { db, model_instance: model_instance.clone() };
+        let mut embeddings_db = Self {
+            db,
+            model_instance: model_instance.clone(),
+            hub_resource_type_embeddings: Arc::new(HashMap::new()),
+            resource_types_fingerprint: None,
+        };
 
         embeddings_db.fill_db(pg_db).await?;
 
@@ -310,13 +365,6 @@ impl EmbeddingsDb {
         self.db
             .create_collection("scripts".to_string(), 384, Distance::Cosine)?;
 
-        if self.db.get_collection("resource_types").is_some() {
-            self.db.delete_collection("resource_types")?;
-        }
-
-        self.db
-            .create_collection("resource_types".to_string(), 384, Distance::Cosine)?;
-
         let hub_base_url = (**HUB_BASE_URL.load()).clone();
 
         let response = match hub_base_url.as_str() {
@@ -425,44 +473,37 @@ impl EmbeddingsDb {
             ));
         }
         let hub_resource_types = response.json::<Vec<HubResourceType>>().await?;
+        self.hub_resource_type_embeddings = Arc::new(
+            hub_resource_types
+                .into_iter()
+                .map(|rt| (rt.name, rt.embedding))
+                .collect(),
+        );
 
-        let resource_types: Vec<ResourceType> =
-            sqlx::query_as!(ResourceType, "SELECT workspace_id, name, schema, description, created_by, edited_at, format_extension, is_fileset, display_name from resource_type ORDER BY name",)
-                .fetch_all(pg_db)
-                .await?;
-
-        for rt in resource_types {
-            let mut hm = HashMap::new();
-            hm.insert("name".to_string(), rt.name.clone());
-            if let Some(schema) = rt.schema.clone() {
-                hm.insert("schema".to_string(), serde_json::to_string(&schema)?);
-            }
-            hm.insert("workspace".to_string(), rt.workspace_id.clone());
-            let hub_rt = hub_resource_types.iter().find(|hrt| hrt.name == rt.name);
-
-            let vector = if let Some(hub_rt) = hub_rt {
-                hub_rt.embedding.clone()
-            } else {
-                self.model_instance
-                    .clone()
-                    .create_embedding(&format!(
-                        "{};{}",
-                        rt.name,
-                        rt.description.unwrap_or_default()
-                    ))
-                    .await?
-            };
-
-            let embedding = Embedding {
-                id: format!("{}_{}", rt.workspace_id, rt.name),
-                vector,
-                metadata: Some(hm),
-            };
+        let (fingerprint, embeddings) = resource_type_embeddings(
+            pg_db,
+            &self.model_instance,
+            &self.hub_resource_type_embeddings,
+        )
+        .await?;
+        self.replace_resource_types(fingerprint, embeddings)
+    }
 
+    fn replace_resource_types(
+        &mut self,
+        fingerprint: Option<String>,
+        embeddings: Vec<Embedding>,
+    ) -> Result<()> {
+        if self.db.get_collection("resource_types").is_some() {
+            self.db.delete_collection("resource_types")?;
+        }
+        self.db
+            .create_collection("resource_types".to_string(), 384, Distance::Cosine)?;
+        for embedding in embeddings {
             self.db
                 .insert_into_collection("resource_types", embedding)?;
         }
-
+        self.resource_types_fingerprint = fingerprint;
         Ok(())
     }
 
@@ -617,6 +658,83 @@ impl EmbeddingsDb {
     }
 }
 
+/// Covers everything an indexed resource type is built from, whichever path wrote it: the
+/// create and update routes, a hub sync's direct upserts, a delete.
+#[cfg(feature = "embedding")]
+async fn resource_types_fingerprint(pg_db: &Pool<Postgres>) -> Result<Option<String>> {
+    Ok(sqlx::query_scalar(
+        "SELECT md5(string_agg(workspace_id || '/' || name || '/' || coalesce(description, '') || '/' || coalesce(schema::text, ''), ',' ORDER BY workspace_id, name)) FROM resource_type",
+    )
+    .fetch_one(pg_db)
+    .await?)
+}
+
+/// The fingerprint is read before the rows: a write landing in between leaves it stale, so
+/// the next check indexes again rather than keeping an index that misses the write.
+#[cfg(feature = "embedding")]
+async fn resource_type_embeddings(
+    pg_db: &Pool<Postgres>,
+    model_instance: &Arc<ModelInstance>,
+    hub_embeddings: &HashMap<String, Vec<f32>>,
+) -> Result<(Option<String>, Vec<Embedding>)> {
+    let fingerprint = resource_types_fingerprint(pg_db).await?;
+    let resource_types: Vec<ResourceType> =
+        sqlx::query_as!(ResourceType, "SELECT workspace_id, name, schema, description, created_by, edited_at, format_extension, is_fileset, display_name from resource_type ORDER BY name",)
+            .fetch_all(pg_db)
+            .await?;
+
+    let mut embeddings = Vec::with_capacity(resource_types.len());
+    for rt in resource_types {
+        let mut hm = HashMap::new();
+        hm.insert("name".to_string(), rt.name.clone());
+        if let Some(schema) = rt.schema.clone() {
+            hm.insert("schema".to_string(), serde_json::to_string(&schema)?);
+        }
+        hm.insert("workspace".to_string(), rt.workspace_id.clone());
+
+        let vector = match hub_embeddings.get(&rt.name) {
+            Some(vector) => vector.clone(),
+            None => {
+                model_instance
+                    .clone()
+                    .create_embedding(&format!(
+                        "{};{}",
+                        rt.name,
+                        rt.description.unwrap_or_default()
+                    ))
+                    .await?
+            }
+        };
+
+        embeddings.push(Embedding {
+            id: format!("{}_{}", rt.workspace_id, rt.name),
+            vector,
+            metadata: Some(hm),
+        });
+    }
+    Ok((fingerprint, embeddings))
+}
+
+/// The full fill runs once a day, so a resource type created or synced in between would stay
+/// unsearchable until then; this re-indexes the resource types alone when the table changed.
+#[cfg(feature = "embedding")]
+pub async fn refresh_resource_type_embeddings(pg_db: &Pool<Postgres>) -> Result<()> {
+    let fingerprint = resource_types_fingerprint(pg_db).await?;
+    let (model_instance, hub_embeddings) = match EMBEDDINGS_DB.read().await.as_ref() {
+        Some(db) if db.resource_types_fingerprint != fingerprint => (
+            db.model_instance.clone(),
+            db.hub_resource_type_embeddings.clone(),
+        ),
+        _ => return Ok(()),
+    };
+    let (fingerprint, embeddings) =
+        resource_type_embeddings(pg_db, &model_instance, &hub_embeddings).await?;
+    if let Some(db) = EMBEDDINGS_DB.write().await.as_mut() {
+        db.replace_resource_types(fingerprint, embeddings)?;
+    }
+    Ok(())
+}
+
 #[cfg(feature = "embedding")]
 fn normalize_l2(v: &Tensor) -> Result<Tensor> {
     Ok(v.broadcast_div(&v.sqr()?.sum_keepdim(1)?.sqrt()?)?)
@@ -630,6 +748,19 @@ pub fn load_embeddings_db(db: &Pool<Postgres>) -> () {
         .unwrap_or(false);
 
     if !disable_embedding {
+        let resource_types_db = db.clone();
+        tokio::spawn(async move {
+            loop {
+                tokio::time::sleep(std::time::Duration::from_secs(
+                    *RESOURCE_TYPE_EMBEDDINGS_CHECK_INTERVAL_SECS,
+                ))
+                .await;
+                if let Err(e) = refresh_resource_type_embeddings(&resource_types_db).await {
+                    tracing::warn!("Failed to refresh resource type embeddings: {e:#}");
+                }
+            }
+        });
+
         let db_clone = db.clone();
         tokio::spawn(async move {
             // Keep retrying model init: a transient failure here must not
@@ -696,7 +827,6 @@ pub async fn update_embeddings_db(db: &Pool<Postgres>) -> bool {
     }
 }
 
-#[cfg(feature = "embedding")]
 pub fn workspaced_service() -> Router {
     Router::new().route("/query_resource_types", get(query_resource_types))
 }
@@ -706,11 +836,6 @@ pub fn global_service() -> Router {
     Router::new().route("/query_hub_scripts", get(query_hub_scripts))
 }
 
-#[cfg(not(feature = "embedding"))]
-pub fn workspaced_service() -> Router {
-    Router::new()
-}
-
 #[cfg(not(feature = "embedding"))]
 pub fn global_service() -> Router {
     Router::new()
@@ -768,3 +893,24 @@ mod tests {
         assert_eq!(trim_to_top_score(vec![0.42f32], 0.05, |s| *s), vec![0.42]);
     }
 }
+
+#[cfg(test)]
+mod text_match_tests {
+    use super::text_match_score;
+
+    // A query names the service, which a type's name only contains ("postgres" in
+    // `postgresql`): an exact comparison finds nothing for "postgres database".
+    #[test]
+    fn ranks_a_name_match_above_a_description_match() {
+        let words = vec!["postgres".to_string(), "database".to_string()];
+        assert_eq!(
+            text_match_score(&words, "postgresql", "PostgreSQL connection"),
+            2.0
+        );
+        assert_eq!(
+            text_match_score(&words, "mysql", "MySQL database credentials"),
+            1.0
+        );
+        assert_eq!(text_match_score(&words, "slack", "Slack bot token"), 0.0);
+    }
+}

frontend/src/lib/components/copilot/chat/script/core.test.ts

@@ -18,48 +18,16 @@ import { ResourceService } from '$lib/gen'
 import { editCodeToolWithDiff, searchResourceTypes } from './core'
 
 describe('searchResourceTypes', () => {
-	const syncedNames = Array.from({ length: 30 }, (_, i) => `type_${i}`)
-
-	// An empty search says nothing about whether a type exists (stale index, or no embeddings):
-	// read as "none", it sends the model to define a type the instance already has.
-	it('lists the instance types when the similarity search finds nothing', async () => {
-		vi.spyOn(ResourceService, 'queryResourceTypes').mockRejectedValue(new Error('500'))
-		vi.spyOn(ResourceService, 'listResourceTypeNames').mockResolvedValue([
-			...syncedNames,
-			'postgresql'
-		])
-
-		const { resourceTypes, note } = await searchResourceTypes('postgres database', 'ws', 5)
-
-		expect(resourceTypes).toEqual([])
-		expect(note).toContain('postgresql')
-		expect(note).not.toContain('never synced')
-	})
-
-	it('returns the schema of a type the query names', async () => {
+	it.each([
+		[['ai_skill', 'ai_instruction'], true],
+		[Array.from({ length: 30 }, (_, i) => `type_${i}`), false]
+	])('notes only an instance that never synced with the hub (%#)', async (names, noted) => {
 		vi.spyOn(ResourceService, 'queryResourceTypes').mockResolvedValue([])
-		vi.spyOn(ResourceService, 'listResourceTypeNames').mockResolvedValue([...syncedNames, 'stripe'])
-		vi.spyOn(ResourceService, 'getResourceType').mockResolvedValue({
-			name: 'stripe',
-			schema: { type: 'object' }
-		})
-
-		const { resourceTypes, note } = await searchResourceTypes('stripe', 'ws', 5)
-
-		expect(resourceTypes.map((rt) => rt.name)).toEqual(['stripe'])
-		expect(note).toBeUndefined()
-	})
-
-	it('notes an instance that never synced with the hub', async () => {
-		vi.spyOn(ResourceService, 'queryResourceTypes').mockResolvedValue([])
-		vi.spyOn(ResourceService, 'listResourceTypeNames').mockResolvedValue([
-			'ai_skill',
-			'ai_instruction'
-		])
+		vi.spyOn(ResourceService, 'listResourceTypeNames').mockResolvedValue(names)
 
 		const { note } = await searchResourceTypes('stripe', 'ws', 5)
 
-		expect(note).toContain('never synced')
+		expect(note?.includes('never synced') ?? false).toBe(noted)
 	})
 })
 

frontend/src/lib/components/copilot/chat/script/core.ts

@@ -72,44 +72,24 @@ export function formatResourceTypes(
 // so fewer than this means the instance never synced with the hub.
 const UNSYNCED_RESOURCE_TYPE_COUNT = 20
 
-/**
- * When the similarity search finds nothing, the instance's type names come back instead: the
- * index is rebuilt from the table once a day and builds without embeddings have none, so an
- * empty search does not mean no type fits. A query that is a type's name gets its schema.
- */
+/** When nothing matches, a note says whether the instance seems never to have synced with the hub. */
 export async function searchResourceTypes(
 	query: string,
 	workspace: string,
 	limit: number
 ): Promise<{ resourceTypes: Pick<ResourceType, 'name' | 'schema'>[]; note?: string }> {
-	const similar = await ResourceService.queryResourceTypes({ workspace, text: query, limit }).catch(
-		() => []
-	)
-	if (similar.length > 0) {
-		return { resourceTypes: similar }
+	const resourceTypes = await ResourceService.queryResourceTypes({ workspace, text: query, limit })
+	if (resourceTypes.length > 0) {
+		return { resourceTypes }
 	}
-
-	const names = [
-		...new Set(await ResourceService.listResourceTypeNames({ workspace }).catch(() => []))
-	]
-	const name = query.trim().toLowerCase()
-	const named = names.includes(name)
-		? await ResourceService.getResourceType({ workspace, path: name }).catch(() => undefined)
-		: undefined
-	if (named) {
-		return { resourceTypes: [{ name: named.name, schema: named.schema }] }
-	}
-	if (names.length === 0) {
-		return { resourceTypes: [] }
-	}
-
-	const unsynced =
-		names.length < UNSYNCED_RESOURCE_TYPE_COUNT
-			? `\n\nOnly ${names.length} resource types exist on this instance, so it has most likely never synced with the Windmill Hub, which has types for most services. ${HUB_SYNC_INSTRUCTIONS}`
-			: ''
+	const count = new Set(await ResourceService.listResourceTypeNames({ workspace }).catch(() => []))
+		.size
 	return {
-		resourceTypes: [],
-		note: `The resource types on this instance are: ${names.join(', ')}. Search again with one of these names to get its schema.${unsynced}`
+		resourceTypes,
+		note:
+			count > 0 && count < UNSYNCED_RESOURCE_TYPE_COUNT
+				? `Only ${count} resource types exist on this instance, so it has most likely never synced with the Windmill Hub, which has types for most services. ${HUB_SYNC_INSTRUCTIONS}`
+				: undefined
 	}
 }