betterdb-retrieval 0.1.0__tar.gz

betterdb_retrieval-0.1.0/.gitignore

@@ -0,0 +1,9 @@
+__pycache__/
+*.pyc
+*.pyo
+.venv/
+dist/
+*.egg-info/
+.cache/
+.ruff_cache/
+.pytest_cache/

betterdb_retrieval-0.1.0/CHANGELOG.md

@@ -0,0 +1,11 @@
+# Changelog
+
+## [0.1.0] - 2026-06-23
+
+### Added
+
+- Initial release. Python equivalent of the TypeScript `@betterdb/retrieval`.
+- `Retriever` — index lifecycle (`create_index` / `drop_index` / `describe_index` / `health`), `upsert` / `delete`, and `query` (vector + filtered + `hybrid="rerank"` KNN search).
+- Typed `RetrievalSchema` (tag / numeric / text fields, HNSW & FLAT vector specs).
+- Shared discovery registry `register` / `unregister`, ownership-checked against the `__betterdb:caches` hash.
+- Observability seams: `RetrievalMetrics` / `RetrievalTracer` protocols and `create_prometheus_metrics` (optional `prometheus` extra).

betterdb_retrieval-0.1.0/PKG-INFO

@@ -0,0 +1,99 @@
+Metadata-Version: 2.4
+Name: betterdb-retrieval
+Version: 0.1.0
+Summary: Developer-facing retrieval SDK over Valkey Search: index lifecycle, upsert, and vector + filtered query.
+Project-URL: Repository, https://github.com/BetterDB-inc/monitor
+License: MIT
+Keywords: embeddings,knn,rag,redis,retrieval,valkey,valkey-search,vector-search
+Requires-Python: >=3.11
+Requires-Dist: betterdb-valkey-search-kit>=0.1.0
+Requires-Dist: prometheus-client>=0.19.0
+Provides-Extra: dev
+Requires-Dist: pytest-asyncio>=0.23.0; extra == 'dev'
+Requires-Dist: pytest>=8.0.0; extra == 'dev'
+Provides-Extra: eval
+Requires-Dist: valkey>=6.0.0; extra == 'eval'
+Description-Content-Type: text/markdown
+
+# @betterdb/retrieval (Python)
+
+`betterdb-retrieval` — developer-facing retrieval SDK over [Valkey Search](https://valkey.io/topics/search/) (`FT.*`): typed index schema, idempotent index lifecycle, upsert/delete, and vector + filtered + hybrid query. This is the Python equivalent of the TypeScript `@betterdb/retrieval` package, built on [`betterdb-valkey-search-kit`](../valkey-search-kit-py/).
+
+## Installation
+
+```bash
+pip install betterdb-retrieval valkey
+```
+
+Requires a Valkey server with the [Valkey Search](https://valkey.io/topics/search/) module loaded.
+
+## Quick start
+
+```python
+from valkey.asyncio import Valkey
+
+from betterdb_retrieval import Retriever, UpsertEntry
+
+client = Valkey.from_url("redis://localhost:6379")
+
+
+async def embed(text: str) -> list[float]:
+    ...  # return an embedding
+
+
+retriever = Retriever(
+    client=client,
+    name="docs",
+    schema={
+        "fields": {
+            "category": {"type": "tag"},
+            "year": {"type": "numeric", "sortable": True},
+        },
+        "vector": {"algorithm": "hnsw", "metric": "cosine"},
+    },
+    embed_fn=embed,
+)
+
+# Create the index if it doesn't exist (idempotent; dims resolved from embed_fn).
+await retriever.create_index()
+
+await retriever.upsert([
+    UpsertEntry(
+        id="doc1",
+        text="Valkey is a high-performance key-value store",
+        fields={"category": "db", "year": 2024},
+    ),
+])
+
+hits = await retriever.query(
+    text="fast in-memory database",
+    k=5,
+    filter={"category": "db"},
+)
+```
+
+## Retriever API
+
+- `create_index()` — create the index if absent (idempotent). Vector dimension is taken from `schema["vector"]["dims"]` or resolved by probing `embed_fn`.
+- `upsert(entries)` — embed each entry's `text` and write it as a hash with its `fields`.
+- `delete(ids)` — delete documents by id.
+- `query(*, k, text=None, vector=None, filter=None, hybrid=None)` — KNN search. Provide `text` (embedded for you) or a precomputed `vector`, a positive `k`, an optional `filter` (tag/numeric fields), and `hybrid="rerank"` to post-process hits through a `rerank_fn`. Returns `list[QueryHit]`.
+- `describe_index()` / `health()` — index stats: doc count, indexing state, dimension, percent indexed, and an optional estimated recall.
+- `drop_index()` — drop the index (no-op if it doesn't exist).
+- `register()` / `unregister()` — publish/remove a discovery marker in the shared `__betterdb:caches` registry, ownership-checked so it never clobbers a foreign cache type.
+
+> `QueryHit.score` is the raw KNN vector **distance** (lower is closer), not a similarity — rank ascending.
+
+## Observability
+
+Pass `metrics` (a `RetrievalMetrics`) and/or `tracer` (a `RetrievalTracer`) to instrument every operation. `create_prometheus_metrics()` provides a ready-made [prometheus-client](https://github.com/prometheus/client_python) implementation.
+
+## Development
+
+```bash
+uv run --extra dev pytest tests -q
+```
+
+## License
+
+MIT

betterdb_retrieval-0.1.0/README.md

@@ -0,0 +1,82 @@
+# @betterdb/retrieval (Python)
+
+`betterdb-retrieval` — developer-facing retrieval SDK over [Valkey Search](https://valkey.io/topics/search/) (`FT.*`): typed index schema, idempotent index lifecycle, upsert/delete, and vector + filtered + hybrid query. This is the Python equivalent of the TypeScript `@betterdb/retrieval` package, built on [`betterdb-valkey-search-kit`](../valkey-search-kit-py/).
+
+## Installation
+
+```bash
+pip install betterdb-retrieval valkey
+```
+
+Requires a Valkey server with the [Valkey Search](https://valkey.io/topics/search/) module loaded.
+
+## Quick start
+
+```python
+from valkey.asyncio import Valkey
+
+from betterdb_retrieval import Retriever, UpsertEntry
+
+client = Valkey.from_url("redis://localhost:6379")
+
+
+async def embed(text: str) -> list[float]:
+    ...  # return an embedding
+
+
+retriever = Retriever(
+    client=client,
+    name="docs",
+    schema={
+        "fields": {
+            "category": {"type": "tag"},
+            "year": {"type": "numeric", "sortable": True},
+        },
+        "vector": {"algorithm": "hnsw", "metric": "cosine"},
+    },
+    embed_fn=embed,
+)
+
+# Create the index if it doesn't exist (idempotent; dims resolved from embed_fn).
+await retriever.create_index()
+
+await retriever.upsert([
+    UpsertEntry(
+        id="doc1",
+        text="Valkey is a high-performance key-value store",
+        fields={"category": "db", "year": 2024},
+    ),
+])
+
+hits = await retriever.query(
+    text="fast in-memory database",
+    k=5,
+    filter={"category": "db"},
+)
+```
+
+## Retriever API
+
+- `create_index()` — create the index if absent (idempotent). Vector dimension is taken from `schema["vector"]["dims"]` or resolved by probing `embed_fn`.
+- `upsert(entries)` — embed each entry's `text` and write it as a hash with its `fields`.
+- `delete(ids)` — delete documents by id.
+- `query(*, k, text=None, vector=None, filter=None, hybrid=None)` — KNN search. Provide `text` (embedded for you) or a precomputed `vector`, a positive `k`, an optional `filter` (tag/numeric fields), and `hybrid="rerank"` to post-process hits through a `rerank_fn`. Returns `list[QueryHit]`.
+- `describe_index()` / `health()` — index stats: doc count, indexing state, dimension, percent indexed, and an optional estimated recall.
+- `drop_index()` — drop the index (no-op if it doesn't exist).
+- `register()` / `unregister()` — publish/remove a discovery marker in the shared `__betterdb:caches` registry, ownership-checked so it never clobbers a foreign cache type.
+
+> `QueryHit.score` is the raw KNN vector **distance** (lower is closer), not a similarity — rank ascending.
+
+## Observability
+
+Pass `metrics` (a `RetrievalMetrics`) and/or `tracer` (a `RetrievalTracer`) to instrument every operation. `create_prometheus_metrics()` provides a ready-made [prometheus-client](https://github.com/prometheus/client_python) implementation.
+
+## Development
+
+```bash
+uv run --extra dev pytest tests -q
+```
+
+## License
+
+MIT

betterdb_retrieval-0.1.0/RELEASE_NOTES.md

@@ -0,0 +1,52 @@
+# betterdb-retrieval v0.1.0
+
+Python port of `@betterdb/retrieval`. Developer-facing retrieval SDK over
+Valkey Search — typed schema, idempotent index lifecycle, upsert/delete, and
+vector + filtered + hybrid query, with built-in observability seams.
+
+Requires Valkey 8+ with the **valkey-search** module (vector index support).
+Works with ElastiCache for Valkey, Memorystore for Valkey, and MemoryDB.
+
+Built on [`betterdb-valkey-search-kit`](https://pypi.org/project/betterdb-valkey-search-kit/).
+
+---
+
+## Installation
+
+```sh
+pip install betterdb-retrieval
+```
+
+---
+
+## What's included
+
+### Retriever
+
+| Method | Description |
+|---|---|
+| `create_index()` | Create or attach to the vector index (idempotent) |
+| `upsert(...)` | Insert or update a document with its vector and fields |
+| `delete(id)` | Delete a document by id |
+| `query(...)` | Vector, filtered, and hybrid (vector + filter) search |
+| `health()` | Index name, doc count, vector dimension |
+
+### Schema & fields
+
+Typed `RetrievalSchema` with TAG / NUMERIC / vector field builders, validated
+against the live `FT.INFO` to tolerate version skew.
+
+### Discovery
+
+Shared discovery registry with atomic register/unregister (EVAL compare-and-set).
+
+### Observability
+
+- `RetrievalMetrics` / `RetrievalTracer` instrumentation seams
+- Prometheus metrics for query latency and result counts
+
+---
+
+## Full changelog
+
+See the repository history for detailed changes.

betterdb_retrieval-0.1.0/betterdb_retrieval/__init__.py

@@ -0,0 +1,96 @@
+"""Developer-facing retrieval SDK over Valkey Search.
+
+Async Python port of the TypeScript ``@betterdb/retrieval`` package: index
+lifecycle, upsert, and vector + filtered query backed by Valkey Search (FT.*).
+"""
+
+from __future__ import annotations
+
+from .discovery import (
+    REGISTRY_KEY,
+    RETRIEVAL_CACHE_TYPE,
+    RetrievalMarker,
+    build_retrieval_marker,
+)
+from .fields import SCORE_FIELD, TEXT_FIELD
+from .ft_create import (
+    build_ft_create_args,
+    index_name,
+    key_prefix,
+    resolve_vector_field_name,
+)
+from .ft_search import QueryFilter, build_ft_search_query
+from .health import IndexHealthSnapshot, RecallEstimator, parse_percent_indexed
+from .prometheus_metrics import (
+    PrometheusRetrievalMetrics,
+    create_prometheus_metrics,
+)
+from .retriever import (
+    EmbedFn,
+    IndexDescription,
+    QueryHit,
+    RerankFn,
+    Retriever,
+    RetrieverClient,
+    UpsertEntry,
+)
+from .schema import (
+    FieldSpec,
+    FtCapabilities,
+    RetrievalSchema,
+    VectorAlgorithm,
+    VectorMetric,
+    VectorSpec,
+)
+from .telemetry import (
+    RetrievalMetrics,
+    RetrievalOperation,
+    RetrievalSpan,
+    RetrievalTracer,
+)
+
+__all__ = [
+    # schema
+    "FieldSpec",
+    "VectorMetric",
+    "VectorAlgorithm",
+    "VectorSpec",
+    "RetrievalSchema",
+    "FtCapabilities",
+    # ft-create
+    "build_ft_create_args",
+    "index_name",
+    "key_prefix",
+    "resolve_vector_field_name",
+    # fields
+    "TEXT_FIELD",
+    "SCORE_FIELD",
+    # ft-search
+    "build_ft_search_query",
+    "QueryFilter",
+    # retriever
+    "Retriever",
+    "RetrieverClient",
+    "IndexDescription",
+    "EmbedFn",
+    "UpsertEntry",
+    "RerankFn",
+    "QueryHit",
+    # discovery
+    "build_retrieval_marker",
+    "REGISTRY_KEY",
+    "RETRIEVAL_CACHE_TYPE",
+    "RetrievalMarker",
+    # health
+    "parse_percent_indexed",
+    "IndexHealthSnapshot",
+    "RecallEstimator",
+    # telemetry
+    "RetrievalMetrics",
+    "RetrievalTracer",
+    "RetrievalSpan",
+    "RetrievalOperation",
+    # prometheus
+    "create_prometheus_metrics",
+    "PrometheusRetrievalMetrics",
+]

betterdb_retrieval-0.1.0/betterdb_retrieval/discovery.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import TypedDict
+
+from .ft_create import index_name
+
+REGISTRY_KEY = "__betterdb:caches"
+RETRIEVAL_PROTOCOL_VERSION = 1
+RETRIEVAL_CACHE_TYPE = "retrieval"
+# TODO: sync with pyproject.toml rather than hardcoding — this drifts on a
+# version bump.
+RETRIEVAL_VERSION = "0.1.0"
+
+
+class RetrievalMarker(TypedDict):
+    type: str
+    prefix: str
+    version: str
+    protocol_version: int
+    capabilities: list[str]
+    index_name: str
+    started_at: str
+
+
+def build_retrieval_marker(name: str, version: str, started_at: str) -> RetrievalMarker:
+    return {
+        "type": RETRIEVAL_CACHE_TYPE,
+        "prefix": name,
+        "version": version,
+        "protocol_version": RETRIEVAL_PROTOCOL_VERSION,
+        "capabilities": ["upsert", "query", "delete"],
+        "index_name": index_name(name),
+        "started_at": started_at,
+    }

betterdb_retrieval-0.1.0/betterdb_retrieval/fields.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+TEXT_FIELD = "__text"
+SCORE_FIELD = "__score"
+
+RESERVED_FIELD_NAMES: tuple[str, ...] = (TEXT_FIELD, SCORE_FIELD)

betterdb_retrieval-0.1.0/betterdb_retrieval/ft_create.py

@@ -0,0 +1,187 @@
+from __future__ import annotations
+
+from .fields import RESERVED_FIELD_NAMES
+from .schema import FieldSpec, FtCapabilities, RetrievalSchema, VectorSpec
+
+_HNSW_DEFAULTS = {"m": 16, "efConstruction": 200, "efRuntime": 10}
+
+_METRIC_MAP = {"cosine": "COSINE", "l2": "L2", "ip": "IP"}
+
+_ALGORITHM_MAP = {"hnsw": "HNSW", "flat": "FLAT"}
+
+
+def _is_positive_int(value: object) -> bool:
+    """Mirror ``Number.isInteger(x) && x > 0`` (accepts integral floats)."""
+    if isinstance(value, bool):
+        return False
+    if isinstance(value, int):
+        return value > 0
+    if isinstance(value, float):
+        return value.is_integer() and value > 0
+    return False
+
+
+def _require_dims(dims: object) -> int:
+    if dims is None or not _is_positive_int(dims):
+        raise ValueError(f"dims must be a positive integer to build FT.CREATE args, got: {dims}")
+    return int(dims)  # type: ignore[arg-type]
+
+
+def _validate_field_names(fields: dict[str, FieldSpec], vector_field_name: str) -> None:
+    for name in fields:
+        if len(name) == 0:
+            raise ValueError("Invalid field name: empty field name is not allowed")
+        if name == vector_field_name:
+            raise ValueError(
+                f"Field name '{name}' collides with the vector field name '{vector_field_name}'"
+            )
+        if name in RESERVED_FIELD_NAMES:
+            raise ValueError(f"Field name '{name}' is reserved and cannot be used in the schema")
+
+
+def _validate_text_field_capabilities(
+    fields: dict[str, FieldSpec], capabilities: FtCapabilities | None
+) -> None:
+    if not (capabilities is not None and capabilities.get("textFields") is False):
+        return
+    text_field_names = [name for name, spec in fields.items() if spec.get("type") == "text"]
+    if text_field_names:
+        raise ValueError(f"Text fields require valkey-search >= 1.2: {', '.join(text_field_names)}")
+
+
+def _validate_flat_hnsw_params(vector: VectorSpec) -> None:
+    if vector.get("algorithm") != "flat":
+        return
+    if vector.get("m") is not None:
+        raise ValueError("FLAT algorithm does not support 'm' parameter")
+    if vector.get("efConstruction") is not None:
+        raise ValueError("FLAT algorithm does not support 'efConstruction' parameter")
+    if vector.get("efRuntime") is not None:
+        raise ValueError("FLAT algorithm does not support 'efRuntime' parameter")
+
+
+def _build_field_args(name: str, spec: FieldSpec) -> list[str]:
+    field_type = spec.get("type")
+    if field_type == "text":
+        return [name, "TEXT"]
+    if field_type == "tag":
+        args = [name, "TAG"]
+        separator = spec.get("separator")
+        if separator is not None:
+            args.extend(["SEPARATOR", separator])
+        return args
+    args = [name, "NUMERIC"]
+    if spec.get("sortable") is True:
+        args.append("SORTABLE")
+    return args
+
+
+def resolve_vector_field_name(vector: VectorSpec) -> str:
+    field_name = vector.get("fieldName")
+    if field_name is None:
+        return "embedding"
+    if field_name.strip() == "":
+        raise ValueError(
+            f"Vector field name must not be empty or whitespace-only, got: '{field_name}'"
+        )
+    return field_name
+
+
+def _build_vector_args(vector: VectorSpec, dims: int) -> list[str]:
+    field_name = resolve_vector_field_name(vector)
+    algorithm = vector.get("algorithm")
+    if algorithm not in _ALGORITHM_MAP:
+        raise ValueError(
+            f"Vector algorithm must be one of {sorted(_ALGORITHM_MAP)}, got: {algorithm!r}"
+        )
+    algo = _ALGORITHM_MAP[algorithm]
+    metric = _METRIC_MAP[vector["metric"]]
+
+    if vector.get("algorithm") == "flat":
+        return [
+            field_name,
+            "VECTOR",
+            algo,
+            "6",
+            "TYPE",
+            "FLOAT32",
+            "DIM",
+            str(dims),
+            "DISTANCE_METRIC",
+            metric,
+        ]
+
+    m = vector.get("m")
+    if m is None:
+        m = _HNSW_DEFAULTS["m"]
+    ef_construction = vector.get("efConstruction")
+    if ef_construction is None:
+        ef_construction = _HNSW_DEFAULTS["efConstruction"]
+    ef_runtime = vector.get("efRuntime")
+    if ef_runtime is None:
+        ef_runtime = _HNSW_DEFAULTS["efRuntime"]
+
+    return [
+        field_name,
+        "VECTOR",
+        algo,
+        "12",
+        "TYPE",
+        "FLOAT32",
+        "DIM",
+        str(dims),
+        "DISTANCE_METRIC",
+        metric,
+        "M",
+        str(m),
+        "EF_CONSTRUCTION",
+        str(ef_construction),
+        "EF_RUNTIME",
+        str(ef_runtime),
+    ]
+
+
+def index_name(name: str) -> str:
+    if name.strip() == "":
+        raise ValueError(f"Index name must not be empty or whitespace-only, got: '{name}'")
+    return f"{name}:idx"
+
+
+def key_prefix(name: str) -> str:
+    if name.strip() == "":
+        raise ValueError(f"Index name must not be empty or whitespace-only, got: '{name}'")
+    return f"{name}:"
+
+
+def build_ft_create_args(
+    name: str,
+    schema: RetrievalSchema,
+    capabilities: FtCapabilities | None = None,
+) -> list[str]:
+    if name.strip() == "":
+        raise ValueError(f"Index name must not be empty or whitespace-only, got: '{name}'")
+
+    dims = _require_dims(schema["vector"].get("dims"))
+    vector_field_name = resolve_vector_field_name(schema["vector"])
+
+    _validate_field_names(schema["fields"], vector_field_name)
+    _validate_text_field_capabilities(schema["fields"], capabilities)
+    _validate_flat_hnsw_params(schema["vector"])
+
+    field_args: list[str] = []
+    for field_name, spec in schema["fields"].items():
+        field_args.extend(_build_field_args(field_name, spec))
+
+    vector_args = _build_vector_args(schema["vector"], dims)
+
+    return [
+        index_name(name),
+        "ON",
+        "HASH",
+        "PREFIX",
+        "1",
+        key_prefix(name),
+        "SCHEMA",
+        *field_args,
+        *vector_args,
+    ]

betterdb_retrieval-0.1.0/betterdb_retrieval/ft_search.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+import math
+
+from betterdb_valkey_search_kit import escape_tag
+
+from .fields import SCORE_FIELD
+from .ft_create import resolve_vector_field_name
+from .schema import RetrievalSchema
+
+QueryFilter = dict[str, "str | int | float"]
+
+
+def _is_number(value: object) -> bool:
+    return isinstance(value, (int, float)) and not isinstance(value, bool)
+
+
+def _build_filter_clause(field: str, value: str | int | float, schema: RetrievalSchema) -> str:
+    spec = schema["fields"].get(field)
+    if spec is None:
+        raise ValueError(f"Cannot filter on unknown field '{field}'")
+    field_type = spec.get("type")
+    if field_type == "tag":
+        return f"@{field}:{{{escape_tag(str(value))}}}"
+    if field_type == "numeric":
+        if not _is_number(value):
+            raise ValueError(
+                f"Numeric filter on field '{field}' requires a number, got: {type(value).__name__}"
+            )
+        if isinstance(value, float) and not math.isfinite(value):
+            raise ValueError(
+                f"Numeric filter on field '{field}' requires a finite number, got: {value!r}"
+            )
+        return f"@{field}:[{value} {value}]"
+    raise ValueError(
+        f"Cannot filter on TEXT field '{field}'; only tag and numeric fields are filterable"
+    )
+
+
+def build_ft_search_query(
+    schema: RetrievalSchema,
+    k: int,
+    filter: QueryFilter | None = None,
+) -> str:
+    vector_field = resolve_vector_field_name(schema["vector"])
+    clauses: list[str] = []
+    if filter is not None:
+        for field, value in filter.items():
+            clauses.append(_build_filter_clause(field, value, schema))
+    filter_expr = f"({' '.join(clauses)})" if clauses else "*"
+    return f"{filter_expr}=>[KNN {k} @{vector_field} $vec AS {SCORE_FIELD}]"

betterdb_retrieval-0.1.0/betterdb_retrieval/health.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable, Optional
+
+
+@dataclass
+class IndexHealthSnapshot:
+    name: str
+    num_docs: int
+    indexing_state: str
+    dims: int
+    percent_indexed: float
+    estimated_recall: Optional[float] = None
+
+
+RecallEstimator = Callable[[IndexHealthSnapshot], float]
+
+_PERCENT_INDEXED_KEYS = ("percent_indexed", "backfill_complete_percent")
+
+
+def _s(x: Any) -> str:
+    if isinstance(x, bytes):
+        try:
+            return x.decode()
+        except UnicodeDecodeError:
+            return ""
+    return str(x)
+
+
+def parse_percent_indexed(info: list[Any]) -> float:
+    """Extract the percent-indexed value from a raw FT.INFO reply.
+
+    valkey-search/RediSearch report either a 0-1 fraction or a 0-100
+    percentage depending on the version; both are normalized to 0-100. Returns
+    0 if the field is absent or unparseable.
+    """
+    for i in range(0, len(info) - 1, 2):
+        if _s(info[i]) not in _PERCENT_INDEXED_KEYS:
+            continue
+        try:
+            value = float(_s(info[i + 1]))
+        except ValueError:
+            return 0.0
+        return value * 100 if value <= 1 else value
+    return 0.0