diff --git a/ai_evals/adapters/frontend/mockBackend.ts b/ai_evals/adapters/frontend/mockBackend.ts index f26f73e1f6..810a137ccf 100644 --- a/ai_evals/adapters/frontend/mockBackend.ts +++ b/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, diff --git a/ai_evals/adapters/frontend/vitestAdapter.test.ts b/ai_evals/adapters/frontend/vitestAdapter.test.ts index e86dd92f79..23a671bd40 100644 --- a/ai_evals/adapters/frontend/vitestAdapter.test.ts +++ b/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 }) => diff --git a/backend/windmill-api-embeddings/src/lib.rs b/backend/windmill-api-embeddings/src/lib.rs index 6e4d6a1c71..461933be49 100644 --- a/backend/windmill-api-embeddings/src/lib.rs +++ b/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::().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::().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, } -#[cfg(feature = "embedding")] #[derive(Serialize)] pub struct ResourceTypeResult { name: String, @@ -138,24 +127,80 @@ fn trim_to_top_score( .take_while(|r| (top_score - score(r)) / top_score <= max_relative_drop) .collect() } -#[cfg(feature = "embedding")] async fn query_resource_types( + Extension(db): Extension, Query(query): Query, Path(w_id): Path, ) -> JsonResult> { - 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, +) -> windmill_common::error::Result> { + let words: Vec = text + .to_lowercase() + .split(|c: char| !c.is_alphanumeric()) + .filter(|word| word.len() >= 3) + .map(str::to_string) + .collect(); + let rows: Vec<(String, Option, Option)> = 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 = 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, + /// 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>>, + /// The `resource_types_fingerprint` the resource types collection was built from. + resource_types_fingerprint: Option, } #[cfg(feature = "embedding")] @@ -295,7 +345,12 @@ impl EmbeddingsDb { pub async fn new(pg_db: &Pool, model_instance: Arc) -> Result { 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::>().await?; + self.hub_resource_type_embeddings = Arc::new( + hub_resource_types + .into_iter() + .map(|rt| (rt.name, rt.embedding)) + .collect(), + ); - let resource_types: Vec = - 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, + embeddings: Vec, + ) -> 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) -> Result> { + 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, + model_instance: &Arc, + hub_embeddings: &HashMap>, +) -> Result<(Option, Vec)> { + let fingerprint = resource_types_fingerprint(pg_db).await?; + let resource_types: Vec = + 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) -> 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 { Ok(v.broadcast_div(&v.sqr()?.sum_keepdim(1)?.sqrt()?)?) @@ -630,6 +748,19 @@ pub fn load_embeddings_db(db: &Pool) -> () { .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) -> 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); + } +} diff --git a/frontend/src/lib/components/copilot/chat/script/core.test.ts b/frontend/src/lib/components/copilot/chat/script/core.test.ts index 2e07925f70..6bfcbc1d21 100644 --- a/frontend/src/lib/components/copilot/chat/script/core.test.ts +++ b/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) }) }) diff --git a/frontend/src/lib/components/copilot/chat/script/core.ts b/frontend/src/lib/components/copilot/chat/script/core.ts index 7ec1d960a7..94272656a0 100644 --- a/frontend/src/lib/components/copilot/chat/script/core.ts +++ b/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[]; 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 } }