diff --git a/mock_tests/test_aggregate_near_image.py b/mock_tests/test_aggregate_near_image.py new file mode 100644 index 000000000..011d01c3f --- /dev/null +++ b/mock_tests/test_aggregate_near_image.py @@ -0,0 +1,48 @@ +from typing import List + +import grpc +import pytest + +import weaviate +from weaviate.classes.query import Metrics +from weaviate.proto.v1 import aggregate_pb2, weaviate_pb2_grpc + + +class MockAggregateService(weaviate_pb2_grpc.WeaviateServicer): + def __init__(self) -> None: + self.requests: List[aggregate_pb2.AggregateRequest] = [] + + def Aggregate( + self, request: aggregate_pb2.AggregateRequest, context: grpc.ServicerContext + ) -> aggregate_pb2.AggregateReply: + self.requests.append(request) + return aggregate_pb2.AggregateReply( + single_result=aggregate_pb2.AggregateReply.Single(objects_count=3) + ) + + +@pytest.fixture(scope="function") +def aggregate_service(start_grpc_server: grpc.Server) -> MockAggregateService: + service = MockAggregateService() + weaviate_pb2_grpc.add_WeaviateServicer_to_server(service, start_grpc_server) + return service + + +def test_aggregate_near_image_sends_the_image( + weaviate_client: weaviate.WeaviateClient, aggregate_service: MockAggregateService +) -> None: + collection = weaviate_client.collections.use("NearImageAggregate") + + res = collection.aggregate.near_image( + "aGVsbG8=", + certainty=0.7, + total_count=True, + return_metrics=Metrics("rating").integer(count=True), + ) + + assert res.total_count == 3 + assert len(aggregate_service.requests) == 1 + request = aggregate_service.requests[0] + assert request.WhichOneof("search") == "near_image" + assert request.near_image.image == "aGVsbG8=" + assert request.near_image.certainty == pytest.approx(0.7) diff --git a/test/collection/test_aggregate_near_media.py b/test/collection/test_aggregate_near_media.py new file mode 100644 index 000000000..e7f49d537 --- /dev/null +++ b/test/collection/test_aggregate_near_media.py @@ -0,0 +1,41 @@ +import pytest + +from weaviate.collections.grpc.aggregate import _AggregateGRPC +from weaviate.util import _ServerVersion + + +@pytest.mark.parametrize( + "type_,field,media_field", + [ + ("audio", "near_audio", "audio"), + ("depth", "near_depth", "depth"), + ("image", "near_image", "image"), + ("imu", "near_imu", "imu"), + ("thermal", "near_thermal", "thermal"), + ("video", "near_video", "video"), + ], +) +def test_near_media_request(type_: str, field: str, media_field: str) -> None: + grpc = _AggregateGRPC(_ServerVersion(1, 36, 0), "Dummy", "tenant1", None, True) + + request = grpc.near_media( + media="aGVsbG8=", + type_=type_, # type: ignore[arg-type] + certainty=None, + distance=0.3, + target_vector=None, + aggregations=[], + filters=None, + group_by=None, + limit=None, + object_limit=5, + objects_count=True, + ) + + assert request.WhichOneof("search") == field + search = getattr(request, field) + assert getattr(search, media_field) == "aGVsbG8=" + assert search.distance == pytest.approx(0.3) + assert request.object_limit == 5 + assert request.objects_count is True + assert request.tenant == "tenant1" diff --git a/weaviate/collections/grpc/aggregate.py b/weaviate/collections/grpc/aggregate.py index 5f9aaa213..c5c9d475d 100644 --- a/weaviate/collections/grpc/aggregate.py +++ b/weaviate/collections/grpc/aggregate.py @@ -259,6 +259,12 @@ def __create_request( near_object: Optional[base_search_pb2.NearObject] = None, near_text: Optional[base_search_pb2.NearTextSearch] = None, near_vector: Optional[base_search_pb2.NearVector] = None, + near_image: Optional[base_search_pb2.NearImageSearch] = None, + near_audio: Optional[base_search_pb2.NearAudioSearch] = None, + near_video: Optional[base_search_pb2.NearVideoSearch] = None, + near_depth: Optional[base_search_pb2.NearDepthSearch] = None, + near_thermal: Optional[base_search_pb2.NearThermalSearch] = None, + near_imu: Optional[base_search_pb2.NearIMUSearch] = None, object_limit: Optional[int] = None, objects_count: bool = False, ) -> aggregate_pb2.AggregateRequest: @@ -272,6 +278,12 @@ def __create_request( near_object=near_object, near_text=near_text, near_vector=near_vector, + near_image=near_image, + near_audio=near_audio, + near_video=near_video, + near_depth=near_depth, + near_thermal=near_thermal, + near_imu=near_imu, object_limit=object_limit, objects_count=objects_count, tenant=self._tenant,