caikit.core.task

Attributes

log
error
ProtoableInputTypes
ValidInputTypes
_InferenceMethodBaseT
_STREAM_OUT_ANNOTATION
_STREAM_PARAMS_ANNOTATION
_UNARY_OUT_ANNOTATION
_UNARY_PARAMS_ANNOTATION
_VISIBLE_ANNOTATION
_METADATA_ANNOTATION

Classes

TaskBase

The TaskBase defines the interface for an abstract AI task

Functions

task(→ Callable[[Type[TaskBase]], Type[TaskBase]])	The decorator for AI Task classes.
_make_keyname_for_module(→ str)

Module Contents

caikit.core.task.log

caikit.core.task.error

caikit.core.task.ProtoableInputTypes

caikit.core.task.ValidInputTypes

caikit.core.task._InferenceMethodBaseT

caikit.core.task._STREAM_OUT_ANNOTATION = '__streaming_output_type'

caikit.core.task._STREAM_PARAMS_ANNOTATION = '__streaming_params'

caikit.core.task._UNARY_OUT_ANNOTATION = '__unary_output_type'

caikit.core.task._UNARY_PARAMS_ANNOTATION = '__unary_params'

caikit.core.task._VISIBLE_ANNOTATION = '__visible'

caikit.core.task._METADATA_ANNOTATION = '__metadata'

class caikit.core.task.TaskBase

The TaskBase defines the interface for an abstract AI task

An AI task is a logical function signature which, when implemented, performs a task in some AI domain. The key property of a task is that the set of required input argument types and the output value type are consistent across all implementations of the task.

class InferenceMethodPtr

Little container class that holds a method name and its flavor of streaming. i.e. the args to a @TaskClass.taskmethod decoration.

method_name: str

input_streaming: bool

output_streaming: bool

context_arg: str | None

deferred_method_decorators: Dict[Type[TaskBase], Dict[str, List[TaskBase]]]

classmethod taskmethod(input_streaming: bool = False, output_streaming: bool = False, context_arg: str | None = None) → Callable[[_InferenceMethodBaseT], _InferenceMethodBaseT]

Decorates a module instancemethod and indicates whether the inputs and outputs should be handled as streams. This will trigger validation that the signature of this method is compatible with the task’s definition of input and output types.

The actual handling of validating the method and registering it is deferred until after the module class is created, which happens outside the context of this decoration.

classmethod deferred_method_decoration(module: Type)

Runs the actual decoration logic that taskmethod would have run if the module class existed during its lifetime.

Validates that all decorated methods match the task’s API expectations, and stores the signatures on the module class for access later.

classmethod has_inference_method_decorators(module_class: Type) → bool

Utility that returns true iff a module has any @TaskClass.taskmethod decorations

classmethod validate_run_signature(signature: caikit.core.signature_parsing.CaikitMethodSignature, input_streaming: bool, output_streaming: bool) → None

Validates that the provided method signature meets the api constraints defined in this task, for the given streaming flavors.

Raises:

ValueError if no type annotations were provided on the method TypeError if the type annotations do not meet the task’s api constraints

classmethod get_required_parameters(input_streaming: bool) → Dict[str, ValidInputTypes | Type[Iterable[ValidInputTypes]]]

Get the set of input types required by this task

classmethod get_output_type(output_streaming: bool) → Type[caikit.core.data_model.base.DataBase]

Get the output type for this task

NOTE: This method is automatically configured by the @task decorator

and should not be overwritten by child classes.

classmethod get_visibility() → bool

Get the visibility for this task.

NOTE: defaults to True even if visibility wasn’t provided

classmethod get_metadata() → Dict[str, Any]

Get any metadata defined for this task

NOTE: defaults to an empty dict if one wasn’t provided

classmethod _raise_on_wrong_output_type(output_type, module, output_streaming: bool)

static _subclass_check(this_type, that_type)

Wrapper around issubclass that first checks if both args are classes. Returns True if the types are the same, or they are both classes and this_type is a subclass of that_type

static _is_iterable_type(typ: Type) → bool

Returns True if typ is an iterable type. Does not work for types like list, tuple, but we’re interested here in List[T] etc.

This is implemented this way to support older python versions where isinstance(typ, typing.Iterable) does not work

caikit.core.task.task(unary_parameters: Dict[str, ValidInputTypes] = None, streaming_parameters: Dict[str, Type[Iterable[ValidInputTypes]]] = None, unary_output_type: Type[caikit.core.data_model.base.DataBase] = None, streaming_output_type: Type[Iterable[Type[caikit.core.data_model.base.DataBase]]] = None, visible: bool = True, metadata: Dict[str, Any] | None = None, **kwargs) → Callable[[Type[TaskBase]], Type[TaskBase]]

The decorator for AI Task classes.

This defines an output data model type for the task, and a minimal set of required inputs that all public models implementing this task must accept.

As an example, the caikit.interfaces.nlp.SentimentTask might look like:

@task(
    unary_parameters={
        "raw_document": caikit.interfaces.nlp.RawDocument
    },
    streaming_parameters={
        "raw_documents": Iterable[caikit.interfaces.nlp.RawDocument]
    }
    unary_output_type=caikit.interfaces.nlp.SentimentPrediction
    streaming_output_type=Iterable[caikit.interfaces.nlp.SentimentPrediction]
)
class SentimentTask(caikit.TaskBase):
    pass

and a module that implements this task might have methods like:

@module(id="b9d98408-84c2-488c-8385-9d698effe60b", task=SentimentTask)
class MyModule(ModuleBase):

    @SentimentTask.taskmethod()
    def run(raw_document: caikit.interfaces.nlp.RawDocument,
            inference_mode: str = "fast") ->
                caikit.interfaces.nlp.SentimentPrediction:
        # impl

    @SentimentTask.taskmethod(input_streaming=True, output_streaming=True)
    def run_bidi_stream(raw_documents: DataStream[caikit.interfaces.nlp.RawDocument])
            -> DataStream[caikit.interfaces.nlp.SentimentPrediction]:
        # impl

Note the run function may include other arguments beyond the minimal required inputs for the task.

Args:

unary_parameters (Dict[str, ValidInputTypes]): The required parameters that all module’s

unary-input inference methods must contain. A dictionary of parameter name to parameter type, where the types can be in the set of:

• Python primitives

• Caikit data models

• Iterable containers of the above

• Caikit model references (maybe?)

streaming_parameters: The same as unary_parameters, but for streaming-input inference

methods. All types must be in the form Iterable[T]

unary_output_type (Type[DataBase]): The unary output type of the task, which all modules’

unary-output inference methods must return. This must be a caikit data model type.

streaming_output_type (Type[Iterable[Type[DataBase]]]): The streaming output type of the

task, which all modules’ streaming-output inference methods must return. This must be in the form Iterable[T].

visible (bool): If this task should be exposed to the end user in documentation or if

it should only be used internally

metadata (Optional[Dict[str, Any]]): Any additional metadata that should

be included in the documentation for this task

Returns:

A decorator function for the task class, registering it with caikit’s core registry of

tasks.

caikit.core.task._make_keyname_for_module(module_class: Type) → str