Source code for chariot.inference_store.metadata

from collections.abc import Mapping
from enum import Enum
from hashlib import sha256
from typing import Any

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

__all__ = [
    "create_metadata",
    "get_metadata",
    "delete_metadata",
    "get_metadata_key_type_counts",
    "get_metadata_statistics",
    "get_data_hash",
    "map_to_inference_store_metadata",
]




def create_metadata(
    model_id: str, inference_id: str, request_body: models.NewExtendedMetadataRequest
) -> models.Metadata | None:
    """Create new metadata for an inference.

    :param model_id: The model id.
    :type model_id: str
    :param inference_id: The id of the inference.
    :type inference_id: str
    :param request_body: The metadata to attach to the inference.
    :type request_body: models.NewMetadataRequest
    :return: The metadata details.
    :rtype: Optional[models.Metadata]Optional
    """
    response = chariot._apis.inferencestore.metadata_api.metadata_model_id_inference_id_post(
        model_id=model_id, inference_id=inference_id, body=request_body.model_dump()
    )
    if not response.data:
        return None
    return _utils.convert_to_dataclass(response.data.model_dump(), models.Metadata)





def get_metadata(model_id: str, inference_id: str) -> list[models.Metadata]:
    """Get all metadata for a particular inference.

    :param model_id: The model id.
    :type model_id: str
    :param inference_id: The id of the inference.
    :type inference_id: str
    :return: The metadata details.
    :rtype: List[models.Metadata]
    """
    response = chariot._apis.inferencestore.metadata_api.metadata_model_id_inference_id_get(
        model_id=model_id, inference_id=inference_id
    )
    if not response.data:
        return []
    return [_utils.convert_to_dataclass(d.model_dump(), models.Metadata) for d in response.data]





def delete_metadata(model_id: str, inference_id: str, key: str) -> str | None:
    """Delete metadata for a particular inference and key pair.

    :param model_id: The model id.
    :type model_id: str
    :param inference_id: The id of the inference.
    :type inference_id: str
    :param key: The key of the metadata to delete.
    :type key: str
    :return: The deleted metadata id.
    :rtype: Optional[str]
    """
    response = chariot._apis.inferencestore.metadata_api.metadata_model_id_inference_id_key_delete(
        model_id=model_id, inference_id=inference_id, key=key
    )
    if not response.data:
        return None
    return response.data





def get_metadata_key_type_counts(model_id: str) -> models.MetadataKeyTypeCounts | None:
    """Get metadata key, type counts.

    :param model_id: The model id.
    :type model_id: str
    :return: The count of each metadata key, type pair.
    :rtype: Optional[models.MetadataKeyTypeCounts]
    """
    response = chariot._apis.inferencestore.metadata_api.metadata_model_id_keys_get(
        model_id=model_id
    )
    if not response.data:
        return None
    return _utils.convert_to_dataclass(response.data.model_dump(), models.MetadataKeyTypeCounts)





def get_metadata_statistics(
    model_id: str, request_body: models.NewGetMetadataStatisticsRequest
) -> models.MetadataStatistics | None:
    """Get metadata statistics for a particular metadata key, type pair.

    :param model_id: The model id.
    :type model_id: str
    :param request_body: The metadata statistics criteria.
    :type request_body: models.NewGetMetadataStatisticsRequest
    :return: The metadata statistics
    :rtype: Optional[models.MetadataStatistics]
    """
    response = chariot._apis.inferencestore.metadata_api.metadata_model_id_statistics_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.MetadataStatistics)





def get_data_hash(data: bytes) -> str:
    """Get the SHA256 hexdigest of the data being inferred upon.

    :param data: The input inference data.
    :type data: bytes
    :return: The data hash.
    :rtype: str
    """
    return sha256(data).hexdigest()



def _map_value_type(vt: str) -> str | None:
    value_type = None
    if vt == "str":
        value_type = "string"
    elif vt == "float":
        value_type = "float"
    elif vt == "int":
        value_type = "int"
    elif vt == "dict":
        value_type = "json"
    return value_type


AllowedCustomMetadataValues = str | float | int | Mapping[str, Any]
CustomMetadata = Mapping[str, AllowedCustomMetadataValues]


class StandardMetadataKeys(str, Enum):
    DataSource = "data_source"
    Latitude = "latitude"
    Longitude = "longitude"




def map_to_inference_store_metadata(data: CustomMetadata) -> dict:
    """Map data from a dictionary into the required inference store format.
    All keys, types, and values must be strings.

    There are a few special keys mapped to the top level inference records
    for advanced querying: {"data_source", "latitude", "longitude"}. If the
    key doesn't belong to that set it's added to extended metadata.

    Valid inference store types are: string, float, int, json

    :param data: The metadata.
    :type data: dict
    :return: The typed metadata.
    :rtype: dict
    """
    standard_metadata: dict = {}
    extended_metadata: list[dict] = []
    for key, value in data.items():
        value_type = type(value).__name__
        vt = _map_value_type(value_type)
        if vt is None:
            continue
        match key:
            case StandardMetadataKeys.DataSource:
                standard_metadata[key] = str(value)
            case StandardMetadataKeys.Latitude | StandardMetadataKeys.Longitude:
                standard_metadata[key] = float(value)
            case _:
                extended_metadata.append({"key": str(key), "type": vt, "value": str(value)})
    out = {"standard_metadata": standard_metadata, "extended_metadata": extended_metadata}
    return out