feat(ai): add vector DB assistant tools
This commit is contained in:
lexmin0412
GitHub
parent
3c7c1924ba
commit
bc409cffc0
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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 = [`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<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,
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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<ToolDefinition> {
|
||||
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<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());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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/);
|
||||
});
|
||||
|
|
|
|||
Loading…
Reference in New Issue