feat(ai): add use_gpu setting

False by default for wider compatibility.

Author: Daniele Briggi

src/sqlite_rag/cli.py

@@ -143,6 +143,9 @@ def configure_settings(
     weight_vec: Optional[float] = typer.Option(
         None, help="Weight for vector search results"
     ),
+    use_gpu: Optional[bool] = typer.Option(
+        None, help="Whether to allow sqlite-ai extension to use the GPU"
+    ),
 ):
     """Configure settings for the RAG system.
 
@@ -165,6 +168,7 @@ def configure_settings(
         "quantize_preload": quantize_preload,
         "weight_fts": weight_fts,
         "weight_vec": weight_vec,
+        "use_gpu": use_gpu,
     }
 
     # Filter out None values (unset options)

src/sqlite_rag/database.py

@@ -21,7 +21,12 @@ def initialize(conn: sqlite3.Connection, settings: Settings) -> sqlite3.Connecti
         conn.enable_load_extension(True)
         try:
             conn.load_extension(
-                str(importlib.resources.files("sqliteai.binaries.cpu") / "ai")
+                str(
+                    importlib.resources.files(
+                        "sqliteai.binaries." + ("gpu" if settings.use_gpu else "cpu")
+                    )
+                    / "ai"
+                )
             )
             conn.load_extension(
                 str(importlib.resources.files("sqlite-vector.binaries") / "vector")

src/sqlite_rag/settings.py

@@ -44,6 +44,10 @@ class Settings:
     weight_fts: float = 1.0
     weight_vec: float = 1.0
 
+    # Allow the sqlite-ai extension to use the GPU
+    # See: https://github.com/sqliteai/sqlite-ai
+    use_gpu = False
+
 
 class SettingsManager:
     def __init__(self, connection: sqlite3.Connection):