Skip to content

feat(ai-guard): enable redaction tests in java #7645

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

Draft
manuel-alvarez-alvarez wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

feat(ai-guard): enable redaction tests in java #7645

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
50 changes: 40 additions & 10 deletions manifests/java.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,16 +33,46 @@ manifest:
- weblog_declaration:
"*": irrelevant (just one weblog is enough to test the SDK)
"spring-boot": v1.57.0
tests/ai_guard/test_ai_guard_sdk.py::Test_NoRedaction: missing_feature
tests/ai_guard/test_ai_guard_sdk.py::Test_RedactedMessagesInSDKResponse: missing_feature
tests/ai_guard/test_ai_guard_sdk.py::Test_Redaction: missing_feature
tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionDisabled: missing_feature
tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionDisabledTelemetry: missing_feature
tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionFailSafe: missing_feature
tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionInSDKResponse: missing_feature
tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionMultiTurnContext: missing_feature
tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionOnBlock: missing_feature
tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionSkipsStructuralFields: missing_feature
tests/ai_guard/test_ai_guard_sdk.py::Test_NoRedaction:
- weblog_declaration:
"*": irrelevant (just one weblog is enough to test the SDK)
"spring-boot": v1.66.0-SNAPSHOT
tests/ai_guard/test_ai_guard_sdk.py::Test_RedactedMessagesInSDKResponse:
- weblog_declaration:
"*": irrelevant (just one weblog is enough to test the SDK)
"spring-boot": v1.66.0-SNAPSHOT
tests/ai_guard/test_ai_guard_sdk.py::Test_Redaction:
- weblog_declaration:
"*": irrelevant (just one weblog is enough to test the SDK)
"spring-boot": v1.66.0-SNAPSHOT
tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionDisabled:
- weblog_declaration:
"*": irrelevant (just one weblog is enough to test the SDK)
"spring-boot": v1.66.0-SNAPSHOT
tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionDisabledTelemetry:
- weblog_declaration:
"*": irrelevant (just one weblog is enough to test the SDK)
"spring-boot": v1.66.0-SNAPSHOT
tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionFailSafe:
- weblog_declaration:
"*": irrelevant (just one weblog is enough to test the SDK)
"spring-boot": v1.66.0-SNAPSHOT
tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionInSDKResponse:
- weblog_declaration:
"*": irrelevant (just one weblog is enough to test the SDK)
"spring-boot": v1.66.0-SNAPSHOT
tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionMultiTurnContext:
- weblog_declaration:
"*": irrelevant (just one weblog is enough to test the SDK)
"spring-boot": v1.66.0-SNAPSHOT
tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionOnBlock:
- weblog_declaration:
"*": irrelevant (just one weblog is enough to test the SDK)
"spring-boot": v1.66.0-SNAPSHOT
tests/ai_guard/test_ai_guard_sdk.py::Test_RedactionSkipsStructuralFields:
- weblog_declaration:
"*": irrelevant (just one weblog is enough to test the SDK)
"spring-boot": v1.66.0-SNAPSHOT
tests/ai_guard/test_ai_guard_sdk.py::Test_RootSpanUserKeep:
- weblog_declaration:
"*": irrelevant (just one weblog is enough to test the SDK)
Expand Down
118 changes: 100 additions & 18 deletions ...-boot/src/main/java/com/datadoghq/system_tests/springboot/ai_guard/AIGuardController.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,10 @@

import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.JsonSerializer;
import com.fasterxml.jackson.databind.SerializerProvider;
import datadog.trace.api.aiguard.AIGuard;
import datadog.trace.api.aiguard.AIGuard.Evaluation;
import datadog.trace.api.interceptor.MutableSpan;
Expand All @@ -18,11 +21,9 @@
import org.springframework.web.bind.annotation.RequestHeader;
import org.springframework.web.bind.annotation.RestController;

import java.lang.reflect.Method;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;


Expand All @@ -35,7 +36,8 @@ public static class JacksonConfig {
public Jackson2ObjectMapperBuilderCustomizer mixInCustomizer() {
return builder -> builder
.mixIn(AIGuard.AIGuardAbortError.class, AIGuardAbortErrorMixIn.class)
.mixIn(AIGuard.Evaluation.class, AIGuardEvaluationMixIn.class);
.mixIn(AIGuard.Evaluation.class, AIGuardEvaluationMixIn.class)
.serializerByType(AIGuard.Message.class, new MessageSerializer());
}
}

Expand Down Expand Up @@ -131,15 +133,13 @@ public void setName(String name) {
}

public AIGuard.Message toAIGuard() {
if (toolCallId != null) {
String contentStr = content != null && content.isTextual() ? content.asText() : null;
return AIGuard.Message.tool(toolCallId, contentStr);
}
if (toolCalls != null && !toolCalls.isEmpty()) {
return AIGuard.Message.assistant(
toolCalls.stream().map(ToolCall::toAIGuard).toArray(AIGuard.ToolCall[]::new));
}
// Handle content parts vs string content
// Every field is carried over independently: an assistant message can hold content
// *and* tool calls at the same time, and dropping either half would change the
// conversation the tracer sends to the AI Guard service.
List<AIGuard.ToolCall> calls = toolCalls == null || toolCalls.isEmpty()
? null
: toolCalls.stream().map(ToolCall::toAIGuard).collect(Collectors.toList());

if (content != null && content.isArray()) {
// Content parts format
List<AIGuard.ContentPart> parts = new ArrayList<>();
Expand All @@ -152,12 +152,11 @@ public AIGuard.Message toAIGuard() {
parts.add(AIGuard.ContentPart.imageUrl(url));
}
}
return AIGuard.Message.message(role, parts);
} else {
// String content format
String contentStr = content != null && content.isTextual() ? content.asText() : null;
return AIGuard.Message.message(role, contentStr);
return new AIGuard.Message(role, parts, calls, toolCallId);
}
// String content format
String contentStr = content != null && content.isTextual() ? content.asText() : null;
return new AIGuard.Message(role, contentStr, calls, toolCallId);
}
}

Expand Down Expand Up @@ -234,6 +233,89 @@ public static abstract class AIGuardEvaluationMixIn {

@JsonProperty("tag_probs")
abstract Object getTagProbabilities();

@JsonProperty("redaction_replacements")
abstract Object getRedactionReplacements();
}

/**
* Writes {@link AIGuard.Message} back in the very shape the request carried it.
*
* <p>Bean serialization would emit the SDK's field names and a null for every field the
* message does not use, which no other tracer's weblog does and which makes the response
* impossible to compare against the messages that were sent. Content parts live under
* {@code content} as an array, exactly like the OpenAI wire format the SDK models.
*/
public static class MessageSerializer extends JsonSerializer<AIGuard.Message> {

@Override
public void serialize(final AIGuard.Message message,
final JsonGenerator gen,
final SerializerProvider serializers) throws IOException {
gen.writeStartObject();
if (message.getRole() != null) {
gen.writeStringField("role", message.getRole());
}
if (message.getContentParts() != null) {
gen.writeFieldName("content");
gen.writeStartArray();
for (final AIGuard.ContentPart part : message.getContentParts()) {
writeContentPart(part, gen);
}
gen.writeEndArray();
} else if (message.getContent() != null) {
// Written on nullness, not emptiness: "" is the redaction remove strategy.
gen.writeStringField("content", message.getContent());
}
if (message.getToolCalls() != null) {
gen.writeFieldName("tool_calls");
gen.writeStartArray();
for (final AIGuard.ToolCall toolCall : message.getToolCalls()) {
writeToolCall(toolCall, gen);
}
gen.writeEndArray();
}
if (message.getToolCallId() != null) {
gen.writeStringField("tool_call_id", message.getToolCallId());
}
gen.writeEndObject();
}

private static void writeContentPart(final AIGuard.ContentPart part,
final JsonGenerator gen) throws IOException {
gen.writeStartObject();
gen.writeStringField("type", part.getType().toString());
if (part.getType() == AIGuard.ContentPart.Type.TEXT) {
gen.writeStringField("text", part.getText());
} else if (part.getType() == AIGuard.ContentPart.Type.IMAGE_URL) {
gen.writeFieldName("image_url");
gen.writeStartObject();
gen.writeStringField("url", part.getImageUrl().getUrl());
gen.writeEndObject();
}
gen.writeEndObject();
}

private static void writeToolCall(final AIGuard.ToolCall toolCall,
final JsonGenerator gen) throws IOException {
gen.writeStartObject();
if (toolCall.getId() != null) {
gen.writeStringField("id", toolCall.getId());
}
final AIGuard.ToolCall.Function function = toolCall.getFunction();
if (function != null) {
gen.writeFieldName("function");
gen.writeStartObject();
if (function.getName() != null) {
gen.writeStringField("name", function.getName());
}
if (function.getArguments() != null) {
gen.writeStringField("arguments", function.getArguments());
}
gen.writeEndObject();
}
gen.writeEndObject();
}
}

}
Loading