From 6e7558f1e094f652d9ba0333c98511f92f1c3877 Mon Sep 17 00:00:00 2001 From: mmprotest Date: Thu, 2 Oct 2025 19:40:08 +1000 Subject: [PATCH] docs: expand README guides --- .github/workflows/ci.yml | 36 +++ .gitignore | 17 ++ .pre-commit-config.yaml | 22 ++ LICENSE | 21 ++ README.md | 198 +++++++++++++- pyproject.toml | 67 +++++ raglite_sqlite/__init__.py | 5 + raglite_sqlite/adapters/__init__.py | 5 + raglite_sqlite/adapters/langchain.py | 45 ++++ raglite_sqlite/adapters/llamaindex.py | 33 +++ raglite_sqlite/api.py | 210 +++++++++++++++ raglite_sqlite/chunking.py | 73 +++++ raglite_sqlite/cli.py | 121 +++++++++ raglite_sqlite/db.py | 250 ++++++++++++++++++ raglite_sqlite/embeddings/__init__.py | 5 + raglite_sqlite/embeddings/base.py | 33 +++ raglite_sqlite/embeddings/openai_backend.py | 30 +++ .../sentence_transformers_backend.py | 26 ++ raglite_sqlite/parsers/__init__.py | 17 ++ raglite_sqlite/parsers/base.py | 12 + raglite_sqlite/parsers/csv.py | 27 ++ raglite_sqlite/parsers/docx.py | 20 ++ raglite_sqlite/parsers/html.py | 40 +++ raglite_sqlite/parsers/md.py | 39 +++ raglite_sqlite/parsers/pdf.py | 20 ++ raglite_sqlite/parsers/txt.py | 14 + raglite_sqlite/rerank.py | 15 ++ raglite_sqlite/schema.sql | 58 ++++ raglite_sqlite/search.py | 140 ++++++++++ raglite_sqlite/typing.py | 36 +++ raglite_sqlite/utils.py | 99 +++++++ tests/conftest.py | 23 ++ tests/data/sample.csv | 3 + tests/data/sample.html | 8 + tests/data/sample.md | 11 + tests/data/sample.txt | 2 + tests/test_cli.py | 44 +++ tests/test_ingest_and_search.py | 36 +++ 38 files changed, 1859 insertions(+), 2 deletions(-) create mode 100644 .github/workflows/ci.yml create mode 100644 .gitignore create mode 100644 .pre-commit-config.yaml create mode 100644 LICENSE create mode 100644 pyproject.toml create mode 100644 raglite_sqlite/__init__.py create mode 100644 raglite_sqlite/adapters/__init__.py create mode 100644 raglite_sqlite/adapters/langchain.py create mode 100644 raglite_sqlite/adapters/llamaindex.py create mode 100644 raglite_sqlite/api.py create mode 100644 raglite_sqlite/chunking.py create mode 100644 raglite_sqlite/cli.py create mode 100644 raglite_sqlite/db.py create mode 100644 raglite_sqlite/embeddings/__init__.py create mode 100644 raglite_sqlite/embeddings/base.py create mode 100644 raglite_sqlite/embeddings/openai_backend.py create mode 100644 raglite_sqlite/embeddings/sentence_transformers_backend.py create mode 100644 raglite_sqlite/parsers/__init__.py create mode 100644 raglite_sqlite/parsers/base.py create mode 100644 raglite_sqlite/parsers/csv.py create mode 100644 raglite_sqlite/parsers/docx.py create mode 100644 raglite_sqlite/parsers/html.py create mode 100644 raglite_sqlite/parsers/md.py create mode 100644 raglite_sqlite/parsers/pdf.py create mode 100644 raglite_sqlite/parsers/txt.py create mode 100644 raglite_sqlite/rerank.py create mode 100644 raglite_sqlite/schema.sql create mode 100644 raglite_sqlite/search.py create mode 100644 raglite_sqlite/typing.py create mode 100644 raglite_sqlite/utils.py create mode 100644 tests/conftest.py create mode 100644 tests/data/sample.csv create mode 100644 tests/data/sample.html create mode 100644 tests/data/sample.md create mode 100644 tests/data/sample.txt create mode 100644 tests/test_cli.py create mode 100644 tests/test_ingest_and_search.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..f729c6a --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,36 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + +jobs: + build: + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [ubuntu-latest, macos-latest, windows-latest] + python-version: ['3.10', '3.11', '3.12'] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Cache pip + uses: actions/cache@v4 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-${{ hashFiles('pyproject.toml') }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -e .[dev] + - name: Lint + run: | + ruff check . + black --check . + - name: Type check + run: mypy raglite_sqlite + - name: Test + run: pytest --cov=raglite_sqlite diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ef55dd0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,17 @@ +__pycache__/ +*.py[cod] +*.egg-info/ +.venv/ +.env +coverage.xml +htmlcov/ +.pytest_cache/ +.mypy_cache/ +.cache/ +.DS_Store +*.db +*.sqlite +*.log +/.ruff_cache +/.pytest_cache +/.coverage diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..f7f0c41 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,22 @@ +repos: + - repo: https://github.com/psf/black + rev: 24.4.2 + hooks: + - id: black + args: ["--line-length", "100"] + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.4.8 + hooks: + - id: ruff + args: ["--fix", "--exit-non-zero-on-fix"] + - id: ruff-format + - repo: https://github.com/pycqa/isort + rev: 5.13.2 + hooks: + - id: isort + args: ["--profile", "black", "--line-length", "100"] + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.6.0 + hooks: + - id: end-of-file-fixer + - id: trailing-whitespace diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..3ff57d3 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 RagLite Contributors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md index f5ebabb..4e06874 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,196 @@ -# raglite-sqlite -Zero-server, zero-docker RAG with SQLite +# RagLite SQLite + +Local-first Retrieval-Augmented Generation (RAG) toolkit built entirely on top of a single SQLite file. RagLite makes it easy to ingest documents, chunk them, embed with pluggable backends, and perform hybrid lexical/vector search—without running servers or Docker images. + +## Why RagLite? + +- **Zero infrastructure** – everything lives inside one SQLite database (`knowledge.db`). +- **Deterministic and offline** – default embedding model is local; no network calls unless explicitly configured. +- **Hybrid retrieval** – combines BM25 via FTS5 with cosine similarity over stored vectors. +- **Python and CLI** – flexible API plus a friendly Typer-based CLI for scripting. +- **Extensible** – pluggable parsers, chunkers, embedding backends, and adapters for LangChain / LlamaIndex. + +## Installation + +```bash +python -m venv .venv && source .venv/bin/activate # (Windows: .venv\Scripts\activate) +pip install -e . +``` + +## Quickstart + +```bash +python -m venv .venv && source .venv/bin/activate # (Windows: .venv\Scripts\activate) +pip install -e . +raglite init --db knowledge.db +raglite index tests/data --db knowledge.db --model sentence-transformers/all-MiniLM-L6-v2 +raglite query "example text" --db knowledge.db -k 8 --hybrid 0.6 +``` + +## CLI Examples + +- Initialize a new database: + ```bash + raglite init --db knowledge.db + ``` +- Index a directory recursively with tags and progress reporting: + ```bash + raglite index docs --db knowledge.db --tags project,internal --recursive --chunk-size 400 --overlap 50 + ``` +- Re-index only changed files by disabling recursion and pointing at a single file: + ```bash + raglite index docs/faq.md --db knowledge.db --no-recursive --skip-unchanged + ``` +- Query for relevant chunks: + ```bash + raglite query "How do I deploy?" --db knowledge.db -k 5 --hybrid 0.7 + ``` +- Filter by tag or doc id: + ```bash + raglite query "release checklist" --db knowledge.db --filter tag=internal --filter doc_id=release-notes + ``` +- Inspect statistics: + ```bash + raglite stats --db knowledge.db + ``` +- Export metadata: + ```bash + raglite export --db knowledge.db --to export.ndjson --include-vectors + ``` +- Vacuum the database to reclaim space: + ```bash + raglite vacuum --db knowledge.db + ``` + +## Python API + +The high-level API is built around `raglite_sqlite.api.RagLite`. Below is a minimal end-to-end script that indexes a directory, +queries it, and prints rich metadata for each hit: + +```python +from raglite_sqlite import RagLite +from raglite_sqlite.embeddings.sentence_transformers_backend import SentenceTransformersBackend + +rag = RagLite("knowledge.db") +backend = SentenceTransformersBackend(model_name="sentence-transformers/all-MiniLM-L6-v2") + +rag.index( + paths=["tests/data"], + tags="demo", + embedding_backend=backend, + chunker="recursive", + chunk_size_tokens=384, + chunk_overlap_tokens=48, +) + +results = rag.search( + query="example text", + embedding_backend=backend, + k=5, + hybrid_weight=0.6, + max_per_doc=2, +) + +for item in results: + print(f"{item['score']:.3f} | {item['doc_id']} | {item['section'] or 'No section'}") + print(item["snippet"]) + print("-" * 40) +``` + +### Advanced ingestion tips + +- **Custom parsers and options** – pass `parser_opts` to `RagLite.index` to control parser behaviour (e.g. CSV column filters). +- **Chunking strategies** – choose between `fixed_tokens` (uniform splits) and `recursive` (paragraph-aware) chunkers. +- **Embedding reuse** – keep `skip_unchanged=True` (default) to leverage hashing + cache table for faster re-index runs. +- **Multiple models** – provide different `model_name` values per run; embeddings are cached per `(sha256, model)` pair. +- **Export/backup** – `raglite export` produces NDJSON that can be restored or version-controlled for auditing. + +### Operating the database + +- SQLite is safe to sync via file-sharing tools (Dropbox, Syncthing) when only one process writes at a time. +- Use `raglite vacuum` periodically after large deletions to compact the file. +- WAL mode is enabled automatically; consider copying the `.db` + `-wal` file pair while the app is running. +- For cloud backups, store the DB file and optionally NDJSON exports in object storage. + +## Ecosystem Integrations + +RagLite ships with lightweight adapters so you can plug the SQLite-backed retriever into popular orchestration frameworks. + +### LangChain + +```python +from langchain.chains import RetrievalQA +from langchain.llms import OpenAI + +from raglite_sqlite import RagLite +from raglite_sqlite.adapters.langchain import RagLiteRetriever +from raglite_sqlite.embeddings.sentence_transformers_backend import SentenceTransformersBackend + +rag = RagLite("knowledge.db") +backend = SentenceTransformersBackend() + +# Ensure documents are indexed before wiring up the retriever. +rag.index(["tests/data"], embedding_backend=backend) + +retriever = RagLiteRetriever(rag, embedding_backend=backend, k=6, hybrid_weight=0.5) + +qa_chain = RetrievalQA.from_chain_type( + llm=OpenAI(), + retriever=retriever, + chain_type="stuff", +) + +response = qa_chain.run("Summarize the sample docs") +print(response) +``` + +The adapter defers LangChain imports until used, keeping the core package lightweight. You can also supply filters (`tag=...`) +via the retriever call (`retriever.get_relevant_documents(query, filters={"tags": "internal"})`). + +### LlamaIndex + +```python +from llama_index.core import ServiceContext, VectorStoreIndex +from llama_index.llms import OpenAI + +from raglite_sqlite import RagLite +from raglite_sqlite.adapters.llamaindex import RagLiteVectorStore +from raglite_sqlite.embeddings.sentence_transformers_backend import SentenceTransformersBackend + +rag = RagLite("knowledge.db") +backend = SentenceTransformersBackend() + +rag.index(["tests/data"], embedding_backend=backend) + +service_context = ServiceContext.from_defaults(llm=OpenAI()) +vector_store = RagLiteVectorStore(rag, embedding_backend=backend) + +index = VectorStoreIndex.from_vector_store(vector_store, service_context=service_context) +query_engine = index.as_query_engine(similarity_top_k=5) + +answer = query_engine.query("What is contained in the sample docs?") +print(answer) +``` + +The vector store wrapper exposes RagLite search semantics to LlamaIndex while leaving indexing/ingest under your control. + +## Design Notes + +- Uses SQLite with WAL mode and FTS5 for zero-config deployment. +- Chunker strategies allow fixed token or recursive splitting with overlaps. +- Embeddings cached via SHA-256 to avoid redundant model calls. +- Hybrid retrieval fuses normalized BM25 and cosine similarity scores. + +## Security & Privacy + +RagLite does not perform any network calls by default. Remote embedding backends (such as OpenAI) are opt-in and require explicit configuration via environment variables. + +## Roadmap + +- Optional REST server for multi-user access. +- Support for additional document formats and OCR. +- Pluggable reranking models. + +## License + +MIT License. See [LICENSE](LICENSE). diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..7b55292 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,67 @@ +[build-system] +requires = ["setuptools>=64", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "raglite-sqlite" +version = "0.1.0" +description = "Local-first RAG toolkit backed by a single SQLite database" +readme = "README.md" +authors = [{ name = "RagLite Contributors" }] +license = { file = "LICENSE" } +requires-python = ">=3.10" +dependencies = [ + "typer[all]>=0.9.0", + "rich>=13.7.0", + "pydantic>=2.7.0", + "beautifulsoup4>=4.12.0", + "python-frontmatter>=1.0.0", + "PyPDF2>=3.0.0", + "python-docx>=1.0.0", + "numpy>=1.24.0", + "tqdm>=4.66.0", + "rapidfuzz>=3.6.0", + "sentence-transformers>=2.5.0", + "orjson>=3.9.0", + "click-spinner>=0.1.10", + "sqlite-fts5", +] + +[project.optional-dependencies] +openai = ["openai>=1.0.0"] +dev = [ + "pytest>=8.1.0", + "pytest-cov>=4.1.0", + "mypy>=1.10.0", + "ruff>=0.4.0", + "black>=24.4.0", + "build>=1.2.1", + "twine>=5.0.0", +] + +[project.urls] +Homepage = "https://github.com/example/raglite-sqlite" +Repository = "https://github.com/example/raglite-sqlite" + +[project.scripts] +raglite = "raglite_sqlite.cli:app" + +[tool.black] +line-length = 100 +target-version = ["py310"] + +[tool.ruff] +line-length = 100 +select = ["E", "F", "I", "B"] + +[tool.ruff.lint.isort] +known-first-party = ["raglite_sqlite"] + +[tool.mypy] +pytorch = false +python_version = "3.10" +warn_unused_ignores = true +warn_redundant_casts = true +warn_unused_configs = true +strict = true +plugins = [] diff --git a/raglite_sqlite/__init__.py b/raglite_sqlite/__init__.py new file mode 100644 index 0000000..42882bb --- /dev/null +++ b/raglite_sqlite/__init__.py @@ -0,0 +1,5 @@ +"""RagLite SQLite package.""" + +from .api import RagLite + +__all__ = ["RagLite"] diff --git a/raglite_sqlite/adapters/__init__.py b/raglite_sqlite/adapters/__init__.py new file mode 100644 index 0000000..d3dbdcd --- /dev/null +++ b/raglite_sqlite/adapters/__init__.py @@ -0,0 +1,5 @@ +"""Adapters for third-party ecosystems.""" + +from .langchain import RagLiteRetriever + +__all__ = ["RagLiteRetriever"] diff --git a/raglite_sqlite/adapters/langchain.py b/raglite_sqlite/adapters/langchain.py new file mode 100644 index 0000000..66f5f11 --- /dev/null +++ b/raglite_sqlite/adapters/langchain.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from typing import Any, Dict, Iterable, List, Optional + +try: + from langchain.schema import Document +except Exception: # pragma: no cover - optional dependency + Document = None # type: ignore[assignment] + +from ..api import RagLite +from ..embeddings.base import EmbeddingBackend +from ..typing import SearchResult + + +class RagLiteRetriever: + """Minimal LangChain retriever wrapper.""" + + def __init__( + self, + rag: RagLite, + *, + backend: Optional[EmbeddingBackend] = None, + k: int = 4, + ) -> None: + self.rag = rag + self.backend = backend + self.k = k + + def get_relevant_documents(self, query: str, **kwargs: Any) -> List[Any]: + results = self.rag.search(query, k=self.k, embedding_backend=self.backend) + if Document is None: + raise RuntimeError("LangChain is not installed") + docs: List[Any] = [] + for item in results: + metadata: Dict[str, Any] = { + "doc_id": item.get("doc_id"), + "section": item.get("section"), + "source_path": item.get("source_path"), + "tags": item.get("tags"), + } + docs.append(Document(page_content=item.get("text", ""), metadata=metadata)) + return docs + + async def aget_relevant_documents(self, query: str, **kwargs: Any) -> List[Any]: + return self.get_relevant_documents(query, **kwargs) diff --git a/raglite_sqlite/adapters/llamaindex.py b/raglite_sqlite/adapters/llamaindex.py new file mode 100644 index 0000000..b17ef47 --- /dev/null +++ b/raglite_sqlite/adapters/llamaindex.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +from ..api import RagLite +from ..embeddings.base import EmbeddingBackend + + +class RagLiteNodeRetriever: + """Minimal adapter compatible with LlamaIndex retriever interface.""" + + def __init__(self, rag: RagLite, *, backend: Optional[EmbeddingBackend] = None, k: int = 4) -> None: + self.rag = rag + self.backend = backend + self.k = k + + def retrieve(self, query: str, **kwargs: Any) -> List[Dict[str, Any]]: + results = self.rag.search(query, k=self.k, embedding_backend=self.backend) + nodes: List[Dict[str, Any]] = [] + for item in results: + nodes.append( + { + "text": item.get("text", ""), + "id": item.get("chunk_id"), + "metadata": { + "doc_id": item.get("doc_id"), + "section": item.get("section"), + "source_path": item.get("source_path"), + "tags": item.get("tags"), + }, + } + ) + return nodes diff --git a/raglite_sqlite/api.py b/raglite_sqlite/api.py new file mode 100644 index 0000000..f9ca192 --- /dev/null +++ b/raglite_sqlite/api.py @@ -0,0 +1,210 @@ +from __future__ import annotations + +import hashlib +from pathlib import Path +from typing import Any, Dict, Iterable, List, Sequence + +from .chunking import chunk_blocks +from .db import Database +from .embeddings.base import EmbeddingBackend +from .parsers.csv import CSVParser +from .parsers.docx import DocxParser +from .parsers.html import HTMLParser +from .parsers.md import MarkdownParser +from .parsers.pdf import PDFParser +from .parsers.txt import TextParser +from .search import assemble_results, bm25_search, hybrid_fuse, vector_search +from .typing import SearchResult +from .utils import detect_mime, iter_files, loads_json, now_ts, normalize_text, sha256_file, sha256_text + +PARSER_REGISTRY = { + "text/plain": TextParser(), + "text/markdown": MarkdownParser(), + "text/html": HTMLParser(), + "application/pdf": PDFParser(), + "application/vnd.openxmlformats-officedocument.wordprocessingml.document": DocxParser(), + "text/csv": CSVParser(), +} + + +def _default_backend() -> EmbeddingBackend: + from .embeddings.sentence_transformers_backend import SentenceTransformersBackend + + return SentenceTransformersBackend() + + +class RagLite: + def __init__(self, db_path: str, create: bool = True) -> None: + self.db_path = Path(db_path) + self.db = Database(self.db_path, create=create) + + def _doc_id(self, path: Path) -> str: + return hashlib.sha1(str(path.resolve()).encode("utf-8")).hexdigest() + + def _select_parser(self, path: Path) -> Any: + mime = detect_mime(path) + return PARSER_REGISTRY.get(mime, TextParser()) + + def index( + self, + paths: Sequence[str] | str, + *, + tags: str | None = None, + parser_opts: dict | None = None, + chunker: str = "recursive", + chunk_size_tokens: int = 512, + chunk_overlap_tokens: int = 64, + embedding_backend: EmbeddingBackend | None = None, + model_name: str | None = None, + skip_unchanged: bool = True, + recurse: bool = True, + glob: str | None = None, + ) -> dict[str, int]: + if isinstance(paths, str): + input_paths = [paths] + else: + input_paths = list(paths) + files = iter_files(input_paths, recurse=recurse, glob=glob) + backend = embedding_backend or _default_backend() + stats = {"files": 0, "chunks": 0, "skipped": 0} + for file_path in files: + stats["files"] += 1 + doc_id = self._doc_id(file_path) + existing = self.db.get_document(doc_id) + doc_sha = sha256_file(file_path) + if skip_unchanged and existing and existing["sha256"] == doc_sha: + stats["skipped"] += 1 + continue + parser = self._select_parser(file_path) + blocks = parser.parse(str(file_path), **(parser_opts or {})) + if isinstance(blocks, Iterable): + blocks_list = list(blocks) + else: + blocks_list = [blocks] + chunks = chunk_blocks(blocks_list, strategy=chunker, size=chunk_size_tokens, overlap=chunk_overlap_tokens) + created_at = now_ts() + self.db.upsert_document( + doc_id, + source_path=str(file_path), + mime=detect_mime(file_path), + tags=tags, + created_at=created_at, + updated_at=created_at, + sha256=doc_sha, + ) + chunk_hashes: list[str] = [] + for idx, chunk in enumerate(chunks): + chunk_id = f"{doc_id}:{idx}" + chunk_hash = sha256_text(chunk.text_norm) + chunk_hashes.append(chunk_hash) + self.db.upsert_chunk( + chunk_id, + doc_id, + idx, + chunk.text, + chunk.text_norm, + chunk.section, + chunk_hash, + ) + stats["chunks"] += len(chunks) + texts = [chunk.text_norm for chunk in chunks] + if texts: + embeddings = self._embed_chunks( + texts, + hashes=chunk_hashes, + backend=backend, + model_name=model_name or getattr(backend, "model_name", None), + ) + for idx, vector in enumerate(embeddings): + chunk_id = f"{doc_id}:{idx}" + self.db.upsert_vector(chunk_id, vector, model_name or getattr(backend, "model_name", "unknown")) + self.db.commit() + return stats + + def _embed_chunks( + self, + texts: Sequence[str], + *, + hashes: Sequence[str], + backend: EmbeddingBackend, + model_name: str | None, + ) -> list[list[float]]: + cached_vectors: list[list[float] | None] = [] + missing_texts: list[str] = [] + missing_indices: list[int] = [] + model = model_name or getattr(backend, "model_name", "default") + for idx, (text, content_sha) in enumerate(zip(texts, hashes)): + cached = self.db.get_embedding_cache(content_sha, model) + if cached is not None: + cached_vectors.append(list(cached)) + else: + cached_vectors.append(None) + missing_texts.append(text) + missing_indices.append(idx) + if missing_texts: + new_vectors = backend.embed_texts(missing_texts, model_name=model_name) + for offset, idx in enumerate(missing_indices): + vector = list(new_vectors[offset]) + cached_vectors[idx] = vector + self.db.upsert_embedding_cache(hashes[idx], model, vector) + return [vec for vec in cached_vectors if vec is not None] + + def search( + self, + query: str, + *, + k: int = 8, + hybrid_weight: float = 0.6, + filters: dict | None = None, + model_name: str | None = None, + embedding_backend: EmbeddingBackend | None = None, + max_per_doc: int = 3, + with_snippets: bool = True, + ) -> List[SearchResult]: + norm_query = normalize_text(query) + lexical = bm25_search(self.db, norm_query, k, filters) + semantic: list[tuple[str, float]] = [] + backend = embedding_backend + if backend is not None: + semantic = vector_search(self.db, backend, norm_query, model_name, k) + fused = hybrid_fuse(lexical, semantic, hybrid_weight, k) + return assemble_results(self.db, fused, k, with_snippets=with_snippets, max_per_doc=max_per_doc) + + def delete(self, doc_id: str) -> None: + self.db.delete_document(doc_id) + self.db.commit() + + def stats(self) -> dict[str, Any]: + return self.db.stats() + + def export(self, to_path: str, include_vectors: bool = False) -> None: + import json + + with open(to_path, "w", encoding="utf-8") as handle: + for doc in self.db.list_documents(): + handle.write(json.dumps({"type": "document", **dict(doc)})) + handle.write("\n") + cur = self.db.conn.execute("SELECT * FROM chunks") + for row in cur.fetchall(): + data = dict(row) + if data.get("extra"): + data["extra"] = loads_json(data["extra"]) + handle.write(json.dumps({"type": "chunk", **data})) + handle.write("\n") + if include_vectors: + cur = self.db.conn.execute("SELECT * FROM vectors") + from array import array + + for row in cur.fetchall(): + data = dict(row) + arr = array("f") + arr.frombytes(data["vec"]) + data["vec"] = list(arr) + handle.write(json.dumps({"type": "vector", **data})) + handle.write("\n") + + def vacuum(self) -> None: + self.db.vacuum() + + def close(self) -> None: + self.db.close() diff --git a/raglite_sqlite/chunking.py b/raglite_sqlite/chunking.py new file mode 100644 index 0000000..fed59f0 --- /dev/null +++ b/raglite_sqlite/chunking.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Iterable, Iterator, List + +from .typing import ParsedBlock +from .utils import normalize_text + + +@dataclass +class Chunk: + position: int + text: str + text_norm: str + section: str | None = None + + +def fixed_tokens(text: str, size: int = 512, overlap: int = 64) -> Iterator[Chunk]: + tokens = text.split() + start = 0 + position = 0 + while start < len(tokens): + end = min(start + size, len(tokens)) + chunk_tokens = tokens[start:end] + raw = " ".join(chunk_tokens) + yield Chunk(position=position, text=raw, text_norm=normalize_text(raw)) + position += 1 + if end == len(tokens): + break + start = max(end - overlap, 0) + + +def recursive(text: str, size: int = 512, overlap: int = 64, break_on: List[str] | None = None) -> Iterator[Chunk]: + if break_on is None: + break_on = ["\n\n", ". "] + + def split(text_block: str) -> Iterable[str]: + for delimiter in break_on: + if delimiter in text_block and len(text_block) > size * 2: + parts = text_block.split(delimiter) + for i, part in enumerate(parts): + suffix = delimiter if i < len(parts) - 1 else "" + yield from split(part + suffix) + return + yield text_block + + pieces = [piece.strip() for piece in split(text) if piece.strip()] + buffer: list[str] = [] + position = 0 + for piece in pieces: + buffer.append(piece) + candidate = " ".join(buffer) + if len(candidate.split()) >= size: + yield Chunk(position=position, text=candidate, text_norm=normalize_text(candidate)) + position += 1 + buffer = buffer[-1:] # keep last element for overlap + if buffer: + remainder = " ".join(buffer) + if remainder: + yield Chunk(position=position, text=remainder, text_norm=normalize_text(remainder)) + + +def chunk_blocks(blocks: Iterable[ParsedBlock], strategy: str = "recursive", size: int = 512, overlap: int = 64) -> list[Chunk]: + chunks: list[Chunk] = [] + for idx, block in enumerate(blocks): + section = block.get("section") + text = block.get("text", "") + chunk_iter = recursive if strategy == "recursive" else fixed_tokens + for chunk in chunk_iter(text, size=size, overlap=overlap): # type: ignore[arg-type] + chunk.section = section + chunk.position += len(chunks) + chunks.append(chunk) + return chunks diff --git a/raglite_sqlite/cli.py b/raglite_sqlite/cli.py new file mode 100644 index 0000000..69e9c55 --- /dev/null +++ b/raglite_sqlite/cli.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Optional + +import typer +from rich.console import Console +from rich.table import Table + +from .api import RagLite +app = typer.Typer(help="RagLite SQLite CLI") +console = Console() + + +def get_rag(db: Path) -> RagLite: + return RagLite(str(db)) + + +def get_backend(model: Optional[str]): + from .embeddings.sentence_transformers_backend import SentenceTransformersBackend + + return SentenceTransformersBackend(model_name=model or "sentence-transformers/all-MiniLM-L6-v2") + + +@app.command() +def init(db: Path = typer.Option(..., help="Database path")) -> None: + rag = get_rag(db) + rag.close() + console.print(f"Initialized database at [bold]{db}[/bold]") + + +@app.command() +def index( + path: Path = typer.Argument(..., exists=True, file_okay=True, dir_okay=True), + db: Path = typer.Option(..., help="Database path"), + tags: Optional[str] = typer.Option(None, help="Comma-separated tags"), + model: Optional[str] = typer.Option(None, help="Embedding model name"), + chunk_size: int = typer.Option(512, help="Chunk size in tokens"), + overlap: int = typer.Option(64, help="Chunk overlap"), + glob: Optional[str] = typer.Option(None, help="Glob pattern"), + recursive: bool = typer.Option(True, help="Recurse into directories"), + skip_unchanged: bool = typer.Option(True, help="Skip unchanged files"), +) -> None: + backend = get_backend(model) + rag = get_rag(db) + result = rag.index( + [str(path)], + tags=tags, + chunk_size_tokens=chunk_size, + chunk_overlap_tokens=overlap, + embedding_backend=backend, + model_name=model, + glob=glob, + recurse=recursive, + skip_unchanged=skip_unchanged, + ) + console.print( + f"Indexed {result['files']} files, {result['chunks']} chunks (skipped {result['skipped']} unchanged)" + ) + + +@app.command() +def query( + text: str = typer.Argument(..., help="Query text"), + db: Path = typer.Option(..., help="Database path"), + k: int = typer.Option(8, help="Number of results"), + hybrid: float = typer.Option(0.6, min=0.0, max=1.0, help="Hybrid weight"), + max_per_doc: int = typer.Option(3, help="Max results per document"), + filters: Optional[list[str]] = typer.Option(None, "--filter", help="Filter key=value"), +) -> None: + rag = get_rag(db) + filter_dict: dict[str, str] | None = None + if filters: + filter_dict = {} + for item in filters: + if "=" not in item: + raise typer.BadParameter("Filters must be in key=value format") + key, value = item.split("=", 1) + filter_dict[key] = value + results = rag.search(text, k=k, hybrid_weight=hybrid, max_per_doc=max_per_doc, filters=filter_dict) + table = Table(show_header=True, header_style="bold magenta") + table.add_column("Score", justify="right") + table.add_column("Doc ID") + table.add_column("Section") + table.add_column("Snippet") + for item in results: + table.add_row( + f"{item['score']:.3f}", + item.get("doc_id", ""), + item.get("section") or "", + (item.get("snippet") or item.get("text", ""))[:120], + ) + console.print(table) + + +@app.command() +def stats(db: Path = typer.Option(..., help="Database path")) -> None: + rag = get_rag(db) + info = rag.stats() + console.print(json.dumps(info, indent=2)) + + +@app.command() +def export( + db: Path = typer.Option(..., help="Database path"), + to: Path = typer.Option(..., help="Destination NDJSON"), + include_vectors: bool = typer.Option( + False, "--include-vectors/--no-include-vectors", help="Include vector blobs" + ), +) -> None: + rag = get_rag(db) + rag.export(str(to), include_vectors=include_vectors) + console.print(f"Exported data to {to}") + + +@app.command() +def vacuum(db: Path = typer.Option(..., help="Database path")) -> None: + rag = get_rag(db) + rag.vacuum() + console.print("VACUUM completed") diff --git a/raglite_sqlite/db.py b/raglite_sqlite/db.py new file mode 100644 index 0000000..fdcc687 --- /dev/null +++ b/raglite_sqlite/db.py @@ -0,0 +1,250 @@ +from __future__ import annotations + +import sqlite3 +from array import array +from pathlib import Path +from typing import Iterable, List, Optional, Sequence + +from .utils import dumps_json, ensure_directory, loads_json + +PRAGMAS = { + "journal_mode": "WAL", + "synchronous": "NORMAL", + "temp_store": "MEMORY", + "mmap_size": 268435456, + "foreign_keys": 1, +} + + +class Database: + def __init__(self, path: Path, create: bool = True) -> None: + self.path = path + ensure_directory(path) + need_init = create and not path.exists() + self.conn = sqlite3.connect(path, check_same_thread=False) + self.conn.row_factory = sqlite3.Row + self._apply_pragmas() + if need_init: + self.run_schema() + self._ensure_fts5() + + def _apply_pragmas(self) -> None: + cur = self.conn.cursor() + for key, value in PRAGMAS.items(): + cur.execute(f"PRAGMA {key}={value}") + cur.close() + + def _ensure_fts5(self) -> None: + cur = self.conn.execute("PRAGMA compile_options") + options = {row[0] for row in cur.fetchall()} + if "ENABLE_FTS5" not in options: + raise RuntimeError("SQLite build does not support FTS5") + + def run_schema(self) -> None: + schema_sql = Path(__file__).with_name("schema.sql").read_text(encoding="utf-8") + self.conn.executescript(schema_sql) + self.conn.commit() + + def close(self) -> None: + self.conn.close() + + # Document helpers + def upsert_document( + self, + doc_id: str, + *, + source_path: str, + mime: str, + tags: str | None, + created_at: int, + updated_at: int, + sha256: str, + ) -> None: + self.conn.execute( + """ + INSERT INTO documents(doc_id, source_path, mime, tags, created_at, updated_at, sha256) + VALUES(?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(doc_id) DO UPDATE SET + source_path=excluded.source_path, + mime=excluded.mime, + tags=excluded.tags, + updated_at=excluded.updated_at, + sha256=excluded.sha256 + """, + (doc_id, source_path, mime, tags, created_at, updated_at, sha256), + ) + + def get_document(self, doc_id: str) -> Optional[sqlite3.Row]: + cur = self.conn.execute("SELECT * FROM documents WHERE doc_id = ?", (doc_id,)) + return cur.fetchone() + + def delete_document(self, doc_id: str) -> None: + self.conn.execute("DELETE FROM documents WHERE doc_id = ?", (doc_id,)) + + def list_documents(self) -> List[sqlite3.Row]: + cur = self.conn.execute("SELECT * FROM documents") + return list(cur.fetchall()) + + # Chunk helpers + def upsert_chunk( + self, + chunk_id: str, + doc_id: str, + position: int, + text: str, + text_norm: str, + section: str | None, + sha256: str, + extra: dict | None = None, + ) -> None: + self.conn.execute( + """ + INSERT INTO chunks(chunk_id, doc_id, position, text, text_norm, section, sha256, extra) + VALUES(?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(chunk_id) DO UPDATE SET + text=excluded.text, + text_norm=excluded.text_norm, + section=excluded.section, + sha256=excluded.sha256, + extra=excluded.extra + """, + (chunk_id, doc_id, position, text, text_norm, section, sha256, dumps_json(extra or {})), + ) + self.conn.execute( + "DELETE FROM chunk_fts WHERE chunk_id = ?", + (chunk_id,), + ) + self.conn.execute( + "INSERT INTO chunk_fts(chunk_id, text_norm) VALUES(?, ?)", + (chunk_id, text_norm), + ) + + def delete_chunk(self, chunk_id: str) -> None: + self.conn.execute("DELETE FROM chunks WHERE chunk_id = ?", (chunk_id,)) + self.conn.execute("DELETE FROM chunk_fts WHERE chunk_id = ?", (chunk_id,)) + + def iter_chunks_for_doc(self, doc_id: str) -> Iterable[sqlite3.Row]: + cur = self.conn.execute("SELECT * FROM chunks WHERE doc_id = ? ORDER BY position", (doc_id,)) + return cur.fetchall() + + # Vector helpers + def upsert_vector( + self, + chunk_id: str, + vector: Sequence[float], + model_name: str, + ) -> None: + blob = array("f", vector).tobytes() + dim = len(vector) + dtype = "float32" + self.conn.execute( + """ + INSERT INTO vectors(chunk_id, dim, dtype, vec, model_name, created_at) + VALUES(?, ?, ?, ?, ?, strftime('%s','now')) + ON CONFLICT(chunk_id) DO UPDATE SET + dim=excluded.dim, + dtype=excluded.dtype, + vec=excluded.vec, + model_name=excluded.model_name, + created_at=excluded.created_at + """, + (chunk_id, dim, dtype, blob, model_name), + ) + + def get_vectors_by_ids(self, chunk_ids: Sequence[str]) -> dict[str, list[float]]: + if not chunk_ids: + return {} + placeholders = ",".join("?" for _ in chunk_ids) + cur = self.conn.execute( + f"SELECT chunk_id, dim, dtype, vec FROM vectors WHERE chunk_id IN ({placeholders})", + tuple(chunk_ids), + ) + vectors: dict[str, list[float]] = {} + for row in cur.fetchall(): + data = array("f") + data.frombytes(row["vec"]) + vectors[row["chunk_id"]] = list(data) + return vectors + + def get_all_vectors(self, model_name: str | None = None) -> tuple[list[list[float]], list[str]]: + if model_name: + cur = self.conn.execute( + "SELECT chunk_id, dim, dtype, vec FROM vectors WHERE model_name = ?", + (model_name,), + ) + else: + cur = self.conn.execute("SELECT chunk_id, dim, dtype, vec FROM vectors") + rows = cur.fetchall() + if not rows: + return [], [] + chunk_ids: list[str] = [] + vectors: list[list[float]] = [] + for row in rows: + data = array("f") + data.frombytes(row["vec"]) + vectors.append(list(data)) + chunk_ids.append(row["chunk_id"]) + return vectors, chunk_ids + + def get_embedding_cache(self, content_sha: str, model_name: str) -> Optional[list[float]]: + cur = self.conn.execute( + "SELECT dim, dtype, vec FROM cache_embeddings WHERE content_sha = ? AND model_name = ?", + (content_sha, model_name), + ) + row = cur.fetchone() + if not row: + return None + data = array("f") + data.frombytes(row["vec"]) + return list(data) + + def upsert_embedding_cache(self, content_sha: str, model_name: str, vector: Sequence[float]) -> None: + self.conn.execute( + """ + INSERT INTO cache_embeddings(content_sha, model_name, dim, dtype, vec) + VALUES(?, ?, ?, ?, ?) + ON CONFLICT(content_sha, model_name) DO UPDATE SET + dim=excluded.dim, + dtype=excluded.dtype, + vec=excluded.vec + """, + (content_sha, model_name, len(vector), "float32", array("f", vector).tobytes()), + ) + + def commit(self) -> None: + self.conn.commit() + + def vacuum(self) -> None: + self.conn.execute("VACUUM") + + def stats(self) -> dict[str, object]: + cur = self.conn.cursor() + counts = { + "documents": cur.execute("SELECT COUNT(*) FROM documents").fetchone()[0], + "chunks": cur.execute("SELECT COUNT(*) FROM chunks").fetchone()[0], + "vectors": cur.execute("SELECT COUNT(*) FROM vectors").fetchone()[0], + } + models = [row[0] for row in cur.execute("SELECT DISTINCT model_name FROM vectors").fetchall()] + return {**counts, "models": models} + + +def cosine_search(matrix: Sequence[Sequence[float]], query_vec: Sequence[float], top_k: int) -> list[tuple[int, float]]: + if not matrix: + return [] + def norm(vec: Sequence[float]) -> float: + return sum(value * value for value in vec) ** 0.5 + + query_norm = norm(query_vec) + if query_norm == 0: + return [] + sims: list[float] = [] + for row in matrix: + row_norm = norm(row) + if row_norm == 0: + sims.append(0.0) + continue + dot = sum(a * b for a, b in zip(row, query_vec)) + sims.append(dot / (row_norm * query_norm)) + indexed = list(enumerate(sims)) + indexed.sort(key=lambda item: item[1], reverse=True) + return [(idx, score) for idx, score in indexed[:top_k]] diff --git a/raglite_sqlite/embeddings/__init__.py b/raglite_sqlite/embeddings/__init__.py new file mode 100644 index 0000000..3cef675 --- /dev/null +++ b/raglite_sqlite/embeddings/__init__.py @@ -0,0 +1,5 @@ +"""Embedding backends for RagLite.""" + +from .sentence_transformers_backend import SentenceTransformersBackend + +__all__ = ["SentenceTransformersBackend"] diff --git a/raglite_sqlite/embeddings/base.py b/raglite_sqlite/embeddings/base.py new file mode 100644 index 0000000..072ce24 --- /dev/null +++ b/raglite_sqlite/embeddings/base.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Sequence + + +class EmbeddingBackend(ABC): + """Abstract embedding backend.""" + + @abstractmethod + def embed_texts(self, texts: Sequence[str], model_name: str | None = None) -> Sequence[Sequence[float]]: + """Return embeddings as float32 numpy array of shape (N, D).""" + + +class DummyBackend(EmbeddingBackend): + """Deterministic embedding backend used for tests.""" + + def __init__(self, dim: int = 8) -> None: + self.dim = dim + self.calls = 0 + + def embed_texts(self, texts: Sequence[str], model_name: str | None = None) -> Sequence[Sequence[float]]: + self.calls += 1 + vectors: list[list[float]] = [] + for text in texts: + h = abs(hash(text)) + vec = [(h >> i) & 0xFF for i in range(self.dim)] + norm = sum(value * value for value in vec) ** 0.5 + if norm == 0: + vec[0] = 1.0 + norm = 1.0 + vectors.append([value / norm for value in vec]) + return vectors diff --git a/raglite_sqlite/embeddings/openai_backend.py b/raglite_sqlite/embeddings/openai_backend.py new file mode 100644 index 0000000..a49343d --- /dev/null +++ b/raglite_sqlite/embeddings/openai_backend.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +import os +from typing import Sequence + +from .base import EmbeddingBackend + + +class OpenAIBackend(EmbeddingBackend): + """Embedding backend using the OpenAI API.""" + + def __init__(self, default_model: str = "text-embedding-3-small", batch_size: int = 64) -> None: + api_key = os.getenv("OPENAI_API_KEY") + if not api_key: + raise RuntimeError("OPENAI_API_KEY is not set") + self.default_model = default_model + self.batch_size = batch_size + from openai import OpenAI + + self._client = OpenAI(api_key=api_key) + + def embed_texts(self, texts: Sequence[str], model_name: str | None = None) -> Sequence[Sequence[float]]: + model = model_name or self.default_model + embeddings: list[list[float]] = [] + for start in range(0, len(texts), self.batch_size): + batch = list(texts[start : start + self.batch_size]) + response = self._client.embeddings.create(model=model, input=batch) + for item in response.data: + embeddings.append([float(value) for value in item.embedding]) + return embeddings diff --git a/raglite_sqlite/embeddings/sentence_transformers_backend.py b/raglite_sqlite/embeddings/sentence_transformers_backend.py new file mode 100644 index 0000000..9630316 --- /dev/null +++ b/raglite_sqlite/embeddings/sentence_transformers_backend.py @@ -0,0 +1,26 @@ +from __future__ import annotations + +from typing import Sequence + +from .base import EmbeddingBackend + + +class SentenceTransformersBackend(EmbeddingBackend): + """Wrapper around sentence-transformers models.""" + + def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2", device: str | None = None) -> None: + self.model_name = model_name + self.device = device + self._model = None + + def _load(self) -> object: + if self._model is None: + from sentence_transformers import SentenceTransformer + + self._model = SentenceTransformer(self.model_name, device=self.device) + return self._model + + def embed_texts(self, texts: Sequence[str], model_name: str | None = None) -> Sequence[Sequence[float]]: + model = self._load() + embeddings = model.encode(list(texts), show_progress_bar=False, convert_to_numpy=True, device=self.device) + return embeddings.astype("float32").tolist() diff --git a/raglite_sqlite/parsers/__init__.py b/raglite_sqlite/parsers/__init__.py new file mode 100644 index 0000000..6415c1a --- /dev/null +++ b/raglite_sqlite/parsers/__init__.py @@ -0,0 +1,17 @@ +"""Document parsers.""" + +from .txt import TextParser +from .md import MarkdownParser +from .html import HTMLParser +from .pdf import PDFParser +from .docx import DocxParser +from .csv import CSVParser + +__all__ = [ + "TextParser", + "MarkdownParser", + "HTMLParser", + "PDFParser", + "DocxParser", + "CSVParser", +] diff --git a/raglite_sqlite/parsers/base.py b/raglite_sqlite/parsers/base.py new file mode 100644 index 0000000..9a8f313 --- /dev/null +++ b/raglite_sqlite/parsers/base.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Iterable + +from ..typing import ParsedBlock + + +class BaseParser(ABC): + @abstractmethod + def parse(self, path: str, **options: object) -> Iterable[ParsedBlock]: + ... diff --git a/raglite_sqlite/parsers/csv.py b/raglite_sqlite/parsers/csv.py new file mode 100644 index 0000000..32bfca3 --- /dev/null +++ b/raglite_sqlite/parsers/csv.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +import csv +from pathlib import Path +from typing import Iterable, Sequence + +from .base import BaseParser +from ..typing import ParsedBlock +from ..utils import normalize_text + + +class CSVParser(BaseParser): + def parse(self, path: str, **options: object) -> Iterable[ParsedBlock]: + columns = options.get("columns") if isinstance(options, dict) else None + selected: Sequence[str] | None = None + if isinstance(columns, Sequence): + selected = list(columns) + rows: list[str] = [] + with Path(path).open("r", encoding="utf-8", newline="") as handle: + reader = csv.DictReader(handle) + for row in reader: + if selected: + values = [str(row.get(col, "")) for col in selected] + else: + values = [str(value) for value in row.values()] + rows.append(" ".join(values)) + yield ParsedBlock(text=normalize_text("\n".join(rows)), section=None) diff --git a/raglite_sqlite/parsers/docx.py b/raglite_sqlite/parsers/docx.py new file mode 100644 index 0000000..5d2756a --- /dev/null +++ b/raglite_sqlite/parsers/docx.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Iterable + +from .base import BaseParser +from ..typing import ParsedBlock +from ..utils import normalize_text + + +class DocxParser(BaseParser): + def parse(self, path: str, **options: object) -> Iterable[ParsedBlock]: + try: + from docx import Document + + document = Document(Path(path)) + texts = [paragraph.text for paragraph in document.paragraphs if paragraph.text.strip()] + except Exception: + texts = [] + yield ParsedBlock(text=normalize_text("\n".join(texts)), section=None) diff --git a/raglite_sqlite/parsers/html.py b/raglite_sqlite/parsers/html.py new file mode 100644 index 0000000..8d59a1c --- /dev/null +++ b/raglite_sqlite/parsers/html.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Iterable, List + +from .base import BaseParser +from ..typing import ParsedBlock +from ..utils import normalize_text + + +class HTMLParser(BaseParser): + def parse(self, path: str, **options: object) -> Iterable[ParsedBlock]: + content = Path(path).read_text(encoding="utf-8") + try: + from bs4 import BeautifulSoup + + soup = BeautifulSoup(content, "html.parser") + extractor = "soup" + except Exception: + soup = None + extractor = "plain" + blocks: List[ParsedBlock] = [] + current_section: str | None = None + if extractor == "soup" and soup is not None: + for element in soup.find_all(["h1", "h2", "h3", "p"]): + if element.name in {"h1", "h2", "h3"}: + current_section = normalize_text(element.get_text(" ")) + else: + text = normalize_text(element.get_text(" ")) + if text: + blocks.append(ParsedBlock(text=text, section=current_section)) + if not blocks: + text = normalize_text(soup.get_text(" ")) + blocks.append(ParsedBlock(text=text, section=None)) + else: + clean = normalize_text( + content.replace("<", " ").replace(">", " ") + ) + blocks.append(ParsedBlock(text=clean, section=None)) + return blocks diff --git a/raglite_sqlite/parsers/md.py b/raglite_sqlite/parsers/md.py new file mode 100644 index 0000000..f056490 --- /dev/null +++ b/raglite_sqlite/parsers/md.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Iterable, List + +from .base import BaseParser +from ..typing import ParsedBlock +from ..utils import normalize_text + + +class MarkdownParser(BaseParser): + def parse(self, path: str, **options: object) -> Iterable[ParsedBlock]: + content = Path(path).read_text(encoding="utf-8") + try: + import frontmatter + + post = frontmatter.loads(content) + body = post.content + except Exception: + body = content + lines = body.splitlines() + blocks: List[ParsedBlock] = [] + current_section: str | None = None + buffer: list[str] = [] + for line in lines: + if line.startswith("#"): + if buffer: + blocks.append(ParsedBlock(text=normalize_text("\n".join(buffer)), section=current_section)) + buffer = [] + current_section = normalize_text(line.lstrip("# ")) + elif line.strip().startswith("```"): + continue + else: + buffer.append(line) + if buffer: + blocks.append(ParsedBlock(text=normalize_text("\n".join(buffer)), section=current_section)) + if not blocks: + blocks.append(ParsedBlock(text=normalize_text(body), section=current_section)) + return blocks diff --git a/raglite_sqlite/parsers/pdf.py b/raglite_sqlite/parsers/pdf.py new file mode 100644 index 0000000..f6e5fc4 --- /dev/null +++ b/raglite_sqlite/parsers/pdf.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Iterable + +from .base import BaseParser +from ..typing import ParsedBlock +from ..utils import normalize_text + + +class PDFParser(BaseParser): + def parse(self, path: str, **options: object) -> Iterable[ParsedBlock]: + try: + from PyPDF2 import PdfReader + + reader = PdfReader(Path(path).open("rb")) + text = "\n".join(page.extract_text() or "" for page in reader.pages) + except Exception: + text = "" + yield ParsedBlock(text=normalize_text(text), section=None) diff --git a/raglite_sqlite/parsers/txt.py b/raglite_sqlite/parsers/txt.py new file mode 100644 index 0000000..b409de6 --- /dev/null +++ b/raglite_sqlite/parsers/txt.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Iterable + +from .base import BaseParser +from ..typing import ParsedBlock +from ..utils import normalize_text + + +class TextParser(BaseParser): + def parse(self, path: str, **options: object) -> Iterable[ParsedBlock]: + content = Path(path).read_text(encoding="utf-8") + yield ParsedBlock(text=normalize_text(content), section=None) diff --git a/raglite_sqlite/rerank.py b/raglite_sqlite/rerank.py new file mode 100644 index 0000000..f2aa880 --- /dev/null +++ b/raglite_sqlite/rerank.py @@ -0,0 +1,15 @@ +from __future__ import annotations + +from typing import Iterable, Protocol + +from .typing import SearchResult + + +class Reranker(Protocol): + def rerank(self, query: str, results: Iterable[SearchResult]) -> Iterable[SearchResult]: + ... + + +class NoopReranker: + def rerank(self, query: str, results: Iterable[SearchResult]) -> Iterable[SearchResult]: + return results diff --git a/raglite_sqlite/schema.sql b/raglite_sqlite/schema.sql new file mode 100644 index 0000000..18ebd59 --- /dev/null +++ b/raglite_sqlite/schema.sql @@ -0,0 +1,58 @@ +PRAGMA foreign_keys = ON; + +CREATE TABLE IF NOT EXISTS documents ( + doc_id TEXT PRIMARY KEY, + source_path TEXT NOT NULL, + mime TEXT, + tags TEXT, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + sha256 TEXT NOT NULL +); + +CREATE TABLE IF NOT EXISTS chunks ( + chunk_id TEXT PRIMARY KEY, + doc_id TEXT NOT NULL, + position INTEGER NOT NULL, + text TEXT NOT NULL, + text_norm TEXT NOT NULL, + section TEXT, + sha256 TEXT NOT NULL, + extra JSON, + FOREIGN KEY(doc_id) REFERENCES documents(doc_id) ON DELETE CASCADE +); + +CREATE VIRTUAL TABLE IF NOT EXISTS chunk_fts USING fts5( + chunk_id, + text_norm, + tokenize = 'unicode61' +); + +CREATE TABLE IF NOT EXISTS vectors ( + chunk_id TEXT PRIMARY KEY, + dim INTEGER NOT NULL, + dtype TEXT NOT NULL, + vec BLOB NOT NULL, + model_name TEXT NOT NULL, + created_at INTEGER NOT NULL, + FOREIGN KEY(chunk_id) REFERENCES chunks(chunk_id) ON DELETE CASCADE +); + +CREATE TABLE IF NOT EXISTS cache_embeddings ( + content_sha TEXT NOT NULL, + model_name TEXT NOT NULL, + dim INTEGER NOT NULL, + dtype TEXT NOT NULL, + vec BLOB NOT NULL, + PRIMARY KEY(content_sha, model_name) +); + +CREATE TABLE IF NOT EXISTS meta ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL +); + +CREATE INDEX IF NOT EXISTS idx_documents_sha ON documents(sha256); +CREATE INDEX IF NOT EXISTS idx_chunks_doc ON chunks(doc_id, position); +CREATE INDEX IF NOT EXISTS idx_chunks_sha ON chunks(sha256); +CREATE INDEX IF NOT EXISTS idx_vectors_model ON vectors(model_name); diff --git a/raglite_sqlite/search.py b/raglite_sqlite/search.py new file mode 100644 index 0000000..add8eff --- /dev/null +++ b/raglite_sqlite/search.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +import math +from typing import Dict, Iterable, List, Tuple + +from .db import Database, cosine_search +from .embeddings.base import EmbeddingBackend +from .typing import SearchResult +from .utils import normalize_text + + +def bm25_search( + db: Database, + query: str, + k: int, + filters: dict | None = None, +) -> list[tuple[str, float]]: + sql = [ + "SELECT c.chunk_id AS chunk_id, bm25(chunk_fts) AS score" + " FROM chunk_fts JOIN chunks c USING(chunk_id)" + " JOIN documents d ON d.doc_id = c.doc_id" + " WHERE chunk_fts MATCH ?" + ] + args: List[object] = [query] + if filters: + if doc_id := filters.get("doc_id"): + sql.append(" AND d.doc_id = ?") + args.append(doc_id) + if tags := filters.get("tags"): + sql.append(" AND d.tags LIKE ?") + args.append(f"%{tags}%") + if source_prefix := filters.get("source_path"): + sql.append(" AND d.source_path LIKE ?") + args.append(f"{source_prefix}%") + sql.append(" ORDER BY score LIMIT ?") + args.append(k * 5) + cur = db.conn.execute("".join(sql), tuple(args)) + results = [(row["chunk_id"], float(row["score"])) for row in cur.fetchall()] + return results + + +def vector_search( + db: Database, + backend: EmbeddingBackend, + query: str, + model_name: str | None, + k: int, +) -> list[tuple[str, float]]: + matrix, chunk_ids = db.get_all_vectors(model_name=model_name) + if not matrix: + return [] + query_vectors = backend.embed_texts([query], model_name=model_name) + query_vec = list(query_vectors[0]) + results = cosine_search(matrix, query_vec, top_k=min(len(chunk_ids), max(k * 5, 10))) + return [(chunk_ids[idx], score) for idx, score in results] + + +def hybrid_fuse( + lexical: list[tuple[str, float]], + semantic: list[tuple[str, float]], + weight: float, + k: int, +) -> dict[str, Dict[str, float]]: + scores: dict[str, Dict[str, float]] = {} + if lexical: + bm25_values = [score for _, score in lexical] + max_bm25 = max(bm25_values) + min_bm25 = min(bm25_values) + else: + max_bm25 = min_bm25 = 0.0 + if semantic: + vector_values = [score for _, score in semantic] + max_vec = max(vector_values) + min_vec = min(vector_values) + else: + max_vec = min_vec = 0.0 + + def normalize(score: float, lo: float, hi: float) -> float: + if math.isclose(hi, lo): + return 0.0 + return (score - lo) / (hi - lo) + + for chunk_id, score in lexical[: k * 5]: + scores.setdefault(chunk_id, {})["bm25"] = normalize(score, min_bm25, max_bm25) + for chunk_id, score in semantic[: k * 5]: + scores.setdefault(chunk_id, {})["vector"] = normalize(score, min_vec, max_vec) + + fused: dict[str, Dict[str, float]] = {} + for chunk_id, components in scores.items(): + bm25_score = components.get("bm25", 0.0) + vector_score = components.get("vector", 0.0) + fused_score = weight * bm25_score + (1 - weight) * vector_score + fused[chunk_id] = { + "fused": fused_score, + "bm25": bm25_score, + "vector": vector_score, + } + return fused + + +def assemble_results(db: Database, scores: dict[str, Dict[str, float]], k: int, with_snippets: bool = True, max_per_doc: int = 3) -> List[SearchResult]: + if not scores: + return [] + placeholders = ",".join("?" for _ in scores) + cur = db.conn.execute( + f"SELECT c.*, d.source_path, d.tags FROM chunks c JOIN documents d ON d.doc_id = c.doc_id WHERE c.chunk_id IN ({placeholders})", + tuple(scores.keys()), + ) + rows = {row["chunk_id"]: row for row in cur.fetchall()} + per_doc: dict[str, int] = {} + sorted_ids = sorted(scores.items(), key=lambda item: item[1]["fused"], reverse=True) + results: List[SearchResult] = [] + for chunk_id, components in sorted_ids: + row = rows.get(chunk_id) + if row is None: + continue + doc_id = row["doc_id"] + count = per_doc.get(doc_id, 0) + if count >= max_per_doc: + continue + per_doc[doc_id] = count + 1 + snippet = row["text"][:200] + result: SearchResult = { + "chunk_id": chunk_id, + "doc_id": doc_id, + "position": row["position"], + "text": row["text"], + "section": row["section"], + "source_path": row["source_path"], + "tags": row["tags"], + "score": components["fused"], + "bm25_score": components.get("bm25", 0.0), + "vector_score": components.get("vector", 0.0), + } + if with_snippets: + result["snippet"] = snippet + results.append(result) + if len(results) >= k: + break + return results diff --git a/raglite_sqlite/typing.py b/raglite_sqlite/typing.py new file mode 100644 index 0000000..f46924e --- /dev/null +++ b/raglite_sqlite/typing.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from typing import Iterable, Optional, Protocol, TypedDict + + +class ParsedBlock(TypedDict, total=False): + text: str + section: Optional[str] + + +class ChunkDict(TypedDict): + chunk_id: str + doc_id: str + position: int + text: str + text_norm: str + section: Optional[str] + sha256: str + + +class SearchResult(TypedDict, total=False): + doc_id: str + score: float + snippet: str + section: Optional[str] + source_path: str + position: int + text: str + tags: Optional[str] + bm25_score: float + vector_score: float + + +class Parser(Protocol): + def parse(self, path: str, **options: object) -> Iterable[ParsedBlock]: + ... diff --git a/raglite_sqlite/utils.py b/raglite_sqlite/utils.py new file mode 100644 index 0000000..097666d --- /dev/null +++ b/raglite_sqlite/utils.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +import hashlib +import json +import os +import time +import unicodedata +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable, Iterator, Sequence + +try: + import orjson # type: ignore[import-not-found] +except Exception: # pragma: no cover - fallback when orjson unavailable + orjson = None + +try: + from tqdm import tqdm +except Exception: # pragma: no cover - fallback when tqdm unavailable + def tqdm(iterable: Iterable[object], total: int | None = None, desc: str = ""): + for item in iterable: + yield item + + +@dataclass +class Progress: + total: int + description: str = "" + + def track(self, iterable: Iterable[object]) -> Iterator[object]: + yield from tqdm(iterable, total=self.total, desc=self.description) + + +def sha256_text(text: str) -> str: + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def sha256_file(path: Path) -> str: + h = hashlib.sha256() + with path.open("rb") as fh: + for chunk in iter(lambda: fh.read(8192), b""): + h.update(chunk) + return h.hexdigest() + + +def ensure_directory(path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + + +def normalize_text(text: str) -> str: + normalized = unicodedata.normalize("NFC", text) + return " ".join(normalized.split()) + + +def detect_mime(path: Path) -> str: + ext = path.suffix.lower() + return { + ".txt": "text/plain", + ".md": "text/markdown", + ".markdown": "text/markdown", + ".html": "text/html", + ".htm": "text/html", + ".pdf": "application/pdf", + ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + ".csv": "text/csv", + }.get(ext, "application/octet-stream") + + +def now_ts() -> int: + return int(time.time()) + + +def iter_files(paths: Sequence[str], recurse: bool = True, glob: str | None = None) -> list[Path]: + candidates: list[Path] = [] + for input_path in paths: + path = Path(input_path) + if path.is_dir(): + pattern = glob or "**/*" if recurse else "*" + for child in path.glob(pattern): + if child.is_file(): + candidates.append(child) + elif path.is_file(): + candidates.append(path) + unique = {p.resolve(): p for p in candidates} + return list(unique.values()) + + +def dumps_json(data: object) -> str: + if orjson is not None: + return orjson.dumps(data).decode("utf-8") + import json + + return json.dumps(data) + + +def loads_json(data: str) -> object: + import json + + return json.loads(data) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..9decda4 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +import sys +from pathlib import Path +from typing import Iterator + +import pytest + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from raglite_sqlite.api import RagLite +from raglite_sqlite.embeddings.base import DummyBackend + + +@pytest.fixture() +def temp_db(tmp_path: Path) -> Iterator[tuple[Path, RagLite, DummyBackend]]: + db_path = tmp_path / "knowledge.db" + rag = RagLite(str(db_path)) + backend = DummyBackend(dim=16) + yield db_path, rag, backend + rag.close() diff --git a/tests/data/sample.csv b/tests/data/sample.csv new file mode 100644 index 0000000..fd84fc2 --- /dev/null +++ b/tests/data/sample.csv @@ -0,0 +1,3 @@ +name,description +alpha,First row content +beta,Second row information diff --git a/tests/data/sample.html b/tests/data/sample.html new file mode 100644 index 0000000..6372c5c --- /dev/null +++ b/tests/data/sample.html @@ -0,0 +1,8 @@ + +Sample HTML + +

Overview

+

This HTML document provides an overview of the project.

+

Additional content to index for search.

+ + diff --git a/tests/data/sample.md b/tests/data/sample.md new file mode 100644 index 0000000..8bd8d89 --- /dev/null +++ b/tests/data/sample.md @@ -0,0 +1,11 @@ +--- +title: Sample Markdown +--- + +# Introduction + +This markdown file discusses examples. + +## Details + +More detailed information lives here. diff --git a/tests/data/sample.txt b/tests/data/sample.txt new file mode 100644 index 0000000..f227828 --- /dev/null +++ b/tests/data/sample.txt @@ -0,0 +1,2 @@ +This is a sample text document. +It contains a few lines about testing RagLite. diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..104086a --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +pytest.importorskip("typer") + +from typer.testing import CliRunner + +from raglite_sqlite.cli import app +from raglite_sqlite.embeddings.base import DummyBackend + + +def test_cli_flow(tmp_path: Path, monkeypatch) -> None: + runner = CliRunner() + db_path = tmp_path / "knowledge.db" + + # init + result = runner.invoke(app, ["init", "--db", str(db_path)]) + assert result.exit_code == 0 + + # monkeypatch backend to avoid heavy model + dummy = DummyBackend(dim=16) + + monkeypatch.setattr("raglite_sqlite.cli.get_backend", lambda model: dummy) + + data_dir = Path(__file__).parent / "data" + result = runner.invoke( + app, + [ + "index", + str(data_dir), + "--db", + str(db_path), + ], + ) + assert result.exit_code == 0 + + result = runner.invoke(app, ["query", "sample", "--db", str(db_path)]) + assert result.exit_code == 0 + + result = runner.invoke(app, ["stats", "--db", str(db_path)]) + assert result.exit_code == 0 diff --git a/tests/test_ingest_and_search.py b/tests/test_ingest_and_search.py new file mode 100644 index 0000000..246b8f1 --- /dev/null +++ b/tests/test_ingest_and_search.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from pathlib import Path + +from raglite_sqlite.embeddings.base import DummyBackend + + +def test_index_and_search(temp_db: tuple[Path, object, object]) -> None: + db_path, rag, backend = temp_db + data_dir = Path(__file__).parent / "data" + result = rag.index([str(data_dir)], embedding_backend=backend) + stats = rag.stats() + assert result["chunks"] > 0 + assert stats["documents"] >= 1 + results = rag.search("sample", embedding_backend=backend) + assert results, "Expected at least one search result" + top = results[0] + assert "sample" in top["text"].lower() + + +def test_idempotent_index(temp_db: tuple[Path, object, object]) -> None: + db_path, rag, backend = temp_db + data_dir = Path(__file__).parent / "data" + rag.index([str(data_dir)], embedding_backend=backend) + result = rag.index([str(data_dir)], embedding_backend=backend) + assert result["skipped"] >= 1 + + +def test_embedding_cache(temp_db: tuple[Path, object, object]) -> None: + db_path, rag, backend = temp_db + data_dir = Path(__file__).parent / "data" + backend.calls = 0 + rag.index([str(data_dir)], embedding_backend=backend) + first_calls = backend.calls + rag.index([str(data_dir)], embedding_backend=backend, skip_unchanged=False) + assert backend.calls == first_calls