Voice et bot modif
This commit is contained in:
pi
parent
189d56026b
commit
7333a22bcd
10774 changed files with 634644 additions and 933308 deletions
|
|
@ -0,0 +1,80 @@
|
|||
import json
|
||||
from typing import Literal
|
||||
|
||||
from redis._parsers.helpers import pairs_to_dict
|
||||
from redis.commands.vectorset.utils import (
|
||||
parse_vemb_result,
|
||||
parse_vlinks_result,
|
||||
parse_vsim_result,
|
||||
)
|
||||
|
||||
from ..helpers import (
|
||||
apply_module_callbacks,
|
||||
get_legacy_responses,
|
||||
get_protocol_version,
|
||||
)
|
||||
from .commands import (
|
||||
VEMB_CMD,
|
||||
VGETATTR_CMD,
|
||||
VINFO_CMD,
|
||||
VLINKS_CMD,
|
||||
VSIM_CMD,
|
||||
VectorSetCommands,
|
||||
)
|
||||
|
||||
|
||||
class _VectorSetBase(VectorSetCommands):
|
||||
"""Base class with shared initialization logic for VectorSet clients."""
|
||||
|
||||
def __init__(self, client, **kwargs):
|
||||
"""Initialize VectorSet client with callbacks."""
|
||||
# Set the module commands' callbacks
|
||||
_MODULE_CALLBACKS = { # noqa: N806
|
||||
VEMB_CMD: parse_vemb_result,
|
||||
VSIM_CMD: parse_vsim_result,
|
||||
VGETATTR_CMD: lambda r: r and json.loads(r) or None,
|
||||
}
|
||||
|
||||
self._RESP2_MODULE_CALLBACKS = {
|
||||
VINFO_CMD: lambda r: r and pairs_to_dict(r) or None,
|
||||
VLINKS_CMD: parse_vlinks_result,
|
||||
}
|
||||
self._RESP3_MODULE_CALLBACKS = {}
|
||||
self._RESP2_UNIFIED_MODULE_CALLBACKS = dict(self._RESP2_MODULE_CALLBACKS)
|
||||
self._RESP3_UNIFIED_MODULE_CALLBACKS = dict(self._RESP3_MODULE_CALLBACKS)
|
||||
self._RESP3_TO_RESP2_LEGACY_MODULE_CALLBACKS = {}
|
||||
|
||||
self.client = client
|
||||
self.execute_command = client.execute_command
|
||||
|
||||
self._MODULE_CALLBACKS = apply_module_callbacks(
|
||||
get_protocol_version(self.client),
|
||||
get_legacy_responses(self.client),
|
||||
common=_MODULE_CALLBACKS,
|
||||
resp2=self._RESP2_MODULE_CALLBACKS,
|
||||
resp3=self._RESP3_MODULE_CALLBACKS,
|
||||
resp2_unified=self._RESP2_UNIFIED_MODULE_CALLBACKS,
|
||||
resp3_unified=self._RESP3_UNIFIED_MODULE_CALLBACKS,
|
||||
resp3_to_resp2_legacy=self._RESP3_TO_RESP2_LEGACY_MODULE_CALLBACKS,
|
||||
)
|
||||
|
||||
for k, v in self._MODULE_CALLBACKS.items():
|
||||
self.client.set_response_callback(k, v)
|
||||
|
||||
|
||||
class VectorSet(_VectorSetBase):
|
||||
"""Sync VectorSet client."""
|
||||
|
||||
_is_async_client: Literal[False] = False
|
||||
|
||||
|
||||
class AsyncVectorSet(_VectorSetBase):
|
||||
"""Async VectorSet client.
|
||||
|
||||
Note: Inherits from _VectorSetBase (not VectorSet) to maintain proper
|
||||
type discrimination. If AsyncVectorSet inherited from VectorSet, the
|
||||
type system would see it as a subtype of SyncClientProtocol, causing
|
||||
@overload resolution to incorrectly infer sync return types.
|
||||
"""
|
||||
|
||||
_is_async_client: Literal[True] = True
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -0,0 +1,600 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from enum import Enum
|
||||
from typing import Any, Awaitable, overload
|
||||
|
||||
from redis.client import NEVER_DECODE
|
||||
from redis.commands.helpers import get_protocol_version
|
||||
from redis.exceptions import DataError
|
||||
from redis.typing import (
|
||||
AsyncClientProtocol,
|
||||
CommandsProtocol,
|
||||
EncodableT,
|
||||
KeyT,
|
||||
Number,
|
||||
SyncClientProtocol,
|
||||
)
|
||||
from redis.utils import check_protocol_version
|
||||
|
||||
VADD_CMD = "VADD"
|
||||
VSIM_CMD = "VSIM"
|
||||
VREM_CMD = "VREM"
|
||||
VDIM_CMD = "VDIM"
|
||||
VCARD_CMD = "VCARD"
|
||||
VEMB_CMD = "VEMB"
|
||||
VLINKS_CMD = "VLINKS"
|
||||
VINFO_CMD = "VINFO"
|
||||
VSETATTR_CMD = "VSETATTR"
|
||||
VGETATTR_CMD = "VGETATTR"
|
||||
VRANDMEMBER_CMD = "VRANDMEMBER"
|
||||
VRANGE_CMD = "VRANGE"
|
||||
|
||||
# Return type for vsim command
|
||||
VSimResult = (
|
||||
list[list[EncodableT] | dict[EncodableT, Number] | dict[EncodableT, dict[str, Any]]]
|
||||
| None
|
||||
)
|
||||
|
||||
# Return type for vemb command
|
||||
VEmbResult = list[EncodableT] | dict[str, EncodableT] | None
|
||||
|
||||
# Return type for vlinks command
|
||||
VLinksResult = list[list[str | bytes] | dict[str | bytes, Number]] | None
|
||||
|
||||
# Return type for vrandmember command
|
||||
VRandMemberResult = list[str] | str | None
|
||||
|
||||
# Return type for vgetattr command
|
||||
VGetAttrResult = dict | None
|
||||
|
||||
|
||||
class QuantizationOptions(Enum):
|
||||
"""Quantization options for the VADD command."""
|
||||
|
||||
NOQUANT = "NOQUANT"
|
||||
BIN = "BIN"
|
||||
Q8 = "Q8"
|
||||
|
||||
|
||||
class CallbacksOptions(Enum):
|
||||
"""Options that can be set for the commands callbacks"""
|
||||
|
||||
RAW = "RAW"
|
||||
WITHSCORES = "WITHSCORES"
|
||||
WITHATTRIBS = "WITHATTRIBS"
|
||||
ALLOW_DECODING = "ALLOW_DECODING"
|
||||
RESP3 = "RESP3"
|
||||
|
||||
|
||||
class VectorSetCommands(CommandsProtocol):
|
||||
"""Redis VectorSet commands"""
|
||||
|
||||
@overload
|
||||
def vadd(
|
||||
self: SyncClientProtocol,
|
||||
key: KeyT,
|
||||
vector: list[float] | bytes,
|
||||
element: str,
|
||||
reduce_dim: int | None = None,
|
||||
cas: bool | None = False,
|
||||
quantization: QuantizationOptions | None = None,
|
||||
ef: Number | None = None,
|
||||
attributes: dict | str | None = None,
|
||||
numlinks: int | None = None,
|
||||
) -> int: ...
|
||||
|
||||
@overload
|
||||
def vadd(
|
||||
self: AsyncClientProtocol,
|
||||
key: KeyT,
|
||||
vector: list[float] | bytes,
|
||||
element: str,
|
||||
reduce_dim: int | None = None,
|
||||
cas: bool | None = False,
|
||||
quantization: QuantizationOptions | None = None,
|
||||
ef: Number | None = None,
|
||||
attributes: dict | str | None = None,
|
||||
numlinks: int | None = None,
|
||||
) -> Awaitable[int]: ...
|
||||
|
||||
def vadd(
|
||||
self,
|
||||
key: KeyT,
|
||||
vector: list[float] | bytes,
|
||||
element: str,
|
||||
reduce_dim: int | None = None,
|
||||
cas: bool | None = False,
|
||||
quantization: QuantizationOptions | None = None,
|
||||
ef: Number | None = None,
|
||||
attributes: dict | str | None = None,
|
||||
numlinks: int | None = None,
|
||||
) -> Awaitable[int] | int:
|
||||
"""
|
||||
Add vector ``vector`` for element ``element`` to a vector set ``key``.
|
||||
|
||||
``reduce_dim`` sets the dimensions to reduce the vector to.
|
||||
If not provided, the vector is not reduced.
|
||||
|
||||
``cas`` is a boolean flag that indicates whether to use CAS (check-and-set style)
|
||||
when adding the vector. If not provided, CAS is not used.
|
||||
|
||||
``quantization`` sets the quantization type to use.
|
||||
If not provided, int8 quantization is used.
|
||||
The options are:
|
||||
- NOQUANT: No quantization
|
||||
- BIN: Binary quantization
|
||||
- Q8: Signed 8-bit quantization
|
||||
|
||||
``ef`` sets the exploration factor to use.
|
||||
If not provided, the default exploration factor is used.
|
||||
|
||||
``attributes`` is a dictionary or json string that contains the attributes to set for the vector.
|
||||
If not provided, no attributes are set.
|
||||
|
||||
``numlinks`` sets the number of links to create for the vector.
|
||||
If not provided, the default number of links is used.
|
||||
|
||||
For more information, see https://redis.io/commands/vadd.
|
||||
"""
|
||||
if not vector or not element:
|
||||
raise DataError("Both vector and element must be provided")
|
||||
|
||||
pieces = []
|
||||
if reduce_dim:
|
||||
pieces.extend(["REDUCE", reduce_dim])
|
||||
|
||||
values_pieces = []
|
||||
if isinstance(vector, bytes):
|
||||
values_pieces.extend(["FP32", vector])
|
||||
else:
|
||||
values_pieces.extend(["VALUES", len(vector)])
|
||||
values_pieces.extend(vector)
|
||||
pieces.extend(values_pieces)
|
||||
|
||||
pieces.append(element)
|
||||
|
||||
if cas:
|
||||
pieces.append("CAS")
|
||||
|
||||
if quantization:
|
||||
pieces.append(quantization.value)
|
||||
|
||||
if ef:
|
||||
pieces.extend(["EF", ef])
|
||||
|
||||
if attributes:
|
||||
if isinstance(attributes, dict):
|
||||
# transform attributes to json string
|
||||
attributes_json = json.dumps(attributes)
|
||||
else:
|
||||
attributes_json = attributes
|
||||
pieces.extend(["SETATTR", attributes_json])
|
||||
|
||||
if numlinks:
|
||||
pieces.extend(["M", numlinks])
|
||||
|
||||
return self.execute_command(VADD_CMD, key, *pieces)
|
||||
|
||||
@overload
|
||||
def vsim(
|
||||
self: SyncClientProtocol,
|
||||
key: KeyT,
|
||||
input: list[float] | bytes | str,
|
||||
with_scores: bool | None = False,
|
||||
with_attribs: bool | None = False,
|
||||
count: int | None = None,
|
||||
ef: Number | None = None,
|
||||
filter: str | None = None,
|
||||
filter_ef: str | None = None,
|
||||
truth: bool | None = False,
|
||||
no_thread: bool | None = False,
|
||||
epsilon: Number | None = None,
|
||||
) -> VSimResult: ...
|
||||
|
||||
@overload
|
||||
def vsim(
|
||||
self: AsyncClientProtocol,
|
||||
key: KeyT,
|
||||
input: list[float] | bytes | str,
|
||||
with_scores: bool | None = False,
|
||||
with_attribs: bool | None = False,
|
||||
count: int | None = None,
|
||||
ef: Number | None = None,
|
||||
filter: str | None = None,
|
||||
filter_ef: str | None = None,
|
||||
truth: bool | None = False,
|
||||
no_thread: bool | None = False,
|
||||
epsilon: Number | None = None,
|
||||
) -> Awaitable[VSimResult]: ...
|
||||
|
||||
def vsim(
|
||||
self,
|
||||
key: KeyT,
|
||||
input: list[float] | bytes | str,
|
||||
with_scores: bool | None = False,
|
||||
with_attribs: bool | None = False,
|
||||
count: int | None = None,
|
||||
ef: Number | None = None,
|
||||
filter: str | None = None,
|
||||
filter_ef: str | None = None,
|
||||
truth: bool | None = False,
|
||||
no_thread: bool | None = False,
|
||||
epsilon: Number | None = None,
|
||||
) -> Awaitable[VSimResult] | VSimResult:
|
||||
"""
|
||||
Compare a vector or element ``input`` with the other vectors in a vector set ``key``.
|
||||
|
||||
``with_scores`` sets if similarity scores should be returned for each element in the result.
|
||||
|
||||
``with_attribs`` ``with_attribs`` sets if the results should be returned with the
|
||||
attributes of the elements in the result, or None when no attributes are present.
|
||||
|
||||
``count`` sets the number of results to return.
|
||||
|
||||
``ef`` sets the exploration factor.
|
||||
|
||||
``filter`` sets the filter that should be applied for the search.
|
||||
|
||||
``filter_ef`` sets the max filtering effort.
|
||||
|
||||
``truth`` when enabled, forces the command to perform a linear scan.
|
||||
|
||||
``no_thread`` when enabled forces the command to execute the search
|
||||
on the data structure in the main thread.
|
||||
|
||||
``epsilon`` floating point between 0 and 1, if specified will return
|
||||
only elements with distance no further than the specified one.
|
||||
|
||||
For more information, see https://redis.io/commands/vsim.
|
||||
"""
|
||||
|
||||
if not input:
|
||||
raise DataError("'input' should be provided")
|
||||
|
||||
pieces = []
|
||||
options = {}
|
||||
|
||||
if isinstance(input, bytes):
|
||||
pieces.extend(["FP32", input])
|
||||
elif isinstance(input, list):
|
||||
pieces.extend(["VALUES", len(input)])
|
||||
pieces.extend(input)
|
||||
else:
|
||||
pieces.extend(["ELE", input])
|
||||
|
||||
if with_scores or with_attribs:
|
||||
if check_protocol_version(get_protocol_version(self.client), 3):
|
||||
options[CallbacksOptions.RESP3.value] = True
|
||||
|
||||
if with_scores:
|
||||
pieces.append("WITHSCORES")
|
||||
options[CallbacksOptions.WITHSCORES.value] = True
|
||||
|
||||
if with_attribs:
|
||||
pieces.append("WITHATTRIBS")
|
||||
options[CallbacksOptions.WITHATTRIBS.value] = True
|
||||
|
||||
if count:
|
||||
pieces.extend(["COUNT", count])
|
||||
|
||||
if epsilon:
|
||||
pieces.extend(["EPSILON", epsilon])
|
||||
|
||||
if ef:
|
||||
pieces.extend(["EF", ef])
|
||||
|
||||
if filter:
|
||||
pieces.extend(["FILTER", filter])
|
||||
|
||||
if filter_ef:
|
||||
pieces.extend(["FILTER-EF", filter_ef])
|
||||
|
||||
if truth:
|
||||
pieces.append("TRUTH")
|
||||
|
||||
if no_thread:
|
||||
pieces.append("NOTHREAD")
|
||||
|
||||
return self.execute_command(VSIM_CMD, key, *pieces, **options)
|
||||
|
||||
@overload
|
||||
def vdim(self: SyncClientProtocol, key: KeyT) -> int: ...
|
||||
|
||||
@overload
|
||||
def vdim(self: AsyncClientProtocol, key: KeyT) -> Awaitable[int]: ...
|
||||
|
||||
def vdim(self, key: KeyT) -> Awaitable[int] | int:
|
||||
"""
|
||||
Get the dimension of a vector set.
|
||||
|
||||
In the case of vectors that were populated using the `REDUCE`
|
||||
option, for random projection, the vector set will report the size of
|
||||
the projected (reduced) dimension.
|
||||
|
||||
Raises `redis.exceptions.ResponseError` if the vector set doesn't exist.
|
||||
|
||||
For more information, see https://redis.io/commands/vdim.
|
||||
"""
|
||||
return self.execute_command(VDIM_CMD, key)
|
||||
|
||||
@overload
|
||||
def vcard(self: SyncClientProtocol, key: KeyT) -> int: ...
|
||||
|
||||
@overload
|
||||
def vcard(self: AsyncClientProtocol, key: KeyT) -> Awaitable[int]: ...
|
||||
|
||||
def vcard(self, key: KeyT) -> Awaitable[int] | int:
|
||||
"""
|
||||
Get the cardinality(the number of elements) of a vector set with key ``key``.
|
||||
|
||||
Raises `redis.exceptions.ResponseError` if the vector set doesn't exist.
|
||||
|
||||
For more information, see https://redis.io/commands/vcard.
|
||||
"""
|
||||
return self.execute_command(VCARD_CMD, key)
|
||||
|
||||
@overload
|
||||
def vrem(self: SyncClientProtocol, key: KeyT, element: str) -> int: ...
|
||||
|
||||
@overload
|
||||
def vrem(self: AsyncClientProtocol, key: KeyT, element: str) -> Awaitable[int]: ...
|
||||
|
||||
def vrem(self, key: KeyT, element: str) -> Awaitable[int] | int:
|
||||
"""
|
||||
Remove an element from a vector set.
|
||||
|
||||
For more information, see https://redis.io/commands/vrem.
|
||||
"""
|
||||
return self.execute_command(VREM_CMD, key, element)
|
||||
|
||||
@overload
|
||||
def vemb(
|
||||
self: SyncClientProtocol,
|
||||
key: KeyT,
|
||||
element: str,
|
||||
raw: bool | None = False,
|
||||
) -> VEmbResult: ...
|
||||
|
||||
@overload
|
||||
def vemb(
|
||||
self: AsyncClientProtocol,
|
||||
key: KeyT,
|
||||
element: str,
|
||||
raw: bool | None = False,
|
||||
) -> Awaitable[VEmbResult]: ...
|
||||
|
||||
def vemb(
|
||||
self, key: KeyT, element: str, raw: bool | None = False
|
||||
) -> Awaitable[VEmbResult] | VEmbResult:
|
||||
"""
|
||||
Get the approximated vector of an element ``element`` from vector set ``key``.
|
||||
|
||||
``raw`` is a boolean flag that indicates whether to return the
|
||||
internal representation used by the vector.
|
||||
|
||||
|
||||
For more information, see https://redis.io/commands/vemb.
|
||||
"""
|
||||
options = {}
|
||||
pieces = []
|
||||
pieces.extend([key, element])
|
||||
|
||||
if check_protocol_version(get_protocol_version(self.client), 3):
|
||||
options[CallbacksOptions.RESP3.value] = True
|
||||
|
||||
if raw:
|
||||
pieces.append("RAW")
|
||||
|
||||
options[NEVER_DECODE] = True
|
||||
if (
|
||||
hasattr(self.client, "connection_pool")
|
||||
and self.client.connection_pool.connection_kwargs["decode_responses"]
|
||||
) or (
|
||||
hasattr(self.client, "nodes_manager")
|
||||
and self.client.nodes_manager.connection_kwargs["decode_responses"]
|
||||
):
|
||||
# allow decoding in the postprocessing callback
|
||||
# if the user set decode_responses=True
|
||||
# in the connection pool
|
||||
options[CallbacksOptions.ALLOW_DECODING.value] = True
|
||||
|
||||
options[CallbacksOptions.RAW.value] = True
|
||||
|
||||
return self.execute_command(VEMB_CMD, *pieces, **options)
|
||||
|
||||
@overload
|
||||
def vlinks(
|
||||
self: SyncClientProtocol,
|
||||
key: KeyT,
|
||||
element: str,
|
||||
with_scores: bool | None = False,
|
||||
) -> VLinksResult: ...
|
||||
|
||||
@overload
|
||||
def vlinks(
|
||||
self: AsyncClientProtocol,
|
||||
key: KeyT,
|
||||
element: str,
|
||||
with_scores: bool | None = False,
|
||||
) -> Awaitable[VLinksResult]: ...
|
||||
|
||||
def vlinks(
|
||||
self, key: KeyT, element: str, with_scores: bool | None = False
|
||||
) -> Awaitable[VLinksResult] | VLinksResult:
|
||||
"""
|
||||
Returns the neighbors for each level the element ``element`` exists in the vector set ``key``.
|
||||
|
||||
The result is a list of lists, where each list contains the neighbors for one level.
|
||||
If the element does not exist, or if the vector set does not exist, None is returned.
|
||||
|
||||
If the ``WITHSCORES`` option is provided, the result is a list of dicts,
|
||||
where each dict contains the neighbors for one level, with the scores as values.
|
||||
|
||||
For more information, see https://redis.io/commands/vlinks
|
||||
"""
|
||||
options = {}
|
||||
pieces = []
|
||||
pieces.extend([key, element])
|
||||
|
||||
if with_scores:
|
||||
pieces.append("WITHSCORES")
|
||||
options[CallbacksOptions.WITHSCORES.value] = True
|
||||
|
||||
return self.execute_command(VLINKS_CMD, *pieces, **options)
|
||||
|
||||
@overload
|
||||
def vinfo(self: SyncClientProtocol, key: KeyT) -> dict | None: ...
|
||||
|
||||
@overload
|
||||
def vinfo(self: AsyncClientProtocol, key: KeyT) -> Awaitable[dict | None]: ...
|
||||
|
||||
def vinfo(self, key: KeyT) -> (dict | None) | Awaitable[dict | None]:
|
||||
"""
|
||||
Get information about a vector set.
|
||||
|
||||
For more information, see https://redis.io/commands/vinfo.
|
||||
"""
|
||||
return self.execute_command(VINFO_CMD, key)
|
||||
|
||||
@overload
|
||||
def vsetattr(
|
||||
self: SyncClientProtocol,
|
||||
key: KeyT,
|
||||
element: str,
|
||||
attributes: dict | str | None = None,
|
||||
) -> int: ...
|
||||
|
||||
@overload
|
||||
def vsetattr(
|
||||
self: AsyncClientProtocol,
|
||||
key: KeyT,
|
||||
element: str,
|
||||
attributes: dict | str | None = None,
|
||||
) -> Awaitable[int]: ...
|
||||
|
||||
def vsetattr(
|
||||
self, key: KeyT, element: str, attributes: dict | str | None = None
|
||||
) -> Awaitable[int] | int:
|
||||
"""
|
||||
Associate or remove JSON attributes ``attributes`` of element ``element``
|
||||
for vector set ``key``.
|
||||
|
||||
For more information, see https://redis.io/commands/vsetattr
|
||||
"""
|
||||
if attributes is None:
|
||||
attributes_json = "{}"
|
||||
elif isinstance(attributes, dict):
|
||||
# transform attributes to json string
|
||||
attributes_json = json.dumps(attributes)
|
||||
else:
|
||||
attributes_json = attributes
|
||||
|
||||
return self.execute_command(VSETATTR_CMD, key, element, attributes_json)
|
||||
|
||||
@overload
|
||||
def vgetattr(
|
||||
self: SyncClientProtocol, key: KeyT, element: str
|
||||
) -> VGetAttrResult: ...
|
||||
|
||||
@overload
|
||||
def vgetattr(
|
||||
self: AsyncClientProtocol, key: KeyT, element: str
|
||||
) -> Awaitable[VGetAttrResult]: ...
|
||||
|
||||
def vgetattr(
|
||||
self, key: KeyT, element: str
|
||||
) -> Awaitable[VGetAttrResult] | VGetAttrResult:
|
||||
"""
|
||||
Retrieve the JSON attributes of an element ``element `` for vector set ``key``.
|
||||
|
||||
If the element does not exist, or if the vector set does not exist, None is
|
||||
returned.
|
||||
|
||||
For more information, see https://redis.io/commands/vgetattr.
|
||||
"""
|
||||
return self.execute_command(VGETATTR_CMD, key, element)
|
||||
|
||||
@overload
|
||||
def vrandmember(
|
||||
self: SyncClientProtocol, key: KeyT, count: int | None = None
|
||||
) -> VRandMemberResult: ...
|
||||
|
||||
@overload
|
||||
def vrandmember(
|
||||
self: AsyncClientProtocol, key: KeyT, count: int | None = None
|
||||
) -> Awaitable[VRandMemberResult]: ...
|
||||
|
||||
def vrandmember(
|
||||
self, key: KeyT, count: int | None = None
|
||||
) -> Awaitable[VRandMemberResult] | VRandMemberResult:
|
||||
"""
|
||||
Returns random elements from a vector set ``key``.
|
||||
|
||||
``count`` is the number of elements to return.
|
||||
If ``count`` is not provided, a single element is returned as a single string.
|
||||
If ``count`` is positive(smaller than the number of elements
|
||||
in the vector set), the command returns a list with up to ``count``
|
||||
distinct elements from the vector set
|
||||
If ``count`` is negative, the command returns a list with ``count`` random elements,
|
||||
potentially with duplicates.
|
||||
If ``count`` is greater than the number of elements in the vector set,
|
||||
only the entire set is returned as a list.
|
||||
|
||||
If the vector set does not exist, ``None`` is returned.
|
||||
|
||||
For more information, see https://redis.io/commands/vrandmember.
|
||||
"""
|
||||
pieces = []
|
||||
pieces.append(key)
|
||||
if count is not None:
|
||||
pieces.append(count)
|
||||
return self.execute_command(VRANDMEMBER_CMD, *pieces)
|
||||
|
||||
@overload
|
||||
def vrange(
|
||||
self: SyncClientProtocol,
|
||||
key: KeyT,
|
||||
start: str,
|
||||
end: str,
|
||||
count: int | None = None,
|
||||
) -> list[str]: ...
|
||||
|
||||
@overload
|
||||
def vrange(
|
||||
self: AsyncClientProtocol,
|
||||
key: KeyT,
|
||||
start: str,
|
||||
end: str,
|
||||
count: int | None = None,
|
||||
) -> Awaitable[list[str]]: ...
|
||||
|
||||
def vrange(
|
||||
self, key: KeyT, start: str, end: str, count: int | None = None
|
||||
) -> Awaitable[list[str]] | list[str]:
|
||||
"""
|
||||
Return elements in a lexicographical range from a vector set ``key``.
|
||||
|
||||
``start`` is the starting point of the lexicographical range. Can be:
|
||||
- A string prefixed with '[' for inclusive range (e.g., '[Redis')
|
||||
- A string prefixed with '(' for exclusive range (e.g., '(a7')
|
||||
- The special symbol '-' to indicate the minimum element
|
||||
|
||||
``end`` is the ending point of the lexicographical range. Can be:
|
||||
- A string prefixed with '[' for inclusive range
|
||||
- A string prefixed with '(' for exclusive range
|
||||
- The special symbol '+' to indicate the maximum element
|
||||
|
||||
``count`` is the maximum number of elements to return.
|
||||
If ``count`` is not provided or negative, all elements in the range are returned.
|
||||
If ``count`` is positive, at most ``count`` elements are returned.
|
||||
|
||||
Returns an array of elements in lexicographical order within the specified range.
|
||||
Returns an empty array if the key doesn't exist.
|
||||
|
||||
For more information, see https://redis.io/commands/vrange.
|
||||
"""
|
||||
pieces = [key, start, end]
|
||||
if count is not None:
|
||||
pieces.append(count)
|
||||
return self.execute_command(VRANGE_CMD, *pieces)
|
||||
|
|
@ -0,0 +1,130 @@
|
|||
import json
|
||||
|
||||
from redis._parsers.helpers import pairs_to_dict
|
||||
from redis.commands.vectorset.commands import CallbacksOptions
|
||||
|
||||
|
||||
def parse_vemb_result(response, **options):
|
||||
"""
|
||||
Handle VEMB result since the command can returning different result
|
||||
structures depending on input options and on quantization type of the vector set.
|
||||
|
||||
Parsing VEMB result into:
|
||||
- List[Union[bytes, Union[int, float]]]
|
||||
- Dict[str, Union[bytes, str, float]]
|
||||
"""
|
||||
if response is None:
|
||||
return response
|
||||
|
||||
if options.get(CallbacksOptions.RAW.value):
|
||||
result = {}
|
||||
result["quantization"] = (
|
||||
response[0].decode("utf-8")
|
||||
if options.get(CallbacksOptions.ALLOW_DECODING.value)
|
||||
else response[0]
|
||||
)
|
||||
result["raw"] = response[1]
|
||||
result["l2"] = float(response[2])
|
||||
if len(response) > 3:
|
||||
result["range"] = float(response[3])
|
||||
return result
|
||||
else:
|
||||
if options.get(CallbacksOptions.RESP3.value):
|
||||
return response
|
||||
|
||||
result = []
|
||||
for i in range(len(response)):
|
||||
try:
|
||||
result.append(int(response[i]))
|
||||
except ValueError:
|
||||
# if the value is not an integer, it should be a float
|
||||
result.append(float(response[i]))
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def parse_vlinks_result(response, **options):
|
||||
"""
|
||||
Handle VLINKS result since the command can be returning different result
|
||||
structures depending on input options.
|
||||
Parsing VLINKS result into:
|
||||
- List[List[str]]
|
||||
- List[Dict[str, Number]]
|
||||
"""
|
||||
if response is None:
|
||||
return response
|
||||
|
||||
if options.get(CallbacksOptions.WITHSCORES.value):
|
||||
result = []
|
||||
# Redis will return a list of list of strings.
|
||||
# This list have to be transformed to list of dicts
|
||||
for level_item in response:
|
||||
level_data_dict = {}
|
||||
for key, value in pairs_to_dict(level_item).items():
|
||||
value = float(value)
|
||||
level_data_dict[key] = value
|
||||
result.append(level_data_dict)
|
||||
return result
|
||||
else:
|
||||
# return the list of elements for each level
|
||||
# list of lists
|
||||
return response
|
||||
|
||||
|
||||
def parse_vsim_result(response, **options):
|
||||
"""
|
||||
Handle VSIM result since the command can be returning different result
|
||||
structures depending on input options.
|
||||
Parsing VSIM result into:
|
||||
- List[List[str]]
|
||||
- List[Dict[str, Number]] - when with_scores is used (without attributes)
|
||||
- List[Dict[str, Mapping[str, Any]]] - when with_attribs is used (without scores)
|
||||
- List[Dict[str, Union[Number, Mapping[str, Any]]]] - when with_scores and with_attribs are used
|
||||
|
||||
"""
|
||||
if response is None:
|
||||
return response
|
||||
|
||||
withscores = bool(options.get(CallbacksOptions.WITHSCORES.value))
|
||||
withattribs = bool(options.get(CallbacksOptions.WITHATTRIBS.value))
|
||||
|
||||
# Exactly one of withscores or withattribs is True
|
||||
if (withscores and not withattribs) or (not withscores and withattribs):
|
||||
# Redis will return a list of list of pairs.
|
||||
# This list have to be transformed to dict
|
||||
result_dict = {}
|
||||
if options.get(CallbacksOptions.RESP3.value):
|
||||
resp_dict = response
|
||||
else:
|
||||
resp_dict = pairs_to_dict(response)
|
||||
for key, value in resp_dict.items():
|
||||
if withscores:
|
||||
value = float(value)
|
||||
else:
|
||||
value = json.loads(value) if value else None
|
||||
|
||||
result_dict[key] = value
|
||||
return result_dict
|
||||
elif withscores and withattribs:
|
||||
it = iter(response)
|
||||
result_dict = {}
|
||||
if options.get(CallbacksOptions.RESP3.value):
|
||||
for elem, data in response.items():
|
||||
if data[1] is not None:
|
||||
attribs_dict = json.loads(data[1])
|
||||
else:
|
||||
attribs_dict = None
|
||||
result_dict[elem] = {"score": data[0], "attributes": attribs_dict}
|
||||
else:
|
||||
for elem, score, attribs in zip(it, it, it):
|
||||
if attribs is not None:
|
||||
attribs_dict = json.loads(attribs)
|
||||
else:
|
||||
attribs_dict = None
|
||||
|
||||
result_dict[elem] = {"score": float(score), "attributes": attribs_dict}
|
||||
return result_dict
|
||||
else:
|
||||
# return the list of elements for each level
|
||||
# list of lists
|
||||
return response
|
||||
Loading…
Add table
Add a link
Reference in a new issue