VectorDB implementation in Python, very simple, uses SQLite

vectordb.py

"""
Vector database module: SQLite-backed store with OpenAI-compatible embeddings.

Provides:
- VectorDB, SQLite-backed store with embeddings from an OpenAI-compatible embedding API
- source_files table, Tracks each file path and mtime so unchanged files skip re-embedding entirely
- embedding_to_blob / blob_to_embedding, Pack/unpack float32 vectors to/from SQLite BLOBs
- cosine_similarity, Computes similarity between two vectors
- SearchResult, NamedTuple for search hits (id, text, score, source)
- SearchResults, Iterable container with ``.threshold(min_score)`` and ``.with_file_contents()``
- chunk_file_semantically, Semantic chunks (headers / defs / paragraphs) with boundary overlap for retrieval
- create_rag_db, One-shot factory: creates DB and syncs from roots
- run_rag_query, Convenience wrapper that returns search results as list of dicts
"""

from __future__ import annotations

import fnmatch
import math
import os
import re
import sqlite3
import struct
from collections import UserList
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any, NamedTuple, Sequence


def embedding_to_blob(embedding: list[float]) -> bytes:
    """Pack a list of floats into a SQLite BLOB (little-endian float32)."""
    return struct.pack(f"<{len(embedding)}f", *embedding)


def blob_to_embedding(blob: bytes | None) -> list[float]:
    """Unpack a BLOB into a list of floats."""
    if not blob:
        return []
    n = len(blob) // 4
    return list(struct.unpack(f"<{n}f", blob))


def cosine_similarity(a: list[float], b: list[float]) -> float:
    """Cosine similarity between two vectors. Returns 0 if either norm is zero."""
    if len(a) != len(b) or len(a) == 0:
        return 0.0
    dot = math.fsum(x * y for x, y in zip(a, b, strict=True))
    norm_a = math.sqrt(math.fsum(x * x for x in a))
    norm_b = math.sqrt(math.fsum(y * y for y in b))
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return dot / (norm_a * norm_b)


def normalize_roots(
    roots: str | Path | Sequence[str] | Sequence[Path],
) -> list[Path]:
    """Normalize roots to a list of Path. Single str or Path becomes a one-element list."""
    if isinstance(roots, (str, Path)):
        return [Path(roots) if isinstance(roots, str) else roots]
    return [Path(p) if not isinstance(p, Path) else p for p in roots]


def normalize_glob_patterns(patterns: str | list[str]) -> list[str]:
    """Normalize to a list of non-empty glob patterns. Empty or whitespace-only yields ['*.md']."""
    if isinstance(patterns, str):
        return [patterns.strip()] if patterns.strip() else ["*.md"]
    out = [p.strip() for p in patterns if p and str(p).strip()]
    return out if out else ["*.md"]


def collect_file_paths(roots: list[Path], glob_pattern: str | list[str]) -> set[str]:
    """Return set of posix-style paths for files under roots matching any glob pattern."""
    patterns = normalize_glob_patterns(glob_pattern)
    out: set[str] = set()
    for p in roots:
        r = p.resolve()
        if r.is_file():
            if any(fnmatch.fnmatch(r.name, pat) for pat in patterns):
                out.add(r.as_posix())
        elif r.is_dir():
            for pattern in patterns:
                for f in r.rglob(pattern):
                    if f.is_file():
                        out.add(f.resolve().as_posix())
    return out


class SearchResult(NamedTuple):
    """A single result from similarity search."""

    id: int
    text: str
    score: float
    source: str | None = None

    def __str__(self) -> str:
        return self.text

    def __repr__(self) -> str:
        return f"SearchResult(id={self.id}, text={self.text}, score={self.score:.3f}, source={self.source!r})"


class SearchResults(UserList[SearchResult]):
    """
    Iterable list of ``SearchResult`` rows from ``VectorDB.search``.

    Chain filters and expansions without re-querying, e.g.::
        db.search(q, k=20).threshold(0.35).with_file_contents()
    """

    def threshold(self, min_score: float) -> SearchResults:
        """Keep hits whose cosine similarity score is ``>= min_score``."""
        return SearchResults([r for r in self.data if r.score >= min_score])

    def with_file_contents(
        self,
        *,
        encoding: str = "utf-8",
        errors: str = "replace",
    ) -> SearchResults:
        """
        Replace chunk text with full file contents for rows that have ``source``.

        One row per distinct ``source`` (``score`` is the max among chunks from that file).
        Rows without ``source`` are kept unchanged and appended after file rows, still sorted by score.
        """
        with_src = [r for r in self.data if r.source]
        no_src = [r for r in self.data if not r.source]
        best_by_path: dict[str, SearchResult] = {}
        for r in with_src:
            prev = best_by_path.get(r.source)
            if prev is None or r.score > prev.score:
                best_by_path[r.source] = r
        file_rows: list[SearchResult] = []
        for path_str, rep in sorted(best_by_path.items(), key=lambda kv: kv[1].score, reverse=True):
            p = Path(path_str)
            try:
                text = p.read_text(encoding=encoding, errors=errors)
            except OSError:
                text = f"<could not read file: {path_str}>"
            file_rows.append(
                SearchResult(id=rep.id, text=text, score=rep.score, source=rep.source),
            )
        combined = file_rows + sorted(no_src, key=lambda r: r.score, reverse=True)
        return SearchResults(combined)

    def to_dicts(self) -> list[dict[str, Any]]:
        """Each dict has id, text, score (rounded), source — for JSON or legacy callers."""
        return [
            {
                "id": r.id,
                "text": r.text,
                "score": round(r.score, 4),
                "source": r.source,
            }
            for r in self.data
        ]

    def __repr__(self) -> str:
        return f"SearchResults({len(self.data)} items)"


def chunk_file_semantically(
    content: str,
    suffix: str,
    *,
    min_chunk_chars: int = 80,
    max_chunk_chars: int = 2000,
    overlap_chars: int = 250,
) -> list[str]:
    """
    Split file content into semantic chunks for focused retrieval, with overlap at boundaries.

    Logical breaks (in order of preference for each format):
    - .md: ATX headers (#–######), horizontal rules (---), then paragraphs
    - .py: top-level ``def`` / ``class`` boundaries, then paragraphs / lines / sentences inside large blocks
    - .txt / other: paragraphs, then lines, then sentences

    After sizing, each chunk (except the first) is prefixed with a tail of the previous *semantic* chunk
    so embeddings share context at boundaries (see overlap_chars).
    """
    content = content.strip()
    if not content:
        return []

    suffix = (suffix or "").lower().lstrip(".")

    def _greedy_pack_units(units: list[str], sep: str, cap: int) -> list[str]:
        """Pack consecutive units into segments of at most cap characters (best-effort logical packing)."""
        out: list[str] = []
        buf = ""
        for u in units:
            u = u.strip()
            if not u:
                continue
            if not buf:
                buf = u
            elif len(buf) + len(sep) + len(u) <= cap:
                buf = f"{buf}{sep}{u}"
            else:
                out.append(buf)
                buf = u
        if buf:
            out.append(buf)
        return out

    def _merge_small(chunks: list[str]) -> list[str]:
        """Merge chunks smaller than min_chunk_chars with the next chunk."""
        out: list[str] = []
        buf = ""
        for c in chunks:
            c = c.strip()
            if not c:
                continue
            if buf and len(buf) < min_chunk_chars and len(buf) + len(c) + 2 <= max_chunk_chars:
                buf = f"{buf}\n\n{c}"
            else:
                if buf:
                    out.append(buf)
                buf = c
        if buf:
            out.append(buf)
        return out

    def _split_oversized_block(block: str) -> list[str]:
        """Split one block that exceeds max_chunk_chars using progressively finer logical boundaries."""
        if len(block) <= max_chunk_chars:
            return [block]
        paras = [p.strip() for p in re.split(r"\n\s*\n+", block) if p.strip()]
        if len(paras) > 1:
            packed = _greedy_pack_units(paras, "\n\n", max_chunk_chars)
            out: list[str] = []
            for p in packed:
                out.extend(_split_oversized_block(p) if len(p) > max_chunk_chars else [p])
            return out
        one = paras[0] if paras else block
        lines = one.split("\n")
        if len(lines) > 1:
            packed = _greedy_pack_units(lines, "\n", max_chunk_chars)
            if not all(len(p) <= max_chunk_chars for p in packed):
                out = []
                for p in packed:
                    out.extend(_split_oversized_block(p) if len(p) > max_chunk_chars else [p])
                return out
            return packed
        sent_split = re.compile(r"(?<=[.!?。！？])\s+")
        sents = [s.strip() for s in sent_split.split(one) if s.strip()]
        if len(sents) > 1:
            packed = _greedy_pack_units(sents, " ", max_chunk_chars)
            out = []
            for p in packed:
                out.extend(_split_oversized_block(p) if len(p) > max_chunk_chars else [p])
            return out
        # Last resort: hard slices on spaces
        out: list[str] = []
        start = 0
        while start < len(one):
            end = min(start + max_chunk_chars, len(one))
            if end < len(one):
                cut = one.rfind(" ", start + max_chunk_chars // 2, end + 1)
                if cut <= start:
                    cut = end
                end = cut
            piece = one[start:end].strip()
            if piece:
                out.append(piece)
            start = max(start + 1, end)
        return out

    def _split_large(chunks: list[str]) -> list[str]:
        """Split chunks larger than max_chunk_chars using logical sub-structure."""
        out: list[str] = []
        for c in chunks:
            if len(c) <= max_chunk_chars:
                out.append(c)
            else:
                out.extend(_split_oversized_block(c))
        return out

    def _tail_for_overlap(prev: str, n: int) -> str:
        """Last up to n characters of prev, aligned to start after a newline when possible."""
        if n <= 0 or not prev:
            return ""
        prev = prev.strip()
        if len(prev) <= n:
            return prev
        tail = prev[-n:]
        nl = tail.find("\n\n")
        if nl != -1 and nl < len(tail) - 32:
            return tail[nl:].strip()
        nl = tail.find("\n")
        if nl != -1 and nl < len(tail) - 24:
            return tail[nl:].strip()
        return tail.strip()

    def _add_overlap(semantic: list[str], n: int) -> list[str]:
        if n <= 0 or len(semantic) < 2:
            return semantic
        out: list[str] = [semantic[0]]
        for i in range(1, len(semantic)):
            prefix = _tail_for_overlap(semantic[i - 1], n)
            body = semantic[i]
            if not prefix:
                out.append(body)
                continue
            if prefix in body[: len(prefix) + 8]:
                out.append(body)
                continue
            out.append(f"{prefix}\n\n{body}")
        return out

    chunks: list[str] = []

    if suffix == "md":
        parts = re.split(r"(?=^(?:#{1,6}\s+.+|[-*_]{3,}\s*)$)", content, flags=re.MULTILINE)
        for p in parts:
            p = p.strip()
            if p:
                chunks.append(p)
        if not chunks:
            chunks = [p.strip() for p in re.split(r"\n\s*\n+", content) if p.strip()]
    elif suffix == "py":
        parts = re.split(r"(?=^(?:def |class )\w)", content, flags=re.MULTILINE)
        for p in parts:
            p = p.strip()
            if p:
                chunks.append(p)
    else:
        chunks = [p.strip() for p in re.split(r"\n\s*\n+", content) if p.strip()]

    chunks = [c.strip() for c in chunks if c.strip()]
    chunks = _merge_small(chunks)
    chunks = _split_large(chunks)
    chunks = _add_overlap(chunks, overlap_chars)
    result = [c for c in chunks if len(c) >= min_chunk_chars]
    if not result and content.strip():
        return [content.strip()]
    return result


class VectorDB:
    """
    Vector database backed by SQLite with OpenAI-compatible embeddings.
    """

    TABLE = "embeddings"
    SOURCE_FILES_TABLE = "source_files"
    EMBED_BATCH_SIZE = 64
    EMBED_MAX_WORKERS = 4
    DEFAULT_BASE_URL = "http://127.0.0.1:1234/v1"
    DEFAULT_API_KEY = "NO_API_KEY"
    DEFAULT_MODEL = "text-embedding-qwen3-embedding-4b"

    def __init__(
        self,
        path: str | Path = "vectors.db",
        *,
        base_url: str | None = None,
        api_key: str | None = None,
        model: str | None = None,
        query_instruction: str | None = None,
    ) -> None:
        """Initialize VectorDB. path: SQLite file or ':memory:'. base_url/api_key: OpenAI-compatible API; fallback to env vars. model: embedding model name. query_instruction: prepended to queries for instruction-tuned embedders."""
        from openai import OpenAI

        self.path = os.fspath(path)
        self.model = model or self.DEFAULT_MODEL
        self._query_instruction = (
            query_instruction
            if query_instruction is not None
            else (
                "Retrieve passages that are semantically relevant to the search query. "
                "Prioritize meaning and context over exact keyword matches."
            )
        )
        self._client = OpenAI(
            base_url=base_url or os.environ.get("OPENAI_BASE_URL", self.DEFAULT_BASE_URL),
            api_key=api_key or os.environ.get("OPENAI_API_KEY", self.DEFAULT_API_KEY),
        )
        self._conn: sqlite3.Connection | None = None

    def _normalize_query(self, query: str) -> str:
        """Normalize query string: collapse whitespace; return single space if empty."""
        if not query or not query.strip():
            return " "
        return " ".join(query.strip().split())

    def _format_query(self, query: str) -> str:
        """Format query with instruction prefix for embedding API."""
        normalized = self._normalize_query(query)
        return f"Instruct: {self._query_instruction}\nQuery: {normalized}"

    def _connection(self) -> sqlite3.Connection:
        """Lazily create and return the SQLite connection; register functions and ensure schema."""
        if self._conn is None:
            self._conn = sqlite3.connect(self.path)
            self._conn.row_factory = sqlite3.Row
            self._register_functions()
            self._ensure_schema()
        return self._conn

    def _register_functions(self) -> None:
        def cosine_similarity_blob(blob_a: bytes | None, blob_b: bytes | None) -> float:
            if blob_a is None or blob_b is None:
                return 0.0
            a = blob_to_embedding(blob_a)
            b = blob_to_embedding(blob_b)
            return cosine_similarity(a, b)

        self._conn.create_function(
            "cosine_similarity",
            2,
            cosine_similarity_blob,
            deterministic=True,
        )

    def _ensure_schema(self) -> None:
        """Create embeddings table if missing; add source and source_mtime columns if absent."""
        self._conn.execute(
            f"""
            CREATE TABLE IF NOT EXISTS {self.TABLE} (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                text TEXT NOT NULL,
                embedding BLOB NOT NULL,
                source TEXT,
                source_mtime REAL,
                created_at TEXT DEFAULT (datetime('now'))
            )
            """
        )
        cur = self._conn.execute(f"SELECT name FROM pragma_table_info('{self.TABLE}') WHERE name = 'source'")
        if cur.fetchone() is None:
            self._conn.execute(f"ALTER TABLE {self.TABLE} ADD COLUMN source TEXT")
        cur = self._conn.execute(
            f"SELECT name FROM pragma_table_info('{self.TABLE}') WHERE name = 'source_mtime'"
        )
        if cur.fetchone() is None:
            self._conn.execute(f"ALTER TABLE {self.TABLE} ADD COLUMN source_mtime REAL")
        self._conn.execute(
            f"""
            CREATE TABLE IF NOT EXISTS {self.SOURCE_FILES_TABLE} (
                path TEXT PRIMARY KEY,
                mtime REAL NOT NULL,
                indexed_at TEXT DEFAULT (datetime('now'))
            )
            """
        )
        self._conn.execute(
            f"""
            INSERT OR IGNORE INTO {self.SOURCE_FILES_TABLE} (path, mtime)
            SELECT source, MAX(source_mtime)
            FROM {self.TABLE}
            WHERE source IS NOT NULL AND source_mtime IS NOT NULL
            GROUP BY source
            """
        )
        self._conn.commit()

    def embed(self, text: str) -> list[float]:
        """Create an embedding for a single string via OpenAI."""
        resp = self._client.embeddings.create(
            model=self.model,
            input=text.strip() or " ",
        )
        return resp.data[0].embedding

    def embed_many(self, texts: list[str]) -> list[list[float]]:
        """Create embeddings for multiple strings in one API call."""
        stripped = [t.strip() or " " for t in texts]
        resp = self._client.embeddings.create(model=self.model, input=stripped)
        by_index = {d.index: d.embedding for d in resp.data}
        return [by_index[i] for i in range(len(texts))]

    def _embed_many_batches_concurrent(self, texts: list[str]) -> list[list[float]]:
        """
        Embed many texts using concurrent API calls (one embed_many per batch).
        DB access must happen on a single thread after this returns.
        """
        if not texts:
            return []
        batch_size = self.EMBED_BATCH_SIZE
        batches = [texts[i : i + batch_size] for i in range(0, len(texts), batch_size)]
        if len(batches) == 1:
            return self.embed_many(batches[0])
        max_workers = min(self.EMBED_MAX_WORKERS, len(batches))

        def _run_batch(idx: int, batch: list[str]) -> tuple[int, list[list[float]]]:
            return idx, self.embed_many(batch)

        results: list[list[list[float]] | None] = [None] * len(batches)
        with ThreadPoolExecutor(max_workers=max_workers) as pool:
            futures = {pool.submit(_run_batch, i, b): i for i, b in enumerate(batches)}
            for fut in as_completed(futures):
                idx, vectors = fut.result()
                results[idx] = vectors
        flat: list[list[float]] = []
        for part in results:
            if part:
                flat.extend(part)
        return flat

    def _get_stored_file_mtime(self, path_posix: str) -> float | None:
        row = self._connection().execute(
            f"SELECT mtime FROM {self.SOURCE_FILES_TABLE} WHERE path = ?",
            (path_posix,),
        ).fetchone()
        if row is None:
            return None
        return float(row[0])

    def _set_source_file_indexed(self, path_posix: str, mtime: float) -> None:
        conn = self._connection()
        with conn:
            conn.execute(
                f"""
                INSERT INTO {self.SOURCE_FILES_TABLE} (path, mtime, indexed_at)
                VALUES (?, ?, datetime('now'))
                ON CONFLICT(path) DO UPDATE SET
                    mtime = excluded.mtime,
                    indexed_at = excluded.indexed_at
                """,
                (path_posix, mtime),
            )

    def add(
        self,
        text: str,
        source: str | None = None,
        source_mtime: float | None = None,
    ) -> int:
        """Embed and store one text; returns row id."""
        embedding = self.embed(text)
        blob = embedding_to_blob(embedding)
        conn = self._connection()
        with conn:
            cur = conn.execute(
                f"INSERT INTO {self.TABLE} (text, embedding, source, source_mtime) VALUES (?, ?, ?, ?)",
                (text, blob, source, source_mtime),
            )
            return cur.lastrowid or 0

    def add_many(
        self,
        texts: list[str],
        sources: list[str] | None = None,
        source_mtimes: list[float] | None = None,
    ) -> list[int]:
        """Embed and store multiple texts; returns row ids."""
        if not texts:
            return []
        if sources is not None and len(sources) != len(texts):
            raise ValueError("sources length must match texts")
        if source_mtimes is not None and len(source_mtimes) != len(texts):
            raise ValueError("source_mtimes length must match texts")
        embeddings = self.embed_many(texts)
        conn = self._connection()
        placeholders = ", ".join(["(?, ?, ?, ?)"] * len(texts))
        params: list[object] = []
        for i, (t, emb) in enumerate(zip(texts, embeddings, strict=True)):
            params.append(t)
            params.append(embedding_to_blob(emb))
            params.append(sources[i] if sources is not None else None)
            params.append(source_mtimes[i] if source_mtimes is not None else None)
        with conn:
            cur = conn.execute(
                f"INSERT INTO {self.TABLE} (text, embedding, source, source_mtime) VALUES {placeholders} RETURNING id",
                params,
            )
            return [row[0] for row in cur.fetchall()]

    def add_many_with_embeddings(
        self,
        texts: list[str],
        embeddings: list[list[float]],
        sources: list[str] | None = None,
        source_mtimes: list[float] | None = None,
    ) -> list[int]:
        """Store rows using precomputed embeddings (no API calls). Returns row ids."""
        if not texts or len(texts) != len(embeddings):
            raise ValueError("texts and embeddings must be non-empty and the same length")
        if sources is not None and len(sources) != len(texts):
            raise ValueError("sources length must match texts")
        if source_mtimes is not None and len(source_mtimes) != len(texts):
            raise ValueError("source_mtimes length must match texts")
        conn = self._connection()
        placeholders = ", ".join(["(?, ?, ?, ?)"] * len(texts))
        params: list[object] = []
        for i, (t, emb) in enumerate(zip(texts, embeddings, strict=True)):
            params.append(t)
            params.append(embedding_to_blob(emb))
            params.append(sources[i] if sources is not None else None)
            params.append(source_mtimes[i] if source_mtimes is not None else None)
        with conn:
            cur = conn.execute(
                f"INSERT INTO {self.TABLE} (text, embedding, source, source_mtime) VALUES {placeholders} RETURNING id",
                params,
            )
            return [row[0] for row in cur.fetchall()]

    def get_source_mtime(self, source: str) -> float | None:
        """Return the stored mtime for this source from source_files, or None."""
        return self._get_stored_file_mtime(source)

    def delete_by_source(self, source: str) -> int:
        """Delete all rows with the given source path and its source_files row. Returns embedding rows deleted."""
        conn = self._connection()
        with conn:
            conn.execute(
                f"DELETE FROM {self.SOURCE_FILES_TABLE} WHERE path = ?",
                (source,),
            )
            cur = conn.execute(
                f"DELETE FROM {self.TABLE} WHERE source = ?",
                (source,),
            )
            return cur.rowcount

    def list_sources(self) -> set[str]:
        """Return source paths that have embedding rows or a source_files row (e.g. empty chunk files)."""
        conn = self._connection()
        from_emb = {
            r[0]
            for r in conn.execute(
                f"SELECT DISTINCT source FROM {self.TABLE} WHERE source IS NOT NULL"
            ).fetchall()
        }
        from_sf = {r[0] for r in conn.execute(f"SELECT path FROM {self.SOURCE_FILES_TABLE}").fetchall()}
        return from_emb | from_sf

    def index_file(self, path: Path, *, encoding: str = "utf-8", force: bool = False) -> int | None:
        """
        (Re)index a single file if its mtime changed (tracked in source_files).
        Skips the file entirely when mtime matches and force is False.
        On change: removes existing embeddings for the file, chunks semantically, embeds in concurrent batches,
        inserts all chunks in one transaction, then updates source_files.

        Returns number of chunks indexed, or None on OS/read error; 0 if skipped (unchanged).
        """
        source = path.resolve().as_posix()
        try:
            mtime = path.stat().st_mtime
        except OSError:
            return None
        if not force and self._get_stored_file_mtime(source) == mtime:
            return 0

        self.delete_by_source(source)
        text = path.read_text(encoding=encoding, errors="replace")
        chunks = chunk_file_semantically(text, path.suffix)
        if not chunks:
            self._set_source_file_indexed(source, mtime)
            return 0

        vectors = self._embed_many_batches_concurrent(chunks)
        sources = [source] * len(chunks)
        mtimes = [mtime] * len(chunks)
        self.add_many_with_embeddings(chunks, vectors, sources=sources, source_mtimes=mtimes)
        self._set_source_file_indexed(source, mtime)
        return len(chunks)

    def search(
        self,
        query: str,
        k: int = 5,
        *,
        min_score: float | None = None,
        file_contents: bool = False,
        encoding: str = "utf-8",
    ) -> SearchResults:
        """
        Embed the query and return the k nearest rows by cosine similarity.

        ``min_score``: drop hits with score strictly below this (cosine similarity in [0, 1]).
        ``file_contents``: after optional filtering, merge to one row per ``source`` with full file text.
        """
        query_embedding = self.embed(self._format_query(query))
        query_blob = embedding_to_blob(query_embedding)
        conn = self._connection()
        rows = conn.execute(
            f"""
            SELECT id, text,
                   cosine_similarity(embedding, ?) AS score,
                   source
            FROM {self.TABLE}
            ORDER BY score DESC
            LIMIT ?
            """,
            (query_blob, k),
        ).fetchall()
        results = SearchResults(
            [
                SearchResult(
                    id=r["id"],
                    text=r["text"],
                    score=r["score"],
                    source=r["source"],
                )
                for r in rows
            ]
        )
        if min_score is not None:
            results = results.threshold(min_score)
        if file_contents:
            results = results.with_file_contents(encoding=encoding)
        return results

    def sync_from_roots(
        self,
        roots: str | Path | list[str] | list[Path],
        *,
        glob_pattern: str | list[str] = "*.md",
        encoding: str = "utf-8",
    ) -> tuple[int, int]:
        """Cleanup and populate from folder roots. Returns (deleted_count, indexed_count)."""
        path_list = normalize_roots(roots)
        current_files = collect_file_paths(path_list, glob_pattern)
        db_sources = self.list_sources()
        to_remove = db_sources - current_files
        deleted = sum(self.delete_by_source(s) for s in to_remove)
        indexed = 0
        for posix_path in current_files:
            n = self.index_file(Path(posix_path), encoding=encoding)
            if n is not None and n > 0:
                indexed += 1
        return deleted, indexed

    def close(self) -> None:
        """Close the database connection."""
        if self._conn is not None:
            self._conn.close()
            self._conn = None

    def __enter__(self) -> VectorDB:
        self._connection()
        return self

    def __exit__(self, *args: object) -> None:
        self.close()


def create_rag_db(
    roots: str | Path | list[str] | list[Path],
    *,
    path: str | Path = ":memory:",
    glob_pattern: str | list[str] = "*.md",
    base_url: str | None = None,
    api_key: str | None = None,
    model: str | None = None,
    query_instruction: str | None = None,
    encoding: str = "utf-8",
) -> VectorDB:
    """Create a VectorDB from files under the given roots, then sync."""
    db = VectorDB(
        path,
        base_url=base_url,
        api_key=api_key,
        model=model,
        query_instruction=query_instruction,
    )
    db.sync_from_roots(roots, glob_pattern=glob_pattern, encoding=encoding)
    return db


def run_rag_query(
    query: str,
    db: VectorDB,
    k: int = 5,
    *,
    min_score: float | None = None,
    file_contents: bool = False,
    encoding: str = "utf-8",
) -> list[dict[str, Any]]:
    """Search the RAG database for the top-k passages; return list of dicts with id, text, score, source."""
    return db.search(
        query,
        k=k,
        min_score=min_score,
        file_contents=file_contents,
        encoding=encoding,
    ).to_dicts()

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--roots",
        nargs="+",
        default=["./"],
        metavar="PATH",
        help="One or more root paths (files or directories) to index. Default: ./",
    )
    parser.add_argument("--glob_pattern", type=str, default="*.md")
    parser.add_argument("--db_path", type=str, default="database.db")
    parser.add_argument(
        "--min_score",
        type=float,
        default=None,
        help="Minimum cosine similarity (0–1); omit for no cutoff.",
    )
    parser.add_argument(
        "--file_contents",
        action="store_true",
        help="Return full file contents per distinct source instead of chunk text.",
    )

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--query", type=str, help="Single query to run (default mode).")
    group.add_argument("--interactive", action="store_true", help="Interactive query loop.")

    args = parser.parse_args()
    db = create_rag_db(args.roots, glob_pattern=args.glob_pattern, path=args.db_path)
    try:
        if args.interactive:
            print("Interactive query mode. Type an empty line to exit.")
            while True:
                try:
                    q = input("Query> ").strip()
                except (EOFError, KeyboardInterrupt):
                    print()
                    break
                if not q:
                    break
                results = run_rag_query(
                    q,
                    db,
                    min_score=args.min_score,
                    file_contents=args.file_contents,
                )
                print(results)
        else:
            results = run_rag_query(
                args.query,
                db,
                min_score=args.min_score,
                file_contents=args.file_contents,
            )
            if args.file_contents:
                print(results)
            else:
                for item in results:
                    print(item["text"])
    finally:
        db.close()