SwiftDBAI: natural language queries for any SQLite database
Drop-in SwiftUI chat view, headless ChatEngine, LLM-agnostic via AnyLanguageModel. Read-only by default with configurable allowlists. Robust SQL parser with 63 tests. Includes demo app with GitHub stars dataset.
This commit is contained in:
Krishna Kumar
254
Tests/SwiftDBAITests/PromptBuilderTests.swift
Normal file
254
Tests/SwiftDBAITests/PromptBuilderTests.swift
Normal file
@@ -0,0 +1,254 @@
|
||||
// PromptBuilderTests.swift
|
||||
// SwiftDBAI
|
||||
|
||||
import Testing
|
||||
@testable import SwiftDBAI
|
||||
|
||||
@Suite("PromptBuilder")
|
||||
struct PromptBuilderTests {
|
||||
|
||||
// MARK: - Helpers
|
||||
|
||||
/// Creates a sample schema for testing.
|
||||
private func makeSampleSchema() -> DatabaseSchema {
|
||||
let usersTable = TableSchema(
|
||||
name: "users",
|
||||
columns: [
|
||||
ColumnSchema(cid: 0, name: "id", type: "INTEGER", isNotNull: true, defaultValue: nil, isPrimaryKey: true),
|
||||
ColumnSchema(cid: 1, name: "name", type: "TEXT", isNotNull: true, defaultValue: nil, isPrimaryKey: false),
|
||||
ColumnSchema(cid: 2, name: "email", type: "TEXT", isNotNull: false, defaultValue: nil, isPrimaryKey: false),
|
||||
ColumnSchema(cid: 3, name: "created_at", type: "TEXT", isNotNull: false, defaultValue: "CURRENT_TIMESTAMP", isPrimaryKey: false),
|
||||
],
|
||||
primaryKey: ["id"],
|
||||
foreignKeys: [],
|
||||
indexes: [
|
||||
IndexSchema(name: "idx_users_email", isUnique: true, columns: ["email"])
|
||||
]
|
||||
)
|
||||
|
||||
let ordersTable = TableSchema(
|
||||
name: "orders",
|
||||
columns: [
|
||||
ColumnSchema(cid: 0, name: "id", type: "INTEGER", isNotNull: true, defaultValue: nil, isPrimaryKey: true),
|
||||
ColumnSchema(cid: 1, name: "user_id", type: "INTEGER", isNotNull: true, defaultValue: nil, isPrimaryKey: false),
|
||||
ColumnSchema(cid: 2, name: "total", type: "REAL", isNotNull: true, defaultValue: nil, isPrimaryKey: false),
|
||||
ColumnSchema(cid: 3, name: "status", type: "TEXT", isNotNull: true, defaultValue: "'pending'", isPrimaryKey: false),
|
||||
],
|
||||
primaryKey: ["id"],
|
||||
foreignKeys: [
|
||||
ForeignKeySchema(fromColumn: "user_id", toTable: "users", toColumn: "id", onUpdate: "NO ACTION", onDelete: "CASCADE")
|
||||
],
|
||||
indexes: []
|
||||
)
|
||||
|
||||
return DatabaseSchema(
|
||||
tables: ["users": usersTable, "orders": ordersTable],
|
||||
tableNames: ["users", "orders"]
|
||||
)
|
||||
}
|
||||
|
||||
private func makeEmptySchema() -> DatabaseSchema {
|
||||
DatabaseSchema(tables: [:], tableNames: [])
|
||||
}
|
||||
|
||||
// MARK: - System Instructions Tests
|
||||
|
||||
@Test("System instructions contain role section")
|
||||
func systemInstructionsContainRole() {
|
||||
let builder = PromptBuilder(schema: makeSampleSchema())
|
||||
let instructions = builder.buildSystemInstructions()
|
||||
|
||||
#expect(instructions.contains("ROLE"))
|
||||
#expect(instructions.contains("SQL assistant"))
|
||||
#expect(instructions.contains("SQLite database"))
|
||||
}
|
||||
|
||||
@Test("System instructions contain schema")
|
||||
func systemInstructionsContainSchema() {
|
||||
let builder = PromptBuilder(schema: makeSampleSchema())
|
||||
let instructions = builder.buildSystemInstructions()
|
||||
|
||||
#expect(instructions.contains("DATABASE SCHEMA"))
|
||||
#expect(instructions.contains("TABLE users"))
|
||||
#expect(instructions.contains("TABLE orders"))
|
||||
#expect(instructions.contains("name TEXT"))
|
||||
#expect(instructions.contains("email TEXT"))
|
||||
}
|
||||
|
||||
@Test("System instructions contain foreign keys from schema")
|
||||
func systemInstructionsContainForeignKeys() {
|
||||
let builder = PromptBuilder(schema: makeSampleSchema())
|
||||
let instructions = builder.buildSystemInstructions()
|
||||
|
||||
#expect(instructions.contains("FOREIGN KEY"))
|
||||
#expect(instructions.contains("REFERENCES users(id)"))
|
||||
}
|
||||
|
||||
@Test("System instructions contain SQL generation rules")
|
||||
func systemInstructionsContainRules() {
|
||||
let builder = PromptBuilder(schema: makeSampleSchema())
|
||||
let instructions = builder.buildSystemInstructions()
|
||||
|
||||
#expect(instructions.contains("SQL GENERATION RULES"))
|
||||
#expect(instructions.contains("Use ONLY the tables and columns"))
|
||||
#expect(instructions.contains("Never generate DDL"))
|
||||
}
|
||||
|
||||
@Test("System instructions contain output format section")
|
||||
func systemInstructionsContainOutputFormat() {
|
||||
let builder = PromptBuilder(schema: makeSampleSchema())
|
||||
let instructions = builder.buildSystemInstructions()
|
||||
|
||||
#expect(instructions.contains("OUTPUT FORMAT"))
|
||||
}
|
||||
|
||||
@Test("Default allowlist is read-only")
|
||||
func defaultAllowlistIsReadOnly() {
|
||||
let builder = PromptBuilder(schema: makeSampleSchema())
|
||||
let instructions = builder.buildSystemInstructions()
|
||||
|
||||
#expect(instructions.contains("ONLY generate SELECT queries"))
|
||||
#expect(instructions.contains("No data modifications"))
|
||||
}
|
||||
|
||||
@Test("Standard allowlist shows correct operations")
|
||||
func standardAllowlistInstructions() {
|
||||
let builder = PromptBuilder(schema: makeSampleSchema(), allowlist: .standard)
|
||||
let instructions = builder.buildSystemInstructions()
|
||||
|
||||
#expect(instructions.contains("INSERT"))
|
||||
#expect(instructions.contains("SELECT"))
|
||||
#expect(instructions.contains("UPDATE"))
|
||||
}
|
||||
|
||||
@Test("Unrestricted allowlist warns about DELETE")
|
||||
func unrestrictedAllowlistWarnsAboutDelete() {
|
||||
let builder = PromptBuilder(schema: makeSampleSchema(), allowlist: .unrestricted)
|
||||
let instructions = builder.buildSystemInstructions()
|
||||
|
||||
#expect(instructions.contains("DELETE"))
|
||||
#expect(instructions.contains("destructive"))
|
||||
#expect(instructions.contains("confirmation"))
|
||||
}
|
||||
|
||||
@Test("Additional context is appended")
|
||||
func additionalContextAppended() {
|
||||
let builder = PromptBuilder(
|
||||
schema: makeSampleSchema(),
|
||||
additionalContext: "All dates are stored in ISO 8601 format."
|
||||
)
|
||||
let instructions = builder.buildSystemInstructions()
|
||||
|
||||
#expect(instructions.contains("ADDITIONAL CONTEXT"))
|
||||
#expect(instructions.contains("ISO 8601"))
|
||||
}
|
||||
|
||||
@Test("No additional context section when nil")
|
||||
func noAdditionalContextWhenNil() {
|
||||
let builder = PromptBuilder(schema: makeSampleSchema())
|
||||
let instructions = builder.buildSystemInstructions()
|
||||
|
||||
#expect(!instructions.contains("ADDITIONAL CONTEXT"))
|
||||
}
|
||||
|
||||
@Test("No additional context section when empty string")
|
||||
func noAdditionalContextWhenEmpty() {
|
||||
let builder = PromptBuilder(schema: makeSampleSchema(), additionalContext: "")
|
||||
let instructions = builder.buildSystemInstructions()
|
||||
|
||||
#expect(!instructions.contains("ADDITIONAL CONTEXT"))
|
||||
}
|
||||
|
||||
@Test("Empty schema produces valid instructions")
|
||||
func emptySchemaProducesValidInstructions() {
|
||||
let builder = PromptBuilder(schema: makeEmptySchema())
|
||||
let instructions = builder.buildSystemInstructions()
|
||||
|
||||
#expect(instructions.contains("ROLE"))
|
||||
#expect(instructions.contains("SQL GENERATION RULES"))
|
||||
// Schema section should still be present, just empty
|
||||
#expect(instructions.contains("DATABASE SCHEMA"))
|
||||
}
|
||||
|
||||
// MARK: - User Prompt Tests
|
||||
|
||||
@Test("User prompt passes through question directly")
|
||||
func userPromptPassesThrough() {
|
||||
let builder = PromptBuilder(schema: makeSampleSchema())
|
||||
let prompt = builder.buildUserPrompt("How many users signed up this week?")
|
||||
|
||||
#expect(prompt == "How many users signed up this week?")
|
||||
}
|
||||
|
||||
// MARK: - Follow-up Prompt Tests
|
||||
|
||||
@Test("Follow-up prompt includes previous context")
|
||||
func followUpPromptIncludesPreviousContext() {
|
||||
let builder = PromptBuilder(schema: makeSampleSchema())
|
||||
let prompt = builder.buildFollowUpPrompt(
|
||||
"Now sort them by name",
|
||||
previousSQL: "SELECT * FROM users WHERE created_at > date('now', '-7 days')",
|
||||
previousResultSummary: "Found 42 users who signed up this week"
|
||||
)
|
||||
|
||||
#expect(prompt.contains("Previous query:"))
|
||||
#expect(prompt.contains("SELECT * FROM users"))
|
||||
#expect(prompt.contains("Previous result:"))
|
||||
#expect(prompt.contains("42 users"))
|
||||
#expect(prompt.contains("Follow-up question:"))
|
||||
#expect(prompt.contains("sort them by name"))
|
||||
}
|
||||
|
||||
// MARK: - Schema Description Quality
|
||||
|
||||
@Test("Schema includes column types and constraints")
|
||||
func schemaIncludesColumnDetails() {
|
||||
let builder = PromptBuilder(schema: makeSampleSchema())
|
||||
let instructions = builder.buildSystemInstructions()
|
||||
|
||||
// Should include type info
|
||||
#expect(instructions.contains("INTEGER"))
|
||||
#expect(instructions.contains("TEXT"))
|
||||
#expect(instructions.contains("REAL"))
|
||||
|
||||
// Should include constraints
|
||||
#expect(instructions.contains("NOT NULL"))
|
||||
#expect(instructions.contains("PRIMARY KEY"))
|
||||
}
|
||||
|
||||
@Test("Schema includes index information")
|
||||
func schemaIncludesIndexes() {
|
||||
let builder = PromptBuilder(schema: makeSampleSchema())
|
||||
let instructions = builder.buildSystemInstructions()
|
||||
|
||||
#expect(instructions.contains("INDEX"))
|
||||
#expect(instructions.contains("idx_users_email"))
|
||||
}
|
||||
|
||||
// MARK: - Sendable Conformance
|
||||
|
||||
@Test("PromptBuilder is Sendable")
|
||||
func promptBuilderIsSendable() async {
|
||||
let builder = PromptBuilder(schema: makeSampleSchema())
|
||||
|
||||
// Verify it can be sent across concurrency boundaries
|
||||
let instructions = await Task.detached {
|
||||
builder.buildSystemInstructions()
|
||||
}.value
|
||||
|
||||
#expect(instructions.contains("ROLE"))
|
||||
}
|
||||
|
||||
// MARK: - Custom Allowlist
|
||||
|
||||
@Test("Custom allowlist with select and delete only")
|
||||
func customAllowlist() {
|
||||
let allowlist = OperationAllowlist([.select, .delete])
|
||||
let builder = PromptBuilder(schema: makeSampleSchema(), allowlist: allowlist)
|
||||
let instructions = builder.buildSystemInstructions()
|
||||
|
||||
#expect(instructions.contains("DELETE"))
|
||||
#expect(instructions.contains("SELECT"))
|
||||
#expect(instructions.contains("destructive"))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user