kiarina-lib-redisearch 1.0.0__py3-none-any.whl

kiarina/lib/redisearch/_core/operations/get.py

@@ -0,0 +1,61 @@
+from typing import Awaitable, Literal, overload
+
+from ..context import RedisearchContext
+from ..utils.unmarshal_mappings import unmarshal_mappings
+from ..views.document import Document
+from .get_key import get_key
+
+
+@overload
+def get(
+    mode: Literal["sync"],
+    ctx: RedisearchContext,
+    id: str,
+) -> Document | None: ...
+
+
+@overload
+def get(
+    mode: Literal["async"],
+    ctx: RedisearchContext,
+    id: str,
+) -> Awaitable[Document | None]: ...
+
+
+def get(
+    mode: Literal["sync", "async"],
+    ctx: RedisearchContext,
+    id: str,
+) -> Document | None | Awaitable[Document | None]:
+    """
+    Get a document from the index.
+    """
+    key = get_key(ctx, id)
+
+    def _after(mapping: dict[bytes, bytes]) -> Document | None:
+        if not mapping:
+            return None
+
+        unmarshaled = unmarshal_mappings(schema=ctx.schema, mapping=mapping)
+
+        return Document(
+            key=key,
+            id=id,
+            mapping=unmarshaled,
+        )
+
+    def _sync() -> Document | None:
+        mapping = ctx.redis.hgetall(key)
+        assert isinstance(mapping, dict)
+        return _after(mapping)
+
+    async def _async() -> Document | None:
+        coro = ctx.redis_async.hgetall(key)
+        assert not isinstance(coro, dict)
+        mapping = await coro
+        return _after(mapping)
+
+    if mode == "sync":
+        return _sync()
+    else:
+        return _async()

kiarina/lib/redisearch/_core/operations/get_info.py

@@ -0,0 +1,155 @@
+from typing import Any, Awaitable, Literal, overload
+
+from ...schema import RedisearchSchema
+from ..context import RedisearchContext
+from ..views.info_result import InfoResult
+
+
+@overload
+def get_info(
+    mode: Literal["sync"],
+    ctx: RedisearchContext,
+) -> InfoResult: ...
+
+
+@overload
+def get_info(
+    mode: Literal["async"],
+    ctx: RedisearchContext,
+) -> Awaitable[InfoResult]: ...
+
+
+def get_info(
+    mode: Literal["sync", "async"],
+    ctx: RedisearchContext,
+) -> InfoResult | Awaitable[InfoResult]:
+    """
+    Get index information using FT.INFO command.
+    """
+
+    def _after(result: dict[str, Any]) -> InfoResult:
+        return InfoResult(
+            index_name=str(result.get("index_name", "")),
+            num_docs=int(result.get("num_docs", 0)),
+            num_terms=int(result.get("num_terms", 0)),
+            num_records=int(result.get("num_records", 0)),
+            index_schema=_parse_schema(ctx.schema, result),
+        )
+
+    def _sync() -> InfoResult:
+        result = ctx.redis.ft(index_name=ctx.settings.index_name).info()  # type: ignore[no-untyped-call]
+        assert isinstance(result, dict)
+        return _after(result)
+
+    async def _async() -> InfoResult:
+        result = await ctx.redis_async.ft(index_name=ctx.settings.index_name).info()  # type: ignore[no-untyped-call]
+        assert isinstance(result, dict)
+        return _after(result)
+
+    if mode == "sync":
+        return _sync()
+    else:
+        return _async()
+
+
+def _parse_schema(schema: RedisearchSchema, result: dict[str, Any]) -> RedisearchSchema:
+    """
+    Parse the schema information from the FT.INFO results
+    """
+    fields: list[dict[str, Any]] = []
+
+    if "attributes" not in result:
+        raise ValueError("The FT.INFO results do not contain attributes.")
+
+    for attr in result["attributes"]:
+        attr_dict = _parse_attribute(attr)
+        field = _parse_field(attr_dict)
+        fields.append(field)
+
+    return RedisearchSchema.from_field_dicts(fields)
+
+
+def _parse_attribute(attr: Any) -> dict[str, Any]:
+    attr_dict = {}
+
+    for i in range(0, len(attr), 2):
+        key = attr[i].decode("utf-8") if isinstance(attr[i], bytes) else attr[i]
+
+        if i + 1 >= len(attr):
+            break
+
+        value = attr[i + 1]
+
+        if isinstance(value, bytes):
+            value = value.decode("utf-8")
+
+        elif isinstance(value, list):
+            value = [v.decode("utf-8") if isinstance(v, bytes) else v for v in value]
+
+        attr_dict[key] = value
+
+    return attr_dict
+
+
+def _parse_field(attr_dict: dict[str, Any]) -> dict[str, Any]:
+    field_type = _get_field_type(attr_dict)
+
+    field_dict: dict[str, Any] = {}
+    field_dict["name"] = str(attr_dict.get("identifier"))
+
+    if field_type == "tag":
+        return _parse_tag_field(field_dict, attr_dict)
+    elif field_type == "numeric":
+        return _parse_numeric_field(field_dict, attr_dict)
+    elif field_type == "text":
+        return _parse_text_field(field_dict, attr_dict)
+    elif field_type == "vector":
+        return _parse_vector_field(field_dict, attr_dict)
+    else:
+        raise ValueError(f"Unknown field type: {field_type}")
+
+
+def _parse_tag_field(
+    field_dict: dict[str, Any], attr_dict: dict[str, Any]
+) -> dict[str, Any]:
+    field_dict["separator"] = str(attr_dict.get("SEPARATOR", ","))
+    field_dict["case_sensitive"] = "CASE_SENSITIVE" in attr_dict
+    field_dict["no_index"] = "NO_INDEX" in attr_dict
+    field_dict["sortable"] = "SORTABLE" in attr_dict
+    return field_dict
+
+
+def _parse_numeric_field(
+    field_dict: dict[str, Any], attr_dict: dict[str, Any]
+) -> dict[str, Any]:
+    field_dict["no_index"] = "NO_INDEX" in attr_dict
+    field_dict["sortable"] = "SORTABLE" in attr_dict
+    return field_dict
+
+
+def _parse_text_field(
+    field_dict: dict[str, Any], attr_dict: dict[str, Any]
+) -> dict[str, Any]:
+    field_dict["weight"] = float(attr_dict.get("WEIGHT", 1.0))
+    field_dict["no_stem"] = "NO_STEM" in attr_dict
+    field_dict["withsuffixtrie"] = "WITHSUFFIX" in attr_dict
+    field_dict["no_index"] = "NO_INDEX" in attr_dict
+    field_dict["sortable"] = "SORTABLE" in attr_dict
+    return field_dict
+
+
+def _parse_vector_field(
+    field_dict: dict[str, Any], attr_dict: dict[str, Any]
+) -> dict[str, Any]:
+    field_dict["dims"] = int(attr_dict.get("dim", 0))
+    field_dict["algorithm"] = str(attr_dict.get("algorithm", ""))
+    field_dict["datatype"] = str(attr_dict.get("data_type", ""))
+    field_dict["distance_metric"] = str(attr_dict.get("distance_metric", ""))
+    return field_dict
+
+
+def _get_field_type(attr_dict: dict[str, Any]) -> str:
+    if "type" not in attr_dict:
+        raise ValueError("The FT.INFO results do not include the field type.")
+
+    return str(attr_dict["type"]).lower()

kiarina/lib/redisearch/_core/operations/get_key.py

@@ -0,0 +1,8 @@
+from ..context import RedisearchContext
+
+
+def get_key(ctx: RedisearchContext, id: str) -> str:
+    """
+    Get the Redis key for a given Redisearch ID.
+    """
+    return f"{ctx.settings.key_prefix}{id}"

kiarina/lib/redisearch/_core/operations/migrate_index.py

@@ -0,0 +1,160 @@
+import logging
+from typing import Any, Awaitable, Literal, overload
+
+from ...schema import RedisearchSchema
+from ..context import RedisearchContext
+from .create_index import create_index
+from .drop_index import drop_index
+from .exists_index import exists_index
+from .get_info import get_info
+
+logger = logging.getLogger(__name__)
+
+
+@overload
+def migrate_index(
+    mode: Literal["sync"],
+    ctx: RedisearchContext,
+) -> bool: ...
+
+
+@overload
+def migrate_index(
+    mode: Literal["async"],
+    ctx: RedisearchContext,
+) -> Awaitable[bool]: ...
+
+
+def migrate_index(
+    mode: Literal["sync", "async"],
+    ctx: RedisearchContext,
+) -> bool | Awaitable[bool]:
+    """
+    Reset the search index.
+    """
+
+    def _log_create_new_index() -> None:
+        logger.info("Createing new index '%s'", ctx.settings.index_name)
+
+    def _log_no_schema_changes() -> None:
+        logger.info("No schema changes detected, migration not needed.")
+
+    def _log_migration_needed(diffs: dict[str, tuple[Any, Any]]) -> None:
+        logger.info("Schema changes detected, migration needed:")
+
+        for path, (old, new) in diffs.items():
+            logger.info(" - %s: %r -> %r", path, old, new)
+
+    def _log_delete_index() -> None:
+        logger.info(
+            "Deleting existing index '%s', data will be re-indexed",
+            ctx.settings.index_name,
+        )
+
+    def _sync() -> bool:
+        if not exists_index(mode="sync", ctx=ctx):
+            _log_create_new_index()
+            create_index(mode="sync", ctx=ctx)
+            return True
+
+        info_result = get_info(mode="sync", ctx=ctx)
+        diffs = _check_schema_changes(current=info_result.index_schema, new=ctx.schema)
+
+        if not diffs:
+            _log_no_schema_changes()
+            return False
+
+        _log_migration_needed(diffs)
+
+        _log_delete_index()
+        drop_index(mode="sync", ctx=ctx, delete_documents=False)
+
+        _log_create_new_index()
+        create_index(mode="sync", ctx=ctx)
+        return True
+
+    async def _async() -> bool:
+        if not await exists_index(mode="async", ctx=ctx):
+            _log_create_new_index()
+            await create_index(mode="async", ctx=ctx)
+            return True
+
+        info_result = await get_info(mode="async", ctx=ctx)
+        diffs = _check_schema_changes(current=info_result.index_schema, new=ctx.schema)
+
+        if not diffs:
+            _log_no_schema_changes()
+            return False
+
+        _log_migration_needed(diffs)
+
+        _log_delete_index()
+        await drop_index(mode="async", ctx=ctx, delete_documents=False)
+
+        _log_create_new_index()
+        await create_index(mode="async", ctx=ctx)
+        return True
+
+    if mode == "sync":
+        return _sync()
+    else:
+        return _async()
+
+
+def _check_schema_changes(
+    current: RedisearchSchema,
+    new: RedisearchSchema,
+) -> dict[str, tuple[Any, Any]]:
+    if current == new:
+        return {}
+
+    return _diff_dict(
+        current.model_dump(),
+        new.model_dump(),
+    )
+
+
+def _diff_dict(
+    d1: dict[str, Any], d2: dict[str, Any], prefix: str = ""
+) -> dict[str, tuple[Any, Any]]:
+    diffs: dict[str, tuple[Any, Any]] = {}
+
+    keys = set(d1.keys()) | set(d2.keys())
+
+    for k in keys:
+        v1, v2 = d1.get(k), d2.get(k)
+        path = f"{prefix}.{k}" if prefix else k
+
+        # Nested dict
+        if isinstance(v1, dict) and isinstance(v2, dict):
+            nested_diff = _diff_dict(v1, v2, prefix=path)
+            diffs.update(nested_diff)
+
+        # Nested list
+        elif isinstance(v1, list) and isinstance(v2, list):
+            max_len = max(len(v1), len(v2))
+
+            for i in range(max_len):
+                p = f"{path}[{i}]"
+
+                try:
+                    item1, item2 = v1[i], v2[i]
+                except IndexError:
+                    diffs[p] = (
+                        v1[i] if i < len(v1) else None,
+                        v2[i] if i < len(v2) else None,
+                    )
+                    continue
+
+                if isinstance(item1, dict) and isinstance(item2, dict):
+                    nested_diff = _diff_dict(item1, item2, prefix=p)
+                    diffs.update(nested_diff)
+
+                elif item1 != item2:
+                    diffs[p] = (item1, item2)
+
+        # Different values
+        elif v1 != v2:
+            diffs[path] = (v1, v2)
+
+    return diffs

kiarina/lib/redisearch/_core/operations/reset_index.py

@@ -0,0 +1,60 @@
+import logging
+from typing import Awaitable, Literal, overload
+
+from ..context import RedisearchContext
+from .create_index import create_index
+from .drop_index import drop_index
+from .exists_index import exists_index
+
+logger = logging.getLogger(__name__)
+
+
+@overload
+def reset_index(
+    mode: Literal["sync"],
+    ctx: RedisearchContext,
+) -> None: ...
+
+
+@overload
+def reset_index(
+    mode: Literal["async"],
+    ctx: RedisearchContext,
+) -> Awaitable[None]: ...
+
+
+def reset_index(
+    mode: Literal["sync", "async"],
+    ctx: RedisearchContext,
+) -> None | Awaitable[None]:
+    """
+    Reset the search index.
+    """
+
+    def _log_delete_index() -> None:
+        logger.info("Deleting existing index '%s'", ctx.settings.index_name)
+
+    def _log_create_index() -> None:
+        logger.info("Creating new index '%s'", ctx.settings.index_name)
+
+    def _sync() -> None:
+        if exists_index(mode="sync", ctx=ctx):
+            _log_delete_index()
+            drop_index(mode="sync", ctx=ctx, delete_documents=True)
+
+        _log_create_index()
+        create_index(mode="sync", ctx=ctx)
+
+    async def _async() -> None:
+        if await exists_index(mode="async", ctx=ctx):
+            _log_delete_index()
+            await drop_index(mode="async", ctx=ctx, delete_documents=True)
+
+        _log_create_index()
+        await create_index(mode="async", ctx=ctx)
+
+    if mode == "sync":
+        _sync()
+        return None
+    else:
+        return _async()

kiarina/lib/redisearch/_core/operations/search.py

@@ -0,0 +1,111 @@
+from typing import Any, Awaitable, Literal, overload
+
+import numpy as np
+from redis.commands.search.query import Query
+from redis.commands.search.result import Result
+
+from ...filter import (
+    RedisearchFilter,
+    RedisearchFilterConditions,
+    create_redisearch_filter,
+)
+from ..context import RedisearchContext
+from ..utils.parse_search_result import parse_search_result
+from ..views.search_result import SearchResult
+from .count import count
+
+
+@overload
+def search(
+    mode: Literal["sync"],
+    ctx: RedisearchContext,
+    vector: list[float],
+    filter: RedisearchFilter | RedisearchFilterConditions | None = None,
+    offset: int | None = None,
+    limit: int | None = None,
+    return_fields: list[str] | None = None,
+) -> SearchResult: ...
+
+
+@overload
+def search(
+    mode: Literal["async"],
+    ctx: RedisearchContext,
+    vector: list[float],
+    filter: RedisearchFilter | RedisearchFilterConditions | None = None,
+    offset: int | None = None,
+    limit: int | None = None,
+    return_fields: list[str] | None = None,
+) -> Awaitable[SearchResult]: ...
+
+
+def search(
+    mode: Literal["sync", "async"],
+    ctx: RedisearchContext,
+    vector: list[float],
+    filter: RedisearchFilter | RedisearchFilterConditions | None = None,
+    offset: int | None = None,
+    limit: int | None = None,
+    return_fields: list[str] | None = None,
+) -> SearchResult | Awaitable[SearchResult]:
+    """
+    Search documents using vector similarity search.
+    """
+    # filter_query
+    if filter is not None:
+        filter = create_redisearch_filter(filter=filter, schema=ctx.schema)
+
+    filter_query = "*" if filter is None else str(filter)
+
+    # vector_field_name
+    vector_field_name = ctx.schema.vector_field.name
+
+    # return_fields
+    return_fields = return_fields or []
+
+    if "distance" not in return_fields:
+        return_fields.append("distance")
+
+    # params
+    params: dict[str, str | int | float | bytes] = {
+        "vector": np.array(vector).astype(ctx.schema.vector_field.dtype).tobytes()
+    }
+
+    def _build_query(limit: int) -> Query:
+        query = Query(
+            f"({filter_query})=>[KNN {limit} @{vector_field_name} $vector AS distance]"
+        )
+
+        if return_fields:
+            query = query.return_fields(*return_fields)
+        else:
+            query = query.no_content()
+
+        query = query.sort_by("distance")
+        query = query.paging(offset or 0, limit)
+
+        return query
+
+    def _parse_search_result(result: Any) -> SearchResult:
+        assert isinstance(result, Result)
+        return parse_search_result(
+            key_prefix=ctx.settings.key_prefix,
+            schema=ctx.schema,
+            return_fields=return_fields,
+            result=result,
+        )
+
+    def _sync() -> SearchResult:
+        query = _build_query(limit or count("sync", ctx, filter).total)
+        result = ctx.redis.ft(ctx.settings.index_name).search(query, params)
+        return _parse_search_result(result)
+
+    async def _async() -> SearchResult:
+        query = _build_query(limit or (await count("async", ctx, filter)).total)
+        result = await ctx.redis_async.ft(ctx.settings.index_name).search(query, params)  # type: ignore
+        return _parse_search_result(result)
+
+    if mode == "sync":
+        return _sync()
+    else:
+        return _async()

kiarina/lib/redisearch/_core/operations/set.py

@@ -0,0 +1,65 @@
+from typing import Any, Awaitable, Literal, overload
+
+from ..context import RedisearchContext
+from ..utils.marshal_mappings import marshal_mappings
+from .get_key import get_key
+
+
+@overload
+def set(
+    mode: Literal["sync"],
+    ctx: RedisearchContext,
+    mapping: dict[str, Any],
+    *,
+    id: str | None = None,
+) -> None: ...
+
+
+@overload
+def set(
+    mode: Literal["async"],
+    ctx: RedisearchContext,
+    mapping: dict[str, Any],
+    *,
+    id: str | None = None,
+) -> Awaitable[None]: ...
+
+
+def set(
+    mode: Literal["sync", "async"],
+    ctx: RedisearchContext,
+    mapping: dict[str, Any],
+    *,
+    id: str | None = None,
+) -> None | Awaitable[None]:
+    """
+    Set a document in the index.
+
+    Fields not present in the schema are saved as they are.
+    Fields present in the schema are converted to the appropriate type and stored.
+    """
+    if id is None:
+        if "id" not in mapping:
+            raise ValueError(
+                'Either "id" parameter or "id" field in mapping must be provided.'
+            )
+
+        id = str(mapping.get("id"))
+
+    key = get_key(ctx, id)
+
+    mapping = marshal_mappings(schema=ctx.schema, mapping=mapping)
+
+    def _sync() -> None:
+        ctx.redis.hset(key, mapping=mapping)
+
+    async def _async() -> None:
+        coro = ctx.redis_async.hset(key, mapping=mapping)
+        assert not isinstance(coro, int)
+        await coro
+
+    if mode == "sync":
+        _sync()
+        return None
+    else:
+        return _async()

kiarina/lib/redisearch/_core/utils/calc_score.py

@@ -0,0 +1,35 @@
+import math
+from typing import Literal
+
+
+def calc_score(
+    distance: float,
+    *,
+    datatype: Literal["FLOAT32", "FLOAT64"],
+    distance_metric: Literal["COSINE", "IP", "L2"],
+) -> float:
+    """
+    Calculate relevance score from distance.
+    """
+    if datatype == "FLOAT32":
+        distance = round(distance, 4)
+    else:
+        distance = round(distance, 7)
+
+    if distance_metric == "COSINE":
+        # Normalise the cosine distance to a score within the range [0, 1]
+        return 1.0 - distance
+
+    elif distance_metric == "IP":
+        # Normalise the inner product distance to a score within the range [0, 1]
+        if distance > 0:
+            return 1.0 - distance
+        else:
+            return -1.0 * distance
+
+    elif distance_metric == "L2":
+        # Convert the Euclidean distance to a similarity score within the range [0, 1]
+        return 1.0 - distance / math.sqrt(2)
+
+    else:
+        raise ValueError(f"Unsupported distance metric: {distance_metric}")