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

Go to download
/*-
 * #%L
 * promptkt-0.1.0-SNAPSHOT
 * %%
 * Copyright (C) 2023 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

/** Stores a list of tasks that can be executed in order. */
class AiTaskList(tasks: List>, val lastTask: AiTask) {

    val plan = tasks + lastTask
    val planner
        get() = object : AiPlanner {
            override fun plan() = plan
        }

    constructor(firstTaskId: String, description: String? = null, op: suspend () -> AiTaskResult): this(listOf(),
        object: AiTask(firstTaskId, description) {
            override suspend fun execute(inputs: Map>, monitor: AiTaskMonitor) = op()
        })

    /** Adds a task to the end of the list. */
    fun  task(id: String, description: String? = null, op: suspend (S) -> T) =
        aitask(id, description) {
            val res = op(it)
            if (res is AiTaskResult<*>) throw IllegalArgumentException("Use aitask() for AiTaskResult")
            AiTaskResult.result(res, modelId = null)
        }

    /** Adds a task to the end of the list. */
    fun  aitask(id: String, description: String? = null, op: suspend (S) -> AiTaskResult): AiTaskList {
        val newTask = object : AiTask(id, description, setOf(lastTask.id)) {
            override suspend fun execute(inputs: Map>, monitor: AiTaskMonitor) =
                op(inputs[lastTask.id]!!.value as S)
        }
        return AiTaskList(plan, newTask)
    }

}

//region BUILDER METHODS

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

/** Creates a task list with a single task. */
fun  aitask(id: String, description: String? = null, op: suspend () -> AiTaskResult): AiTaskList =
    AiTaskList(id, description, op)

//endregion