diff --git a/.claude/CLAUDE.md b/.claude/CLAUDE.md index 162effc..45702e6 100644 --- a/.claude/CLAUDE.md +++ b/.claude/CLAUDE.md @@ -44,11 +44,24 @@ backends. - `providers.py` + `data/providers.yaml` — the ~21-provider registry + `recommend_backend` decision framework. `requirements.py` — `check_requirements` / `setup_guide` / `install_backend`. -- Modules: `analytics`, `config`, `health`, `io`, `migration`, `search`, - `text`, `time_indexed`, `util`, `cli`. (`compare.py` was folded into - `providers.py`.) -- 5 user-facing skills in `vd/data/skills/`; `vd-add-backend` dev skill in - `.claude/skills/`. +- `asynchronous.py` — async surface: `connect_async`, the universal + `asyncio.to_thread` wrappers, and **native** async bases + (`AsyncAbstractClient` / `AsyncAbstractCollection`) plus a registry + (`register_async_backend`, `list_async_backends`). Native today: `qdrant` + with `url=` (embedded Qdrant's async client blocks the loop, so its factory + returns the wrapper there). + The I/O-free embed / dimension / query policy is shared with the sync base + through `base._CollectionPolicy`. +- Hybrid: `vd.hybrid_search` works everywhere (client-side BM25 + RRF + fallback, `search.BM25Index`); `weaviate`, `elasticsearch`, `redis` and + `lancedb` satisfy `SupportsHybrid` with a native lexical side. +- Modules: `analytics`, `config`, `filters`, `health`, `io`, `migration`, + `search`, `text`, `time_indexed`, `util`, `cli`. (`compare.py` was folded + into `providers.py`.) `providers.install_command` returns + `pip install "vd[]"` — the `pyproject.toml` extras are the SSOT + for what each adapter needs. +- 7 skills in `vd/data/skills/` (6 user-facing incl. `vd-setup-backend`, plus + the `vd-add-backend` dev skill), each symlinked from `.claude/skills/`. **Embedding is external.** The core operates on vectors; an `embedder` passed to `connect` is an optional convenience. `vd` has **no `dol`/`imbed` @@ -57,6 +70,24 @@ dependency** — the core is stdlib + `pyyaml` only. The first consumer, `ef`, was adapted on its `adapt-to-vd-0.2` branch (it dropped its dummy-embedder workaround); merge that only after vd 0.2 publishes. +### Running the tests + +```bash +uv venv .venv && . .venv/bin/activate +uv pip install -e ".[test,dev]" # dev = the embedded backends the suite sweeps +python -m pytest --doctest-modules -o doctest_optionflags='ELLIPSIS IGNORE_EXCEPTION_DETAIL' +``` + +This is CI's command. `testpaths = ["tests", "vd"]`, so package doctests run +too; the root `conftest.py` skips backend modules whose SDK is missing. +Parametrized suites (`tests/test_core.py`, `tests/test_hybrid.py`) run every +reachable backend via the `client` fixture in `tests/conftest.py`: embedded +backends always, server backends only when their port answers. Bring the +servers up with `docker compose -f tests/docker-compose.yml up -d` and install +their clients (`uv pip install -e ".[pgvector,redis,elasticsearch,weaviate,mongodb,milvus]"`). +CI installs only the `test` extra, so backend-specific tests skip there — +run them locally before changing an adapter. + ## 3. Core contracts (the design the refactor should converge on) ### 3.1 The hierarchy: `Client` → `Collection` @@ -133,7 +164,10 @@ fallback engine (`reciprocal_rank_fusion` already exists in `search.py`). ## 4. Refactor priorities (gaps to fill) -In rough order: +In rough order. Status as of 2026-09: done — 2, 3, 4, 7; partly done — 1 +(the `F(...)` builder is missing), 6 (native async: `qdrant` only; issue #20 +tracks the rest), 9 (dimension checks are loud and early; `model_id` is not +stored); open — 5, 8. Native hybrid for the remaining backends is issue #17. 1. **`UnsupportedFilterError`** + per-adapter documented filter subset + a `_compile_filter` translator per adapter (memory evaluates in Python; diff --git a/.claude/skills/vd-setup-backend b/.claude/skills/vd-setup-backend new file mode 120000 index 0000000..9860c24 --- /dev/null +++ b/.claude/skills/vd-setup-backend @@ -0,0 +1 @@ +../../vd/data/skills/vd-setup-backend \ No newline at end of file diff --git a/README.md b/README.md index aef4ce5..02d1532 100644 --- a/README.md +++ b/README.md @@ -1,82 +1,89 @@ # vd -**A facade over vector databases — one Pythonic interface, ~15 backends.** +**One Pythonic interface to ~15 vector databases.** Store vectors, search them, filter by metadata, and switch between databases (memory, Chroma, Qdrant, pgvector, LanceDB, Pinecone, …) by changing one word. -`vd` lets you operate on any vector database and switch between them with a -one-word change, while keeping each backend's particular power one escape hatch -away. It also helps you *choose* the right backend and *set it up*. +*Still typing code with your own ten fingers? [Your artisanal corner is at the bottom](#for-artisanal-hand-typed-contributions).* + +## For AI agents + +**Skills.** `vd` ships agent skills in [`vd/data/skills/`](vd/data/skills/). They install into any agent host with [`gh skill`](https://cli.github.com/manual/gh_skill): + +```bash +gh skill install i2mint/vd vd-quickstart --agent claude-code +gh skill install i2mint/vd vd-backend-choose +gh skill install i2mint/vd vd-setup-backend +gh skill install i2mint/vd vd-ingest +gh skill install i2mint/vd vd-search +gh skill install i2mint/vd vd-ops +gh skill install i2mint/vd vd-add-backend +``` + +They also ship inside the wheel, so after `pip install vd` you can link them without GitHub: + +```bash +mkdir -p ~/.claude/skills +ln -s "$(python -c 'import vd; print(vd.skills_dir())')"/vd-* ~/.claude/skills/ +``` + +| Skill | Use it when… | +|-------|--------------| +| `vd-quickstart` | doing basic vector search: connect, create a collection, add docs, query | +| `vd-backend-choose` | picking a vector DB and weighing trade-offs | +| `vd-setup-backend` | installing, starting and verifying a backend (pip, Docker, API keys) | +| `vd-ingest` | loading documents: cleaning, chunking, metadata, bulk insert | +| `vd-search` | filters, hybrid (keyword + vector) search, multi-query, RRF, similar-to, dedup | +| `vd-ops` | export/import, migration, stats, integrity, health checks, benchmarks | +| `vd-add-backend` | (developer) implementing or reviewing a backend adapter | + +**Project instructions** for coding agents working *on* `vd` live in [`.claude/CLAUDE.md`](.claude/CLAUDE.md): architecture, test commands, conventions and the refactor roadmap. + +**What an agent can do with `vd`:** recommend and set up a vector database for a user, index their documents, run filtered or hybrid searches, and move a collection between databases. The whole surface is two `Mapping`s plus `search`: ```python import vd client = vd.connect("memory") # switch DB = change this one word col = client.create_collection("docs") -col["a"] = vd.Document(id="a", text="cats", vector=[0.1, 0.9, 0.0]) -col["b"] = vd.Document(id="b", text="pizza", vector=[0.9, 0.0, 0.1]) +col["a"] = vd.Document(id="a", text="cats", vector=[0.1, 0.9, 0.0], metadata={"kind": "pet"}) +col["b"] = vd.Document(id="b", text="pizza", vector=[0.9, 0.0, 0.1], metadata={"kind": "food"}) for hit in col.search([0.1, 0.8, 0.0], limit=2): - print(hit["id"], hit["score"]) + print(hit["id"], round(hit["score"], 3)) +# a 1.0 +# b 0.123 ``` ## Install ```bash -pip install vd # core (zero heavy deps) + the memory backend -pip install vd[chroma] # + a specific backend's client -pip install vd[embedded] # + all embedded backends (chroma, qdrant, faiss, …) -pip install vd[all-backends] # + every backend client +pip install vd # core (stdlib + pyyaml) and the memory backend +pip install "vd[chroma]" # plus one backend's client +pip install "vd[embedded]" # plus chroma, qdrant, faiss, lancedb, sqlite_vec, duckdb +pip install "vd[all-backends]" # plus every backend client ``` -The core is near-zero-dependency. Each backend's client library is an optional -extra named after the backend. +Each backend's client is an optional extra named after the backend. Keep the quotes: zsh treats `[...]` as a glob. ## The mental model -`vd` stores and searches **vectors**. Turning text into vectors — *embedding* — -is deliberately **external**: `vd` never embeds on its own. This keeps the -facade honest (most vector DBs do not embed for you) and lightweight. - -- **Vector-first.** You hold the embedding model. Hand `vd` `Document`s that - already carry a `vector`; search with a pre-computed query vector. -- **Text convenience.** Pass an `embedder` (`text -> vector`) to `connect`, and - then raw text works: `col["k"] = "some text"`, `col.search("a query")`. +`vd` stores and searches **vectors**. Turning text into vectors, *embedding*, is deliberately **external**: `vd` never embeds on its own. Hand it `Document`s that carry a `vector`, and search with a query vector. -With no embedder, passing text raises `EmbeddingRequiredError` — loud, never a -silent wrong-model embedding. +For convenience, pass any `text -> vector` function as `embedder`, and raw text then works too: ```python -client = vd.connect("chroma", persist_directory="./db", embedder=my_embed_fn) -col = client.create_collection("docs") -col["a"] = "cats and kittens" # embedded for you -hits = list(col.search("pets", limit=5)) # query embedded for you -``` - -## Choosing a backend - -`vd` ships a provider registry distilled from a practitioner report -(`misc/docs/11 -- VectorDB Selection & Setup Guide ...md`) and a recommender: - -```python -vd.print_recommendation( - corpus_size="medium", persistence=True, can_run_docker=True, - cloud_ok=True, budget="free", needs_hybrid=False, -) -vd.print_backends_table() # the whole landscape -vd.compare_backends(["chroma", "qdrant", "pgvector"]) -``` +import vd -## Setting a backend up +def embed(text): # stand-in; use a real embedding model + return [text.count(c) / (len(text) or 1) for c in "aeiou"] -```python -vd.check_requirements("qdrant") # diagnoses readiness, prints the next step -vd.setup_guide("qdrant") # full pip / docker / env-var playbook -vd.install_backend("qdrant") # the pip command (run=True to install) +col = vd.connect("memory", embedder=embed).create_collection("notes") +col["k1"] = "cats and kittens" # embedded for you +col["k2"] = ("dogs and puppies", {"kind": "dog"}) # text + metadata +print([h["id"] for h in col.search("a kitten", limit=1)]) +# ['k1'] ``` -`check_requirements` is deployment-aware: it checks the pip package for -embedded backends, whether a server answers for self-hosted ones, and the -required environment variables for managed ones — always ending with one -concrete next action. +Without an embedder, passing text raises `EmbeddingRequiredError`: loud, never a silent wrong-model embedding. ## The API @@ -87,101 +94,137 @@ concrete next action. | `Document` | dataclass | `id`, `text`, `vector`, `metadata` | ```python -col["k"] = vd.Document(id="k", text="…", vector=[...], metadata={"y": 2024}) -doc = col["k"] # get -del col["k"] # delete -"k" in col, len(col), list(col) +import vd -col.search(query, *, limit=10, filter=None, egress=None, **backend_kwargs) +col = vd.connect("memory").create_collection("api") +col["k"] = vd.Document(id="k", text="hello", vector=[1.0, 0.0], metadata={"year": 2024}) +doc = col["k"] # get +print("k" in col, len(col), list(col)) # membership, count, keys +# True 1 ['k'] +del col["k"] # delete ``` -`search` yields dicts `{"id", "text", "score", "metadata"}` (`score` is -higher-is-better). Transform results with an `egress`: `vd.id_only`, -`vd.id_and_score`, `vd.text_only`, `vd.id_text_score`, or your own. +`search(query, *, limit=10, filter=None, egress=None)` yields dicts `{"id", "text", "score", "metadata"}` where `score` is higher-is-better. Transform results with an `egress` such as `vd.id_only`, `vd.id_and_score`, `vd.text_only`, `vd.id_text_score`, or your own function. ### Metadata filtering -One backend-agnostic, MongoDB-style filter language — `$eq $ne $gt $gte $lt -$lte $in $nin $exists $and $or $not`: +One backend-agnostic, MongoDB-style filter language: `$eq $ne $gt $gte $lt $lte $in $nin $exists $and $or $not`. ```python -col.search(qvec, filter={"year": {"$gte": 2020}, "kind": {"$in": ["news", "blog"]}}) +import vd + +col = vd.connect("memory").create_collection("posts") +col["a"] = vd.Document(id="a", text="", vector=[1.0, 0.0], metadata={"year": 2023, "kind": "news"}) +col["b"] = vd.Document(id="b", text="", vector=[0.9, 0.1], metadata={"year": 2019, "kind": "blog"}) +hits = col.search([1.0, 0.0], filter={"year": {"$gte": 2020}, "kind": {"$in": ["news", "blog"]}}) +print([h["id"] for h in hits]) +# ['a'] ``` -Each backend declares which operators it honors natively; an unsupported one -raises `UnsupportedFilterError` rather than silently mis-filtering. Backends -with rich native filtering (Qdrant, Pinecone, MongoDB) translate the filter; -the rest apply it client-side with the same semantics. +Each backend declares which operators it honors. An unsupported one raises `UnsupportedFilterError` instead of silently mis-filtering. Backends with rich native filtering (Qdrant, Pinecone, MongoDB) translate the filter; the rest apply it client-side with the same semantics. + +### Hybrid search + +`vd.hybrid_search` fuses a vector search with a keyword (BM25) search using Reciprocal Rank Fusion, on any backend: + +```python +import vd + +col = vd.connect("memory").create_collection("hybrid") +col["a"] = vd.Document(id="a", text="error code E1234 in the parser", vector=[0.2, 0.8]) +col["b"] = vd.Document(id="b", text="parser overview", vector=[0.9, 0.1]) +hits = vd.hybrid_search(col, [0.9, 0.1], query_text="E1234", limit=2) +print([h["id"] for h in hits]) +# ['a', 'b'] +``` + +Weaviate, Elasticsearch, Redis and LanceDB run the keyword side on their own text index (`isinstance(col, vd.SupportsHybrid)`); other backends use a client-side BM25 scan. + +### Async + +```python +import asyncio +import vd + +async def main(): + async with await vd.connect_async("memory") as client: + col = await client.create_collection("docs", dimension=2) + await col.set("a", vd.Document(id="a", text="x", vector=[1.0, 0.0])) + return [h["id"] async for h in col.search([1.0, 0.0], limit=1)] + +print(asyncio.run(main())) +# ['a'] +``` + +Every backend works through a thread-pool wrapper. Backends listed by `vd.list_async_backends()` can return a native async client that does real non-blocking I/O: currently `qdrant` when connected to a server with `url=`. Check `client.native_async`. ### Escape hatches -The facade never traps you. `client.client` is the raw backend client; -`collection.native` is the raw backend collection — both supported, documented -API for reaching backend-specific features. +The facade never traps you. `client.client` is the raw backend client and `collection.native` is the raw backend collection. Both are supported, documented API for reaching backend-specific features. + +## Choosing and setting up a backend + +```python +import vd + +vd.print_recommendation(corpus_size="medium", persistence=True, can_run_docker=True, cloud_ok=True, budget="free", needs_hybrid=False) +vd.print_backends_table() # the whole landscape +vd.compare_backends(["chroma", "qdrant", "pgvector"]) + +vd.check_requirements("qdrant") # diagnoses readiness and prints the next step +print(vd.setup_guide("qdrant")) # pip / docker / env-var playbook +vd.install_command("qdrant") # 'pip install "vd[qdrant]"' +``` -## Backends +`check_requirements` is deployment-aware. It checks the client library for embedded backends, whether a server answers for self-hosted ones, and the required environment variables for managed ones, and always ends with one concrete next action. | Archetype | Backends | |-----------|----------| | **Embedded** (pip-only) | `memory`, `chroma`, `lancedb`, `sqlite_vec`, `duckdb`, `faiss` | -| **Server** (also embedded) | `qdrant`, `weaviate`, `milvus` | -| **Server** | `redis`, `elasticsearch`, `pgvector` | +| **Server** (also embedded) | `qdrant`, `milvus` | +| **Server** | `weaviate`, `redis`, `elasticsearch`, `pgvector` | | **Managed** | `pinecone`, `mongodb` (Atlas), `turbopuffer` | -`vd.list_backends()` shows what is installed and ready now. +`vd.list_backends()` shows which adapters are installed and ready now. ## The toolkit Beyond the facade, `vd` bundles the composite operations people actually do: -- **`vd.search`** — `multi_query_search`, `reciprocal_rank_fusion`, - `search_similar_to_document`, `deduplicate_results`. -- **`vd.io`** — `export_collection` / `import_collection` (JSONL, JSON, - directory). -- **`vd.migration`** — `migrate_collection`, `migrate_client`, - `copy_collection` — move data between *any* two backends. -- **`vd.analytics`** — `collection_stats`, `find_duplicates`, `find_outliers`, - `validate_collection`. -- **`vd.health`** — `health_check_backend`, `benchmark_search`. -- **`vd.text`** — convenience text cleaning / chunking. -- **`vd.TimeIndexedCollection`** — a time-windowed wrapper over any collection. -- **CLI** — `vd backends`, `vd install`, `vd export/import`, `vd migrate`, … - -## Skills - -This package ships agent skills you can install into any agent host with -[`gh skill`](https://cli.github.com/manual/gh_skill) (don't have it? -[install gh](https://cli.github.com/)): +- **`vd.search`**: `multi_query_search`, `reciprocal_rank_fusion`, `hybrid_search`, `BM25Index`, `search_similar_to_document`, `deduplicate_results`. +- **`vd.io`**: `export_collection` / `import_collection` (JSONL, JSON, directory). +- **`vd.migration`**: `migrate_collection`, `migrate_client`, `copy_collection` to move data between any two backends. +- **`vd.analytics`**: `collection_stats`, `find_duplicates`, `find_outliers`, `validate_collection`. +- **`vd.health`**: `health_check_backend`, `benchmark_search`. +- **`vd.text`**: convenience text cleaning and chunking. +- **`vd.TimeIndexedCollection`**: a time-windowed wrapper over any collection. +- **CLI**: `vd backends`, `vd install`, `vd export`, `vd import`, `vd migrate`, … + +## For artisanal, hand-typed contributions + +Welcome, fellow keyboard enthusiast. Here is what the sections above don't already cover. + +**Dev setup and tests.** ```bash -gh skill install i2mint/vd vd-quickstart --agent claude-code -gh skill install i2mint/vd vd-backend-choose -gh skill install i2mint/vd vd-ingest -gh skill install i2mint/vd vd-search -gh skill install i2mint/vd vd-ops -gh skill install i2mint/vd vd-add-backend +git clone https://github.com/i2mint/vd && cd vd +uv venv .venv && . .venv/bin/activate +uv pip install -e ".[test,dev]" +python -m pytest --doctest-modules -o doctest_optionflags='ELLIPSIS IGNORE_EXCEPTION_DETAIL' ``` -| Skill | Use it when… | -|-------|--------------| -| `vd-quickstart` | doing basic semantic/vector search — connect, create a collection, add docs, query | -| `vd-backend-choose` | picking a vector DB, weighing trade-offs, or installing/starting a backend | -| `vd-ingest` | loading documents/files into a collection — cleaning, chunking, metadata, bulk insert | -| `vd-search` | going beyond a basic `.search()` — filters, multi-query, RRF, similar-to, dedup, by-vector | -| `vd-ops` | managing a collection — export/import, migrate, stats, integrity, health checks, benchmarks | -| `vd-add-backend` | (developer) implementing or reviewing a new vd backend adapter | - -## Design - -- **Embedding is external.** The core operates on vectors; an `embedder` is an - injected, optional convenience — never a hard dependency. -- **Two mappings.** A `Client` is a `Mapping` of collections; a `Collection` is - a `MutableMapping` of documents plus `search`. Idiomatic, minimal, familiar. -- **Thin adapters.** `AbstractClient` / `AbstractCollection` implement - everything users see; a backend supplies a handful of raw primitives. Adding - a backend is ~150 lines — see the `vd-add-backend` skill. -- **Capabilities, not a fat base.** Optional features (`SupportsBatch`, - `SupportsHybrid`) are `@runtime_checkable` protocols you feature-discover. +That is exactly what CI runs; package doctests are included. The backend-parametrized suites skip server backends that aren't running. To exercise them, start the containers with `docker compose -f tests/docker-compose.yml up -d` and install their clients with `uv pip install -e ".[pgvector,redis,elasticsearch,weaviate,mongodb,milvus]"`. + +**Design rationale.** + +- **Embedding is external.** The core operates on vectors; an `embedder` is an injected, optional convenience, never a hard dependency. +- **Two mappings.** A `Client` is a `Mapping` of collections; a `Collection` is a `MutableMapping` of documents plus `search`. Idiomatic, minimal, familiar. +- **Thin adapters.** `AbstractClient` / `AbstractCollection` implement everything users see; a backend supplies a handful of raw primitives. Adding a backend is about 150 lines, and the `vd-add-backend` skill walks through it. +- **Capabilities, not a fat base.** Optional features (`SupportsBatch`, `SupportsHybrid`, `SupportsNativeAsync`) are `@runtime_checkable` protocols you feature-discover. + +The longer design notes are in [`misc/docs/vd_design_notes.md`](misc/docs/vd_design_notes.md), and the backend-selection research report is in [`misc/docs/`](misc/docs/). + +**Contributing.** Open an issue or a pull request at [github.com/i2mint/vd](https://github.com/i2mint/vd/issues). Releases are automated: merging to `master` bumps the version and publishes to PyPI. ## License diff --git a/conftest.py b/conftest.py new file mode 100644 index 0000000..7b9a853 --- /dev/null +++ b/conftest.py @@ -0,0 +1,32 @@ +""" +Root pytest configuration: decide which package modules doctest collection skips. + +``pytest --doctest-modules`` imports every module under ``vd/``. Each backend +adapter raises :class:`ImportError` at import time when its optional client +SDK is absent (``pip install vd[]`` installs it), which would abort +collection. Those modules are skipped here, so the rest of the package's +doctests always run and a backend's doctests run whenever its SDK is present. + +``misc/`` holds demo scripts and design notes, not tests. +""" + +import importlib +import pathlib + +_HERE = pathlib.Path(__file__).parent + + +def _unimportable_backend_modules() -> list[str]: + """Return the paths of backend modules whose optional SDK is not installed.""" + skipped = [] + for path in sorted((_HERE / "vd" / "backends").glob("*.py")): + if path.name.startswith("_"): + continue + try: + importlib.import_module(f"vd.backends.{path.stem}") + except ImportError: + skipped.append(str(path.relative_to(_HERE))) + return skipped + + +collect_ignore = ["misc", *_unimportable_backend_modules()] diff --git a/pyproject.toml b/pyproject.toml index c5d7d3d..cd20eaa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ keywords = [ "similarity search", "RAG", ] -authors = [] +authors = [{ name = "Thor Whalen" }] classifiers = [ "Development Status :: 4 - Beta", "Intended Audience :: Developers", @@ -84,12 +84,28 @@ dev = [ # Lightweight test-runner deps for CI (no backend extras — those live in # `dev`). Installed in CI via [tool.wads.ci.install].extras so the async # suite (tests/test_async.py) has pytest-asyncio. -test = ["pytest>=7.0", "pytest-cov>=4.0", "pytest-asyncio>=0.23"] +# hatchling: tests/test_package.py checks the sdist file list (#28). +test = ["pytest>=7.0", "pytest-cov>=4.0", "pytest-asyncio>=0.23", "hatchling"] docs = ["sphinx>=6.0", "sphinx-rtd-theme>=1.0"] [tool.hatch.build.targets.wheel] packages = ["vd"] +[tool.hatch.build.targets.sdist] +# Walk only these roots. `.claude/skills/*` are symlinks into `vd/data/skills/` +# and hatch skips files whose real path it has already seen, so walking +# `.claude` first (even with an `exclude`) dropped `vd/data/skills/` from the +# sdist, and the wheel built from it shipped no skills. +only-include = [ + "vd", + "tests", + "conftest.py", + "misc", + "README.md", + "CHANGELOG.md", + "LICENSE", +] + [tool.ruff] line-length = 88 target-version = "py310" @@ -118,7 +134,9 @@ convention = "google" [tool.pytest.ini_options] minversion = "6.0" -testpaths = ["tests"] +# `vd` is listed so the package's own doctests run under `--doctest-modules`; +# the root conftest.py skips backend modules whose optional SDK is missing. +testpaths = ["tests", "vd"] doctest_optionflags = ["NORMALIZE_WHITESPACE", "ELLIPSIS"] # Make every `async def test_*` run automatically — no per-test decorator # needed. Used by tests/test_async.py. diff --git a/tests/test_async.py b/tests/test_async.py index 5db4f5c..8450bca 100644 --- a/tests/test_async.py +++ b/tests/test_async.py @@ -16,8 +16,7 @@ - :func:`vd.hybrid_search_async` dispatches correctly through the wrapper; - the async context-manager protocol works. -Phase 2 follow-ups will add per-backend native async adapters; each backend -will get its own parametrized test entries at that time. +Native async adapters (Phase 2, #20) are tested in ``tests/test_async_native.py``. """ import pytest diff --git a/tests/test_async_native.py b/tests/test_async_native.py new file mode 100644 index 0000000..2e2a1f7 --- /dev/null +++ b/tests/test_async_native.py @@ -0,0 +1,395 @@ +""" +Tests for native async backends (Phase 2 of the async work, #20). + +Covers the dispatch layer and the first native adapter: + +- :func:`vd.register_async_backend` + :func:`vd.connect_async` dispatch: a + registered native factory is returned instead of the ``to_thread`` wrapper, + and ``native=False`` forces the wrapper; +- the native qdrant adapter (``qdrant_client.AsyncQdrantClient``), exercised + in Qdrant's embedded ``:memory:`` mode so no server is needed. Its results + must match the wrapped sync adapter's exactly; +- :func:`vd.hybrid_search_async` over a native collection (no ``.sync`` to + fall back on). +""" + +import pytest + +import vd +from vd.asynchronous import AsyncClientWrapper + +from tests.conftest import make_embedder + +pytestmark = pytest.mark.asyncio + + +# --------------------------------------------------------------------------- # +# Registry / dispatch +# --------------------------------------------------------------------------- # + + +async def test_registered_native_factory_is_used_and_can_be_bypassed(): + made = {} + + class FakeNativeClient: + native_async = True + + def __init__(self, **kwargs): + made.update(kwargs) + + vd.register_async_backend("memory_fake_native")(FakeNativeClient) + try: + client = await vd.connect_async("memory_fake_native", flavor="x") + assert isinstance(client, FakeNativeClient) + assert made == {"flavor": "x"} + assert "memory_fake_native" in vd.list_async_backends() + finally: + from vd.asynchronous import _async_backends + + _async_backends.pop("memory_fake_native", None) + + +async def test_async_factory_coroutine_is_awaited(): + class FakeNativeClient: + native_async = True + + async def factory(**kwargs): + return FakeNativeClient() + + vd.register_async_backend("coro_fake_native", factory) + try: + assert isinstance( + await vd.connect_async("coro_fake_native"), FakeNativeClient + ) + finally: + from vd.asynchronous import _async_backends + + _async_backends.pop("coro_fake_native", None) + + +async def test_backend_without_native_gets_wrapper(): + client = await vd.connect_async("memory") + assert isinstance(client, AsyncClientWrapper) + assert client.native_async is False + + +# --------------------------------------------------------------------------- # +# Native qdrant +# --------------------------------------------------------------------------- # + + +@pytest.fixture +def qdrant_available(): + pytest.importorskip("qdrant_client") + if "qdrant" not in vd.list_backends(): + pytest.skip("qdrant backend not installed") + + +def _native(**kwargs): + """The native async qdrant client, in embedded mode (no server needed). + + ``connect_async`` only picks it for a server (``url=``): qdrant-client's + embedded async client runs synchronous code and would block the event + loop. Constructing it directly is how these tests exercise its code + without a server. + """ + from vd.backends.qdrant import NativeAsyncQdrantClient + + return NativeAsyncQdrantClient(**kwargs) + + +async def test_qdrant_connect_async_dispatch(qdrant_available): + from qdrant_client import AsyncQdrantClient + + from vd.backends.qdrant import NativeAsyncQdrantClient + + # Embedded mode: the thread-pool wrapper (keeps the event loop free). + embedded = await vd.connect_async("qdrant") + assert isinstance(embedded, AsyncClientWrapper) + await embedded.close() + # Server mode: the native client (construction does not connect). + remote = await vd.connect_async("qdrant", url="http://localhost:6399", + check_compatibility=False) + assert isinstance(remote, NativeAsyncQdrantClient) + assert remote.native_async is True + assert isinstance(remote.client, AsyncQdrantClient) + await remote.close() + remote2 = await vd.connect_async("qdrant", location="http://localhost:6399", + check_compatibility=False) + assert isinstance(remote2, NativeAsyncQdrantClient) + await remote2.close() + # native=False always gives the wrapper + wrapped = await vd.connect_async("qdrant", url="http://localhost:6399", check_compatibility=False, + native=False) + assert isinstance(wrapped, AsyncClientWrapper) + await wrapped.close() + + +async def test_qdrant_native_client_surface_types(qdrant_available): + async with _native() as client: + assert isinstance(client, vd.AsyncClient) + assert isinstance(client, vd.SupportsNativeAsync) + col = await client.create_collection("docs", dimension=2) + assert isinstance(col, vd.AsyncCollection) + assert col.native_async is True + assert col.native is client.client + + +async def _populate(col): + await col.set( + "a", vd.Document(id="a", text="cats purr", vector=[1.0, 0.0], + metadata={"k": 1, "tag": "pet"}) + ) + await col.set( + "b", vd.Document(id="b", text="dogs bark", vector=[0.0, 1.0], + metadata={"k": 2, "tag": "pet"}) + ) + await col.set( + "c", vd.Document(id="c", text="cats and dogs", vector=[0.6, 0.8], + metadata={"k": 3}) + ) + + +async def test_qdrant_native_crud(qdrant_available): + async with _native() as client: + col = await client.create_collection("crud", dimension=2) + assert await col.count() == 0 + assert [k async for k in col.keys()] == [] + with pytest.raises(KeyError): + await col.get("a") + + await _populate(col) + assert await col.count() == 3 + assert sorted([k async for k in col.keys()]) == ["a", "b", "c"] + + doc = await col.get("a") + assert (doc.id, doc.text, doc.metadata) == ("a", "cats purr", + {"k": 1, "tag": "pet"}) + assert doc.vector == pytest.approx([1.0, 0.0]) + + # set is an idempotent replace + await col.set("a", vd.Document(id="a", text="new", vector=[1.0, 0.0])) + assert (await col.get("a")).text == "new" + assert await col.count() == 3 + + await col.delete("b") + assert await col.count() == 2 + with pytest.raises(KeyError): + await col.get("b") + with pytest.raises(KeyError): + await col.delete("b") + + +async def test_qdrant_native_search_filter_egress(qdrant_available): + async with _native() as client: + col = await client.create_collection("srch", dimension=2) + await _populate(col) + hits = [h async for h in col.search([0.9, 0.1], limit=2)] + assert [h["id"] for h in hits] == ["a", "c"] + assert hits[0]["score"] >= hits[1]["score"] + assert {"id", "text", "score", "metadata"} <= set(hits[0]) + + ids = [h["id"] async for h in col.search( + [0.9, 0.1], filter={"k": {"$gte": 2}, "tag": "pet"})] + assert ids == ["b"] + + only_ids = [r async for r in col.search([0.9, 0.1], limit=3, + egress=lambda h: h["id"])] + assert only_ids == ["a", "c", "b"] + + +async def test_qdrant_native_matches_wrapped_sync(qdrant_available): + """Native and wrapped-sync adapters return identical results.""" + results = [] + for native in (True, False): + client = _native() if native else await vd.connect_async("qdrant") + async with client: + col = await client.create_collection("parity", dimension=2, + metric="l2") + await _populate(col) + hits = [ + (h["id"], round(h["score"], 6), h["text"], h["metadata"]) + async for h in col.search([0.7, 0.3], limit=3, + filter={"k": {"$in": [1, 3]}}) + ] + results.append(hits) + assert results[0] == results[1] + assert [r[0] for r in results[0]] == ["a", "c"] + + +async def test_qdrant_native_embedder_and_dimension_checks(qdrant_available): + embed = make_embedder() + async with _native(embedder=embed) as client: + col = await client.create_collection("emb") + await col.set("a", "cats and kittens") + await col.set("b", ("dogs and puppies", {"kind": "dog"})) + assert (await col.get("b")).metadata == {"kind": "dog"} + assert col.dimension == len(embed("x")) + hits = [h["id"] async for h in col.search("cats and kittens", limit=1)] + assert hits == ["a"] + with pytest.raises(ValueError, match="dimension mismatch"): + await col.set("c", vd.Document(id="c", text="x", vector=[1.0, 2.0])) + with pytest.raises(ValueError, match="dimension mismatch"): + async for _ in col.search([1.0, 2.0]): + pass + + async with _native() as client: + col = await client.create_collection("noemb", dimension=2) + with pytest.raises(vd.EmbeddingRequiredError): + await col.set("a", "raw text needs an embedder") + with pytest.raises(vd.UnsupportedFilterError): + async for _ in col.search([1.0, 0.0], filter={"x": {"$regex": "a"}}): + pass + + +async def test_qdrant_native_batch_ops(qdrant_available): + async with _native() as client: + col = await client.create_collection("batch", dimension=2) + await col.upsert(vd.Document(id="x", text="x", vector=[1.0, 0.0])) + await col.add_documents( + [vd.Document(id=str(i), text=f"t{i}", vector=[1.0, float(i)]) + for i in range(5)], + batch_size=2, + ) + assert await col.count() == 6 + assert sorted([k async for k in col.keys()]) == [ + "0", "1", "2", "3", "4", "x" + ] + + +async def test_qdrant_native_keys_paginate(qdrant_available): + async with _native() as client: + col = await client.create_collection("many", dimension=2) + await col.add_documents( + [vd.Document(id=f"d{i}", text="t", vector=[1.0, float(i)]) + for i in range(600)], + batch_size=250, + ) + keys = [k async for k in col.keys()] + assert len(keys) == len(set(keys)) == 600 + + +async def test_qdrant_native_client_surface(qdrant_available): + async with _native() as client: + assert [n async for n in client.list_collections()] == [] + await client.create_collection("one", dimension=2) + lazy = await client.create_collection("lazy") # no dimension yet + assert sorted([n async for n in client.list_collections()]) == [ + "lazy", "one" + ] + with pytest.raises(ValueError): + await client.create_collection("one", dimension=2) + with pytest.raises(KeyError): + await client.get_collection("nope") + with pytest.raises(KeyError): + await client.delete_collection("nope") + + # the lazy collection learns its dimension on first write + await lazy.set("a", vd.Document(id="a", text="t", vector=[0.0, 1.0])) + again = await client.get_collection("lazy") + assert (await again.get("a")).text == "t" + + same = await client.get_or_create_collection("one", dimension=2) + assert await same.count() == 0 + fresh = await client.get_or_create_collection("two", dimension=2) + assert isinstance(fresh, vd.AsyncCollection) + + await client.delete_collection("one") + assert "one" not in [n async for n in client.list_collections()] + + +async def test_hybrid_search_async_over_native_collection(qdrant_available): + async with _native() as client: + col = await client.create_collection("hyb", dimension=2) + await _populate(col) + hits = [ + h async for h in vd.hybrid_search_async( + col, [0.9, 0.1], query_text="cats", limit=3 + ) + ] + ids = [h["id"] for h in hits] + assert ids[0] == "a" # top on both the dense and the lexical side + assert set(ids) == {"a", "b", "c"} + filtered = [ + h["id"] async for h in vd.hybrid_search_async( + col, [0.9, 0.1], query_text="cats", limit=3, + filter={"tag": "pet"}, + ) + ] + assert set(filtered) == {"a", "b"} + with pytest.raises(ValueError, match="query_text"): + async for _ in vd.hybrid_search_async(col, [0.9, 0.1], limit=2): + pass + + +async def test_hybrid_search_async_native_custom_lexical_and_errors( + qdrant_available, +): + async with _native() as client: + col = await client.create_collection("hyb2", dimension=2) + await _populate(col) + + async def my_lex(collection, text, *, limit, filter): + return [{"id": "b", "text": "dogs bark", "score": 9.0, "metadata": {}}] + + ids = [ + h["id"] async for h in vd.hybrid_search_async( + col, [1.0, 0.0], query_text="zzz", limit=3, lexical_search=my_lex, + egress=lambda h: {**h, "seen": True}, + ) + ] + assert "b" in ids[:2] + with pytest.raises(ValueError, match="non-empty"): + async for _ in vd.hybrid_search_async(col, [1.0, 0.0], query_text=""): + pass + + +async def test_hybrid_search_async_delegates_to_native_hybrid(): + calls = {} + + class FakeNativeHybrid: + native_async = True + + async def hybrid_search(self, query, **kwargs): + calls.update(kwargs, query=query) + yield {"id": "x", "text": "", "score": 1.0, "metadata": {}} + + hits = [ + h async for h in vd.hybrid_search_async( + FakeNativeHybrid(), "q", limit=2, alpha=0.5 + ) + ] + assert [h["id"] for h in hits] == ["x"] + assert calls["query"] == "q" and calls["alpha"] == 0.5 + assert calls["k_dense"] == calls["k_lexical"] == 50 + + +async def test_qdrant_native_lazy_collection_before_first_write(qdrant_available): + async with _native() as client: + col = await client.create_collection("lazy") # no dimension → not created + assert await col.count() == 0 + assert [k async for k in col.keys()] == [] + assert [h async for h in col.search([1.0, 0.0])] == [] + with pytest.raises(KeyError): + await col.get("a") + with pytest.raises(KeyError): + await col.delete("a") + await client.delete_collection("lazy") # registered but never created + assert [n async for n in client.list_collections()] == [] + + +async def test_hybrid_search_async_native_accepts_sync_lexical_callable( + qdrant_available, +): + """A sync lexical_search (e.g. vd.bm25_lexical_search) gets a mapping of docs.""" + async with _native() as client: + col = await client.create_collection("hyb3", dimension=2) + await _populate(col) + ids = [ + h["id"] async for h in vd.hybrid_search_async( + col, [0.0, 1.0], query_text="cats", limit=3, + lexical_search=vd.bm25_lexical_search, + ) + ] + assert set(ids) == {"a", "b", "c"} + assert ids.index("a") < ids.index("b") or ids[0] == "b" diff --git a/tests/test_hybrid.py b/tests/test_hybrid.py index 5d85866..2f0f0fd 100644 --- a/tests/test_hybrid.py +++ b/tests/test_hybrid.py @@ -272,7 +272,7 @@ def test_hybrid_search_honors_filter(populated_collection): def test_native_vs_fallback_path_is_observable(populated_collection, backend_name): """isinstance(c, SupportsHybrid) cleanly splits the two paths.""" is_native = isinstance(populated_collection, vd.SupportsHybrid) - if backend_name in {"weaviate", "elasticsearch", "redis"}: + if backend_name in {"weaviate", "elasticsearch", "redis", "lancedb"}: assert is_native, ( f"{backend_name} should be SupportsHybrid in this PR; got {is_native}" ) @@ -314,3 +314,121 @@ def my_lex(collection, query_text, *, limit, filter, **kwargs): assert calls["count"] == 1 ids = [h["id"] for h in hits] assert "noise_1" in ids, "the custom lexical search's hit should fuse in" + + +# ---------- LanceDB: native lexical side via its built-in FTS index -------- # + + +@pytest.fixture +def lance_col(): + pytest.importorskip("lancedb") + col = vd.connect("lancedb").create_collection("lance_hybrid", dimension=2) + col["a"] = vd.Document(id="a", text="the quick brown fox", vector=[1.0, 0.0]) + col["b"] = vd.Document(id="b", text="lazy dog sleeps", vector=[0.0, 1.0]) + return col + + +def _fts_indexed_columns(col) -> list: + return [ + list(i.columns) for i in col.native.list_indices() if i.index_type == "FTS" + ] + + +def test_lancedb_hybrid_uses_native_fts_index(lance_col): + """The first hybrid call builds a LanceDB FTS index on ``text`` and uses it.""" + assert isinstance(lance_col, vd.SupportsHybrid) + assert _fts_indexed_columns(lance_col) == [] + hits = list( + vd.hybrid_search(lance_col, [0.0, 1.0], query_text="quick fox", limit=2) + ) + assert {h["id"] for h in hits} == {"a", "b"} + assert _fts_indexed_columns(lance_col) == [["text"]] + # The lexical primitive alone finds only the term match. + assert [h["id"] for h in lance_col._lexical_query("quick fox", limit=5, + filter=None)] == ["a"] + + +def test_lancedb_lexical_sees_writes_after_index_build(lance_col): + """Upserts and deletes after the FTS index exists are reflected.""" + list(vd.hybrid_search(lance_col, [1.0, 0.0], query_text="fox", limit=1)) + lance_col["c"] = vd.Document(id="c", text="a fox runs", vector=[0.5, 0.5]) + lance_col["a"] = vd.Document(id="a", text="nothing here", vector=[1.0, 0.0]) + ids = [h["id"] for h in lance_col._lexical_query("fox", limit=5, filter=None)] + assert ids == ["c"] + del lance_col["c"] + assert lance_col._lexical_query("fox", limit=5, filter=None) == [] + + +def test_lancedb_lexical_filter_and_odd_queries(lance_col): + """Metadata filters apply; query-syntax characters don't raise.""" + lance_col["c"] = vd.Document( + id="c", text="fox on the hill", vector=[0.5, 0.5], metadata={"k": 1} + ) + hits = lance_col._lexical_query("fox", limit=5, filter={"k": 1}) + assert [h["id"] for h in hits] == ["c"] + assert {"id", "text", "score", "metadata"} <= set(hits[0]) + for odd in ['a "quote', "(", "title:x", "C++?"]: + lance_col._lexical_query(odd, limit=5, filter=None) + + +def test_lancedb_hybrid_before_first_write(): + """A created-but-empty collection returns no hits instead of raising.""" + pytest.importorskip("lancedb") + col = vd.connect("lancedb").create_collection("empty_hybrid", dimension=2) + assert list(vd.hybrid_search(col, [1.0, 0.0], query_text="x", limit=3)) == [] + + +def test_lexical_search_on_native_path_warns(lance_col): + """A custom lexical_search can't be honoured natively; say so instead of ignoring it.""" + calls = [] + + def my_lex(collection, text, *, limit, filter, **kwargs): + calls.append(text) + return [] + + with pytest.warns(UserWarning, match="lexical_search"): + list(vd.hybrid_search(lance_col, [1.0, 0.0], query_text="fox", + limit=1, lexical_search=my_lex)) + assert calls == [] + + +def test_lancedb_first_hybrid_calls_concurrently(): + """Concurrent first hybrid calls must not race on creating the FTS index.""" + from concurrent.futures import ThreadPoolExecutor + + pytest.importorskip("lancedb") + for round_ in range(3): # a fresh table each round hits the creation window + client = vd.connect("lancedb") + col = client.create_collection(f"race{round_}", dimension=2) + col["a"] = vd.Document(id="a", text="the quick brown fox", vector=[1.0, 0.0]) + col["b"] = vd.Document(id="b", text="lazy dog sleeps", vector=[0.0, 1.0]) + # Each worker gets its own collection object, as separate requests would. + cols = [client.get_collection(f"race{round_}") for _ in range(8)] + + def run(c): + return [h["id"] for h in vd.hybrid_search( + c, [1.0, 0.0], query_text="quick fox", limit=2)] + + with ThreadPoolExecutor(8) as pool: + results = list(pool.map(run, cols)) + assert all(r and r[0] == "a" for r in results), results + + +def test_lancedb_hybrid_after_table_dropped_and_recreated(lance_col): + """A dropped-and-recreated table gets a fresh FTS index (no stale cache).""" + client_like = lance_col._db + list(vd.hybrid_search(lance_col, [1.0, 0.0], query_text="fox", limit=1)) + client_like.drop_table(lance_col.name) # dropped behind vd's back + lance_col["c"] = vd.Document(id="c", text="a fox again", vector=[1.0, 0.0]) + hits = list(vd.hybrid_search(lance_col, [1.0, 0.0], query_text="fox", limit=1)) + assert [h["id"] for h in hits] == ["c"] + # and through the client API, with a fresh collection object + client = vd.connect("lancedb") + col = client.create_collection("recreated", dimension=2) + col["a"] = vd.Document(id="a", text="red fox", vector=[1.0, 0.0]) + list(vd.hybrid_search(col, [1.0, 0.0], query_text="fox", limit=1)) + client.delete_collection("recreated") + col2 = client.create_collection("recreated", dimension=2) + col2["b"] = vd.Document(id="b", text="blue fox", vector=[1.0, 0.0]) + assert [h["id"] for h in vd.hybrid_search( + col2, [1.0, 0.0], query_text="fox", limit=1)] == ["b"] diff --git a/tests/test_package.py b/tests/test_package.py new file mode 100644 index 0000000..3ab35b6 --- /dev/null +++ b/tests/test_package.py @@ -0,0 +1,47 @@ +"""Package-level metadata tests.""" + +import importlib.metadata + +import vd + + +def test_version_matches_installed_distribution(): + """``vd.__version__`` must track the released version, not a stale literal.""" + assert vd.__version__ == importlib.metadata.version("vd") + + +def test_bundled_skills_are_spec_clean(): + """Each bundled skill's frontmatter name matches its folder; audience in metadata.""" + import re + + skills = sorted(vd.skills_dir().glob("*/SKILL.md")) + assert len(skills) >= 7 + for path in skills: + front = path.read_text().split("---")[1] + assert re.search(rf"^name: {path.parent.name}$", front, re.M), path + assert not re.search(r"^audience:", front, re.M), path + + +def test_sdist_ships_the_bundled_skills(): + """The sdist (which the wheel is built from) must carry vd/data/skills (#28). + + ``.claude/skills/*`` symlink into ``vd/data/skills/``; if hatch walks + ``.claude`` it keeps those copies and silently drops the real ones. + """ + import pathlib + + import pytest + + sdist = pytest.importorskip("hatchling.builders.sdist") + root = pathlib.Path(__file__).resolve().parent.parent + paths = { # hatchling joins with os.sep; compare POSIX-style (Windows CI) + f.distribution_path.replace("\\", "/") + for f in sdist.SdistBuilder(str(root)).recurse_included_files() + } + expected = { + f"vd/data/skills/{p.parent.name}/SKILL.md" + for p in vd.skills_dir().glob("*/SKILL.md") + } + assert expected <= paths + assert "vd/data/providers.yaml" in paths + assert not any(p.startswith(".claude") for p in paths) diff --git a/tests/test_pgvector_unit.py b/tests/test_pgvector_unit.py new file mode 100644 index 0000000..527b33d --- /dev/null +++ b/tests/test_pgvector_unit.py @@ -0,0 +1,28 @@ +""" +Server-free unit tests for the pgvector adapter's value conversions. + +The live contract tests (``test_core.py`` etc.) only run when a Postgres + +pgvector server is reachable; these run wherever the ``pgvector`` client is +installed. +""" + +import pytest + +pytest.importorskip("pgvector") +pytest.importorskip("psycopg") + +from vd.backends.pgvector import _embedding_to_list # noqa: E402 + + +def test_embedding_to_list_accepts_pgvector_vector(): + """pgvector-python >= 0.5 returns its own (non-iterable) Vector type (#26).""" + from pgvector import Vector + + assert _embedding_to_list(Vector([1.0, 2.5])) == [1.0, 2.5] + + +def test_embedding_to_list_accepts_numpy_and_sequences(): + np = pytest.importorskip("numpy") + assert _embedding_to_list(np.array([1.0, 2.0], dtype="float32")) == [1.0, 2.0] + assert _embedding_to_list((1, 2)) == [1.0, 2.0] + assert all(isinstance(x, float) for x in _embedding_to_list([1, 2])) diff --git a/tests/test_providers.py b/tests/test_providers.py index 79675ee..3deaedd 100644 --- a/tests/test_providers.py +++ b/tests/test_providers.py @@ -55,10 +55,52 @@ def test_unknown_provider_is_none(): def test_install_command(): - assert install_command("qdrant") == "pip install qdrant-client" + # Backends with a vd adapter install through vd's own extra, so the + # printed command installs exactly what the adapter needs. + assert install_command("qdrant") == 'pip install "vd[qdrant]"' + assert install_command("pgvector") == 'pip install "vd[pgvector]"' + # Providers vd has no adapter for fall back to their raw client package. + assert install_command("opensearch") == "pip install opensearch-py" assert "no installation" in install_command("memory") +def _normalize_extra(name: str) -> str: + import re + + return re.sub(r"[-_.]+", "-", name).lower() + + +def test_every_adapter_has_a_matching_vd_extra(): + """install_command's `vd[]` must name a real extra of this package.""" + import importlib.metadata + + extras = { + _normalize_extra(e) + for e in importlib.metadata.metadata("vd").get_all("Provides-Extra") or [] + } + for name in provider_names(): + meta = provider(name) + if meta.get("adapter") and meta.get("pip_packages"): + assert _normalize_extra(meta["adapter"]) in extras, name + + +def test_install_backend_runs_pip_without_shell_quotes(monkeypatch): + import subprocess + import sys + + calls = [] + monkeypatch.setattr(subprocess, "run", lambda args, **kw: calls.append(args)) + cmd = vd.install_backend("qdrant", run=True) + assert cmd == 'pip install "vd[qdrant]"' + assert calls == [[sys.executable, "-m", "pip", "install", "vd[qdrant]"]] + calls.clear() + vd.install_backend("opensearch", run=True) + assert calls == [[sys.executable, "-m", "pip", "install", "opensearch-py"]] + calls.clear() + vd.install_backend("memory", run=True) + assert calls == [] + + def test_memory_always_installed(): assert is_installed("memory") is True diff --git a/tests/test_readme.py b/tests/test_readme.py new file mode 100644 index 0000000..c0e5267 --- /dev/null +++ b/tests/test_readme.py @@ -0,0 +1,40 @@ +""" +The README's Python examples must run and print what they claim. + +Each ```python block is executed in a fresh namespace. Full-line ``# ...`` +comments directly after a ``print(...)`` line are the expected output. +""" + +import contextlib +import io +import pathlib +import re + +import pytest + +README = pathlib.Path(__file__).resolve().parent.parent / "README.md" +BLOCKS = re.findall(r"```python\n(.*?)```", README.read_text(), re.S) + + +def _expected_output(block: str) -> list[str]: + lines = block.splitlines() + expected = [] + for i, line in enumerate(lines): + if "print(" in line: + j = i + 1 + while j < len(lines) and lines[j].startswith("# "): + expected.append(lines[j][2:]) + j += 1 + return expected + + +@pytest.mark.parametrize("block", BLOCKS, ids=[f"block{i}" for i in range(len(BLOCKS))]) +def test_readme_example(block, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + out = io.StringIO() + with contextlib.redirect_stdout(out): + exec(compile(block, str(README), "exec"), {}) + expected = _expected_output(block) + if expected: + got = [line.strip() for line in out.getvalue().splitlines()] + assert got[-len(expected):] == expected diff --git a/vd/__init__.py b/vd/__init__.py index f96e388..d8f812a 100644 --- a/vd/__init__.py +++ b/vd/__init__.py @@ -33,9 +33,14 @@ from __future__ import annotations +from importlib.metadata import PackageNotFoundError as _PackageNotFoundError +from importlib.metadata import version as _version from pathlib import Path as _Path -__version__ = "0.2.0" +try: # the installed distribution is the single source of truth + __version__ = _version("vd") +except _PackageNotFoundError: # pragma: no cover - running from an uninstalled tree + __version__ = "unknown" def skills_dir() -> _Path: @@ -73,10 +78,14 @@ def skills_dir() -> _Path: # ----- async support ------------------------------------------------------- # from vd.asynchronous import ( # noqa: E402 + AsyncAbstractClient, + AsyncAbstractCollection, AsyncClientWrapper, AsyncCollectionWrapper, connect_async, hybrid_search_async, + list_async_backends, + register_async_backend, ) # ----- the entry point & registry ------------------------------------------ # @@ -235,8 +244,12 @@ def skills_dir() -> _Path: "AsyncCollection", "AsyncClientWrapper", "AsyncCollectionWrapper", + "AsyncAbstractClient", + "AsyncAbstractCollection", "connect_async", "hybrid_search_async", + "register_async_backend", + "list_async_backends", # filter language "matches_filter", "validate_filter", diff --git a/vd/asynchronous.py b/vd/asynchronous.py index 1c247aa..b6ca5ef 100644 --- a/vd/asynchronous.py +++ b/vd/asynchronous.py @@ -2,22 +2,26 @@ Async support for ``vd``: universal wrapper + opt-in native implementations. This module gives every ``vd`` backend an ``async``/``await`` surface day one, -without forking the adapter hierarchy. Two pieces: +without forking the adapter hierarchy. Three pieces: - :class:`AsyncCollectionWrapper` / :class:`AsyncClientWrapper` — thin adapters that take any sync :class:`vd.Collection` / :class:`vd.Client` and dispatch every method to :func:`asyncio.to_thread`. This is the **universal fallback**: every backend works through it. -- :func:`connect_async` — the entry point. Mirrors :func:`vd.connect`. If a - backend ships a native async client (Phase 2 follow-ups: chroma, qdrant, - weaviate, elasticsearch, redis, mongodb, lancedb, milvus, pinecone, - turbopuffer), :func:`connect_async` returns *that*; otherwise it returns - the wrapper. - -The asyncio.to_thread wrapper does **not** unblock the event loop — it just -moves blocking calls off the main thread, freeing the loop. For real -non-blocking I/O against a network backend, use a client that satisfies -:class:`vd.SupportsNativeAsync`. +- :class:`AsyncAbstractCollection` / :class:`AsyncAbstractClient` — bases for + **native** async adapters, which do real non-blocking I/O through a + backend's own async SDK. A backend implements a few ``async`` raw + primitives and registers its client with :func:`register_async_backend`. + Native today: ``qdrant`` against a server (``url=``), on + ``qdrant_client.AsyncQdrantClient``. +- :func:`connect_async` — the entry point. Mirrors :func:`vd.connect`. It + returns the backend's registered native client when there is one, and the + wrapper otherwise (or when called with ``native=False``). + +The asyncio.to_thread wrapper does **not** make I/O non-blocking — it moves +blocking calls off the event loop onto a worker thread. For real +non-blocking I/O against a network backend, use a client whose +``native_async`` attribute is ``True`` (see :class:`vd.SupportsNativeAsync`). The module name is ``vd.asynchronous`` (not ``vd.async``) because ``async`` is a Python keyword. @@ -26,16 +30,22 @@ from __future__ import annotations import asyncio +import inspect +from abc import ABCMeta, abstractmethod from typing import Any, AsyncIterator, Callable, Iterable, Optional, Union from vd.base import ( AsyncClient, AsyncCollection, Document, + DocumentInput, Filter, SearchResult, + StaticIndexError, SupportsHybrid, Vector, + _coerce_document, + _CollectionPolicy, ) # --------------------------------------------------------------------------- # @@ -261,26 +271,329 @@ async def __aexit__(self, *exc) -> None: await self.close() +# --------------------------------------------------------------------------- # +# Native async bases — for backends whose SDK ships an async client +# --------------------------------------------------------------------------- # + + +class AsyncAbstractCollection(_CollectionPolicy, metaclass=ABCMeta): + """ + Base class for **native** async collections (the async sibling of + :class:`vd.AbstractCollection`). + + A backend subclasses this and implements ``async`` raw primitives; the + user-facing :class:`~vd.AsyncCollection` surface is provided here, with + the same input coercion, embedding, dimension checks, filter validation + and ``egress`` handling as the sync base (both share one policy mixin). + + Subclass responsibilities (async raw primitives) + ------------------------------------------------ + ``async _write_many(docs)`` + Upsert documents; each ``vector`` is set and dimension-checked. + ``async _read(key) -> Document`` + Fetch one document; raise ``KeyError`` if absent. + ``async _drop(key)`` + Delete one document; raise ``KeyError`` if absent. + ``_keys() -> AsyncIterator[str]`` + An async generator of document ids. + ``async _count() -> int`` + Number of documents. + ``async _query(vector, *, limit, filter, **kwargs) -> list[SearchResult]`` + Raw nearest-neighbor search; ``filter`` is the canonical AST. + """ + + #: Real non-blocking I/O through the backend's async SDK. + native_async: bool = True + + # ----- escape hatch --------------------------------------------------- # + + @property + def native(self) -> Any: + """The raw backend handle (escape hatch), or ``None``.""" + return getattr(self, "_native", None) + + # ----- AsyncCollection contract --------------------------------------- # + + async def get(self, key: str) -> Document: + """Fetch one document; raises ``KeyError`` if absent.""" + return await self._read(key) + + async def set(self, key: str, value: Union[str, tuple, Document]) -> None: + """Insert or replace a document (idempotent upsert).""" + self._check_writable() + doc = self._ensure_vector(_coerce_document(key, value)) + await self._write_many([doc]) + + async def delete(self, key: str) -> None: + """Delete a document; raises ``KeyError`` if absent.""" + self._check_writable() + await self._drop(key) + + async def keys(self) -> AsyncIterator[str]: + """Yield document ids.""" + async for key in self._keys(): + yield key + + async def count(self) -> int: + """Return the number of documents.""" + return await self._count() + + async def search( + self, + query: Union[str, Vector], + *, + limit: int = 10, + filter: Optional[Filter] = None, + egress: Optional[Callable[[SearchResult], Any]] = None, + **kwargs, + ) -> AsyncIterator[SearchResult]: + """Yield the ``limit`` documents most similar to ``query``. + + Same contract as :meth:`vd.AbstractCollection.search`. + """ + from vd.filters import validate_filter + + validate_filter(filter, supported=self.supported_filter_operators) + vector = self._resolve_query(query) + for result in await self._query(vector, limit=limit, filter=filter, **kwargs): + yield egress(result) if egress is not None else result + + # ----- batch convenience ---------------------------------------------- # + + async def add_documents( + self, + documents: Iterable[DocumentInput], + *, + batch_size: int = 100, + ) -> None: + """Batch upsert — mirrors :meth:`vd.AbstractCollection.add_documents`.""" + from vd.util import normalize_document_input + + self._check_writable() + batch: list[Document] = [] + for item in documents: + doc = normalize_document_input(item, auto_id=True) + self._ensure_vector(doc) + batch.append(doc) + if len(batch) >= batch_size: + await self._write_many(batch) + batch = [] + if batch: + await self._write_many(batch) + + async def upsert(self, document: Document) -> None: + """Insert or replace ``document``.""" + await self.set(document.id, document) + + def _check_writable(self) -> None: + if not self.supports_incremental_writes: + raise StaticIndexError( + f"Collection {self.name!r} uses a static index and cannot " + f"accept writes after creation. Rebuild it instead." + ) + + # ----- raw primitives — adapters MUST implement ----------------------- # + + @abstractmethod + async def _write_many(self, docs: list[Document]) -> None: + """Upsert documents (vectors set and dimension-checked).""" + + @abstractmethod + async def _read(self, key: str) -> Document: + """Fetch one document; raise ``KeyError`` if absent.""" + + @abstractmethod + async def _drop(self, key: str) -> None: + """Delete one document; raise ``KeyError`` if absent.""" + + @abstractmethod + def _keys(self) -> AsyncIterator[str]: + """Async-iterate document ids.""" + + @abstractmethod + async def _count(self) -> int: + """Number of documents.""" + + @abstractmethod + async def _query( + self, + vector: Vector, + *, + limit: int, + filter: Optional[Filter], + **kwargs, + ) -> list[SearchResult]: + """Raw nearest-neighbor search.""" + + +class AsyncAbstractClient(metaclass=ABCMeta): + """ + Base class for **native** async clients (the async sibling of + :class:`vd.AbstractClient`). + + A backend implements :meth:`create_collection`, :meth:`get_collection`, + :meth:`delete_collection` and :meth:`list_collections` as coroutines / + async generators; :meth:`get_or_create_collection`, the ``client`` escape + hatch, :meth:`close` and ``async with`` support come for free. Register + the class with :func:`register_async_backend` so :func:`connect_async` + returns it. + + Parameters + ---------- + embedder : callable, optional + A ``text -> vector`` function handed to every collection. + **config + Backend-specific connection configuration. + """ + + #: Real non-blocking I/O through the backend's async SDK. + native_async: bool = True + + #: The registry name of this backend (set by :func:`register_async_backend`). + backend_name: str = "" + + def __init__( + self, + *, + embedder: Optional[Callable[[str], Vector]] = None, + **config, + ): + self._embedder = embedder + self.config = config + + @property + def client(self) -> Any: + """The raw async backend client — a supported, documented escape hatch.""" + return getattr(self, "_client", None) + + @abstractmethod + async def create_collection( + self, + name: str, + *, + dimension: Optional[int] = None, + metric: str = "cosine", + **index_config, + ) -> AsyncAbstractCollection: + """Create a new collection; raise ``ValueError`` if it exists.""" + + @abstractmethod + async def get_collection(self, name: str) -> AsyncAbstractCollection: + """Return an existing collection; raise ``KeyError`` if absent.""" + + @abstractmethod + async def delete_collection(self, name: str) -> None: + """Drop a collection; raise ``KeyError`` if absent.""" + + @abstractmethod + def list_collections(self) -> AsyncIterator[str]: + """Async-iterate collection names.""" + + async def get_or_create_collection( + self, + name: str, + *, + dimension: Optional[int] = None, + metric: str = "cosine", + **index_config, + ) -> AsyncAbstractCollection: + """Return collection ``name``, creating it if missing.""" + try: + return await self.get_collection(name) + except KeyError: + return await self.create_collection( + name, dimension=dimension, metric=metric, **index_config + ) + + async def close(self) -> None: + """Release backend resources (closes the raw client if it can).""" + close = getattr(self.client, "close", None) + if close is not None: + result = close() + if inspect.isawaitable(result): + await result + + async def __aenter__(self) -> "AsyncAbstractClient": + return self + + async def __aexit__(self, *exc) -> None: + await self.close() + + +# --------------------------------------------------------------------------- # +# Native async backend registry +# --------------------------------------------------------------------------- # + +#: name -> factory returning a native async client (or an awaitable of one). +_async_backends: dict[str, Callable[..., Any]] = {} + + +def register_async_backend(name: str, factory: Optional[Callable[..., Any]] = None): + """ + Register a native async client factory for backend ``name``. + + Use as a class decorator on an :class:`AsyncAbstractClient` subclass, or + call it with a ``factory`` (a class or a function, sync or ``async``, + taking the :func:`connect_async` keyword arguments). Once registered, + :func:`connect_async` returns the factory's client instead of the + ``to_thread`` wrapper. + + Examples + -------- + >>> @register_async_backend('example') # doctest: +SKIP + ... class ExampleAsyncClient(AsyncAbstractClient): + ... ... + """ + + def decorator(factory: Callable[..., Any]) -> Callable[..., Any]: + if isinstance(factory, type): + factory.backend_name = name + _async_backends[name] = factory + return factory + + return decorator(factory) if factory is not None else decorator + + +def list_async_backends() -> list[str]: + """ + Return the names of backends with a registered native async client. + + Every other backend still works with :func:`connect_async`, through the + universal ``to_thread`` wrapper. + """ + return sorted(_async_backends) + + # --------------------------------------------------------------------------- # # Entry point # --------------------------------------------------------------------------- # -async def connect_async(backend: str, **kwargs) -> AsyncClient: +async def connect_async(backend: str, *, native: bool = True, **kwargs) -> AsyncClient: """ Async sibling of :func:`vd.connect`. - Returns an :class:`~vd.AsyncClient`. Today every backend goes through - the universal :class:`AsyncClientWrapper` (built on - :func:`asyncio.to_thread`); Phase 2 follow-ups will plug in native async - clients per backend, which :func:`connect_async` will return instead. + Returns an :class:`~vd.AsyncClient`. When the backend has a native async + client (see :func:`list_async_backends`; today ``qdrant``), its registered + factory decides: it may return a native client doing real non-blocking + I/O (qdrant with ``url=``) or the wrapper (embedded qdrant, whose async + client would block the loop). Every other backend goes + through the universal :class:`AsyncClientWrapper`, built on + :func:`asyncio.to_thread`. Check ``client.native_async`` to tell them + apart. Parameters ---------- backend : str Backend name — same vocabulary as :func:`vd.connect`. + native : bool + Use the backend's native async client when one is registered + (default). ``False`` forces the ``to_thread`` wrapper around the + sync adapter. **kwargs - Forwarded to :func:`vd.connect`. + Forwarded to the native client's constructor, or to + :func:`vd.connect` for the wrapper. Both take the same arguments + (``embedder``, ``url``, ``path``, ...). Returns ------- @@ -304,9 +617,12 @@ async def connect_async(backend: str, **kwargs) -> AsyncClient: # but keep it lazy so this module is safe to import standalone). from vd.util import connect - # Per-backend native async adapters can be wired here in Phase 2 by - # checking a registry for an async constructor before falling back. For - # Phase 1 every backend uses the universal wrapper. + factory = _async_backends.get(backend) if native else None + if factory is not None: + client = factory(**kwargs) + if inspect.isawaitable(client): + client = await client + return client sync_client = await asyncio.to_thread(connect, backend, **kwargs) return AsyncClientWrapper(sync_client) @@ -333,11 +649,16 @@ async def hybrid_search_async( """ Async sibling of :func:`vd.hybrid_search`. - If the wrapped sync collection's class supports native hybrid (i.e. - satisfies :class:`~vd.SupportsHybrid`), dispatches the whole fused call - to a worker thread. Otherwise runs the universal client-side BM25 + RRF - fallback in a worker thread too. In both cases the awaitable + async - iterator interface stays uniform. + For a wrapped sync collection, runs :func:`vd.hybrid_search` (native + hybrid if the backend has it, else the client-side BM25 + RRF fallback) + on a worker thread. For a native async collection it awaits the + collection's own ``hybrid_search`` if it has one, and otherwise fuses the + collection's async dense search with a client-side BM25 scan (O(N): it + reads every document) via RRF. On a native collection, an ``async def`` + ``lexical_search`` receives the async collection; a sync one receives a + materialized ``{id: Document}`` dict (every document is read per call) + and runs on a worker thread. Either way the awaitable + async iterator + interface stays uniform. Parameters mirror :func:`vd.hybrid_search` exactly; see that function for the full docs. @@ -367,6 +688,23 @@ async def hybrid_search_async( """ from vd.search import hybrid_search as sync_hybrid_search + if getattr(collection, "native_async", False): + async for hit in _native_hybrid_search_async( + collection, + query, + query_text=query_text, + limit=limit, + filter=filter, + k_dense=k_dense, + k_lexical=k_lexical, + rrf_k=rrf_k, + lexical_search=lexical_search, + egress=egress, + **kwargs, + ): + yield hit + return + sync_collection = getattr(collection, "sync", collection) def _run() -> list[SearchResult]: @@ -391,12 +729,96 @@ def _run() -> list[SearchResult]: yield hit +async def _native_hybrid_search_async( + collection: Any, + query: Union[str, Vector], + *, + query_text: Optional[str], + limit: int, + filter: Optional[Filter], + k_dense: Optional[int], + k_lexical: Optional[int], + rrf_k: int, + lexical_search: Optional[Callable[..., Any]], + egress: Optional[Callable[[SearchResult], Any]], + **kwargs, +) -> AsyncIterator[SearchResult]: + """Hybrid search over a native async collection (see :func:`hybrid_search_async`).""" + from vd.search import _HYBRID_OVERFETCH_FLOOR, BM25Index, _rrf_fuse + + k_dense_eff = ( + k_dense if k_dense is not None else max(4 * limit, _HYBRID_OVERFETCH_FLOOR) + ) + k_lexical_eff = ( + k_lexical if k_lexical is not None else max(4 * limit, _HYBRID_OVERFETCH_FLOOR) + ) + native_hybrid = getattr(collection, "hybrid_search", None) + if native_hybrid is not None: + if lexical_search is not None: + from vd.search import _warn_lexical_search_ignored + + _warn_lexical_search_ignored(collection) + async for hit in native_hybrid( + query, + query_text=query_text, + limit=limit, + filter=filter, + k_dense=k_dense_eff, + k_lexical=k_lexical_eff, + rrf_k=rrf_k, + egress=egress, + **kwargs, + ): + yield hit + return + + if isinstance(query, str): + text = query_text if query_text is not None else query + else: + if query_text is None: + raise ValueError( + "hybrid_search needs a `query_text` for the lexical side when " + "`query` is a vector. Either pass query_text=..., or pass " + "`query` as a string and let the embedder handle both." + ) + text = query_text + if not text: + raise ValueError("hybrid_search needs a non-empty lexical query string.") + + dense = [ + hit async for hit in collection.search(query, limit=k_dense_eff, filter=filter) + ] + if lexical_search is not None and inspect.iscoroutinefunction(lexical_search): + # An ``async def`` lexical search gets the async collection itself. + lexical = await lexical_search( + collection, text, limit=k_lexical_eff, filter=filter + ) + else: + # The default BM25 scan and sync callables (e.g. vd.bm25_lexical_search) + # expect a sync ``id -> Document`` mapping: materialize one. + docs = {key: await collection.get(key) async for key in collection.keys()} + if lexical_search is None: + lexical = BM25Index(docs, filter=filter).search(text, limit=k_lexical_eff) + else: + lexical = await asyncio.to_thread( + lexical_search, docs, text, limit=k_lexical_eff, filter=filter + ) + if inspect.isawaitable(lexical): + lexical = await lexical + for hit in _rrf_fuse([dense, list(lexical)], rrf_k=rrf_k, limit=limit): + yield egress(hit) if egress is not None else hit + + # Re-export SupportsHybrid so users importing from vd.asynchronous have the # whole hybrid surface in one place — even if their native-async adapter # decides to also satisfy SupportsHybrid directly. __all__ = [ "AsyncCollectionWrapper", "AsyncClientWrapper", + "AsyncAbstractCollection", + "AsyncAbstractClient", + "register_async_backend", + "list_async_backends", "connect_async", "hybrid_search_async", "SupportsHybrid", diff --git a/vd/backends/lancedb.py b/vd/backends/lancedb.py index 64aba6a..dbe3321 100644 --- a/vd/backends/lancedb.py +++ b/vd/backends/lancedb.py @@ -8,13 +8,19 @@ ``(id, text, vector, metadata)``; metadata travels as a JSON string and the canonical filter is applied client-side (over-fetching candidates first). +Collections satisfy :class:`vd.SupportsHybrid`: the lexical side runs on +LanceDB's built-in BM25 full-text index over ``text`` (no extra package), +created on the first :meth:`~LanceDBCollection.hybrid_search` call. + Requires: ``pip install lancedb`` """ from __future__ import annotations +import inspect import json import tempfile +import threading from typing import Callable, Iterable, Iterator, Optional try: @@ -40,6 +46,13 @@ ) from vd.util import register_backend +#: Serializes FTS-index creation within the process (see ``_ensure_fts_index``). +_FTS_INDEX_LOCK = threading.Lock() + +#: ``(database uri, table name)`` pairs known to have the FTS index on ``text``. +#: Shared by every collection object; entries are dropped when the table is. +_FTS_READY: set = set() + #: vd metric -> LanceDB distance-type name. _METRIC = {"cosine": "cosine", "l2": "l2", "dot": "dot"} @@ -173,6 +186,144 @@ def _query( ) return apply_client_filter(results, filter, limit=limit) + # ----- native hybrid: LanceDB FTS (BM25) + dense, fused with RRF ------ # + + @staticmethod + def _has_fts_index(table) -> bool: + """Whether ``table`` has LanceDB's full-text index on ``text``.""" + return any( + getattr(index, "index_type", None) == "FTS" + and list(getattr(index, "columns", [])) == ["text"] + for index in table.list_indices() + ) + + def _fts_key(self) -> tuple: + return (getattr(self._db, "uri", id(self._db)), self.name) + + def _ensure_fts_index(self, table) -> None: + """ + Build LanceDB's native full-text index on ``text`` if it is missing. + + Known-indexed tables are cached per ``(database, table)``, so the + check costs nothing after the first call. Creation is serialized by a + process-wide lock (taken only when the index is missing), and a + failed create is tolerated when the index turns out to exist + afterwards (a concurrent creator in another process). + """ + key = self._fts_key() + if key in _FTS_READY: + return + if not self._has_fts_index(table): + with _FTS_INDEX_LOCK: + table = self._table or table + if not self._has_fts_index(table): + try: + if "config" in inspect.signature(table.create_index).parameters: + from lancedb.index import FTS + + table.create_index("text", config=FTS(), replace=False) + else: # pragma: no cover - lancedb before the FTS() config + table.create_fts_index( + "text", replace=False, use_tantivy=False + ) + except Exception: + fresh = self._table + if fresh is None or not self._has_fts_index(fresh): + raise + _FTS_READY.add(key) + + def _lexical_query( + self, + text: str, + *, + limit: int, + filter: Optional[Filter], + **kwargs, + ) -> list[SearchResult]: + """ + BM25 lexical search through LanceDB's built-in full-text index. + + The lexical side of :meth:`hybrid_search`. The index on ``text`` is + built on the first call; rows written afterwards are still searched + (LanceDB scans unindexed rows), so no rebuild is needed for + correctness. Building the index is a write, so on read-only storage + create it beforehand (``collection.native.create_index("text", + config=lancedb.index.FTS())``). Call ``collection.native.optimize()`` + after large ingests to fold new rows into the index for speed. Metadata is + filtered client-side, as in :meth:`_query`. + """ + del kwargs + for attempt in range(2): + table = self._table + if table is None: + return [] + self._ensure_fts_index(table) + table = self._table # reopen: a handle from before the index misses it + if table is None: # dropped in the meantime + return [] + try: + hits = ( + table.search(text, query_type="fts") + .limit(overfetch_limit(limit, filter)) + .to_list() + ) + break + except (ValueError, RuntimeError) as error: + # The table was dropped and recreated since the index was + # cached: forget it and build the index again, once. (Newer + # lancedb raises ValueError here, older releases RuntimeError.) + if attempt or "INVERTED index" not in str(error): + raise + _FTS_READY.discard(self._fts_key()) + results = [] + for hit in hits: + doc = self._to_document(hit) + results.append( + { + "id": doc.id, + "text": doc.text, + "score": float(hit.get("_score", 0.0)), + "metadata": doc.metadata, + } + ) + return apply_client_filter(results, filter, limit=limit) + + def hybrid_search( + self, + query, + *, + query_text=None, + limit: int = 10, + filter: Optional[Filter] = None, + k_dense: Optional[int] = None, + k_lexical: Optional[int] = None, + rrf_k: int = 60, + egress=None, + **kwargs, + ): + """ + Hybrid (dense + native FTS) search, fused client-side with RRF. + + See :class:`vd.SupportsHybrid` for the canonical contract. The dense + side is :meth:`_query`; the lexical side is :meth:`_lexical_query` + (LanceDB's BM25 full-text index, built on first use). Fusion is done + by ``vd`` rather than LanceDB's own ``query_type="hybrid"`` so the + fused score is the same RRF score on every backend. Pass + ``query_text=...`` when ``query`` is a vector. + """ + return self._hybrid_via_rrf( + query, + self._lexical_query, + query_text=query_text, + limit=limit, + filter=filter, + k_dense=k_dense, + k_lexical=k_lexical, + rrf_k=rrf_k, + egress=egress, + **kwargs, + ) + @staticmethod def _to_document(row: dict) -> Document: vector = row.get("vector") @@ -247,6 +398,7 @@ def delete_collection(self, name: str) -> None: if name in set(_table_names(self._client)): self._client.drop_table(name) self._metrics.pop(name, None) + _FTS_READY.discard((getattr(self._client, "uri", id(self._client)), name)) def list_collections(self) -> Iterator[str]: names = set(_table_names(self._client)) | set(self._metrics) diff --git a/vd/backends/pgvector.py b/vd/backends/pgvector.py index 1a4c35f..690a57e 100644 --- a/vd/backends/pgvector.py +++ b/vd/backends/pgvector.py @@ -153,6 +153,21 @@ def _quoted(name: str) -> str: # --------------------------------------------------------------------------- # +def _embedding_to_list(embedding) -> list[float]: + """ + Convert an ``embedding`` column value to ``list[float]``. + + pgvector-python returns a numpy array (< 0.5) or its own ``pgvector.Vector`` + (>= 0.5, not iterable); both expose a list conversion. Plain sequences are + accepted too. + """ + for attr in ("to_list", "tolist"): + convert = getattr(embedding, attr, None) + if convert is not None: + return [float(x) for x in convert()] + return [float(x) for x in embedding] + + class PgvectorCollection(AbstractCollection): """ A ``vd`` collection backed by one Postgres table with a ``vector(N)`` column. @@ -294,7 +309,7 @@ def _read(self, key: str) -> Document: return Document( id=key, text=text or "", - vector=list(embedding) if embedding is not None else None, + vector=_embedding_to_list(embedding) if embedding is not None else None, metadata=metadata if isinstance(metadata, dict) else json.loads(metadata or "{}"), diff --git a/vd/backends/qdrant.py b/vd/backends/qdrant.py index 6c50fa0..09101ef 100644 --- a/vd/backends/qdrant.py +++ b/vd/backends/qdrant.py @@ -12,22 +12,33 @@ unsigned ints, so each id is mapped to a deterministic UUID5 and the original is kept in the point payload. +The backend is also **natively async** against a server: +``vd.connect_async("qdrant", url=...)`` returns a :class:`NativeAsyncQdrantClient` +built on ``qdrant_client.AsyncQdrantClient``. Embedded mode (no ``url``) gets +the thread-pool wrapper, because qdrant-client's embedded async client would +block the event loop. + Requires: ``pip install qdrant-client`` """ from __future__ import annotations import uuid -from typing import Any, Callable, Iterable, Iterator, Optional +from typing import Any, AsyncIterator, Callable, Iterable, Iterator, Optional try: - from qdrant_client import QdrantClient, models + from qdrant_client import AsyncQdrantClient, QdrantClient, models except ImportError as e: # pragma: no cover raise ImportError( "The qdrant backend needs the 'qdrant-client' package. " "Install it with: pip install qdrant-client" ) from e +from vd.asynchronous import ( + AsyncAbstractClient, + AsyncAbstractCollection, + register_async_backend, +) from vd.base import ( AbstractClient, AbstractCollection, @@ -123,6 +134,55 @@ def _compile_field(qkey: str, cond: dict, must: list, must_not: list) -> None: must.append(models.FieldCondition(key=qkey, range=models.Range(**range_kw))) +def _to_point(doc: Document) -> "models.PointStruct": + """Build the Qdrant point for ``doc`` (id mapped to a UUID, payload nested).""" + return models.PointStruct( + id=_point_id(doc.id), + vector=doc.vector, + payload={_ID_KEY: doc.id, _TEXT_KEY: doc.text, "metadata": doc.metadata or {}}, + ) + + +def _point_to_result(point, metric: str) -> SearchResult: + """Convert a scored Qdrant point to a ``vd`` result dict.""" + payload = point.payload or {} + score = point.score + # Qdrant `point.score` per metric (see vd.base "Score semantics"): + # - cosine: cosine similarity in [-1, 1] → matches vd canonical + # - dot: raw inner product → matches vd canonical + # - euclid: a *distance* value (lower-is-better); Qdrant's + # own sort orders ascending in that case. The existing + # transform 1/(1+d) matches vd's canonical l2 score directly + # (no un-negation), so leave it as-is. If a future Qdrant + # client version switches Euclid to higher-is-better, this + # branch must be revisited. + return { + "id": payload.get(_ID_KEY, str(point.id)), + "text": payload.get(_TEXT_KEY, ""), + "score": 1.0 / (1.0 + score) if metric == "l2" else score, + "metadata": payload.get("metadata", {}), + } + + +def _to_document(point) -> Document: + """Convert a retrieved Qdrant point (with payload and vector) to a Document.""" + payload = point.payload or {} + vector = point.vector + return Document( + id=payload.get(_ID_KEY, str(point.id)), + text=payload.get(_TEXT_KEY, ""), + vector=list(vector) if vector is not None else None, + metadata=payload.get("metadata", {}), + ) + + +def _vectors_config(dimension: int, metric: str) -> "models.VectorParams": + """The Qdrant vector config for a collection of ``dimension`` and ``metric``.""" + return models.VectorParams( + size=dimension, distance=_DISTANCE.get(metric, models.Distance.COSINE) + ) + + class QdrantCollection(AbstractCollection): """A collection backed by one Qdrant collection. Native payload filtering.""" @@ -154,30 +214,18 @@ def _ensure_collection(self) -> None: if not self._client.collection_exists(self.name): self._client.create_collection( collection_name=self.name, - vectors_config=models.VectorParams( - size=self.dimension, - distance=_DISTANCE.get(self.metric, models.Distance.COSINE), - ), + vectors_config=_vectors_config(self.dimension, self.metric), ) - @staticmethod - def _payload(doc: Document) -> dict: - return {_ID_KEY: doc.id, _TEXT_KEY: doc.text, "metadata": doc.metadata or {}} - - def _point(self, doc: Document) -> "models.PointStruct": - return models.PointStruct( - id=_point_id(doc.id), vector=doc.vector, payload=self._payload(doc) - ) - # ----- raw primitives ------------------------------------------------- # def _write(self, doc: Document) -> None: self._ensure_collection() - self._client.upsert(self.name, points=[self._point(doc)]) + self._client.upsert(self.name, points=[_to_point(doc)]) def _write_many(self, docs: list[Document]) -> None: self._ensure_collection() - self._client.upsert(self.name, points=[self._point(d) for d in docs]) + self._client.upsert(self.name, points=[_to_point(d) for d in docs]) def _read(self, key: str) -> Document: if not self._client.collection_exists(self.name): @@ -187,7 +235,7 @@ def _read(self, key: str) -> Document: ) if not points: raise KeyError(key) - return self._to_document(points[0]) + return _to_document(points[0]) def _drop(self, key: str) -> None: if not self._client.collection_exists(self.name) or not self._client.retrieve( @@ -235,39 +283,7 @@ def _query( with_payload=True, **kwargs, ) - results = [] - for point in response.points: - payload = point.payload or {} - score = point.score - # Qdrant `point.score` per metric (see vd.base "Score semantics"): - # - cosine: cosine similarity in [-1, 1] → matches vd canonical - # - dot: raw inner product → matches vd canonical - # - euclid: a *distance* value (lower-is-better); Qdrant's - # own sort orders ascending in that case. The existing - # transform 1/(1+d) matches vd's canonical l2 score directly - # (no un-negation), so leave it as-is. If a future Qdrant - # client version switches Euclid to higher-is-better, this - # branch must be revisited. - results.append( - { - "id": payload.get(_ID_KEY, str(point.id)), - "text": payload.get(_TEXT_KEY, ""), - "score": 1.0 / (1.0 + score) if self.metric == "l2" else score, - "metadata": payload.get("metadata", {}), - } - ) - return results - - @staticmethod - def _to_document(point) -> Document: - payload = point.payload or {} - vector = point.vector - return Document( - id=payload.get(_ID_KEY, str(point.id)), - text=payload.get(_TEXT_KEY, ""), - vector=list(vector) if vector is not None else None, - metadata=payload.get("metadata", {}), - ) + return [_point_to_result(point, self.metric) for point in response.points] @register_backend("qdrant") @@ -301,12 +317,11 @@ def __init__( **config, ): super().__init__(embedder=embedder, **config) - if url is not None: - self._client = QdrantClient(url=url, api_key=api_key, **config) - elif path is not None: - self._client = QdrantClient(path=path) - else: - self._client = QdrantClient(location=location or ":memory:") + self._client = QdrantClient( + **_qdrant_client_kwargs( + path=path, url=url, api_key=api_key, location=location, config=config + ) + ) self._metrics: dict[str, str] = {} def create_collection( @@ -356,3 +371,246 @@ def list_collections(self) -> Iterator[str]: def close(self) -> None: """Close the underlying Qdrant client.""" self._client.close() + + +# --------------------------------------------------------------------------- # +# Native async — qdrant_client.AsyncQdrantClient +# --------------------------------------------------------------------------- # + + +def _qdrant_client_kwargs( + *, + path: Optional[str], + url: Optional[str], + api_key: Optional[str], + location: Optional[str], + config: dict, +) -> dict: + """Constructor kwargs shared by the sync and async Qdrant clients.""" + if url is not None: + return {"url": url, "api_key": api_key, **config} + if path is not None: + return {"path": path} + if location and location.startswith(("http://", "https://")): + return {"location": location, "api_key": api_key, **config} + return {"location": location or ":memory:"} + + +class NativeAsyncQdrantCollection(AsyncAbstractCollection): + """ + A Qdrant collection driven by ``AsyncQdrantClient`` — non-blocking I/O. + + Same storage layout, filter translation and scores as + :class:`QdrantCollection`, so data written by either is readable by both. + """ + + supported_filter_operators = SUPPORTED_FILTER_OPERATORS + + def __init__( + self, + name: str, + client: AsyncQdrantClient, + *, + embedder: Optional[Callable[[str], Vector]] = None, + dimension: Optional[int] = None, + metric: str = "cosine", + ): + self.name = name + self._client = client + self._embedder = embedder + self.dimension = dimension + self.metric = metric + + @property + def native(self) -> AsyncQdrantClient: + """The raw ``AsyncQdrantClient`` (escape hatch).""" + return self._client + + async def _ensure_collection(self) -> None: + """Create the Qdrant collection lazily, once the dimension is known.""" + if not await self._client.collection_exists(self.name): + await self._client.create_collection( + collection_name=self.name, + vectors_config=_vectors_config(self.dimension, self.metric), + ) + + # ----- async raw primitives ------------------------------------------ # + + async def _write_many(self, docs: list[Document]) -> None: + await self._ensure_collection() + await self._client.upsert(self.name, points=[_to_point(d) for d in docs]) + + async def _read(self, key: str) -> Document: + if not await self._client.collection_exists(self.name): + raise KeyError(key) + points = await self._client.retrieve( + self.name, ids=[_point_id(key)], with_payload=True, with_vectors=True + ) + if not points: + raise KeyError(key) + return _to_document(points[0]) + + async def _drop(self, key: str) -> None: + if not await self._client.collection_exists( + self.name + ) or not await self._client.retrieve(self.name, ids=[_point_id(key)]): + raise KeyError(key) + await self._client.delete( + self.name, points_selector=models.PointIdsList(points=[_point_id(key)]) + ) + + async def _keys(self) -> AsyncIterator[str]: + if not await self._client.collection_exists(self.name): + return + offset = None + while True: + points, offset = await self._client.scroll( + self.name, limit=256, offset=offset, with_payload=[_ID_KEY] + ) + for point in points: + yield point.payload[_ID_KEY] + if offset is None: + break + + async def _count(self) -> int: + if not await self._client.collection_exists(self.name): + return 0 + return (await self._client.count(self.name)).count + + async def _query( + self, + vector: Vector, + *, + limit: int, + filter: Optional[Filter], + **kwargs, + ) -> list[SearchResult]: + if not await self._client.collection_exists(self.name): + return [] + response = await self._client.query_points( + self.name, + query=vector, + limit=limit, + query_filter=_to_qdrant_filter(filter), + with_payload=True, + **kwargs, + ) + return [_point_to_result(point, self.metric) for point in response.points] + + +class NativeAsyncQdrantClient(AsyncAbstractClient): + """ + Native async Qdrant client — what ``await vd.connect_async("qdrant", url=...)`` + returns for a Qdrant server or cloud cluster. + + Takes the same arguments as :class:`QdrantClientAdapter` (``path``, + ``url``, ``api_key``, ``location``, ``embedder``). It also works embedded + (no ``url``), but there qdrant-client's async client runs synchronous code + inside its coroutines and blocks the event loop, so + :func:`vd.connect_async` returns the thread-pool wrapper for embedded mode + instead (see :func:`_connect_async_qdrant`). + + Examples + -------- + >>> import asyncio, vd + >>> async def go(): + ... async with NativeAsyncQdrantClient() as client: # embedded, for the demo + ... col = await client.create_collection("docs", dimension=2) + ... await col.set("a", vd.Document(id="a", text="x", vector=[1.0, 0.0])) + ... return client.native_async, await col.count() + >>> asyncio.run(go()) + (True, 1) + """ + + backend_name = "qdrant" + + def __init__( + self, + *, + embedder: Optional[Callable[[str], Vector]] = None, + path: Optional[str] = None, + url: Optional[str] = None, + api_key: Optional[str] = None, + location: Optional[str] = None, + **config, + ): + super().__init__(embedder=embedder, **config) + self._client = AsyncQdrantClient( + **_qdrant_client_kwargs( + path=path, url=url, api_key=api_key, location=location, config=config + ) + ) + self._metrics: dict[str, str] = {} + + def _collection( + self, name: str, *, dimension: Optional[int], metric: str + ) -> NativeAsyncQdrantCollection: + return NativeAsyncQdrantCollection( + name, + self._client, + embedder=self._embedder, + dimension=dimension, + metric=metric, + ) + + async def create_collection( + self, + name: str, + *, + dimension: Optional[int] = None, + metric: str = "cosine", + **index_config, + ) -> NativeAsyncQdrantCollection: + if await self._client.collection_exists(name) or name in self._metrics: + raise ValueError(f"Collection {name!r} already exists") + self._metrics[name] = metric + collection = self._collection(name, dimension=dimension, metric=metric) + if dimension is not None: # eager create when the dimension is known + await collection._ensure_collection() + return collection + + async def get_collection(self, name: str) -> NativeAsyncQdrantCollection: + if not await self._client.collection_exists(name) and name not in self._metrics: + raise KeyError(f"Collection {name!r} does not exist") + return self._collection( + name, dimension=None, metric=self._metrics.get(name, "cosine") + ) + + async def delete_collection(self, name: str) -> None: + exists = await self._client.collection_exists(name) + if not exists and name not in self._metrics: + raise KeyError(f"Collection {name!r} does not exist") + if exists: + await self._client.delete_collection(name) + self._metrics.pop(name, None) + + async def list_collections(self) -> AsyncIterator[str]: + response = await self._client.get_collections() + names = {c.name for c in response.collections} | set(self._metrics) + for name in sorted(names): + yield name + + +async def _connect_async_qdrant(**kwargs) -> Any: + """ + The :func:`vd.connect_async` factory for ``qdrant``. + + With ``url=`` (or an ``http(s)://`` ``location=``: a Qdrant server or + cloud cluster) it returns + :class:`NativeAsyncQdrantClient`, which does real non-blocking network I/O. + Without it (embedded ``:memory:`` / ``path=`` mode), qdrant-client's local + async client would run blocking code on the event loop, so it returns the + ``asyncio.to_thread`` wrapper around the sync adapter instead. + """ + location = str(kwargs.get("location") or "") + if kwargs.get("url") is not None or location.startswith(("http://", "https://")): + return NativeAsyncQdrantClient(**kwargs) + import asyncio + + from vd.asynchronous import AsyncClientWrapper + from vd.util import connect + + return AsyncClientWrapper(await asyncio.to_thread(connect, "qdrant", **kwargs)) + + +register_async_backend("qdrant", _connect_async_qdrant) diff --git a/vd/base.py b/vd/base.py index 5ce346d..cf2d9e2 100644 --- a/vd/base.py +++ b/vd/base.py @@ -529,50 +529,15 @@ def _coerce_document(key: str, value: Union[str, tuple, Document]) -> Document: ) -# --------------------------------------------------------------------------- # -# AbstractCollection — adapter-author base -# --------------------------------------------------------------------------- # - - -class AbstractCollection(MutableMapping): +class _CollectionPolicy: """ - Base class implementing the :class:`Collection` contract for adapters. - - A backend subclasses this and implements the *raw primitives* below; - everything users see is provided here, once, uniformly: - - - flexible ``__setitem__`` inputs (text / tuple / :class:`Document`), - - optional text embedding when a ``Document`` arrives without a vector, - - text-query embedding in :meth:`search`, - - central filter validation against :attr:`supported_filter_operators`, - - ``egress`` result transforms, - - batch helpers (:meth:`add_documents`, :meth:`upsert`), - - eager dimension-mismatch detection. - - Subclass responsibilities (raw primitives) - ------------------------------------------ - ``_write(doc)`` - Upsert one document. Its ``vector`` is guaranteed non-``None`` and - dimension-checked. - ``_read(key) -> Document`` - Fetch one document; raise ``KeyError`` if absent. - ``_drop(key)`` - Delete one document; raise ``KeyError`` if absent. - ``_keys() -> Iterator[str]`` - Iterate document ids. - ``_count() -> int`` - Number of documents. - ``_query(vector, *, limit, filter, **kwargs) -> Iterable[SearchResult]`` - Raw nearest-neighbor search. ``filter`` is the canonical AST — the - adapter translates it. Each result is a dict with at least ``id``, - ``text``, ``score``, ``metadata``. + The I/O-free document and query policy shared by every collection base. - Optional overrides - ------------------ - ``_write_many(docs)`` - Efficient bulk upsert. Defaults to a loop over ``_write``. - ``native`` (property) - The raw backend collection handle (escape hatch). + Holds what :class:`AbstractCollection` (sync) and + :class:`vd.asynchronous.AsyncAbstractCollection` (native async) must do + identically: embedding text through the injected embedder, learning and + enforcing the collection dimension, and resolving search / hybrid query + inputs. Nothing here touches a backend, so both bases reuse it as is. """ #: Filter operators this backend can honor. Default: the full language. @@ -677,6 +642,53 @@ def _resolve_hybrid_inputs( raise ValueError("hybrid_search needs a non-empty lexical query string.") return vec, text + +# --------------------------------------------------------------------------- # +# AbstractCollection — adapter-author base +# --------------------------------------------------------------------------- # + + +class AbstractCollection(_CollectionPolicy, MutableMapping): + """ + Base class implementing the :class:`Collection` contract for adapters. + + A backend subclasses this and implements the *raw primitives* below; + everything users see is provided here, once, uniformly: + + - flexible ``__setitem__`` inputs (text / tuple / :class:`Document`), + - optional text embedding when a ``Document`` arrives without a vector, + - text-query embedding in :meth:`search`, + - central filter validation against :attr:`supported_filter_operators`, + - ``egress`` result transforms, + - batch helpers (:meth:`add_documents`, :meth:`upsert`), + - eager dimension-mismatch detection. + + Subclass responsibilities (raw primitives) + ------------------------------------------ + ``_write(doc)`` + Upsert one document. Its ``vector`` is guaranteed non-``None`` and + dimension-checked. + ``_read(key) -> Document`` + Fetch one document; raise ``KeyError`` if absent. + ``_drop(key)`` + Delete one document; raise ``KeyError`` if absent. + ``_keys() -> Iterator[str]`` + Iterate document ids. + ``_count() -> int`` + Number of documents. + ``_query(vector, *, limit, filter, **kwargs) -> Iterable[SearchResult]`` + Raw nearest-neighbor search. ``filter`` is the canonical AST — the + adapter translates it. Each result is a dict with at least ``id``, + ``text``, ``score``, ``metadata``. + + Optional overrides + ------------------ + ``_write_many(docs)`` + Efficient bulk upsert. Defaults to a loop over ``_write``. + ``native`` (property) + The raw backend collection handle (escape hatch). + """ + # ----- MutableMapping interface --------------------------------------- # def __setitem__(self, key: str, value: Union[str, tuple, Document]) -> None: diff --git a/vd/data/providers.yaml b/vd/data/providers.yaml index 38ff847..43b11fd 100644 --- a/vd/data/providers.yaml +++ b/vd/data/providers.yaml @@ -240,7 +240,7 @@ providers: requires_server: false # Milvus Lite is embedded embedded_mode: true managed_free_tier: true # Zilliz free cluster - pip_packages: [pymilvus] + pip_packages: [pymilvus, milvus-lite] # milvus-lite = embedded mode (not native Windows) extras: "[milvus-lite]" license: Apache-2.0 hybrid_search: true @@ -512,7 +512,7 @@ providers: requires_server: true embedded_mode: false managed_free_tier: false # depends on the Postgres host - pip_packages: [pgvector, psycopg] + pip_packages: [pgvector, "psycopg[binary]"] extras: "psycopg[binary]" license: PostgreSQL-permissive hybrid_search: false # manual SQL hybrid diff --git a/vd/data/skills/vd-add-backend/SKILL.md b/vd/data/skills/vd-add-backend/SKILL.md index 32a56ce..d6afe3a 100644 --- a/vd/data/skills/vd-add-backend/SKILL.md +++ b/vd/data/skills/vd-add-backend/SKILL.md @@ -7,7 +7,8 @@ description: >- AbstractClient/AbstractCollection raw-primitive contract, filter handling, capability protocols, the escape hatch, and the provider registry. Trigger on "add a backend to vd", "implement the X backend", "vd adapter". -audience: developers +metadata: + audience: developers --- # Implementing a `vd` backend @@ -88,11 +89,74 @@ class Client(AbstractClient): archetype, pip package, license, docs URLs, `verify_command`, notes) and set its `adapter:` field to your backend name. Add the name to `_BACKEND_MODULES` in `vd/backends/__init__.py`. -8. **Tests:** add the backend to `TESTABLE_BACKENDS` in `tests/conftest.py` if - it can run in plain CI — the parametrized `test_core.py` suite then exercises - it automatically. Otherwise it is "correct-by-construction" (no server here). +8. **Tests:** add the backend to `EMBEDDED_BACKENDS` in `tests/conftest.py` if + it needs no server, or to `SERVER_BACKENDS` (with a TCP probe and connect + kwargs) if it does — the parametrized `test_core.py` and `test_hybrid.py` + suites then exercise it automatically, and server backends are skipped + when their container is down. The root `conftest.py` needs no edit: it + skips doctest collection for any backend module that raises + `ImportError`, so a missing SDK never breaks the run. 9. **Update user skills** in `vd/data/skills/` if the happy path changed. +## Optional: native hybrid search + +If the backend has a real text index (BM25 / full-text), implement a lexical +primitive and let the shared orchestration fuse it with the dense side: + +```python +def _lexical_query(self, text, *, limit, filter, **kwargs): # -> list[dict] + ... # same result shape as _query; apply_client_filter if filtering client-side + +def hybrid_search(self, query, *, query_text=None, limit=10, filter=None, + k_dense=None, k_lexical=None, rrf_k=60, egress=None, **kwargs): + return self._hybrid_via_rrf(query, self._lexical_query, query_text=query_text, + limit=limit, filter=filter, k_dense=k_dense, + k_lexical=k_lexical, rrf_k=rrf_k, egress=egress, + **kwargs) +``` + +Defining `hybrid_search` makes the collection satisfy `SupportsHybrid`, so +`vd.hybrid_search` takes the native path. References: `lancedb.py` (embedded +FTS index), `elasticsearch.py`, `redis.py`, `weaviate.py`. Then add the +backend to the native set in `test_native_vs_fallback_path_is_observable` +(`tests/test_hybrid.py`). + +## Optional: native async + +If the backend's SDK ships an async client, add async siblings in the same +module on the bases from `vd/asynchronous.py`, and register the client: + +```python +from vd.asynchronous import (AsyncAbstractClient, AsyncAbstractCollection, + register_async_backend) + +class NativeAsyncCollection(AsyncAbstractCollection): + async def _write_many(self, docs): ... + async def _read(self, key): ... # raise KeyError if absent + async def _drop(self, key): ... # raise KeyError if absent + async def _keys(self): ... # async generator of ids + async def _count(self): ... + async def _query(self, vector, *, limit, filter, **kwargs): ... # -> list[dict] + +@register_async_backend("") +class NativeAsyncClient(AsyncAbstractClient): + async def create_collection(self, name, *, dimension=None, metric="cosine", **index_config): ... + async def get_collection(self, name): ... + async def delete_collection(self, name): ... + async def list_collections(self): ... # async generator of names +``` + +`AsyncAbstractCollection` shares the sync base's embedding, dimension and +filter-validation policy, so only I/O differs. `vd.connect_async("")` +then returns the native client (`native_async is True`); `native=False` still +gives the `to_thread` wrapper. `register_async_backend(name, factory)` also accepts a +function (sync or `async`) that picks per connection mode: qdrant's factory +returns the native client only for `url=`, because qdrant-client's embedded +async client runs blocking code on the event loop. Measure before claiming +non-blocking. Share pure helpers (filter compiler, row ↔ +`Document` converters) between the sync and async classes rather than copying +them — see `qdrant.py`, and its parity test in `tests/test_async_native.py`. + ## Hard don'ts - Don't embed inside the adapter. Embedding is external — the `embedder` is diff --git a/vd/data/skills/vd-backend-choose/SKILL.md b/vd/data/skills/vd-backend-choose/SKILL.md index edb2590..875fae3 100644 --- a/vd/data/skills/vd-backend-choose/SKILL.md +++ b/vd/data/skills/vd-backend-choose/SKILL.md @@ -1,18 +1,20 @@ --- name: vd-backend-choose description: >- - Backend-selection and setup tooling for the vd package. Use this skill when - the user is picking a vector database with vd, asks "which backend should I - use", weighs persistence / cloud / cost / hybrid / scale trade-offs, hits a - "backend not installed" error, or needs help installing and starting a - vectorDB (pip packages, Docker, API keys, env vars). -audience: users + Backend-selection tooling for the vd package. Use this skill when the user + is picking a vector database with vd, asks "which backend should I use", + compares backends, or weighs persistence / cloud / cost / hybrid search / + async / scale / license trade-offs. For installing, starting and verifying + the chosen backend (pip, Docker, API keys), use vd-setup-backend. +metadata: + audience: users --- -# vd — choosing and setting up a backend +# vd — choosing a backend `vd` knows ~21 vector databases and ships facade adapters for 15. This skill -has two jobs: **choose** the right one, then **set it up**. +**chooses** the right one; **vd-setup-backend** then installs, starts and +verifies it. ## 1. Choose @@ -38,9 +40,23 @@ vd.print_recommendation( ``` Key heuristics it applies: tiny + no persistence → `memory`; already running -Postgres → `pgvector`; no Docker → embedded (`chroma`/`lancedb`); air-gapped → -self-hostable Apache/BSD backends; hybrid wanted → `weaviate`; huge scale → -`milvus`; free managed → `qdrant`. +Postgres → `pgvector`; no Docker → embedded (`chroma`, or `lancedb` when hybrid +search is wanted); air-gapped → self-hostable Apache/BSD backends; hybrid +wanted → `weaviate`; huge scale → `milvus`; free managed → `qdrant`. + +### Capabilities that differ by backend in vd + +`vd.hybrid_search` works on every backend (a client-side BM25 + RRF fallback), +but these backends run the lexical side natively, which scales far better: +`weaviate`, `elasticsearch`, `redis`, and `lancedb` (the only embedded one). +Check with `isinstance(collection, vd.SupportsHybrid)`. + +`vd.connect_async` also works on every backend (a thread-pool wrapper), but +only these do real non-blocking I/O through a native async SDK: +`vd.list_async_backends()` → currently `qdrant`, when connected to a server +(`url=`); embedded Qdrant uses the wrapper. `client.native_async` tells you +which one you got. Prefer a native client for +high-concurrency async apps (FastAPI, Starlette). ### Browse the landscape @@ -71,20 +87,13 @@ and stores **URLs** to live pricing/docs (never cached prices — they drift). ## 2. Set up +Hand over to **vd-setup-backend**. In short: + ```python vd.check_requirements("qdrant") # diagnoses readiness, prints the NEXT STEP -vd.setup_guide("qdrant") # full copy-pasteable playbook (pip/docker/env) -vd.install_backend("qdrant") # returns the pip command; run=True to install +print(vd.setup_guide("qdrant")) # full copy-pasteable playbook (pip/docker/env) ``` -`check_requirements` is archetype-aware: for embedded backends it checks the -pip package (and quirks like sqlite-vec needing SQLite ≥3.41); for server -backends it checks whether something answers on the default port (non-fatal if -the backend also runs embedded); for managed backends it checks the required -env vars (`PINECONE_API_KEY`, `QDRANT_URL`, `MONGODB_URI`, …). It always ends -with one concrete **next step** — a pip command, a `docker run` one-liner, or -an `export VAR=...`. - ## 3. Connect ```python diff --git a/vd/data/skills/vd-ingest/SKILL.md b/vd/data/skills/vd-ingest/SKILL.md index 960b05c..27a6767 100644 --- a/vd/data/skills/vd-ingest/SKILL.md +++ b/vd/data/skills/vd-ingest/SKILL.md @@ -7,7 +7,8 @@ description: >- metadata, and adding the result in batches. Trigger on requests like "load these docs into a vector DB", "chunk this text", "preprocess before embedding", or "bulk insert". -audience: users +metadata: + audience: users --- # Ingesting a corpus into vd diff --git a/vd/data/skills/vd-ops/SKILL.md b/vd/data/skills/vd-ops/SKILL.md index ad8c98d..2a5da84 100644 --- a/vd/data/skills/vd-ops/SKILL.md +++ b/vd/data/skills/vd-ops/SKILL.md @@ -7,7 +7,8 @@ description: >- statistics, finding duplicates / outliers, validating integrity, running health checks, benchmarking search or insert latency, or driving any of the above from the `vd` CLI. -audience: users +metadata: + audience: users --- # Operational tasks on vd collections diff --git a/vd/data/skills/vd-quickstart/SKILL.md b/vd/data/skills/vd-quickstart/SKILL.md index 490c407..fd66574 100644 --- a/vd/data/skills/vd-quickstart/SKILL.md +++ b/vd/data/skills/vd-quickstart/SKILL.md @@ -6,7 +6,8 @@ description: >- a backend, creating a collection, adding documents, and running a query. Also trigger on mentions of "vector database", "embeddings + search", or any time the user imports `vd` for a new task and needs the basic happy-path setup. -audience: users +metadata: + audience: users --- # vd — quickstart diff --git a/vd/data/skills/vd-search/SKILL.md b/vd/data/skills/vd-search/SKILL.md index d12e600..880effe 100644 --- a/vd/data/skills/vd-search/SKILL.md +++ b/vd/data/skills/vd-search/SKILL.md @@ -3,10 +3,11 @@ name: vd-search description: >- Advanced-search tooling for the vd package. Use this skill when the user goes beyond a single basic .search() call with vd — metadata filters with - MongoDB-style operators, multi-query searches, reciprocal rank fusion, - finding documents similar to an existing one, deduplicating result sets, or - searching by a pre-computed query vector. -audience: users + MongoDB-style operators, hybrid (keyword + vector) search, BM25, multi-query + searches, reciprocal rank fusion, finding documents similar to an existing + one, deduplicating result sets, or searching by a pre-computed query vector. +metadata: + audience: users --- # Advanced search with vd @@ -131,6 +132,40 @@ merged = vd.reciprocal_rank_fusion([list_a, list_b], k=60) is rank-based, so it doesn't care that scores from different lists aren't comparable. Prefer RRF over a hand-rolled score average. +## Hybrid search (keyword + vector) + +Dense vectors miss exact terms (product codes, names, rare jargon); keyword +search misses paraphrases. `vd.hybrid_search` runs both and fuses the two +rankings with RRF. It works on **every** backend: + +```python +# Text query: the collection's embedder handles the dense side, +# the same text drives the keyword side. +hits = list(vd.hybrid_search(docs, "neural networks", limit=10)) + +# Pre-computed query vector: say what the keyword side should match. +hits = list(vd.hybrid_search( + docs, query_vec, query_text="neural networks", limit=10, + filter={'year': {'$gte': 2017}}, +)) +``` + +Result dicts have the usual shape; `score` is the fused RRF score. + +- **Native vs fallback.** `weaviate`, `elasticsearch`, `redis` and `lancedb` + run the keyword side on their own text index + (`isinstance(docs, vd.SupportsHybrid)` is `True`). Every other backend uses + a client-side BM25 scan that reads the whole collection on each call: + fine up to ~100k documents, slow beyond. +- **Many queries, fallback backend?** Build the keyword index once: + + ```python + index = vd.BM25Index(docs) # O(N) once + index.search("neural networks", limit=5) + ``` + +- Async code: `vd.hybrid_search_async` takes the same arguments. + ## Finding documents similar to an existing document ```python @@ -230,9 +265,10 @@ final = list(vd.deduplicate_results(merged, key='id', keep='first'))[:10] - **Vector queries must match dimension.** A vector from a different embedding model than the collection's will either error or silently return garbage matches. -- **Don't mistake "score" for "probability".** It's a similarity score, scale - depends on the backend (cosine ∈ [-1, 1] for memory; chroma may return a - distance instead of a similarity). Don't threshold across backends without +- **Don't mistake "score" for "probability".** It's a higher-is-better + similarity on every backend, but its scale depends on the metric (cosine in + [-1, 1], dot unbounded, l2 squashed to (0, 1]) and hybrid results carry an + RRF score instead. Don't threshold across metrics or search types without checking. ## See also diff --git a/vd/data/skills/vd-setup-backend/SKILL.md b/vd/data/skills/vd-setup-backend/SKILL.md new file mode 100644 index 0000000..60f3c94 --- /dev/null +++ b/vd/data/skills/vd-setup-backend/SKILL.md @@ -0,0 +1,173 @@ +--- +name: vd-setup-backend +description: >- + Install, start and verify a vector database for the vd package. Use this + skill when the user needs to get a backend running — pip-installing a + client, starting a server with Docker, setting cloud credentials (API keys, + URIs, env vars), connecting vd to it, and smoke-testing the connection. Also + trigger on BackendNotInstalledError, "connection refused" on a vector DB + port, "how do I run qdrant / pgvector / redis / elasticsearch / milvus + locally", or "vd.connect fails". For choosing which backend to use, see + vd-backend-choose. +metadata: + audience: users +--- + +# vd — install, start and verify a backend + +The loop is always the same: **diagnose → act on the next step → re-diagnose → +connect → smoke-test.** `vd` carries the per-backend facts (pip extra, Docker +one-liner, env vars, docs links), so read them from `vd` instead of guessing. + +## 1. Diagnose + +```python +import vd + +report = vd.check_requirements("qdrant") # prints a report; returns a dict +report["ok"], report["next_step"] +``` + +It checks, depending on the backend's archetype: + +- **embedded** (`memory`, `chroma`, `lancedb`, `sqlite_vec`, `duckdb`, `faiss`): + the client library is importable, plus quirks (sqlite-vec needs SQLite ≥ 3.41 + and extension loading; Milvus Lite is not native-Windows). +- **server** (`qdrant`, `weaviate`, `milvus`, `redis`, `elasticsearch`, + `pgvector`): the client is installed and something answers on the default + port. For `qdrant` and `milvus`, which also run embedded, a missing server + is reported but does not make the backend "not ready". +- **managed** (`pinecone`, `mongodb`, `turbopuffer`): the client is installed + and the required environment variables are set. + +`next_step` is always one concrete action: a pip command, a `docker run` +one-liner, or an `export VAR=...`. Do it, then call `check_requirements` +again until `report["ok"]` is `True`. + +## 2. Get the full playbook + +```python +print(vd.setup_guide("pgvector")) +``` + +prints the install command, the Docker one-liner (server backends), the +credentials to set (managed backends), a verify command and docs links. + +## 3. Install the client + +Every backend with a `vd` adapter installs through `vd`'s own extra, which +pins exactly the libraries the adapter imports: + +```bash +pip install "vd[qdrant]" # one backend (keep the quotes: zsh globs []) +pip install "vd[embedded]" # chroma, qdrant, faiss, lancedb, sqlite_vec, duckdb +pip install "vd[all-backends]" # every client +``` + +`vd.install_command("pgvector")` returns the command as a string, and +`vd.install_backend("pgvector", run=True)` runs it with the current +interpreter's pip. Ask the user before running installs on their machine. + +## 4. Start the backend + +**Embedded — nothing to start.** Pass a path to persist, omit it for a +throwaway store: + +| Backend | Persist with | Default when omitted | +|---|---|---| +| `memory` | (never persists) | in-process dict | +| `chroma` | `persist_directory="./db"` | in-memory, shared by every chroma client in the process | +| `lancedb` | `path="./lance"` (or `s3://…`) | temp directory | +| `sqlite_vec` | `path="./vd.sqlite"` | `":memory:"` | +| `duckdb` | `path="./vd.duckdb"` | `":memory:"` | +| `faiss` | `path="./faiss_dir"` | in-memory | +| `qdrant` | `path="./qdrant_data"` | `":memory:"` | +| `milvus` | `path="./milvus.db"` (Milvus Lite) | temp `.db` | + +**Server — run the container** from `vd.setup_guide(name)`, e.g.: + +```bash +docker run -p 6379:6379 redis:8 +docker run -p 9200:9200 -e discovery.type=single-node -e xpack.security.enabled=false docker.elastic.co/elasticsearch/elasticsearch:8.18.0 +docker run -p 5432:5432 -e POSTGRES_PASSWORD=pw pgvector/pgvector:pg17 +``` + +Wait until the port answers (Elasticsearch takes ~20 s). Disabling security, +as above, is for local development only. + +**Managed — create the account and export credentials** (never commit them): + +| Backend | Credentials | Read from the environment by vd? | +|---|---|---| +| `pinecone` | `PINECONE_API_KEY` | yes (`api_key=` overrides) | +| `mongodb` (Atlas) | `MONGODB_URI` | yes (`uri=` overrides) | +| `turbopuffer` | `TURBOPUFFER_API_KEY` | yes (`api_key=` overrides) | +| `pgvector` | `DATABASE_URL` or `POSTGRES_DSN` | yes (`dsn=` / `url=` override) | +| `qdrant` cloud | `QDRANT_URL`, `QDRANT_API_KEY` | no — pass `url=`, `api_key=` | +| `weaviate` cloud | `WEAVIATE_URL`, `WEAVIATE_API_KEY` | no — pass `url=`, `api_key=` | +| `milvus` / Zilliz | `MILVUS_URI`, `MILVUS_TOKEN` | no — pass `uri=`, `token=` | +| `elasticsearch` | `ELASTICSEARCH_URL`, `ELASTIC_API_KEY` | no — pass `url=`, `api_key=` | + +## 5. Connect + +```python +vd.connect("redis", host="localhost", port=6379) # or url="redis://..." +vd.connect("elasticsearch", url="http://localhost:9200") +vd.connect("pgvector", dsn="postgresql://user:pw@localhost:5432/db") +vd.connect("qdrant", url="http://localhost:6333") # server +vd.connect("qdrant", url=os.environ["QDRANT_URL"], api_key=os.environ["QDRANT_API_KEY"]) +vd.connect("weaviate") # localhost:8080 + gRPC 50051 +vd.connect("milvus", uri="http://localhost:19530") # server; path= for Lite +vd.connect("pinecone") # PINECONE_API_KEY from env +vd.connect("mongodb") # MONGODB_URI from env +``` + +Add `embedder=my_fn` only if you want to pass raw text (see **vd-quickstart**). + +## 6. Smoke-test + +Run this against the new client before building on it. It uses a unique +collection name and cleans up after itself: + +```python +import uuid +import vd + +def smoke_test(client, dim=3): + """Create, write, read, search and drop a scratch collection.""" + name = f"vd_smoke_{uuid.uuid4().hex[:8]}" + col = client.create_collection(name, dimension=dim) + try: + col["a"] = vd.Document(id="a", text="alpha", vector=[1.0, 0.0, 0.0]) + col["b"] = vd.Document(id="b", text="beta", vector=[0.0, 1.0, 0.0]) + assert col["a"].text == "alpha" and len(col) == 2 + top = next(iter(col.search([0.9, 0.1, 0.0], limit=1))) + assert top["id"] == "a", top + finally: + client.delete_collection(name) + return "ok" + +smoke_test(vd.connect("memory")) +``` + +## Troubleshooting + +- **`BackendNotInstalledError`** — the message contains the exact + `pip install "vd[...]"` command. +- **Connection refused** — the server isn't up yet or is on another port; + `check_requirements` probes the default port. Weaviate also needs gRPC port + `50051` published. +- **sqlite-vec: "extension loading disabled"** — this Python's `sqlite3` was + built without it (common with macOS system Python). Use a Homebrew/pyenv + Python or `pysqlite3-binary`. +- **Milvus on Windows** — Milvus Lite has no native-Windows build; use WSL2 + or a Milvus server via `uri=`. +- **pgvector: an error about the `vector` extension** — the Postgres server + lacks pgvector; use the `pgvector/pgvector` image or install the + `postgresql--pgvector` package on the server. +- **Free-tier limits and prices drift.** `vd.provider(name)["docs"]` links to + the live pricing pages; re-check there rather than quoting numbers. + +For contributors, the repo's `tests/docker-compose.yml` brings up pgvector, +Redis Stack, Elasticsearch, Weaviate and MongoDB Atlas Local together for the +live test suite. diff --git a/vd/providers.py b/vd/providers.py index 6f88d7a..22b271b 100644 --- a/vd/providers.py +++ b/vd/providers.py @@ -124,10 +124,18 @@ def install_command(name: str) -> str: """ Return the ``pip install`` command that makes ``name`` usable. + A backend ``vd`` has an adapter for installs through ``vd``'s own extra + (``vd[]``), which pins exactly the client libraries that adapter + imports. The extra is double-quoted so the command is safe to paste into + zsh, bash, PowerShell and cmd. Providers without an adapter get their raw + client package(s). + Examples -------- >>> install_command('qdrant') - 'pip install qdrant-client' + 'pip install "vd[qdrant]"' + >>> install_command('opensearch') + 'pip install opensearch-py' >>> install_command('memory') 'memory needs no installation (built into vd)' """ @@ -137,7 +145,11 @@ def install_command(name: str) -> str: packages = meta.get("pip_packages") or [] if not packages: return f"{name} needs no installation (built into vd)" - return "pip install " + " ".join(packages) + if meta.get("adapter"): + return f'pip install "vd[{meta["adapter"]}]"' + return "pip install " + " ".join( + f'"{p}"' if "[" in p else p for p in packages + ) # --------------------------------------------------------------------------- # diff --git a/vd/requirements.py b/vd/requirements.py index 3e6c54a..67b899a 100644 --- a/vd/requirements.py +++ b/vd/requirements.py @@ -334,6 +334,8 @@ def setup_guide(backend: str) -> str: lines.append("") lines.append("1. Install the client:") lines.append(f" {install_command(backend)}") + if meta.get("adapter") and meta.get("pip_packages"): + lines.append(f" (client libraries: {', '.join(meta['pip_packages'])})") if backend in _DOCKER_COMMANDS: lines += [ @@ -384,8 +386,9 @@ def install_backend(backend: str, *, run: bool = False) -> str: """ cmd = install_command(backend) if run and cmd.startswith("pip install "): + import shlex import subprocess - packages = cmd.removeprefix("pip install ").split() + packages = shlex.split(cmd.removeprefix("pip install ")) subprocess.run([sys.executable, "-m", "pip", "install", *packages], check=True) return cmd diff --git a/vd/search.py b/vd/search.py index f60ded7..3f90dc8 100644 --- a/vd/search.py +++ b/vd/search.py @@ -383,7 +383,7 @@ class BM25Index: Construction is **O(N)** in the collection size; each :meth:`search` is O(matching documents). Fine for prototypes and collections up to ~100k documents; for larger workloads switch to a backend with a native text - index (weaviate, elasticsearch, redis, …). + index (weaviate, elasticsearch, redis, lancedb, …). Parameters ---------- @@ -516,7 +516,7 @@ def bm25_lexical_search( :meth:`BM25Index.search` **per query** instead of calling this function in a loop — the term statistics are then computed once rather than on every call. For larger workloads, switch to a backend with native hybrid search - (weaviate, elasticsearch, redis, …) or pass a custom ``lexical_search`` + (weaviate, elasticsearch, redis, lancedb, …) or pass a custom ``lexical_search`` callable to :func:`hybrid_search` that consults a real text index. Parameters @@ -582,6 +582,21 @@ def _rrf_fuse( return fused[:limit] +def _warn_lexical_search_ignored(collection: Any) -> None: + """Warn that a custom ``lexical_search`` is ignored on a native hybrid path.""" + import warnings + + warnings.warn( + f"lexical_search= is ignored: {type(collection).__name__} runs hybrid " + f"search natively (it satisfies vd.SupportsHybrid), so the backend's own " + f"text index is the lexical side. Call vd.bm25_lexical_search / your " + f"callable and vd.reciprocal_rank_fusion yourself to force a custom " + f"lexical side.", + UserWarning, + stacklevel=3, + ) + + def hybrid_search( collection: Collection, query: Union[str, Vector], @@ -633,7 +648,8 @@ def hybrid_search( lexical_search : callable, optional Custom ``lexical_search(collection, query_text, *, limit, filter, **kwargs) -> list[SearchResult]``. Defaults to - :func:`bm25_lexical_search`. Used only on the fallback path. + :func:`bm25_lexical_search`. Used only on the fallback path; on the + native path it is ignored with a ``UserWarning``. egress : callable, optional Per-result transform applied before yielding. **kwargs @@ -670,6 +686,8 @@ def hybrid_search( # Native path. if isinstance(collection, SupportsHybrid): + if lexical_search is not None: + _warn_lexical_search_ignored(collection) for hit in collection.hybrid_search( query, query_text=query_text, diff --git a/vd/text.py b/vd/text.py index 7bb3496..4e5ad42 100644 --- a/vd/text.py +++ b/vd/text.py @@ -50,7 +50,7 @@ def clean_text( >>> clean_text(text, remove_urls=True) 'Hello World! Visit' >>> clean_text(text, lowercase=True, remove_punctuation=True) - 'hello world visit https examplecom' + 'hello world visit httpsexamplecom' """ if remove_urls: text = re.sub(r"https?://\S+|www\.\S+", "", text) @@ -296,7 +296,7 @@ def extract_metadata( >>> meta['title'] 'My Title' >>> meta['char_count'] - 28 + 30 """ metadata = {}