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

kiarina/lib/redisearch/filter/_types.py

@@ -0,0 +1,32 @@
+from typing import Any, TypeAlias
+
+RedisearchFilterConditions: TypeAlias = list[list[Any]]
+"""
+RedisearchFilter に変換可能なリスト
+
+Examples:
+    >>> [
+    ...     ["color", "in", ["blue", "red"]],
+    ...     ["price", "<", 1000],
+    ...     ["title", "like", "*hello*"]
+    ... ]
+
+Tag conditions:
+    - ["color", "==", "blue"]
+    - ["color", "!=", "blue"]
+    - ["color", "in", ["blue", "red"]]
+    - ["color", "not in", ["blue", "red"]]
+
+Numeric conditions:
+    - ["price", "==", 1000]
+    - ["price", "!=", 1000]
+    - ["price", ">", 1000]
+    - ["price", "<", 1000]
+    - ["price", ">=", 1000]
+    - ["price", "<=", 1000]
+
+Text conditions:
+    - ["title", "==", "hello"]
+    - ["title", "!=", "hello"]
+    - ["title", "like", "*hello*"]
+"""

kiarina/lib/redisearch/filter/_utils.py

@@ -0,0 +1,18 @@
+import re
+
+ESCAPED_CHARACTERS: str = r"[,.<>{}\[\]\\\"\':;!@#$%^&*()\-+=~\/ ]"
+"""
+Regular expression patterns
+
+for characters requiring escaping in Redisearch queries
+"""
+
+ESCAPED_CHARACTERS_RE: re.Pattern[str] = re.compile(ESCAPED_CHARACTERS)
+"""Compiled regular expression for escaped characters"""
+
+
+def escape_token(value: str) -> str:
+    """
+    Escape special characters in a Redisearch query token.
+    """
+    return ESCAPED_CHARACTERS_RE.sub(r"\\\g<0>", value)

kiarina/lib/redisearch/schema/__init__.py

@@ -0,0 +1,25 @@
+"""
+フィールド名に `payload`, `distance` は使用できません。
+フィールド名に `id` を使用する場合、ドキュメントの ID と同じ値として扱われます。
+"""
+
+from ._field.numeric import NumericFieldSchema
+from ._field.tag import TagFieldSchema
+from ._field.text import TextFieldSchema
+from ._field.vector.flat import FlatVectorFieldSchema
+from ._field.vector.hnsw import HNSWVectorFieldSchema
+from ._model import RedisearchSchema
+from ._types import FieldSchema
+
+__all__ = [
+    # ._field
+    "FlatVectorFieldSchema",
+    "HNSWVectorFieldSchema",
+    "NumericFieldSchema",
+    "TagFieldSchema",
+    "TextFieldSchema",
+    # ._model
+    "RedisearchSchema",
+    # ._types
+    "FieldSchema",
+]

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

@@ -0,0 +1,20 @@
+from pydantic import BaseModel, Field, field_validator
+
+
+class BaseFieldSchema(BaseModel):
+    """
+    Base class of the field schema
+    """
+
+    name: str = Field(...)
+    """Redisearch field name"""
+
+    @field_validator("name")
+    @classmethod
+    def forbid_reserved_names(cls, v: str) -> str:
+        reserved = {"payload", "distance"}
+
+        if v in reserved:
+            raise ValueError(f'"{v}" is a reserved name and cannot be used.')
+
+        return v

kiarina/lib/redisearch/schema/_field/numeric.py

@@ -0,0 +1,33 @@
+from typing import Literal
+
+from redis.commands.search.field import NumericField
+
+from .base import BaseFieldSchema
+
+
+class NumericFieldSchema(BaseFieldSchema):
+    """
+    Schema for numeric fields
+    """
+
+    type: Literal["numeric"] = "numeric"
+
+    no_index: bool = False
+    """Flag to prevent index creation"""
+
+    sortable: bool | None = False
+    """Flag to indicate if the field is sortable"""
+
+    # --------------------------------------------------
+    # Public Methods
+    # --------------------------------------------------
+
+    def to_field(self) -> NumericField:
+        """
+        Convert the field schema to a Redisearch field
+        """
+        return NumericField(
+            self.name,
+            sortable=self.sortable,
+            no_index=self.no_index,
+        )

kiarina/lib/redisearch/schema/_field/tag.py

@@ -0,0 +1,46 @@
+from typing import Literal
+
+from pydantic import Field
+from redis.commands.search.field import TagField
+
+from .base import BaseFieldSchema
+
+
+class TagFieldSchema(BaseFieldSchema):
+    """
+    Schema for tag fields
+    """
+
+    type: Literal["tag"] = "tag"
+
+    separator: str = ","
+    """Tag separator"""
+
+    case_sensitive: bool = False
+    """Flag to indicate if case sensitivity is enabled"""
+
+    no_index: bool = False
+    """Flag to prevent index creation"""
+
+    sortable: bool | None = False
+    """Flag to indicate if the field is sortable"""
+
+    multiple: bool = Field(False, exclude=True)
+    """
+    Flag to indicate if multiple tags are allowed
+
+    This field is not a feature of Redisearch and is only used within this library.
+    Therefore, it does not affect the migration.
+    """
+
+    def to_field(self) -> TagField:
+        """
+        Convert the field schema to a Redisearch field
+        """
+        return TagField(
+            self.name,
+            separator=self.separator,
+            case_sensitive=self.case_sensitive,
+            sortable=self.sortable,
+            no_index=self.no_index,
+        )

kiarina/lib/redisearch/schema/_field/text.py

@@ -0,0 +1,44 @@
+from typing import Literal
+
+from redis.commands.search.field import TextField
+
+from .base import BaseFieldSchema
+
+
+class TextFieldSchema(BaseFieldSchema):
+    """
+    Schema for text fields
+    """
+
+    type: Literal["text"] = "text"
+
+    weight: float = 1
+    """Weight"""
+
+    no_stem: bool = False
+    """Flag to indicate if stemming is disabled"""
+
+    phonetic_matcher: str | None = None
+    """Phonetic matcher"""
+
+    withsuffixtrie: bool = False
+    """Flag to indicate if suffix trie is used"""
+
+    no_index: bool = False
+    """Flag to prevent index creation"""
+
+    sortable: bool | None = False
+    """Flag to indicate if the field is sortable"""
+
+    def to_field(self) -> TextField:
+        """
+        Convert the field schema to a Redisearch field
+        """
+        return TextField(
+            self.name,
+            weight=self.weight,
+            no_stem=self.no_stem,
+            phonetic_matcher=self.phonetic_matcher,  # type: ignore
+            sortable=self.sortable,
+            no_index=self.no_index,
+        )

kiarina/lib/redisearch/schema/_field/vector/base.py

@@ -0,0 +1,61 @@
+from typing import Any, Literal
+
+import numpy as np
+from pydantic import Field
+
+from ..base import BaseFieldSchema
+
+
+class VectorFieldSchema(BaseFieldSchema):
+    """
+    Base class for vector field schemas
+    """
+
+    type: Literal["vector"] = "vector"
+
+    dims: int = Field(...)
+    """Dimensionality"""
+
+    datatype: Literal["FLOAT32", "FLOAT64"] = "FLOAT32"
+    """Data type"""
+
+    distance_metric: Literal["L2", "COSINE", "IP"] = "COSINE"
+    """Distance metric"""
+
+    initial_cap: int | None = None
+    """Initial capacity"""
+
+    # --------------------------------------------------
+    # Properties
+    # --------------------------------------------------
+
+    @property
+    def dtype(self) -> Any:
+        """
+        Get the numpy data type
+        """
+        if self.datatype == "FLOAT32":
+            return np.float32
+        elif self.datatype == "FLOAT64":
+            return np.float64
+        else:
+            raise ValueError(f"Unsupported datatype: {self.datatype}")
+
+    # --------------------------------------------------
+    # Protected Methods
+    # --------------------------------------------------
+
+    def _get_attributes(self) -> dict[str, Any]:
+        """
+        Get attributes for the vector field
+        """
+        attributes = {
+            "TYPE": self.datatype,
+            "DIM": self.dims,
+            "DISTANCE_METRIC": self.distance_metric,
+        }
+
+        if self.initial_cap is not None:
+            attributes["INITIAL_CAP"] = self.initial_cap
+
+        return attributes

kiarina/lib/redisearch/schema/_field/vector/flat.py

@@ -0,0 +1,40 @@
+from typing import Any, Literal
+
+from redis.commands.search.field import VectorField
+
+from .base import VectorFieldSchema
+
+
+class FlatVectorFieldSchema(VectorFieldSchema):
+    """
+    Schema for FLAT vector fields
+    """
+
+    algorithm: Literal["FLAT"] = "FLAT"
+
+    block_size: int | None = None
+
+    # --------------------------------------------------
+    # Public Methods
+    # --------------------------------------------------
+
+    def to_field(self) -> VectorField:
+        """
+        Convert field schema to Redisearch field
+        """
+        return VectorField(self.name, self.algorithm, self._get_attributes())
+
+    # --------------------------------------------------
+    # Protected Methods
+    # --------------------------------------------------
+
+    def _get_attributes(self) -> dict[str, Any]:
+        """
+        Get attributes for the vector field
+        """
+        attributes = super()._get_attributes()
+
+        if self.block_size is not None:
+            attributes["BLOCK_SIZE"] = self.block_size
+
+        return attributes

kiarina/lib/redisearch/schema/_field/vector/hnsw.py

@@ -0,0 +1,53 @@
+from typing import Any, Literal
+
+from pydantic import Field
+from redis.commands.search.field import VectorField
+
+from .base import VectorFieldSchema
+
+
+class HNSWVectorFieldSchema(VectorFieldSchema):
+    """
+    Schema for HNSW vector fields
+    """
+
+    algorithm: Literal["HNSW"] = "HNSW"
+
+    m: int = Field(default=16)
+
+    ef_construction: int = Field(default=200)
+
+    ef_runtime: int = Field(default=10)
+
+    epsilon: float = Field(default=0.01)
+
+    # --------------------------------------------------
+    # Public Methods
+    # --------------------------------------------------
+
+    def to_field(self) -> VectorField:
+        """
+        Convert field schema to Redisearch field
+        """
+        return VectorField(self.name, self.algorithm, self._get_attributes())
+
+    # --------------------------------------------------
+    # Protected Methods
+    # --------------------------------------------------
+
+    def _get_attributes(self) -> dict[str, Any]:
+        """
+        Get attributes for the vector field
+        """
+        attributes = super()._get_attributes()
+
+        attributes.update(
+            {
+                "M": self.m,
+                "EF_CONSTRUCTION": self.ef_construction,
+                "EF_RUNTIME": self.ef_runtime,
+                "EPSILON": self.epsilon,
+            }
+        )
+
+        return attributes

kiarina/lib/redisearch/schema/_model.py

@@ -0,0 +1,98 @@
+from typing import Any, Self
+
+from pydantic import BaseModel
+from pydantic import Field as PydanticField
+from redis.commands.search.field import Field as RedisearchField
+
+from ._field.vector.flat import FlatVectorFieldSchema
+from ._field.vector.hnsw import HNSWVectorFieldSchema
+from ._types import FieldSchema
+
+
+class RedisearchSchema(BaseModel):
+    """
+    Redisearch index schema
+    """
+
+    fields: list[FieldSchema] = PydanticField(default_factory=list)
+    """All fields"""
+
+    # --------------------------------------------------
+    # Properties
+    # --------------------------------------------------
+
+    @property
+    def field_names(self) -> list[str]:
+        """
+        Obtain the names of all fields in the schema
+        """
+        return [field.name for field in self.fields if field.name]
+
+    @property
+    def vector_field(self) -> FlatVectorFieldSchema | HNSWVectorFieldSchema:
+        """
+        Obtain the vector field from the schema
+
+        Even if there are two or more vector fields, the first one is returned.
+        The vector field used for vector search is assumed to exist only once,
+        if at all, within the schema.
+        """
+        for field in self.fields:
+            if field.type == "vector":
+                return field
+
+        raise ValueError("No vector field found")
+
+    # --------------------------------------------------
+    # Magic Methods
+    # --------------------------------------------------
+
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, RedisearchSchema):
+            return False
+
+        # Ignore fields not output by model_dump() when comparing
+        return self.model_dump() == other.model_dump()
+
+    # --------------------------------------------------
+    # Public Methods
+    # --------------------------------------------------
+
+    def get_field(self, name: str) -> FieldSchema | None:
+        """
+        Get a field by its name
+
+        Args:
+            name (str): Field name
+
+        Returns:
+            FieldSchema | None: The field if found, otherwise None
+        """
+        for field in self.fields:
+            if field.name == name:
+                return field
+
+        return None
+
+    def to_fields(self) -> list[RedisearchField]:
+        """
+        Convert the schema model to a list of Redisearch fields
+        """
+        return [field.to_field() for field in self.fields]
+
+    # --------------------------------------------------
+    # Class Methods
+    # --------------------------------------------------
+
+    @classmethod
+    def from_field_dicts(cls, field_dicts: list[dict[str, Any]]) -> Self:
+        """
+        Create a schema model from a list of field dictionaries
+
+        Args:
+            field_dicts (list[dict[str, Any]]): List of field dictionaries
+
+        Returns:
+            RedisearchSchema: The schema model
+        """
+        return cls.model_validate({"fields": field_dicts})

kiarina/lib/redisearch/schema/_types.py

@@ -0,0 +1,16 @@
+from typing import TypeAlias
+
+from ._field.numeric import NumericFieldSchema
+from ._field.tag import TagFieldSchema
+from ._field.text import TextFieldSchema
+from ._field.vector.flat import FlatVectorFieldSchema
+from ._field.vector.hnsw import HNSWVectorFieldSchema
+
+FieldSchema: TypeAlias = (
+    NumericFieldSchema
+    | TagFieldSchema
+    | TextFieldSchema
+    | FlatVectorFieldSchema
+    | HNSWVectorFieldSchema
+)
+"""Type of the field schema"""

kiarina/lib/redisearch/settings.py

@@ -0,0 +1,47 @@
+from typing import Any
+
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic_settings_manager import SettingsManager
+
+
+class RedisearchSettings(BaseSettings):
+    """
+    Redisearch settings
+    """
+
+    model_config = SettingsConfigDict(env_prefix="KIARINA_LIB_REDISEARCH_")
+
+    key_prefix: str = ""
+    """
+    Redis key prefix
+
+    The prefix for keys of documents registered with Redisearch.
+    Specify a string ending with a colon. e.g. "myapp:"
+    """
+
+    index_name: str = "default"
+    """
+    Redisearch index name
+
+    Only alphanumeric characters, underscores, hyphens, and periods.
+    The beginning consists solely of letters.
+    """
+
+    index_schema: list[dict[str, Any]] | None = None
+    """
+    Redisearch index schema
+
+    RedisearchSchema.from_field_dicts can be used to
+    create the schema from a list of field dictionaries.
+    """
+
+    protect_index_deletion: bool = False
+    """
+    Protect index deletion
+
+    When set to True, the delete_index operation is protected,
+    preventing the index from being accidentally deleted.
+    """
+
+
+settings_manager = SettingsManager(RedisearchSettings, multi=True)