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

kiarina/lib/redisearch/__init__.py

@@ -0,0 +1,35 @@
+import logging
+from importlib import import_module
+from importlib.metadata import version
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ._sync.client import RedisearchClient
+    from ._sync.registry import create_redisearch_client
+    from .settings import RedisearchSettings, settings_manager
+
+__version__ = version("kiarina-lib-redisearch")
+
+__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": "._sync.registry",
+        "RedisearchClient": "._sync.client",
+        "RedisearchSettings": ".settings",
+        "settings_manager": ".settings",
+    }
+
+    globals()[name] = getattr(import_module(module_map[name], __name__), name)
+    return globals()[name]

kiarina/lib/redisearch/_async/client.py

@@ -0,0 +1,181 @@
+from typing import Any
+
+from redis.asyncio 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:  # type: ignore[no-untyped-call]
+            # 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_async=redis
+        )
+        """Redisearch context"""
+
+    # --------------------------------------------------
+    # Index operations
+    # --------------------------------------------------
+
+    async def exists_index(self) -> bool:
+        """
+        Check if the index exists.
+        """
+        return await exists_index("async", self.ctx)
+
+    async def create_index(self) -> None:
+        """
+        Create the search index.
+        """
+        await create_index("async", self.ctx)
+
+    async def drop_index(self, *, delete_documents: bool = False) -> bool:
+        """
+        Delete the index.
+        """
+        return await drop_index("async", self.ctx, delete_documents=delete_documents)
+
+    async def reset_index(self) -> None:
+        """
+        Reset the search index.
+        """
+        await reset_index("async", self.ctx)
+
+    async def migrate_index(self) -> None:
+        """
+        Migrate the search index.
+        """
+        await migrate_index("async", self.ctx)
+
+    async def get_info(self) -> InfoResult:
+        """
+        Get index information using FT.INFO command.
+        """
+        return await get_info("async", self.ctx)
+
+    # --------------------------------------------------
+    # Data operations
+    # --------------------------------------------------
+
+    async 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.
+        """
+        await set("async", self.ctx, mapping, id=id)
+
+    async def delete(self, id: str) -> None:
+        """
+        Delete a document from the index.
+        """
+        await delete("async", self.ctx, id)
+
+    async def get(self, id: str) -> Document | None:
+        """
+        Get a document from the index.
+        """
+        return await get("async", self.ctx, id)
+
+    # --------------------------------------------------
+    # Search operations
+    # --------------------------------------------------
+
+    async def count(
+        self,
+        *,
+        filter: RedisearchFilter | RedisearchFilterConditions | None = None,
+    ) -> SearchResult:
+        """
+        Count documents matching the filter.
+        """
+        return await count("async", self.ctx, filter=filter)
+
+    async 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 await find(
+            "async",
+            self.ctx,
+            filter=filter,
+            sort_by=sort_by,
+            sort_desc=sort_desc,
+            offset=offset,
+            limit=limit,
+            return_fields=return_fields,
+        )
+
+    async 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 await search(
+            "async",
+            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/_async/registry.py

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

kiarina/lib/redisearch/_core/context.py

@@ -0,0 +1,69 @@
+from dataclasses import dataclass
+
+from redis import Redis
+from redis.asyncio import Redis as AsyncRedis
+
+from ..schema import RedisearchSchema
+from ..settings import RedisearchSettings
+
+
+@dataclass
+class RedisearchContext:
+    """
+    Redisearch context
+    """
+
+    settings: RedisearchSettings
+
+    _schema: RedisearchSchema | None = None
+
+    _redis: Redis | None = None
+
+    _redis_async: AsyncRedis | None = None
+
+    # --------------------------------------------------
+    # Properties
+    # --------------------------------------------------
+
+    @property
+    def schema(self) -> RedisearchSchema:
+        """
+        Redisearch index schema
+        """
+        if self._schema is None:
+            if not self.settings.index_schema:
+                raise ValueError("Index schema is not set in RedisearchSettings")
+
+            self._schema = RedisearchSchema.from_field_dicts(self.settings.index_schema)
+
+        return self._schema
+
+    @schema.setter
+    def schema(self, value: RedisearchSchema) -> None:
+        self._schema = value
+
+    # --------------------------------------------------
+
+    @property
+    def redis(self) -> Redis:
+        if self._redis is None:
+            raise ValueError("Redis client is not set in RedisearchContext")
+
+        return self._redis
+
+    @redis.setter
+    def redis(self, value: Redis) -> None:
+        self._redis = value
+
+    # --------------------------------------------------
+
+    @property
+    def redis_async(self) -> AsyncRedis:
+        if self._redis_async is None:
+            raise ValueError("Async Redis client is not set in RedisearchContext")
+
+        return self._redis_async
+
+    @redis_async.setter
+    def redis_async(self, value: AsyncRedis) -> None:
+        self._redis_async = value

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

@@ -0,0 +1,55 @@
+from typing import Awaitable, Literal, overload
+
+from redis.commands.search.query import Query
+
+from ...filter import (
+    RedisearchFilter,
+    RedisearchFilterConditions,
+    create_redisearch_filter,
+)
+from ..context import RedisearchContext
+from ..views.search_result import SearchResult
+
+
+@overload
+def count(
+    mode: Literal["sync"],
+    ctx: RedisearchContext,
+    filter: RedisearchFilter | RedisearchFilterConditions | None = None,
+) -> SearchResult: ...
+
+
+@overload
+def count(
+    mode: Literal["async"],
+    ctx: RedisearchContext,
+    filter: RedisearchFilter | RedisearchFilterConditions | None = None,
+) -> Awaitable[SearchResult]: ...
+
+
+def count(
+    mode: Literal["sync", "async"],
+    ctx: RedisearchContext,
+    filter: RedisearchFilter | RedisearchFilterConditions | None = None,
+) -> SearchResult | Awaitable[SearchResult]:
+    """
+    Count documents matching the filter.
+    """
+    if filter is not None:
+        filter = create_redisearch_filter(filter=filter, schema=ctx.schema)
+
+    filter_query = "*" if filter is None else str(filter)
+    query = Query(filter_query).no_content().paging(0, 0)
+
+    def _sync() -> SearchResult:
+        result = ctx.redis.ft(ctx.settings.index_name).search(query)
+        return SearchResult(total=result.total, duration=result.duration)  # pyright: ignore[reportAttributeAccessIssue]
+
+    async def _async() -> SearchResult:
+        result = await ctx.redis_async.ft(ctx.settings.index_name).search(query)
+        return SearchResult(total=result.total, duration=result.duration)  # pyright: ignore[reportAttributeAccessIssue]
+
+    if mode == "sync":
+        return _sync()
+    else:
+        return _async()

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

@@ -0,0 +1,52 @@
+from typing import Awaitable, Literal, overload
+
+from redis.commands.search.index_definition import IndexDefinition, IndexType
+
+from ..context import RedisearchContext
+
+
+@overload
+def create_index(
+    mode: Literal["sync"],
+    ctx: RedisearchContext,
+) -> None: ...
+
+
+@overload
+def create_index(
+    mode: Literal["async"],
+    ctx: RedisearchContext,
+) -> Awaitable[None]: ...
+
+
+def create_index(
+    mode: Literal["sync", "async"],
+    ctx: RedisearchContext,
+) -> None | Awaitable[None]:
+    """
+    Create the index.
+    """
+    fields = ctx.schema.to_fields()
+
+    definition = IndexDefinition(  # type: ignore[no-untyped-call]
+        prefix=[ctx.settings.key_prefix],
+        index_type=IndexType.HASH,
+    )
+
+    def _sync() -> None:
+        ctx.redis.ft(ctx.settings.index_name).create_index(
+            fields=fields,
+            definition=definition,
+        )
+
+    async def _async() -> None:
+        await ctx.redis_async.ft(ctx.settings.index_name).create_index(
+            fields=fields,
+            definition=definition,
+        )
+
+    if mode == "sync":
+        _sync()
+        return None
+    else:
+        return _async()

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

@@ -0,0 +1,43 @@
+from typing import Awaitable, Literal, overload
+
+from ..context import RedisearchContext
+from .get_key import get_key
+
+
+@overload
+def delete(
+    mode: Literal["sync"],
+    ctx: RedisearchContext,
+    id: str,
+) -> None: ...
+
+
+@overload
+def delete(
+    mode: Literal["async"],
+    ctx: RedisearchContext,
+    id: str,
+) -> Awaitable[None]: ...
+
+
+def delete(
+    mode: Literal["sync", "async"],
+    ctx: RedisearchContext,
+    id: str,
+) -> None | Awaitable[None]:
+    """
+    Delete a document from the index.
+    """
+    key = get_key(ctx, id)
+
+    def _sync() -> None:
+        ctx.redis.delete(key)
+
+    async def _async() -> None:
+        await ctx.redis_async.delete(key)
+
+    if mode == "sync":
+        _sync()
+        return None
+    else:
+        return _async()

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

@@ -0,0 +1,59 @@
+from typing import Awaitable, Literal, overload
+
+from ..context import RedisearchContext
+
+
+@overload
+def drop_index(
+    mode: Literal["sync"],
+    ctx: RedisearchContext,
+    *,
+    delete_documents: bool = False,
+) -> bool: ...
+
+
+@overload
+def drop_index(
+    mode: Literal["async"],
+    ctx: RedisearchContext,
+    *,
+    delete_documents: bool = False,
+) -> Awaitable[bool]: ...
+
+
+def drop_index(
+    mode: Literal["sync", "async"],
+    ctx: RedisearchContext,
+    *,
+    delete_documents: bool = False,
+) -> bool | Awaitable[bool]:
+    """
+    Delete the index.
+    """
+
+    def _sync() -> bool:
+        if ctx.settings.protect_index_deletion:
+            return False
+
+        try:
+            ctx.redis.ft(ctx.settings.index_name).dropindex(delete_documents)
+            return True
+        except Exception:
+            return False
+
+    async def _async() -> bool:
+        if ctx.settings.protect_index_deletion:
+            return False
+
+        try:
+            await ctx.redis_async.ft(ctx.settings.index_name).dropindex(
+                delete_documents
+            )
+            return True
+        except Exception:
+            return False
+
+    if mode == "sync":
+        return _sync()
+    else:
+        return _async()

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

@@ -0,0 +1,56 @@
+from typing import Awaitable, Literal, overload
+
+from ..context import RedisearchContext
+from .get_info import get_info
+
+
+@overload
+def exists_index(
+    mode: Literal["sync"],
+    ctx: RedisearchContext,
+) -> bool: ...
+
+
+@overload
+def exists_index(
+    mode: Literal["async"],
+    ctx: RedisearchContext,
+) -> Awaitable[bool]: ...
+
+
+def exists_index(
+    mode: Literal["sync", "async"],
+    ctx: RedisearchContext,
+) -> bool | Awaitable[bool]:
+    """
+    Check if the index exists.
+    """
+
+    def _handle_exception(e: Exception) -> bool:
+        # <= Redis 7
+        if str(e) == "Unknown index name":
+            return False
+        # Redis 8
+        elif "no such index" in str(e):
+            return False
+
+        raise
+
+    def _sync() -> bool:
+        try:
+            get_info(mode="sync", ctx=ctx)
+            return True
+        except Exception as e:
+            return _handle_exception(e)
+
+    async def _async() -> bool:
+        try:
+            await get_info(mode="async", ctx=ctx)
+            return True
+        except Exception as e:
+            return _handle_exception(e)
+
+    if mode == "sync":
+        return _sync()
+    else:
+        return _async()

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

@@ -0,0 +1,105 @@
+from typing import Any, Awaitable, Literal, overload
+
+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 find(
+    mode: Literal["sync"],
+    ctx: RedisearchContext,
+    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: ...
+
+
+@overload
+def find(
+    mode: Literal["async"],
+    ctx: RedisearchContext,
+    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,
+) -> Awaitable[SearchResult]: ...
+
+
+def find(
+    mode: Literal["sync", "async"],
+    ctx: RedisearchContext,
+    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 | Awaitable[SearchResult]:
+    """
+    Find documents matching the filter criteria.
+
+    Args:
+        limit (int | None):
+            Number of documents to retrieve. If None, retrieves all documents
+        return_fields (list[str] | None):
+            Fields to return. If None, returns no content, only IDs.
+    """
+    if filter is not None:
+        filter = create_redisearch_filter(filter=filter, schema=ctx.schema)
+
+    filter_query = "*" if filter is None else str(filter)
+
+    def _build_query(limit: int) -> Query:
+        query = Query(filter_query)
+
+        if return_fields:
+            query = query.return_fields(*return_fields)
+        else:
+            query = query.no_content()
+
+        if sort_by:
+            query = query.sort_by(sort_by, asc=not sort_desc)
+
+        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(mode="sync", ctx=ctx, filter=filter).total)
+        result = ctx.redis.ft(ctx.settings.index_name).search(query)
+        return _parse_search_result(result)
+
+    async def _async() -> SearchResult:
+        query = _build_query(
+            limit or (await count(mode="async", ctx=ctx, filter=filter)).total
+        )
+        result = await ctx.redis_async.ft(ctx.settings.index_name).search(query)
+        return _parse_search_result(result)
+
+    if mode == "sync":
+        return _sync()
+    else:
+        return _async()