Source code for chariot.inference_store.retention_policy

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

__all__ = [
    "create_retention_policy",
    "get_retention_policy",
    "delete_retention_policy",
]




def create_retention_policy(
    model_id: str, request_body: models.NewRetentionPolicyRequest
) -> models.RetentionPolicy | None:
    """Create a new retention policy for a particular model.

    :param model_id: The model id.
    :type model_id: str
    :param request_body: The retention policy to attach to the model.
    :type request_body: models.NewRetentionPolicyRequest
    :return: Retention policy details.
    :rtype: Optional[models.RetentionPolicy]
    """
    response = chariot._apis.inferencestore.retention_policy_api.retention_policy_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.RetentionPolicy)





def get_retention_policy(model_id: str) -> models.RetentionPolicy | None:
    """Get the current retention policy for a model.

    :param model_id: The model id.
    :type model_id: str
    :return: Retention policy details.
    :rtype: Optional[models.RetentionPolicy]
    """
    response = chariot._apis.inferencestore.retention_policy_api.retention_policy_model_id_get(
        model_id=model_id
    )
    if not response.data:
        return None
    return _utils.convert_to_dataclass(response.data.model_dump(), models.RetentionPolicy)





def delete_retention_policy(model_id: str, retention_policy_id: str) -> str | None:
    """Delete a retention policy for a model.

    :param model_id: The model id.
    :type model_id: str
    :param retention_policy_id: The retention policy id.
    :type retention_policy_id: str
    :return: The deleted retention policy id.
    :rtype: Optional[str]
    """
    response = chariot._apis.inferencestore.retention_policy_api.retention_policy_model_id_policy_id_delete(
        model_id=model_id, policy_id=retention_policy_id
    )
    if not response.data:
        return None
    return response.data