feat(ai): add vector DB assistant tools

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"
             />
             <LightDropdown
+              v-if="showActionButtons"
               :model-value="activeAction"
               :items="actionMenuItems"
               content-class="w-max min-w-0"

apps/desktop/src/lib/ai.ts

@@ -10,6 +10,28 @@ import { effectiveDatabaseTypeForConnection } from "@/lib/jdbcDialect";
 
 import type { AgentEvent } from "@/lib/tauri";
 
+const VECTOR_DB_TYPES: ReadonlySet<DatabaseType> = 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<Record<DatabaseType, string>> = {
+    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 = buildUserPrompt(input.action, input.context, input.instruction, isZh);
-  const userPrompt = [`Action: ${input.action}`, instruction, "", "User request:", input.instruction.trim() || "(No extra instruction provided.)"].join("\n");
   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<string>();
   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,

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<AiMessage> = messages.to_vec();
     let mut final_text = String::new();

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).
+/// Returns true for vector database types (Qdrant, Milvus, Weaviate, ChromaDb).
-pub fn read_only_tools() -> Vec<ToolDefinition> {
+/// If modifying this, also update VECTOR_DB_TYPES in apps/desktop/src/lib/ai.ts.
-    vec![list_tables_tool(), get_columns_tool()]
+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<ToolDefinition> {
+    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<ToolDefinition> {
+    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<AppState>,
+    connection_id: &str,
+    database: &str,
+    _db_type: &DatabaseType,
+) -> Result<String, String> {
+    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<String> = 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<AppState>,
+    connection_id: &str,
+    database: &str,
+    db_type: &DatabaseType,
+) -> Result<String, String> {
+    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<String, String> {
+    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<AppState>,
+    connection_id: &str,
+    database: &str,
+    name: &str,
+) -> Result<String, String> {
+    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());
+    }
+}

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<Vec<CollectionInfo>, 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<Vec<CollectionInfo>, 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<Vec<Collection
     Ok(infos)
 }
 
-async fn list_milvus_collections(client: &VectorClient) -> Result<Vec<CollectionInfo>, String> {
+async fn list_milvus_collections(client: &VectorClient, database: &str) -> Result<Vec<CollectionInfo>, 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()
 }
 

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<Vec<db::vector_driver::CollectionInfo>, 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,

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> = {}): 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> = {}): 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/);
+});