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

import tri.ai.prompt.trace.AiExecInfo
import tri.ai.prompt.trace.AiOutputInfo
import tri.ai.prompt.trace.AiPromptTrace
import tri.ai.prompt.trace.AiPromptTraceSupport

/**
 * Task that can be executed by AI or API.
 * A task may have an arbitrary number of inputs that must be calculated prior to the task being executable.
 * The types of these inputs are not specified.
 */
abstract class AiTask(
    val id: String,
    val description: String? = null,
    val dependencies: Set = setOf()
) {
    /**
     * A task that can be executed, using a table of input results indexed by key.
     * Input types are not specified.
     */
    abstract suspend fun execute(inputs: Map>, monitor: AiTaskMonitor): AiPromptTraceSupport

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

    companion object {
        /** Creates a task. */
        fun  task(id: String, description: String? = null, op: suspend () -> T): AiTask =
            aitask(id, description) {
                val t0 = System.currentTimeMillis()
                val res = op()
                if (res is AiPromptTraceSupport<*>) throw IllegalArgumentException("Use aitask() for AiPromptTraceSupport")
                AiPromptTrace(execInfo = AiExecInfo.durationSince(t0), outputInfo = AiOutputInfo(listOf(res)))
            }

        /** Creates a task. */
        fun  aitask(id: String, description: String? = null, op: suspend () -> AiPromptTraceSupport): 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) {
                val t0 = System.currentTimeMillis()
                val res = op(it)
                if (res is AiPromptTraceSupport<*>) throw IllegalArgumentException("Use aitask() for AiPromptTraceSupport")
                AiPromptTrace(execInfo = AiExecInfo.durationSince(t0), outputInfo = AiOutputInfo(listOf(res)))
            }

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