Source code for chariot.inference_store.retention_task

import chariot._apis
from chariot.inference_store import _utils, models

__all__ = [
    "create_retention_task",
    "filter_retention_tasks",
    "get_retention_task",
]




def create_retention_task(
    model_id: str, request_body: models.NewRetentionTaskRequest
) -> models.RetentionTask | None:
    """Create a retention task associated with a non-automated retention policy.

    :param model_id: The model id.
    :type model_id: str
    :param request_body: The retention policy to execute.
    :type request_body: models.NewRetentionTaskRequest
    :return: The retention task.
    :rtype: Optional[models.RetentionTask]
    """
    response = chariot._apis.inferencestore.retention_task_api.retention_task_model_id_post(
        model_id=model_id, body=request_body.model_dump()
    )
    if not response.data:
        return None
    return _utils.convert_to_dataclass(response.data.model_dump(), models.RetentionTask)





def filter_retention_tasks(
    model_id: str, request_body: models.NewGetRetentionTasksRequest
) -> list[models.RetentionTask]:
    """Filter retention tasks.

    :param model_id: The model id.
    :type model_id: str
    :param request_body: The retention task filter.
    :type request_body: models.NewGetRetentionTasksRequest
    :return: The collection of retention tasks that meet the filter criteria.
    :rtype: List[models.RetentionTask]
    """
    response = chariot._apis.inferencestore.retention_task_api.retention_task_model_id_filter_post(
        model_id=model_id, body=request_body.model_dump()
    )
    if not response.data:
        return []
    return [
        _utils.convert_to_dataclass(d.model_dump(), models.RetentionTask) for d in response.data
    ]





def get_retention_task(model_id: str, retention_task_id: str) -> models.RetentionTask | None:
    """Get a retention task.

    :param model_id: The model id.
    :type model_id: str
    :param retention_task_id: The retention task id.
    :type retention_task_id: str
    :return: The retention task.
    :rtype: Optional[models.RetentionTask]
    """
    response = chariot._apis.inferencestore.retention_task_api.retention_task_model_id_task_id_get(
        model_id=model_id, task_id=retention_task_id
    )
    if not response.data:
        return None
    return _utils.convert_to_dataclass(response.data.model_dump(), models.RetentionTask)