diff --git a/Sources/SQLiteVecCLI/CLI.swift b/Sources/SQLiteVecCLI/CLI.swift index 5cd3c27..679e922 100644 --- a/Sources/SQLiteVecCLI/CLI.swift +++ b/Sources/SQLiteVecCLI/CLI.swift @@ -5,16 +5,25 @@ import SQLiteVec enum CLI { static func main() async throws { try SQLiteVec.initialize() - let data: [(index: Int, vector: [Float])] = [ - (1, [0.1, 0.1, 0.1, 0.1]), - (2, [0.2, 0.2, 0.2, 0.2]), - (3, [0.3, 0.3, 0.3, 0.3]), - (4, [0.4, 0.4, 0.4, 0.4]), - (5, [0.5, 0.5, 0.5, 0.5]), + 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 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' + ) + """) for row in data { try await db.execute( """ @@ -23,6 +32,13 @@ enum CLI { """, params: [row.index, row.vector] ) + try await db.execute( + """ + INSERT INTO docs(rowid, title, content) + VALUES (?, ?, ?) + """, + params: [row.index, row.title, row.content] + ) } let result = try await db.query( """ @@ -35,5 +51,17 @@ enum CLI { params: [query] ) print(result) + + let textResults = try await db.query( + """ + SELECT rowid, title, content, rank + FROM docs + WHERE docs MATCH ? + ORDER BY rank + LIMIT 3 + """, + params: [textQuery] + ) + print(textResults) } }