Apply PR feedback (3)

Author: manuel-alvarez-alvarez

dd-java-agent/agent-aiguard/build.gradle

@@ -1,3 +1,5 @@
+import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar
+
 plugins {
   id 'com.gradleup.shadow'
 }
@@ -20,19 +22,15 @@ dependencies {
   implementation project(':communication')
 
   testImplementation project(':utils:test-utils')
-  testImplementation('org.skyscreamer:jsonassert:1.5.1')
+  testImplementation('org.skyscreamer:jsonassert:1.5.3')
+  testImplementation('com.fasterxml.jackson.core:jackson-databind:2.20.0')
 }
 
-shadowJar {
+tasks.named("shadowJar", ShadowJar) {
   dependencies deps.excludeShared
 }
 
-jar {
+tasks.named("jar", Jar) {
   archiveClassifier = 'unbundled'
 }
 
-spotless {
-  java {
-    target 'src/**/*.java'
-  }
-}

dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java

@@ -1,8 +1,12 @@
 package com.datadog.aiguard;
 
+import static java.util.Collections.singletonMap;
+
+import com.squareup.moshi.JsonAdapter;
 import com.squareup.moshi.JsonReader;
 import com.squareup.moshi.JsonWriter;
 import com.squareup.moshi.Moshi;
+import com.squareup.moshi.Types;
 import datadog.communication.http.OkHttpUtils;
 import datadog.trace.api.Config;
 import datadog.trace.api.aiguard.AIGuard;
@@ -20,11 +24,13 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import java.io.IOException;
+import java.lang.annotation.Annotation;
+import java.lang.reflect.Type;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import okhttp3.HttpUrl;
@@ -36,6 +42,12 @@
 import okhttp3.ResponseBody;
 import okio.BufferedSink;
 
+/**
+ * Concrete implementation of the SDK used to interact with the AIGuard REST API.
+ *
+ * <p>An instance of this class is initialized and configured automatically during agent startup
+ * through {@link AIGuardSystem#start()}.
+ */
 public class AIGuardInternal implements Evaluator {
 
   public static class BadConfigurationException extends RuntimeException {
@@ -87,7 +99,7 @@ static void uninstall() {
     this.url = url;
     this.headers = headers;
     this.client = client;
-    this.moshi = new Moshi.Builder().build();
+    this.moshi = new Moshi.Builder().add(new AIGuardFactory()).build();
     final Config config = Config.get();
     this.meta = mapOf("service", config.getServiceName(), "env", config.getEnv());
   }
@@ -126,21 +138,20 @@ private static String getToolName(final Message current, final List<Message> mes
           .map(ToolCall::getFunction)
           .map(Function::getName)
           .collect(Collectors.joining(","));
-    } else {
-      // assistant message with tool output (search the linked tool call in reverse order)
-      final String id = current.getToolCallId();
-      for (int i = messages.size() - 1; i >= 0; i--) {
-        final Message message = messages.get(i);
-        if (message.getToolCalls() != null) {
-          for (final ToolCall toolCall : message.getToolCalls()) {
-            if (toolCall.getId().equals(id)) {
-              return toolCall.getFunction() == null ? null : toolCall.getFunction().getName();
-            }
+    }
+    // assistant message with tool output (search the linked tool call in reverse order)
+    final String id = current.getToolCallId();
+    for (int i = messages.size() - 1; i >= 0; i--) {
+      final Message message = messages.get(i);
+      if (message.getToolCalls() != null) {
+        for (final ToolCall toolCall : message.getToolCalls()) {
+          if (toolCall.getId().equals(id)) {
+            return toolCall.getFunction() == null ? null : toolCall.getFunction().getName();
           }
         }
       }
-      return null;
     }
+    return null;
   }
 
   private boolean isBlockingEnabled(final Object isBlockingEnabled) {
@@ -155,18 +166,17 @@ public Evaluation evaluate(final List<Message> messages, final Options options)
     final AgentTracer.TracerAPI tracer = AgentTracer.get();
     final AgentSpan span = tracer.buildSpan(SPAN_NAME, SPAN_NAME).start();
     try (final AgentScope scope = tracer.activateSpan(span)) {
-      final Message current = messages.get(messages.size() - 1);
-      if (isToolCall(current)) {
+      final Message last = messages.get(messages.size() - 1);
+      if (isToolCall(last)) {
         span.setTag(TARGET_TAG, "tool");
-        final String toolName = getToolName(current, messages);
+        final String toolName = getToolName(last, messages);
         if (toolName != null) {
           span.setTag(TOOL_TAG, toolName);
         }
       } else {
         span.setTag(TARGET_TAG, "prompt");
       }
-      final Map<String, Object> metaStruct =
-          Collections.singletonMap(META_STRUCT_KEY, truncate(messages));
+      final Map<String, Object> metaStruct = singletonMap(META_STRUCT_KEY, truncate(messages));
       span.setMetaStruct(META_STRUCT_TAG, metaStruct);
       final Request.Builder request =
           new Request.Builder()
@@ -243,6 +253,69 @@ public static void install(final Evaluator evaluator) {
     }
   }
 
+  static class AIGuardFactory implements JsonAdapter.Factory {
+
+    @Nullable
+    @Override
+    public JsonAdapter<?> create(
+        final Type type, final Set<? extends Annotation> annotations, final Moshi moshi) {
+      final Class<?> rawType = Types.getRawType(type);
+      if (rawType != AIGuard.Message.class) {
+        return null;
+      }
+      return new MessageAdapter(moshi.adapter(AIGuard.ToolCall.class));
+    }
+  }
+
+  static class MessageAdapter extends JsonAdapter<Message> {
+
+    private final JsonAdapter<AIGuard.ToolCall> toolCallAdapter;
+
+    MessageAdapter(final JsonAdapter<ToolCall> toolCallAdapter) {
+      this.toolCallAdapter = toolCallAdapter;
+    }
+
+    @Nullable
+    @Override
+    public Message fromJson(JsonReader reader) throws IOException {
+      throw new UnsupportedOperationException("Serializing only adapter");
+    }
+
+    @Override
+    public void toJson(final JsonWriter writer, @Nullable final Message value) throws IOException {
+      if (value == null) {
+        writer.nullValue();
+        return;
+      }
+      writer.beginObject();
+      writeValue(writer, "role", value.getRole());
+      writeValue(writer, "content", value.getContent());
+      writeArray(writer, "tool_calls", value.getToolCalls());
+      writeValue(writer, "tool_call_id", value.getToolCallId());
+      writer.endObject();
+    }
+
+    private void writeValue(final JsonWriter writer, final String name, final Object value)
+        throws IOException {
+      if (value != null) {
+        writer.name(name);
+        writer.jsonValue(value);
+      }
+    }
+
+    private void writeArray(final JsonWriter writer, final String name, final List<ToolCall> value)
+        throws IOException {
+      if (value != null) {
+        writer.name(name);
+        writer.beginArray();
+        for (final ToolCall toolCall : value) {
+          toolCallAdapter.toJson(writer, toolCall);
+        }
+        writer.endArray();
+      }
+    }
+  }
+
   static class MoshiJsonRequestBody extends RequestBody {
 
     private static final MediaType JSON = MediaType.parse("application/json");

dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy

@@ -1,5 +1,8 @@
 package com.datadog.aiguard
 
+import com.fasterxml.jackson.annotation.JsonInclude
+import com.fasterxml.jackson.databind.ObjectMapper
+import com.fasterxml.jackson.databind.PropertyNamingStrategies
 import com.squareup.moshi.Moshi
 import datadog.trace.api.Config
 import datadog.trace.api.aiguard.AIGuard
@@ -40,6 +43,13 @@ class AIGuardInternalTests extends DDSpecification {
   @Shared
   protected static final MOSHI = new Moshi.Builder().build()
 
+  @Shared
+  protected static final MAPPER = new ObjectMapper()
+  .setPropertyNamingStrategy(PropertyNamingStrategies.SNAKE_CASE)
+  .setDefaultPropertyInclusion(
+  JsonInclude.Value.construct(JsonInclude.Include.NON_NULL, JsonInclude.Include.NON_NULL)
+  )
+
   @Shared
   protected static final TOOL_CALL = [
     AIGuard.Message.message('system', 'You are a beautiful AI assistant'),
@@ -423,11 +433,15 @@ class AIGuardInternalTests extends DDSpecification {
     }
     assert request.body().contentType().toString().contains('application/json')
     final receivedBody = readRequestBody(request.body())
-    final expectedBody = MOSHI.adapter(Map).toJson([data: [attributes: [messages: messages, meta: [service: 'ai_guard_test', env: 'test']]]])
-    JSONAssert.assertEquals(expectedBody, receivedBody, JSONCompareMode.LENIENT)
+    final expectedBody = snakeCaseJson([data: [attributes: [messages: messages, meta: [service: 'ai_guard_test', env: 'test']]]])
+    JSONAssert.assertEquals(expectedBody, receivedBody, JSONCompareMode.NON_EXTENSIBLE)
     return true
   }
 
+  private static String snakeCaseJson(final Object value) {
+    MAPPER.writeValueAsString(value)
+  }
+
   private static String readRequestBody(final RequestBody body) {
     final output = new ByteArrayOutputStream()
     final buffer = Okio.buffer(Okio.sink(output))

dd-trace-api/build.gradle.kts

@@ -25,8 +25,8 @@ val excludedClassesCoverage by extra(
     "datadog.trace.api.internal.TraceSegment",
     "datadog.trace.api.internal.TraceSegment.NoOp",
     "datadog.trace.api.aiguard.AIGuard",
-    "datadog.trace.api.aiguard.AIGuard.AIGuardClientError",
     "datadog.trace.api.aiguard.AIGuard.AIGuardAbortError",
+    "datadog.trace.api.aiguard.AIGuard.AIGuardClientError",
     "datadog.trace.api.aiguard.AIGuard.Options",
     "datadog.trace.api.civisibility.CIVisibility",
     "datadog.trace.api.civisibility.DDTestModule",

dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java

@@ -191,8 +191,8 @@ public static class Message {
 
     private final String role;
     private final String content;
-    private final List<ToolCall> tool_calls;
-    private final String tool_call_id;
+    private final List<ToolCall> toolCalls;
+    private final String toolCallId;
 
     /**
      * Creates a new message with the specified parameters.
@@ -211,8 +211,8 @@ public Message(
         final String toolCallId) {
       this.role = role;
       this.content = content;
-      this.tool_calls = toolCalls;
-      this.tool_call_id = toolCallId;
+      this.toolCalls = toolCalls;
+      this.toolCallId = toolCallId;
     }
 
     /**
@@ -239,7 +239,7 @@ public String getContent() {
      * @return list of tool calls, or null if this message has no tool calls
      */
     public List<ToolCall> getToolCalls() {
-      return tool_calls;
+      return toolCalls;
     }
 
     /**
@@ -248,7 +248,7 @@ public List<ToolCall> getToolCalls() {
      * @return the tool call ID, or null if this is not a tool response message
      */
     public String getToolCallId() {
-      return tool_call_id;
+      return toolCallId;
     }
 
     /**