diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md index 9d3f862fd5..2dc891b830 100644 --- a/python/CHANGELOG.md +++ b/python/CHANGELOG.md @@ -5,6 +5,8 @@ This project adheres to [Semantic Versioning](http://semver.org/). ## [Unreleased] +- Added `client.annotations` for data reviews, phases, and review history. + ### What's New #### MCAP imports diff --git a/python/lib/sift_client/_internal/low_level_wrappers/annotations.py b/python/lib/sift_client/_internal/low_level_wrappers/annotations.py new file mode 100644 index 0000000000..bbd4da3804 --- /dev/null +++ b/python/lib/sift_client/_internal/low_level_wrappers/annotations.py @@ -0,0 +1,352 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, cast + +from sift.annotation_logs.v1.annotation_logs_pb2 import ( + AnnotationLogAssignedProperties, + AnnotationLogCommentProperties, + AnnotationLogStateUpdateProperties, + CreateAnnotationLogRequest, + CreateAnnotationLogResponse, + ListAnnotationLogsRequest, + ListAnnotationLogsResponse, +) +from sift.annotation_logs.v1.annotation_logs_pb2_grpc import AnnotationLogServiceStub +from sift.annotations.v1.annotations_pb2 import ( + ArchiveAnnotationRequest, + ArchiveAnnotationResponse, + BatchArchiveAnnotationsRequest, + BatchArchiveAnnotationsResponse, + BatchUnarchiveAnnotationsRequest, + BatchUnarchiveAnnotationsResponse, + CreateAnnotationResponse, + GetAnnotationRequest, + GetAnnotationResponse, + ListAnnotationsRequest, + ListAnnotationsResponse, + UnarchiveAnnotationRequest, + UnarchiveAnnotationResponse, + UpdateAnnotationRequest, + UpdateAnnotationResponse, +) +from sift.annotations.v1.annotations_pb2_grpc import AnnotationServiceStub + +from sift_client._internal.low_level_wrappers.base import DEFAULT_PAGE_SIZE, LowLevelClientBase +from sift_client.sift_types.annotation import ( + Annotation, + AnnotationCommentElement, + AnnotationCreateBase, + AnnotationLog, + AnnotationLogKind, + AnnotationLogState, + AnnotationUpdate, +) +from sift_client.transport import WithGrpcClient + +if TYPE_CHECKING: + from sift_client.transport.grpc_transport import GrpcClient + +# Configure logging +logger = logging.getLogger(__name__) + + +class AnnotationsLowLevelClient(LowLevelClientBase, WithGrpcClient): + """Low-level client for the AnnotationsAPI. + + A thin wrapper around the autogenerated bindings for the AnnotationsAPI and the + AnnotationLogsAPI. + """ + + def __init__(self, grpc_client: GrpcClient): + """Initialize the AnnotationsLowLevelClient. + + Args: + grpc_client: The gRPC client to use for making API calls. + """ + super().__init__(grpc_client) + + async def get_annotation(self, annotation_id: str) -> Annotation: + """Get an annotation by annotation_id. + + Args: + annotation_id: The annotation ID to get. + + Returns: + The Annotation. + """ + request = GetAnnotationRequest(annotation_id=annotation_id) + response = await self._grpc_client.get_stub(AnnotationServiceStub).GetAnnotation(request) + return Annotation._from_proto(cast("GetAnnotationResponse", response).annotation) + + async def list_annotations( + self, + *, + page_size: int | None = DEFAULT_PAGE_SIZE, + page_token: str | None = None, + query_filter: str | None = None, + order_by: str | None = None, + ) -> tuple[list[Annotation], str]: + """List one page of annotations. + + Args: + page_size: The maximum number of annotations to return. + page_token: A page token for pagination. + query_filter: A CEL filter string. + order_by: How to order the retrieved annotations. + + Returns: + A tuple of (annotations, next_page_token). + """ + request_kwargs: dict[str, Any] = {} + if page_size is not None: + request_kwargs["page_size"] = page_size + if page_token is not None: + request_kwargs["page_token"] = page_token + if query_filter is not None: + request_kwargs["filter"] = query_filter + if order_by is not None: + request_kwargs["order_by"] = order_by + + request = ListAnnotationsRequest(**request_kwargs) + response = await self._grpc_client.get_stub(AnnotationServiceStub).ListAnnotations(request) + response = cast("ListAnnotationsResponse", response) + + annotations = [Annotation._from_proto(a) for a in response.annotations] + return annotations, response.next_page_token + + async def list_all_annotations( + self, + *, + query_filter: str | None = None, + order_by: str | None = None, + page_size: int | None = DEFAULT_PAGE_SIZE, + max_results: int | None = None, + ) -> list[Annotation]: + """List every matching annotation. + + Args: + query_filter: A CEL filter string. + order_by: How to order the retrieved annotations. + page_size: The number of results to fetch per request. + max_results: Maximum number of results to return. + + Returns: + A list of all matching annotations. + """ + return await self._handle_pagination( + self.list_annotations, + kwargs={"query_filter": query_filter}, + order_by=order_by, + max_results=max_results, + page_size=page_size, + ) + + async def create_annotation(self, *, create: AnnotationCreateBase) -> Annotation: + """Create an annotation. + + Args: + create: The annotation definition. + + Returns: + The created Annotation. + """ + request_proto = create.to_proto() + request_proto.linked_channels.extend(create.linked_channels_to_proto()) + response = await self._grpc_client.get_stub(AnnotationServiceStub).CreateAnnotation( + request_proto + ) + return Annotation._from_proto(cast("CreateAnnotationResponse", response).annotation) + + async def update_annotation(self, update: AnnotationUpdate) -> Annotation: + """Update select fields of an annotation. + + Args: + update: The updates to apply. Its `resource_id` must be set. + + Returns: + The updated Annotation. + """ + grpc_annotation, update_mask = update.to_proto_with_mask() + if "linked_channels" in update.model_fields_set: + grpc_annotation.linked_channels.extend(update.linked_channels_to_proto()) + update_mask.paths.append("linked_channels") + request = UpdateAnnotationRequest(annotation=grpc_annotation, update_mask=update_mask) + response = await self._grpc_client.get_stub(AnnotationServiceStub).UpdateAnnotation(request) + return Annotation._from_proto(cast("UpdateAnnotationResponse", response).annotation) + + async def archive_annotation(self, annotation_id: str) -> Annotation: + """Archive an annotation. + + Args: + annotation_id: The ID of the annotation to archive. + + Returns: + The archived Annotation. + """ + request = ArchiveAnnotationRequest(annotation_id=annotation_id) + response = await self._grpc_client.get_stub(AnnotationServiceStub).ArchiveAnnotation( + request + ) + return Annotation._from_proto(cast("ArchiveAnnotationResponse", response).annotation) + + async def unarchive_annotation(self, annotation_id: str) -> Annotation: + """Unarchive an annotation. + + Args: + annotation_id: The ID of the annotation to unarchive. + + Returns: + The unarchived Annotation. + """ + request = UnarchiveAnnotationRequest(annotation_id=annotation_id) + response = await self._grpc_client.get_stub(AnnotationServiceStub).UnarchiveAnnotation( + request + ) + return Annotation._from_proto(cast("UnarchiveAnnotationResponse", response).annotation) + + async def batch_archive_annotations(self, annotation_ids: list[str]) -> list[Annotation]: + """Archive many annotations in one call. + + Args: + annotation_ids: The IDs of the annotations to archive. + + Returns: + The archived Annotations. + """ + request = BatchArchiveAnnotationsRequest(annotation_ids=annotation_ids) + response = await self._grpc_client.get_stub(AnnotationServiceStub).BatchArchiveAnnotations( + request + ) + response = cast("BatchArchiveAnnotationsResponse", response) + return [Annotation._from_proto(a) for a in response.annotations] + + async def batch_unarchive_annotations(self, annotation_ids: list[str]) -> list[Annotation]: + """Unarchive many annotations in one call. + + Args: + annotation_ids: The IDs of the annotations to unarchive. + + Returns: + The unarchived Annotations. + """ + request = BatchUnarchiveAnnotationsRequest(annotation_ids=annotation_ids) + response = await self._grpc_client.get_stub( + AnnotationServiceStub + ).BatchUnarchiveAnnotations(request) + response = cast("BatchUnarchiveAnnotationsResponse", response) + return [Annotation._from_proto(a) for a in response.annotations] + + async def list_annotation_logs( + self, + *, + annotation_id: str | None = None, + page_size: int | None = DEFAULT_PAGE_SIZE, + page_token: str | None = None, + query_filter: str | None = None, + order_by: str | None = None, + ) -> tuple[list[AnnotationLog], str]: + """List one page of annotation logs. + + Args: + annotation_id: Restrict results to this annotation. + page_size: The maximum number of logs to return. + page_token: A page token for pagination. + query_filter: A CEL filter string. + order_by: How to order the retrieved logs. + + Returns: + A tuple of (logs, next_page_token). + """ + request_kwargs: dict[str, Any] = {} + if annotation_id is not None: + request_kwargs["annotation_id"] = annotation_id + if page_size is not None: + request_kwargs["page_size"] = page_size + if page_token is not None: + request_kwargs["page_token"] = page_token + if query_filter is not None: + request_kwargs["filter"] = query_filter + if order_by is not None: + request_kwargs["order_by"] = order_by + + request = ListAnnotationLogsRequest(**request_kwargs) + response = await self._grpc_client.get_stub(AnnotationLogServiceStub).ListAnnotationLogs( + request + ) + response = cast("ListAnnotationLogsResponse", response) + + logs = [AnnotationLog._from_proto(log) for log in response.annotation_logs] + return logs, response.next_page_token + + async def list_all_annotation_logs( + self, + *, + annotation_id: str | None = None, + query_filter: str | None = None, + order_by: str | None = None, + page_size: int | None = DEFAULT_PAGE_SIZE, + max_results: int | None = None, + ) -> list[AnnotationLog]: + """List every matching annotation log. + + Args: + annotation_id: Restrict results to this annotation. + query_filter: A CEL filter string. + order_by: How to order the retrieved logs. + page_size: The number of results to fetch per request. + max_results: Maximum number of results to return. + + Returns: + A list of all matching logs. + """ + return await self._handle_pagination( + self.list_annotation_logs, + kwargs={"annotation_id": annotation_id, "query_filter": query_filter}, + order_by=order_by, + max_results=max_results, + page_size=page_size, + ) + + async def create_annotation_log( + self, + *, + annotation_id: str, + kind: AnnotationLogKind, + assigned_to_user_id: str | None = None, + state: AnnotationLogState | None = None, + comment: list[AnnotationCommentElement] | None = None, + ) -> AnnotationLog: + """Create an annotation log. + + Args: + annotation_id: The annotation to attach the log to. + kind: Which kind of log to create. + assigned_to_user_id: The assignee, for ASSIGNED logs. + state: The new state, for STATE_UPDATE logs. + comment: The comment body, for COMMENT logs. + + Returns: + The created AnnotationLog. + """ + request_kwargs: dict[str, Any] = {"annotation_id": annotation_id, "kind": kind.value} + if kind is AnnotationLogKind.ASSIGNED: + request_kwargs["assigned"] = AnnotationLogAssignedProperties( + assigned_to_user_id=assigned_to_user_id or "" + ) + elif kind is AnnotationLogKind.STATE_UPDATE: + request_kwargs["state_update"] = AnnotationLogStateUpdateProperties( + state=state.value if state else 0 # type: ignore[arg-type] + ) + elif kind is AnnotationLogKind.COMMENT: + request_kwargs["comment"] = AnnotationLogCommentProperties( + body=[e._to_proto() for e in comment or []] + ) + + request = CreateAnnotationLogRequest(**request_kwargs) + response = await self._grpc_client.get_stub(AnnotationLogServiceStub).CreateAnnotationLog( + request + ) + return AnnotationLog._from_proto( + cast("CreateAnnotationLogResponse", response).annotation_log + ) diff --git a/python/lib/sift_client/_tests/conftest.py b/python/lib/sift_client/_tests/conftest.py index 23f2dd7a62..28d3ee11d2 100644 --- a/python/lib/sift_client/_tests/conftest.py +++ b/python/lib/sift_client/_tests/conftest.py @@ -79,6 +79,7 @@ def mock_client(): """Create a mock SiftClient for unit testing.""" client = MagicMock(spec=SiftClient) # Configure the mock to have the necessary API attributes + client.annotations = MagicMock() client.assets = MagicMock() client.reports = MagicMock() client.runs = MagicMock() diff --git a/python/lib/sift_client/_tests/resources/test_annotations.py b/python/lib/sift_client/_tests/resources/test_annotations.py new file mode 100644 index 0000000000..19ccb1cfb7 --- /dev/null +++ b/python/lib/sift_client/_tests/resources/test_annotations.py @@ -0,0 +1,377 @@ +"""Pytest tests for the Annotations API. + +These tests demonstrate and validate the usage of the Annotations API including: +- Basic annotation operations (get, list, find) +- Annotation filtering and searching +- Annotation creation, updates, and archiving +- The review workflow: assignment, state, and annotation logs +""" + +from datetime import datetime, timedelta, timezone + +import pytest + +from sift_client import SiftClient +from sift_client.resources import AnnotationLogsAPI, AnnotationsAPI, AnnotationsAPIAsync +from sift_client.sift_types import Annotation +from sift_client.sift_types.annotation import ( + AnnotationCreate, + AnnotationLog, + AnnotationLogKind, + AnnotationState, + AnnotationType, + AnnotationUpdate, + PhaseCreate, +) + +pytestmark = pytest.mark.integration + + +def test_client_binding(sift_client): + assert isinstance(sift_client.annotations, AnnotationsAPI) + assert isinstance(sift_client.annotations.logs, AnnotationLogsAPI) + assert isinstance(sift_client.async_.annotations, AnnotationsAPIAsync) + + +@pytest.fixture(scope="session") +def test_timestamp_str(): + """A per-session suffix so annotation names stay unique across runs.""" + return datetime.now(timezone.utc).isoformat() + + +@pytest.fixture(scope="session") +def window(): + """A short, fixed time range for the annotations these tests create.""" + end = datetime.now(timezone.utc) + return end - timedelta(minutes=5), end + + +@pytest.fixture +def annotations_api_async(sift_client: SiftClient): + """Get the async annotations API instance.""" + return sift_client.async_.annotations + + +@pytest.fixture(scope="session") +def new_annotation(sift_client, test_timestamp_str, window, nostromo_asset): + """Create an annotation for the session and archive it on teardown.""" + start, end = window + created = sift_client.annotations.create( + AnnotationCreate( + name=f"test_annotation_{test_timestamp_str}", + description="Created by Sift Client pytest", + start_time=start, + end_time=end, + assets=[nostromo_asset], + state=AnnotationState.OPEN, + ) + ) + yield created + sift_client.annotations.archive(created) + + +class TestAnnotations: + """Tests for the Annotations API.""" + + def test_create(self, new_annotation, test_timestamp_str, nostromo_asset): + """Test that create returns a fully populated annotation.""" + assert isinstance(new_annotation, Annotation) + assert new_annotation.id_ is not None + assert new_annotation.name == f"test_annotation_{test_timestamp_str}" + assert new_annotation.annotation_type is AnnotationType.DATA_REVIEW + assert new_annotation.state is AnnotationState.OPEN + assert new_annotation.is_archived is False + assert nostromo_asset.id_ in new_annotation.asset_ids + + def test_create_phase(self, sift_client, test_timestamp_str, window, nostromo_asset): + """Test creating a phase annotation, which carries no review state.""" + start, end = window + phase = sift_client.annotations.create( + PhaseCreate( + name=f"test_phase_{test_timestamp_str}", + start_time=start, + end_time=end, + assets=[nostromo_asset], + ) + ) + + assert phase.annotation_type is AnnotationType.PHASE + sift_client.annotations.archive(phase) + + def test_get(self, sift_client, new_annotation): + """Test getting an annotation by ID.""" + fetched = sift_client.annotations.get(annotation_id=new_annotation._id_or_error) + + assert fetched.id_ == new_annotation.id_ + assert fetched.name == new_annotation.name + + def test_basic_list(self, sift_client, new_annotation): + """Test basic annotation listing functionality.""" + annotations = sift_client.annotations.list_(limit=5) + + assert isinstance(annotations, list) + for annotation in annotations: + assert isinstance(annotation, Annotation) + assert annotation.id_ is not None + + def test_list_with_name_filter(self, sift_client, new_annotation): + """Test annotation listing with name filtering.""" + by_name = sift_client.annotations.list_(name=new_annotation.name) + + assert len(by_name) == 1 + assert by_name[0].id_ == new_annotation.id_ + + def test_list_with_id_filter(self, sift_client, new_annotation): + """Test annotation listing filtered to specific IDs.""" + annotations = sift_client.annotations.list_(annotation_ids=[new_annotation._id_or_error]) + + assert [a.id_ for a in annotations] == [new_annotation.id_] + + def test_list_with_type_and_state_filters(self, sift_client, new_annotation): + """Test annotation listing filtered by type and review state.""" + annotations = sift_client.annotations.list_( + annotation_ids=[new_annotation._id_or_error], + annotation_type=AnnotationType.DATA_REVIEW, + state=AnnotationState.OPEN, + ) + + assert [a.id_ for a in annotations] == [new_annotation.id_] + + def test_list_with_asset_filter(self, sift_client, new_annotation, nostromo_asset): + """Test annotation listing filtered by asset.""" + annotations = sift_client.annotations.list_(assets=[nostromo_asset], limit=20) + + assert new_annotation.id_ in {a.id_ for a in annotations} + + def test_find(self, sift_client, new_annotation): + """Test finding a single annotation.""" + found = sift_client.annotations.find(name=new_annotation.name) + + assert found is not None + assert found.id_ == new_annotation.id_ + + def test_find_nonexistent(self, sift_client): + """Test finding a non-existent annotation returns None.""" + found = sift_client.annotations.find( + name=f"nonexistent_annotation_{datetime.now(timezone.utc).timestamp()}" + ) + assert found is None + + def test_update(self, sift_client, new_annotation): + """Test updating an annotation's description and state.""" + updated = sift_client.annotations.update( + new_annotation, + AnnotationUpdate(description="updated", state=AnnotationState.FAILED), + ) + + assert updated.description == "updated" + assert updated.state is AnnotationState.FAILED + # The name was not in the mask, so it is unchanged. + assert updated.name == new_annotation.name + + sift_client.annotations.update(new_annotation, {"state": AnnotationState.OPEN}) + + def test_update_accepts_dict(self, sift_client, new_annotation, test_timestamp_str): + """Test that update accepts a plain dict.""" + renamed = f"test_annotation_renamed_{test_timestamp_str}" + updated = sift_client.annotations.update(new_annotation._id_or_error, {"name": renamed}) + + assert updated.name == renamed + sift_client.annotations.update(new_annotation, {"name": new_annotation.name}) + + def test_archive_and_unarchive(self, sift_client, test_timestamp_str, window, nostromo_asset): + """Test archiving and unarchiving an annotation.""" + start, end = window + annotation = sift_client.annotations.create( + AnnotationCreate( + name=f"test_annotation_archive_{test_timestamp_str}", + start_time=start, + end_time=end, + assets=[nostromo_asset], + ) + ) + + archived = sift_client.annotations.archive(annotation) + assert archived.is_archived is True + + # Archived annotations are excluded from list_ by default. + assert sift_client.annotations.find(name=annotation.name) is None + assert ( + sift_client.annotations.find(name=annotation.name, include_archived=True).id_ + == annotation.id_ + ) + + unarchived = sift_client.annotations.unarchive(annotation) + assert unarchived.is_archived is False + + sift_client.annotations.archive(annotation) + + def test_instance_methods(self, sift_client, test_timestamp_str, window, nostromo_asset): + """Test the update, resolve, and archive methods on the instance itself.""" + start, end = window + annotation = sift_client.annotations.create( + AnnotationCreate( + name=f"test_annotation_instance_{test_timestamp_str}", + start_time=start, + end_time=end, + assets=[nostromo_asset], + state=AnnotationState.OPEN, + ) + ) + + annotation.update({"description": "from-instance"}) + assert annotation.description == "from-instance" + + annotation.set_accepted() + assert annotation.state is AnnotationState.ACCEPTED + + annotation.archive() + assert annotation.is_archived is True + + +class TestAnnotationLogs: + """Tests for the nested annotation logs API.""" + + def test_comment(self, sift_client, new_annotation): + """Test adding a comment and reading it back.""" + log = sift_client.annotations.logs.add_comment(new_annotation, "looks fine to me") + + assert isinstance(log, AnnotationLog) + assert log.kind is AnnotationLogKind.COMMENT + assert log.text == "looks fine to me" + + logs = sift_client.annotations.logs.list_(annotation=new_annotation) + assert log.id_ in {entry.id_ for entry in logs} + + def test_list_filtered_by_kind(self, sift_client, new_annotation): + """Test filtering the history by log kind.""" + sift_client.annotations.logs.add_comment(new_annotation, "another comment") + + logs = sift_client.annotations.logs.list_( + annotation=new_annotation, kind=AnnotationLogKind.COMMENT + ) + + assert logs + assert all(entry.kind is AnnotationLogKind.COMMENT for entry in logs) + + def test_annotation_logs_property(self, sift_client, new_annotation): + """Test the logs property on the Annotation instance.""" + sift_client.annotations.logs.add_comment(new_annotation, "via property") + + assert any(entry.kind is AnnotationLogKind.COMMENT for entry in new_annotation.logs) + + +class TestAnnotationWorkflow: + """Tests for the task-shaped API: named creation verbs and review verbs.""" + + @pytest.fixture(scope="session") + def nostromo_channels(self, sift_client, nostromo_asset): + """The channels to draw test annotations on.""" + return sift_client.channels.list_(asset=nostromo_asset) + + def test_create_derives_the_asset_from_channels( + self, sift_client, test_timestamp_str, window, nostromo_asset, nostromo_channels + ): + """Passing Channel objects links them and fills in the asset.""" + start, end = window + review = sift_client.annotations.create( + AnnotationCreate( + name=f"test_review_{test_timestamp_str}", + start_time=start, + end_time=end, + linked_channels=nostromo_channels, + description="needs a look", + tags=["sift-client-pytest"], + ) + ) + + assert review.annotation_type is AnnotationType.DATA_REVIEW + # The asset came from the channels, not from an `assets` argument. + assert review.asset_ids == [nostromo_asset.id_] + assert {c.id_ for c in review.linked_channels} == {c.id_ for c in nostromo_channels} + assert review.tags == ["sift-client-pytest"] + + sift_client.annotations.archive(review) + + def test_create_phase_model( + self, sift_client, test_timestamp_str, window, nostromo_asset, nostromo_channels + ): + """PhaseCreate makes a phase with no review state.""" + start, end = window + phase = sift_client.annotations.create( + PhaseCreate( + name=f"test_phase_verb_{test_timestamp_str}", + start_time=start, + end_time=end, + linked_channels=nostromo_channels, + ) + ) + + assert phase.annotation_type is AnnotationType.PHASE + assert phase.state is None + + sift_client.annotations.archive(phase) + + def test_phase_model_rejects_a_state(self): + """PhaseCreate has no `state` field and forbids extras, so it fails validation.""" + assert "state" not in PhaseCreate.model_fields + with pytest.raises(ValueError, match="state"): + PhaseCreate.model_validate( + { + "name": "a", + "start_time": datetime.now(timezone.utc), + "end_time": datetime.now(timezone.utc), + "state": AnnotationState.OPEN, + } + ) + + def test_models_pin_their_annotation_type(self): + """Neither model accepts the other's type, so a phase with a state can't be built.""" + now = datetime.now(timezone.utc) + with pytest.raises(ValueError, match="annotation_type"): + AnnotationCreate( + name="a", start_time=now, end_time=now, annotation_type=AnnotationType.PHASE + ) + with pytest.raises(ValueError, match="annotation_type"): + PhaseCreate( + name="a", start_time=now, end_time=now, annotation_type=AnnotationType.DATA_REVIEW + ) + + def test_create_with_explicit_assets( + self, sift_client, test_timestamp_str, window, nostromo_asset + ): + """`assets` accepts Asset objects and skips the channel lookup.""" + start, end = window + review = sift_client.annotations.create( + AnnotationCreate( + name=f"test_review_assets_{test_timestamp_str}", + start_time=start, + end_time=end, + assets=[nostromo_asset], + ) + ) + + assert review.asset_ids == [nostromo_asset.id_] + assert review.linked_channels == [] + + sift_client.annotations.archive(review) + + def test_state_verbs(self, sift_client, new_annotation): + """Each review verb changes the state.""" + assert ( + sift_client.annotations.set_accepted(new_annotation).state is AnnotationState.ACCEPTED + ) + assert sift_client.annotations.set_failed(new_annotation).state is AnnotationState.FAILED + assert sift_client.annotations.set_open(new_annotation).state is AnnotationState.OPEN + + def test_instance_state_verbs(self, sift_client, new_annotation): + """Test the review verbs on the Annotation instance.""" + assert new_annotation.set_accepted().state is AnnotationState.ACCEPTED + assert new_annotation.set_failed().state is AnnotationState.FAILED + assert new_annotation.set_open().state is AnnotationState.OPEN + + def test_instance_comment(self, sift_client, new_annotation): + """Test commenting from the Annotation instance.""" + log = new_annotation.add_comment("from the instance") + + assert log.kind is AnnotationLogKind.COMMENT diff --git a/python/lib/sift_client/_tests/sift_types/test_annotation.py b/python/lib/sift_client/_tests/sift_types/test_annotation.py new file mode 100644 index 0000000000..6ac3a751cb --- /dev/null +++ b/python/lib/sift_client/_tests/sift_types/test_annotation.py @@ -0,0 +1,251 @@ +"""Tests for sift_types.Annotation model.""" + +from datetime import datetime, timezone +from unittest.mock import MagicMock + +import pytest +from sift.channels.v3.channels_pb2 import Channel as ChannelProto +from sift.common.type.v1.channel_data_type_pb2 import ChannelDataType as ChannelDataTypeProto + +from sift_client.sift_types import Annotation +from sift_client.sift_types.annotation import ( + AnnotationCommentElement, + AnnotationCreate, + AnnotationState, + AnnotationType, + AnnotationUpdate, + PhaseCreate, +) +from sift_client.sift_types.channel import Channel + +START = datetime(2026, 1, 1, tzinfo=timezone.utc) +END = datetime(2026, 1, 2, tzinfo=timezone.utc) + + +class TestAnnotationEnums: + """Unit tests for the CEL filter representations.""" + + def test_type_filter_str(self): + assert AnnotationType.PHASE.to_filter_str() == "ANNOTATION_TYPE_PHASE" + + def test_state_filter_str(self): + assert AnnotationState.ACCEPTED.to_filter_str() == "ANNOTATION_STATE_RESOLVED" + + +class TestAnnotationCreate: + """Unit tests for AnnotationCreate - tests _to_proto_helpers and validators.""" + + def test_minimal_create(self): + proto = AnnotationCreate(name="a", start_time=START, end_time=END).to_proto() + + assert proto.name == "a" + assert proto.annotation_type == AnnotationType.DATA_REVIEW.value + + def test_run_reaches_the_proto(self): + """Both create models inherit the `run` helper, which maps it to `run_id`.""" + run_id = "d961e543-0dde-4069-b046-8ae124f48d9b" + + for model in (AnnotationCreate, PhaseCreate): + proto = model(name="a", start_time=START, end_time=END, run=run_id).to_proto() + + assert proto.run_id == run_id + + def test_assets_and_tags_are_names(self): + create = AnnotationCreate( + name="a", start_time=START, end_time=END, assets=["Nostromo"], tags=["review"] + ) + proto = create.to_proto() + + assert list(proto.assets) == ["Nostromo"] + assert list(proto.tags) == ["review"] + + def test_metadata_converter(self): + proto = AnnotationCreate( + name="a", start_time=START, end_time=END, metadata={"k": "v", "n": 1.5} + ).to_proto() + + by_key = {m.key.name: m for m in proto.metadata} + assert by_key["k"].string_value == "v" + assert by_key["n"].number_value == 1.5 + + def test_linked_channels_takes_channel_objects(self): + """linked_channels takes Channel objects, wrapping them for the proto.""" + channel = Channel._from_proto( + ChannelProto( + channel_id="ch-1", + name="pressure", + data_type=ChannelDataTypeProto.CHANNEL_DATA_TYPE_DOUBLE, + ) + ) + + create = AnnotationCreate( + name="a", start_time=START, end_time=END, linked_channels=[channel] + ) + linked = create.linked_channels_to_proto() + + assert len(linked) == 1 + assert linked[0].channel.channel_id == "ch-1" + + def test_phase_has_no_state(self): + """PhaseCreate carries no review state, so this fails before the call.""" + with pytest.raises(ValueError, match="state"): + PhaseCreate.model_validate( + { + "name": "a", + "start_time": START, + "end_time": END, + "state": AnnotationState.OPEN, + } + ) + + def test_rejects_inverted_time_range(self): + with pytest.raises(ValueError, match="start_time must not be after end_time"): + AnnotationCreate(name="a", start_time=END, end_time=START) + + +class TestAnnotationUpdate: + """Unit tests for AnnotationUpdate - tests field masks.""" + + def test_update_mask_only_includes_set_fields(self): + update = AnnotationUpdate(name="renamed", state=AnnotationState.FAILED) + update.resource_id = "an-1" + + proto, mask = update.to_proto_with_mask() + + assert proto.annotation_id == "an-1" + assert proto.name == "renamed" + assert proto.state == AnnotationState.FAILED.value + assert set(mask.paths) == {"name", "state"} + + def test_assignment(self): + update = AnnotationUpdate(assigned_to_user_id="user-1") + update.resource_id = "an-1" + + proto, mask = update.to_proto_with_mask() + + assert proto.assigned_to_user_id == "user-1" + assert mask.paths == ["assigned_to_user_id"] + + def test_requires_resource_id(self): + with pytest.raises(ValueError, match="Resource ID must be set"): + AnnotationUpdate(name="renamed").to_proto_with_mask() + + +class TestAnnotationCommentElement: + """Unit tests for comment body elements.""" + + def test_text_element(self): + proto = AnnotationCommentElement(text="looks fine")._to_proto() + assert proto.text == "looks fine" + assert AnnotationCommentElement._from_proto(proto).text == "looks fine" + + def test_mention_element(self): + proto = AnnotationCommentElement(user_id="u-1")._to_proto() + assert proto.user_mention.user_id == "u-1" + assert AnnotationCommentElement._from_proto(proto).user_id == "u-1" + + def test_rejects_both(self): + with pytest.raises(ValueError, match="exactly one"): + AnnotationCommentElement(text="hi", user_id="u-1") + + +@pytest.fixture +def mock_annotation(mock_client): + """Create a mock Annotation instance for testing.""" + annotation = Annotation( + proto=MagicMock(), + id_="test_annotation_id", + name="test_annotation", + description="test", + start_time=START, + end_time=END, + annotation_type=AnnotationType.DATA_REVIEW, + organization_id="org1", + created_date=START, + modified_date=START, + created_by_user_id="user1", + modified_by_user_id="user1", + tags=[], + asset_ids=["asset1"], + linked_channel_ids=[], + linked_calculated_channel_version_ids=[], + metadata={}, + is_archived=False, + pending=False, + state=AnnotationState.OPEN, + run_id=None, + assigned_to_user_id=None, + archived_date=None, + ) + annotation._apply_client_to_instance(mock_client) + return annotation + + +class TestAnnotation: + """Unit tests for Annotation model - tests properties and methods.""" + + def test_logs_property_calls_client(self, mock_annotation, mock_client): + mock_client.annotations.logs.list_.return_value = [] + + _ = mock_annotation.logs + + mock_client.annotations.logs.list_.assert_called_once_with(annotation="test_annotation_id") + + def test_update_calls_client_and_updates_self(self, mock_annotation, mock_client): + updated = MagicMock() + mock_client.annotations.update.return_value = updated + + with MagicMock() as mock_update: + mock_annotation._update = mock_update + + update = AnnotationUpdate(name="renamed") + result = mock_annotation.update(update) + + mock_client.annotations.update.assert_called_once_with( + annotation=mock_annotation, update=update + ) + mock_update.assert_called_once_with(updated) + assert result is mock_annotation + + def test_archive_calls_client(self, mock_annotation, mock_client): + mock_client.annotations.archive.return_value = MagicMock() + + with MagicMock() as mock_update: + mock_annotation._update = mock_update + result = mock_annotation.archive() + + mock_client.annotations.archive.assert_called_once_with(annotation=mock_annotation) + assert result is mock_annotation + + def test_unarchive_calls_client(self, mock_annotation, mock_client): + mock_client.annotations.unarchive.return_value = MagicMock() + + with MagicMock() as mock_update: + mock_annotation._update = mock_update + result = mock_annotation.unarchive() + + mock_client.annotations.unarchive.assert_called_once_with(annotation=mock_annotation) + assert result is mock_annotation + + def test_assign_forwards_to_the_resource(self, mock_annotation, mock_client): + mock_client.annotations.assign_to_user.return_value = MagicMock() + + with MagicMock() as mock_update: + mock_annotation._update = mock_update + mock_annotation.assign_to_user("user-2") + + mock_client.annotations.assign_to_user.assert_called_once_with( + annotation=mock_annotation, user="user-2" + ) + + def test_state_verbs_forward_to_the_resource(self, mock_annotation, mock_client): + for verb in ("set_open", "set_failed", "set_accepted"): + getattr(mock_client.annotations, verb).return_value = MagicMock() + + with MagicMock() as mock_update: + mock_annotation._update = mock_update + getattr(mock_annotation, verb)() + + getattr(mock_client.annotations, verb).assert_called_once_with( + annotation=mock_annotation + ) diff --git a/python/lib/sift_client/client.py b/python/lib/sift_client/client.py index 5c1c49551c..dc1a46e8ba 100644 --- a/python/lib/sift_client/client.py +++ b/python/lib/sift_client/client.py @@ -9,6 +9,8 @@ from sift_client.credentials import ResolvedCredentials, resolve_credentials from sift_client.errors import SiftWarning from sift_client.resources import ( + AnnotationsAPI, + AnnotationsAPIAsync, AssetsAPI, AssetsAPIAsync, CalculatedChannelsAPI, @@ -101,6 +103,9 @@ class SiftClient( ping: PingAPI """Instance of the Ping API for making synchronous requests.""" + annotations: AnnotationsAPI + """Instance of the Annotations API for making synchronous requests.""" + assets: AssetsAPI """Instance of the Assets API for making synchronous requests.""" @@ -234,6 +239,7 @@ def __init__( self.cache = CacheNamespace(self) self.ping = PingAPI(self) + self.annotations = AnnotationsAPI(self) self.assets = AssetsAPI(self) self.calculated_channels = CalculatedChannelsAPI(self) self.channels = ChannelsAPI(self) @@ -255,6 +261,7 @@ def __init__( # Accessor for the asynchronous APIs self.async_ = AsyncAPIs( ping=PingAPIAsync(self), + annotations=AnnotationsAPIAsync(self), assets=AssetsAPIAsync(self), calculated_channels=CalculatedChannelsAPIAsync(self), channels=ChannelsAPIAsync(self), diff --git a/python/lib/sift_client/resources/__init__.py b/python/lib/sift_client/resources/__init__.py index efa97fba3d..99861cd669 100644 --- a/python/lib/sift_client/resources/__init__.py +++ b/python/lib/sift_client/resources/__init__.py @@ -150,6 +150,7 @@ async def main(): ``` """ +from sift_client.resources.annotations import AnnotationLogsAPIAsync, AnnotationsAPIAsync from sift_client.resources.assets import AssetsAPIAsync from sift_client.resources.calculated_channels import CalculatedChannelsAPIAsync from sift_client.resources.channels import ChannelsAPIAsync @@ -186,6 +187,8 @@ async def main(): # ruff: noqa All imports needs to be imported before sync_stubs to avoid circular import from sift_client.resources.sync_stubs import ( + AnnotationLogsAPI, + AnnotationsAPI, AssetsAPI, CalculatedChannelsAPI, ChannelsAPI, @@ -220,6 +223,10 @@ async def main(): TestResultsAPIAsync.__test__ = False # type: ignore __all__ = [ + "AnnotationLogsAPI", + "AnnotationLogsAPIAsync", + "AnnotationsAPI", + "AnnotationsAPIAsync", "AssetsAPI", "AssetsAPIAsync", "AutoRegisterStreamingClient", diff --git a/python/lib/sift_client/resources/annotations.py b/python/lib/sift_client/resources/annotations.py new file mode 100644 index 0000000000..18740293dc --- /dev/null +++ b/python/lib/sift_client/resources/annotations.py @@ -0,0 +1,556 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, cast + +from sift_client._internal.low_level_wrappers.annotations import AnnotationsLowLevelClient +from sift_client.resources._base import ResourceBase +from sift_client.sift_types.annotation import ( + Annotation, + AnnotationCommentElement, + AnnotationCreate, + AnnotationCreateBase, + AnnotationLog, + AnnotationLogKind, + AnnotationState, + AnnotationType, + AnnotationUpdate, + PhaseCreate, +) +from sift_client.sift_types.asset import Asset +from sift_client.sift_types.channel import Channel +from sift_client.sift_types.report import Report +from sift_client.sift_types.rule import Rule +from sift_client.sift_types.run import Run +from sift_client.sift_types.user import User +from sift_client.util import cel_utils as cel + +if TYPE_CHECKING: + import re + from collections.abc import Iterable + from datetime import datetime + + from sift_client.client import SiftClient + from sift_client.sift_types.calculated_channel import CalculatedChannel + from sift_client.sift_types.tag import Tag + + +class AnnotationLogsAPIAsync(ResourceBase): + """High-level API for an annotation's history. + + Each log records one event: an assignment, a state change, or a comment. + Reachable as `client.annotations.logs`. + """ + + def __init__(self, sift_client: SiftClient): + """Initialize the AnnotationLogsAPI. + + Args: + sift_client: The Sift client to use. + """ + super().__init__(sift_client) + self._low_level_client = AnnotationsLowLevelClient(grpc_client=self.client.grpc_client) + + async def list_( + self, + *, + annotation: str | Annotation, + # self ids + annotation_logs: list[str | AnnotationLog] | None = None, + # created/modified ranges + created_after: datetime | None = None, + created_before: datetime | None = None, + modified_after: datetime | None = None, + modified_before: datetime | None = None, + # created/modified users + created_by: Any | str | None = None, + # log specific + kind: AnnotationLogKind | None = None, + # common filters + filter_query: str | None = None, + order_by: str | None = None, + limit: int | None = None, + page_size: int | None = None, + ) -> list[AnnotationLog]: + """List annotation logs. + + Args: + annotation: The Annotation or annotation ID whose history to list. + annotation_logs: Filter to these AnnotationLogs or log IDs. + created_after: Filter logs created after this datetime. + created_before: Filter logs created before this datetime. + modified_after: Filter logs modified after this datetime. + modified_before: Filter logs modified before this datetime. + created_by: Filter logs created by this user ID. + kind: Filter to comments, state updates, or assignments. + filter_query: Explicit CEL query to filter logs. + order_by: Field and direction to order results by. + limit: Maximum number of logs to return. If None, returns all matches. + page_size: Number of results to fetch per request. + + Returns: + A list of AnnotationLog objects that match the filter criteria. + """ + filter_parts = [ + *self._build_time_cel_filters( + created_after=created_after, + created_before=created_before, + modified_after=modified_after, + modified_before=modified_before, + created_by=created_by, + ), + ] + if annotation_logs: + log_ids = [ + x._id_or_error if isinstance(x, AnnotationLog) else x for x in annotation_logs + ] + filter_parts.append(cel.in_("annotation_log_id", log_ids)) + if kind: + filter_parts.append(cel.equals("kind", kind.to_filter_str())) + if filter_query: + filter_parts.append(filter_query) + query_filter = cel.and_(*filter_parts) + + logs = await self._low_level_client.list_all_annotation_logs( + annotation_id=annotation._id_or_error + if isinstance(annotation, Annotation) + else annotation, + query_filter=query_filter or None, + order_by=order_by, + max_results=limit, + **({"page_size": page_size} if page_size is not None else {}), + ) + return self._apply_client_to_instances(logs) + + async def add_comment( + self, annotation: str | Annotation, text: str | list[AnnotationCommentElement] + ) -> AnnotationLog: + """Add a comment to an annotation. + + Args: + annotation: The Annotation or annotation ID to comment on. + text: Plain text, or a list of elements to mix text with user mentions. + + Returns: + The created AnnotationLog. + """ + body = [AnnotationCommentElement(text=text)] if isinstance(text, str) else text + log = await self._low_level_client.create_annotation_log( + annotation_id=annotation._id_or_error + if isinstance(annotation, Annotation) + else annotation, + kind=AnnotationLogKind.COMMENT, + comment=body, + ) + return self._apply_client_to_instance(log) + + +BATCH_LIMIT = 1000 +"""Annotations per call to BatchArchiveAnnotations, per the service.""" + + +def _create_from_dict(create: dict) -> AnnotationCreate | PhaseCreate: + """Pick the create model from a dict's `annotation_type`. + + A dict from a config file holds a name or a number, never a live enum member, and + omits the key entirely for a data review. + """ + value = create.get("annotation_type") + if value is None: + resolved = AnnotationType.DATA_REVIEW + elif isinstance(value, AnnotationType): + resolved = value + else: + try: + resolved = AnnotationType[value] if isinstance(value, str) else AnnotationType(value) + except (KeyError, ValueError): + names = ", ".join(t.name for t in AnnotationType) + raise ValueError( + f"Unknown annotation_type {value!r}. Expected one of: {names}" + ) from None + payload = {**create, "annotation_type": resolved} + if resolved is AnnotationType.PHASE: + return PhaseCreate.model_validate(payload) + return AnnotationCreate.model_validate(payload) + + +class AnnotationsAPIAsync(ResourceBase): + """High-level API for interacting with annotations. + + An annotation marks a time range on one or more assets. A data review annotation + carries a review state and an assignee. A phase annotation marks a segment of a run + and carries no state. + """ + + def __init__(self, sift_client: SiftClient): + """Initialize the AnnotationsAPI. + + Args: + sift_client: The Sift client to use. + """ + super().__init__(sift_client) + self._low_level_client = AnnotationsLowLevelClient(grpc_client=self.client.grpc_client) + self.logs = AnnotationLogsAPIAsync(sift_client) + + async def get(self, *, annotation_id: str) -> Annotation: + """Get an Annotation. + + Args: + annotation_id: The ID of the annotation. + + Returns: + The Annotation. + """ + annotation = await self._low_level_client.get_annotation(annotation_id=annotation_id) + return self._apply_client_to_instance(annotation) + + async def list_( + self, + *, + name: str | Iterable[str] | None = None, + name_contains: str | None = None, + name_regex: str | re.Pattern | None = None, + # self ids + annotation_ids: list[str] | None = None, + # created/modified ranges + created_after: datetime | None = None, + created_before: datetime | None = None, + modified_after: datetime | None = None, + modified_before: datetime | None = None, + # created/modified users + created_by: Any | str | None = None, + modified_by: Any | str | None = None, + # tags and metadata + tags: list[str] | list[Tag] | None = None, + metadata: dict[str, Any] | None = None, + # annotation specific + annotation_type: AnnotationType | None = None, + state: AnnotationState | None = None, + assigned_to: Any | str | None = None, + pending: bool | None = None, + assets: list[Asset] | list[str] | None = None, + runs: list[Run] | list[str] | None = None, + rules: list[str | Rule] | None = None, + reports: list[str | Report] | None = None, + start_time_after: datetime | None = None, + start_time_before: datetime | None = None, + end_time_after: datetime | None = None, + end_time_before: datetime | None = None, + # common filters + description_contains: str | None = None, + include_archived: bool = False, + filter_query: str | None = None, + order_by: str | None = None, + limit: int | None = None, + page_size: int | None = None, + ) -> list[Annotation]: + """List annotations. + + Args: + name: Exact name, or any iterable of names to match against. + name_contains: Partial name of the annotation. + name_regex: Regular expression to filter annotations by name. + annotation_ids: Filter to annotations with any of these IDs. + created_after: Filter annotations created after this datetime. + created_before: Filter annotations created before this datetime. + modified_after: Filter annotations modified after this datetime. + modified_before: Filter annotations modified before this datetime. + created_by: Filter annotations created by this user ID. + modified_by: Filter annotations last modified by this user ID. + tags: Filter annotations with any of these Tags or tag names. + metadata: Filter annotations by metadata criteria. + annotation_type: Filter to DATA_REVIEW or PHASE annotations. + state: Filter to a review state. + assigned_to: Filter to annotations assigned to this user's name. + pending: Filter to annotations from an ongoing rule violation. + rules: Filter to annotations created by any of these Rules or rule IDs. + reports: Filter to annotations in any of these Reports or report IDs. + assets: Filter annotations on any of these Assets or asset IDs. + runs: Filter annotations on any of these Runs or run IDs. + start_time_after: Filter annotations that start after this datetime. + start_time_before: Filter annotations that start before this datetime. + end_time_after: Filter annotations that end after this datetime. + end_time_before: Filter annotations that end before this datetime. + description_contains: Partial description of the annotation. + include_archived: If True, include archived annotations in results. + filter_query: Explicit CEL query to filter annotations. + order_by: Field and direction to order results by. + limit: Maximum number of annotations to return. If None, returns all matches. + page_size: Number of results to fetch per request. Lower this if you hit gRPC + message size limits on responses. If None, uses the server default. + + Returns: + A list of Annotation objects that match the filter criteria. + """ + filter_parts = [ + *self._build_name_cel_filters( + name=name if isinstance(name, str) else None, + names=None if name is None or isinstance(name, str) else list(name), + name_contains=name_contains, + name_regex=name_regex, + ), + *self._build_time_cel_filters( + created_after=created_after, + created_before=created_before, + modified_after=modified_after, + modified_before=modified_before, + created_by=created_by, + modified_by=modified_by, + ), + *self._build_tags_metadata_cel_filters(tag_names=tags, metadata=metadata), + *self._build_common_cel_filters( + description_contains=description_contains, + include_archived=include_archived, + filter_query=filter_query, + ), + ] + if annotation_ids: + filter_parts.append(cel.in_("annotation_id", annotation_ids)) + if annotation_type: + filter_parts.append(cel.equals("annotation_type", annotation_type.to_filter_str())) + if state: + filter_parts.append(cel.equals("state", state.to_filter_str())) + if assigned_to: + filter_parts.append(cel.equals("assignee", assigned_to)) + if pending is not None: + filter_parts.append(cel.equals("pending", pending)) + if assets: + if all(isinstance(a, str) for a in assets): + filter_parts.append(cel.in_("asset_id", cast("list[str]", assets))) + else: + asset_objs = cast("list[Asset]", assets) + filter_parts.append(cel.in_("asset_id", [a._id_or_error for a in asset_objs])) + if runs: + run_ids = [r._id_or_error if isinstance(r, Run) else r for r in runs] + filter_parts.append(cel.in_("run_id", run_ids)) + if rules: + filter_parts.append( + cel.in_("rule_id", [r._id_or_error if isinstance(r, Rule) else r for r in rules]) + ) + if reports: + filter_parts.append( + cel.in_( + "report_id", [r._id_or_error if isinstance(r, Report) else r for r in reports] + ) + ) + if start_time_after: + filter_parts.append(cel.greater_than("start_time", start_time_after)) + if start_time_before: + filter_parts.append(cel.less_than("start_time", start_time_before)) + if end_time_after: + filter_parts.append(cel.greater_than("end_time", end_time_after)) + if end_time_before: + filter_parts.append(cel.less_than("end_time", end_time_before)) + query_filter = cel.and_(*filter_parts) + + annotations = await self._low_level_client.list_all_annotations( + query_filter=query_filter or None, + order_by=order_by, + max_results=limit, + **({"page_size": page_size} if page_size is not None else {}), + ) + return self._apply_client_to_instances(annotations) + + async def find(self, **kwargs) -> Annotation | None: + """Find one annotation. Takes the same arguments as `list_`. + + Raises if more than one matches. + + Args: + **kwargs: Keyword arguments to pass to `list_`. + + Returns: + The Annotation found or None. + """ + annotations = await self.list_(**kwargs) + if len(annotations) > 1: + raise ValueError(f"Multiple ({len(annotations)}) annotations found for query") + elif len(annotations) == 1: + return annotations[0] + return None + + async def create(self, create: AnnotationCreateBase | dict) -> Annotation: + """Create an annotation. + + Pass an `AnnotationCreate` for a data review or a `PhaseCreate` for a phase. A + dict picks the model from its `annotation_type`, which may be a name or a number. + + Args: + create: The annotation definition. `assets` takes Assets or asset IDs, + `tags` takes Tags or tag names. + + Returns: + The created Annotation. + """ + if isinstance(create, dict): + create = _create_from_dict(create) + if not create.assets and create.linked_channels: + create.assets = cast( + "list[str | Asset]", await self._assets_for_channels(create.linked_channels) + ) + elif create.assets: + create.assets = cast("list[str | Asset]", await self._asset_names(create.assets)) + created = await self._low_level_client.create_annotation(create=create) + return self._apply_client_to_instance(created) + + async def _asset_names(self, assets: list[str | Asset]) -> list[str]: + names = [a.name for a in assets if isinstance(a, Asset)] + ids = [a for a in assets if isinstance(a, str)] + if ids: + names.extend(a.name for a in await self.client.async_.assets.list_(asset_ids=ids)) + return names + + async def _assets_for_channels( + self, channels: list[Channel] | list[CalculatedChannel] | list[Channel | CalculatedChannel] + ) -> list[str]: + asset_ids = {c.asset_id for c in channels if isinstance(c, Channel) and c.asset_id} + if not asset_ids: + return [] + assets = await self.client.async_.assets.list_(asset_ids=list(asset_ids)) + return [a.name for a in assets] + + async def update( + self, annotation: str | Annotation, update: AnnotationUpdate | dict + ) -> Annotation: + """Update an Annotation. + + `tags`, `linked_channels`, and `metadata` are replaced, not merged. + + Args: + annotation: The Annotation or annotation ID to update. + update: Updates to apply to the Annotation. + + Returns: + The updated Annotation. + """ + annotation_id = ( + annotation._id_or_error if isinstance(annotation, Annotation) else annotation + ) + if isinstance(update, dict): + update = AnnotationUpdate.model_validate(update) + update.resource_id = annotation_id + updated = await self._low_level_client.update_annotation(update) + return self._apply_client_to_instance(updated) + + async def archive(self, annotation: str | Annotation) -> Annotation: + """Archive an annotation. + + Args: + annotation: The Annotation or annotation ID to archive. + + Returns: + The archived Annotation. + """ + annotation_id = ( + annotation._id_or_error if isinstance(annotation, Annotation) else annotation + ) + archived = await self._low_level_client.archive_annotation(annotation_id=annotation_id) + return self._apply_client_to_instance(archived) + + async def unarchive(self, annotation: str | Annotation) -> Annotation: + """Unarchive an annotation. + + Args: + annotation: The Annotation or annotation ID to unarchive. + + Returns: + The unarchived Annotation. + """ + annotation_id = ( + annotation._id_or_error if isinstance(annotation, Annotation) else annotation + ) + unarchived = await self._low_level_client.unarchive_annotation(annotation_id=annotation_id) + return self._apply_client_to_instance(unarchived) + + async def batch_archive(self, annotations: list[str | Annotation]) -> list[Annotation]: + """Archive many annotations, one call per `BATCH_LIMIT` of them. + + Args: + annotations: The Annotations or annotation IDs to archive. + + Returns: + The archived Annotations. + """ + archived = [] + for batch in self._batches(annotations): + archived.extend( + await self._low_level_client.batch_archive_annotations(annotation_ids=batch) + ) + return self._apply_client_to_instances(archived) + + async def batch_unarchive(self, annotations: list[str | Annotation]) -> list[Annotation]: + """Unarchive many annotations, one call per `BATCH_LIMIT` of them. + + Args: + annotations: The Annotations or annotation IDs to unarchive. + + Returns: + The unarchived Annotations. + """ + unarchived = [] + for batch in self._batches(annotations): + unarchived.extend( + await self._low_level_client.batch_unarchive_annotations(annotation_ids=batch) + ) + return self._apply_client_to_instances(unarchived) + + @staticmethod + def _batches(annotations: list[str | Annotation]) -> list[list[str]]: + """Split annotations into ID batches the service will accept.""" + ids = [a._id_or_error if isinstance(a, Annotation) else a for a in annotations] + return [ids[i : i + BATCH_LIMIT] for i in range(0, len(ids), BATCH_LIMIT)] + + async def assign_to_user(self, annotation: str | Annotation, user: str | User) -> Annotation: + """Assign an annotation to a user for review. + + Args: + annotation: The Annotation or annotation ID to assign. + user: The User or user ID to assign to. + + Returns: + The updated Annotation. + """ + user_id = user._id_or_error if isinstance(user, User) else user + return await self.update(annotation, AnnotationUpdate(assigned_to_user_id=user_id)) + + async def _set_state(self, annotation: str | Annotation, state: AnnotationState) -> Annotation: + """Set the review state, skipping the call if already there. + + The server rejects a redundant state change with INVALID_ARGUMENT. The state is + only known without a fetch when an Annotation was passed rather than an ID. + """ + if isinstance(annotation, Annotation) and annotation.state is state: + return annotation + return await self.update(annotation, AnnotationUpdate(state=state)) + + async def set_accepted(self, annotation: str | Annotation) -> Annotation: + """Set the review state to Accepted. + + Args: + annotation: The Annotation or annotation ID. + + Returns: + The updated Annotation. + """ + return await self._set_state(annotation, AnnotationState.ACCEPTED) + + async def set_failed(self, annotation: str | Annotation) -> Annotation: + """Set the review state to Failed. + + Args: + annotation: The Annotation or annotation ID. + + Returns: + The updated Annotation. + """ + return await self._set_state(annotation, AnnotationState.FAILED) + + async def set_open(self, annotation: str | Annotation) -> Annotation: + """Set the review state to Open. + + Args: + annotation: The Annotation or annotation ID. + + Returns: + The updated Annotation. + """ + return await self._set_state(annotation, AnnotationState.OPEN) diff --git a/python/lib/sift_client/resources/sync_stubs/__init__.py b/python/lib/sift_client/resources/sync_stubs/__init__.py index b9544706d2..64206b780c 100644 --- a/python/lib/sift_client/resources/sync_stubs/__init__.py +++ b/python/lib/sift_client/resources/sync_stubs/__init__.py @@ -4,6 +4,8 @@ from sift_client._internal.sync_wrapper import generate_sync_api from sift_client.resources import ( + AnnotationLogsAPIAsync, + AnnotationsAPIAsync, AssetsAPIAsync, CalculatedChannelsAPIAsync, ChannelsAPIAsync, @@ -30,6 +32,11 @@ ) PingAPI = generate_sync_api(PingAPIAsync, "PingAPI") +# AnnotationLogsAPI must be generated before AnnotationsAPI so it can be nested under it. +AnnotationLogsAPI = generate_sync_api(AnnotationLogsAPIAsync, "AnnotationLogsAPI") +AnnotationsAPI = generate_sync_api( + AnnotationsAPIAsync, "AnnotationsAPI", nested_resources={"logs": AnnotationLogsAPI} +) AssetsAPI = generate_sync_api(AssetsAPIAsync, "AssetsAPI") CalculatedChannelsAPI = generate_sync_api(CalculatedChannelsAPIAsync, "CalculatedChannelsAPI") ChannelsAPI = generate_sync_api(ChannelsAPIAsync, "ChannelsAPI") @@ -87,6 +94,8 @@ DataImportAPI = generate_sync_api(DataImportAPIAsync, "DataImportAPI") __all__ = [ + "AnnotationLogsAPI", + "AnnotationsAPI", "AssetsAPI", "CalculatedChannelsAPI", "ChannelsAPI", diff --git a/python/lib/sift_client/resources/sync_stubs/__init__.pyi b/python/lib/sift_client/resources/sync_stubs/__init__.pyi index 9528097604..8275d1c63a 100644 --- a/python/lib/sift_client/resources/sync_stubs/__init__.pyi +++ b/python/lib/sift_client/resources/sync_stubs/__init__.pyi @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: import re + from collections.abc import Iterable from datetime import datetime, timedelta from pathlib import Path from typing import TYPE_CHECKING, Any, Sequence @@ -17,6 +18,16 @@ if TYPE_CHECKING: ReplayResult, ) from sift_client.client import SiftClient + from sift_client.sift_types.annotation import ( + Annotation, + AnnotationCommentElement, + AnnotationCreateBase, + AnnotationLog, + AnnotationLogKind, + AnnotationState, + AnnotationType, + AnnotationUpdate, + ) from sift_client.sift_types.asset import Asset, AssetUpdate from sift_client.sift_types.calculated_channel import ( CalculatedChannel, @@ -86,6 +97,314 @@ if TYPE_CHECKING: ) from sift_client.sift_types.user import User +class AnnotationLogsAPI: + """Sync counterpart to `AnnotationLogsAPIAsync`. + + High-level API for an annotation's history. + + Each log records one event: an assignment, a state change, or a comment. + Reachable as `client.annotations.logs`. + """ + + def __init__(self, sift_client: SiftClient): + """Initialize the AnnotationLogsAPI. + + Args: + sift_client: The Sift client to use. + """ + ... + + def _run(self, coro): ... + def add_comment( + self, annotation: str | Annotation, text: str | list[AnnotationCommentElement] + ) -> AnnotationLog: + """Add a comment to an annotation. + + Args: + annotation: The Annotation or annotation ID to comment on. + text: Plain text, or a list of elements to mix text with user mentions. + + Returns: + The created AnnotationLog. + """ + ... + + def list_( + self, + *, + annotation: str | Annotation, + annotation_logs: list[str | AnnotationLog] | None = None, + created_after: datetime | None = None, + created_before: datetime | None = None, + modified_after: datetime | None = None, + modified_before: datetime | None = None, + created_by: Any | str | None = None, + kind: AnnotationLogKind | None = None, + filter_query: str | None = None, + order_by: str | None = None, + limit: int | None = None, + page_size: int | None = None, + ) -> list[AnnotationLog]: + """List annotation logs. + + Args: + annotation: The Annotation or annotation ID whose history to list. + annotation_logs: Filter to these AnnotationLogs or log IDs. + created_after: Filter logs created after this datetime. + created_before: Filter logs created before this datetime. + modified_after: Filter logs modified after this datetime. + modified_before: Filter logs modified before this datetime. + created_by: Filter logs created by this user ID. + kind: Filter to comments, state updates, or assignments. + filter_query: Explicit CEL query to filter logs. + order_by: Field and direction to order results by. + limit: Maximum number of logs to return. If None, returns all matches. + page_size: Number of results to fetch per request. + + Returns: + A list of AnnotationLog objects that match the filter criteria. + """ + ... + +class AnnotationsAPI: + """Sync counterpart to `AnnotationsAPIAsync`. + + High-level API for interacting with annotations. + + An annotation marks a time range on one or more assets. A data review annotation + carries a review state and an assignee. A phase annotation marks a segment of a run + and carries no state. + """ + + def __init__(self, sift_client: SiftClient): + """Initialize the AnnotationsAPI. + + Args: + sift_client: The Sift client to use. + """ + ... + + def _run(self, coro): ... + def archive(self, annotation: str | Annotation) -> Annotation: + """Archive an annotation. + + Args: + annotation: The Annotation or annotation ID to archive. + + Returns: + The archived Annotation. + """ + ... + + def assign_to_user(self, annotation: str | Annotation, user: str | User) -> Annotation: + """Assign an annotation to a user for review. + + Args: + annotation: The Annotation or annotation ID to assign. + user: The User or user ID to assign to. + + Returns: + The updated Annotation. + """ + ... + + def batch_archive(self, annotations: list[str | Annotation]) -> list[Annotation]: + """Archive many annotations, one call per `BATCH_LIMIT` of them. + + Args: + annotations: The Annotations or annotation IDs to archive. + + Returns: + The archived Annotations. + """ + ... + + def batch_unarchive(self, annotations: list[str | Annotation]) -> list[Annotation]: + """Unarchive many annotations, one call per `BATCH_LIMIT` of them. + + Args: + annotations: The Annotations or annotation IDs to unarchive. + + Returns: + The unarchived Annotations. + """ + ... + + def create(self, create: AnnotationCreateBase | dict) -> Annotation: + """Create an annotation. + + Pass an `AnnotationCreate` for a data review or a `PhaseCreate` for a phase. A + dict picks the model from its `annotation_type`, which may be a name or a number. + + Args: + create: The annotation definition. `assets` takes Assets or asset IDs, + `tags` takes Tags or tag names. + + Returns: + The created Annotation. + """ + ... + + def find(self, **kwargs) -> Annotation | None: + """Find one annotation. Takes the same arguments as `list_`. + + Raises if more than one matches. + + Args: + **kwargs: Keyword arguments to pass to `list_`. + + Returns: + The Annotation found or None. + """ + ... + + def get(self, *, annotation_id: str) -> Annotation: + """Get an Annotation. + + Args: + annotation_id: The ID of the annotation. + + Returns: + The Annotation. + """ + ... + + def list_( + self, + *, + name: str | Iterable[str] | None = None, + name_contains: str | None = None, + name_regex: str | re.Pattern | None = None, + annotation_ids: list[str] | None = None, + created_after: datetime | None = None, + created_before: datetime | None = None, + modified_after: datetime | None = None, + modified_before: datetime | None = None, + created_by: Any | str | None = None, + modified_by: Any | str | None = None, + tags: list[str] | list[Tag] | None = None, + metadata: dict[str, Any] | None = None, + annotation_type: AnnotationType | None = None, + state: AnnotationState | None = None, + assigned_to: Any | str | None = None, + pending: bool | None = None, + assets: list[Asset] | list[str] | None = None, + runs: list[Run] | list[str] | None = None, + rules: list[str | Rule] | None = None, + reports: list[str | Report] | None = None, + start_time_after: datetime | None = None, + start_time_before: datetime | None = None, + end_time_after: datetime | None = None, + end_time_before: datetime | None = None, + description_contains: str | None = None, + include_archived: bool = False, + filter_query: str | None = None, + order_by: str | None = None, + limit: int | None = None, + page_size: int | None = None, + ) -> list[Annotation]: + """List annotations. + + Args: + name: Exact name, or any iterable of names to match against. + name_contains: Partial name of the annotation. + name_regex: Regular expression to filter annotations by name. + annotation_ids: Filter to annotations with any of these IDs. + created_after: Filter annotations created after this datetime. + created_before: Filter annotations created before this datetime. + modified_after: Filter annotations modified after this datetime. + modified_before: Filter annotations modified before this datetime. + created_by: Filter annotations created by this user ID. + modified_by: Filter annotations last modified by this user ID. + tags: Filter annotations with any of these Tags or tag names. + metadata: Filter annotations by metadata criteria. + annotation_type: Filter to DATA_REVIEW or PHASE annotations. + state: Filter to a review state. + assigned_to: Filter to annotations assigned to this user's name. + pending: Filter to annotations from an ongoing rule violation. + rules: Filter to annotations created by any of these Rules or rule IDs. + reports: Filter to annotations in any of these Reports or report IDs. + assets: Filter annotations on any of these Assets or asset IDs. + runs: Filter annotations on any of these Runs or run IDs. + start_time_after: Filter annotations that start after this datetime. + start_time_before: Filter annotations that start before this datetime. + end_time_after: Filter annotations that end after this datetime. + end_time_before: Filter annotations that end before this datetime. + description_contains: Partial description of the annotation. + include_archived: If True, include archived annotations in results. + filter_query: Explicit CEL query to filter annotations. + order_by: Field and direction to order results by. + limit: Maximum number of annotations to return. If None, returns all matches. + page_size: Number of results to fetch per request. Lower this if you hit gRPC + message size limits on responses. If None, uses the server default. + + Returns: + A list of Annotation objects that match the filter criteria. + """ + ... + + def set_accepted(self, annotation: str | Annotation) -> Annotation: + """Set the review state to Accepted. + + Args: + annotation: The Annotation or annotation ID. + + Returns: + The updated Annotation. + """ + ... + + def set_failed(self, annotation: str | Annotation) -> Annotation: + """Set the review state to Failed. + + Args: + annotation: The Annotation or annotation ID. + + Returns: + The updated Annotation. + """ + ... + + def set_open(self, annotation: str | Annotation) -> Annotation: + """Set the review state to Open. + + Args: + annotation: The Annotation or annotation ID. + + Returns: + The updated Annotation. + """ + ... + + def unarchive(self, annotation: str | Annotation) -> Annotation: + """Unarchive an annotation. + + Args: + annotation: The Annotation or annotation ID to unarchive. + + Returns: + The unarchived Annotation. + """ + ... + + def update(self, annotation: str | Annotation, update: AnnotationUpdate | dict) -> Annotation: + """Update an Annotation. + + `tags`, `linked_channels`, and `metadata` are replaced, not merged. + + Args: + annotation: The Annotation or annotation ID to update. + update: Updates to apply to the Annotation. + + Returns: + The updated Annotation. + """ + ... + @property + def logs(self) -> AnnotationLogsAPI: + """Nested AnnotationLogsAPI for making synchronous requests.""" + ... + class AssetsAPI: """Sync counterpart to `AssetsAPIAsync`. diff --git a/python/lib/sift_client/sift_types/__init__.py b/python/lib/sift_client/sift_types/__init__.py index 45589c9f0e..f844cbc25f 100644 --- a/python/lib/sift_client/sift_types/__init__.py +++ b/python/lib/sift_client/sift_types/__init__.py @@ -131,6 +131,19 @@ import sys +from sift_client.sift_types.annotation import ( + Annotation, + AnnotationCommentElement, + AnnotationCreate, + AnnotationCreateBase, + AnnotationLog, + AnnotationLogKind, + AnnotationLogState, + AnnotationState, + AnnotationType, + AnnotationUpdate, + PhaseCreate, +) from sift_client.sift_types.asset import Asset, AssetUpdate from sift_client.sift_types.calculated_channel import ( CalculatedChannel, @@ -234,6 +247,16 @@ TestReportUpdate.__test__ = False # type: ignore __all__ = [ + "Annotation", + "AnnotationCommentElement", + "AnnotationCreate", + "AnnotationCreateBase", + "AnnotationLog", + "AnnotationLogKind", + "AnnotationLogState", + "AnnotationState", + "AnnotationType", + "AnnotationUpdate", "Asset", "AssetUpdate", "CalculatedChannel", @@ -260,6 +283,7 @@ "JobStatus", "JobStatusDetails", "JobType", + "PhaseCreate", "PrincipalAttributeAssignment", "PrincipalAttributeEnumValue", "PrincipalAttributeKey", diff --git a/python/lib/sift_client/sift_types/annotation.py b/python/lib/sift_client/sift_types/annotation.py new file mode 100644 index 0000000000..86f3d23c49 --- /dev/null +++ b/python/lib/sift_client/sift_types/annotation.py @@ -0,0 +1,561 @@ +from __future__ import annotations + +from datetime import datetime, timezone +from enum import Enum +from typing import TYPE_CHECKING, ClassVar, Literal, cast + +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from sift.annotation_logs.v1.annotation_logs_pb2 import ( + AnnotationCommentBodyElement as AnnotationCommentBodyElementProto, +) +from sift.annotation_logs.v1.annotation_logs_pb2 import ( + AnnotationCommentBodyElementType as AnnotationCommentBodyElementTypeProto, +) +from sift.annotation_logs.v1.annotation_logs_pb2 import ( + AnnotationCommentUserMention as AnnotationCommentUserMentionProto, +) +from sift.annotation_logs.v1.annotation_logs_pb2 import ( + AnnotationLogKind as AnnotationLogKindProto, +) +from sift.annotation_logs.v1.annotation_logs_pb2 import ( + AnnotationLogSearchResult as AnnotationLogProto, +) +from sift.annotation_logs.v1.annotation_logs_pb2 import ( + AnnotationLogState as AnnotationLogStateProto, +) +from sift.annotations.v1.annotations_pb2 import ( + Annotation as AnnotationProto, +) +from sift.annotations.v1.annotations_pb2 import ( + AnnotationLinkedCalculatedChannel, + AnnotationLinkedChannelsChannel, +) +from sift.annotations.v1.annotations_pb2 import ( + AnnotationLinkedChannel as AnnotationLinkedChannelProto, +) +from sift.annotations.v1.annotations_pb2 import ( + AnnotationState as AnnotationStateProto, +) +from sift.annotations.v1.annotations_pb2 import ( + AnnotationType as AnnotationTypeProto, +) +from sift.annotations.v1.annotations_pb2 import ( + CreateAnnotationRequest as CreateAnnotationRequestProto, +) + +from sift_client.sift_types._base import ( + BaseType, + MappingHelper, + ModelCreate, + ModelCreateUpdateBase, + ModelUpdate, +) +from sift_client.sift_types.asset import Asset # noqa: TC001 +from sift_client.sift_types.calculated_channel import CalculatedChannel +from sift_client.sift_types.channel import Channel # noqa: TC001 +from sift_client.sift_types.run import Run +from sift_client.sift_types.tag import Tag +from sift_client.sift_types.user import User +from sift_client.util.metadata import metadata_dict_to_proto, metadata_proto_to_dict + +if TYPE_CHECKING: + from sift_client.client import SiftClient + + +class AnnotationType(Enum): + """Enum for annotation types.""" + + UNSPECIFIED = AnnotationTypeProto.ANNOTATION_TYPE_UNSPECIFIED # 0 + DATA_REVIEW = AnnotationTypeProto.ANNOTATION_TYPE_DATA_REVIEW # 1 + PHASE = AnnotationTypeProto.ANNOTATION_TYPE_PHASE # 2 + + def to_filter_str(self) -> str: + """Convert to the string used in CEL filters.""" + return f"ANNOTATION_TYPE_{self.name}" + + +class AnnotationState(Enum): + """Review state of a data review annotation, named as Sift shows it. + + The proto spells these UNSPECIFIED, OPEN, FLAGGED, and RESOLVED. + """ + + UNSPECIFIED = AnnotationStateProto.ANNOTATION_STATE_UNSPECIFIED # 0 + OPEN = AnnotationStateProto.ANNOTATION_STATE_OPEN # 1 + FAILED = AnnotationStateProto.ANNOTATION_STATE_FLAGGED # 2 + ACCEPTED = AnnotationStateProto.ANNOTATION_STATE_RESOLVED # 3 + + def to_filter_str(self) -> str: + """Convert to the string used in CEL filters, which uses the proto spelling.""" + proto_name = {"FAILED": "FLAGGED", "ACCEPTED": "RESOLVED"}.get(self.name, self.name) + return f"ANNOTATION_STATE_{proto_name}" + + +def _linked_channel_to_proto(channel: Channel | CalculatedChannel) -> AnnotationLinkedChannelProto: + """Wrap a Channel or CalculatedChannel in the proto's oneof.""" + if isinstance(channel, CalculatedChannel): + return AnnotationLinkedChannelProto( + calculated_channel=AnnotationLinkedCalculatedChannel( + calculated_channel_version_id=channel.version_id or "" + ) + ) + return AnnotationLinkedChannelProto( + channel=AnnotationLinkedChannelsChannel(channel_id=channel._id_or_error) + ) + + +class Annotation(BaseType[AnnotationProto, "Annotation"]): + """Annotation model representing a time range on one or more assets.""" + + # Required fields + name: str + description: str + start_time: datetime + end_time: datetime + annotation_type: AnnotationType + organization_id: str + created_date: datetime + modified_date: datetime + created_by_user_id: str + modified_by_user_id: str + tags: list[str] + asset_ids: list[str] + linked_channel_ids: list[str] + linked_calculated_channel_version_ids: list[str] + metadata: dict[str, str | float | bool] + is_archived: bool + pending: bool + + # Optional fields + state: AnnotationState | None + run_id: str | None + assigned_to_user_id: str | None + archived_date: datetime | None + + @classmethod + def _from_proto( + cls, proto: AnnotationProto, sift_client: SiftClient | None = None + ) -> Annotation: + return cls( + proto=proto, + id_=proto.annotation_id, + name=proto.name, + description=proto.description, + start_time=proto.start_time.ToDatetime(tzinfo=timezone.utc), + end_time=proto.end_time.ToDatetime(tzinfo=timezone.utc), + annotation_type=AnnotationType(proto.annotation_type), + organization_id=proto.organization_id, + created_date=proto.created_date.ToDatetime(tzinfo=timezone.utc), + modified_date=proto.modified_date.ToDatetime(tzinfo=timezone.utc), + created_by_user_id=proto.created_by_user_id, + modified_by_user_id=proto.modified_by_user_id, + tags=list(proto.tags), + asset_ids=list(proto.asset_ids), + linked_channel_ids=[ + c.bit_field_element.channel_id + if c.WhichOneof("type") == "bit_field_element" + else c.channel.channel_id + for c in proto.linked_channels + if c.WhichOneof("type") != "calculated_channel" + ], + linked_calculated_channel_version_ids=[ + c.calculated_channel.calculated_channel_version_id + for c in proto.linked_channels + if c.WhichOneof("type") == "calculated_channel" + ], + metadata=metadata_proto_to_dict(proto.metadata), # type: ignore + is_archived=proto.is_archived, + pending=proto.pending, + state=AnnotationState(proto.state) if proto.HasField("state") else None, + run_id=proto.run_id if proto.HasField("run_id") else None, + assigned_to_user_id=proto.assigned_to_user_id or None, + archived_date=( + proto.archived_date.ToDatetime(tzinfo=timezone.utc) + if proto.HasField("archived_date") + else None + ), + _client=sift_client, + ) + + @property + def logs(self) -> list[AnnotationLog]: + """Return this annotation's history.""" + return self.client.annotations.logs.list_(annotation=self._id_or_error) + + def update(self, update: AnnotationUpdate | dict) -> Annotation: + """Update the Annotation. + + Args: + update: The update to apply. See AnnotationUpdate for updatable fields. + + Returns: + The updated annotation. + """ + updated = self.client.annotations.update(annotation=self, update=update) + self._update(updated) + return self + + def archive(self) -> Annotation: + """Archive the annotation.""" + updated = self.client.annotations.archive(annotation=self) + self._update(updated) + return self + + def unarchive(self) -> Annotation: + """Unarchive the annotation.""" + updated = self.client.annotations.unarchive(annotation=self) + self._update(updated) + return self + + @property + def assets(self) -> list[Asset]: + """Fetch the Assets this annotation is on.""" + return self.client.assets.list_(asset_ids=self.asset_ids) if self.asset_ids else [] + + @property + def linked_channels(self) -> list[Channel | CalculatedChannel]: + """Fetch the Channels and CalculatedChannels this annotation is drawn on.""" + linked: list[Channel | CalculatedChannel] = [] + if self.linked_channel_ids: + linked.extend(self.client.channels.list_(channel_ids=self.linked_channel_ids)) + if self.linked_calculated_channel_version_ids: + quoted = ", ".join(f'"{v}"' for v in self.linked_calculated_channel_version_ids) + linked.extend( + self.client.calculated_channels.list_versions( + filter_query=f"calculated_channel_version_id in [{quoted}]" + ) + ) + return linked + + @property + def run(self) -> Run | None: + """Fetch the Run this annotation belongs to.""" + return self.client.runs.get(run_id=self.run_id) if self.run_id else None + + @property + def created_by(self) -> User: + """Fetch the User that created this annotation.""" + return self.client.users.get(user_id=self.created_by_user_id) + + @property + def modified_by(self) -> User: + """Fetch the User that last modified this annotation.""" + return self.client.users.get(user_id=self.modified_by_user_id) + + @property + def assigned_to(self) -> User | None: + """Fetch the User this annotation is assigned to.""" + return ( + self.client.users.get(user_id=self.assigned_to_user_id) + if self.assigned_to_user_id + else None + ) + + def assign_to_user(self, user: str | User) -> Annotation: + """Assign the annotation to a user for review. + + Args: + user: The User or user ID to assign to. + + Returns: + The updated annotation. + """ + updated = self.client.annotations.assign_to_user(annotation=self, user=user) + self._update(updated) + return self + + def set_open(self) -> Annotation: + """Set the review state to Open.""" + updated = self.client.annotations.set_open(annotation=self) + self._update(updated) + return self + + def set_failed(self) -> Annotation: + """Set the review state to Failed.""" + updated = self.client.annotations.set_failed(annotation=self) + self._update(updated) + return self + + def set_accepted(self) -> Annotation: + """Set the review state to Accepted.""" + updated = self.client.annotations.set_accepted(annotation=self) + self._update(updated) + return self + + def add_comment(self, text: str | list[AnnotationCommentElement]) -> AnnotationLog: + """Add a comment to the annotation. + + Args: + text: Plain text, or a list of elements to mix text with user mentions. + + Returns: + The created AnnotationLog. + """ + return self.client.annotations.logs.add_comment(self, text) + + +class AnnotationBase(ModelCreateUpdateBase): + """Base class for Annotation create and update models with shared fields and validation.""" + + description: str | None = None + tags: list[str] | list[Tag] | None = None + linked_channels: ( + list[Channel] | list[CalculatedChannel] | list[Channel | CalculatedChannel] | None + ) = Field(default=None, exclude=True) + metadata: dict[str, str | float | bool] | None = None + + _to_proto_helpers: ClassVar[dict[str, MappingHelper]] = { + "metadata": MappingHelper( + proto_attr_path="metadata", + update_field="metadata", + converter=metadata_dict_to_proto, + ), + "tags": MappingHelper( + proto_attr_path="tags", + update_field="tags", + converter=lambda tags: [tag.name if isinstance(tag, Tag) else tag for tag in tags], + ), + } + + def linked_channels_to_proto(self) -> list[AnnotationLinkedChannelProto]: + """Build the proto form of `linked_channels`.""" + return [_linked_channel_to_proto(c) for c in self.linked_channels or []] + + @model_validator(mode="after") + def _validate_time_range(self): + start = getattr(self, "start_time", None) + end = getattr(self, "end_time", None) + if start is not None and end is not None and start > end: + raise ValueError("start_time must not be after end_time") + return self + + +class AnnotationCreateBase(AnnotationBase, ModelCreate[CreateAnnotationRequestProto]): + """Shared fields for creating an annotation. + + `assets` takes Assets or asset IDs, and `tags` takes Tags or tag names. Omit + `assets` and the asset is taken from `linked_channels`. + """ + + name: str + start_time: datetime + end_time: datetime + assets: list[str | Asset] | None = None + """Assets or asset IDs. Resolved to names, which is what the proto takes.""" + run: str | Run | None = None + organization_id: str | None = None + + _to_proto_helpers: ClassVar[dict[str, MappingHelper]] = { + **AnnotationBase._to_proto_helpers, + "run": MappingHelper(proto_attr_path="run_id", update_field="run_id"), + } + + def _get_proto_class(self) -> type[CreateAnnotationRequestProto]: + return CreateAnnotationRequestProto + + @field_validator("run", mode="after") + @classmethod + def _run_to_id(cls, value): + return value._id_or_error if isinstance(value, Run) else value + + @model_validator(mode="after") + def _mark_annotation_type_set(self): + """Assign here, not as a field default. + + `to_proto` excludes unset fields, so a field default would send UNSPECIFIED. + """ + self.annotation_type = self.annotation_type + return self + + +class AnnotationCreate(AnnotationCreateBase): + """Create a data review annotation, which carries a review state and an assignee.""" + + annotation_type: Literal[AnnotationType.DATA_REVIEW] = AnnotationType.DATA_REVIEW + state: AnnotationState | None = None + assign_to_user: str | User | None = None + + _to_proto_helpers: ClassVar[dict[str, MappingHelper]] = { + **AnnotationCreateBase._to_proto_helpers, + "assign_to_user": MappingHelper( + proto_attr_path="assign_to_user_id", + update_field="assign_to_user_id", + ), + } + + @field_validator("assign_to_user", mode="after") + @classmethod + def _user_to_id(cls, value): + return value._id_or_error if isinstance(value, User) else value + + +class PhaseCreate(AnnotationCreateBase): + """Create a phase annotation, which labels a segment of a run and has no state. + + `extra="forbid"` so that passing `state` raises rather than being dropped. + """ + + model_config = ConfigDict(extra="forbid") + + annotation_type: Literal[AnnotationType.PHASE] = AnnotationType.PHASE + + +class AnnotationUpdate(AnnotationBase, ModelUpdate[AnnotationProto]): + """Update model for Annotation.""" + + name: str | None = None + start_time: datetime | None = None + end_time: datetime | None = None + assigned_to_user_id: str | None = None + state: AnnotationState | None = None + is_archived: bool | None = None + + def _get_proto_class(self) -> type[AnnotationProto]: + return AnnotationProto + + def _add_resource_id_to_proto(self, proto_msg: AnnotationProto): + if self._resource_id is None: + raise ValueError("Resource ID must be set before adding to proto") + proto_msg.annotation_id = self._resource_id + + +class AnnotationLogKind(Enum): + """Enum for the kind of entry in an annotation's history.""" + + UNSPECIFIED = AnnotationLogKindProto.ANNOTATION_LOG_KIND_UNSPECIFIED # 0 + COMMENT = AnnotationLogKindProto.ANNOTATION_LOG_KIND_COMMENT # 1 + STATE_UPDATE = AnnotationLogKindProto.ANNOTATION_LOG_KIND_STATE_UPDATE # 2 + ASSIGNED = AnnotationLogKindProto.ANNOTATION_LOG_KIND_ASSIGNED # 3 + + def to_filter_str(self) -> str: + """Convert to the string used in CEL filters. + + Log kinds filter on the bare name, unlike annotation type and state which + take the full enum value name. + """ + return self.name + + +class AnnotationLogState(Enum): + """Enum for the state recorded by a state update log.""" + + UNSPECIFIED = AnnotationLogStateProto.ANNOTATION_LOG_STATE_UNSPECIFIED # 0 + CREATED = AnnotationLogStateProto.ANNOTATION_LOG_STATE_CREATED # 1 + OPEN = AnnotationLogStateProto.ANNOTATION_LOG_STATE_OPEN # 2 + FAILED = AnnotationLogStateProto.ANNOTATION_LOG_STATE_FLAGGED # 3 + ACCEPTED = AnnotationLogStateProto.ANNOTATION_LOG_STATE_RESOLVED # 4 + + +class AnnotationCommentElement(BaseModel): + """One element of a comment body: literal text or a user mention. + + Attributes: + text: The literal text, when this element is text. + user_id: The mentioned user's ID, when this element is a mention. + """ + + text: str | None = None + user_id: str | User | None = None + + @field_validator("user_id", mode="after") + @classmethod + def _user_to_id(cls, value): + return value._id_or_error if isinstance(value, User) else value + + @model_validator(mode="after") + def _validate_exactly_one(self) -> AnnotationCommentElement: + if bool(self.text) == bool(self.user_id): + raise ValueError("AnnotationCommentElement requires exactly one of text or user_id") + return self + + @classmethod + def _from_proto(cls, proto: AnnotationCommentBodyElementProto) -> AnnotationCommentElement: + if ( + proto.type + == AnnotationCommentBodyElementTypeProto.ANNOTATION_COMMENT_BODY_ELEMENT_TYPE_USER_MENTION + ): + return cls(user_id=proto.user_mention.user_id) + return cls(text=proto.text) + + def _to_proto(self) -> AnnotationCommentBodyElementProto: + if self.user_id: + return AnnotationCommentBodyElementProto( + type=AnnotationCommentBodyElementTypeProto.ANNOTATION_COMMENT_BODY_ELEMENT_TYPE_USER_MENTION, + user_mention=AnnotationCommentUserMentionProto(user_id=cast("str", self.user_id)), + ) + return AnnotationCommentBodyElementProto( + type=AnnotationCommentBodyElementTypeProto.ANNOTATION_COMMENT_BODY_ELEMENT_TYPE_TEXT, + text=self.text or "", + ) + + +class AnnotationLog(BaseType[AnnotationLogProto, "AnnotationLog"]): + """One entry in an annotation's history: an assignment, a state change, or a comment.""" + + annotation_id: str + kind: AnnotationLogKind + created_date: datetime + modified_date: datetime + created_by_user_id: str + + # Set on ASSIGNED logs + assigned_to_user_id: str | None + # Set on STATE_UPDATE logs + state: AnnotationLogState | None + # Set on COMMENT logs + comment: list[AnnotationCommentElement] | None + + @classmethod + def _from_proto( + cls, proto: AnnotationLogProto, sift_client: SiftClient | None = None + ) -> AnnotationLog: + which = proto.WhichOneof("properties") + return cls( + proto=proto, + id_=proto.annotation_log_id, + annotation_id=proto.annotation_id, + kind=AnnotationLogKind(proto.kind), + created_date=proto.created_date.ToDatetime(tzinfo=timezone.utc), + modified_date=proto.modified_date.ToDatetime(tzinfo=timezone.utc), + created_by_user_id=proto.created_by_user_id, + assigned_to_user_id=( + proto.assigned.assigned_to_user_id if which == "assigned" else None + ), + state=( + AnnotationLogState(proto.state_update.state) if which == "state_update" else None + ), + comment=( + [AnnotationCommentElement._from_proto(e) for e in proto.comment.body] + if which == "comment" + else None + ), + _client=sift_client, + ) + + @property + def created_by(self) -> User: + """Fetch the User that created this log entry.""" + return self.client.users.get(user_id=self.created_by_user_id) + + @property + def assigned_to(self) -> User | None: + """Fetch the User this log entry assigned the annotation to.""" + if not self.assigned_to_user_id: + return None + return self.client.users.get(user_id=self.assigned_to_user_id) + + @property + def mentioned_users(self) -> list[User]: + """Fetch the Users mentioned in this comment.""" + return [ + self.client.users.get(user_id=cast("str", e.user_id)) + for e in self.comment or [] + if e.user_id + ] + + @property + def text(self) -> str: + """The comment body as plain text, with mentions rendered as user IDs.""" + if not self.comment: + return "" + return "".join(e.text or f"@{e.user_id}" for e in self.comment) diff --git a/python/lib/sift_client/sift_types/run.py b/python/lib/sift_client/sift_types/run.py index 516230153c..b61968d7b8 100644 --- a/python/lib/sift_client/sift_types/run.py +++ b/python/lib/sift_client/sift_types/run.py @@ -24,6 +24,7 @@ from sift_stream_bindings import RunFormPy from sift_client.client import SiftClient + from sift_client.sift_types.annotation import Annotation from sift_client.sift_types.asset import Asset from sift_client.sift_types.data_import import DataImport, DataTypeKey, ImportConfig from sift_client.sift_types.job import Job @@ -101,6 +102,11 @@ def assets(self) -> list[Asset]: return [] return self.client.assets.list_(asset_ids=self.asset_ids) + @property + def annotations(self) -> list[Annotation]: + """Fetch the Annotations on this run.""" + return self.client.annotations.list_(runs=[self]) + @property def data_imports(self) -> list[DataImport]: """Return all data imports that ingested into this run.""" diff --git a/python/lib/sift_client/util/util.py b/python/lib/sift_client/util/util.py index 80da283528..73253bcd3b 100644 --- a/python/lib/sift_client/util/util.py +++ b/python/lib/sift_client/util/util.py @@ -4,6 +4,7 @@ if TYPE_CHECKING: from sift_client.resources import ( + AnnotationsAPIAsync, AssetsAPIAsync, CalculatedChannelsAPIAsync, ChannelsAPIAsync, @@ -29,6 +30,9 @@ class AsyncAPIs(NamedTuple): ping: PingAPIAsync """Instance of the Ping API for making asynchronous requests.""" + annotations: AnnotationsAPIAsync + """Instance of the Annotations API for making asynchronous requests.""" + assets: AssetsAPIAsync """Instance of the Assets API for making asynchronous requests."""