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

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

@@ -0,0 +1,57 @@
+from typing import Any
+
+import numpy as np
+
+from ...schema import RedisearchSchema
+
+
+def marshal_mappings(
+    *,
+    schema: RedisearchSchema,
+    mapping: dict[str, Any],
+) -> dict[str, Any]:
+    """
+    Convert the mapping to the appropriate format based on the schema
+    """
+    marshaled: dict[str, Any] = {}
+
+    for key, value in mapping.items():
+        field = schema.get_field(key)
+
+        if not field:
+            marshaled[key] = value
+            continue
+
+        if field.type == "tag":
+            if isinstance(value, (list, tuple)):
+                if not field.multiple:
+                    raise ValueError(
+                        f"Field '{key}' does not allow multiple tags. Got: {value}"
+                    )
+
+                marshaled[key] = field.separator.join(str(v) for v in value)
+            else:
+                marshaled[key] = str(value)
+
+        elif field.type == "numeric":
+            try:
+                marshaled[key] = float(value)
+            except (TypeError, ValueError):
+                raise ValueError(
+                    f"Field '{key}' requires a numeric value. Got: {value}"
+                )
+
+        elif field.type == "text":
+            marshaled[key] = str(value)
+
+        elif field.type == "vector":
+            if not isinstance(value, list) or not all(
+                isinstance(v, (float, int)) for v in value
+            ):
+                raise ValueError(
+                    f"Field '{key}' requires a list of floats. Got: {value}"
+                )
+
+            marshaled[key] = np.array(value).astype(field.dtype).tobytes()
+
+    return marshaled

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

@@ -0,0 +1,57 @@
+from redis.commands.search.result import Result
+
+from ...schema import RedisearchSchema
+from ..views.document import Document
+from ..views.search_result import SearchResult
+from .calc_score import calc_score
+
+
+def parse_search_result(
+    *,
+    key_prefix: str,
+    schema: RedisearchSchema,
+    return_fields: list[str] | None,
+    result: Result,
+) -> SearchResult:
+    documents: list[Document] = []
+
+    for doc in result.docs:
+        # key
+        key: str = doc.id
+
+        if not key.startswith(key_prefix):
+            raise ValueError(
+                f"Document ID {doc.id} does not start with key prefix {key_prefix}"
+            )
+
+        # id
+        id = key[len(key_prefix) :]
+
+        # mapping
+        mapping = doc.__dict__
+
+        if "id" in mapping:
+            if "id" in (return_fields or []):
+                mapping["id"] = id
+            else:
+                mapping.pop("id")
+
+        mapping.pop("payload", None)
+
+        # score
+        score = 0.0
+
+        if distance := mapping.pop("distance", None):
+            score = calc_score(
+                float(distance),
+                datatype=schema.vector_field.datatype,
+                distance_metric=schema.vector_field.distance_metric,
+            )
+
+        documents.append(Document(key=key, id=id, score=score, mapping=mapping))
+
+    return SearchResult(
+        total=result.total,
+        duration=result.duration,
+        documents=documents,
+    )

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

@@ -0,0 +1,57 @@
+from typing import Any
+
+import numpy as np
+
+from ...schema import RedisearchSchema
+
+
+def unmarshal_mappings(
+    *,
+    schema: RedisearchSchema,
+    mapping: dict[bytes, bytes],
+) -> dict[str, Any]:
+    """
+    Convert the mapping from the appropriate format based on the schema
+    """
+    unmarshaled: dict[str, Any] = {}
+
+    for bkey, value in mapping.items():
+        key = bkey.decode("utf-8")
+
+        field = schema.get_field(key)
+
+        if not field:
+            try:
+                unmarshaled[key] = _decode_numeric(value)
+            except ValueError:
+                unmarshaled[key] = value.decode("utf-8")
+
+        elif field.type == "tag":
+            if field.multiple:
+                unmarshaled[key] = value.decode("utf-8").split(field.separator)
+            else:
+                unmarshaled[key] = value.decode("utf-8").split(field.separator).pop(0)
+
+        elif field.type == "numeric":
+            unmarshaled[key] = _decode_numeric(value)
+
+        elif field.type == "text":
+            unmarshaled[key] = value.decode("utf-8")
+
+        elif field.type == "vector":
+            unmarshaled[key] = np.frombuffer(value, dtype=field.dtype).tolist()
+
+    return unmarshaled
+
+
+def _decode_numeric(value: bytes) -> float | int:
+    try:
+        return int(value)
+    except ValueError:
+        pass
+    try:
+        return float(value)
+    except ValueError:
+        pass
+
+    raise ValueError(f"Cannot decode numeric value: {value.decode('utf-8')}")

kiarina/lib/redisearch/_core/views/document.py

@@ -0,0 +1,25 @@
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class Document(BaseModel):
+    """
+    Redisearch Document
+    """
+
+    key: str = ""
+    """
+    Redis key
+
+    {key_prefix}:{id} is the key
+    """
+
+    id: str = ""
+    """Redisearch document ID"""
+
+    score: float = 0.0
+    """Redisearch document score"""
+
+    mapping: dict[str, Any] = Field(default_factory=dict)
+    """Redisearch document mapping"""

kiarina/lib/redisearch/_core/views/info_result.py

@@ -0,0 +1,24 @@
+from pydantic import BaseModel
+
+from ...schema import RedisearchSchema
+
+
+class InfoResult(BaseModel):
+    """
+    Model representing FT.INFO results
+    """
+
+    index_name: str
+    """Index name"""
+
+    num_docs: int
+    """Number of documents"""
+
+    num_terms: int
+    """Number of terms"""
+
+    num_records: int
+    """Number of records"""
+
+    index_schema: RedisearchSchema
+    """Index schema"""

kiarina/lib/redisearch/_core/views/search_result.py

@@ -0,0 +1,31 @@
+from pydantic import BaseModel, Field
+
+from .document import Document
+
+
+class SearchResult(BaseModel):
+    """
+    Model representing FT.SEARCH results
+    """
+
+    total: int = 0
+    """
+    Total number of results
+
+    For searches not using KNN (RedisearchClient.find):
+        This indicates the total number of documents matching the query,
+        regardless of the limit parameter.
+
+    For vector search using KNN (RedisearchClient.search):
+        When using KNN, the total in FT.SEARCH depends on the value of k.
+        If limit is not specified,
+        the number of records in the population after filtering will be total.
+        When a limit is specified, k becomes equal to the limit,
+        so the value of total is less than or equal to the limit.
+    """
+
+    duration: float = 0.0
+    """The execution time of the query in milliseconds"""
+
+    documents: list[Document] = Field(default_factory=list)
+    """List of documents"""

kiarina/lib/redisearch/_sync/client.py

@@ -0,0 +1,179 @@
+from typing import Any
+
+from redis import Redis
+
+from .._core.context import RedisearchContext
+from .._core.operations.count import count
+from .._core.operations.create_index import create_index
+from .._core.operations.delete import delete
+from .._core.operations.drop_index import drop_index
+from .._core.operations.exists_index import exists_index
+from .._core.operations.find import find
+from .._core.operations.get import get
+from .._core.operations.get_info import get_info
+from .._core.operations.get_key import get_key
+from .._core.operations.migrate_index import migrate_index
+from .._core.operations.reset_index import reset_index
+from .._core.operations.search import search
+from .._core.operations.set import set
+from .._core.views.document import Document
+from .._core.views.info_result import InfoResult
+from .._core.views.search_result import SearchResult
+from ..filter import RedisearchFilter, RedisearchFilterConditions
+from ..settings import RedisearchSettings
+
+
+class RedisearchClient:
+    """
+    Redisearch client
+    """
+
+    def __init__(
+        self,
+        settings: RedisearchSettings,
+        *,
+        redis: Redis,
+    ) -> None:
+        """
+        Initialize the Redisearch client.
+        """
+        if redis.get_encoder().decode_responses:
+            # As the vector field in Redisearch is expected to be handled as bytes in redis-py,
+            raise ValueError("Redis client must have decode_responses=False")
+
+        self.ctx: RedisearchContext = RedisearchContext(settings=settings, _redis=redis)
+        """Redisearch context"""
+
+    # --------------------------------------------------
+    # Index operations
+    # --------------------------------------------------
+
+    def exists_index(self) -> bool:
+        """
+        Check if the index exists.
+        """
+        return exists_index("sync", self.ctx)
+
+    def create_index(self) -> None:
+        """
+        Create the search index.
+        """
+        create_index("sync", self.ctx)
+
+    def drop_index(self, *, delete_documents: bool = False) -> bool:
+        """
+        Delete the index.
+        """
+        return drop_index("sync", self.ctx, delete_documents=delete_documents)
+
+    def reset_index(self) -> None:
+        """
+        Reset the search index.
+        """
+        reset_index("sync", self.ctx)
+
+    def migrate_index(self) -> None:
+        """
+        Migrate the search index.
+        """
+        migrate_index("sync", self.ctx)
+
+    def get_info(self) -> InfoResult:
+        """
+        Get index information using FT.INFO command.
+        """
+        return get_info("sync", self.ctx)
+
+    # --------------------------------------------------
+    # Data operations
+    # --------------------------------------------------
+
+    def set(self, mapping: dict[str, Any], *, id: str | None = None) -> None:
+        """
+        Set a hash value.
+
+        If no id is specified, the mapping must contain an "id" field.
+        """
+        set("sync", self.ctx, mapping, id=id)
+
+    def delete(self, id: str) -> None:
+        """
+        Delete a document from the index.
+        """
+        delete("sync", self.ctx, id)
+
+    def get(self, id: str) -> Document | None:
+        """
+        Get a document from the index.
+        """
+        return get("sync", self.ctx, id)
+
+    # --------------------------------------------------
+    # Search operations
+    # --------------------------------------------------
+
+    def count(
+        self,
+        *,
+        filter: RedisearchFilter | RedisearchFilterConditions | None = None,
+    ) -> SearchResult:
+        """
+        Count documents matching the filter.
+        """
+        return count("sync", self.ctx, filter=filter)
+
+    def find(
+        self,
+        *,
+        filter: RedisearchFilter | RedisearchFilterConditions | None = None,
+        sort_by: str | None = None,
+        sort_desc: bool = False,
+        offset: int | None = None,
+        limit: int | None = None,
+        return_fields: list[str] | None = None,
+    ) -> SearchResult:
+        """
+        Find documents matching the filter criteria.
+        """
+        return find(
+            "sync",
+            self.ctx,
+            filter=filter,
+            sort_by=sort_by,
+            sort_desc=sort_desc,
+            offset=offset,
+            limit=limit,
+            return_fields=return_fields,
+        )
+
+    def search(
+        self,
+        *,
+        vector: list[float],
+        filter: RedisearchFilter | RedisearchFilterConditions | None = None,
+        offset: int | None = None,
+        limit: int | None = None,
+        return_fields: list[str] | None = None,
+    ) -> SearchResult:
+        """
+        Search documents using vector similarity search.
+        """
+        return search(
+            "sync",
+            self.ctx,
+            vector=vector,
+            filter=filter,
+            offset=offset,
+            limit=limit,
+            return_fields=return_fields,
+        )
+
+    # --------------------------------------------------
+    # Utilities
+    # --------------------------------------------------
+
+    def get_key(self, id: str) -> str:
+        """
+        Get the Redis key for a given Redisearch ID.
+        """
+        return get_key(self.ctx, id)

kiarina/lib/redisearch/_sync/registry.py

@@ -0,0 +1,16 @@
+import redis
+
+from ..settings import settings_manager
+from .client import RedisearchClient
+
+
+def create_redisearch_client(
+    config_key: str | None = None,
+    *,
+    redis: redis.Redis,
+) -> RedisearchClient:
+    """
+    Create a Redisearch client.
+    """
+    settings = settings_manager.get_settings_by_key(config_key)
+    return RedisearchClient(settings, redis=redis)

kiarina/lib/redisearch/asyncio.py

@@ -0,0 +1,33 @@
+import logging
+from importlib import import_module
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ._async.client import RedisearchClient
+    from ._async.registry import create_redisearch_client
+    from .settings import RedisearchSettings, settings_manager
+
+__all__ = [
+    "create_redisearch_client",
+    "RedisearchClient",
+    "RedisearchSettings",
+    "settings_manager",
+]
+
+logging.getLogger(__name__).addHandler(logging.NullHandler())
+
+
+def __getattr__(name: str) -> object:
+    if name not in __all__:
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+    module_map = {
+        "create_redisearch_client": "._async.registry",
+        "RedisearchClient": "._async.client",
+        "RedisearchSettings": ".settings",
+        "settings_manager": ".settings",
+    }
+
+    parent = __name__.rsplit(".", 1)[0]
+    globals()[name] = getattr(import_module(module_map[name], parent), name)
+    return globals()[name]

kiarina/lib/redisearch/filter/__init__.py

@@ -0,0 +1,61 @@
+"""
+Module for constructing filter queries for Redisearch
+
+RedisearchFilter can be combined using & and | operators to create
+complex logical expressions that are evaluated in the Redis Query language.
+
+This interface allows users to construct complex queries without needing to know
+the Redis Query language.
+
+Filter-based fields are not initialised directly.
+Instead, they are constructed by combining RedisFilterFields
+using the & and | operators.
+
+Examples:
+    >>> import kiarina.lib.redisearch.filter as rf
+    >>> filter = (rf.Tag("color") == "blue") & (rf.Numeric("price") < 100)
+    >>> print(str(filter))
+    (@color:{blue} @price:[-inf (100)])
+
+All examples:
+    >>> import kiarina.lib.redisearch.filter as rf
+    >>>
+    >>> rf.Tag("color") == "blue"
+    >>> rf.Tag("color") == ["blue", "red"]
+    >>> rf.Tag("color") != "blue"
+    >>> rf.Tag("color") != ["blue", "red"]
+    >>>
+    >>> rf.Numeric("price") == 100
+    >>> rf.Numeric("price") != 100
+    >>> rf.Numeric("price") > 100
+    >>> rf.Numeric("price") < 100
+    >>> rf.Numeric("price") >= 100
+    >>> rf.Numeric("price") <= 100
+    >>>
+    >>> rf.Text("title") == "hello"
+    >>> rf.Text("title") != "hello"
+    >>> rf.Text("title") % "*hello*"
+    >>>
+    >>> (rf.Tag("color") == "blue") & (rf.Numeric("price") < 100)
+    >>> (rf.Tag("color") == "blue") | (rf.Numeric("price")
+"""
+
+from ._field.numeric import Numeric
+from ._field.tag import Tag
+from ._field.text import Text
+from ._model import RedisearchFilter
+from ._registry import create_redisearch_filter
+from ._types import RedisearchFilterConditions
+
+__all__ = [
+    # ._field
+    "Numeric",
+    "Tag",
+    "Text",
+    # ._model
+    "RedisearchFilter",
+    # ._registry
+    "create_redisearch_filter",
+    # ._types
+    "RedisearchFilterConditions",
+]

kiarina/lib/redisearch/filter/_decorators.py

@@ -0,0 +1,28 @@
+from functools import wraps
+from typing import Any, Callable
+
+
+def check_operator_misuse(func: Callable[..., Any]) -> Callable[..., Any]:
+    """
+    Decorator to check misuse of the equality operator
+    """
+
+    @wraps(func)
+    def wrapper(instance: Any, *args: Any, **kwargs: Any) -> Any:
+        other = kwargs.get("other") if "other" in kwargs else None
+
+        if not other:
+            for arg in args:
+                if isinstance(arg, type(instance)):
+                    other = arg
+                    break
+
+        if isinstance(other, type(instance)):
+            raise ValueError(
+                "The equality operator is overridden when creating a FilterExpression."
+                "Use .equals() for equality checks."
+            )
+
+        return func(instance, *args, **kwargs)
+
+    return wrapper

kiarina/lib/redisearch/filter/_enums.py

@@ -0,0 +1,28 @@
+from enum import Enum
+
+
+class RedisearchFilterOperator(Enum):
+    """
+    Filter operator enumeration type
+    """
+
+    EQ = 1
+    """Equal"""
+    NE = 2
+    """Not Equal"""
+    LT = 3
+    """Less Than"""
+    GT = 4
+    """Greater Than"""
+    LE = 5
+    """Less Than or Equal"""
+    GE = 6
+    """Greater Than or Equal"""
+    OR = 7
+    """Logical OR"""
+    AND = 8
+    """Logical AND"""
+    LIKE = 9
+    """LIKE"""
+    IN = 10
+    """IN"""

kiarina/lib/redisearch/filter/_field/__init__.py

@@ -0,0 +1,5 @@
+from .numeric import Numeric
+from .tag import Tag
+from .text import Text
+
+__all__ = ["Numeric", "Tag", "Text"]

kiarina/lib/redisearch/filter/_field/base.py

@@ -0,0 +1,67 @@
+from typing import Any, Self
+
+from .._enums import RedisearchFilterOperator
+
+
+class RedisearchFieldFilter:
+    """
+    Base class for field filters.
+    """
+
+    OPERATORS: dict[RedisearchFilterOperator, str] = {}
+    """Supported operators"""
+
+    def __init__(self, field_name: str):
+        """
+        Initialization
+        """
+        self._field_name: str = field_name
+        """Field name"""
+
+        self._operator: RedisearchFilterOperator = RedisearchFilterOperator.EQ
+        """Filter operator"""
+
+        self._value: Any = None
+        """Filter value"""
+
+    # --------------------------------------------------
+    # Public Methods
+    # --------------------------------------------------
+
+    def equals(self, other: Self) -> bool:
+        """
+        Check if another filter field is equal to this one.
+        """
+        if not isinstance(other, type(self)):
+            return False
+
+        return self._field_name == other._field_name and self._value == other._value
+
+    # --------------------------------------------------
+    # Protected Methods
+    # --------------------------------------------------
+
+    def _set(
+        self,
+        *,
+        operator: RedisearchFilterOperator,
+        value: Any,
+        value_type: tuple[Any],
+    ) -> None:
+        # Check if the operator is supported by this class
+        if operator not in self.OPERATORS:
+            raise ValueError(
+                f"Operator {operator} not supported by {self.__class__.__name__}. "
+                f"Supported operators are {self.OPERATORS.values()}."
+            )
+
+        if not isinstance(value, value_type):
+            raise TypeError(
+                f"Right side argument passed to operator {self.OPERATORS[operator]} "
+                f"with left side "
+                f"argument {self.__class__.__name__} must be of type {value_type}, "
+                f"received value {value}"
+            )
+
+        self._operator = operator
+        self._value = value