Merge pull request #10 from sqliteai/highlight-sentence-in-results

Highlight sentence in results

Author: danielebriggi

.devcontainer/devcontainer.json

@@ -1,18 +1,6 @@
 {
     "name": "Python 3.10",
     "image": "mcr.microsoft.com/devcontainers/python:3.10",
-    "runArgs": [
-        "--runtime",
-        "nvidia",
-        "--gpus",
-        "all",
-        // optional but make sure CUDA workloads are available
-        "--env",
-        "NVIDIA_VISIBLE_DEVICES=all",
-        // optional but make sure CUDA workloads are available
-        "--env",
-        "NVIDIA_DRIVER_CAPABILITIES=compute,utility"
-    ],
     "customizations": {
         "vscode": {
             "extensions": [
@@ -26,13 +14,5 @@
                 "hbenl.vscode-test-explorer"
             ]
         }
-    },
-    "hostRequirements": {
-        "gpu": "optional"
-    },
-    "remoteEnv": {
-        // optional but make sure CUDA workloads are available
-        "NVIDIA_VISIBLE_DEVICES": "all",
-        "NVIDIA_DRIVER_CAPABILITIES": "compute,utility"
     }
 }

.devcontainer/py3.11/devcontainer.json

@@ -0,0 +1,38 @@
+{
+    "name": "Python 3.11",
+    "image": "mcr.microsoft.com/devcontainers/python:3.11",
+    "runArgs": [
+        "--runtime",
+        "nvidia",
+        "--gpus",
+        "all",
+        // optional but make sure CUDA workloads are available
+        "--env",
+        "NVIDIA_VISIBLE_DEVICES=all",
+        // optional but make sure CUDA workloads are available
+        "--env",
+        "NVIDIA_DRIVER_CAPABILITIES=compute,utility"
+    ],
+    "customizations": {
+        "vscode": {
+            "extensions": [
+                "ms-python.black-formatter",
+                "ms-python.flake8",
+                "ms-python.isort",
+                "ms-python.vscode-pylance",
+                "ms-python.python",
+                "ms-python.debugpy",
+                "ms-python.vscode-python-envs",
+                "hbenl.vscode-test-explorer"
+            ]
+        }
+    },
+    "hostRequirements": {
+        "gpu": "optional"
+    },
+    "remoteEnv": {
+        // optional but make sure CUDA workloads are available
+        "NVIDIA_VISIBLE_DEVICES": "all",
+        "NVIDIA_DRIVER_CAPABILITIES": "compute,utility"
+    }
+}

.github/workflows/test.yaml

@@ -49,6 +49,10 @@ jobs:
 
     - name: Test
       # Using default directory for models
+      # COVERAGE_CORE=pytrace: Workaround for Python 3.11 segfault with SQLite extensions + C tracer
+      # See: https://github.com/nedbat/coveragepy/issues/1665
+      env:
+        COVERAGE_CORE: ${{ matrix.python-version == '3.11' && 'pytrace' || '' }}
       run: |
         pytest --cov --cov-branch --cov-report=xml -v -m "not slow" ./tests
 

src/sqlite_rag/cli.py

@@ -439,12 +439,12 @@ def reset(
 def search(
     ctx: typer.Context,
     query: str,
-    limit: int = typer.Option(10, help="Number of results to return"),
+    limit: int = typer.Option(5, help="Number of results to return"),
     debug: bool = typer.Option(
         False,
         "-d",
         "--debug",
-        help="Print extra debug information with modern formatting",
+        help="Print extra debug information with sentence-level details",
     ),
     peek: bool = typer.Option(
         False, "--peek", help="Print debug information using compact table format"

src/sqlite_rag/database.py

@@ -76,28 +76,50 @@ def _create_schema(conn: sqlite3.Connection, settings: Settings):
         )
 
         # TODO: this table is not ready for sqlite-sync, it uses the id AUTOINCREMENT
-        cursor.execute(
+        cursor.executescript(
             """
             CREATE TABLE IF NOT EXISTS chunks (
                 id INTEGER PRIMARY KEY AUTOINCREMENT,
                 document_id TEXT,
                 content TEXT,
-                embedding BLOB,
-                FOREIGN KEY (document_id) REFERENCES documents (id) ON DELETE CASCADE
+                embedding BLOB
             );
+            CREATE INDEX IF NOT EXISTS idx_chunks_document_id ON chunks (document_id);
         """
         )
 
+        cursor.executescript(
+            """
+            CREATE TABLE IF NOT EXISTS sentences (
+                id TEXT PRIMARY KEY,
+                chunk_id INTEGER,
+                content TEXT,
+                embedding BLOB,
+                start_offset INTEGER,
+                end_offset INTEGER
+            );
+            CREATE INDEX IF NOT EXISTS idx_sentences_chunk_id ON sentences (chunk_id);
+            """
+        )
+
         cursor.execute(
             """
             CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(content, content='chunks', content_rowid='id');
         """
         )
 
         cursor.execute(
-            f"""
-            SELECT vector_init('chunks', 'embedding', 'type={settings.vector_type},dimension={settings.embedding_dim},{settings.other_vector_options}');
-        """
+            """
+            SELECT vector_init('chunks', 'embedding', ?);
+        """,
+            (settings.get_vector_init_options(),),
+        )
+        # TODO: same configuration as chunks (or different options?)
+        cursor.execute(
+            """
+            SELECT vector_init('sentences', 'embedding', ?);
+        """,
+            (settings.get_vector_init_options(),),
         )
 
         conn.commit()