feat(workflow): typed produces slots — StageConfig uses TypedArtifactSlot, TOML accepts {name, kind} objects

This commit is contained in:
2026-05-18 21:39:08 +04:00
parent c2267a58a0
commit a5dc983a57
4 changed files with 34 additions and 13 deletions
@@ -1,11 +1,10 @@
package com.correx.core.transitions.graph
import com.correx.core.artifacts.kind.TypedArtifactSlot
import com.correx.core.events.types.ArtifactId
import com.correx.core.inference.GenerationConfig
import com.correx.core.inference.ModelCapability
import kotlinx.serialization.Serializable
@Serializable
data class StageConfig(
val requiredCapabilities: Set<ModelCapability> = emptySet(),
val tokenBudget: Int = 4096,
@@ -16,7 +15,7 @@ data class StageConfig(
),
val allowedTools: Set<String> = emptySet(),
val maxRetries: Int = 3,
val produces: Set<ArtifactId> = emptySet(),
val produces: List<TypedArtifactSlot> = emptyList(),
val needs: Set<ArtifactId> = emptySet(),
val metadata: Map<String, String> = emptyMap(),
)
+1
View File
@@ -7,6 +7,7 @@ dependencies {
implementation(project(":core:transitions"))
implementation(project(":core:inference"))
implementation(project(":core:events"))
implementation(project(":core:artifacts"))
implementation("com.fasterxml.jackson.dataformat:jackson-dataformat-toml:2.17.0")
implementation("com.fasterxml.jackson.module:jackson-module-kotlin:2.17.0")
testImplementation "org.junit.jupiter:junit-jupiter"
@@ -1,5 +1,8 @@
package com.correx.infrastructure.workflow
import com.correx.core.artifacts.kind.ArtifactKindRegistry
import com.correx.core.artifacts.kind.DefaultArtifactKindRegistry
import com.correx.core.artifacts.kind.TypedArtifactSlot
import com.correx.core.events.types.ArtifactId
import com.correx.core.events.types.StageId
import com.correx.core.events.types.TransitionId
@@ -21,11 +24,16 @@ private data class WorkflowFile(
val transitions: List<TransitionSection> = emptyList(),
)
private data class ProducesEntry(
val name: String = "",
val kind: String = "",
)
private data class StageSection(
val id: String = "",
val prompt: String? = null,
val systemPrompt: String? = null,
val produces: List<String> = emptyList(),
@JsonProperty("system_prompt") val systemPrompt: String? = null,
val produces: List<ProducesEntry> = emptyList(),
val needs: List<String> = emptyList(),
@JsonProperty("allowed_tools") val allowedTools: List<String> = emptyList(),
@JsonProperty("token_budget") val tokenBudget: Int = 4096,
@@ -50,7 +58,9 @@ private val mapper = TomlMapper.builder()
.disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
.build()
class TomlWorkflowLoader : WorkflowLoader {
class TomlWorkflowLoader(
private val registry: ArtifactKindRegistry = DefaultArtifactKindRegistry(),
) : WorkflowLoader {
override fun load(path: Path): WorkflowGraph {
val raw = path.readText()
val file = runCatching { mapper.readValue<WorkflowFile>(raw) }
@@ -62,14 +72,18 @@ class TomlWorkflowLoader : WorkflowLoader {
val startId = StageId(start)
val stageMap = stages.associate { s ->
StageId(s.id) to StageConfig(
produces = s.produces.map { ArtifactId(it) }.toSet(),
produces = s.produces.map { entry ->
val resolvedKind = registry.get(entry.kind)
?: error("Unknown artifact kind '${entry.kind}' in stage '${s.id}'")
TypedArtifactSlot(name = ArtifactId(entry.name), kind = resolvedKind)
},
needs = s.needs.map { ArtifactId(it) }.toSet(),
allowedTools = s.allowedTools.toSet(),
tokenBudget = s.tokenBudget,
maxRetries = s.maxRetries,
metadata = buildMap {
s.prompt?.let { put("prompt", it) }
?: s.systemPrompt?.let { put("systemPrompt", it) }
s.systemPrompt?.let { put("systemPrompt", it) }
},
)
}
@@ -77,7 +91,7 @@ class TomlWorkflowLoader : WorkflowLoader {
val declaredIds = stageMap.keys.map { it.value }.toSet()
validate(start, declaredIds, transitions)
val allProduced = stageMap.values.flatMap { it.produces }.toSet()
val allProduced = stageMap.values.flatMap { it.produces }.map { it.name }.toSet()
stageMap.forEach { (stageId, config) ->
config.needs.forEach { needed ->
if (needed !in allProduced) {
@@ -17,7 +17,7 @@ class TomlWorkflowLoaderTest {
[[stages]]
id = "collect"
prompt = "prompts/collect.md"
produces = ["system_stats"]
produces = [{ name = "system_stats", kind = "file_written" }]
allowed_tools = ["ShellTool"]
token_budget = 2048
max_retries = 2
@@ -26,7 +26,7 @@ class TomlWorkflowLoaderTest {
id = "report"
prompt = "prompts/report.md"
needs = ["system_stats"]
produces = ["report"]
produces = [{ name = "report", kind = "file_written" }]
token_budget = 2048
[[transitions]]
@@ -53,10 +53,10 @@ class TomlWorkflowLoaderTest {
assertEquals(2, graph.transitions.size)
val collectStage = graph.stages[graph.start]!!
assertEquals(setOf("system_stats"), collectStage.produces.map { it.value }.toSet())
assertEquals(setOf("system_stats"), collectStage.produces.map { it.name.value }.toSet())
assertEquals("prompts/collect.md", collectStage.metadata["prompt"])
val reportStage = graph.stages.values.first { it.produces.any { a -> a.value == "report" } }
val reportStage = graph.stages.values.first { it.produces.any { a -> a.name.value == "report" } }
assertEquals(setOf("system_stats"), reportStage.needs.map { it.value }.toSet())
}
@@ -90,4 +90,11 @@ class TomlWorkflowLoaderTest {
val path = Files.createTempFile("workflow", ".toml").also { it.writeText(bad) }
assertThrows<WorkflowValidationException> { loader.load(path) }
}
@Test
fun `unknown artifact kind throws error`() {
val bad = validToml.replace("kind = \"file_written\"", "kind = \"unknown_kind\"")
val path = Files.createTempFile("workflow", ".toml").also { it.writeText(bad) }
assertThrows<IllegalStateException> { loader.load(path) }
}
}