Custom registration

genai_monitor.registration.api.register_function

register_function(
    func: Callable,
    model_output_to_bytes: Callable[[Any], bytes],
    bytes_to_model_output: Callable[[bytes], Any],
    model_output_to_base_type: Optional[
        Callable[[Any], BaseType]
    ] = None,
    parse_inference_method_arguments: Optional[
        Callable[[Dict[str, Any]], Jsonable]
    ] = None,
    model_hashing_function: Optional[
        Callable[[object], str]
    ] = None,
    max_unique_instances: int = 1,
)

Registers a function.

PARAMETER	DESCRIPTION
func	The function to register. TYPE: Callable
model_output_to_bytes	The function to convert the model output to bytes. TYPE: Callable[[Any], bytes]
bytes_to_model_output	The function to convert bytes to the model output. TYPE: Callable[[bytes], Any]
model_output_to_base_type	The function to convert the model output to a base type. TYPE: Optional[Callable[[Any], BaseType]] DEFAULT: None
parse_inference_method_arguments	The function to parse the inference method arguments. TYPE: Optional[Callable[[Dict[str, Any]], Jsonable]] DEFAULT: None
model_hashing_function	The function to hash the model. TYPE: Optional[Callable[[object], str]] DEFAULT: None
max_unique_instances	The maximum number of unique sample instances for each conditioning. TYPE: int DEFAULT: 1

Source code in src/genai_monitor/registration/api.py

def register_function(
    func: Callable,
    model_output_to_bytes: Callable[[Any], bytes],
    bytes_to_model_output: Callable[[bytes], Any],
    model_output_to_base_type: Optional[Callable[[Any], BaseType]] = None,
    parse_inference_method_arguments: Optional[Callable[[Dict[str, Any]], Jsonable]] = None,
    model_hashing_function: Optional[Callable[[object], str]] = None,
    max_unique_instances: int = 1,
):
    """Registers a function.

    Args:
        func: The function to register.
        model_output_to_bytes: The function to convert the model output to bytes.
        bytes_to_model_output: The function to convert bytes to the model output.
        model_output_to_base_type: The function to convert the model output to a base type.
        parse_inference_method_arguments: The function to parse the inference method arguments.
        model_hashing_function: The function to hash the model.
        max_unique_instances: The maximum number of unique sample instances for each conditioning.
    """
    func_name = f"{func.__module__}.{func.__name__}"

    if parse_inference_method_arguments is None:
        conditioning_parser = DefaultConditioningParser()
    else:
        conditioning_parser = _make_cls(
            cls_name=_make_conditioning_parser_name(func_name),
            base=BaseConditioningParser,
            method_mapper={"parse_func_arguments": parse_inference_method_arguments},
        )()

    conditioning_parser.max_unique_instances = max_unique_instances  # type: ignore

    output_parser_method_mapper = {
        "model_output_to_bytes": model_output_to_bytes,
        "bytes_to_model_output": bytes_to_model_output,
    }

    if model_output_to_base_type is not None:
        output_parser_method_mapper["model_output_to_base_type"] = model_output_to_base_type

    output_parser = _make_cls(
        cls_name=_make_output_parser_name(func.__name__),
        base=BaseModelOutputParser,
        method_mapper=output_parser_method_mapper,
    )()

    if not model_hashing_function:
        model_hashing_function = default_model_hashing_function

    register_inference_method(
        inference_method=func,
        output_parser=output_parser,
        conditioning_parser=conditioning_parser,
        hashing_function=model_hashing_function,
        max_unique_instances=max_unique_instances,
    )

genai_monitor.registration.api.register_class

register_class(
    cls: Type,
    inference_methods: List[str],
    model_output_to_bytes: Callable[[Any], bytes],
    bytes_to_model_output: Callable[[bytes], Any],
    model_output_to_base_type: Optional[
        Callable[[Any], BaseType]
    ] = None,
    parse_inference_method_arguments: Optional[
        Callable[[Dict[str, Any]], Jsonable]
    ] = None,
    model_hashing_function: Optional[
        Callable[[object], str]
    ] = None,
    max_unique_instances: int = 1,
)

Registers a class with inference methods.

PARAMETER	DESCRIPTION
cls	The class to register. TYPE: Type
inference_methods	The names of the inference methods. TYPE: List[str]
model_output_to_bytes	The function to convert the model output to bytes. TYPE: Callable[[Any], bytes]
bytes_to_model_output	The function to convert bytes to the model output. TYPE: Callable[[bytes], Any]
model_output_to_base_type	The function to convert the model output to a base type. TYPE: Optional[Callable[[Any], BaseType]] DEFAULT: None
parse_inference_method_arguments	The function to parse the inference method arguments. TYPE: Optional[Callable[[Dict[str, Any]], Jsonable]] DEFAULT: None
model_hashing_function	The function to hash the model. TYPE: Optional[Callable[[object], str]] DEFAULT: None
max_unique_instances	The maximum number of unique sample instances for each conditioning. TYPE: int DEFAULT: 1

Source code in src/genai_monitor/registration/api.py

def register_class(
    cls: Type,
    inference_methods: List[str],
    model_output_to_bytes: Callable[[Any], bytes],
    bytes_to_model_output: Callable[[bytes], Any],
    model_output_to_base_type: Optional[Callable[[Any], BaseType]] = None,
    parse_inference_method_arguments: Optional[Callable[[Dict[str, Any]], Jsonable]] = None,
    model_hashing_function: Optional[Callable[[object], str]] = None,
    max_unique_instances: int = 1,
):
    """Registers a class with inference methods.

    Args:
        cls: The class to register.
        inference_methods: The names of the inference methods.
        model_output_to_bytes: The function to convert the model output to bytes.
        bytes_to_model_output: The function to convert bytes to the model output.
        model_output_to_base_type: The function to convert the model output to a base type.
        parse_inference_method_arguments: The function to parse the inference method arguments.
        model_hashing_function: The function to hash the model.
        max_unique_instances: The maximum number of unique sample instances for each conditioning.
    """
    cls_name = cls.__name__
    methods_to_wrap = [getattr(cls, method_name) for method_name in inference_methods]
    if parse_inference_method_arguments is None:
        conditioning_parser = DefaultConditioningParser()
    else:
        conditioning_parser = _make_cls(
            cls_name=_make_conditioning_parser_name(cls_name),
            base=BaseConditioningParser,
            method_mapper={"parse_func_arguments": parse_inference_method_arguments},
        )()

    output_parser_method_mapper = {
        "model_output_to_bytes": model_output_to_bytes,
        "bytes_to_model_output": bytes_to_model_output,
    }

    if model_output_to_base_type is not None:
        output_parser_method_mapper["model_output_to_base_type"] = model_output_to_base_type

    output_parser = _make_cls(
        cls_name=_make_output_parser_name(cls_name),
        base=BaseModelOutputParser,
        method_mapper=output_parser_method_mapper,
    )()

    if not model_hashing_function:
        model_hashing_function = default_model_hashing_function

    for inference_method in methods_to_wrap:
        register_inference_method(
            inference_method=inference_method,
            output_parser=output_parser,
            conditioning_parser=conditioning_parser,
            hashing_function=model_hashing_function,
            max_unique_instances=max_unique_instances,
        )