AI update based on proto changes (#1005)

Co-authored-by: Ethan <[email protected]>
Co-authored-by: Naveed Jooma <[email protected]>

Author: 3 people

src/viam/app/data_client.py

@@ -27,12 +27,14 @@
     CaptureInterval,
     CaptureMetadata,
     ConfigureDatabaseUserRequest,
+    CreateIndexRequest,
     DataRequest,
     DataServiceStub,
     DeleteBinaryDataByFilterRequest,
     DeleteBinaryDataByFilterResponse,
     DeleteBinaryDataByIDsRequest,
     DeleteBinaryDataByIDsResponse,
+    DeleteIndexRequest,
     DeleteTabularDataRequest,
     DeleteTabularDataResponse,
     ExportTabularDataRequest,
@@ -42,6 +44,10 @@
     GetDatabaseConnectionResponse,
     GetLatestTabularDataRequest,
     GetLatestTabularDataResponse,
+    Index,
+    IndexableCollection,
+    ListIndexesRequest,
+    ListIndexesResponse,
     Order,
     RemoveBinaryDataFromDatasetByIDsRequest,
     RemoveBoundingBoxFromImageByIDRequest,
@@ -499,6 +505,7 @@ async def tabular_data_by_mql(
         use_recent_data: Optional[bool] = None,
         tabular_data_source_type: TabularDataSourceType.ValueType = TabularDataSourceType.TABULAR_DATA_SOURCE_TYPE_STANDARD,
         pipeline_id: Optional[str] = None,
+        query_prefix_name: Optional[str] = None,
     ) -> List[Dict[str, Union[ValueTypes, datetime]]]:
         """Obtain unified tabular data and metadata, queried with MQL.
 
@@ -525,6 +532,7 @@ async def tabular_data_by_mql(
                 Defaults to `TABULAR_DATA_SOURCE_TYPE_STANDARD`.
             pipeline_id (str): The ID of the data pipeline to query. Defaults to `None`.
                 Required if `tabular_data_source_type` is `TABULAR_DATA_SOURCE_TYPE_PIPELINE_SINK`.
+            query_prefix_name (str): Optional field that can be used to specify a saved query to run.
 
         Returns:
             List[Dict[str, Union[ValueTypes, datetime]]]: An array of decoded BSON data objects.
@@ -535,7 +543,12 @@ async def tabular_data_by_mql(
         data_source = TabularDataSource(type=tabular_data_source_type, pipeline_id=pipeline_id)
         if use_recent_data:
             data_source.type = TabularDataSourceType.TABULAR_DATA_SOURCE_TYPE_HOT_STORAGE
-        request = TabularDataByMQLRequest(organization_id=organization_id, mql_binary=binary, data_source=data_source)
+        request = TabularDataByMQLRequest(
+            organization_id=organization_id,
+            mql_binary=binary,
+            data_source=data_source,
+            query_prefix_name=query_prefix_name,
+        )
         response: TabularDataByMQLResponse = await self._data_client.TabularDataByMQL(request, metadata=self._metadata)
         return [bson.decode(bson_bytes) for bson_bytes in response.raw_data]
 
@@ -1131,8 +1144,8 @@ async def add_bounding_box_to_image_by_id(
 
         Args:
             binary_id (Union[~viam.proto.app.data.BinaryID, str]): The binary data ID or :class:`BinaryID` of the image to add the bounding
-                box to. *DEPRECATED:* :class:`BinaryID` *is deprecated and will be removed in a future release. Instead, pass binary data
-                IDs as a list of strings.*
+                box to. *DEPRECATED:* :class:`BinaryID` *is deprecated and will be removed in a future release. Instead, pass binary data IDs as a
+                list of strings.*
             label (str): A label for the bounding box.
             x_min_normalized (float): Min X value of the bounding box normalized from 0 to 1.
             y_min_normalized (float): Min Y value of the bounding box normalized from 0 to 1.
@@ -2061,6 +2074,86 @@ async def _list_data_pipeline_runs(self, id: str, page_size: int, page_token: st
         response: ListDataPipelineRunsResponse = await self._data_pipelines_client.ListDataPipelineRuns(request, metadata=self._metadata)
         return DataClient.DataPipelineRunsPage.from_proto(response, self, page_size)
 
+    async def create_index(
+        self,
+        organization_id: str,
+        collection_type: IndexableCollection.ValueType,
+        index_spec: Dict[str, Any],
+        pipeline_name: Optional[str] = None,
+    ) -> None:
+        """Starts a custom index build.
+
+        Args:
+            organization_id (str): The ID of the organization that owns the data.
+                To find your organization ID, visit the organization settings page.
+            collection_type (IndexableCollection.ValueType): The type of collection the index is on.
+            index_spec (List[Dict[str, Any]]): The MongoDB index specification defined in JSON format.
+            pipeline_name (Optional[str]): The name of the pipeline if the collection type is PIPELINE_SINK.
+
+        For more information, see `Data Client API <https://docs.viam.com/dev/reference/apis/data-client/#createindex>`_.
+        """
+        index_spec_bytes = [bson.encode(index_spec)]
+        request = CreateIndexRequest(
+            organization_id=organization_id,
+            collection_type=collection_type,
+            index_spec=index_spec_bytes,
+            pipeline_name=pipeline_name,
+        )
+        await self._data_client.CreateIndex(request, metadata=self._metadata)
+
+    async def list_indexes(
+        self,
+        organization_id: str,
+        collection_type: IndexableCollection.ValueType,
+        pipeline_name: Optional[str] = None,
+    ) -> Sequence[Index]:
+        """Returns all the indexes for a given collection.
+
+        Args:
+            organization_id (str): The ID of the organization that owns the data.
+                To find your organization ID, visit the organization settings page.
+            collection_type (IndexableCollection.ValueType): The type of collection the index is on.
+            pipeline_name (Optional[str]): The name of the pipeline if the collection type is PIPELINE_SINK.
+
+        Returns:
+            List[Index]: A list of indexes.
+
+        For more information, see `Data Client API <https://docs.viam.com/dev/reference/apis/data-client/#listindexes>`_.
+        """
+        request = ListIndexesRequest(
+            organization_id=organization_id,
+            collection_type=collection_type,
+            pipeline_name=pipeline_name,
+        )
+        response: ListIndexesResponse = await self._data_client.ListIndexes(request, metadata=self._metadata)
+        return response.indexes
+
+    async def delete_index(
+        self,
+        organization_id: str,
+        collection_type: IndexableCollection.ValueType,
+        index_name: str,
+        pipeline_name: Optional[str] = None,
+    ) -> None:
+        """Drops the specified custom index from a collection.
+
+        Args:
+            organization_id (str): The ID of the organization that owns the data.
+                To find your organization ID, visit the organization settings page.
+            collection_type (IndexableCollection.ValueType): The type of collection the index is on.
+            index_name (str): The name of the index to delete.
+            pipeline_name (Optional[str]): The name of the pipeline if the collection type is PIPELINE_SINK.
+
+        For more information, see `Data Client API <https://docs.viam.com/dev/reference/apis/data-client/#deleteindex>`_.
+        """
+        request = DeleteIndexRequest(
+            organization_id=organization_id,
+            collection_type=collection_type,
+            index_name=index_name,
+            pipeline_name=pipeline_name,
+        )
+        await self._data_client.DeleteIndex(request, metadata=self._metadata)
+
     @staticmethod
     def create_filter(
         component_name: Optional[str] = None,

src/viam/services/worldstatestore/worldstatestore.py

@@ -85,6 +85,6 @@ async def stream_transform_changes(
                 print(f"Transform {change.transform.uuid} {change.change_type}")
 
         Returns:
-            AsyncIterator[StreamTransformChangesResponse]: A stream of transform changes
+            AsyncGenerator[StreamTransformChangesResponse, None]: A stream of transform changes
         """
         ...

tests/mocks/services.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Any, AsyncIterator, Dict, List, Mapping, Optional, Sequence, Union
+from typing import Any, AsyncGenerator, Dict, List, Mapping, Optional, Sequence, Union
 
 import bson
 import numpy as np
@@ -213,10 +213,14 @@
     BoundingBoxLabelsByFilterResponse,
     ConfigureDatabaseUserRequest,
     ConfigureDatabaseUserResponse,
+    CreateIndexRequest,
+    CreateIndexResponse,
     DeleteBinaryDataByFilterRequest,
     DeleteBinaryDataByFilterResponse,
     DeleteBinaryDataByIDsRequest,
     DeleteBinaryDataByIDsResponse,
+    DeleteIndexRequest,
+    DeleteIndexResponse,
     DeleteTabularDataRequest,
     DeleteTabularDataResponse,
     ExportTabularDataRequest,
@@ -225,6 +229,9 @@
     GetDatabaseConnectionResponse,
     GetLatestTabularDataRequest,
     GetLatestTabularDataResponse,
+    Index,
+    ListIndexesRequest,
+    ListIndexesResponse,
     RemoveBinaryDataFromDatasetByIDsRequest,
     RemoveBinaryDataFromDatasetByIDsResponse,
     RemoveBoundingBoxFromImageByIDRequest,
@@ -379,7 +386,7 @@
 class MockVision(Vision):
     def __init__(
         self,
-        name: str,
+        name,
         detectors: List[str],
         detections: List[Detection],
         classifiers: List[str],
@@ -845,6 +852,7 @@ def __init__(
         bbox_labels_response: List[str],
         hostname_response: str,
         additional_params: Mapping[str, ValueTypes],
+        list_indexes_response: List[Index],
     ):
         self.tabular_response = tabular_response
         self.tabular_export_response = tabular_export_response
@@ -857,6 +865,7 @@ def __init__(
         self.was_tabular_data_requested = False
         self.was_binary_data_requested = False
         self.additional_params = additional_params
+        self.list_indexes_response = list_indexes_response
 
     async def TabularDataByFilter(self, stream: Stream[TabularDataByFilterRequest, TabularDataByFilterResponse]) -> None:
         request = await stream.recv_message()
@@ -1043,6 +1052,7 @@ async def TabularDataBySQL(self, stream: Stream[TabularDataBySQLRequest, Tabular
     async def TabularDataByMQL(self, stream: Stream[TabularDataByMQLRequest, TabularDataByMQLResponse]) -> None:
         request = await stream.recv_message()
         assert request is not None
+        self.query_prefix_name = request.query_prefix_name if request.HasField("query_prefix_name") else None
         await stream.send_message(TabularDataByMQLResponse(raw_data=[bson.encode(dict) for dict in self.tabular_query_response]))
 
     async def GetLatestTabularData(self, stream: Stream[GetLatestTabularDataRequest, GetLatestTabularDataResponse]) -> None:
@@ -1068,6 +1078,24 @@ async def ExportTabularData(self, stream: Stream[ExportTabularDataRequest, Expor
         for tabular_data in self.tabular_export_response:
             await stream.send_message(tabular_data)
 
+    async def CreateIndex(self, stream: Stream[CreateIndexRequest, CreateIndexResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        self.create_index_request = request
+        await stream.send_message(CreateIndexResponse())
+
+    async def ListIndexes(self, stream: Stream[ListIndexesRequest, ListIndexesResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        self.list_indexes_request = request
+        await stream.send_message(ListIndexesResponse(indexes=self.list_indexes_response))
+
+    async def DeleteIndex(self, stream: Stream[DeleteIndexRequest, DeleteIndexResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        self.delete_index_request = request
+        await stream.send_message(DeleteIndexResponse())
+
 
 class MockDataset(DatasetServiceBase):
     def __init__(self, create_response: str, datasets_response: Sequence[Dataset], merged_response: Optional[str] = None):
@@ -1931,7 +1959,7 @@ async def stream_transform_changes(
         *,
         extra: Optional[Mapping[str, Any]] = None,
         timeout: Optional[float] = None,
-    ) -> AsyncIterator["StreamTransformChangesResponse"]:
+    ) -> AsyncGenerator[StreamTransformChangesResponse, None]:
         self.extra = extra
         self.timeout = timeout
 

tests/test_data_client.py

@@ -1,6 +1,7 @@
 from datetime import datetime
 from typing import List
 
+import bson
 import pytest
 from google.protobuf.timestamp_pb2 import Timestamp
 from grpclib.testing import ChannelFor
@@ -16,6 +17,9 @@
     CaptureMetadata,
     ExportTabularDataResponse,
     Filter,
+    Index,
+    IndexableCollection,
+    IndexCreator,
     Order,
 )
 from viam.utils import create_filter, dict_to_struct, struct_to_dict
@@ -142,6 +146,23 @@
 TAGS_RESPONSE = ["tag"]
 HOSTNAME_RESPONSE = "host"
 
+INDEX_NAME = "my_index"
+INDEX_SPEC = {"key": 1}
+INDEX_SPEC_BYTES = [bson.encode(INDEX_SPEC)]
+COLLECTION_TYPE = IndexableCollection.INDEXABLE_COLLECTION_PIPELINE_SINK
+PIPELINE_NAME = "my_pipeline"
+INDEX_CREATED_BY = IndexCreator.INDEX_CREATOR_CUSTOMER
+LIST_INDEXES_RESPONSE = [
+    Index(
+        collection_type=COLLECTION_TYPE,
+        pipeline_name=PIPELINE_NAME,
+        index_name=INDEX_NAME,
+        index_spec=INDEX_SPEC_BYTES,
+        created_by=INDEX_CREATED_BY,
+    )
+]
+QUERY_PREFIX_NAME = "my_saved_query"
+
 AUTH_TOKEN = "auth_token"
 DATA_SERVICE_METADATA = {"authorization": f"Bearer {AUTH_TOKEN}"}
 
@@ -158,6 +179,7 @@ def service() -> MockData:
         bbox_labels_response=BBOX_LABELS,
         hostname_response=HOSTNAME_RESPONSE,
         additional_params=ADDITIONAL_PARAMS,
+        list_indexes_response=LIST_INDEXES_RESPONSE,
     )
 
 
@@ -195,9 +217,10 @@ async def test_tabular_data_by_sql(self, service: MockData):
     async def test_tabular_data_by_mql(self, service: MockData):
         async with ChannelFor([service]) as channel:
             client = DataClient(channel, DATA_SERVICE_METADATA)
-            response = await client.tabular_data_by_mql(ORG_ID, MQL_BINARY)
+            response = await client.tabular_data_by_mql(ORG_ID, MQL_BINARY, query_prefix_name=QUERY_PREFIX_NAME)
             assert isinstance(response[0]["key1"], datetime)
             assert response == TABULAR_QUERY_RESPONSE
+            assert service.query_prefix_name == QUERY_PREFIX_NAME
             response = await client.tabular_data_by_mql(ORG_ID, mql_binary=[b"mql_binary"])
             assert isinstance(response[0]["key1"], datetime)
             assert response == TABULAR_QUERY_RESPONSE
@@ -429,6 +452,55 @@ async def test_remove_binary_data_to_dataset_by_ids(self, service: MockData):
             assert service.removed_binary_data_ids == BINARY_DATA_IDS
             assert service.dataset_id == DATASET_ID
 
+    async def test_create_index(self, service: MockData):
+        async with ChannelFor([service]) as channel:
+            client = DataClient(channel, DATA_SERVICE_METADATA)
+            await client.create_index(
+                organization_id=ORG_ID,
+                collection_type=COLLECTION_TYPE,
+                index_spec=INDEX_SPEC,
+                pipeline_name=PIPELINE_NAME,
+            )
+            assert service.create_index_request.organization_id == ORG_ID
+            assert service.create_index_request.collection_type == COLLECTION_TYPE
+            assert service.create_index_request.index_spec == INDEX_SPEC_BYTES
+            assert service.create_index_request.pipeline_name == PIPELINE_NAME
+
+    async def test_list_indexes(self, service: MockData):
+        async with ChannelFor([service]) as channel:
+            client = DataClient(channel, DATA_SERVICE_METADATA)
+            indexes = await client.list_indexes(
+                organization_id=ORG_ID,
+                collection_type=COLLECTION_TYPE,
+                pipeline_name=PIPELINE_NAME,
+            )
+            assert service.list_indexes_request.organization_id == ORG_ID
+            assert service.list_indexes_request.collection_type == COLLECTION_TYPE
+            assert service.list_indexes_request.pipeline_name == PIPELINE_NAME
+
+            assert len(indexes) == len(LIST_INDEXES_RESPONSE)
+            for i, index in enumerate(indexes):
+                expected_index = LIST_INDEXES_RESPONSE[i]
+                assert index.collection_type == expected_index.collection_type
+                assert index.pipeline_name == expected_index.pipeline_name
+                assert index.index_name == expected_index.index_name
+                assert index.index_spec == expected_index.index_spec
+                assert index.created_by == expected_index.created_by
+
+    async def test_delete_index(self, service: MockData):
+        async with ChannelFor([service]) as channel:
+            client = DataClient(channel, DATA_SERVICE_METADATA)
+            await client.delete_index(
+                organization_id=ORG_ID,
+                collection_type=COLLECTION_TYPE,
+                index_name=INDEX_NAME,
+                pipeline_name=PIPELINE_NAME,
+            )
+            assert service.delete_index_request.organization_id == ORG_ID
+            assert service.delete_index_request.collection_type == COLLECTION_TYPE
+            assert service.delete_index_request.index_name == INDEX_NAME
+            assert service.delete_index_request.pipeline_name == PIPELINE_NAME
+
     def assert_filter(self, filter: Filter) -> None:
         assert filter.component_name == COMPONENT_NAME
         assert filter.component_type == COMPONENT_TYPE