diff --git a/studio/backend/core/rag/db.py b/studio/backend/core/rag/db.py new file mode 100644 index 0000000000..4fbd238843 --- /dev/null +++ b/studio/backend/core/rag/db.py @@ -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 `, + 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 diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 60f1b562fe..299190fab3 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -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 diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index bc24c1c84a..b43fc849de 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -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 diff --git a/studio/backend/core/rag/vector_store.py b/studio/backend/core/rag/vector_store.py index 51553d8d77..9517179715 100644 --- a/studio/backend/core/rag/vector_store.py +++ b/studio/backend/core/rag/vector_store.py @@ -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_`` for standalone -knowledge bases and ``thread_`` for per-thread document sets. +A "scope" is `kb_` for standalone knowledge bases or +`thread_` 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() diff --git a/studio/backend/requirements/rag.txt b/studio/backend/requirements/rag.txt index 3614a969e9..5ac1840b41 100644 --- a/studio/backend/requirements/rag.txt +++ b/studio/backend/requirements/rag.txt @@ -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 diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index fc59c508aa..504eb50b9e 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -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. """ diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index 13e1c65ac4..46673dacd3 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -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()) diff --git a/tests/python/test_rag_vector_store.py b/tests/python/test_rag_vector_store.py index 26bbdefd26..f269b5ccfc 100644 --- a/tests/python/test_rag_vector_store.py +++ b/tests/python/test_rag_vector_store.py @@ -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"