Skip to content

feat: Add template config methods to AI SDK #184

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mattrmc1 wants to merge 1 commit into
base: main
Choose a base branch
from
+258 −16
Open

feat: Add template config methods to AI SDK #184

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 46 additions & 0 deletions lib/sdk/server-ai/src/main/java/com/launchdarkly/sdk/server/ai/LDAIClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,52 @@ AIJudgeConfig judgeConfig(
AIJudgeConfigDefault defaultValue,
Map<String, Object> variables);

/**
* Retrieves a completion (chat/prompt) AI Config with Mustache placeholders left intact
* (no interpolation). Useful for displaying prompt previews or storing templates for later
* rendering.
*
* @param key the AI Config key
* @param context the context to evaluate the configuration in
* @param defaultValue the default returned when the flag is absent or cannot be evaluated; when
* {@code null}, a disabled default is used
* @return the completion config with raw (non-interpolated) message content, never {@code null}
*/
AICompletionConfig completionConfigTemplate(
String key,
LDContext context,
AICompletionConfigDefault defaultValue);

/**
* Retrieves an agent AI Config with Mustache placeholders left intact (no interpolation). Useful
* for auditing instruction templates or building UI previews.
*
* @param key the AI Config key
* @param context the context to evaluate the configuration in
* @param defaultValue the default returned when the flag is absent or cannot be evaluated; when
* {@code null}, a disabled default is used
* @return the agent config with raw (non-interpolated) instructions, never {@code null}
*/
AIAgentConfig agentConfigTemplate(
String key,
LDContext context,
AIAgentConfigDefault defaultValue);

/**
* Retrieves a judge AI Config with Mustache placeholders left intact (no interpolation). Useful
* for auditing judge prompt templates.
*
* @param key the AI Config key
* @param context the context to evaluate the configuration in
* @param defaultValue the default returned when the flag is absent or cannot be evaluated; when
* {@code null}, a disabled default is used
* @return the judge config with raw (non-interpolated) message content, never {@code null}
*/
AIJudgeConfig judgeConfigTemplate(
String key,
LDContext context,
AIJudgeConfigDefault defaultValue);

/**
* Reconstructs a tracker from a resumption token, preserving the original run's identity.
* <p>
Expand Down
78 changes: 62 additions & 16 deletions lib/sdk/server-ai/src/main/java/com/launchdarkly/sdk/server/ai/LDAIClientImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,9 @@ public final class LDAIClientImpl implements LDAIClient {
private static final String TRACK_USAGE_AGENT_CONFIG = "$ld:ai:usage:agent-config";
private static final String TRACK_USAGE_AGENT_CONFIGS = "$ld:ai:usage:agent-configs";
private static final String TRACK_USAGE_JUDGE_CONFIG = "$ld:ai:usage:judge-config";
private static final String TRACK_USAGE_COMPLETION_CONFIG_TEMPLATE = "$ld:ai:usage:completion-config-template";
private static final String TRACK_USAGE_AGENT_CONFIG_TEMPLATE = "$ld:ai:usage:agent-config-template";
private static final String TRACK_USAGE_JUDGE_CONFIG_TEMPLATE = "$ld:ai:usage:judge-config-template";

private static final LDContext INIT_TRACK_CONTEXT = LDContext
.builder("ld-internal-tracking")
Expand Down Expand Up @@ -94,7 +97,7 @@ public AICompletionConfig completionConfig(
client.trackMetric(TRACK_USAGE_COMPLETION_CONFIG, context, LDValue.of(key), 1);
AICompletionConfigDefault effectiveDefault =
defaultValue != null ? defaultValue : AICompletionConfigDefault.disabled();
return (AICompletionConfig) evaluate(key, context, effectiveDefault, Mode.COMPLETION, variables);
return (AICompletionConfig) evaluate(key, context, effectiveDefault, Mode.COMPLETION, variables, true);
}

@Override
Expand Down Expand Up @@ -137,14 +140,47 @@ public AIJudgeConfig judgeConfig(
client.trackMetric(TRACK_USAGE_JUDGE_CONFIG, context, LDValue.of(key), 1);
AIJudgeConfigDefault effectiveDefault =
defaultValue != null ? defaultValue : AIJudgeConfigDefault.disabled();
return (AIJudgeConfig) evaluate(key, context, effectiveDefault, Mode.JUDGE, variables);
return (AIJudgeConfig) evaluate(key, context, effectiveDefault, Mode.JUDGE, variables, true);
}

@Override
public AICompletionConfig completionConfigTemplate(
String key,
LDContext context,
AICompletionConfigDefault defaultValue) {
client.trackMetric(TRACK_USAGE_COMPLETION_CONFIG_TEMPLATE, context, LDValue.of(key), 1);
AICompletionConfigDefault effectiveDefault =
defaultValue != null ? defaultValue : AICompletionConfigDefault.disabled();
return (AICompletionConfig) evaluate(key, context, effectiveDefault, Mode.COMPLETION, null, false);
}

@Override
public AIAgentConfig agentConfigTemplate(
String key,
LDContext context,
AIAgentConfigDefault defaultValue) {
client.trackMetric(TRACK_USAGE_AGENT_CONFIG_TEMPLATE, context, LDValue.of(key), 1);
AIAgentConfigDefault effectiveDefault =
defaultValue != null ? defaultValue : AIAgentConfigDefault.disabled();
return (AIAgentConfig) evaluate(key, context, effectiveDefault, Mode.AGENT, null, false);
}

@Override
public AIJudgeConfig judgeConfigTemplate(
String key,
LDContext context,
AIJudgeConfigDefault defaultValue) {
client.trackMetric(TRACK_USAGE_JUDGE_CONFIG_TEMPLATE, context, LDValue.of(key), 1);
AIJudgeConfigDefault effectiveDefault =
defaultValue != null ? defaultValue : AIJudgeConfigDefault.disabled();
return (AIJudgeConfig) evaluate(key, context, effectiveDefault, Mode.JUDGE, null, false);
}

private AIAgentConfig evaluateAgent(
String key, LDContext context, AIAgentConfigDefault defaultValue, Map<String, Object> variables) {
AIAgentConfigDefault effectiveDefault =
defaultValue != null ? defaultValue : AIAgentConfigDefault.disabled();
return (AIAgentConfig) evaluate(key, context, effectiveDefault, Mode.AGENT, variables);
return (AIAgentConfig) evaluate(key, context, effectiveDefault, Mode.AGENT, variables, true);
}

/**
Expand All @@ -157,14 +193,15 @@ private AIConfig evaluate(
LDContext context,
AIConfigDefault defaultValue,
Mode mode,
Map<String, Object> variables) {
Map<String, Object> variables,
boolean interpolate) {
LDValue value = client.jsonValueVariation(key, context, LDValue.ofNull());

// A valid AI Config variation is always a JSON object (it carries the _ldMeta block). When the
// flag is absent or cannot be evaluated the base SDK hands back our null sentinel; in that case
// we return the caller's typed default directly rather than serializing it and parsing it back.
if (value == null || value.getType() != LDValueType.OBJECT) {
return buildConfigFromDefault(key, mode, defaultValue, context, variables);
return buildConfigFromDefault(key, mode, defaultValue, context, variables, interpolate);
}

AIConfigFlagValue parsed = AIConfigParser.parse(value);
Expand All @@ -174,18 +211,19 @@ private AIConfig evaluate(
logger.warn(
"AI Config mode mismatch for {}: expected {}, got {}. Returning default config.",
key, mode.getWireValue(), flagMode.getWireValue());
return buildConfigFromDefault(key, mode, defaultValue, context, variables);
return buildConfigFromDefault(key, mode, defaultValue, context, variables, interpolate);
}

return buildConfig(key, mode, parsed, context, variables);
return buildConfig(key, mode, parsed, context, variables, interpolate);
}

private AIConfig buildConfig(
String key,
Mode mode,
AIConfigFlagValue parsed,
LDContext context,
Map<String, Object> variables) {
Map<String, Object> variables,
boolean interpolate) {
Supplier<LDAIConfigTracker> factory = trackerFactory(
key, parsed.getVariationKey(), parsed.getVersion(),
parsed.getModel(), parsed.getProvider(), context);
Expand All @@ -196,7 +234,8 @@ private AIConfig buildConfig(
parsed.isEnabled(),
parsed.getModel(),
parsed.getProvider(),
interpolate(parsed.getInstructions(), variables, context),
interpolate ? interpolate(parsed.getInstructions(), variables, context)
: parsed.getInstructions(),
parsed.getJudgeConfiguration(),
parsed.getTools(),
factory);
Expand All @@ -206,7 +245,8 @@ private AIConfig buildConfig(
parsed.isEnabled(),
parsed.getModel(),
parsed.getProvider(),
interpolateMessages(parsed.getMessages(), variables, context),
interpolate ? interpolateMessages(parsed.getMessages(), variables, context)
: parsed.getMessages(),
parsed.getEvaluationMetricKey(),
factory);
case COMPLETION:
Expand All @@ -216,7 +256,8 @@ private AIConfig buildConfig(
parsed.isEnabled(),
parsed.getModel(),
parsed.getProvider(),
interpolateMessages(parsed.getMessages(), variables, context),
interpolate ? interpolateMessages(parsed.getMessages(), variables, context)
: parsed.getMessages(),
parsed.getJudgeConfiguration(),
parsed.getTools(),
factory);
Expand All @@ -225,14 +266,16 @@ private AIConfig buildConfig(

/**
* Builds the typed config straight from the caller-supplied default, used when the flag is absent
* or cannot be evaluated. Prompt content is interpolated exactly as it is for an evaluated flag.
* or cannot be evaluated. Prompt content is interpolated exactly as it is for an evaluated flag,
* unless {@code interpolate} is {@code false} (template mode).
*/
private AIConfig buildConfigFromDefault(
String key,
Mode mode,
AIConfigDefault defaultValue,
LDContext context,
Map<String, Object> variables) {
Map<String, Object> variables,
boolean interpolate) {
// Default configs still get real trackers — the configKey was requested even if no flag was found.
// variationKey is null because no flag evaluation occurred.
Supplier<LDAIConfigTracker> factory = trackerFactory(key, null, null, null, null, context);
Expand All @@ -244,7 +287,8 @@ private AIConfig buildConfigFromDefault(
agent.isEnabled(),
agent.getModel(),
agent.getProvider(),
interpolate(agent.getInstructions(), variables, context),
interpolate ? interpolate(agent.getInstructions(), variables, context)
: agent.getInstructions(),
agent.getJudgeConfiguration(),
agent.getTools(),
factory);
Expand All @@ -256,7 +300,8 @@ private AIConfig buildConfigFromDefault(
judge.isEnabled(),
judge.getModel(),
judge.getProvider(),
interpolateMessages(judge.getMessages(), variables, context),
interpolate ? interpolateMessages(judge.getMessages(), variables, context)
: judge.getMessages(),
judge.getEvaluationMetricKey(),
factory);
}
Expand All @@ -268,7 +313,8 @@ private AIConfig buildConfigFromDefault(
completion.isEnabled(),
completion.getModel(),
completion.getProvider(),
interpolateMessages(completion.getMessages(), variables, context),
interpolate ? interpolateMessages(completion.getMessages(), variables, context)
: completion.getMessages(),
completion.getJudgeConfiguration(),
completion.getTools(),
factory);
Expand Down
150 changes: 150 additions & 0 deletions lib/sdk/server-ai/src/test/java/com/launchdarkly/sdk/server/ai/LDAIClientImplTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -289,6 +289,156 @@ public void agentConfigsUsageCountExcludesNullEntries() {
verify(client).trackMetric(eq("$ld:ai:usage:agent-configs"), eq(context), eq(LDValue.of(2)), eq(2.0));
}

// ---- Template config methods ----------------------------------------------

// Tracking events

@Test
public void completionConfigTemplateFiresTemplateUsageEvent() {
when(client.jsonValueVariation(anyString(), any(), any())).thenReturn(LDValue.ofNull());
ai.completionConfigTemplate("my-key", context, null);
verify(client).trackMetric(
eq("$ld:ai:usage:completion-config-template"), eq(context), eq(LDValue.of("my-key")), eq(1.0));
}

@Test
public void agentConfigTemplateFiresTemplateUsageEvent() {
when(client.jsonValueVariation(anyString(), any(), any())).thenReturn(LDValue.ofNull());
ai.agentConfigTemplate("agent-key", context, null);
verify(client).trackMetric(
eq("$ld:ai:usage:agent-config-template"), eq(context), eq(LDValue.of("agent-key")), eq(1.0));
}

@Test
public void judgeConfigTemplateFiresTemplateUsageEvent() {
when(client.jsonValueVariation(anyString(), any(), any())).thenReturn(LDValue.ofNull());
ai.judgeConfigTemplate("judge-key", context, null);
verify(client).trackMetric(
eq("$ld:ai:usage:judge-config-template"), eq(context), eq(LDValue.of("judge-key")), eq(1.0));
}

// Placeholder preservation

@Test
public void completionConfigTemplatePreservesPlaceholders() {
String json = "{\"_ldMeta\":{\"enabled\":true,\"mode\":\"completion\"},"
+ "\"messages\":[{\"role\":\"system\",\"content\":\"Hello {{name}}\"}]}";
when(client.jsonValueVariation(anyString(), any(), any())).thenReturn(LDValue.parse(json));

AICompletionConfig config = ai.completionConfigTemplate("key", context, null);
assertThat(config.getMessages().get(0).getContent(), is("Hello {{name}}"));
}

@Test
public void agentConfigTemplatePreservesPlaceholders() {
String json = "{\"_ldMeta\":{\"enabled\":true,\"mode\":\"agent\"},"
+ "\"instructions\":\"You research {{topic}}\"}";
when(client.jsonValueVariation(anyString(), any(), any())).thenReturn(LDValue.parse(json));

AIAgentConfig config = ai.agentConfigTemplate("key", context, null);
assertThat(config.getInstructions(), is("You research {{topic}}"));
}

@Test
public void judgeConfigTemplatePreservesPlaceholders() {
String json = "{\"_ldMeta\":{\"enabled\":true,\"mode\":\"judge\"},"
+ "\"messages\":[{\"role\":\"user\",\"content\":\"Rate {{response}}\"}]}";
when(client.jsonValueVariation(anyString(), any(), any())).thenReturn(LDValue.parse(json));

AIJudgeConfig config = ai.judgeConfigTemplate("key", context, null);
assertThat(config.getMessages().get(0).getContent(), is("Rate {{response}}"));
}

// ldctx non-interpolation

@Test
public void completionConfigTemplateDoesNotInterpolateLdctx() {
String json = "{\"_ldMeta\":{\"enabled\":true,\"mode\":\"completion\"},"
+ "\"messages\":[{\"role\":\"user\",\"content\":\"{{ldctx.key}}\"}]}";
when(client.jsonValueVariation(anyString(), any(), any())).thenReturn(LDValue.parse(json));

AICompletionConfig config = ai.completionConfigTemplate("key", LDContext.create("ctx-123"), null);
assertThat(config.getMessages().get(0).getContent(), is("{{ldctx.key}}"));
}

@Test
public void agentConfigTemplateDoesNotInterpolateLdctx() {
String json = "{\"_ldMeta\":{\"enabled\":true,\"mode\":\"agent\"},"
+ "\"instructions\":\"Hello {{ldctx.key}}\"}";
when(client.jsonValueVariation(anyString(), any(), any())).thenReturn(LDValue.parse(json));

AIAgentConfig config = ai.agentConfigTemplate("key", LDContext.create("ctx-123"), null);
assertThat(config.getInstructions(), is("Hello {{ldctx.key}}"));
}

@Test
public void judgeConfigTemplateDoesNotInterpolateLdctx() {
String json = "{\"_ldMeta\":{\"enabled\":true,\"mode\":\"judge\"},"
+ "\"messages\":[{\"role\":\"user\",\"content\":\"{{ldctx.key}}\"}]}";
when(client.jsonValueVariation(anyString(), any(), any())).thenReturn(LDValue.parse(json));

AIJudgeConfig config = ai.judgeConfigTemplate("key", LDContext.create("ctx-123"), null);
assertThat(config.getMessages().get(0).getContent(), is("{{ldctx.key}}"));
}

// Null default yields disabled config

@Test
public void completionConfigTemplateNullDefaultYieldsDisabled() {
when(client.jsonValueVariation(anyString(), any(), any()))
.thenAnswer(inv -> inv.getArgument(2, LDValue.class));
AICompletionConfig config = ai.completionConfigTemplate("key", context, null);
assertThat(config.isEnabled(), is(false));
}

@Test
public void agentConfigTemplateNullDefaultYieldsDisabled() {
when(client.jsonValueVariation(anyString(), any(), any()))
.thenAnswer(inv -> inv.getArgument(2, LDValue.class));
AIAgentConfig config = ai.agentConfigTemplate("key", context, null);
assertThat(config.isEnabled(), is(false));
}

@Test
public void judgeConfigTemplateNullDefaultYieldsDisabled() {
when(client.jsonValueVariation(anyString(), any(), any()))
.thenAnswer(inv -> inv.getArgument(2, LDValue.class));
AIJudgeConfig config = ai.judgeConfigTemplate("key", context, null);
assertThat(config.isEnabled(), is(false));
}

// Tracker non-null

@Test
public void completionConfigTemplateHasTracker() {
String json = "{\"_ldMeta\":{\"enabled\":true,\"mode\":\"completion\"},"
+ "\"messages\":[{\"role\":\"system\",\"content\":\"Hello {{name}}\"}]}";
when(client.jsonValueVariation(anyString(), any(), any())).thenReturn(LDValue.parse(json));

AICompletionConfig config = ai.completionConfigTemplate("key", context, null);
assertThat(config.createTracker(), is(notNullValue()));
}

@Test
public void agentConfigTemplateHasTracker() {
String json = "{\"_ldMeta\":{\"enabled\":true,\"mode\":\"agent\"},"
+ "\"instructions\":\"You research {{topic}}\"}";
when(client.jsonValueVariation(anyString(), any(), any())).thenReturn(LDValue.parse(json));

AIAgentConfig config = ai.agentConfigTemplate("key", context, null);
assertThat(config.createTracker(), is(notNullValue()));
}

@Test
public void judgeConfigTemplateHasTracker() {
String json = "{\"_ldMeta\":{\"enabled\":true,\"mode\":\"judge\"},"
+ "\"messages\":[{\"role\":\"user\",\"content\":\"Rate {{response}}\"}]}";
when(client.jsonValueVariation(anyString(), any(), any())).thenReturn(LDValue.parse(json));

AIJudgeConfig config = ai.judgeConfigTemplate("key", context, null);
assertThat(config.createTracker(), is(notNullValue()));
}

private static Map<String, Object> variables() {
return new HashMap<>();
}
Expand Down
Loading