diff --git a/src/main/java/it/cnr/isti/workflow/manager/assistant/FlowAssistantPromptService.java b/src/main/java/it/cnr/isti/workflow/manager/assistant/FlowAssistantPromptService.java index d3771d6..a561ea2 100644 --- a/src/main/java/it/cnr/isti/workflow/manager/assistant/FlowAssistantPromptService.java +++ b/src/main/java/it/cnr/isti/workflow/manager/assistant/FlowAssistantPromptService.java @@ -71,7 +71,8 @@ public class FlowAssistantPromptService { { "blockId": "b1", "blockType": "LLMBlock", - "purpose": "..." + "purpose": "...", + "operation": "ADD" } ] } @@ -86,7 +87,13 @@ public class FlowAssistantPromptService { - Prefer HTTPServerCall for fetching files, APIs, or URLs before processing them. - If a task can be completed with a single stateless step, do not introduce MCPAgent. - If a later step must reuse data from an earlier step, make the earlier step produce it and the later step consume it explicitly. + - For DRAFT with a non-empty user request, "blocks" must contain at least one block. - blockId must be stable symbolic ids like b1, b2, b3. + - operation must be one of KEEP, ADD, UPDATE, REMOVE. + - For DRAFT, use ADD for every block. + - For REFINE/FIX with a current flow, mark unchanged existing blocks as KEEP. + - Use ADD only for new blocks, UPDATE only for existing blocks whose configuration must change, and REMOVE only for existing blocks to delete. + - For KEEP blocks, reuse the current block id/name when available and do not imply configuration changes. - Do not return block configuration yet. - Do not return connections yet. - Return valid JSON with no markdown fences. @@ -210,6 +217,7 @@ public class FlowAssistantPromptService { - Prefer canonical names such as response, output, input, and prompt when they exist. - Do not leave processing blocks disconnected when one step depends on another step's output or execution order. - For MCPAgent blocks that reuse a shared session, still connect the producer response to the consumer's upstream ordering placeholder, for example state_ready, so execution order is explicit. + - In REFINE/FIX, existing connections between KEEP/UPDATE blocks are preserved by the backend; return only connections that are new or intentionally changed. - Allowed block ids are listed in the flow plan and configured blocks. Use only those exact values. - Return the minimal set of connections required by the user request. - If no connections are needed, still return JSON with "connections": []. diff --git a/src/main/java/it/cnr/isti/workflow/manager/assistant/FlowAssistantService.java b/src/main/java/it/cnr/isti/workflow/manager/assistant/FlowAssistantService.java index 60254a6..0c5d01b 100644 --- a/src/main/java/it/cnr/isti/workflow/manager/assistant/FlowAssistantService.java +++ b/src/main/java/it/cnr/isti/workflow/manager/assistant/FlowAssistantService.java @@ -79,7 +79,14 @@ public class FlowAssistantService { private record AssistantFlowPlan(String name, String description, List blocks) { } - private record AssistantBlockPlan(String blockId, String blockType, String purpose) { + private record AssistantBlockPlan(String blockId, String blockType, String purpose, String operation) { + } + + private enum PlanOperation { + KEEP, + ADD, + UPDATE, + REMOVE } private record AssistantConfiguredBlockDraft(String blockId, String name, JsonNode config) { @@ -232,8 +239,8 @@ public class FlowAssistantService { String planPrompt = promptService.buildPlanPrompt(mode, userPrompt, currentFlow, errors, catalog); ParsedPlan parsedPlan = invokeStructuredAndValidate(provider, model, planPrompt, "plan", rawResponse -> { ParsedPlan plan = parsePlan(rawResponse); - AssistantFlowPlan normalizedPlan = validateAndNormalizePlan(plan.plan(), userPrompt, currentFlow, - catalogByType.keySet()); + AssistantFlowPlan normalizedPlan = validateAndNormalizePlan(plan.plan(), mode, userPrompt, currentFlow, + errors, catalogByType.keySet()); return new ParsedPlan(normalizedPlan, plan.rationale()); }); boolean requireSharedMemorySemantics = isSharedMemoryContext(userPrompt, currentFlow, parsedPlan.plan()); @@ -241,34 +248,61 @@ public class FlowAssistantService { List> assembledBlocks = new ArrayList<>(); Map> blocksByPlanId = new LinkedHashMap<>(); Map> blocksByAlias = new LinkedHashMap<>(); + Map> oldBlockIdToAssembledBlock = new LinkedHashMap<>(); + Set removedExistingBlockIds = new LinkedHashSet<>(); + Set usedExistingBlockIds = new LinkedHashSet<>(); List configuredBlocks = new ArrayList<>(); List rationaleParts = new ArrayList<>(); appendRationale(rationaleParts, parsedPlan.rationale()); - for (AssistantBlockPlan blockPlan : parsedPlan.plan().blocks()) { - progressListener.onProgress("configuring_blocks", "Configuring block " + blockPlan.blockId()); + List blockPlans = parsedPlan.plan().blocks(); + for (int blockIndex = 0; blockIndex < blockPlans.size(); blockIndex++) { + AssistantBlockPlan blockPlan = blockPlans.get(blockIndex); + PlanOperation operation = parsePlanOperation(blockPlan.operation()); + Block existingBlock = resolveExistingBlockForPlan(blockPlan, currentFlow, blockIndex, usedExistingBlockIds); + if (existingBlock != null) { + usedExistingBlockIds.add(existingBlock.getId()); + } + + if (operation == PlanOperation.REMOVE) { + if (existingBlock != null) { + removedExistingBlockIds.add(existingBlock.getId()); + } + continue; + } + BlockCatalogService.AssistantPromptBlockDescriptor descriptor = catalogByType.get(blockPlan.blockType()); if (descriptor == null) { throw new ResponseStatusException(HttpStatus.BAD_GATEWAY, "Assistant selected an unknown block type: " + blockPlan.blockType()); } - String blockPrompt = promptService.buildBlockConfigurationPrompt(mode, userPrompt, descriptor, - parsedPlan.plan(), blockPlan, currentFlow, errors, model); - ConfiguredBlockResult configuredBlock = invokeStructuredAndValidate(provider, model, blockPrompt, - "block configuration for " + blockPlan.blockId(), rawResponse -> { - ParsedBlockDraft parsedBlock = parseBlockDraft(rawResponse); - AssistantConfiguredBlockDraft normalizedDraft = normalizeBlockDraft(blockPlan, - parsedBlock.block()); - Block assembledBlock = buildBlock(descriptor, blockPlan, normalizedDraft, model, - requireSharedMemorySemantics); - return new ConfiguredBlockResult(parsedBlock, assembledBlock); - }); - appendRationale(rationaleParts, configuredBlock.parsedBlock().rationale()); + Block assembledBlock; + if (operation == PlanOperation.KEEP && existingBlock != null) { + progressListener.onProgress("configuring_blocks", "Keeping block " + blockPlan.blockId()); + assembledBlock = existingBlock; + } else { + progressListener.onProgress("configuring_blocks", "Configuring block " + blockPlan.blockId()); + String blockPrompt = promptService.buildBlockConfigurationPrompt(mode, userPrompt, descriptor, + parsedPlan.plan(), blockPlan, currentFlow, errors, model); + ConfiguredBlockResult configuredBlock = invokeStructuredAndValidate(provider, model, blockPrompt, + "block configuration for " + blockPlan.blockId(), rawResponse -> { + ParsedBlockDraft parsedBlock = parseBlockDraft(rawResponse); + AssistantConfiguredBlockDraft normalizedDraft = normalizeBlockDraft(blockPlan, + parsedBlock.block()); + Block newBlock = buildBlock(descriptor, blockPlan, normalizedDraft, model, + requireSharedMemorySemantics); + return new ConfiguredBlockResult(parsedBlock, newBlock); + }); + appendRationale(rationaleParts, configuredBlock.parsedBlock().rationale()); + assembledBlock = configuredBlock.block(); + } - Block assembledBlock = configuredBlock.block(); assembledBlocks.add(assembledBlock); blocksByPlanId.put(blockPlan.blockId(), assembledBlock); + if (existingBlock != null) { + oldBlockIdToAssembledBlock.put(existingBlock.getId(), assembledBlock); + } registerBlockAlias(blocksByAlias, blockPlan.blockId(), assembledBlock); registerBlockAlias(blocksByAlias, assembledBlock.getName(), assembledBlock); registerBlockAlias(blocksByAlias, blockPlan.purpose(), assembledBlock); @@ -300,13 +334,16 @@ public class FlowAssistantService { }); appendRationale(rationaleParts, parsedConnections.rationale()); - List connections = parsedConnections.connections().stream() + List generatedConnections = parsedConnections.connections().stream() .map(connection -> toConnection(connection, blocksByPlanId, blocksByAlias)) .toList(); + List connections = mergeConnections( + preserveCurrentConnections(currentFlow, oldBlockIdToAssembledBlock, removedExistingBlockIds), + generatedConnections); FlowCreateRequest flow = new FlowCreateRequest( - defaultIfBlank(parsedPlan.plan().name(), "Assistant flow"), - parsedPlan.plan().description(), + defaultIfBlank(parsedPlan.plan().name(), currentFlow == null ? "Assistant flow" : currentFlow.name()), + defaultIfBlank(parsedPlan.plan().description(), currentFlow == null ? null : currentFlow.description()), FlowData.builder() .blocks(assembledBlocks) .connections(connections) @@ -615,6 +652,124 @@ public class FlowAssistantService { return (Block) factory.create(configuration); } + private PlanOperation parsePlanOperation(String rawOperation) { + if (rawOperation == null || rawOperation.isBlank()) { + return PlanOperation.ADD; + } + try { + return PlanOperation.valueOf(rawOperation.trim().toUpperCase(Locale.ROOT)); + } catch (IllegalArgumentException e) { + throw new ResponseStatusException(HttpStatus.BAD_GATEWAY, + "Assistant returned an unknown block operation: " + rawOperation); + } + } + + private Block resolveExistingBlockForPlan(AssistantBlockPlan blockPlan, FlowCreateRequest currentFlow, + int planIndex, Set usedExistingBlockIds) { + List> existingBlocks = currentFlow == null || currentFlow.flow() == null + || currentFlow.flow().getBlocks() == null + ? List.of() + : currentFlow.flow().getBlocks(); + if (existingBlocks.isEmpty()) { + return null; + } + + Block exact = findExistingBlock(existingBlocks, usedExistingBlockIds, + block -> Objects.equals(block.getId(), blockPlan.blockId()) + || normalizeBlockReference(blockPlan.blockId()) != null + && normalizeBlockReference(blockPlan.blockId()) + .equals(normalizeBlockReference(block.getName()))); + if (exact != null) { + return exact; + } + + Block byPurpose = findExistingBlock(existingBlocks, usedExistingBlockIds, + block -> normalizeBlockReference(blockPlan.purpose()) != null + && normalizeBlockReference(blockPlan.purpose()).equals(normalizeBlockReference(block.getName()))); + if (byPurpose != null) { + return byPurpose; + } + + if (planIndex >= 0 && planIndex < existingBlocks.size()) { + Block byPosition = existingBlocks.get(planIndex); + if (!usedExistingBlockIds.contains(byPosition.getId()) + && Objects.equals(blockPlan.blockType(), blockTypeName(byPosition))) { + return byPosition; + } + } + + List> sameType = existingBlocks.stream() + .filter(block -> !usedExistingBlockIds.contains(block.getId())) + .filter(block -> Objects.equals(blockPlan.blockType(), blockTypeName(block))) + .toList(); + return sameType.size() == 1 ? sameType.getFirst() : null; + } + + private Block findExistingBlock(List> existingBlocks, Set usedExistingBlockIds, + java.util.function.Predicate> predicate) { + for (Block block : existingBlocks) { + if (block == null || usedExistingBlockIds.contains(block.getId())) { + continue; + } + if (predicate.test(block)) { + return block; + } + } + return null; + } + + private String blockTypeName(Block block) { + return block == null || block.getType() == null ? null : block.getType().getName(); + } + + private List preserveCurrentConnections(FlowCreateRequest currentFlow, + Map> oldBlockIdToAssembledBlock, Set removedExistingBlockIds) { + if (currentFlow == null || currentFlow.flow() == null || currentFlow.flow().getConnections() == null + || currentFlow.flow().getConnections().isEmpty()) { + return List.of(); + } + + List preserved = new ArrayList<>(); + for (Connection connection : currentFlow.flow().getConnections()) { + if (connection == null + || removedExistingBlockIds.contains(connection.getSourceId()) + || removedExistingBlockIds.contains(connection.getTargetId())) { + continue; + } + Block source = oldBlockIdToAssembledBlock.get(connection.getSourceId()); + Block target = oldBlockIdToAssembledBlock.get(connection.getTargetId()); + if (source == null || target == null) { + continue; + } + preserved.add(Connection.builder() + .sourceId(source.getId()) + .sourceName(connection.getSourceName()) + .targetId(target.getId()) + .targetName(connection.getTargetName()) + .build()); + } + return preserved; + } + + private List mergeConnections(List preservedConnections, List generatedConnections) { + Map merged = new LinkedHashMap<>(); + for (Connection connection : preservedConnections == null ? List.of() : preservedConnections) { + merged.put(connectionKey(connection), connection); + } + for (Connection connection : generatedConnections == null ? List.of() : generatedConnections) { + merged.putIfAbsent(connectionKey(connection), connection); + } + return List.copyOf(merged.values()); + } + + private String connectionKey(Connection connection) { + if (connection == null) { + return ""; + } + return connection.getSourceId() + "|" + connection.getSourceName() + "|" + + connection.getTargetId() + "|" + connection.getTargetName(); + } + private Connection toConnection(AssistantConnectionDraft connection, Map> blocksByPlanId, Map> blocksByAlias) { Block source = resolveConnectionBlock(connection.fromBlockId(), blocksByPlanId, blocksByAlias); @@ -925,8 +1080,8 @@ public class FlowAssistantService { } } - private AssistantFlowPlan validateAndNormalizePlan(AssistantFlowPlan plan, String userPrompt, - FlowCreateRequest currentFlow, + private AssistantFlowPlan validateAndNormalizePlan(AssistantFlowPlan plan, OperationMode mode, String userPrompt, + FlowCreateRequest currentFlow, List errors, Set availableBlockTypes) { if (plan == null) { throw new ResponseStatusException(HttpStatus.BAD_GATEWAY, "Assistant returned an empty plan"); @@ -936,12 +1091,15 @@ public class FlowAssistantService { return new AssistantFlowPlan( defaultIfBlank(plan.name(), "Assistant flow"), defaultIfBlank(plan.description(), userPrompt.trim()), - List.of(new AssistantBlockPlan("b1", "LLMBlock", userPrompt.trim()))); + List.of(new AssistantBlockPlan("b1", "LLMBlock", userPrompt.trim(), PlanOperation.ADD.name()))); } throw new ResponseStatusException(HttpStatus.BAD_GATEWAY, "Assistant returned a plan with no blocks"); } Set ids = new LinkedHashSet<>(); - for (AssistantBlockPlan block : plan.blocks()) { + Set usedExistingBlockIds = new LinkedHashSet<>(); + List normalizedBlocks = new ArrayList<>(); + for (int blockIndex = 0; blockIndex < plan.blocks().size(); blockIndex++) { + AssistantBlockPlan block = plan.blocks().get(blockIndex); if (block.blockId() == null || block.blockId().isBlank() || !ids.add(block.blockId())) { throw new ResponseStatusException(HttpStatus.BAD_GATEWAY, "Assistant returned invalid or duplicate block ids in the plan"); @@ -963,8 +1121,56 @@ public class FlowAssistantService { throw new ResponseStatusException(HttpStatus.BAD_GATEWAY, "Assistant returned an inconsistent plan: workflows that require shared memory or reusable state between steps must use MCPAgent blocks with a shared MCP session instead of standalone LLMBlock nodes"); } + PlanOperation operation = normalizePlanOperation(block, mode, currentFlow, errors, blockIndex, + usedExistingBlockIds); + normalizedBlocks.add(new AssistantBlockPlan( + block.blockId(), + block.blockType(), + block.purpose(), + operation.name())); } - return plan; + return new AssistantFlowPlan(plan.name(), plan.description(), normalizedBlocks); + } + + private PlanOperation normalizePlanOperation(AssistantBlockPlan block, OperationMode mode, + FlowCreateRequest currentFlow, List errors, int blockIndex, + Set usedExistingBlockIds) { + if (block.operation() != null && !block.operation().isBlank()) { + PlanOperation explicitOperation = parsePlanOperation(block.operation()); + if (explicitOperation != PlanOperation.ADD) { + Block matched = resolveExistingBlockForPlan(block, currentFlow, blockIndex, usedExistingBlockIds); + if (matched != null) { + usedExistingBlockIds.add(matched.getId()); + } + } + return explicitOperation; + } + if (mode == OperationMode.DRAFT || !hasCurrentFlowBlocks(currentFlow)) { + return PlanOperation.ADD; + } + + Block matched = resolveExistingBlockForPlan(block, currentFlow, blockIndex, usedExistingBlockIds); + if (matched == null) { + return PlanOperation.ADD; + } + usedExistingBlockIds.add(matched.getId()); + if (mode == OperationMode.FIX && hasValidationErrorForBlock(matched, errors)) { + return PlanOperation.UPDATE; + } + return PlanOperation.KEEP; + } + + private boolean hasValidationErrorForBlock(Block block, List errors) { + if (block == null || block.getId() == null || errors == null || errors.isEmpty()) { + return false; + } + for (ValidationError error : errors) { + if (Objects.equals(block.getId(), error.id()) + || error.relatedNodeIds() != null && error.relatedNodeIds().contains(block.getId())) { + return true; + } + } + return false; } private boolean canUseMinimalDraftFallback(String userPrompt, FlowCreateRequest currentFlow, diff --git a/src/test/java/it/cnr/isti/workflow/manager/controllers/AssistantControllerTest.java b/src/test/java/it/cnr/isti/workflow/manager/controllers/AssistantControllerTest.java index 8b98d2b..5e740ba 100644 --- a/src/test/java/it/cnr/isti/workflow/manager/controllers/AssistantControllerTest.java +++ b/src/test/java/it/cnr/isti/workflow/manager/controllers/AssistantControllerTest.java @@ -271,6 +271,27 @@ public class AssistantControllerTest { assertEquals(1, response.flow().flow().getConnections().size()); } + @Test + public void refineKeepsUnchangedBlocksWithoutReconfiguringThem() { + mockRefineKeepAndAddResponses(); + FlowCreateRequest initialFlow = TestAssistantResponses.singleBlockFlow(); + Block initialBlock = initialFlow.flow().getBlocks().getFirst(); + + AssistantFlowResponse response = assistantController.refine( + new AssistantRefineRequest( + "Add a human review after the LLM step", + initialFlow, + MODEL, + 1)); + + assertTrue(response.valid()); + assertEquals(2, response.flow().flow().getBlocks().size()); + assertEquals(initialBlock.getId(), response.flow().flow().getBlocks().getFirst().getId()); + assertEquals(initialBlock.getSpecificConfiguration(), + response.flow().flow().getBlocks().getFirst().getSpecificConfiguration()); + assertEquals(1, response.flow().flow().getConnections().size()); + } + @Test public void fixRepairsInvalidFlow() { mockAssistantResponses(); @@ -620,6 +641,60 @@ public class AssistantControllerTest { .thenAnswer(answer); } + private void mockRefineKeepAndAddResponses() { + Answer answer = invocation -> { + String prompt = invocation.getArgument(1, String.class); + if (prompt.contains("TASK: PLAN")) { + return TestAssistantResponses.wrap(java.util.Map.of( + "rationale", "Kept the classifier and planned a new human review block.", + "plan", java.util.Map.of( + "name", "Ticket classification with review", + "description", "Classify incoming tickets and review urgent ones.", + "blocks", java.util.List.of( + java.util.Map.of( + "blockId", "b1", + "blockType", "LLMBlock", + "purpose", "Classify incoming ticket", + "operation", "KEEP"), + java.util.Map.of( + "blockId", "b2", + "blockType", "HumanInteractionBlock", + "purpose", "Review urgent tickets", + "operation", "ADD"))))); + } + if (prompt.contains("TASK: BLOCK_CONFIG") + && prompt.contains("Current block to configure:\n{\n \"blockId\" : \"b1\"")) { + throw new AssertionError("Unchanged block b1 should not be reconfigured"); + } + if (prompt.contains("TASK: BLOCK_CONFIG") + && prompt.contains("Current block to configure:\n{\n \"blockId\" : \"b2\"")) { + return TestAssistantResponses.wrap(java.util.Map.of( + "rationale", "Configured the human review block.", + "block", java.util.Map.of( + "blockId", "b2", + "name", "Human review", + "config", java.util.Map.of( + "actionDescription", "Review high-risk tickets")))); + } + if (prompt.contains("TASK: CONNECTIONS")) { + return TestAssistantResponses.wrap(java.util.Map.of( + "rationale", "Connected classification to review.", + "connections", java.util.List.of( + java.util.Map.of( + "fromBlockId", "b1", + "fromOutput", "response", + "toBlockId", "b2", + "toInput", "input")))); + } + throw new IllegalStateException("Unexpected assistant prompt:\n" + prompt); + }; + + Mockito.when(internalOllamaLLMProvider.generate(Mockito.eq(MODEL), Mockito.anyString())) + .thenAnswer(answer); + Mockito.when(internalOllamaLLMProvider.generateJson(Mockito.eq(MODEL), Mockito.anyString())) + .thenAnswer(answer); + } + private void mockIndexCreateSessionResponsesWithConsumerFlags() { Answer answer = invocation -> { String prompt = invocation.getArgument(1, String.class);