import uuid as uuid_lib
from dataclasses import dataclass
from typing import (
Dict,
List,
Literal,
Optional,
Sequence,
Set,
Tuple,
TypeVar,
Union,
cast,
)
from typing_extensions import TypeAlias
from weaviate.collections.classes.config import ConsistencyLevel
from weaviate.collections.classes.filters import _Filters
from weaviate.collections.classes.grpc import (
PROPERTIES,
PROPERTY,
REFERENCE,
REFERENCES,
BM25OperatorOptions,
BM25OperatorOr,
HybridFusion,
HybridVectorType,
Move,
NearVectorInputType,
QueryNested,
Rerank,
TargetVectorJoinType,
_MetadataQuery,
_QueryReference,
_QueryReferenceMultiTarget,
_Sorting,
)
from weaviate.collections.classes.internal import (
_Generative,
_GroupBy,
)
from weaviate.collections.filters import _FilterToGRPC
from weaviate.collections.grpc.shared import _BaseGRPC
from weaviate.proto.v1 import base_search_pb2, search_get_pb2
from weaviate.types import NUMBER, UUID
from weaviate.util import _ServerVersion
from weaviate.validator import _validate_input, _ValidateArgument
# Can be found in the google.protobuf.internal.well_known_types.pyi stub file but is defined explicitly here for clarity.
_PyValue: TypeAlias = Union[
Dict[str, "_PyValue"],
List["_PyValue"],
str,
float,
bool,
None,
List[float],
List[int],
List[str],
List[bool],
List[UUID],
]
[docs]
@dataclass
class _Move:
force: float
concepts: List[str]
objects: List[uuid_lib.UUID]
A = TypeVar("A")
[docs]
class _QueryGRPC(_BaseGRPC):
def __init__(
self,
weaviate_version: _ServerVersion,
name: str,
tenant: Optional[str],
consistency_level: Optional[ConsistencyLevel],
validate_arguments: bool,
uses_125_api: bool,
uses_127_api: bool,
):
super().__init__(weaviate_version, consistency_level, validate_arguments)
self._name: str = name
self._tenant = tenant
self._validate_arguments = validate_arguments
self.__uses_125_api = uses_125_api
self.__uses_127_api = uses_127_api
def __parse_near_options(
self,
certainty: Optional[NUMBER] = None,
distance: Optional[NUMBER] = None,
) -> Tuple[Optional[float], Optional[float]]:
if self._validate_arguments:
_validate_input(
[
_ValidateArgument([float, int, None], "certainty", certainty),
_ValidateArgument([float, int, None], "distance", distance),
]
)
return (
float(certainty) if certainty is not None else None,
float(distance) if distance is not None else None,
)
[docs]
def get(
self,
*,
limit: Optional[int] = None,
offset: Optional[int] = None,
after: Optional[UUID] = None,
filters: Optional[_Filters] = None,
sort: Optional[_Sorting] = None,
return_metadata: Optional[_MetadataQuery] = None,
return_properties: Union[PROPERTIES, bool, None] = None,
return_references: Optional[REFERENCES] = None,
generative: Optional[_Generative] = None,
rerank: Optional[Rerank] = None,
) -> search_get_pb2.SearchRequest:
if self._validate_arguments:
_validate_input(_ValidateArgument([_Sorting, None], "sort", sort))
if sort is not None:
sort_by: Optional[List[search_get_pb2.SortBy]] = [
search_get_pb2.SortBy(ascending=sort.ascending, path=[sort.prop])
for sort in sort.sorts
]
else:
sort_by = None
return self.__create_request(
after=after,
limit=limit,
offset=offset,
filters=filters,
metadata=return_metadata,
return_properties=return_properties,
return_references=return_references,
generative=generative,
rerank=rerank,
sort_by=sort_by,
)
[docs]
def hybrid(
self,
*,
query: Optional[str],
alpha: Optional[float] = None,
vector: Optional[HybridVectorType] = None,
properties: Optional[List[str]] = None,
fusion_type: Optional[HybridFusion] = None,
distance: Optional[NUMBER] = None,
limit: Optional[int] = None,
offset: Optional[int] = None,
bm25_operator: Optional[BM25OperatorOptions] = None,
autocut: Optional[int] = None,
filters: Optional[_Filters] = None,
group_by: Optional[_GroupBy] = None,
return_metadata: Optional[_MetadataQuery] = None,
return_properties: Union[PROPERTIES, bool, None] = None,
return_references: Optional[REFERENCES] = None,
generative: Optional[_Generative] = None,
rerank: Optional[Rerank] = None,
target_vector: Optional[TargetVectorJoinType] = None,
) -> search_get_pb2.SearchRequest:
return self.__create_request(
limit=limit,
offset=offset,
filters=filters,
group_by=group_by,
metadata=return_metadata,
return_properties=return_properties,
return_references=return_references,
generative=generative,
rerank=rerank,
autocut=autocut,
hybrid_search=self._parse_hybrid(
query,
alpha,
vector,
properties,
bm25_operator,
fusion_type,
distance,
target_vector,
),
)
[docs]
def bm25(
self,
*,
query: Optional[str],
properties: Optional[List[str]] = None,
limit: Optional[int] = None,
offset: Optional[int] = None,
operator: Optional[BM25OperatorOptions] = None,
autocut: Optional[int] = None,
filters: Optional[_Filters] = None,
group_by: Optional[_GroupBy] = None,
return_metadata: Optional[_MetadataQuery] = None,
return_properties: Union[PROPERTIES, bool, None] = None,
return_references: Optional[REFERENCES] = None,
generative: Optional[_Generative] = None,
rerank: Optional[Rerank] = None,
) -> search_get_pb2.SearchRequest:
if self._validate_arguments:
_validate_input(
[
_ValidateArgument([None, str], "query", query),
_ValidateArgument([List, None], "properties", properties),
]
)
return self.__create_request(
limit=limit,
offset=offset,
filters=filters,
group_by=group_by,
metadata=return_metadata,
return_properties=return_properties,
return_references=return_references,
generative=generative,
rerank=rerank,
autocut=autocut,
bm25=(
base_search_pb2.BM25(
query=query,
properties=properties if properties is not None else [],
search_operator=base_search_pb2.SearchOperatorOptions(
operator=operator.operator,
minimum_or_tokens_match=operator.minimum_should_match
if isinstance(operator, BM25OperatorOr)
else None,
)
if operator is not None
else None,
)
if query is not None
else None
),
)
[docs]
def near_vector(
self,
*,
near_vector: NearVectorInputType,
certainty: Optional[NUMBER] = None,
distance: Optional[NUMBER] = None,
limit: Optional[int] = None,
offset: Optional[int] = None,
autocut: Optional[int] = None,
filters: Optional[_Filters] = None,
group_by: Optional[_GroupBy] = None,
generative: Optional[_Generative] = None,
rerank: Optional[Rerank] = None,
target_vector: Optional[TargetVectorJoinType] = None,
return_metadata: Optional[_MetadataQuery] = None,
return_properties: Union[PROPERTIES, bool, None] = None,
return_references: Optional[REFERENCES] = None,
) -> search_get_pb2.SearchRequest:
return self.__create_request(
limit=limit,
offset=offset,
filters=filters,
metadata=return_metadata,
return_properties=return_properties,
return_references=return_references,
generative=generative,
rerank=rerank,
autocut=autocut,
group_by=group_by,
near_vector=self._parse_near_vector(
near_vector, certainty, distance, target_vector=target_vector
),
)
[docs]
def near_object(
self,
*,
near_object: UUID,
certainty: Optional[NUMBER] = None,
distance: Optional[NUMBER] = None,
limit: Optional[int] = None,
offset: Optional[int] = None,
autocut: Optional[int] = None,
filters: Optional[_Filters] = None,
group_by: Optional[_GroupBy] = None,
generative: Optional[_Generative] = None,
rerank: Optional[Rerank] = None,
target_vector: Optional[TargetVectorJoinType] = None,
return_metadata: Optional[_MetadataQuery] = None,
return_properties: Union[PROPERTIES, bool, None] = None,
return_references: Optional[REFERENCES] = None,
) -> search_get_pb2.SearchRequest:
return self.__create_request(
limit=limit,
offset=offset,
filters=filters,
metadata=return_metadata,
return_properties=return_properties,
return_references=return_references,
generative=generative,
rerank=rerank,
autocut=autocut,
group_by=group_by,
near_object=self._parse_near_object(near_object, certainty, distance, target_vector),
)
[docs]
def near_text(
self,
*,
near_text: Union[List[str], str],
certainty: Optional[NUMBER] = None,
distance: Optional[NUMBER] = None,
move_to: Optional[Move] = None,
move_away: Optional[Move] = None,
limit: Optional[int] = None,
offset: Optional[int] = None,
autocut: Optional[int] = None,
filters: Optional[_Filters] = None,
group_by: Optional[_GroupBy] = None,
generative: Optional[_Generative] = None,
rerank: Optional[Rerank] = None,
target_vector: Optional[TargetVectorJoinType] = None,
return_metadata: Optional[_MetadataQuery] = None,
return_properties: Union[PROPERTIES, bool, None] = None,
return_references: Optional[REFERENCES] = None,
) -> search_get_pb2.SearchRequest:
return self.__create_request(
limit=limit,
offset=offset,
filters=filters,
metadata=return_metadata,
return_properties=return_properties,
return_references=return_references,
generative=generative,
rerank=rerank,
autocut=autocut,
group_by=group_by,
near_text=self._parse_near_text(
near_text,
certainty,
distance,
move_away=move_away,
move_to=move_to,
target_vector=target_vector,
),
)
def __create_request(
self,
limit: Optional[int] = None,
offset: Optional[int] = None,
after: Optional[UUID] = None,
filters: Optional[_Filters] = None,
metadata: Optional[_MetadataQuery] = None,
return_properties: Union[PROPERTIES, bool, None] = None,
return_references: Optional[REFERENCES] = None,
generative: Optional[_Generative] = None,
rerank: Optional[Rerank] = None,
autocut: Optional[int] = None,
group_by: Optional[_GroupBy] = None,
near_vector: Optional[base_search_pb2.NearVector] = None,
sort_by: Optional[Sequence[search_get_pb2.SortBy]] = None,
hybrid_search: Optional[base_search_pb2.Hybrid] = None,
bm25: Optional[base_search_pb2.BM25] = None,
near_object: Optional[base_search_pb2.NearObject] = None,
near_text: Optional[base_search_pb2.NearTextSearch] = None,
near_audio: Optional[base_search_pb2.NearAudioSearch] = None,
near_depth: Optional[base_search_pb2.NearDepthSearch] = None,
near_image: Optional[base_search_pb2.NearImageSearch] = None,
near_imu: Optional[base_search_pb2.NearIMUSearch] = None,
near_thermal: Optional[base_search_pb2.NearThermalSearch] = None,
near_video: Optional[base_search_pb2.NearVideoSearch] = None,
) -> search_get_pb2.SearchRequest:
if self._validate_arguments:
_validate_input(
[
_ValidateArgument([int, None], "limit", limit),
_ValidateArgument([int, None], "offset", offset),
_ValidateArgument([uuid_lib.UUID, str, None], "after", after),
_ValidateArgument([_Filters, None], "filters", filters),
_ValidateArgument([_MetadataQuery, None], "metadata", metadata),
_ValidateArgument([_Generative, None], "generative", generative),
_ValidateArgument([Rerank, None], "rerank", rerank),
_ValidateArgument([int, None], "autocut", autocut),
_ValidateArgument([_GroupBy, None], "group_by", group_by),
_ValidateArgument(
[str, bool, QueryNested, Sequence, None],
"return_properties",
return_properties,
),
_ValidateArgument(
[_QueryReference, Sequence, None],
"return_references",
return_references,
),
]
)
if isinstance(return_properties, Sequence):
for prop in return_properties:
_validate_input(
_ValidateArgument(
expected=[str, QueryNested],
name="return_properties",
value=prop,
)
)
if isinstance(return_references, Sequence):
for ref in return_references:
_validate_input(
_ValidateArgument(
expected=[_QueryReference],
name="return_references",
value=ref,
)
)
if return_references is not None:
return_references_parsed: Optional[Set[REFERENCE]] = self.__convert_to_set(
return_references
)
else:
return_references_parsed = None
return_properties_parsed = self.__parse_return_properties(return_properties)
return search_get_pb2.SearchRequest(
uses_123_api=True,
uses_125_api=self.__uses_125_api,
uses_127_api=self.__uses_127_api,
collection=self._name,
limit=limit,
offset=offset,
after=str(after) if after is not None else "",
autocut=autocut,
properties=self._translate_properties_from_python_to_grpc(
return_properties_parsed, return_references_parsed
),
metadata=(self._metadata_to_grpc(metadata) if metadata is not None else None),
consistency_level=self._consistency_level,
tenant=self._tenant,
filters=_FilterToGRPC.convert(filters),
generative=(
generative.to_grpc(self._weaviate_version) if generative is not None else None
),
group_by=group_by.to_grpc() if group_by is not None else None,
rerank=(
search_get_pb2.Rerank(property=rerank.prop, query=rerank.query)
if rerank is not None
else None
),
near_vector=near_vector,
sort_by=sort_by,
hybrid_search=hybrid_search,
bm25_search=bm25,
near_object=near_object,
near_text=near_text,
near_audio=near_audio,
near_depth=near_depth,
near_image=near_image,
near_imu=near_imu,
near_thermal=near_thermal,
near_video=near_video,
)
def __resolve_property(self, prop: QueryNested) -> search_get_pb2.ObjectPropertiesRequest:
props = prop.properties if isinstance(prop.properties, list) else [prop.properties]
return search_get_pb2.ObjectPropertiesRequest(
prop_name=prop.name,
primitive_properties=[p for p in props if isinstance(p, str)],
object_properties=[
self.__resolve_property(p) for p in props if isinstance(p, QueryNested)
],
)
def __parse_return_properties(
self, props: Union[PROPERTIES, bool, None]
) -> Optional[Set[PROPERTY]]:
if props is None or props is True:
return None
return self.__convert_to_set([] if props is False else props)
[docs]
def _translate_properties_from_python_to_grpc(
self, properties: Optional[Set[PROPERTY]], references: Optional[Set[REFERENCE]]
) -> Optional[search_get_pb2.PropertiesRequest]:
if properties is None and references is None:
return None
return search_get_pb2.PropertiesRequest(
return_all_nonref_properties=properties is None,
non_ref_properties=(
None
if properties is None
else [prop for prop in properties if isinstance(prop, str)]
),
ref_properties=(
None
if references is None
else [
search_get_pb2.RefPropertiesRequest(
reference_property=ref.link_on,
properties=self._translate_properties_from_python_to_grpc(
self.__parse_return_properties(ref.return_properties),
(
None
if ref.return_references is None
else self.__convert_to_set(ref.return_references)
),
),
metadata=(
self._metadata_to_grpc(ref._return_metadata)
if ref._return_metadata is not None
else None
),
target_collection=(
ref.target_collection
if isinstance(ref, _QueryReferenceMultiTarget)
else None
),
)
for ref in references
]
),
object_properties=(
None
if properties is None
else [
self.__resolve_property(prop)
for prop in properties
if isinstance(prop, QueryNested)
]
),
)
@staticmethod
def __convert_to_set(args: Union[A, Sequence[A]]) -> Set[A]:
if isinstance(args, list):
return set(args)
else:
return {cast(A, args)}