Problem Statement
Design an in-memory Prompt Management and Template Engine (similar to prompt registry databases in LangChain or prompt-ops engines). The engine must allow developers to register multi-turn conversational templates with message roles (SYSTEM, USER, ASSISTANT), extract required variables (using double-curly brackets e.g., {{user_name}}) using regular expressions, validate and compile prompts by dynamically injecting variable maps, and support prompt versioning registries to retrieve specific configurations at runtime.
Design Decisions & Patterns Used
Working with Large Language Models (LLMs) requires structured prompt formatting. Instead of concatenating raw strings, modern architectures use structured message lists with roles. Additionally, prompts are treated as code artifacts: they are versioned, validated, and compiled. We separate the template parsing logic (extracting tokens like {{var}}) from the registry logic (managing versions) to keep the engine highly modular and thread-safe.
We will utilize the following Design Patterns:
- Builder Pattern: Providing a fluent interface to construct multi-message conversation templates easily.
- Strategy Pattern: Decoupling the templating engine (e.g., swapping double-curly bracket replacement with a different regex parsing format if needed).
- Factory Pattern: Standardizing the creation of structured `Message` formats based on role types.
Functional Requirements
- Model structured messages using roles:
SYSTEM,USER,ASSISTANT. - Parse dynamic variable tags (e.g.,
{{user}}) inside template text using regex. - Validate and compile template variables. Throw an exception if a required variable is missing from the injection map.
- Maintain a registry of versioned prompts, allowing retrieval of specific versions or the latest version automatically.
Objects Required
Role(Enum mapping chat message categories)Message(Value object tracking the final compiled role and text payload)MessageTemplate(Template representation of a single message, managing local variable extractions)PromptTemplate(Composite entity wrapping list templates, version tags, and compilation triggers)PromptManager(Core registry storing version history maps)
Role Enum & Message Class
The Role enum specifies chat participants, and the Message class acts as the final compiled output payload sent to the LLM api client.
public enum Role {
SYSTEM,
USER,
ASSISTANT
}
Let's define the Message class:
public class Message {
private final Role role;
private final String content;
public Message(Role role, String content) {
this.role = role;
this.content = content;
}
public Role getRole() { return role; }
public String getContent() { return content; }
@Override
public String toString() {
return String.format("[%s]: %s", role, content);
}
}
The constructor assigns immutable values, and the toString() method format helps output logs directly for tracing.
MessageTemplate Class
The MessageTemplate represents a single template line. It parses required variables from its text using regular expressions and compiles the string on request.
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
public class MessageTemplate {
private final Role role;
private final String templateText;
private final Set<String> variables;
// Pattern to match variables inside double curly brackets: {{variable}}
private static final Pattern VAR_PATTERN = Pattern.compile("\\{\\{([^}]+)\\}\\}");
public MessageTemplate(Role role, String templateText) {
this.role = role;
this.templateText = templateText;
this.variables = extractVariables(templateText);
}
private Set<String> extractVariables(String text) {
Set<String> vars = new HashSet<>();
Matcher matcher = VAR_PATTERN.matcher(text);
while (matcher.find()) {
vars.add(matcher.group(1).trim());
}
return vars;
}
public Message compile(Map<String, String> variableValues) {
String compiledText = templateText;
for (String var : variables) {
String value = variableValues.get(var);
if (value == null) {
throw new IllegalArgumentException("Validation Error: Missing value for prompt variable: " + var);
}
// Replace all instances of {{var}} with the value
compiledText = compiledText.replace("{{" + var + "}}", value);
}
return new Message(role, compiledText);
}
public Role getRole() { return role; }
public String getTemplateText() { return templateText; }
public Set<String> getVariables() { return variables; }
}
Let's break down the logic of every method in the MessageTemplate class:
MessageTemplate(role, templateText): The constructor sets up properties and immediately callsextractVariables()to parse and register required variables.extractVariables(text): Evaluates the text using regular expressions to capture terms wrapped in double-curly braces, saving them in a local set to define the template's parameters.compile(variableValues): Replaces variables with their values from the map, throwing anIllegalArgumentExceptionif any required variable is missing. It then instantiates and returns a compiledMessage.
PromptTemplate Class
The PromptTemplate class represents a full multi-turn prompt sequence, collecting variables across all messages and compiling them into a final message list.
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
public class PromptTemplate {
private final String id;
private final int version;
private final List<MessageTemplate> messageTemplates;
public PromptTemplate(String id, int version) {
this.id = id;
this.version = version;
this.messageTemplates = new ArrayList<>();
}
public void addMessageTemplate(Role role, String templateText) {
messageTemplates.add(new MessageTemplate(role, templateText));
}
public Set<String> getRequiredVariables() {
Set<String> allVars = new HashSet<>();
for (MessageTemplate mt : messageTemplates) {
allVars.addAll(mt.getVariables());
}
return allVars;
}
public List<Message> compile(Map<String, String> variableValues) {
// Run validation check across all required variables first
Set<String> required = getRequiredVariables();
for (String req : required) {
if (!variableValues.containsKey(req)) {
throw new IllegalArgumentException("Compile Error: Missing required prompt variable: '" + req + "'");
}
}
List<Message> compiledMessages = new ArrayList<>();
for (MessageTemplate mt : messageTemplates) {
compiledMessages.add(mt.compile(variableValues));
}
return compiledMessages;
}
public String getId() { return id; }
public int getVersion() { return version; }
}
Let's break down the logic of every method in the PromptTemplate class:
PromptTemplate(id, version): The constructor sets up the prompt ID and version number.addMessageTemplate(role, templateText): Adds a new message template to the conversational sequence.getRequiredVariables(): Aggregates the required variables from all message templates in the sequence.compile(variableValues): Validates that all required variables are present in the map, compiles the message templates, and returns the final compiled message list.
PromptManager Class
The PromptManager acts as the registry, storing prompts by ID and version number.
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
public class PromptManager {
// Maps promptId -> version -> PromptTemplate
private final Map<String, Map<Integer, PromptTemplate>> registry;
public PromptManager() {
this.registry = new ConcurrentHashMap<>();
}
public void registerPrompt(PromptTemplate template) {
registry.computeIfAbsent(template.getId(), k -> new HashMap<>())
.put(template.getVersion(), template);
System.out.println("Registered Prompt: '" + template.getId() + "' v" + template.getVersion());
}
public PromptTemplate getPrompt(String id, int version) {
Map<Integer, PromptTemplate> versions = registry.get(id);
if (versions == null || !versions.containsKey(version)) {
throw new IllegalArgumentException("No prompt found for ID '" + id + "' v" + version);
}
return versions.get(version);
}
public PromptTemplate getLatestPrompt(String id) {
Map<Integer, PromptTemplate> versions = registry.get(id);
if (versions == null || versions.isEmpty()) {
throw new IllegalArgumentException("No prompt found for ID '" + id + "'");
}
// Find maximum version key in the map
int latestVersion = versions.keySet().stream().max(Integer::compareTo).orElse(1);
return versions.get(latestVersion);
}
}
Let's break down the logic of every method in the PromptManager class:
PromptManager(): The constructor initializes the thread-safe concurrent registry map.registerPrompt(template): Adds a prompt template to the registry, nesting it under its ID and version number.getPrompt(id, version): Retrieves the requested version of a prompt, throwing an exception if the ID or version is not registered.getLatestPrompt(id): Finds and returns the highest version number registered for the prompt ID, helping manage default template configurations.
Main Driver Class
This class tests our prompt template engine. It registers templates, compiles variable injections, and checks the validation errors.
import java.util.HashMap;
import java.util.List;
import java.util.Map;
public class Main {
public static void main(String[] args) {
PromptManager manager = new PromptManager();
System.out.println("==========================================");
System.out.println("Scenario 1: Registering and Compiling Prompts");
System.out.println("==========================================");
// Define Prompt v1: Customer service assistant
PromptTemplate customerSupportV1 = new PromptTemplate("customer_support", 1);
customerSupportV1.addMessageTemplate(Role.SYSTEM, "You are an assistant specialized in {{domain}}.");
customerSupportV1.addMessageTemplate(Role.USER, "Explain how {{topic}} works.");
manager.registerPrompt(customerSupportV1);
// Compile prompt by injecting variable values
Map<String, String> vars1 = new HashMap<>();
vars1.put("domain", "Fintech Payment Gateways");
vars1.put("topic", "Tokenization API");
PromptTemplate fetched = manager.getLatestPrompt("customer_support");
List<Message> compiled1 = fetched.compile(vars1);
System.out.println("\n--- Compiled Conversation ---");
for (Message msg : compiled1) {
System.out.println(msg);
}
System.out.println("\n==========================================");
System.out.println("Scenario 2: Validation Check (Missing Variables)");
System.out.println("==========================================");
// Leave out the required variable "topic"
Map<String, String> invalidVars = new HashMap<>();
invalidVars.put("domain", "Healthtech EMR");
try {
System.out.println("Attempting to compile with missing variables:");
fetched.compile(invalidVars);
} catch (Exception e) {
System.out.println("Caught Expected Exception:\n" + e.getMessage());
}
}
}
The main() driver configures the prompt manager, registers templates, compiles conversational prompts with variable maps, and verifies that validation checks catch missing parameters correctly.
Comments
Post a Comment