Skip to content

Commit

Permalink
Added FTS5 tests, formatting
Browse files Browse the repository at this point in the history
  • Loading branch information
jkrukowski committed Nov 27, 2024
1 parent 7dd6b5c commit eaa734a
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 18 deletions.
69 changes: 51 additions & 18 deletions Sources/SQLiteVecCLI/CLI.swift
Original file line number Diff line number Diff line change
@@ -1,33 +1,66 @@
import Foundation
import SQLiteVec

struct Item {
let index: Int
let vector: [Float]
let title: String
let content: String
}

@main
enum CLI {
static func main() async throws {
try SQLiteVec.initialize()
let data: [(index: Int, vector: [Float], title: String, content: String)] = [
(1, [0.1, 0.1, 0.1, 0.1], "Introduction to Machine Learning", "Machine learning is a subset of artificial intelligence..."),
(2, [0.2, 0.2, 0.2, 0.2], "Deep Learning Basics", "Deep learning uses neural networks to learn from data..."),
(3, [0.3, 0.3, 0.3, 0.3], "Natural Language Processing", "NLP combines linguistics and machine learning..."),
(4, [0.4, 0.4, 0.4, 0.4], "Computer Vision", "Computer vision enables machines to understand visual data..."),
(5, [0.5, 0.5, 0.5, 0.5], "Reinforcement Learning", "Reinforcement learning involves agents learning through interaction...")
let data = [
Item(
index: 1,
vector: [0.1, 0.1, 0.1, 0.1],
title: "Introduction to Machine Learning",
content: "Machine learning is a subset of artificial intelligence..."
),
Item(
index: 2,
vector: [0.2, 0.2, 0.2, 0.2],
title: "Deep Learning Basics",
content: "Deep learning uses neural networks to learn from data..."
),
Item(
index: 3,
vector: [0.3, 0.3, 0.3, 0.3],
title: "Natural Language Processing",
content: "NLP combines linguistics and machine learning..."
),
Item(
index: 4,
vector: [0.4, 0.4, 0.4, 0.4],
title: "Computer Vision",
content: "Computer vision enables machines to understand visual data..."
),
Item(
index: 5,
vector: [0.5, 0.5, 0.5, 0.5],
title: "Reinforcement Learning",
content: "Reinforcement learning involves agents learning through interaction..."
),
]
let query: [Float] = [0.3, 0.3, 0.3, 0.3]
let textQuery = "learning"

let db = try Database(.inMemory)
try await db.execute("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4])")
try await db.execute("""
CREATE VIRTUAL TABLE docs USING fts5(
title,
content,
tokenize='porter'
)
""")
try await db.execute(
"""
CREATE VIRTUAL TABLE docs USING fts5(
title,
content,
tokenize='porter'
)
""")
for row in data {
try await db.execute(
"""
INSERT INTO vec_items(rowid, embedding)
INSERT INTO vec_items(rowid, embedding)
VALUES (?, ?)
""",
params: [row.index, row.vector]
Expand All @@ -42,10 +75,10 @@ enum CLI {
}
let result = try await db.query(
"""
SELECT rowid, distance
FROM vec_items
WHERE embedding MATCH ?
ORDER BY distance
SELECT rowid, distance
FROM vec_items
WHERE embedding MATCH ?
ORDER BY distance
LIMIT 3
""",
params: [query]
Expand Down
9 changes: 9 additions & 0 deletions Tests/SQLiteVecTests/CoreTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -16,4 +16,13 @@ final class CoreTests: XCTestCase {
XCTAssertTrue(extensionNames.contains("vec_each"), "vec_each should be loaded")
XCTAssertTrue(extensionNames.contains("vec0"), "vec0 should be loaded")
}

func testCompileFlags() async throws {
try SQLiteVec.initialize()
let db = try Database(.inMemory)
let result = try await db.query("PRAGMA compile_options")
let compileFlags = result.compactMap { $0["compile_options"] as? String }

XCTAssertTrue(compileFlags.contains("ENABLE_FTS5"), "FTS5 should be enabled")
}
}

0 comments on commit eaa734a

Please sign in to comment.