tri.ai.pips.AiTask.kt Maven / Gradle / Ivy

Go to download
/*-
 * #%L
 * tri.promptfx:promptkt
 * %%
 * Copyright (C) 2023 - 2024 Johns Hopkins University Applied Physics Laboratory
 * %%
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *      http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * #L%
 */
package tri.ai.pips

/** Task that can be executed by AI or API. */
abstract class AiTask(
    val id: String,
    val description: String? = null,
    val dependencies: Set = setOf()
) {
    abstract suspend fun execute(inputs: Map>, monitor: AiTaskMonitor): AiTaskResult

    /** Wrap this in a task that monitors and informs a callback when result is obtained. */
    fun monitor(callback: (T) -> Unit): AiTask = object : AiTask(id) {
        override suspend fun execute(inputs: Map>, monitor: AiTaskMonitor): AiTaskResult {
            val res = [email protected](inputs, monitor)
            res.value?.let { callback(it) }
            return res
        }
    }

    companion object {
        /** Creates a task. */
        fun  task(id: String, description: String? = null, op: suspend () -> T): AiTask =
            aitask(id, description) {
                val res = op()
                if (res is AiTaskResult<*>) throw IllegalArgumentException("Use aitask() for AiTaskResult")
                AiTaskResult.result(res, modelId = null)
            }

        /** Creates a task. */
        fun  aitask(id: String, description: String? = null, op: suspend () -> AiTaskResult): AiTask =
            object: AiTask(id, description) {
                override suspend fun execute(inputs: Map>, monitor: AiTaskMonitor) = op()
            }

        /** Creates a task that depends on a provided list of tasks. */
        fun  List>.task(id: String, description: String? = null, op: suspend (Map>) -> T): AiTask =
            aitask(id, description) { AiTaskResult.result(op(it)) }

        /** Creates a task that depends on a provided list of tasks. */
        fun  List>.aitask(id: String, description: String? = null, op: suspend (Map>) -> AiTaskResult): AiTask =
            object : AiTask(id, description, map { it.id }.toSet()) {
                override suspend fun execute(inputs: Map>, monitor: AiTaskMonitor) =
                    op(inputs)
            }
    }
}