diff --git a/apps/desktop/src/components/editor/AiAssistant.vue b/apps/desktop/src/components/editor/AiAssistant.vue index fe5e25ca0..8d0754f31 100644 --- a/apps/desktop/src/components/editor/AiAssistant.vue +++ b/apps/desktop/src/components/editor/AiAssistant.vue @@ -17,7 +17,7 @@ import { connectionIconType } from "@/lib/connectionPresentation"; import DatabaseIcon from "@/components/icons/DatabaseIcon.vue"; import { useQueryStore } from "@/stores/queryStore"; import { useToast } from "@/composables/useToast"; -import { buildAiContext, runAgentStream, type AiAction } from "@/lib/ai"; +import { buildAiContext, runAgentStream, isVectorDbType, type AiAction } from "@/lib/ai"; import { formatAiModelOption } from "@/lib/aiModelPresentation"; import type { AgentEvent } from "@/lib/tauri"; import { buildAiAgentPlan } from "@/lib/aiAgentPlan"; @@ -287,6 +287,11 @@ const actionMenuItems = computed(() => ); const aiCodeAppearance = computed(() => (isDark.value ? "dark" : "light")); +const showActionButtons = computed(() => { + if (!props.connection) return true; + return !isVectorDbType(props.connection.db_type); +}); + const { databaseOptions: allDbOptions, loadDatabaseOptions } = useDatabaseOptions(); const dbOptions = computed(() => { @@ -1312,6 +1317,7 @@ async function openExternalUrl(url: string) { item-class="text-xs px-2" /> = new Set([ + "qdrant", + "milvus", + "weaviate", + "chromadb", + // If modifying this, also update is_vector_db() in crates/dbx-core/src/agent_tools.rs. +]); + +export function isVectorDbType(dbType: DatabaseType): boolean { + return VECTOR_DB_TYPES.has(dbType); +} + +function dbLabel(dbType: DatabaseType): string { + const labels: Partial> = { + qdrant: "Qdrant", + milvus: "Milvus", + weaviate: "Weaviate", + chromadb: "ChromaDB", + }; + return labels[dbType] || dbType; +} + export type AiAction = "generate" | "explain" | "optimize" | "fix" | "convert" | "sampleData"; export type AiAssistantMode = "ask" | "agent"; @@ -50,10 +72,8 @@ export interface AiRequestInput { function buildAgentRequest(input: AiRequestInput, history?: api.AiMessage[]): { messages: api.AiMessage[]; systemPrompt: string; taskContract: api.AiTaskContract; maxTokens: number; temperature: number } { const isZh = isChineseLocale(currentLocale()); - const skill = aiSkillForAction(input.action); const systemPrompt = buildSystemPrompt(input.action, input.context, input.mode); - const instruction = isZh ? skill.userInstruction.zh : skill.userInstruction.en; - const userPrompt = [`Action: ${input.action}`, instruction, "", "User request:", input.instruction.trim() || "(No extra instruction provided.)"].join("\n"); + const userPrompt = buildUserPrompt(input.action, input.context, input.instruction, isZh); const taskContract: api.AiTaskContract = { action: input.action, mode: input.mode || "ask", @@ -124,6 +144,17 @@ export async function runAgentStream(input: AiRequestInput, history: api.AiMessa ); } +export function buildUserPrompt(action: AiAction, context: AiContext, instruction: string, isZh: boolean): string { + const userRequest = instruction.trim() || (isZh ? "（无额外说明）" : "(No extra instruction provided.)"); + if (isVectorDbType(context.databaseType)) { + // Vector databases: skip SQL action instructions, only send the user's request + return userRequest; + } + const skill = aiSkillForAction(action); + const skillInstruction = isZh ? skill.userInstruction.zh : skill.userInstruction.en; + return [`Action: ${action}`, skillInstruction, "", "User request:", userRequest].join("\n"); +} + function actionParams(action: AiAction): { maxTokens: number; temperature: number } { switch (action) { case "explain": @@ -142,6 +173,9 @@ export function extractSql(text: string): string { } export function buildSystemPrompt(action: AiAction, context: AiContext, mode: AiAssistantMode = "ask"): string { + if (isVectorDbType(context.databaseType)) { + return buildVectorSystemPrompt(context, mode); + } const schema = formatSchema(context); const resultPreview = context.lastResultPreview ? `\nLast result preview:\n${context.lastResultPreview}\n` : ""; const lastError = context.lastError ? `\nLast error:\n${context.lastError}\n` : ""; @@ -207,6 +241,50 @@ function buildBasePromptLines(isZh: boolean): string[] { ]; } +function buildVectorSystemPrompt(context: AiContext, mode: AiAssistantMode): string { + const isZh = isChineseLocale(currentLocale()); + const schema = formatSchema(context); + const resultPreview = context.lastResultPreview ? `\nLast result preview:\n${context.lastResultPreview}\n` : ""; + const lastError = context.lastError ? `\nLast error:\n${context.lastError}\n` : ""; + const lines: string[] = [ + isZh ? `你是 DBX 内置的向量数据库助手。当前连接的是 ${dbLabel(context.databaseType)} 数据库。用中文回复。` : `You are DBX's vector database assistant. Connected to ${dbLabel(context.databaseType)}. Reply in English.`, + isZh ? "数据存储在集合（collections）中，每条记录包含唯一标识及可选的元数据负载（payload/metadata）。" : "Data is stored in collections. Each record has a unique identifier and optional metadata payload.", + ...buildVectorModePromptLines(context, mode, isZh), + "", + `Database type: ${context.databaseType}`, + `Connection: ${context.connectionName}`, + `Database: ${context.database}`, + schemaCoverageLine(context, isZh), + "", + `Current collection:\n${context.currentSql.trim() || "(none)"}`, + lastError, + resultPreview, + "", + `Schema:\n${schema}`, + ]; + + if (context.schemaScope === "focused_table") { + lines.push( + isZh + ? "Schema 上下文只覆盖当前打开的集合；数据库中可能还有其他集合。用户询问当前有哪些集合或提到上下文中不存在的集合时，不要直接断言不存在，先用 list_collections 工具确认。" + : "Schema context covers only the currently opened collection; the database may contain other collections. When the user asks what collections exist or mentions a collection absent from context, do not conclude it is missing; use list_collections to verify first.", + ); + } + + return lines.filter(Boolean).join("\n"); +} + +function buildVectorModePromptLines(context: AiContext, mode: AiAssistantMode, isZh: boolean): string[] { + if (mode === "agent") { + return [isZh ? "你处于 Agent 模式。你有以下工具可用：list_collections、browse_collection。" : "You are in Agent mode. You have the following tools available: list_collections, browse_collection."]; + } + return [ + isZh + ? `你处于 Ask 模式。你只能使用 list_collections 确认集合清单；不要浏览集合数据。${dbLabel(context.databaseType)} 的查询格式为 REST API（METHOD /path + JSON body），具体格式因数据库类型而异。只生成查询请求文本和说明，不要暗示已经执行。` + : `You are in Ask mode. You may only use list_collections to inspect collection names; do not browse collection data. ${dbLabel(context.databaseType)} uses a REST API query format (METHOD /path + JSON body) that varies by database type. Generate query strings and explanations only; do not imply execution.`, + ]; +} + function buildModePromptLines(mode: AiAssistantMode, isZh: boolean): string[] { if (mode === "agent") { return [ @@ -226,7 +304,10 @@ function buildModePromptLines(mode: AiAssistantMode, isZh: boolean): string[] { function schemaCoverageLine(context: AiContext, isZh: boolean): string { if (context.schemaScope === "focused_table") { - return isZh ? "Schema context scope: focused table only; not a complete database table list." : "Schema context scope: focused table only; not a complete database table list."; + if (isVectorDbType(context.databaseType)) { + return isZh ? "Schema 上下文只覆盖当前打开的集合，不是完整的集合列表。" : "Schema context scope: focused collection only; not a complete collection list."; + } + return "Schema context scope: focused table only; not a complete database table list."; } return context.truncated ? "Schema context is truncated." : "Schema context is complete for the loaded database scope."; } @@ -279,6 +360,7 @@ export async function buildAiContext(tab: QueryTab, connection: ConnectionConfig const tableKeys = new Set(); let truncated = false; let schemaScope: AiContext["schemaScope"] = "database"; + let currentCollectionName: string | undefined; if (tab.tableMeta) { schemaScope = "focused_table"; @@ -308,7 +390,43 @@ export async function buildAiContext(tab: QueryTab, connection: ConnectionConfig tables.push(entry); } - if (!tab.tableMeta && !["redis", "mongodb"].includes(connection.db_type)) { + // Vector databases: load collections instead of SQL tables + if (isVectorDbType(databaseType)) { + try { + const collections = await api.vectorListCollections(tab.connectionId); + + // Find the currently opened collection (tab.sql is UUID for ChromaDB, name for others) + const currentCollection = collections.find((c) => c.id === tab.sql || c.name === tab.sql); + if (currentCollection) { + schemaScope = "focused_table"; + currentCollectionName = currentCollection.name; + tables.push({ + name: currentCollection.name, + tableType: "COLLECTION", + comment: currentCollection.dimension ? `${currentCollection.dimension}d vector` : undefined, + columns: [], + }); + tableKeys.add(aiTableMentionKey(undefined, currentCollection.name)); + } + + for (const col of collections.slice(0, maxTables)) { + const key = aiTableMentionKey(undefined, col.name); + if (tableKeys.has(key)) continue; + tables.push({ + name: col.name, + tableType: "COLLECTION", + comment: col.dimension ? `${col.dimension}d vector` : undefined, + columns: [], + }); + tableKeys.add(key); + } + if (collections.length > maxTables) truncated = true; + } catch { + truncated = true; + } + } + + if (!tab.tableMeta && !["redis", "mongodb"].includes(connection.db_type) && !isVectorDbType(databaseType)) { try { const schemas = await loadCandidateSchemas(tab, connection); for (const schema of schemas) { @@ -355,7 +473,7 @@ export async function buildAiContext(tab: QueryTab, connection: ConnectionConfig connectionName: connection.name, databaseType, database: tab.database, - currentSql: tab.sql, + currentSql: currentCollectionName ?? tab.sql, lastError: extractLastError(tab.result), lastResultPreview: formatResultPreview(tab.result), tables, diff --git a/crates/dbx-core/src/agent_loop.rs b/crates/dbx-core/src/agent_loop.rs index 02d8007e8..6c4fbf5ef 100644 --- a/crates/dbx-core/src/agent_loop.rs +++ b/crates/dbx-core/src/agent_loop.rs @@ -129,7 +129,11 @@ pub async fn run_agent_loop( ) .await; } - let tools = if is_agent_mode { agent_tools::all_tools(agent_ctx.db_type) } else { agent_tools::read_only_tools() }; + let tools = if is_agent_mode { + agent_tools::all_tools(agent_ctx.db_type) + } else { + agent_tools::read_only_tools(agent_ctx.db_type) + }; let task_contract = task_contract.cloned(); let mut conversation_messages: Vec = messages.to_vec(); let mut final_text = String::new(); diff --git a/crates/dbx-core/src/agent_tools.rs b/crates/dbx-core/src/agent_tools.rs index 034f3a23c..1d51a2dda 100644 --- a/crates/dbx-core/src/agent_tools.rs +++ b/crates/dbx-core/src/agent_tools.rs @@ -4,6 +4,7 @@ use serde_json::json; use crate::agent_events::{ToolCall, ToolDefinition, ToolResult}; use crate::connection::AppState; +use crate::db::vector_driver; use crate::models::connection::DatabaseType; use crate::query::QueryExecutionOptions; use crate::query_execution_sql::{build_explain_sql, supports_explain_plan, supports_sql_query, ExplainSqlOptions}; @@ -19,18 +20,35 @@ const EXECUTE_QUERY_LIMIT: usize = 50; /// Maximum number of rows returned by get_sample_data tool. const SAMPLE_DATA_LIMIT: usize = 20; +/// Maximum number of rows returned by browse_collection tool. +const BROWSE_COLLECTION_LIMIT: usize = 20; + /// Absolute maximum rows any query tool may request. const MAX_ALLOWED_ROWS: usize = 100; -/// Get read-only tool definitions (list_tables + get_columns). -pub fn read_only_tools() -> Vec { - vec![list_tables_tool(), get_columns_tool()] +/// Returns true for vector database types (Qdrant, Milvus, Weaviate, ChromaDb). +/// If modifying this, also update VECTOR_DB_TYPES in apps/desktop/src/lib/ai.ts. +pub fn is_vector_db(db_type: DatabaseType) -> bool { + matches!(db_type, DatabaseType::Qdrant | DatabaseType::Milvus | DatabaseType::Weaviate | DatabaseType::ChromaDb) +} + +/// Get read-only tool definitions for the given database type. +/// Returns vector tools for vector DBs, SQL tools otherwise. +pub fn read_only_tools(db_type: DatabaseType) -> Vec { + if is_vector_db(db_type) { + vec![list_collections_tool()] + } else { + vec![list_tables_tool(), get_columns_tool()] + } } /// Get all available tool definitions for the given database type. /// Includes read-only tools plus execute_query, get_sample_data, and /// explain_query for database types that support them. pub fn all_tools(db_type: DatabaseType) -> Vec { + if is_vector_db(db_type) { + return vec![list_collections_tool(), browse_collection_tool()]; + } let mut tools = vec![list_tables_tool(), get_columns_tool()]; if supports_sql_query(db_type) { tools.push(execute_query_tool()); @@ -167,6 +185,45 @@ fn explain_query_tool() -> ToolDefinition { } } +/// list_collections tool definition (vector databases). +fn list_collections_tool() -> ToolDefinition { + ToolDefinition { + name: "list_collections", + description: "List all collections in the current vector database. Returns collection names and dimensions.", + parameters: json!({ + "type": "object", + "properties": {}, + "required": [] + }), + read_only: true, + parallel_ok: true, + } +} + +/// browse_collection tool definition (vector databases). +fn browse_collection_tool() -> ToolDefinition { + ToolDefinition { + name: "browse_collection", + description: "Browse documents in a collection. Returns up to 20 items with payload/metadata (vectors excluded for compactness). For ChromaDB, use the collection id (UUID from list_collections) instead of the collection name.", + parameters: json!({ + "type": "object", + "properties": { + "collection": { + "type": "string", + "description": "Collection name" + }, + "limit": { + "type": "number", + "description": "Max items to return (default 20, max 100)" + } + }, + "required": ["collection"] + }), + read_only: true, + parallel_ok: true, + } +} + /// Execute a tool call and return the result. pub async fn execute_tool( tool_call: &ToolCall, @@ -180,6 +237,8 @@ pub async fn execute_tool( "get_columns" => execute_get_columns(tool_call, state, connection_id, database, db_type).await, "execute_query" => execute_execute_query(tool_call, state, connection_id, database, db_type).await, "get_sample_data" => execute_get_sample_data(tool_call, state, connection_id, database, db_type).await, + "list_collections" => execute_list_collections(tool_call, state, connection_id, database, db_type).await, + "browse_collection" => execute_browse_collection(tool_call, state, connection_id, database, db_type).await, "explain_query" => { let (text_result, explain_data) = execute_explain_query(tool_call, state, connection_id, database, db_type).await; @@ -552,3 +611,233 @@ async fn execute_explain_query( (Ok(text), explain_data) } + +/// Execute list_collections tool (vector databases). +async fn execute_list_collections( + _tool_call: &ToolCall, + state: &Arc, + connection_id: &str, + database: &str, + _db_type: &DatabaseType, +) -> Result { + let collections = crate::schema::list_vector_collections_core(state, connection_id, database) + .await + .map_err(|e| format!("Failed to list collections: {e}"))?; + + if collections.is_empty() { + return Ok("No collections found.".to_string()); + } + + let mut lines: Vec = collections + .iter() + .map(|c| { + let mut line = format!("- {} (COLLECTION)", c.name); + if let Some(dim) = c.dimension { + line.push_str(&format!(" -- {}d", dim)); + } + line.push_str(&format!(" [id: {}]", c.id)); + line + }) + .collect(); + + if lines.len() > LIST_TABLES_LIMIT { + lines.truncate(LIST_TABLES_LIMIT); + lines.push(format!("... (showing {LIST_TABLES_LIMIT} of {} collections)", collections.len())); + } + + Ok(lines.join("\n")) +} + +/// Execute browse_collection tool (vector databases). +/// Generates a database-specific REST query and executes it. +async fn execute_browse_collection( + tool_call: &ToolCall, + state: &Arc, + connection_id: &str, + database: &str, + db_type: &DatabaseType, +) -> Result { + let collection = tool_call + .arguments + .get("collection") + .and_then(|v| v.as_str()) + .ok_or("Missing required parameter: collection")? + .trim(); + + if collection.is_empty() { + return Err("Collection name cannot be empty".to_string()); + } + + let limit = tool_call + .arguments + .get("limit") + .and_then(|v| v.as_u64()) + .map(|l| (l as usize).min(MAX_ALLOWED_ROWS)) + .unwrap_or(BROWSE_COLLECTION_LIMIT); + + // ChromaDB requires UUID in URL path, not collection name. + // If the collection param is already a UUID (from list_collections output), use it directly. + let collection_id = if *db_type == DatabaseType::ChromaDb && !is_uuid(collection) { + resolve_chroma_collection_uuid(state, connection_id, database, collection).await? + } else { + collection.to_string() + }; + + let query = build_browse_query(db_type, &collection_id, database, limit)?; + + let options = QueryExecutionOptions { max_rows: Some(limit), timeout_secs: Some(30), ..Default::default() }; + let result = + crate::query::execute_sql_statement_with_options(state, connection_id, database, &query, None, None, options) + .await?; + + format_query_result_as_text(&result, limit) +} + +/// Build a browse query for the given vector database type. +/// Intentionally omits offset/pagination — Agent browse only fetches the first N items. +fn build_browse_query( + db_type: &DatabaseType, + collection: &str, + database: &str, + limit: usize, +) -> Result { + let collection = collection.trim(); + if collection.is_empty() { + return Err("Collection name cannot be empty".to_string()); + } + let limit = limit.max(1) as u64; + + match db_type { + DatabaseType::Qdrant => Ok(format!( + "POST /collections/{}/points/scroll\n{}", + vector_driver::path_segment(collection), + serde_json::json!({ "limit": limit, "with_payload": true, "with_vector": false }) + )), + // Milvus v2 omitting outputFields defaults to returning only scalar fields (no vectors). + DatabaseType::Milvus => Ok(format!( + "POST /v2/vectordb/entities/query\n{}", + serde_json::json!({ + "dbName": if database.is_empty() { "default" } else { database }, + "collectionName": collection, + "filter": "", "limit": limit + }) + )), + DatabaseType::Weaviate => { + Ok(format!("GET /v1/objects?class={}&limit={}", vector_driver::query_value(collection), limit)) + } + // TODO: ChromaDB Cloud 支持自定义租户和数据库，当前只实现了本地部署 + // （固定 default_tenant / default_database），后续支持云服务时需改为可配置。 + DatabaseType::ChromaDb => Ok(format!( + "POST /api/v2/tenants/default_tenant/databases/default_database/collections/{}/get\n{}", + collection, + serde_json::json!({ "limit": limit, "include": ["documents", "metadatas"] }) + )), + _ => Err(format!("Unsupported database type: {:?}", db_type)), + } +} + +/// Check if a string looks like a UUID (simple check — 36 chars with 4 hyphens). +fn is_uuid(s: &str) -> bool { + s.len() == 36 + && s.chars().filter(|&c| c == '-').count() == 4 + && s.chars().all(|c| c.is_ascii_hexdigit() || c == '-') +} + +/// Resolve a ChromaDB collection name to its UUID by listing all collections. +async fn resolve_chroma_collection_uuid( + state: &Arc, + connection_id: &str, + database: &str, + name: &str, +) -> Result { + let collections = crate::schema::list_vector_collections_core(state, connection_id, database).await?; + collections + .into_iter() + .find(|c| c.name == name) + .map(|c| c.id) + .ok_or_else(|| format!("Collection '{name}' not found")) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn vector_read_only_tools_do_not_include_collection_browsing() { + let tools = read_only_tools(DatabaseType::Qdrant); + let names: Vec<&str> = tools.iter().map(|tool| tool.name).collect(); + + assert_eq!(names, vec!["list_collections"]); + } + + #[test] + fn vector_agent_tools_include_collection_browsing() { + let tools = all_tools(DatabaseType::Qdrant); + let names: Vec<&str> = tools.iter().map(|tool| tool.name).collect(); + + assert_eq!(names, vec!["list_collections", "browse_collection"]); + } + + #[test] + fn build_browse_query_qdrant() { + let q = build_browse_query(&DatabaseType::Qdrant, "articles", "", 10).unwrap(); + assert!(q.starts_with("POST /collections/articles/points/scroll")); + assert!(q.contains("\"limit\":10")); + assert!(q.contains("\"with_payload\":true")); + } + + #[test] + fn build_browse_query_qdrant_encodes_url_chars() { + let q = build_browse_query(&DatabaseType::Qdrant, "my collection", "", 10).unwrap(); + assert!(q.starts_with("POST /collections/my%20collection/points/scroll")); + } + + #[test] + fn build_browse_query_milvus() { + let q = build_browse_query(&DatabaseType::Milvus, "articles", "custom_db", 20).unwrap(); + assert!(q.starts_with("POST /v2/vectordb/entities/query")); + assert!(q.contains("\"dbName\":\"custom_db\"")); + assert!(q.contains("\"collectionName\":\"articles\"")); + assert!(q.contains("\"limit\":20")); + assert!(!q.contains("outputFields")); + } + + #[test] + fn build_browse_query_milvus_default_db() { + let q = build_browse_query(&DatabaseType::Milvus, "articles", "", 10).unwrap(); + assert!(q.contains("\"dbName\":\"default\"")); + } + + #[test] + fn build_browse_query_weaviate() { + let q = build_browse_query(&DatabaseType::Weaviate, "Articles", "", 5).unwrap(); + assert_eq!(q, "GET /v1/objects?class=Articles&limit=5"); + } + + #[test] + fn build_browse_query_weaviate_encodes_query_param() { + let q = build_browse_query(&DatabaseType::Weaviate, "A&B", "", 5).unwrap(); + assert!(q.contains("class=A%26B")); + } + + #[test] + fn build_browse_query_chromadb() { + let q = build_browse_query(&DatabaseType::ChromaDb, "uuid-123", "", 15).unwrap(); + assert!( + q.starts_with("POST /api/v2/tenants/default_tenant/databases/default_database/collections/uuid-123/get") + ); + assert!(q.contains("\"limit\":15")); + } + + #[test] + fn build_browse_query_rejects_empty_collection() { + let result = build_browse_query(&DatabaseType::Qdrant, " ", "", 10); + assert!(result.is_err()); + } + + #[test] + fn build_browse_query_rejects_unsupported_type() { + let result = build_browse_query(&DatabaseType::Postgres, "articles", "", 10); + assert!(result.is_err()); + } +} diff --git a/crates/dbx-core/src/db/vector_driver.rs b/crates/dbx-core/src/db/vector_driver.rs index 4dbd8e151..58946d2d4 100644 --- a/crates/dbx-core/src/db/vector_driver.rs +++ b/crates/dbx-core/src/db/vector_driver.rs @@ -152,9 +152,17 @@ pub async fn test_connection(client: &VectorClient, timeout: Duration) -> Result } pub async fn list_collections(client: &VectorClient) -> Result, String> { + list_collections_with_db(client, "").await +} + +/// List collections, passing an optional database name (used by Milvus). +pub(crate) async fn list_collections_with_db( + client: &VectorClient, + database: &str, +) -> Result, String> { match client.kind { VectorDbKind::Qdrant => list_qdrant_collections(client).await, - VectorDbKind::Milvus => list_milvus_collections(client).await, + VectorDbKind::Milvus => list_milvus_collections(client, database).await, VectorDbKind::Weaviate => list_weaviate_collections(client).await, VectorDbKind::ChromaDb => list_chroma_collections(client).await, } @@ -176,9 +184,10 @@ async fn list_qdrant_collections(client: &VectorClient) -> Result Result, String> { +async fn list_milvus_collections(client: &VectorClient, database: &str) -> Result, String> { + let db_name = if database.is_empty() { "default" } else { database }; let body = send_json( - client.post("/v2/vectordb/collections/list").json(&serde_json::json!({ "dbName": "default" })), + client.post("/v2/vectordb/collections/list").json(&serde_json::json!({ "dbName": db_name })), "Milvus", ) .await?; @@ -424,11 +433,11 @@ fn starts_with_http_method(input: &str) -> bool { ["GET ", "POST ", "PUT ", "DELETE "].iter().any(|prefix| input.to_ascii_uppercase().starts_with(prefix)) } -fn path_segment(value: &str) -> String { +pub(crate) fn path_segment(value: &str) -> String { utf8_percent_encode(value, PATH_SEGMENT_ENCODE_SET).to_string() } -fn query_value(value: &str) -> String { +pub(crate) fn query_value(value: &str) -> String { utf8_percent_encode(value, QUERY_VALUE_ENCODE_SET).to_string() } diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index ba898834b..a557a23cf 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -816,6 +816,25 @@ pub async fn list_tables_core( .await } +/// List vector database collections, returning structured info (name, id, dimension). +/// Only works for PoolKind::VectorDb connections; returns an error for other types. +pub async fn list_vector_collections_core( + state: &AppState, + connection_id: &str, + database: &str, +) -> Result, String> { + let pool_key = + state.get_or_create_pool(connection_id, if database.is_empty() { None } else { Some(database) }).await?; + let client = { + let connections = state.connections.read().await; + match connections.get(&pool_key) { + Some(PoolKind::VectorDb(client)) => client.clone(), + _ => return Err("Not a vector database connection".to_string()), + } + }; + db::vector_driver::list_collections_with_db(&client, database).await +} + pub async fn get_table_comment_core( state: &AppState, connection_id: &str, diff --git a/packages/app-tests/aiPrompt.test.ts b/packages/app-tests/aiPrompt.test.ts index 1d5e2324b..d9c052f26 100644 --- a/packages/app-tests/aiPrompt.test.ts +++ b/packages/app-tests/aiPrompt.test.ts @@ -30,7 +30,7 @@ Object.defineProperty(globalThis, "localStorage", { configurable: true, }); -const { buildSystemPrompt } = await import("../../apps/desktop/src/lib/ai.ts"); +const { buildSystemPrompt, isVectorDbType, buildUserPrompt } = await import("../../apps/desktop/src/lib/ai.ts"); function context(overrides: Partial = {}): AiContext { return { @@ -98,3 +98,100 @@ test("prompt enforces database dialect and single executable statement safety", assert.match(prompt, /不要生成多语句 SQL/); assert.match(prompt, /不要在同一个回答里混合 SELECT 和写操作/); }); + +// Vector database tests + +function vectorContext(overrides: Partial = {}): AiContext { + return { + connectionName: "my-qdrant", + databaseType: "qdrant", + database: "default", + currentSql: "articles", + tables: [ + { + name: "articles", + tableType: "COLLECTION", + comment: "384d vector", + columns: [], + }, + ], + truncated: false, + ...overrides, + }; +} + +test("isVectorDbType returns true for vector databases", () => { + assert.equal(isVectorDbType("qdrant"), true); + assert.equal(isVectorDbType("milvus"), true); + assert.equal(isVectorDbType("weaviate"), true); + assert.equal(isVectorDbType("chromadb"), true); +}); + +test("isVectorDbType returns false for SQL databases", () => { + assert.equal(isVectorDbType("mysql"), false); + assert.equal(isVectorDbType("postgres"), false); + assert.equal(isVectorDbType("sqlserver"), false); +}); + +test("vector system prompt does not contain SQL references", () => { + const prompt = buildSystemPrompt("generate", vectorContext(), "ask"); + + assert.doesNotMatch(prompt, /```sql/); + assert.doesNotMatch(prompt, /execute_query/); + assert.match(prompt, /collection/); + assert.match(prompt, /REST API/); +}); + +test("vector agent mode lists vector tools", () => { + const prompt = buildSystemPrompt("generate", vectorContext(), "agent"); + + assert.match(prompt, /list_collections/); + assert.match(prompt, /browse_collection/); + assert.doesNotMatch(prompt, /execute_query/); + assert.doesNotMatch(prompt, /list_tables/); +}); + +test("vector focused table prompt warns about unknown collections", () => { + const prompt = buildSystemPrompt("generate", vectorContext({ schemaScope: "focused_table" }), "agent"); + + assert.match(prompt, /不是完整的集合列表/); + assert.match(prompt, /当前打开的集合/); + assert.match(prompt, /list_collections/); +}); + +test("vector ask mode mentions REST API format", () => { + const prompt = buildSystemPrompt("generate", vectorContext(), "ask"); + + assert.match(prompt, /REST API/); + assert.match(prompt, /Qdrant/); + assert.match(prompt, /list_collections/); + assert.match(prompt, /do not browse collection data|不要浏览集合数据/); + assert.doesNotMatch(prompt, /```sql/); + assert.doesNotMatch(prompt, /execute_query/); +}); + +test("vector system prompt preserves last error and result preview", () => { + const prompt = buildSystemPrompt( + "generate", + vectorContext({ + lastError: "Qdrant error", + lastResultPreview: "id | payload\n1 | {}", + }), + "ask", + ); + + assert.match(prompt, /Last error:\nQdrant error/); + assert.match(prompt, /Last result preview:\nid \| payload/); +}); + +test("buildUserPrompt skips action instruction for vector databases", () => { + const vectorCtx = vectorContext(); + const sqlCtx = context(); + + const vectorPrompt = buildUserPrompt("generate", vectorCtx, "show me articles", true); + assert.equal(vectorPrompt, "show me articles"); + + const sqlPrompt = buildUserPrompt("generate", sqlCtx, "show me users", true); + assert.match(sqlPrompt, /Action: generate/); + assert.match(sqlPrompt, /生成 SQL/); +});