ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: swap RAG vector store from Qdrant to sqlite-vec

Browse source 
asg017/sqlite-vec is Apache-2.0 and OSI-approved. Replaces
qdrant-client (~30 MB) with a small SQLite extension loaded into a
dedicated rag.db file. Single file holds RAG vectors; bm25s indexes
and chat-side studio.db are unaffected.

- New core/rag/db.py owns the rag.db connection and sqlite-vec load.
  Extension load runs once at first open. Process-wide singleton
  protected by a lock; check_same_thread=False + WAL handles the
  FastAPI thread pool.
- core/rag/vector_store.py keeps the same public API
  (ensure_collection / upsert_chunks / search / collection_exists /
  delete_scope / delete_document) so callers in routes/rag.py,
  core/rag/ingestion.py, core/rag/tool.py, and core/rag/retrieval.py
  don't change. ensure_collection is now a no-op; collection_exists
  returns True iff the scope has at least one indexed vector.
- search uses sqlite-vec's vec_distance_cosine and converts distance
  to similarity in [0, 1] so the per-scope min_score threshold
  semantics stay identical.
- Mixed-dim scopes coexist behind WHERE scope = ? — the per-scope
  embedder resolver guarantees one embedder per scope.
- requirements/rag.txt swaps qdrant-client for sqlite-vec.
- utils/paths/storage_roots.py drops rag_vectordb_root() (the old
  qdrant directory); rag.db lives directly under rag_root().
- Rewritten tests/python/test_rag_vector_store.py for the new
  semantics (collection_exists tracks populated scopes; new tests
  for filtered search and upsert conflict resolution).

Python build requirement: connection.enable_load_extension(True)
must be available. install.sh creates the venv via uv-managed
python-build-standalone, which is compiled with
--enable-loadable-sqlite-extensions, so this works on standard
installs. core/rag/db.py raises an actionable error on the rare
custom-interpreter case.
This commit is contained in:
Roland Tannous 2026-05-25 15:13:56 +04:00
commit 2093fb1608
8 changed files with 364 additions and 165 deletions
Download patch file Download diff file Expand all files Collapse all files

112
studio/backend/core/rag/db.py Normal file
View file

@ -0,0 +1,112 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""sqlite-vec backed connection helper for RAG vectors.
Single process-wide connection, opened lazily on first use. The
extension-load step runs once at open time. studio.db (chat history,
RAG metadata) stays untouched, so the extension-load surface is scoped
to the RAG code path only — chat code keeps its plain sqlite handle.
"""
from __future__ import annotations
import logging
import sqlite3
import threading
from pathlib import Path
from utils.paths.storage_roots import ensure_dir, rag_root
logger = logging.getLogger(__name__)
_conn: sqlite3.Connection | None = None
_conn_lock = threading.Lock()
def rag_db_path() -> Path:
return rag_root() / "rag.db"
def _load_sqlite_vec(conn: sqlite3.Connection) -> None:
"""Enable extension loading and pull in sqlite-vec.
install.sh creates the studio venv via `uv venv --python <ver>`,
which uses uv's managed python-build-standalone build. That CPython
is compiled with --enable-loadable-sqlite-extensions, so this path
succeeds on standard installs. The actionable error message is
here for the rare custom-interpreter case.
"""
try:
conn.enable_load_extension(True)
except AttributeError as exc:
raise RuntimeError(
"This Python build cannot load SQLite extensions "
"(connection.enable_load_extension is unavailable). RAG "
"requires sqlite-vec, which loads as a SQLite extension. "
"Re-install studio via install.sh so the venv uses uv's "
"managed Python (python-build-standalone), compiled with "
"--enable-loadable-sqlite-extensions."
) from exc
import sqlite_vec
sqlite_vec.load(conn)
conn.enable_load_extension(False)
def _ensure_schema(conn: sqlite3.Connection) -> None:
conn.executescript(
"""
CREATE TABLE IF NOT EXISTS rag_vectors (
chunk_id TEXT PRIMARY KEY,
scope TEXT NOT NULL,
document_id TEXT NOT NULL,
chunk_index INTEGER NOT NULL,
kind TEXT NOT NULL DEFAULT 'text',
dim INTEGER NOT NULL,
vector BLOB NOT NULL,
payload_json TEXT NOT NULL DEFAULT '{}'
);
CREATE INDEX IF NOT EXISTS idx_rag_vectors_scope
ON rag_vectors(scope);
CREATE INDEX IF NOT EXISTS idx_rag_vectors_scope_doc
ON rag_vectors(scope, document_id);
"""
)
conn.commit()
def get_rag_connection() -> sqlite3.Connection:
"""Lazy process-wide sqlite connection to rag.db with sqlite-vec loaded.
Returns the cached connection on subsequent calls. FastAPI's thread
pool plus check_same_thread=False + WAL mode handles concurrent
reads; writes are serialized by SQLite itself.
"""
global _conn
with _conn_lock:
if _conn is None:
ensure_dir(rag_root())
conn = sqlite3.connect(
str(rag_db_path()),
check_same_thread = False,
)
conn.row_factory = sqlite3.Row
conn.execute("PRAGMA journal_mode = WAL")
_load_sqlite_vec(conn)
_ensure_schema(conn)
_conn = conn
logger.info("RAG vector store: opened %s", rag_db_path())
return _conn
def _reset_for_tests() -> None:
"""Drop the cached connection. Test-only — production never calls."""
global _conn
with _conn_lock:
if _conn is not None:
try:
_conn.close()
except Exception:
pass
_conn = None

4
studio/backend/core/rag/embeddings.py
View file

@ -260,8 +260,8 @@ def encode_images(
Works with CLIP-family models (BGE-VL, openai/clip-*) whose
`encode` accepts PIL.Image objects in the same call as text. The
returned vectors live in the same 512-d (or model-specific) space
as text vectors from this model, so a single Qdrant collection
holds both kinds.
as text vectors from this model, so a single scope's vector rows
hold both kinds.
"""
from io import BytesIO

11
studio/backend/core/rag/ingestion.py
View file

@ -8,11 +8,12 @@ spawn a fresh subprocess per job with ``mp.get_context("spawn")`` and stream
progress events back over a queue. The subprocess does the heavy work
(parse → chunk → load embedder → embed in batches) and ships
``(chunks, vectors)`` batches back. The parent persists everything:
sqlite rows, Qdrant points, and (at job completion) a rebuilt BM25 index.
sqlite rows, vector_store rows in rag.db, and (at job completion) a
rebuilt BM25 index.
Only the parent process holds the Qdrant local-mode file lock — the
subprocess never opens it directly. This keeps search available
throughout the lifetime of an ingestion job.
Only the parent process owns the rag.db connection (sqlite-vec loaded
there); subprocesses never open it directly. This keeps search
available throughout the lifetime of an ingestion job.
"""
from __future__ import annotations
@ -449,7 +450,7 @@ def _insert_chunks_and_collect_for_bm25(
chunks_meta: list[dict],
vectors: list[list[float]],
) -> list[dict]:
"""Insert chunks into sqlite + Qdrant; return [{id, text}] for BM25.
"""Insert chunks into sqlite + vector_store; return [{id, text}] for BM25.
Image-kind chunks ship a stable image_path and skip BM25 (no text
body to tokenise). Paired image/caption chunks share a pair_group

252
studio/backend/core/rag/vector_store.py
View file

@ -1,44 +1,30 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Qdrant local-mode vector store.
"""SQLite-backed vector store for RAG, distance math via sqlite-vec.
Qdrant's local mode (``QdrantClient(path=...)``) acquires a file lock on
the storage directory, so only one process at a time can hold the
client. That means *all* vector reads/writes funnel through this module
in the FastAPI parent process. Ingestion subprocesses do not open
Qdrant directly — they compute vectors and send them back over a queue
for the parent to persist.
Replaces the previous Qdrant local-mode store. Vectors are stored as
BLOBs in a single `rag_vectors` table inside rag.db, keyed by
chunk_id; cosine distance is computed by sqlite-vec's
`vec_distance_cosine(blob, blob)` scalar function.
A "scope" is the collection name: ``kb_<uuid>`` for standalone
knowledge bases and ``thread_<uuid>`` for per-thread document sets.
A "scope" is `kb_<uuid>` for standalone knowledge bases or
`thread_<uuid>` for per-thread document sets. The scope column is
indexed; queries filter by scope before computing distances so
different scopes can hold vectors of different dimensions without
breaking the cosine math. The per-scope embedder resolver
(routes/rag.py:_resolve_scope_embedder) guarantees one embedder per
scope, so dims within a scope are always consistent.
"""
from __future__ import annotations
import json
import logging
import threading
from typing import Any, Iterable
from utils.paths.storage_roots import ensure_dir, rag_vectordb_root
from typing import Iterable
logger = logging.getLogger(__name__)
_client: Any | None = None
_client_lock = threading.Lock()
def get_qdrant() -> Any:
"""Lazy singleton; parent process only."""
global _client
with _client_lock:
if _client is None:
from qdrant_client import QdrantClient
path = ensure_dir(rag_vectordb_root())
_client = QdrantClient(path = str(path))
return _client
def kb_scope(kb_id: str) -> str:
return f"kb_{kb_id}"
@ -49,45 +35,77 @@ def thread_scope(thread_id: str) -> str:
def collection_exists(scope: str) -> bool:
client = get_qdrant()
try:
client.get_collection(collection_name = scope)
return True
except Exception:
return False
"""Whether the scope has any indexed vectors.
Used by callers (notably retrieval.retrieve_dense) to short-circuit
when the scope was never populated. Cheap — a covering index hit.
"""
from core.rag.db import get_rag_connection
conn = get_rag_connection()
row = conn.execute(
"SELECT 1 FROM rag_vectors WHERE scope = ? LIMIT 1",
(scope,),
).fetchone()
return row is not None
def ensure_collection(scope: str, dim: int) -> None:
client = get_qdrant()
if collection_exists(scope):
return
from qdrant_client.models import Distance, VectorParams
client.create_collection(
collection_name = scope,
vectors_config = VectorParams(size = dim, distance = Distance.COSINE),
)
def upsert_chunks(
scope: str,
points: Iterable[dict],
) -> None:
"""Insert/update chunk vectors.
Each point must have keys ``id`` (str), ``vector`` (list[float]) and
``payload`` (dict with at least ``document_id`` and ``chunk_index``).
"""No-op for sqlite-vec — vectors get inserted directly into the
shared table. Kept for API parity with the previous Qdrant store
so ingestion callers don't need conditional logic.
"""
from qdrant_client.models import PointStruct
_ = scope, dim # unused; signature preserved
client = get_qdrant()
structured = [
PointStruct(id = p["id"], vector = p["vector"], payload = p["payload"])
for p in points
]
if not structured:
def upsert_chunks(scope: str, points: Iterable[dict]) -> None:
"""Insert/update vectors. Each point: {id, vector, payload}.
Conflict resolution is per chunk_id (the primary key): re-ingesting
overwrites in place. Payload is round-tripped as JSON so the
Qdrant-shaped {filename, page_number, kind, ...} dicts callers
already build can be reused unchanged.
"""
import sqlite_vec
from core.rag.db import get_rag_connection
rows = []
for p in points:
payload = p.get("payload") or {}
vec = list(p["vector"])
rows.append(
(
p["id"],
scope,
str(payload.get("document_id") or ""),
int(payload.get("chunk_index") or 0),
str(payload.get("kind") or "text"),
len(vec),
sqlite_vec.serialize_float32(vec),
json.dumps(payload, default = str),
)
)
if not rows:
return
client.upsert(collection_name = scope, points = structured)
conn = get_rag_connection()
conn.executemany(
"""
INSERT INTO rag_vectors
(chunk_id, scope, document_id, chunk_index, kind, dim, vector, payload_json)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(chunk_id) DO UPDATE SET
scope = excluded.scope,
document_id = excluded.document_id,
chunk_index = excluded.chunk_index,
kind = excluded.kind,
dim = excluded.dim,
vector = excluded.vector,
payload_json = excluded.payload_json
""",
rows,
)
conn.commit()
def search(
@ -97,73 +115,69 @@ def search(
top_k: int,
document_ids: list[str] | None = None,
) -> list[dict]:
from qdrant_client.models import FieldCondition, Filter, MatchAny
"""Cosine-distance search filtered to a single scope.
Returns rows in the same shape as the old Qdrant path —
{chunk_id, score, payload} — where ``score`` is cosine similarity
in [0, 1] (vec_distance_cosine returns 1 - similarity, so we
invert). Filtered-by-document_ids variant for the per-thread
"search only these uploads" case.
"""
import sqlite_vec
from core.rag.db import get_rag_connection
client = get_qdrant()
query_filter = None
if document_ids:
query_filter = Filter(
must = [
FieldCondition(
key = "document_id",
match = MatchAny(any = document_ids),
)
]
)
if not collection_exists(scope):
return []
# qdrant-client 1.10 deprecated `search()` in favor of
# `query_points()`. Fall back to `search()` on older clients so
# the call works against any pinned version in the install matrix.
if hasattr(client, "query_points"):
response = client.query_points(
collection_name = scope,
query = query_vector,
limit = top_k,
query_filter = query_filter,
with_payload = True,
serialized = sqlite_vec.serialize_float32(list(query_vector))
sql = (
"SELECT chunk_id, payload_json, "
" vec_distance_cosine(vector, ?) AS distance "
"FROM rag_vectors WHERE scope = ?"
)
params: list = [serialized, scope]
if document_ids:
placeholders = ",".join("?" for _ in document_ids)
sql += f" AND document_id IN ({placeholders})"
params.extend(document_ids)
sql += " ORDER BY distance ASC LIMIT ?"
params.append(int(top_k))
conn = get_rag_connection()
rows = conn.execute(sql, params).fetchall()
out: list[dict] = []
for row in rows:
score = 1.0 - float(row["distance"])
try:
payload = json.loads(row["payload_json"] or "{}")
except json.JSONDecodeError:
payload = {}
out.append(
{
"chunk_id": row["chunk_id"],
"score": score,
"payload": payload,
}
)
results = response.points
else:
results = client.search(
collection_name = scope,
query_vector = query_vector,
limit = top_k,
query_filter = query_filter,
)
return [
{
"chunk_id": str(r.id),
"score": float(r.score),
"payload": dict(r.payload or {}),
}
for r in results
]
return out
def delete_scope(scope: str) -> None:
client = get_qdrant()
if not collection_exists(scope):
return
client.delete_collection(collection_name = scope)
from core.rag.db import get_rag_connection
conn = get_rag_connection()
conn.execute("DELETE FROM rag_vectors WHERE scope = ?", (scope,))
conn.commit()
def delete_document(scope: str, document_id: str) -> None:
from qdrant_client.models import FieldCondition, Filter, FilterSelector, MatchValue
from core.rag.db import get_rag_connection
if not collection_exists(scope):
return
client = get_qdrant()
client.delete(
collection_name = scope,
points_selector = FilterSelector(
filter = Filter(
must = [
FieldCondition(
key = "document_id",
match = MatchValue(value = document_id),
)
]
)
),
conn = get_rag_connection()
conn.execute(
"DELETE FROM rag_vectors WHERE scope = ? AND document_id = ?",
(scope, document_id),
)
conn.commit()

14
studio/backend/requirements/rag.txt
View file

@ -3,11 +3,15 @@
# path. Skipped in NO_TORCH (Intel Mac GGUF-only) mode because RAG
# embedding relies on sentence-transformers, which requires torch.
# Vector store + lexical index. qdrant-client supports a pure-Python
# local mode (QdrantClient(path=...)) that we use to keep the studio
# install self-contained — no separate server. bm25s persists per-scope
# indices to disk.
qdrant-client>=1.12
# Vector store + lexical index.
#
# sqlite-vec is an Apache-2.0 SQLite extension (asg017/sqlite-vec) that
# adds vector functions (vec_distance_cosine, serialize_float32, vec0
# virtual tables). The studio loads it into a dedicated rag.db file and
# stores vectors as BLOB columns alongside the RAG metadata — no separate
# vector server, no second client library. bm25s persists per-scope
# lexical indexes to disk.
sqlite-vec>=0.1.5
bm25s>=0.2
# Image preprocessing helpers required by Qwen3-VL-Embedding-2B (the

4
studio/backend/routes/rag.py
View file

@ -182,7 +182,7 @@ def _now_ms() -> int:
def _resolve_scope_embedder(scope: str) -> str | None:
"""Look up the embedder that populated `scope`'s Qdrant collection.
"""Look up the embedder that populated `scope`'s vector rows.
Returns the model name to use for query-side embedding so the
similarity math doesn't mix vector spaces (Qwen3-VL 2048-d
@ -985,7 +985,7 @@ def clear_thread_documents(
) -> dict:
"""Purge every RAG document attached to ``thread_id``.
Removes the per-thread Qdrant collection, the bm25 index, the
Removes the per-thread vector rows, the bm25 index, the
rag_documents/rag_chunks rows, and the uploaded files. The chat
thread itself is untouched.
"""

5
studio/backend/utils/paths/storage_roots.py
View file

@ -128,10 +128,6 @@ def rag_uploads_root() -> Path:
return rag_root() / "uploads"
def rag_vectordb_root() -> Path:
return rag_root() / "qdrant"
def rag_bm25_root() -> Path:
return rag_root() / "bm25"
@ -279,7 +275,6 @@ def ensure_studio_directories() -> None:
tensorboard_root,
rag_root,
rag_uploads_root,
rag_vectordb_root,
rag_bm25_root,
):
ensure_dir(dir_fn())

127
tests/python/test_rag_vector_store.py
View file

@ -1,4 +1,4 @@
"""Qdrant local-mode vector store tests (skipped if qdrant-client is unavailable)."""
"""sqlite-vec backed RAG vector store tests."""
import sys
from pathlib import Path
@ -10,58 +10,73 @@ STUDIO_BACKEND = REPO_ROOT / "studio" / "backend"
if str(STUDIO_BACKEND) not in sys.path:
sys.path.insert(0, str(STUDIO_BACKEND))
pytest.importorskip("qdrant_client")
pytest.importorskip("sqlite_vec")
@pytest.fixture
def isolated_qdrant(tmp_path, monkeypatch):
def isolated_rag_db(tmp_path, monkeypatch):
"""Point rag.db at tmp_path and reset the cached connection so
each test gets a fresh database.
"""
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
from core.rag import vector_store
from core.rag import db as rag_db
# Reset client cache so the fixture's tmp path is used.
vector_store._client = None
rag_db._reset_for_tests()
yield tmp_path
vector_store._client = None
rag_db._reset_for_tests()
def test_ensure_and_upsert_and_search(isolated_qdrant):
def test_upsert_and_search_returns_nearest_first(isolated_rag_db):
from core.rag import vector_store
scope = "kb_test"
vector_store.ensure_collection(scope, dim = 4)
points = [
{
"id": "p1",
"vector": [1.0, 0.0, 0.0, 0.0],
"payload": {"document_id": "doc1", "chunk_index": 0, "text": "first"},
},
{
"id": "p2",
"vector": [0.0, 1.0, 0.0, 0.0],
"payload": {"document_id": "doc1", "chunk_index": 1, "text": "second"},
},
]
vector_store.upsert_chunks(scope, points)
vector_store.ensure_collection(scope, dim = 4) # no-op under sqlite-vec
vector_store.upsert_chunks(
scope,
[
{
"id": "p1",
"vector": [1.0, 0.0, 0.0, 0.0],
"payload": {"document_id": "doc1", "chunk_index": 0, "text": "first"},
},
{
"id": "p2",
"vector": [0.0, 1.0, 0.0, 0.0],
"payload": {"document_id": "doc1", "chunk_index": 1, "text": "second"},
},
],
)
results = vector_store.search(scope, [1.0, 0.0, 0.0, 0.0], top_k = 2)
assert results
assert len(results) == 2
assert results[0]["chunk_id"] == "p1"
# Cosine similarity converted to [0, 1]; closer = higher.
assert results[0]["score"] > results[1]["score"]
def test_delete_scope_removes_collection(isolated_qdrant):
def test_collection_exists_tracks_populated_scope(isolated_rag_db):
from core.rag import vector_store
scope = "kb_to_delete"
vector_store.ensure_collection(scope, dim = 3)
assert not vector_store.collection_exists(scope)
vector_store.upsert_chunks(
scope,
[
{
"id": "sole",
"vector": [1.0, 0.0, 0.0],
"payload": {"document_id": "d", "chunk_index": 0},
}
],
)
assert vector_store.collection_exists(scope)
vector_store.delete_scope(scope)
assert not vector_store.collection_exists(scope)
def test_delete_document_removes_only_its_points(isolated_qdrant):
def test_delete_document_removes_only_its_points(isolated_rag_db):
from core.rag import vector_store
scope = "kb_doc_del"
vector_store.ensure_collection(scope, dim = 3)
vector_store.upsert_chunks(
scope,
[
@ -82,3 +97,61 @@ def test_delete_document_removes_only_its_points(isolated_qdrant):
doc_ids = {r["payload"]["document_id"] for r in results}
assert "drop" not in doc_ids
assert "keep" in doc_ids
def test_search_filtered_by_document_ids(isolated_rag_db):
from core.rag import vector_store
scope = "kb_filter"
vector_store.upsert_chunks(
scope,
[
{
"id": "a",
"vector": [1.0, 0.0, 0.0],
"payload": {"document_id": "alpha", "chunk_index": 0},
},
{
"id": "b",
"vector": [1.0, 0.0, 0.0],
"payload": {"document_id": "beta", "chunk_index": 0},
},
],
)
results = vector_store.search(
scope,
[1.0, 0.0, 0.0],
top_k = 5,
document_ids = ["alpha"],
)
doc_ids = {r["payload"]["document_id"] for r in results}
assert doc_ids == {"alpha"}
def test_upsert_overwrites_on_conflicting_chunk_id(isolated_rag_db):
from core.rag import vector_store
scope = "kb_overwrite"
vector_store.upsert_chunks(
scope,
[
{
"id": "same",
"vector": [1.0, 0.0, 0.0],
"payload": {"document_id": "d", "chunk_index": 0, "v": "v1"},
}
],
)
vector_store.upsert_chunks(
scope,
[
{
"id": "same",
"vector": [0.0, 1.0, 0.0],
"payload": {"document_id": "d", "chunk_index": 0, "v": "v2"},
}
],
)
results = vector_store.search(scope, [0.0, 1.0, 0.0], top_k = 5)
assert len(results) == 1
assert results[0]["payload"]["v"] == "v2"