Build Custom Modules

Comprehensive guide to developing custom ExecModules that extend ValkyrAI workflow capabilities.

What is an ExecModule?

An ExecModule is a pluggable worker that performs a specific operation in a workflow. Each module:

• Extends VModule base class
• Implements execute() method
• Receives Map<String, Object> input
• Returns Map<String, Object> output
• Never throws exceptions (catches and returns errors)
• Is reusable across workflows

Basic Structure

Minimal Example: Echo Module

package com.valkyrlabs.workflow.modules.custom;

import com.valkyrlabs.workflow.core.VModule;
import com.valkyrlabs.workflow.core.annotations.VModuleAware;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.Map;

@Component("echoModule")
@VModuleAware
public class EchoModule extends VModule {

    private static final Logger logger = LoggerFactory.getLogger(EchoModule.class);

    @Override
    public Map<String, Object> execute(
        Workflow workflow,
        Task task,
        ExecModule module,
        Map<String, Object> input
    ) {
        try {
            // Get input
            String message = (String) input.get("message");
            logger.info("Echo received: {}", message);

            // Process
            String result = message.toUpperCase();

            // Return output
            return Map.of(
                "status", "success",
                "original", message,
                "echoed", result
            );
        } catch (Exception e) {
            logger.error("Echo failed", e);
            return Map.of(
                "status", "error",
                "error", e.getMessage()
            );
        }
    }
}

Workflow YAML Usage

workflows:
  - id: echo-demo
    name: Echo Demonstration
    tasks:
      - id: task-1
        moduleId: echoModule
        inputMapping:
          message: workflow.state.userMessage
        outputMapping:
          echoed: workflow.state.echoResult

---

Input & Output Handling

Extracting Input Values

@Override
public Map<String, Object> execute(
    Workflow workflow,
    Task task,
    ExecModule module,
    Map<String, Object> input
) {
    // Get strings
    String email = (String) input.get("email");

    // Get integers
    Integer count = ((Number) input.get("count")).intValue();

    // Get booleans
    Boolean isActive = (Boolean) input.get("isActive");

    // Get objects
    Map<String, Object> user = (Map) input.get("user");
    String userName = (String) user.get("name");

    // Get arrays
    List<String> emails = (List) input.get("emailList");
    emails.forEach(e -> logger.info("Email: {}", e));

    // Get with defaults
    String param = (String) input.getOrDefault("optional", "default");

    // Get and cast safely
    Object value = input.get("uncertain");
    if (value instanceof String) {
        String str = (String) value;
    }

    return Map.of("status", "ok");
}

Returning Output

// Simple key-value map
return Map.of(
    "status", "success",
    "result", value
);

// Complex nested structure
Map<String, Object> response = new HashMap<>();
response.put("status", "success");
response.put("data", Map.of(
    "id", "123",
    "items", List.of("a", "b", "c")
));
return response;

// Error response
return Map.of(
    "status", "error",
    "code", "INVALID_INPUT",
    "message", "Email is required"
);

---

Common Module Patterns

Pattern 1: HTTP Call Module

@Component("httpCallModule")
@VModuleAware
public class HttpCallModule extends VModule {

    @Autowired
    private RestTemplate restTemplate;

    @Override
    public Map<String, Object> execute(
        Workflow workflow,
        Task task,
        ExecModule module,
        Map<String, Object> input
    ) {
        try {
            String url = (String) input.get("url");
            String method = (String) input.getOrDefault("method", "GET");
            Map<String, Object> body = (Map) input.get("body");

            ResponseEntity<Map> response = null;

            switch (method.toUpperCase()) {
                case "GET":
                    response = restTemplate.getForEntity(url, Map.class);
                    break;
                case "POST":
                    response = restTemplate.postForEntity(url, body, Map.class);
                    break;
                case "PUT":
                    restTemplate.put(url, body);
                    response = new ResponseEntity<>(body, HttpStatus.OK);
                    break;
                default:
                    return Map.of("status", "error", "error", "Unsupported method");
            }

            return Map.of(
                "status", "success",
                "statusCode", response.getStatusCode().value(),
                "body", response.getBody()
            );
        } catch (HttpClientErrorException e) {
            logger.error("HTTP {} error: {}", e.getStatusCode(), e.getMessage());
            return Map.of(
                "status", "error",
                "code", e.getStatusCode().value(),
                "error", e.getMessage()
            );
        } catch (Exception e) {
            logger.error("HTTP request failed", e);
            return Map.of("status", "error", "error", e.getMessage());
        }
    }
}

---

Pattern 2: Database Query Module

@Component("queryModule")
@VModuleAware
public class QueryModule extends VModule {

    @Autowired
    private JdbcTemplate jdbcTemplate;

    @Override
    public Map<String, Object> execute(
        Workflow workflow,
        Task task,
        ExecModule module,
        Map<String, Object> input
    ) {
        try {
            String query = (String) input.get("sql");
            List<Map<String, Object>> params = (List) input.getOrDefault("params", List.of());

            List<Map<String, Object>> results = jdbcTemplate.queryForList(
                query,
                params.toArray()
            );

            return Map.of(
                "status", "success",
                "rows", results,
                "count", results.size()
            );
        } catch (DataAccessException e) {
            logger.error("Query failed", e);
            return Map.of(
                "status", "error",
                "error", "Database error: " + e.getMessage()
            );
        }
    }
}

---

Pattern 3: Data Transform Module

@Component("transformModule")
@VModuleAware
public class TransformModule extends VModule {

    @Override
    public Map<String, Object> execute(
        Workflow workflow,
        Task task,
        ExecModule module,
        Map<String, Object> input
    ) {
        try {
            List<Map<String, Object>> items = (List) input.get("items");
            String operation = (String) input.get("operation");

            List<Map<String, Object>> result = switch (operation) {
                case "uppercase" -> items.stream()
                    .map(item -> Map.of(
                        "name", ((String) item.get("name")).toUpperCase(),
                        "value", item.get("value")
                    ))
                    .collect(Collectors.toList());

                case "filter_active" -> items.stream()
                    .filter(item -> (Boolean) item.getOrDefault("active", false))
                    .collect(Collectors.toList());

                case "sum_values" -> {
                    double sum = items.stream()
                        .mapToDouble(item -> ((Number) item.get("value")).doubleValue())
                        .sum();
                    yield List.of(Map.of("total", sum));
                }

                default -> items;
            };

            return Map.of(
                "status", "success",
                "items", result,
                "count", result.size()
            );
        } catch (Exception e) {
            logger.error("Transform failed", e);
            return Map.of("status", "error", "error", e.getMessage());
        }
    }
}

---

Pattern 4: Async/Callback Module

@Component("slackNotifyModule")
@VModuleAware
public class SlackNotifyModule extends VModule {

    @Autowired
    private RestTemplate restTemplate;

    @Autowired
    private WorkflowStateService workflowStateService;

    @Override
    public Map<String, Object> execute(
        Workflow workflow,
        Task task,
        ExecModule module,
        Map<String, Object> input
    ) {
        try {
            String webhookUrl = (String) module.getConfig().get("webhookUrl");
            String message = (String) input.get("message");
            UUID workflowId = workflow.getId();

            // Send async, don't wait for response
            CompletableFuture.runAsync(() -> {
                try {
                    Map<String, Object> payload = Map.of(
                        "text", message,
                        "ts", System.currentTimeMillis()
                    );
                    restTemplate.postForEntity(webhookUrl, payload, String.class);
                } catch (Exception e) {
                    logger.error("Slack notification failed", e);
                }
            });

            // Return immediately
            return Map.of(
                "status", "queued",
                "message", "Notification sent to Slack"
            );
        } catch (Exception e) {
            logger.error("Slack module failed", e);
            return Map.of("status", "error", "error", e.getMessage());
        }
    }
}

---

Configuration & Secrets

Access Module Configuration

// Module config comes from ExecModule.config (encrypted)
@Override
public Map<String, Object> execute(...) {
    Map<String, Object> config = module.getConfig();

    // Get encrypted values (auto-decrypted)
    String apiKey = (String) config.get("apiKey");
    String password = (String) config.get("dbPassword");

    // Get non-sensitive config
    String endpoint = (String) config.get("endpoint");
    Integer timeout = ((Number) config.get("timeout")).intValue();

    return Map.of("status", "ok");
}

Store Sensitive Config

// When creating ExecModule, mark sensitive fields
@SecureField  // Auto-encrypts at rest
private String apiKey;

---

Error Handling Best Practices

Do's ✅

// ✅ Catch all exceptions
try {
    // work
} catch (Exception e) {
    logger.error("Operation failed", e);
    return errorResponse("Operation failed");
}

// ✅ Return structured error
return Map.of(
    "status", "error",
    "code", "RESOURCE_NOT_FOUND",
    "message", "User 123 not found",
    "details", Map.of("userId", 123)
);

// ✅ Log with context
logger.error("Email send failed for userId={} attempt={}",
    userId, attemptNumber, e);

// ✅ Graceful degradation
try {
    return enrichedData;
} catch (EnrichmentException e) {
    logger.warn("Could not enrich, returning base data");
    return baseData;  // Partial success
}

Don'ts ❌

// ❌ Throw exceptions from execute()
throw new RuntimeException("Failed!");

// ❌ Return null
return null;

// ❌ Silent failures
try {
    doSomething();
} catch (Exception e) {
    // No logging, no error response!
}

// ❌ Swallow important context
catch (Exception e) {
    return Map.of("status", "error");  // Lost message
}

---

Testing Your Module

Unit Test Example

@ExtendWith(MockitoExtension.class)
class EchoModuleTest {

    @InjectMocks
    private EchoModule module;

    @Test
    void testExecuteSuccess() {
        // Arrange
        Workflow workflow = new Workflow();
        Task task = new Task();
        ExecModule execModule = new ExecModule();
        Map<String, Object> input = Map.of("message", "hello");

        // Act
        Map<String, Object> result = module.execute(workflow, task, execModule, input);

        // Assert
        assertThat(result).containsEntry("status", "success");
        assertThat(result).containsEntry("echoed", "HELLO");
    }

    @Test
    void testExecuteHandlesNullInput() {
        // Arrange
        Map<String, Object> input = Map.of();  // Missing "message"

        // Act
        Map<String, Object> result = module.execute(
            new Workflow(), new Task(), new ExecModule(), input
        );

        // Assert
        assertThat(result).containsEntry("status", "error");
    }
}

Integration Test Example

@SpringBootTest
@DirtiesContext
class EchoModuleIntegrationTest {

    @Autowired
    private WorkflowService workflowService;

    @Autowired
    private EchoModule echoModule;

    @Test
    void testModuleInWorkflow() throws Exception {
        // Create workflow with EchoModule
        Workflow workflow = createTestWorkflow(echoModule);

        // Execute
        Workflow result = workflowService.executeWorkflow(workflow).get();

        // Assert
        assertThat(result.getStatus()).isEqualTo(WorkflowStatus.COMPLETED);
        assertThat(result.getState())
            .containsEntry("echoResult", "HELLO WORLD");
    }
}

---

Module Lifecycle

Registration

// Automatically registered via @Component + @VModuleAware
@Component("myModule")
@VModuleAware
public class MyModule extends VModule {
    // ...
}

// Manually verify registration
@Autowired
private ApplicationContext context;

public void checkRegistration() {
    MyModule bean = context.getBean("myModule", MyModule.class);
    logger.info("Module registered: {}", bean.getClass().getSimpleName());
}

Discovery

// Modules discoverable via API
GET /v1/workflow/modules?search=my

// Response includes:
{
  "id": "myModule",
  "name": "My Custom Module",
  "description": "Does custom work",
  "inputSchema": { ... },
  "outputSchema": { ... }
}

Lifecycle Hooks

@Component("myModule")
@VModuleAware
public class MyModule extends VModule {

    @PostConstruct
    public void initialize() {
        logger.info("Module initialized");
    }

    @PreDestroy
    public void cleanup() {
        logger.info("Module cleaning up");
    }

    @Override
    public Map<String, Object> execute(...) {
        // ...
    }
}

---

Deployment Checklist

• Module extends VModule
• Annotated with @Component and @VModuleAware
• execute() method catches all exceptions
• Input/output validated and logged
• Error responses well-structured
• Sensitive data marked with @SecureField
• Unit tests with 80%+ coverage
• Integration test with sample workflow
• Documentation with example YAML
• No hardcoded secrets or endpoints
• Performance acceptable (< 5 seconds)
• Thread-safe if using shared state

✅ READY FOR PRODUCTION