diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml new file mode 100644 index 0000000..f18fa60 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -0,0 +1,19 @@ +name: Bug report +about: Report a bug in raglite +labels: bug +body: + - type: textarea + id: description + attributes: + label: Description + description: What happened? + validations: + required: true + - type: input + id: version + attributes: + label: raglite version + - type: textarea + id: steps + attributes: + label: Steps to reproduce diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml new file mode 100644 index 0000000..f18c804 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.yml @@ -0,0 +1,19 @@ +name: Feature request +about: Suggest an idea for raglite +labels: enhancement +body: + - type: textarea + id: summary + attributes: + label: Summary + description: Briefly describe the feature. + validations: + required: true + - type: textarea + id: motivation + attributes: + label: Motivation + - type: textarea + id: alternatives + attributes: + label: Alternatives considered diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f729c6a..32780a1 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -7,30 +7,30 @@ on: jobs: build: - runs-on: ${{ matrix.os }} + runs-on: ubuntu-latest strategy: matrix: - os: [ubuntu-latest, macos-latest, windows-latest] - python-version: ['3.10', '3.11', '3.12'] + python-version: ["3.10", "3.11"] 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] + pip install -e .[dev,server] - name: Lint run: | - ruff check . - black --check . + ruff check src tests + black --check src tests + isort --check-only src tests - name: Type check - run: mypy raglite_sqlite - - name: Test - run: pytest --cov=raglite_sqlite + run: | + mypy src + - name: Tests + run: | + pytest -q + - name: Self-test + run: | + raglite self-test diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index f7f0c41..e187a59 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,22 +1,19 @@ repos: - repo: https://github.com/psf/black - rev: 24.4.2 + rev: 23.12.1 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 + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.1.9 + hooks: + - id: ruff + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.8.0 hooks: - - id: end-of-file-fixer - - id: trailing-whitespace + - id: mypy + additional_dependencies: + - types-requests diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..0215826 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,4 @@ +# Changelog + +## 0.1.0 - 2024-05-01 +- Initial release of raglite with SQLite-backed RAG pipeline, Typer CLI, FastAPI server, and demo corpus. diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md new file mode 100644 index 0000000..3fe3aed --- /dev/null +++ b/CODE_OF_CONDUCT.md @@ -0,0 +1,29 @@ +# Contributor Covenant Code of Conduct + +## Our Pledge + +We as members, contributors, and leaders pledge to make participation in our community a harassment-free experience for everyone. + +## Our Standards + +Examples of behaviour that contributes to a positive environment include: + +- Demonstrating empathy and kindness toward other people +- Being respectful of differing opinions, viewpoints, and experiences +- Giving and gracefully accepting constructive feedback + +## Enforcement Responsibilities + +Project maintainers are responsible for clarifying and enforcing our standards of acceptable behaviour. + +## Scope + +This Code of Conduct applies within all community spaces. + +## Enforcement + +Instances of abusive, harassing, or otherwise unacceptable behaviour may be reported by contacting the maintainers at `opensource@example.com`. + +## Attribution + +This Code of Conduct is adapted from the [Contributor Covenant](https://www.contributor-covenant.org/version/2/1/code_of_conduct.html). diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..1783136 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,25 @@ +# Contributing to raglite + +Thank you for helping improve raglite! The project aims to provide a practical local-first RAG stack powered by SQLite. + +## Getting started + +1. Fork and clone the repository. +2. Install dependencies: `pip install -e .[dev,server]`. +3. Enable pre-commit hooks: `pre-commit install`. +4. Run the demo self-test: `raglite self-test`. + +## Development workflow + +- Make sure `pytest -q` passes before submitting a pull request. +- Run `ruff`, `mypy`, and `black` using the provided pre-commit configuration. +- Update documentation and add tests for new features. +- Follow the MIT license and keep dependencies minimal. + +## Reporting issues + +Open an issue with detailed reproduction steps, expected behaviour, and environment information. Attach logs if available. + +## Code of conduct + +Please review `CODE_OF_CONDUCT.md` for community guidelines. diff --git a/README.md b/README.md index 77090cc..a453eb5 100644 --- a/README.md +++ b/README.md @@ -1,221 +1,51 @@ -# RagLite SQLite +# raglite -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. +Local-first retrieval augmented generation toolkit built on SQLite. Raglite bundles ingestion, chunking, embeddings, hybrid BM25/vector search, a Typer CLI, and a FastAPI microservice. Everything runs on CPU and stores state in a single SQLite database. -## Why RagLite? +## Features -- **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, rerankers, and adapters for LangChain / LlamaIndex. -- **Multi-user ready** – optional FastAPI server turns the SQLite database into a shared retrieval service. - -## Installation - -```bash -python -m venv .venv && source .venv/bin/activate # (Windows: .venv\Scripts\activate) -pip install -e . -``` +- SQLite FTS5 BM25 search with optional vector reranking. +- Pluggable embedding backends with automatic fallback to a deterministic hash model for testing. +- Python API, Typer CLI, and FastAPI server. +- Demo corpus and integration examples for LangChain and LlamaIndex. ## 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 +pip install -e .[server,dev] +raglite init-db --db demo.db +raglite ingest --db demo.db --path src/demo/mini_corpus --embed-model debug +raglite query --db demo.db --text "quick start guide" ``` -## 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 - ``` -- Apply a cross-encoder reranker at query time: - ```bash - raglite query "release checklist" --db knowledge.db --reranker cross-encoder --reranker-option model_name=cross-encoder/ms-marco-MiniLM-L-6-v2 - ``` -- 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 - ``` -- Serve a REST API for multi-user access (requires `pip install raglite-sqlite[server]`): - ```bash - raglite serve --db knowledge.db --host 0.0.0.0 --port 8080 - ``` - -## REST API +## Architecture -RagLite can expose a FastAPI-powered REST server for shared deployments. Install the optional extras first: - -```bash -pip install raglite-sqlite[server] ``` - -Then launch the server with `raglite serve`. Key endpoints: - -- `GET /health` – basic health check. -- `GET /stats` – retrieve database statistics. -- `POST /index` – index new files on disk. -- `POST /query` – perform hybrid search with optional reranking. -- `POST /delete` – remove a document by `doc_id`. - -All endpoints operate directly on the shared SQLite database, making it easy to expose retrieval to multiple users without extra infrastructure. - -## Document formats & OCR - -Beyond plain text and Markdown, RagLite understands PDF, HTML, DOCX, CSV, JSON, PPTX, and image files. Image ingestion uses OCR via `pytesseract`/`Pillow`—install them with `pip install raglite-sqlite[ocr]`. Custom parser options can be passed via `RagLite.index(..., parser_opts={...})` or the REST API. - -## 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) ++-----------------+ +| Typer CLI | ++-----------------+ + | ++-----------------+ +-----------------------+ +| API Layer |<--->| FastAPI Service | ++-----------------+ +-----------------------+ + | ++-----------------+ +-----------------------+ +| Search Engine |<--->| Vector Backends | ++-----------------+ +-----------------------+ + | ++-----------------------------------------------+ +| SQLite DB | ++-----------------------------------------------+ ``` -### 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 +See `src/raglite/schema.sql` for the full schema. -RagLite ships with lightweight adapters so you can plug the SQLite-backed retriever into popular orchestration frameworks. +## Tests -### 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) +```bash +pytest -q ``` -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. - ## License -MIT License. See [LICENSE](LICENSE). +MIT diff --git a/pyproject.toml b/pyproject.toml index 105ba34..7f4ff49 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ requires = ["setuptools>=64", "wheel"] build-backend = "setuptools.build_meta" [project] -name = "raglite-sqlite" +name = "raglite" version = "0.1.0" description = "Local-first RAG toolkit backed by a single SQLite database" readme = "README.md" @@ -11,65 +11,64 @@ 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", + "typer>=0.9.0", + "readability-lxml>=0.8.1", "beautifulsoup4>=4.12.0", - "python-frontmatter>=1.0.0", - "PyPDF2>=3.0.0", - "python-docx>=1.0.0", - "tqdm>=4.66.0", - "orjson>=3.9.0", - "sqlite-fts5", + "pypdf>=4.0.0", + "sentence-transformers>=2.5.0", ] [project.optional-dependencies] -embeddings = [ - "numpy>=1.24.0", - "sentence-transformers>=2.5.0", +server = [ + "fastapi>=0.110.0", + "uvicorn>=0.29.0", ] -openai = ["openai>=1.0.0"] ocr = [ "pytesseract>=0.3.10", "Pillow>=10.0.0", ] -server = [ - "fastapi>=0.110.0", - "uvicorn[standard]>=0.29.0", +rerank = [ + "sentence-transformers>=2.5.0", ] dev = [ "pytest>=8.1.0", "pytest-cov>=4.1.0", - "mypy>=1.10.0", + "mypy>=1.8.0", "ruff>=0.4.0", "black>=24.4.0", - "build>=1.2.1", - "twine>=5.0.0", + "isort>=5.13.2", + "pre-commit>=3.6.0", ] [project.urls] -Homepage = "https://github.com/example/raglite-sqlite" -Repository = "https://github.com/example/raglite-sqlite" +Homepage = "https://github.com/mmprotest/raglite-sqlite" +Repository = "https://github.com/mmprotest/raglite-sqlite" [project.scripts] -raglite = "raglite_sqlite.cli:app" +raglite = "raglite.cli:app" + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.setuptools.package-data] +raglite = ["schema.sql"] [tool.black] line-length = 100 target-version = ["py310"] +[tool.isort] +profile = "black" +known_first_party = ["raglite"] + [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 = [] +mypy_path = ["src"] diff --git a/raglite_sqlite/__init__.py b/raglite_sqlite/__init__.py deleted file mode 100644 index db07740..0000000 --- a/raglite_sqlite/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -"""RagLite SQLite package.""" - -from .api import RagLite - -try: # pragma: no cover - optional dependency - from .server import create_app -except Exception: # pragma: no cover - FastAPI not installed - create_app = None # type: ignore[assignment] - -__all__ = ["RagLite", "create_app"] diff --git a/raglite_sqlite/adapters/__init__.py b/raglite_sqlite/adapters/__init__.py deleted file mode 100644 index d3dbdcd..0000000 --- a/raglite_sqlite/adapters/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""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 deleted file mode 100644 index 66f5f11..0000000 --- a/raglite_sqlite/adapters/langchain.py +++ /dev/null @@ -1,45 +0,0 @@ -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 deleted file mode 100644 index b17ef47..0000000 --- a/raglite_sqlite/adapters/llamaindex.py +++ /dev/null @@ -1,33 +0,0 @@ -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 deleted file mode 100644 index dc63239..0000000 --- a/raglite_sqlite/api.py +++ /dev/null @@ -1,233 +0,0 @@ -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 .embeddings.hash_backend import HashingBackend -from .parsers.csv import CSVParser -from .parsers.docx import DocxParser -from .parsers.html import HTMLParser -from .parsers.image import ImageParser -from .parsers.json import JSONParser -from .parsers.md import MarkdownParser -from .parsers.pdf import PDFParser -from .parsers.pptx import PptxParser -from .parsers.txt import TextParser -from .search import assemble_results, bm25_search, hybrid_fuse, vector_search -from .typing import SearchResult -from .rerank import Reranker, get_reranker -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(), - "application/json": JSONParser(), - "application/vnd.openxmlformats-officedocument.presentationml.presentation": PptxParser(), - "image/png": ImageParser(), - "image/jpeg": ImageParser(), - "image/tiff": ImageParser(), -} - - -def _default_backend() -> EmbeddingBackend: - return HashingBackend() - - -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, - reranker: Reranker | str | None = None, - reranker_options: dict[str, object] | None = None, - ) -> 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) - results = assemble_results( - self.db, fused, k, with_snippets=with_snippets, max_per_doc=max_per_doc - ) - reranker_instance: Reranker | None - if isinstance(reranker, str): - try: - reranker_instance = get_reranker(reranker, **(reranker_options or {})) - except KeyError as exc: - raise ValueError(f"Unknown reranker '{reranker}'") from exc - else: - reranker_instance = reranker - if reranker_instance is not None: - results = list(reranker_instance.rerank(query, results)) - return results - - 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 deleted file mode 100644 index fed59f0..0000000 --- a/raglite_sqlite/chunking.py +++ /dev/null @@ -1,73 +0,0 @@ -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 deleted file mode 100644 index 033825e..0000000 --- a/raglite_sqlite/cli.py +++ /dev/null @@ -1,185 +0,0 @@ -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(backend: str, model: Optional[str]): - backend_name = (backend or "hash").lower() - if backend_name == "hash": - from .embeddings.hash_backend import HashingBackend - - return HashingBackend() - if backend_name in {"sentence-transformers", "st"}: - try: - from .embeddings.sentence_transformers_backend import SentenceTransformersBackend - except ImportError as exc: # pragma: no cover - optional dependency missing - raise typer.BadParameter( - "Sentence Transformers backend requires raglite-sqlite[embeddings]" - ) from exc - - return SentenceTransformersBackend( - model_name=model or "sentence-transformers/all-MiniLM-L6-v2" - ) - raise typer.BadParameter(f"Unknown embedding backend '{backend}'") - - -@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"), - backend: str = typer.Option("hash", help="Embedding backend to use (hash or sentence-transformers)"), - 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_impl = get_backend(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_impl, - 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"), - reranker: Optional[str] = typer.Option(None, help="Name of reranker to apply"), - reranker_option: Optional[list[str]] = typer.Option( - None, "--reranker-option", help="Reranker option 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 - reranker_opts: dict[str, object] | None = None - if reranker_option: - reranker_opts = {} - for item in reranker_option: - if "=" not in item: - raise typer.BadParameter("Reranker options must be in key=value format") - key, value = item.split("=", 1) - reranker_opts[key] = value - try: - results = rag.search( - text, - k=k, - hybrid_weight=hybrid, - max_per_doc=max_per_doc, - filters=filter_dict, - reranker=reranker, - reranker_options=reranker_opts, - ) - except ValueError as exc: - raise typer.BadParameter(str(exc)) from exc - 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") - - -@app.command() -def serve( - db: Path = typer.Option(..., help="Database path"), - host: str = typer.Option("127.0.0.1", help="Host to bind"), - port: int = typer.Option(8000, help="Port to bind"), - reload: bool = typer.Option(False, help="Enable auto-reload"), -) -> None: - """Run the optional REST server.""" - - try: - from .server import create_app - except ImportError as exc: # pragma: no cover - optional dependency missing - raise typer.BadParameter( - "The REST server requires FastAPI. Install raglite-sqlite[server]." - ) from exc - try: - import uvicorn - except ImportError as exc: # pragma: no cover - optional dependency missing - raise typer.BadParameter( - "Running the REST server requires uvicorn. Install raglite-sqlite[server]." - ) from exc - - app_instance = create_app(str(db)) - uvicorn.run(app_instance, host=host, port=port, reload=reload) diff --git a/raglite_sqlite/db.py b/raglite_sqlite/db.py deleted file mode 100644 index ec9686a..0000000 --- a/raglite_sqlite/db.py +++ /dev/null @@ -1,263 +0,0 @@ -from __future__ import annotations - -import sqlite3 -from array import array -from pathlib import Path -from typing import Iterable, Iterator, 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 iter_vectors(self, model_name: str | None = None) -> Iterator[tuple[str, list[float]]]: - 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") - for row in cur: - data = array("f") - data.frombytes(row["vec"]) - yield row["chunk_id"], list(data) - - 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 deleted file mode 100644 index 90db5b3..0000000 --- a/raglite_sqlite/embeddings/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Embedding backends for RagLite.""" - -from .hash_backend import HashingBackend -from .sentence_transformers_backend import SentenceTransformersBackend - -__all__ = ["HashingBackend", "SentenceTransformersBackend"] diff --git a/raglite_sqlite/embeddings/base.py b/raglite_sqlite/embeddings/base.py deleted file mode 100644 index 072ce24..0000000 --- a/raglite_sqlite/embeddings/base.py +++ /dev/null @@ -1,33 +0,0 @@ -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/hash_backend.py b/raglite_sqlite/embeddings/hash_backend.py deleted file mode 100644 index 8915786..0000000 --- a/raglite_sqlite/embeddings/hash_backend.py +++ /dev/null @@ -1,38 +0,0 @@ -from __future__ import annotations - -import hashlib -import math -from collections import Counter -from typing import Sequence - -from ..utils import normalize_text -from .base import EmbeddingBackend - - -class HashingBackend(EmbeddingBackend): - """Lightweight hashing-based embedding backend.""" - - def __init__(self, dim: int = 384) -> None: - self.dim = dim - self.model_name = f"hashing-{dim}" - - def embed_texts(self, texts: Sequence[str], model_name: str | None = None) -> Sequence[Sequence[float]]: - return [self._embed(text) for text in texts] - - def _embed(self, text: str) -> list[float]: - tokens = [token for token in normalize_text(text).split() if token] - if not tokens: - return [0.0] * self.dim - counts = Counter(tokens) - vector = [0.0] * self.dim - for token, weight in counts.items(): - index = self._bucket(token) - vector[index] += float(weight) - norm = math.sqrt(sum(value * value for value in vector)) - if norm == 0.0: - return vector - return [value / norm for value in vector] - - def _bucket(self, token: str) -> int: - digest = hashlib.blake2b(token.encode("utf-8"), digest_size=16).digest() - return int.from_bytes(digest, "big") % self.dim diff --git a/raglite_sqlite/embeddings/openai_backend.py b/raglite_sqlite/embeddings/openai_backend.py deleted file mode 100644 index a49343d..0000000 --- a/raglite_sqlite/embeddings/openai_backend.py +++ /dev/null @@ -1,30 +0,0 @@ -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 deleted file mode 100644 index 9630316..0000000 --- a/raglite_sqlite/embeddings/sentence_transformers_backend.py +++ /dev/null @@ -1,26 +0,0 @@ -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 deleted file mode 100644 index 195426d..0000000 --- a/raglite_sqlite/parsers/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -"""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 -from .json import JSONParser -from .pptx import PptxParser -from .image import ImageParser - -__all__ = [ - "TextParser", - "MarkdownParser", - "HTMLParser", - "PDFParser", - "DocxParser", - "CSVParser", - "JSONParser", - "PptxParser", - "ImageParser", -] diff --git a/raglite_sqlite/parsers/base.py b/raglite_sqlite/parsers/base.py deleted file mode 100644 index 9a8f313..0000000 --- a/raglite_sqlite/parsers/base.py +++ /dev/null @@ -1,12 +0,0 @@ -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 deleted file mode 100644 index 32bfca3..0000000 --- a/raglite_sqlite/parsers/csv.py +++ /dev/null @@ -1,27 +0,0 @@ -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 deleted file mode 100644 index 5d2756a..0000000 --- a/raglite_sqlite/parsers/docx.py +++ /dev/null @@ -1,20 +0,0 @@ -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 deleted file mode 100644 index 8d59a1c..0000000 --- a/raglite_sqlite/parsers/html.py +++ /dev/null @@ -1,40 +0,0 @@ -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/image.py b/raglite_sqlite/parsers/image.py deleted file mode 100644 index 8c47d1f..0000000 --- a/raglite_sqlite/parsers/image.py +++ /dev/null @@ -1,36 +0,0 @@ -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 ImageParser(BaseParser): - """Parse image files using pytesseract-based OCR.""" - - def __init__(self, default_lang: str = "eng") -> None: - self.default_lang = default_lang - - def parse(self, path: str, **options: object) -> Iterable[ParsedBlock]: - lang = str(options.get("lang", self.default_lang)) - config = options.get("tesseract_config") - try: - from PIL import Image - except ImportError as exc: # pragma: no cover - dependency missing - raise RuntimeError("Pillow is required for OCR parsing") from exc - try: - import pytesseract - except ImportError as exc: # pragma: no cover - dependency missing - raise RuntimeError("pytesseract is required for OCR parsing") from exc - image_path = Path(path) - with Image.open(image_path) as image: - try: - text = pytesseract.image_to_string(image, lang=lang, config=config) - except pytesseract.pytesseract.TesseractNotFoundError as exc: # pragma: no cover - environment specific - raise RuntimeError( - "Tesseract OCR binary not found. Install it or adjust TESSDATA_PREFIX." - ) from exc - yield ParsedBlock(text=normalize_text(text), section=None) diff --git a/raglite_sqlite/parsers/json.py b/raglite_sqlite/parsers/json.py deleted file mode 100644 index be159df..0000000 --- a/raglite_sqlite/parsers/json.py +++ /dev/null @@ -1,20 +0,0 @@ -from __future__ import annotations - -import json -from pathlib import Path -from typing import Iterable - -from .base import BaseParser -from ..typing import ParsedBlock -from ..utils import normalize_text - - -class JSONParser(BaseParser): - """Parse JSON documents by rendering them as a stable, human-readable string.""" - - def parse(self, path: str, **options: object) -> Iterable[ParsedBlock]: - indent = int(options.get("indent", 2) or 0) - sort_keys = bool(options.get("sort_keys", True)) - data = json.loads(Path(path).read_text(encoding="utf-8")) - text = json.dumps(data, indent=indent or None, sort_keys=sort_keys, ensure_ascii=False) - yield ParsedBlock(text=normalize_text(text), section=None) diff --git a/raglite_sqlite/parsers/md.py b/raglite_sqlite/parsers/md.py deleted file mode 100644 index f056490..0000000 --- a/raglite_sqlite/parsers/md.py +++ /dev/null @@ -1,39 +0,0 @@ -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 deleted file mode 100644 index f6e5fc4..0000000 --- a/raglite_sqlite/parsers/pdf.py +++ /dev/null @@ -1,20 +0,0 @@ -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/pptx.py b/raglite_sqlite/parsers/pptx.py deleted file mode 100644 index 20dac7c..0000000 --- a/raglite_sqlite/parsers/pptx.py +++ /dev/null @@ -1,45 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from typing import Iterable -from xml.etree import ElementTree -from zipfile import ZipFile - -from .base import BaseParser -from ..typing import ParsedBlock -from ..utils import normalize_text - - -class PptxParser(BaseParser): - """Lightweight PPTX parser that extracts text from slide XML payloads.""" - - SLIDE_PREFIX = "ppt/slides/" - SLIDE_SUFFIX = ".xml" - - def parse(self, path: str, **options: object) -> Iterable[ParsedBlock]: - pptx_path = Path(path) - texts: list[str] = [] - with ZipFile(pptx_path) as archive: - slide_names = sorted( - name - for name in archive.namelist() - if name.startswith(self.SLIDE_PREFIX) and name.endswith(self.SLIDE_SUFFIX) - ) - for slide_name in slide_names: - with archive.open(slide_name) as handle: - xml_bytes = handle.read() - try: - root = ElementTree.fromstring(xml_bytes) - except ElementTree.ParseError: - continue - namespaces = { - "a": "http://schemas.openxmlformats.org/drawingml/2006/main", - } - slide_text: list[str] = [] - for node in root.findall('.//a:t', namespaces): - if node.text: - slide_text.append(node.text) - if slide_text: - texts.append(" ".join(slide_text)) - combined = normalize_text("\n\n".join(texts)) - yield ParsedBlock(text=combined, section=None) diff --git a/raglite_sqlite/parsers/txt.py b/raglite_sqlite/parsers/txt.py deleted file mode 100644 index b409de6..0000000 --- a/raglite_sqlite/parsers/txt.py +++ /dev/null @@ -1,14 +0,0 @@ -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 deleted file mode 100644 index 85cae6a..0000000 --- a/raglite_sqlite/rerank.py +++ /dev/null @@ -1,65 +0,0 @@ -from __future__ import annotations - -from typing import Callable, 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 - - -RerankerFactory = Callable[..., Reranker] - -_RERANKER_REGISTRY: dict[str, RerankerFactory] = { - "none": lambda **_: NoopReranker(), -} - - -def register_reranker(name: str, factory: RerankerFactory) -> None: - """Register a reranker factory.""" - - _RERANKER_REGISTRY[name] = factory - - -def get_reranker(name: str, **options: object) -> Reranker: - """Instantiate a reranker by name.""" - - factory = _RERANKER_REGISTRY.get(name) - if factory is None: - raise KeyError(name) - return factory(**options) - - -class CrossEncoderReranker: - """Rerank using a sentence-transformers CrossEncoder.""" - - def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2", batch_size: int = 16) -> None: - try: - from sentence_transformers import CrossEncoder - except ImportError as exc: # pragma: no cover - dependency missing - raise RuntimeError( - "sentence-transformers is required for the cross-encoder reranker" - ) from exc - - self.model = CrossEncoder(model_name) - self.batch_size = batch_size - - def rerank(self, query: str, results: Iterable[SearchResult]) -> Iterable[SearchResult]: - result_list = list(results) - if not result_list: - return result_list - pairs = [(query, item["text"]) for item in result_list] - scores = self.model.predict(pairs, batch_size=self.batch_size) - for item, score in zip(result_list, scores): - item["rerank_score"] = float(score) - return sorted(result_list, key=lambda item: item.get("rerank_score", 0.0), reverse=True) - - -register_reranker("cross-encoder", lambda **opts: CrossEncoderReranker(**opts)) diff --git a/raglite_sqlite/schema.sql b/raglite_sqlite/schema.sql deleted file mode 100644 index 18ebd59..0000000 --- a/raglite_sqlite/schema.sql +++ /dev/null @@ -1,58 +0,0 @@ -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 deleted file mode 100644 index 74281f2..0000000 --- a/raglite_sqlite/search.py +++ /dev/null @@ -1,161 +0,0 @@ -from __future__ import annotations - -import heapq -import math -from typing import Dict, Iterable, List, Sequence, Tuple - -from .db import Database -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]]: - query_vectors = backend.embed_texts([query], model_name=model_name) - if not query_vectors: - return [] - query_vec = list(query_vectors[0]) - - def norm(vec: Sequence[float]) -> float: - return math.sqrt(sum(value * value for value in vec)) - - query_norm = norm(query_vec) - if math.isclose(query_norm, 0.0): - return [] - - limit = max(k * 5, 10) - heap: list[tuple[float, str]] = [] - for chunk_id, vector in db.iter_vectors(model_name=model_name): - row_norm = norm(vector) - if math.isclose(row_norm, 0.0): - score = 0.0 - else: - score = sum(a * b for a, b in zip(query_vec, vector)) / (row_norm * query_norm) - if len(heap) < limit: - heapq.heappush(heap, (score, chunk_id)) - continue - if score > heap[0][0]: - heapq.heapreplace(heap, (score, chunk_id)) - heap.sort(reverse=True) - return [(chunk_id, score) for score, chunk_id in heap] - - -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/server.py b/raglite_sqlite/server.py deleted file mode 100644 index 4e5b80f..0000000 --- a/raglite_sqlite/server.py +++ /dev/null @@ -1,130 +0,0 @@ -from __future__ import annotations - -from threading import RLock -from typing import Any, Dict, Optional - -from fastapi import FastAPI, HTTPException -from pydantic import BaseModel, Field - -from .api import RagLite -from .embeddings.base import EmbeddingBackend -from .rerank import Reranker, get_reranker - - -def _default_backend() -> EmbeddingBackend: - from .embeddings.hash_backend import HashingBackend - - return HashingBackend() - - -class QueryPayload(BaseModel): - query: str - k: int = 8 - hybrid_weight: float = Field(default=0.6, ge=0.0, le=1.0) - max_per_doc: int = 3 - filters: Optional[Dict[str, str]] = None - model_name: Optional[str] = None - with_snippets: bool = True - use_semantic: bool = True - reranker: Optional[str] = None - reranker_options: Optional[Dict[str, Any]] = None - - -class IndexPayload(BaseModel): - paths: list[str] - tags: Optional[str] = None - parser_opts: Optional[Dict[str, Any]] = None - chunker: str = "recursive" - chunk_size_tokens: int = 512 - chunk_overlap_tokens: int = 64 - model_name: Optional[str] = None - skip_unchanged: bool = True - recurse: bool = True - glob: Optional[str] = None - - -class DeletePayload(BaseModel): - doc_id: str - - -def create_app( - db_path: str, - *, - embedding_backend: EmbeddingBackend | None = None, - rerankers: dict[str, Reranker] | None = None, -) -> FastAPI: - """Create a FastAPI application exposing the RagLite API.""" - - rag = RagLite(db_path) - backend = embedding_backend or _default_backend() - reranker_cache = rerankers or {} - lock = RLock() - - app = FastAPI(title="RagLite SQLite", version="1.0") - - @app.get("/health") - def health() -> dict[str, str]: - return {"status": "ok"} - - @app.get("/stats") - def stats() -> dict[str, Any]: - with lock: - return rag.stats() - - @app.post("/query") - def query(payload: QueryPayload) -> dict[str, Any]: - with lock: - resolved_backend = backend if payload.use_semantic else None - reranker_instance: Reranker | None = None - if payload.reranker: - reranker_instance = reranker_cache.get(payload.reranker) - if reranker_instance is None: - try: - reranker_instance = get_reranker( - payload.reranker, **(payload.reranker_options or {}) - ) - except KeyError as exc: - raise HTTPException(status_code=400, detail=str(exc)) from exc - reranker_cache[payload.reranker] = reranker_instance - results = rag.search( - payload.query, - k=payload.k, - hybrid_weight=payload.hybrid_weight, - filters=payload.filters, - model_name=payload.model_name, - embedding_backend=resolved_backend, - max_per_doc=payload.max_per_doc, - with_snippets=payload.with_snippets, - reranker=reranker_instance, - ) - return {"results": results} - - @app.post("/index") - def index(payload: IndexPayload) -> dict[str, Any]: - with lock: - stats = rag.index( - payload.paths, - tags=payload.tags, - parser_opts=payload.parser_opts, - chunker=payload.chunker, - chunk_size_tokens=payload.chunk_size_tokens, - chunk_overlap_tokens=payload.chunk_overlap_tokens, - embedding_backend=backend, - model_name=payload.model_name, - skip_unchanged=payload.skip_unchanged, - recurse=payload.recurse, - glob=payload.glob, - ) - return stats - - @app.post("/delete") - def delete(payload: DeletePayload) -> dict[str, str]: - with lock: - rag.delete(payload.doc_id) - return {"status": "deleted", "doc_id": payload.doc_id} - - @app.on_event("shutdown") - def shutdown() -> None: - rag.close() - - return app diff --git a/raglite_sqlite/typing.py b/raglite_sqlite/typing.py deleted file mode 100644 index 84c7f94..0000000 --- a/raglite_sqlite/typing.py +++ /dev/null @@ -1,37 +0,0 @@ -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 - rerank_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 deleted file mode 100644 index c7aaee4..0000000 --- a/raglite_sqlite/utils.py +++ /dev/null @@ -1,110 +0,0 @@ -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", - ".json": "application/json", - ".yml": "application/x-yaml", - ".yaml": "application/x-yaml", - ".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation", - ".png": "image/png", - ".jpg": "image/jpeg", - ".jpeg": "image/jpeg", - ".tif": "image/tiff", - ".tiff": "image/tiff", - }.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: dict[Path, Path] = {} - for input_path in paths: - path = Path(input_path) - if path.is_dir(): - if glob: - iterator = path.rglob(glob) if recurse else path.glob(glob) - else: - iterator = path.rglob("*") if recurse else path.glob("*") - for child in iterator: - if child.is_file(): - candidates.setdefault(child.resolve(), child) - elif path.is_file(): - candidates.setdefault(path.resolve(), path) - return list(candidates.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/src/demo/demo.gif b/src/demo/demo.gif new file mode 100644 index 0000000..8f11eb9 --- /dev/null +++ b/src/demo/demo.gif @@ -0,0 +1 @@ +Demo GIF placeholder \ No newline at end of file diff --git a/src/demo/mini_corpus/doc1.md b/src/demo/mini_corpus/doc1.md new file mode 100644 index 0000000..2500755 --- /dev/null +++ b/src/demo/mini_corpus/doc1.md @@ -0,0 +1,3 @@ +# Document 1 + +This is sample document number 1 for the raglite demo. It contains a quick start guide paragraph mentioning SQLite and retrieval augmented generation. diff --git a/src/demo/mini_corpus/doc2.md b/src/demo/mini_corpus/doc2.md new file mode 100644 index 0000000..64cebf6 --- /dev/null +++ b/src/demo/mini_corpus/doc2.md @@ -0,0 +1,3 @@ +# Document 2 + +This is sample document number 2 for the raglite demo. It contains a quick start guide paragraph mentioning SQLite and retrieval augmented generation. diff --git a/src/demo/mini_corpus/doc3.md b/src/demo/mini_corpus/doc3.md new file mode 100644 index 0000000..7b2a7d3 --- /dev/null +++ b/src/demo/mini_corpus/doc3.md @@ -0,0 +1,3 @@ +# Document 3 + +This is sample document number 3 for the raglite demo. It contains a quick start guide paragraph mentioning SQLite and retrieval augmented generation. diff --git a/src/demo/mini_corpus/doc4.md b/src/demo/mini_corpus/doc4.md new file mode 100644 index 0000000..6a57c1e --- /dev/null +++ b/src/demo/mini_corpus/doc4.md @@ -0,0 +1,3 @@ +# Document 4 + +This is sample document number 4 for the raglite demo. It contains a quick start guide paragraph mentioning SQLite and retrieval augmented generation. diff --git a/src/demo/mini_corpus/doc5.md b/src/demo/mini_corpus/doc5.md new file mode 100644 index 0000000..5a1b5c9 --- /dev/null +++ b/src/demo/mini_corpus/doc5.md @@ -0,0 +1,3 @@ +# Document 5 + +This is sample document number 5 for the raglite demo. It contains a quick start guide paragraph mentioning SQLite and retrieval augmented generation. diff --git a/src/demo/mini_corpus/doc6.md b/src/demo/mini_corpus/doc6.md new file mode 100644 index 0000000..625dceb --- /dev/null +++ b/src/demo/mini_corpus/doc6.md @@ -0,0 +1,3 @@ +# Document 6 + +This is sample document number 6 for the raglite demo. It contains a quick start guide paragraph mentioning SQLite and retrieval augmented generation. diff --git a/src/demo/mini_corpus/faq.txt b/src/demo/mini_corpus/faq.txt new file mode 100644 index 0000000..e600793 --- /dev/null +++ b/src/demo/mini_corpus/faq.txt @@ -0,0 +1,5 @@ +Frequently Asked Questions about Raglite +======================================= + +Q: How do I run the quick start guide? +A: Use raglite init-db, raglite ingest, and raglite query. diff --git a/src/demo/mini_corpus/guide.html b/src/demo/mini_corpus/guide.html new file mode 100644 index 0000000..90a9777 --- /dev/null +++ b/src/demo/mini_corpus/guide.html @@ -0,0 +1 @@ +

Raglite Quickstart

The quick start guide demonstrates indexing markdown, text, and html files into SQLite.

diff --git a/src/demo/mini_corpus/placeholder.pdf b/src/demo/mini_corpus/placeholder.pdf new file mode 100644 index 0000000..9df6ba1 Binary files /dev/null and b/src/demo/mini_corpus/placeholder.pdf differ diff --git a/src/demo/run_demo.ps1 b/src/demo/run_demo.ps1 new file mode 100644 index 0000000..fa2b141 --- /dev/null +++ b/src/demo/run_demo.ps1 @@ -0,0 +1,8 @@ +param( + [string]$Db = "demo.db" +) + +raglite init-db --db $Db +raglite ingest --db $Db --path (Join-Path $PSScriptRoot 'mini_corpus') --embed-model debug +raglite query --db $Db --text "quick start guide" --k 3 +raglite serve --db $Db --host 127.0.0.1 --port 8000 diff --git a/src/demo/run_demo.sh b/src/demo/run_demo.sh new file mode 100755 index 0000000..9faa1f3 --- /dev/null +++ b/src/demo/run_demo.sh @@ -0,0 +1,8 @@ +#!/usr/bin/env bash +set -euo pipefail +DB=${1:-demo.db} + +raglite init-db --db "$DB" +raglite ingest --db "$DB" --path "$(dirname "$0")/mini_corpus" --embed-model debug +raglite query --db "$DB" --text "quick start guide" --k 3 +raglite serve --db "$DB" --host 127.0.0.1 --port 8000 diff --git a/src/examples/langchain_integration.py b/src/examples/langchain_integration.py new file mode 100644 index 0000000..1fe0c58 --- /dev/null +++ b/src/examples/langchain_integration.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from pathlib import Path +from typing import List + +from langchain.schema import Document +from langchain.schema.retriever import BaseRetriever + +from raglite.api import RagliteAPI, RagliteConfig + + +class RagliteRetriever(BaseRetriever): + def __init__(self, db_path: Path) -> None: + self.api = RagliteAPI(RagliteConfig(db_path)) + self.api.init_db() + + def _get_relevant_documents(self, query: str) -> List[Document]: + results = self.api.query(query, top_k=5) + return [Document(page_content=r.text, metadata=r.metadata) for r in results] + + async def _aget_relevant_documents(self, query: str) -> List[Document]: # pragma: no cover - async wrapper + return self._get_relevant_documents(query) + + +if __name__ == "__main__": + retriever = RagliteRetriever(Path("langchain.db")) + demo_dir = Path(__file__).resolve().parents[1] / "demo" / "mini_corpus" + retriever.api.index(demo_dir, strategy="fixed") + docs = retriever.get_relevant_documents("quick start guide") + for doc in docs: + print(doc) diff --git a/src/examples/llamaindex_integration.py b/src/examples/llamaindex_integration.py new file mode 100644 index 0000000..ef95d6e --- /dev/null +++ b/src/examples/llamaindex_integration.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any, List + +from llama_index.core.schema import TextNode +from llama_index.core.vector_stores.types import VectorStore + +from raglite.api import RagliteAPI, RagliteConfig + + +class RagliteVectorStore(VectorStore): + def __init__(self, db_path: Path) -> None: + self.api = RagliteAPI(RagliteConfig(db_path)) + self.api.init_db() + + def add(self, nodes: List[TextNode]) -> List[str]: # pragma: no cover - integration shim + temp_dir = Path("llamaindex_tmp") + temp_dir.mkdir(exist_ok=True) + ids = [] + for node in nodes: + path = temp_dir / f"node_{node.node_id}.txt" + path.write_text(node.get_content(), encoding="utf-8") + self.api.index(path.parent, strategy="fixed") + ids.append(node.node_id) + return ids + + def query(self, query: str, top_k: int = 5) -> List[TextNode]: # pragma: no cover - integration shim + results = self.api.query(query, top_k=top_k) + return [TextNode(text=r.text, id_=str(r.chunk_id), metadata=r.metadata) for r in results] + + +if __name__ == "__main__": + store = RagliteVectorStore(Path("llamaindex.db")) + demo_dir = Path(__file__).resolve().parents[1] / "demo" / "mini_corpus" + store.api.index(demo_dir, strategy="fixed") + print(store.query("quick start guide")[0]) diff --git a/src/examples/python_api_minimal.py b/src/examples/python_api_minimal.py new file mode 100644 index 0000000..db7d546 --- /dev/null +++ b/src/examples/python_api_minimal.py @@ -0,0 +1,17 @@ +from pathlib import Path + +from raglite.api import RagliteAPI, RagliteConfig + + +def main() -> None: + db = Path("example.db") + api = RagliteAPI(RagliteConfig(db)) + api.init_db() + demo_dir = Path(__file__).resolve().parents[1] / "demo" / "mini_corpus" + api.index(demo_dir, strategy="fixed", ocr=False) + for result in api.query("quick start guide", top_k=3): + print(f"score={result.score:.3f} text={result.text[:60]}...") + + +if __name__ == "__main__": + main() diff --git a/src/raglite/__init__.py b/src/raglite/__init__.py new file mode 100644 index 0000000..1f04fb1 --- /dev/null +++ b/src/raglite/__init__.py @@ -0,0 +1,14 @@ +"""raglite - local-first RAG toolkit built on SQLite.""" + +from .api import RagliteAPI, add_tags, index_corpus, init_db, query, stats +from .config import RagliteConfig + +__all__ = [ + "RagliteAPI", + "RagliteConfig", + "add_tags", + "index_corpus", + "init_db", + "query", + "stats", +] diff --git a/src/raglite/api.py b/src/raglite/api.py new file mode 100644 index 0000000..0a77c19 --- /dev/null +++ b/src/raglite/api.py @@ -0,0 +1,112 @@ +"""High level Python API.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional + +from .config import RagliteConfig +from .db import apply_migrations, temp_connection +from .ingest import IngestResult, ingest_path +from .search import SearchResult, hybrid_search + + +@dataclass +class RagliteAPI: + config: RagliteConfig + + @property + def db_path(self) -> Path: + return self.config.db_path + + def init_db(self) -> None: + with temp_connection(self.db_path) as conn: + apply_migrations(conn) + + def index(self, corpus_path: Path, *, strategy: str = "recursive", ocr: bool = False) -> IngestResult: + return ingest_path(self.db_path, corpus_path, config=self.config, strategy=strategy, ocr=ocr) + + def query( + self, + text: str, + *, + top_k: int = 10, + alpha: Optional[float] = None, + rerank: bool = False, + tags: Optional[Dict[str, str]] = None, + ) -> List[SearchResult]: + with temp_connection(self.db_path) as conn: + return hybrid_search( + conn, + text, + alpha=alpha if alpha is not None else self.config.alpha, + top_k=top_k, + embed_model=self.config.embed_model, + rerank=rerank, + tags=tags, + ) + + def add_tags(self, document_id: int, tags: Dict[str, str]) -> None: + with temp_connection(self.db_path) as conn: + conn.execute( + "UPDATE chunks SET tags_json = json_patch(COALESCE(tags_json, '{}'), ?) WHERE document_id = ?", + (json.dumps(tags), document_id), + ) + + def stats(self) -> Dict[str, int]: + with temp_connection(self.db_path) as conn: + doc_count = conn.execute("SELECT COUNT(*) FROM documents").fetchone()[0] + chunk_count = conn.execute("SELECT COUNT(*) FROM chunks").fetchone()[0] + embed_count = conn.execute("SELECT COUNT(*) FROM embeddings").fetchone()[0] + return {"documents": doc_count, "chunks": chunk_count, "embeddings": embed_count} + + +def init_db(db_path: Path | str) -> None: + RagliteAPI(RagliteConfig(Path(db_path))).init_db() + + +def index_corpus( + db_path: Path | str, + corpus_path: Path | str, + *, + strategy: str = "recursive", + ocr: bool = False, + embed_model: Optional[str] = None, +) -> IngestResult: + config = RagliteConfig(Path(db_path)) + if embed_model: + config.embed_model = embed_model + api = RagliteAPI(config) + api.init_db() + return api.index(Path(corpus_path), strategy=strategy, ocr=ocr) + + +def query( + db_path: Path | str, + text: str, + *, + top_k: int = 10, + alpha: float = None, + rerank: bool = False, + tags: Optional[Dict[str, str]] = None, + embed_model: Optional[str] = None, +) -> List[SearchResult]: + config = RagliteConfig(Path(db_path)) + if embed_model: + config.embed_model = embed_model + if alpha is not None: + config.alpha = alpha + api = RagliteAPI(config) + return api.query(text, top_k=top_k, alpha=alpha, rerank=rerank, tags=tags) + + +def add_tags(db_path: Path | str, document_id: int, tags: Dict[str, str]) -> None: + api = RagliteAPI(RagliteConfig(Path(db_path))) + api.add_tags(document_id, tags) + + +def stats(db_path: Path | str) -> Dict[str, int]: + api = RagliteAPI(RagliteConfig(Path(db_path))) + return api.stats() diff --git a/src/raglite/chunk.py b/src/raglite/chunk.py new file mode 100644 index 0000000..116ec60 --- /dev/null +++ b/src/raglite/chunk.py @@ -0,0 +1,81 @@ +"""Utilities for chunking text.""" + +from __future__ import annotations + +import re +from typing import List + +TOKEN_PATTERN = re.compile(r"\w+|[^\w\s]", re.UNICODE) + + +def estimate_tokens(text: str) -> int: + return len(TOKEN_PATTERN.findall(text)) + + +def split_fixed_tokens(text: str, *, max_tokens: int = 350, overlap: int = 50) -> List[str]: + tokens = TOKEN_PATTERN.findall(text) + if not tokens: + return [] + if max_tokens <= 0: + raise ValueError("max_tokens must be positive") + if overlap >= max_tokens: + raise ValueError("overlap must be smaller than max_tokens") + + chunks: List[str] = [] + start = 0 + while start < len(tokens): + end = min(len(tokens), start + max_tokens) + chunk_tokens = tokens[start:end] + chunks.append(" ".join(chunk_tokens).strip()) + if end == len(tokens): + break + start = end - overlap + return [c for c in chunks if c] + + +def split_recursive( + text: str, + *, + max_tokens: int = 350, + overlap: int = 50, + prefer_headings: bool = True, +) -> List[str]: + if not text.strip(): + return [] + sections = _split_by_headings(text) if prefer_headings else [text] + results: List[str] = [] + for section in sections: + tokens = TOKEN_PATTERN.findall(section) + if len(tokens) <= max_tokens: + results.append(section.strip()) + continue + results.extend(split_fixed_tokens(section, max_tokens=max_tokens, overlap=overlap)) + return [c for c in results if c] + + +def _split_by_headings(text: str) -> List[str]: + parts: List[str] = [] + current: List[str] = [] + for line in text.splitlines(): + if line.strip().startswith(("#", "=", "-")) and current: + parts.append("\n".join(current).strip()) + current = [line] + else: + current.append(line) + if current: + parts.append("\n".join(current).strip()) + return parts or [text] + + +def chunk_text( + text: str, + *, + strategy: str = "recursive", + max_tokens: int = 350, + overlap: int = 50, +) -> List[str]: + if strategy == "recursive": + return split_recursive(text, max_tokens=max_tokens, overlap=overlap) + if strategy == "fixed": + return split_fixed_tokens(text, max_tokens=max_tokens, overlap=overlap) + raise ValueError(f"Unknown strategy: {strategy}") diff --git a/src/raglite/cli.py b/src/raglite/cli.py new file mode 100644 index 0000000..1a2743a --- /dev/null +++ b/src/raglite/cli.py @@ -0,0 +1,101 @@ +"""Typer CLI for raglite.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path +from typing import Optional + +import typer + +from .api import RagliteAPI +from .config import RagliteConfig + +app = typer.Typer(help="Local-first RAG toolkit on SQLite") + + +def get_api(db: Path, embed_model: Optional[str] = None) -> RagliteAPI: + config = RagliteConfig(db_path=db) + if embed_model: + config.embed_model = embed_model + api = RagliteAPI(config) + api.init_db() + return api + + +@app.command() +def init_db(db: Path = typer.Option(Path("raglite.db"), help="Database path")) -> None: + api = get_api(db) + typer.echo(f"Initialised database at {api.db_path}") + + +@app.command() +def ingest( + path: Path = typer.Option(..., exists=True, file_okay=True, dir_okay=True), + db: Path = typer.Option(Path("raglite.db")), + strategy: str = typer.Option("recursive"), + embed_model: Optional[str] = typer.Option(None), + ocr: bool = typer.Option(False, help="Enable OCR for PDFs"), +) -> None: + api = get_api(db, embed_model) + result = api.index(path, strategy=strategy, ocr=ocr) + typer.echo(json.dumps(result.__dict__, indent=2)) + + +@app.command() +def query( + text: str, + db: Path = typer.Option(Path("raglite.db")), + k: int = typer.Option(5), + alpha: Optional[float] = typer.Option(None), + rerank: bool = typer.Option(False), +) -> None: + api = get_api(db) + results = api.query(text, top_k=k, alpha=alpha, rerank=rerank) + typer.echo(json.dumps([r.__dict__ for r in results], indent=2)) + + +@app.command() +def serve( + db: Path = typer.Option(Path("raglite.db")), + host: str = typer.Option("127.0.0.1"), + port: int = typer.Option(8000), +) -> None: + env = dict(os.environ) + env["RAGLITE_DB"] = str(db) + subprocess.run([sys.executable, "-m", "uvicorn", "raglite.server.app:app", "--host", host, "--port", str(port)], check=True, env=env) + + +@app.command() +def stats(db: Path = typer.Option(Path("raglite.db"))) -> None: + api = get_api(db) + typer.echo(json.dumps(api.stats(), indent=2)) + + +@app.command("self-test") +def self_test() -> None: + temp_db = Path("selftest.db") + if temp_db.exists(): + temp_db.unlink() + api = get_api(temp_db, embed_model="debug") + demo_path = Path(__file__).resolve().parents[1] / "demo" / "mini_corpus" + api.index(demo_path, strategy="fixed") + results = api.query("quick start guide", top_k=3) + typer.echo(json.dumps({"backend": "python", "results": [r.text for r in results]}, indent=2)) + + +@app.command() +def benchmark(db: Path = typer.Option(Path("raglite.db"))) -> None: + subprocess.run([sys.executable, "scripts/bench_basic.py", str(db)], check=True) + + +@app.command() +def eval(db: Path = typer.Option(Path("raglite.db"))) -> None: + subprocess.run([sys.executable, "scripts/eval_beir.py", str(db)], check=True) + + +if __name__ == "__main__": + app() diff --git a/src/raglite/config.py b/src/raglite/config.py new file mode 100644 index 0000000..91b020d --- /dev/null +++ b/src/raglite/config.py @@ -0,0 +1,64 @@ +"""Configuration utilities for raglite.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, Optional + +DEFAULT_EMBED_MODEL = "all-MiniLM-L6-v2" +DEFAULT_CHUNK_TOKENS = 350 +DEFAULT_CHUNK_OVERLAP = 50 +DEFAULT_ALPHA = 0.6 + + +def _default_cache_dir() -> Path: + return Path.home() / ".cache" / "raglite" + + +@dataclass(slots=True) +class RagliteConfig: + """Runtime configuration for raglite components.""" + + db_path: Path + embed_model: str = DEFAULT_EMBED_MODEL + cache_dir: Path = field(default_factory=_default_cache_dir) + chunk_tokens: int = DEFAULT_CHUNK_TOKENS + chunk_overlap: int = DEFAULT_CHUNK_OVERLAP + alpha: float = DEFAULT_ALPHA + rerank_model: Optional[str] = None + extra_metadata: Dict[str, str] = field(default_factory=dict) + + def ensure_cache_dir(self) -> Path: + self.cache_dir.mkdir(parents=True, exist_ok=True) + return self.cache_dir + + @classmethod + def from_env( + cls, + db_path: str | Path, + *, + embed_model: Optional[str] = None, + cache_dir: Optional[str | Path] = None, + alpha: Optional[float] = None, + ) -> "RagliteConfig": + path = Path(db_path) + cfg = cls(db_path=path) + if embed_model: + cfg.embed_model = embed_model + if cache_dir: + cfg.cache_dir = Path(cache_dir) + if alpha is not None: + cfg.alpha = alpha + return cfg + + def model_cache_path(self, model_name: Optional[str] = None) -> Path: + name = model_name or self.embed_model + return self.ensure_cache_dir() / "models" / name.replace("/", "_") + + def chunk_options(self) -> Dict[str, int]: + return {"max_tokens": self.chunk_tokens, "overlap": self.chunk_overlap} + + +def clamp_alpha(value: float) -> float: + return max(0.0, min(1.0, value)) diff --git a/src/raglite/db.py b/src/raglite/db.py new file mode 100644 index 0000000..7c18d2f --- /dev/null +++ b/src/raglite/db.py @@ -0,0 +1,55 @@ +"""SQLite database helpers for raglite.""" + +from __future__ import annotations + +import sqlite3 +from contextlib import contextmanager +from pathlib import Path +from typing import Iterator, Optional + +SCHEMA_PATH = Path(__file__).with_name("schema.sql") + + +PRAGMAS = { + "journal_mode": "WAL", + "synchronous": "NORMAL", + "temp_store": "MEMORY", + "mmap_size": 268435456, +} + + +class RagliteDatabaseError(RuntimeError): + """Raised for database specific errors.""" + + +def connect(db_path: Path | str, *, read_only: bool = False) -> sqlite3.Connection: + path = Path(db_path) + if read_only: + uri = f"file:{path}?mode=ro" + conn = sqlite3.connect(uri, uri=True) + else: + conn = sqlite3.connect(path) + conn.row_factory = sqlite3.Row + _apply_pragmas(conn) + return conn + + +def _apply_pragmas(conn: sqlite3.Connection) -> None: + for pragma, value in PRAGMAS.items(): + conn.execute(f"PRAGMA {pragma}={value}") + + +def apply_migrations(conn: sqlite3.Connection, schema_path: Optional[Path] = None) -> None: + path = schema_path or SCHEMA_PATH + sql = path.read_text(encoding="utf-8") + with conn: + conn.executescript(sql) + + +@contextmanager +def temp_connection(db_path: Path | str) -> Iterator[sqlite3.Connection]: + conn = connect(db_path) + try: + yield conn + finally: + conn.close() diff --git a/src/raglite/embed.py b/src/raglite/embed.py new file mode 100644 index 0000000..419164d --- /dev/null +++ b/src/raglite/embed.py @@ -0,0 +1,73 @@ +"""Embedding utilities.""" + +from __future__ import annotations + +import hashlib +import math +import struct +from array import array +from dataclasses import dataclass +from functools import lru_cache +from typing import Iterable, List, Sequence + + +@dataclass +class EmbeddingStore: + model_name: str + dimension: int + + def embed_many(self, texts: Sequence[str]) -> List[bytes]: # pragma: no cover - overridden + raise NotImplementedError + + +class DebugEmbeddingStore(EmbeddingStore): + def __init__(self, dimension: int = 256): + super().__init__(model_name="debug", dimension=dimension) + + def embed_many(self, texts: Sequence[str]) -> List[bytes]: + vectors: List[bytes] = [] + for text in texts: + vec = array("f", [0.0] * self.dimension) + for token in text.split(): + digest = hashlib.sha256(token.encode("utf-8")).digest() + for i in range(0, len(digest), 4): + idx = (digest[i] + i) % self.dimension + value = struct.unpack("!f", digest[i : i + 4])[0] + vec[idx] += value + norm = math.sqrt(sum(v * v for v in vec)) + if norm: + for i in range(self.dimension): + vec[i] /= norm + vectors.append(vec.tobytes()) + return vectors + + +class SentenceTransformerStore(EmbeddingStore): + def __init__(self, model_name: str) -> None: + self._model = _load_sentence_transformer(model_name) + super().__init__(model_name=model_name, dimension=self._model.get_sentence_embedding_dimension()) + + def embed_many(self, texts: Sequence[str]) -> List[bytes]: + embeddings = self._model.encode(list(texts), convert_to_numpy=False, normalize_embeddings=True) + return [array("f", vec).tobytes() for vec in embeddings] + + +@lru_cache(maxsize=4) +def get_embedding_store(model_name: str) -> EmbeddingStore: + if model_name.lower() in {"debug", "hash"}: + return DebugEmbeddingStore() + return SentenceTransformerStore(model_name) + + +def embedding_from_bytes(blob: bytes) -> array: + vec = array("f") + vec.frombytes(blob) + return vec + + +def _load_sentence_transformer(model_name: str): # pragma: no cover - heavy load + try: + from sentence_transformers import SentenceTransformer + except Exception as exc: # pragma: no cover + raise RuntimeError("sentence-transformers is required for embedding") from exc + return SentenceTransformer(model_name) diff --git a/src/raglite/ingest.py b/src/raglite/ingest.py new file mode 100644 index 0000000..9f3ccc4 --- /dev/null +++ b/src/raglite/ingest.py @@ -0,0 +1,201 @@ +"""Corpus ingestion utilities.""" + +from __future__ import annotations + +import mimetypes +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Iterator, List, Sequence, Tuple + +import sqlite3 + +from . import chunk as chunk_utils +from .config import RagliteConfig +from .db import apply_migrations, connect +from .embed import get_embedding_store + +try: + from readability import Document # type: ignore +except Exception: # pragma: no cover - optional dependency + Document = None # type: ignore + +try: + from bs4 import BeautifulSoup # type: ignore +except Exception: # pragma: no cover - optional dependency + BeautifulSoup = None # type: ignore + +try: + from pypdf import PdfReader # type: ignore +except Exception: # pragma: no cover - optional dependency + PdfReader = None # type: ignore + +try: # optional OCR + import pytesseract # type: ignore + from PIL import Image # type: ignore +except Exception: # pragma: no cover - optional dependency + pytesseract = None # type: ignore + Image = None # type: ignore + + +@dataclass +class IngestedChunk: + id: int + document_id: int + chunk_idx: int + text: str + tokens: int + + +@dataclass +class IngestResult: + documents: int + chunks: int + embeddings: int + + +TEXT_MIME_TYPES = { + "text/plain", + "text/markdown", + "text/html", + "application/json", +} + + +class UnsupportedDocument(RuntimeError): + pass + + +def discover_files(path: Path) -> Iterator[Path]: + if path.is_file(): + yield path + return + for root, _, files in os.walk(path): + for name in files: + yield Path(root) / name + + +def read_text_file(path: Path) -> str: + return path.read_text(encoding="utf-8", errors="ignore") + + +def read_html_file(path: Path) -> str: + html = path.read_text(encoding="utf-8", errors="ignore") + if Document is None or BeautifulSoup is None: + return html + doc = Document(html) + summary = doc.summary() + soup = BeautifulSoup(summary, "html.parser") + return soup.get_text(" ") + + +def read_pdf_file(path: Path, *, ocr: bool = False) -> str: + if PdfReader is None: + raise UnsupportedDocument("pypdf is required to read PDF files") + reader = PdfReader(str(path)) + texts: List[str] = [] + for page in reader.pages: + text = page.extract_text() or "" + if not text.strip() and ocr and pytesseract and Image: + try: + images = page.images + except AttributeError: # pragma: no cover + images = [] + for image in images: + try: + img = Image.open(image.data) # type: ignore[arg-type] + except Exception: # pragma: no cover + continue + texts.append(pytesseract.image_to_string(img)) + texts.append(text) + return "\n".join(filter(None, texts)) + + +def ingest_path( + db_path: Path, + corpus_path: Path, + *, + config: RagliteConfig, + strategy: str = "recursive", + ocr: bool = False, +) -> IngestResult: + conn = connect(db_path) + apply_migrations(conn) + embedding_store = get_embedding_store(config.embed_model) + total_docs = 0 + total_chunks = 0 + total_embeddings = 0 + + with conn: + for file_path in discover_files(corpus_path): + try: + doc_text, mime = load_file(file_path, ocr=ocr) + except UnsupportedDocument: + continue + if not doc_text.strip(): + continue + total_docs += 1 + doc_id = insert_document(conn, file_path, mime) + chunk_texts = chunk_utils.chunk_text( + doc_text, + strategy=strategy, + max_tokens=config.chunk_tokens, + overlap=config.chunk_overlap, + ) + chunks = insert_chunks(conn, doc_id, chunk_texts) + total_chunks += len(chunks) + embeddings = embedding_store.embed_many([c.text for c in chunks]) + total_embeddings += len(embeddings) + insert_embeddings(conn, chunks, embeddings, embedding_store.model_name, embedding_store.dimension) + conn.close() + return IngestResult(total_docs, total_chunks, total_embeddings) + + +def load_file(path: Path, *, ocr: bool = False) -> Tuple[str, str]: + mime, _ = mimetypes.guess_type(path.name) + mime = mime or "application/octet-stream" + if path.suffix.lower() in {".md", ".markdown"}: + text = read_text_file(path) + return text, "text/markdown" + if mime in {"text/plain", "application/json"}: + return read_text_file(path), mime + if mime == "text/html": + return read_html_file(path), mime + if path.suffix.lower() == ".pdf": + return read_pdf_file(path, ocr=ocr), "application/pdf" + raise UnsupportedDocument(f"Unsupported file: {path}") + + +def insert_document(conn: sqlite3.Connection, path: Path, mime: str) -> int: + cur = conn.execute( + "INSERT INTO documents(path, title, mime) VALUES (?, ?, ?)", + (str(path), path.stem, mime), + ) + return int(cur.lastrowid) + + +def insert_chunks(conn: sqlite3.Connection, document_id: int, chunk_texts: Sequence[str]) -> List[IngestedChunk]: + chunks: List[IngestedChunk] = [] + for idx, text in enumerate(chunk_texts): + tokens = chunk_utils.estimate_tokens(text) + cur = conn.execute( + "INSERT INTO chunks(document_id, chunk_idx, text, tokens) VALUES (?, ?, ?, ?)", + (document_id, idx, text, tokens), + ) + chunk_id = int(cur.lastrowid) + chunks.append(IngestedChunk(chunk_id, document_id, idx, text, tokens)) + return chunks + + +def insert_embeddings( + conn: sqlite3.Connection, + chunks: Sequence[IngestedChunk], + vectors: Sequence[bytes], + model: str, + dim: int, +) -> None: + for chunk, vector in zip(chunks, vectors): + conn.execute( + "INSERT INTO embeddings(chunk_id, model, dim, embedding) VALUES (?, ?, ?, ?)", + (chunk.id, model, dim, vector), + ) diff --git a/src/raglite/schema.sql b/src/raglite/schema.sql new file mode 100644 index 0000000..e09f089 --- /dev/null +++ b/src/raglite/schema.sql @@ -0,0 +1,49 @@ +PRAGMA foreign_keys = ON; + +CREATE TABLE IF NOT EXISTS documents ( + id INTEGER PRIMARY KEY, + path TEXT NOT NULL, + title TEXT, + mime TEXT, + created_at TEXT DEFAULT (datetime('now')), + meta_json TEXT DEFAULT '{}' +); + +CREATE TABLE IF NOT EXISTS chunks ( + id INTEGER PRIMARY KEY, + document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, + chunk_idx INTEGER NOT NULL, + text TEXT NOT NULL, + tokens INTEGER DEFAULT 0, + tags_json TEXT DEFAULT '{}' +); + +CREATE TABLE IF NOT EXISTS embeddings ( + id INTEGER PRIMARY KEY, + chunk_id INTEGER NOT NULL REFERENCES chunks(id) ON DELETE CASCADE, + model TEXT NOT NULL, + dim INTEGER NOT NULL, + embedding BLOB NOT NULL +); + +CREATE VIRTUAL TABLE IF NOT EXISTS chunk_fts USING fts5( + text, + content='chunks', + content_rowid='id' +); + +CREATE INDEX IF NOT EXISTS idx_chunks_document_id ON chunks(document_id); +CREATE INDEX IF NOT EXISTS idx_embeddings_chunk_model ON embeddings(chunk_id, model); + +CREATE TRIGGER IF NOT EXISTS chunks_ai AFTER INSERT ON chunks BEGIN + INSERT INTO chunk_fts(rowid, text) VALUES (new.id, new.text); +END; + +CREATE TRIGGER IF NOT EXISTS chunks_ad AFTER DELETE ON chunks BEGIN + INSERT INTO chunk_fts(chunk_fts, rowid, text) VALUES('delete', old.id, old.text); +END; + +CREATE TRIGGER IF NOT EXISTS chunks_au AFTER UPDATE ON chunks BEGIN + INSERT INTO chunk_fts(chunk_fts, rowid, text) VALUES('delete', old.id, old.text); + INSERT INTO chunk_fts(rowid, text) VALUES(new.id, new.text); +END; diff --git a/src/raglite/search.py b/src/raglite/search.py new file mode 100644 index 0000000..ae05936 --- /dev/null +++ b/src/raglite/search.py @@ -0,0 +1,111 @@ +"""Hybrid search implementation.""" + +from __future__ import annotations + +import json +import sqlite3 +from dataclasses import dataclass +from typing import Dict, List, Optional + +from .config import clamp_alpha +from .embed import embedding_from_bytes, get_embedding_store +from .vector import get_backend + + +@dataclass +class SearchResult: + chunk_id: int + document_id: int + score: float + text: str + metadata: Dict[str, str] + + +@dataclass +class RankedChunk: + chunk_id: int + score: float + + +def bm25(conn: sqlite3.Connection, query: str, *, k: int = 200) -> List[RankedChunk]: + cur = conn.execute( + "SELECT rowid, bm25(chunk_fts) as score FROM chunk_fts WHERE chunk_fts MATCH ? ORDER BY score LIMIT ?", + (query, k), + ) + return [RankedChunk(int(row[0]), float(row[1])) for row in cur.fetchall()] + + +def normalize_scores(scores: List[RankedChunk]) -> Dict[int, float]: + if not scores: + return {} + values = [c.score for c in scores] + max_score = max(values) + min_score = min(values) + if max_score == min_score: + return {c.chunk_id: 1.0 for c in scores} + return {c.chunk_id: (max_score - c.score) / (max_score - min_score) for c in scores} + + +def hybrid_search( + conn: sqlite3.Connection, + query: str, + *, + alpha: float = 0.6, + top_k: int = 10, + embed_model: str, + rerank: bool = False, + tags: Optional[Dict[str, str]] = None, +) -> List[SearchResult]: + alpha = clamp_alpha(alpha) + candidates = bm25(conn, query) + bm25_norm = normalize_scores(candidates) + + backend = get_backend(conn) + embedding_store = get_embedding_store(embed_model) + query_vec = embedding_from_bytes(embedding_store.embed_many([query])[0]) + vector_results = backend.search( + conn, + query_vec, + top_n=max(top_k, len(candidates)) or top_k, + prefilter_ids=[c.chunk_id for c in candidates] if candidates else None, + ) + vector_scores = {c.chunk_id: c.score for c in vector_results} + + combined: List[SearchResult] = [] + ordered_ids = [c.chunk_id for c in candidates] or [c.chunk_id for c in vector_results] + seen = set() + for chunk_id in ordered_ids: + if chunk_id in seen: + continue + seen.add(chunk_id) + bm_score = bm25_norm.get(chunk_id, 0.0) + vec_score = vector_scores.get(chunk_id, 0.0) + score = alpha * bm_score + (1 - alpha) * vec_score + chunk_row = conn.execute( + "SELECT c.id, c.document_id, c.text, c.tags_json, d.meta_json FROM chunks c JOIN documents d ON d.id = c.document_id WHERE c.id = ?", + (chunk_id,), + ).fetchone() + if not chunk_row: + continue + tags_json = json.loads(chunk_row[3] or "{}") + if tags and not _tags_match(tags, tags_json): + continue + metadata = json.loads(chunk_row[4] or "{}") + combined.append( + SearchResult( + chunk_id=int(chunk_row[0]), + document_id=int(chunk_row[1]), + score=score, + text=str(chunk_row[2]), + metadata=metadata | {"tags": tags_json}, + ) + ) + combined.sort(key=lambda item: item.score, reverse=True) + return combined[:top_k] + + +def _tags_match(required: Dict[str, str], existing: Dict[str, str]) -> bool: + for key, value in required.items(): + if existing.get(key) != value: + return False + return True diff --git a/src/raglite/server/__init__.py b/src/raglite/server/__init__.py new file mode 100644 index 0000000..76e4a1d --- /dev/null +++ b/src/raglite/server/__init__.py @@ -0,0 +1,5 @@ +"""Server package for raglite.""" + +from .app import app, create_app + +__all__ = ["app", "create_app"] diff --git a/src/raglite/server/app.py b/src/raglite/server/app.py new file mode 100644 index 0000000..9998747 --- /dev/null +++ b/src/raglite/server/app.py @@ -0,0 +1,84 @@ +"""FastAPI application for raglite.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Dict, Optional + +try: + from fastapi import FastAPI, HTTPException + from pydantic import BaseModel +except Exception as exc: # pragma: no cover - optional dependency + raise RuntimeError("Install raglite[server] to use the FastAPI app") from exc + +from ..api import RagliteAPI +from ..config import RagliteConfig + + +class QueryRequest(BaseModel): + text: str + k: int = 10 + alpha: Optional[float] = None + rerank: bool = False + tags: Optional[Dict[str, str]] = None + + +class IngestRequest(BaseModel): + path: str + strategy: str = "recursive" + ocr: bool = False + + +def create_app(db_path: str | Path) -> FastAPI: + config = RagliteConfig(Path(db_path)) + api = RagliteAPI(config) + api.init_db() + + app = FastAPI(title="raglite", version="0.1.0") + + @app.get("/health") + def health() -> Dict[str, str]: + return {"status": "ok"} + + @app.get("/stats") + def stats() -> Dict[str, int]: + return api.stats() + + @app.post("/query") + def query(request: QueryRequest): + results = api.query( + request.text, + top_k=request.k, + alpha=request.alpha, + rerank=request.rerank, + tags=request.tags, + ) + return { + "results": [ + { + "chunk_id": r.chunk_id, + "document_id": r.document_id, + "score": r.score, + "text": r.text, + "metadata": r.metadata, + } + for r in results + ] + } + + @app.post("/ingest") + def ingest(request: IngestRequest): + corpus_path = Path(request.path) + if not corpus_path.exists(): + raise HTTPException(status_code=404, detail="Path not found") + result = api.index(corpus_path, strategy=request.strategy, ocr=request.ocr) + return { + "documents": result.documents, + "chunks": result.chunks, + "embeddings": result.embeddings, + } + + return app + + +app = create_app(Path("raglite.db")) diff --git a/src/raglite/vector/__init__.py b/src/raglite/vector/__init__.py new file mode 100644 index 0000000..5e106c6 --- /dev/null +++ b/src/raglite/vector/__init__.py @@ -0,0 +1,5 @@ +"""Vector backend selection.""" + +from .backend import Backend, get_backend + +__all__ = ["Backend", "get_backend"] diff --git a/src/raglite/vector/backend.py b/src/raglite/vector/backend.py new file mode 100644 index 0000000..7fdf187 --- /dev/null +++ b/src/raglite/vector/backend.py @@ -0,0 +1,33 @@ +"""Vector backend protocol and factory.""" + +from __future__ import annotations + +import sqlite3 +from typing import Iterable, List, Optional, Protocol + +from array import array + +from .python_fallback import PythonFallbackBackend +from .sqlite_ext import SQLiteExtensionBackend +from .types import Candidate + + +class Backend(Protocol): + name: str + + def search( + self, + conn: sqlite3.Connection, + query_vector: array, + *, + top_n: int, + prefilter_ids: Optional[Iterable[int]] = None, + ) -> List[Candidate]: + ... + + +def get_backend(conn: sqlite3.Connection) -> Backend: + backend = SQLiteExtensionBackend.create(conn) + if backend is not None: + return backend + return PythonFallbackBackend() diff --git a/src/raglite/vector/python_fallback.py b/src/raglite/vector/python_fallback.py new file mode 100644 index 0000000..4ad4c9e --- /dev/null +++ b/src/raglite/vector/python_fallback.py @@ -0,0 +1,57 @@ +"""Python fallback vector backend.""" + +from __future__ import annotations + +import sqlite3 +from dataclasses import dataclass +from typing import Iterable, List, Optional + +from ..embed import embedding_from_bytes +from .types import Candidate + + +@dataclass +class PythonFallbackBackend: + name: str = "python" + + def search( + self, + conn: sqlite3.Connection, + query_vector, + *, + top_n: int, + prefilter_ids: Optional[Iterable[int]] = None, + ) -> List[Candidate]: + ids = list(prefilter_ids) if prefilter_ids is not None else self._all_chunk_ids(conn) + if not ids: + return [] + placeholders = ",".join("?" for _ in ids) + cur = conn.execute( + f"SELECT chunk_id, embedding FROM embeddings WHERE chunk_id IN ({placeholders})", + ids, + ) + query_vec = query_vector + query_norm = self._norm(query_vec) + if not query_norm: + return [] + scored: List[Candidate] = [] + for row in cur.fetchall(): + vector = embedding_from_bytes(row[1]) + norm = self._norm(vector) + if not norm: + continue + score = float(self._dot(query_vec, vector) / (query_norm * norm)) + scored.append(Candidate(int(row[0]), score)) + scored.sort(key=lambda c: c.score, reverse=True) + return scored[:top_n] + + def _all_chunk_ids(self, conn: sqlite3.Connection) -> List[int]: + cur = conn.execute("SELECT DISTINCT chunk_id FROM embeddings") + return [int(row[0]) for row in cur.fetchall()] + + def _dot(self, a, b) -> float: + length = min(len(a), len(b)) + return float(sum(a[i] * b[i] for i in range(length))) + + def _norm(self, vec) -> float: + return float(sum(v * v for v in vec) ** 0.5) diff --git a/src/raglite/vector/sqlite_ext.py b/src/raglite/vector/sqlite_ext.py new file mode 100644 index 0000000..8368029 --- /dev/null +++ b/src/raglite/vector/sqlite_ext.py @@ -0,0 +1,64 @@ +"""Backend powered by sqlite extensions.""" + +from __future__ import annotations + +import sqlite3 +from dataclasses import dataclass, field +from typing import Iterable, List, Optional + +from .python_fallback import PythonFallbackBackend +from .types import Candidate + + +EXTENSION_NAMES = ["sqlite_vec", "vec0", "sqlite_vss"] + + +@dataclass +class SQLiteExtensionBackend: + name: str = "sqlite-extension" + _fallback: PythonFallbackBackend = field(default_factory=PythonFallbackBackend) + + @classmethod + def create(cls, conn: sqlite3.Connection) -> Optional["SQLiteExtensionBackend"]: + if not hasattr(conn, "enable_load_extension"): + return None + for ext in EXTENSION_NAMES: + try: + conn.enable_load_extension(True) + conn.load_extension(ext) + conn.enable_load_extension(False) + return cls() + except sqlite3.OperationalError: + continue + finally: + if hasattr(conn, "enable_load_extension"): + try: + conn.enable_load_extension(False) + except sqlite3.OperationalError: + pass + return None + + def search( + self, + conn: sqlite3.Connection, + query_vector, + *, + top_n: int, + prefilter_ids: Optional[Iterable[int]] = None, + ) -> List[Candidate]: + try: + cur = conn.execute( + "SELECT chunk_id, score FROM embedding_search(?, ?) ORDER BY score DESC LIMIT ?", + (getattr(query_vector, "tobytes", lambda: bytes(query_vector))(), len(query_vector), top_n), + ) + rows = cur.fetchall() + if rows: + return [Candidate(int(row[0]), float(row[1])) for row in rows] + except sqlite3.OperationalError: + pass + return self._fallback.search( + conn, + query_vector, + top_n=top_n, + prefilter_ids=prefilter_ids, + ) diff --git a/src/raglite/vector/types.py b/src/raglite/vector/types.py new file mode 100644 index 0000000..42be132 --- /dev/null +++ b/src/raglite/vector/types.py @@ -0,0 +1,7 @@ +from dataclasses import dataclass + + +@dataclass +class Candidate: + chunk_id: int + score: float diff --git a/src/scripts/bench_basic.py b/src/scripts/bench_basic.py new file mode 100644 index 0000000..bbf942b --- /dev/null +++ b/src/scripts/bench_basic.py @@ -0,0 +1,54 @@ +"""Basic benchmarking harness for raglite.""" + +from __future__ import annotations + +import json +import time +from pathlib import Path + +from raglite.api import RagliteAPI, RagliteConfig + + +def main(db_path: str | None = None) -> None: + db = Path(db_path or "bench.db") + if db.exists(): + db.unlink() + api = RagliteAPI(RagliteConfig(db)) + api.init_db() + corpus_dir = db.parent / "bench_corpus" + corpus_dir.mkdir(exist_ok=True) + doc = "Raglite benchmarking quick start guide with SQLite hybrid search." + for i in range(200): + (corpus_dir / f"doc_{i}.txt").write_text(f"{doc} repetition {i}", encoding="utf-8") + start = time.time() + api.index(corpus_dir, strategy="fixed") + ingest_time = time.time() - start + query_latencies = [] + for _ in range(20): + q_start = time.time() + api.query("quick start guide", top_k=3) + query_latencies.append(time.time() - q_start) + stats = { + "ingest_seconds": ingest_time, + "query_latency_p50": percentile(query_latencies, 50), + "query_latency_p95": percentile(query_latencies, 95), + } + print(json.dumps(stats, indent=2)) + + +def percentile(values, percent): + values = sorted(values) + k = (len(values) - 1) * percent / 100 + f = int(k) + c = min(f + 1, len(values) - 1) + if f == c: + return values[int(k)] + d0 = values[f] * (c - k) + d1 = values[c] * (k - f) + return d0 + d1 + + +if __name__ == "__main__": + import sys + + main(sys.argv[1] if len(sys.argv) > 1 else None) diff --git a/src/scripts/eval_beir.py b/src/scripts/eval_beir.py new file mode 100644 index 0000000..32af1f4 --- /dev/null +++ b/src/scripts/eval_beir.py @@ -0,0 +1,45 @@ +"""Evaluate raglite on a small synthetic dataset.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import List + +from raglite.api import RagliteAPI, RagliteConfig + +DATA = [ + ("doc1", "Neural retrieval quick start guide"), + ("doc2", "SQLite hybrid search tutorial"), + ("doc3", "Benchmarking local rag systems"), +] + +QUERIES = { + "quickstart": "quick start guide", + "sqlite": "sqlite hybrid", + "benchmark": "benchmark rag", +} + + +def main(db_path: str | None = None) -> None: + db = Path(db_path or "eval.db") + if db.exists(): + db.unlink() + api = RagliteAPI(RagliteConfig(db)) + api.init_db() + corpus_dir = db.parent / "eval_corpus" + corpus_dir.mkdir(exist_ok=True) + for name, text in DATA: + (corpus_dir / f"{name}.txt").write_text(text, encoding="utf-8") + api.index(corpus_dir, strategy="fixed") + results = {} + for name, query in QUERIES.items(): + hits = api.query(query, top_k=3) + results[name] = [hit.chunk_id for hit in hits] + print(json.dumps(results, indent=2)) + + +if __name__ == "__main__": + import sys + + main(sys.argv[1] if len(sys.argv) > 1 else None) diff --git a/src/scripts/load_test.sh b/src/scripts/load_test.sh new file mode 100755 index 0000000..ddba25a --- /dev/null +++ b/src/scripts/load_test.sh @@ -0,0 +1,8 @@ +#!/usr/bin/env bash +set -euo pipefail +HOST=${1:-http://127.0.0.1:8000} +for i in {1..5}; do + curl -s -X POST "$HOST/query" -H 'Content-Type: application/json' -d '{"text": "quick start guide", "k": 3}' >/dev/null + echo "Completed query $i" + sleep 1 +done diff --git a/tests/conftest.py b/tests/conftest.py index 9decda4..511da64 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,23 +1,7 @@ -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() +PROJECT_ROOT = Path(__file__).resolve().parents[1] +SRC = PROJECT_ROOT / "src" +if str(SRC) not in sys.path: + sys.path.insert(0, str(SRC)) diff --git a/tests/data/sample.csv b/tests/data/sample.csv deleted file mode 100644 index fd84fc2..0000000 --- a/tests/data/sample.csv +++ /dev/null @@ -1,3 +0,0 @@ -name,description -alpha,First row content -beta,Second row information diff --git a/tests/data/sample.html b/tests/data/sample.html deleted file mode 100644 index 6372c5c..0000000 --- a/tests/data/sample.html +++ /dev/null @@ -1,8 +0,0 @@ - -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 deleted file mode 100644 index 8bd8d89..0000000 --- a/tests/data/sample.md +++ /dev/null @@ -1,11 +0,0 @@ ---- -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 deleted file mode 100644 index f227828..0000000 --- a/tests/data/sample.txt +++ /dev/null @@ -1,2 +0,0 @@ -This is a sample text document. -It contains a few lines about testing RagLite. diff --git a/tests/test_chunk.py b/tests/test_chunk.py new file mode 100644 index 0000000..4a533fe --- /dev/null +++ b/tests/test_chunk.py @@ -0,0 +1,14 @@ +from raglite import chunk + + +def test_split_fixed_tokens_overlap(): + text = " ".join(str(i) for i in range(100)) + parts = chunk.split_fixed_tokens(text, max_tokens=20, overlap=5) + assert parts + assert all(len(chunk.TOKEN_PATTERN.findall(p)) <= 20 for p in parts) + + +def test_split_recursive_handles_headings(): + text = "# Title\ncontent\n# Heading\nmore content" + parts = chunk.split_recursive(text, max_tokens=10) + assert len(parts) == 2 diff --git a/tests/test_cli.py b/tests/test_cli.py deleted file mode 100644 index 792a80d..0000000 --- a/tests/test_cli.py +++ /dev/null @@ -1,44 +0,0 @@ -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 backend, 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_e2e_demo.py b/tests/test_e2e_demo.py new file mode 100644 index 0000000..a17a978 --- /dev/null +++ b/tests/test_e2e_demo.py @@ -0,0 +1,14 @@ +from pathlib import Path + +from raglite.api import RagliteAPI, RagliteConfig + + +def test_demo_corpus_query(tmp_path: Path): + db = tmp_path / "demo.db" + api = RagliteAPI(RagliteConfig(db_path=db, embed_model="debug")) + api.init_db() + demo_dir = Path(__file__).resolve().parents[1] / "src" / "demo" / "mini_corpus" + api.index(demo_dir, strategy="fixed") + results = api.query("quick start guide", top_k=5) + assert results + assert any("quick start" in r.text.lower() for r in results) diff --git a/tests/test_embed.py b/tests/test_embed.py new file mode 100644 index 0000000..9e842cc --- /dev/null +++ b/tests/test_embed.py @@ -0,0 +1,14 @@ +from raglite.embed import DebugEmbeddingStore, embedding_from_bytes, get_embedding_store + + +def test_debug_embedding_repeatable(): + store = DebugEmbeddingStore(dimension=32) + vec1 = embedding_from_bytes(store.embed_many(["hello world"])[0]) + vec2 = embedding_from_bytes(store.embed_many(["hello world"])[0]) + assert len(vec1) == 32 + assert list(vec1) == list(vec2) + + +def test_get_embedding_store_debug(): + store = get_embedding_store("debug") + assert isinstance(store, DebugEmbeddingStore) diff --git a/tests/test_ingest_and_search.py b/tests/test_ingest_and_search.py deleted file mode 100644 index 246b8f1..0000000 --- a/tests/test_ingest_and_search.py +++ /dev/null @@ -1,36 +0,0 @@ -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 diff --git a/tests/test_parsers.py b/tests/test_parsers.py deleted file mode 100644 index dcb91fc..0000000 --- a/tests/test_parsers.py +++ /dev/null @@ -1,65 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from zipfile import ZipFile - -import pytest - -from raglite_sqlite.parsers.image import ImageParser -from raglite_sqlite.parsers.json import JSONParser -from raglite_sqlite.parsers.pptx import PptxParser - - -def test_json_parser(tmp_path: Path) -> None: - data_path = tmp_path / "sample.json" - data_path.write_text("{" "\"title\": \"Hello\", \"items\": [1, 2]}", encoding="utf-8") - parser = JSONParser() - blocks = list(parser.parse(str(data_path))) - assert blocks - assert "Hello" in blocks[0]["text"] - assert "items" in blocks[0]["text"] - - -def test_pptx_parser(tmp_path: Path) -> None: - pptx_path = tmp_path / "sample.pptx" - slide_xml = """ - - - - - - Hello PPTX - - - - - - """ - with ZipFile(pptx_path, "w") as archive: - archive.writestr("ppt/slides/slide1.xml", slide_xml) - parser = PptxParser() - blocks = list(parser.parse(str(pptx_path))) - assert blocks - assert "Hello PPTX" in blocks[0]["text"] - - -def test_image_parser(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: - pytest.importorskip("PIL") - pytest.importorskip("pytesseract") - from PIL import Image - - image_path = tmp_path / "hello.png" - image = Image.new("RGB", (10, 10), color="white") - image.save(image_path) - - import pytesseract - - def fake_ocr(image, lang="eng", config=None): # type: ignore[no-untyped-def] - return "Hello OCR" - - monkeypatch.setattr(pytesseract, "image_to_string", fake_ocr) - parser = ImageParser() - blocks = list(parser.parse(str(image_path))) - assert blocks - assert blocks[0]["text"] == "Hello OCR" diff --git a/tests/test_rerank.py b/tests/test_rerank.py deleted file mode 100644 index 97b5eb3..0000000 --- a/tests/test_rerank.py +++ /dev/null @@ -1,40 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -from raglite_sqlite.rerank import register_reranker - - -class ReverseReranker: - def rerank(self, query: str, results): # type: ignore[override] - return list(reversed(list(results))) - - -def test_custom_reranker(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) - - class PrefixReranker: - def rerank(self, query: str, results): # type: ignore[override] - ordered = sorted(results, key=lambda item: item["doc_id"], reverse=True) - for idx, item in enumerate(ordered): - item["rerank_score"] = float(len(ordered) - idx) - return ordered - - results = rag.search("sample", embedding_backend=backend, reranker=PrefixReranker()) - assert results - assert results[0]["rerank_score"] >= results[-1]["rerank_score"] - - -def test_registry_reranker(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) - - register_reranker("reverse", lambda **_: ReverseReranker()) - baseline = rag.search("sample", embedding_backend=backend) - results = rag.search("sample", embedding_backend=backend, reranker="reverse") - assert results - if len(baseline) > 1: - assert results[0]["chunk_id"] == baseline[-1]["chunk_id"] diff --git a/tests/test_server.py b/tests/test_server.py deleted file mode 100644 index fc52d9f..0000000 --- a/tests/test_server.py +++ /dev/null @@ -1,24 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -import pytest - -pytest.importorskip("fastapi", reason="FastAPI is required for server tests") -from fastapi.testclient import TestClient - -from raglite_sqlite.server import create_app - - -def test_server_query_endpoint(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) - - app = create_app(str(db_path), embedding_backend=backend) - client = TestClient(app) - response = client.post("/query", json={"query": "sample", "use_semantic": False}) - assert response.status_code == 200 - data = response.json() - assert "results" in data - assert data["results"], "Expected at least one result from the REST API" diff --git a/tests/test_vector_fallback.py b/tests/test_vector_fallback.py new file mode 100644 index 0000000..a1927cf --- /dev/null +++ b/tests/test_vector_fallback.py @@ -0,0 +1,28 @@ +import sqlite3 +from pathlib import Path + +from raglite.embed import DebugEmbeddingStore, embedding_from_bytes +from raglite.vector.python_fallback import PythonFallbackBackend + + +def test_python_fallback_similarity(tmp_path: Path): + conn = sqlite3.connect(tmp_path / "vec.db") + conn.executescript( + """ + CREATE TABLE embeddings(chunk_id INTEGER PRIMARY KEY, embedding BLOB); + """ + ) + conn.execute( + "INSERT INTO embeddings VALUES (?, ?)", + (1, DebugEmbeddingStore(4).embed_many(["hello"])[0]), + ) + conn.execute( + "INSERT INTO embeddings VALUES (?, ?)", + (2, DebugEmbeddingStore(4).embed_many(["another"])[0]), + ) + conn.commit() + backend = PythonFallbackBackend() + store = DebugEmbeddingStore(4) + query = embedding_from_bytes(store.embed_many(["hello"])[0]) + results = backend.search(conn, query, top_n=2) + assert results and results[0].chunk_id == 1