import ast
from datetime import datetime
from enum import Enum
import numpy as np
from pydantic import BaseModel, ConfigDict, Field
[docs]
class TaskType(str, Enum):
"""The selection of task types supported by the inference store."""
IMAGE_CLASSIFICATION = "Image Classification"
OBJECT_DETECTION = "Object Detection"
IMAGE_SEGMENTATION = "Image Segmentation"
ORIENTED_OBJECT_DETECTION = "Oriented Object Detection"
IMAGE_EMBEDDING = "Image Embedding"
TEXT_EMBEDDING = "Text Embedding"
[docs]
class InferenceAction(str, Enum):
"""The selection of inference actions relative to task types
supported by the inference store.
"""
PREDICT = "predict"
PREDICT_PROBA = "predict_proba"
DETECT = "detect"
EMBED = "embed"
[docs]
class RetentionTaskState(str, Enum):
"""Defines the retention task state options."""
PENDING = "pending"
SCHEDULED = "scheduled"
RUNNING = "running"
FAILED = "failed"
STOPPED = "stopped"
COMPLETE = "complete"
[docs]
class DeleteAction(str, Enum):
"""The selection of delete actions."""
SOFT = "soft"
HARD = "hard"
[docs]
class EmbeddingDistanceMetric(str, Enum):
"""The selection of distance metrics supported by the inference store."""
L2 = "l2"
NEGATIVE_DOT_PRODUCT = "negative_dot_product"
COSINE = "cosine"
[docs]
class TimeWindowFilter(BaseModel):
"""Get inferences within a time window expressed in the
following datetime format: YYYY-MM-DDTHH:MM:SSZ. To format
an existing datetime object: dt.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
:param start: The left time bound. Defaults to 1970-01-01T00:00:00Z if not supplied.
:type start: Optional[str]
:param end: The right time bound: Defaults to now if not supplied.
:type end: Optional[str]
"""
start: str | None = None
end: str | None = None
[docs]
class EmbeddingFilter(BaseModel):
"""Get embeddings using a query embedding and distance metric.
:param query_embedding: The query embedding to search against.
:type query_embedding: List[float]
:param distance_metric: The distance metric operator (defaults to cosine)
:type end: str
"""
query_embedding: list[float]
distance_metric: str | None = EmbeddingDistanceMetric.COSINE
model_config = ConfigDict(use_enum_values=True)
[docs]
class GeoPoint(BaseModel):
"""Defines a point on the globe.
:param latitude: A latitude value in decimal format between -90 and 90.
:type latitude: float
:param longitude: A longitude value in decimal format between -180 and 180.
:type longitude: float
"""
latitude: float
longitude: float
[docs]
class GeoCircle(BaseModel):
"""Defines a circular search area using a centerpoint and radius.
:param center: A latitude value in decimal format between -90 and 90.
:type center: GeoPoint
:param radius: The radius in meters to expand from the center
:type radius: float
"""
center: GeoPoint
radius: float
[docs]
class GeoRectangle(BaseModel):
"""Defines a rectangular search area using two points on the globe.
:param p1: The first point on the globe.
:type p1: GeoPoint
:param p2: The second point on the globe.
:type p2: GeoPoint
"""
p1: GeoPoint
p2: GeoPoint
[docs]
class GeolocationFilter(BaseModel):
"""Defines a rectangular search area using two points on the globe.
:param gps_coordinates_circle: The circular area on the globe to search for inferences.
:type gps_coordinates_circle: models.GeoCircle
:param rectangular_search_area: The rectangular area on the globe to search for inferences.
:type rectangular_search_area: models.GeoRectangle
"""
gps_coordinates_circle: GeoCircle | None = None
gps_coordinates_rectangle: GeoRectangle | None = None
[docs]
class DeconstructedInferenceFilter(BaseModel):
"""Helper object for filtering inferences by labels and scores.
:param labels: The list of labels to filter on.
:type labels: Optional[List[str]]
:param minimum_score: The minimum confidence score.
:type minimum_score: Optional[float]
:param maximum_score: The maximum confidence score.
:type maximum_score: Optional[float]
"""
labels: list[str] | None = None
minimum_score: float | None = None
maximum_score: float | None = None
[docs]
class BaseInferenceFilter(BaseModel):
"""Helper object for filtering inferences.
:param inference_action_filter: Get inferences with a given inference action.
:type inference_action_filter: str
:param data_hash_filter: Get inferences with a given data hash.
:type data_hash_filter: str
:param data_source_filter: Get inferences with a given data source.
:type data_source_filter: str
:param time_window_filter: Get inferences within a time window.
:type time_window_filter: models.TimeWindowFilter
:param metadata_filter: Get inferences matching the intersection of one or more metadata filters.
:type metadata_filter: List[models.MetadataFilter]
:param location_filter: Get inferences in a circular or rectangular area on the globe.
:type location_filter: models.GeolocationFilter
:param deconstructed_inference_filter: Get inferences with certain labels or confidence scores.
:type deconstructed_inference_filter: models.DeconstructedInferenceFilter
"""
inference_action_filter: str | None = None
data_hash_filter: str | None = None
data_source_filter: str | None = None
time_window_filter: TimeWindowFilter | None = None
metadata_filter: list[MetadataFilter] | None = None
location_filter: GeolocationFilter | None = None
deconstructed_inference_filter: DeconstructedInferenceFilter | None = None
[docs]
class NewRegisterModelRequest(BaseModel):
"""Helper object to register a model for inference storage.
:param model_id: The model id.
:type model_id: str
:param project_id: The project id.
:type project_id: str
:param task_type: The model task type.
:type task_type: TaskType
:param embedding_size: The dimension of the embeddings produced.
:type embedding_size: int
"""
model_id: str
project_id: str
task_type: TaskType
embedding_size: int = 0
model_config = ConfigDict(protected_namespaces=(), use_enum_values=True)
[docs]
class NewInferenceStorageRequest(BaseModel):
"""Helper object to store an inference.
:param model_id: The id of the model.
:type model_id: str
:param inference_id: The id of the inference.
:type inference_id: str
:param data: The collection of inference and metadata.
:type data: models.NewInferenceAndMetadataCollection
:param data_storage_key: An optional data storage key returned from the upload.upload_data function.
:type data_storage_key: str
:param is_protected: Whether the inference and its associated data should be protected from deletion by retention policy
:type is_protected: bool
"""
model_id: str
inference_id: str
data: NewInferenceAndMetadataCollection
data_storage_key: str | None = None
is_protected: bool = False
model_config = ConfigDict(protected_namespaces=())
[docs]
class NewGetInferencesRequest(BaseModel):
"""Get inferences matching a series of filters.
:param filters: A collection of filters.
:type filters: models.BaseInferenceFilter
:param embedding_filters: An extended set of filters for embedding models.
:type embedding_filters: models.EmbeddingFilter
:param pagination: Get inferences matching pagination constraints.
:type pagination: models.Pagination
:param presign: Whether to presign data_storage_key(s).
:type presign: bool
"""
filters: BaseInferenceFilter = BaseInferenceFilter()
embedding_filters: EmbeddingFilter | None = None
pagination: Pagination | None = None
presign: bool = False
[docs]
class NewRetentionTaskRequest(BaseModel):
"""Helper object to create a retention task.
:param dry_run: If true, returns the number of inferences that would be deleted if the retention policy was fully executed.
:type dry_run: bool
:param retention_policy_id: The id of the retention policy to run.
:type retention_policy_id: str
"""
dry_run: bool
retention_policy_id: str
[docs]
class NewGetRetentionTasksRequest(BaseModel):
"""Helper object to filter retention tasks.
:param retention_policy_id: Filter by retention policy id.
:type retention_policy_id: str
:param state_filter: Filter by retention task state.
:type state_filter: models.RetentionTaskState
:param time_window_filter: Get retention tasks within a time window.
:type time_window_filter: models.TimeWindowFilter
:param pagination: Get retention tasks matching pagination constraints.
:type pagination: models.Pagination
"""
retention_policy_id: str | None = None
state_filter: RetentionTaskState | None = None
time_window_filter: TimeWindowFilter | None = None
pagination: Pagination | None = None
model_config = ConfigDict(use_enum_values=True)
[docs]
class NewRetentionPolicyRequest(BaseModel):
"""Helper object to create a retention policy. A value of -1 for maximum_record_age indicates
that inferences should never be deleted. The maximum_blob_age must be greater than or equal to 0.
The maximum_blob_age must also be equal to or less than the maximum_record_age.
:param automated_interval: Interval to automatically run retention policy in hours. Set to 0 if manual executions are desired.
:type automated_interval: int
:param maximum_record_age: The maximum age (in hours) since now in which an inference and its associated data is safe from deletion.
:type maximum_record_age: int
:param maximum_blob_age: The maximum age (in hours) since now in which a blob is safe from deletion.
:type maximum_blob_age: int
:param delete_action: Whether to 'soft' or 'hard' delete the database records
:type delete_action: models.DeleteAction
"""
automated_interval: int
maximum_record_age: int
maximum_blob_age: int
delete_action: DeleteAction = DeleteAction.SOFT
model_config = ConfigDict(use_enum_values=True)
[docs]
class NewExportTaskRequest(BaseModel):
"""Helper object to create a retention task.
:param filters: A collection of filters.
:type filters: models.BaseInferenceFilter
:param include_inference_ids: The list of additional inference ids to include.
:type include_inference_ids: list[str]
:param exclude_inference_ids: The list of additional inference ids to exclude.
:type exclude_inference_ids: list[str]
:param include_inferences_as_annotations: Whether to include the inferences as annotations (only eligible for version=v2 inferences)".
:type include_inferences_as_annotations: bool
:param include_custom_metadata: Whether to include custom metadata attached to the inference.
:type include_custom_metadata: bool
"""
filters: BaseInferenceFilter = BaseInferenceFilter()
include_inference_ids: list[str] = Field(default_factory=list)
exclude_inference_ids: list[str] = Field(default_factory=list)
include_inferences_as_annotations: bool = False
include_custom_metadata: bool = False
[docs]
class Model(BaseModel):
"""Defines an inference store model resource.
:param model_id: The model id.
:type model_id: str
:param created_at: A timestamp of when the model was created
:type created_at: datetime
:param updated_at: A timestamp of when the model was updated
:type updated_at: datetime
:param embedding_size: The dimension of the embeddings produced.
:type embedding_size: int
"""
model_id: str
created_at: datetime
updated_at: datetime
embedding_size: int
model_config = ConfigDict(protected_namespaces=())
[docs]
class Inference(BaseModel):
"""Defines an inference store inference resource.
:param model_id: The model id.
:type model_id: str
:param inference_id: The inference id.
:type inference_id: str
:param created_at: A timestamp of when the inference was created
:type created_at: datetime
:param updated_at: A timestamp of when the inference was updated
:type updated_at: datetime
:param inference_action: The inference action.
:type inference_action: str
:param data: The inference data. Returned as an arbitrary string when the associated model is not registered as an embedding model.
:type data: Optional[str]
:param embedding_distance: The embedding distance from a query embedding.
:type embedding_distance: Optional[float]
:param metadata: The collection of metadata associated to the inference.
:type metadata: Optional[List[Metadata]]
:param data_hash: The hash of the inference input data.
:type data_hash: str
:param data_source: The data source of the inference.
:type data_source: Optional[str]
:param data_coordinates: A set of geospatial coordinates defining where the inference occurred.
:type data_coordinates: Optional[GeoPoint]
:param data_storage_key: The internal data storage key for the inference input data.
:type data_storage_key: Optional[str]
:param presigned_url: A presigned url to download the inference input data.
:type presigned_url: Optional[str]
:param is_protected: Whether the inference and its associated data should be protected from deletion by retention policy
:type is_protected: bool
:param version: The inference version.
:type version: str
"""
model_id: str
inference_id: str
created_at: datetime
updated_at: datetime
inference_action: str
data: str | None = None
embedding_distance: float | None = None
metadata: list[Metadata] | None = None
data_hash: str
data_source: str | None = None
data_coordinates: GeoPoint | None = None
data_storage_key: str | None = None
presigned_url: str | None = None
is_protected: bool = False
version: str
model_config = ConfigDict(protected_namespaces=())
@property
def structured_embedding(self) -> list[int | float] | None:
if self.data is None:
return []
return ast.literal_eval(self.data)
[docs]
class RetentionTask(BaseModel):
"""Defines an inference store retentiont task resource.
:param id: The retention task id.
:type id: str
:param retention_policy_id: The retention policy id.
:type retention_policy_id: str
:param deleted_record_count: The count of all records currently deleted.
:type deleted_record_count: int
:param deleted_blob_count: The count of all blobs currently deleted.
:type deleted_blob_count: int
:param total_record_count: The expected total count of all records to be deleted.
:type total_record_count: int
:param total_blob_count: The expected total count of all blobs to be deleted.
:type total_blob_count: int
:param state: The state of the retention task.
:type state: str
:param maximum_record_age_limit: The timestamp in which records are safe from deletion.
:type maximum_record_age_limit: datetime
:param maximum_blob_age_limit: The timestamp in which blobs are safe from deletion.
:type maximum_blob_age_limit: datetime
"""
id: str
retention_policy_id: str
deleted_record_count: int
deleted_blob_count: int
total_record_count: int
total_blob_count: int
state: str
maximum_record_age_limit: datetime
maximum_blob_age_limit: datetime
[docs]
class RetentionPolicy(BaseModel):
"""Defines an inference store retention policy resource.
:param id: The retention policy id.
:type id: str
:param model_id: The model id.
:type model_id: str
:param delete_action: Whether records are soft deleted or hard deleted.
:type delete_action: bool
:param maximum_record_age: The maximum age (in hours) in which an inference and its associated data is safe from deletion.
:type maximum_record_age: int
:param maximum_blob_age: The maximum age (in hours) in which a blob is safe from deletion.
:type maximum_blob_age: int
:param automated_interval: The interval (in hours) in which this retention policy will be automatically run.
:type automated_interval: int
:param last_scheduled_at: The last time a retention task for this retention policy was run.
:type last_scheduled_at: datetime
"""
id: str
model_id: str
delete_action: str
maximum_record_age: int
maximum_blob_age: int
automated_interval: int | None = None
last_scheduled_at: datetime | None = None
model_config = ConfigDict(protected_namespaces=())
[docs]
class ExportTask(BaseModel):
"""Defines an inference store retentiont task resource.
:param id: The export task id.
:type id: str
:param progress_count: The current number of inferences added to the archive file.
:type progress_count: int
:param expected_count: The expected number of inferences to be added to archive file.
:type expected_count: int
:param state: The export task state.
:type state: str
:param presigned_url: A presigned url to the archive if the export task is complete.
:type presigned_url: Optional[str]
"""
id: str
progress_count: int
expected_count: int
state: str
presigned_url: str | None = None
[docs]
class DataUpload(BaseModel):
"""Defines an inference store data upload resource.
:param data_presigned_url: A presigned url to upload the inference input data to.
:type data_presigned_url: str
:param data_storage_key: The internal storage key to the inference input data.
:type data_storage_key: str
"""
data_presigned_url: str
data_storage_key: str