Source code for chariot.inference_store.register

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

__all__ = [
    "register_model_for_inference_storage",
    "get_registered_model",
    "delete_registered_model",
]




@mcp_setting(mutating=True)
def register_model_for_inference_storage(
    request_body: models.NewRegisterModelRequest,
) -> str | None:
    """Register a model for inference storage.

    :param request_body: The model to register for inference storage.
    :type request_body: models.NewRegisterModelRequest
    :return: The model details.
    :rtype: Optional[str]
    """
    response = chariot._apis.inferencestore.register_api.model_post(body=request_body.model_dump())
    if not response.data:
        return None
    return response.data





def get_registered_model(model_id: str) -> models.Model | None:
    """Get details about a model registered for inference storage.

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





def delete_registered_model(model_id: str):
    """Delete a model registered for inference storage.

    :param model_id: The model id.
    :type model_id: str
    :return: The model id.
    :rtype: Optional[str]
    """
    response = chariot._apis.inferencestore.register_api.model_model_id_delete(model_id)
    if not response.data:
        return None
    return response.data