Improve message writer tests

Author: manuel-alvarez-alvarez

communication/src/main/java/datadog/communication/serialization/Codec.java

@@ -1,33 +1,36 @@
 package datadog.communication.serialization;
 
+import datadog.communication.serialization.custom.aiguard.FunctionWriter;
 import datadog.communication.serialization.custom.aiguard.MessageWriter;
 import datadog.communication.serialization.custom.aiguard.ToolCallWriter;
 import datadog.communication.serialization.custom.stacktrace.StackTraceEventFrameWriter;
 import datadog.communication.serialization.custom.stacktrace.StackTraceEventWriter;
+import datadog.trace.api.Config;
 import datadog.trace.api.aiguard.AIGuard;
 import datadog.trace.util.stacktrace.StackTraceEvent;
 import datadog.trace.util.stacktrace.StackTraceFrame;
 import java.nio.ByteBuffer;
 import java.nio.CharBuffer;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
 
 public final class Codec extends ClassValue<ValueWriter<?>> {
 
-  private static final Map<Class<?>, ValueWriter<?>> defaultConfig =
-      Stream.of(
-              new Object[][] {
-                {StackTraceEvent.class, new StackTraceEventWriter()},
-                {StackTraceFrame.class, new StackTraceEventFrameWriter()},
-                {AIGuard.Message.class, new MessageWriter()},
-                {AIGuard.ToolCall.class, new ToolCallWriter()},
-              })
-          .collect(Collectors.toMap(data -> (Class<?>) data[0], data -> (ValueWriter<?>) data[1]));
-
-  public static final Codec INSTANCE = new Codec(defaultConfig);
+  public static final Codec INSTANCE;
+
+  static {
+    final Map<Class<?>, ValueWriter<?>> writers = new HashMap<>(1 << 3);
+    writers.put(StackTraceEvent.class, new StackTraceEventWriter());
+    writers.put(StackTraceFrame.class, new StackTraceEventFrameWriter());
+    if (Config.get().isAiGuardEnabled()) {
+      writers.put(AIGuard.Message.class, new MessageWriter());
+      writers.put(AIGuard.ToolCall.class, new ToolCallWriter());
+      writers.put(AIGuard.ToolCall.Function.class, new FunctionWriter());
+    }
+    INSTANCE = new Codec(writers);
+  }
 
   private final Map<Class<?>, ValueWriter<?>> config;
 

communication/src/main/java/datadog/communication/serialization/custom/aiguard/FunctionWriter.java

@@ -0,0 +1,21 @@
+package datadog.communication.serialization.custom.aiguard;
+
+import datadog.communication.serialization.EncodingCache;
+import datadog.communication.serialization.ValueWriter;
+import datadog.communication.serialization.Writable;
+import datadog.trace.api.aiguard.AIGuard;
+
+public class FunctionWriter implements ValueWriter<AIGuard.ToolCall.Function> {
+
+  @Override
+  public void write(
+      final AIGuard.ToolCall.Function function,
+      final Writable writable,
+      final EncodingCache encodingCache) {
+    writable.startMap(2);
+    writable.writeString("name", encodingCache);
+    writable.writeString(function.getName(), encodingCache);
+    writable.writeString("arguments", encodingCache);
+    writable.writeString(function.getArguments(), encodingCache);
+  }
+}

communication/src/main/java/datadog/communication/serialization/custom/aiguard/MessageWriter.java

@@ -44,10 +44,7 @@ private static void writeToolCallArray(
       final EncodingCache encodingCache) {
     if (present) {
       writable.writeString(key, encodingCache);
-      writable.startArray(values.size());
-      for (final AIGuard.ToolCall toolCall : values) {
-        writable.writeObject(toolCall, encodingCache);
-      }
+      writable.writeObject(values, encodingCache);
     }
   }
 

communication/src/main/java/datadog/communication/serialization/custom/aiguard/ToolCallWriter.java

@@ -14,14 +14,6 @@ public void write(
     writable.writeString("id", encodingCache);
     writable.writeString(value.getId(), encodingCache);
     writable.writeString("function", encodingCache);
-
-    final AIGuard.ToolCall.Function function = value.getFunction();
-    if (function != null) {
-      writable.startMap(2);
-      writable.writeString("name", encodingCache);
-      writable.writeString(function.getName(), encodingCache);
-      writable.writeString("arguments", encodingCache);
-      writable.writeString(function.getArguments(), encodingCache);
-    }
+    writable.writeObject(value.getFunction(), encodingCache);
   }
 }

communication/src/test/groovy/datadog/communication/serialization/aiguard/MessageWriterTest.groovy

@@ -0,0 +1,119 @@
+package datadog.communication.serialization.aiguard
+
+import datadog.communication.serialization.EncodingCache
+import datadog.communication.serialization.GrowableBuffer
+import datadog.communication.serialization.msgpack.MsgPackWriter
+import datadog.trace.api.aiguard.AIGuard
+import datadog.trace.test.util.DDSpecification
+import org.msgpack.core.MessagePack
+import org.msgpack.value.Value
+
+import java.nio.charset.StandardCharsets
+import java.util.function.Function
+
+class MessageWriterTest extends DDSpecification {
+
+  private EncodingCache encodingCache
+  private GrowableBuffer buffer
+  private MsgPackWriter writer
+
+  void setup() {
+    injectSysConfig('ai_guard.enabled', 'true')
+    final HashMap<CharSequence, byte[]> cache = new HashMap<>()
+    encodingCache = new EncodingCache() {
+      @Override
+      byte[] encode(CharSequence chars) {
+        cache.computeIfAbsent(chars, s -> s.toString().getBytes(StandardCharsets.UTF_8))
+      }
+    }
+    buffer = new GrowableBuffer(1024)
+    writer = new MsgPackWriter(buffer)
+  }
+
+  void 'test write message'() {
+    given:
+    final message = AIGuard.Message.message('user', 'What day is today?')
+
+    when:
+    writer.writeObject(message, encodingCache)
+
+    then:
+    try (final unpacker = MessagePack.newDefaultUnpacker(buffer.slice())) {
+      final value = asStringValueMap(unpacker.unpackValue())
+      value.size() == 2
+      value.role == 'user'
+      value.content == 'What day is today?'
+    }
+  }
+
+  void 'test write tool call'()  {
+    given:
+    final message =
+    AIGuard.Message.assistant(
+    AIGuard.ToolCall.toolCall('call_1', 'function_1', 'args_1'),
+    AIGuard.ToolCall.toolCall('call_2', 'function_2', 'args_2'))
+
+    when:
+    writer.writeObject(message, encodingCache)
+
+    then:
+    try (final unpacker = MessagePack.newDefaultUnpacker(buffer.slice())) {
+      final value = asStringKeyMap(unpacker.unpackValue())
+      value.size() == 2
+      asString(value.role) == 'assistant'
+
+      final toolCalls = value.get('tool_calls').asArrayValue().list()
+      toolCalls.size() == 2
+
+      final firstCall = asStringKeyMap(toolCalls[0])
+      asString(firstCall.id) == 'call_1'
+      final firstFunction = asStringValueMap(firstCall.function)
+      firstFunction.name == 'function_1'
+      firstFunction.arguments == 'args_1'
+
+      final secondCall = asStringKeyMap(toolCalls[1])
+      asString(secondCall.id) == 'call_2'
+      final secondFunction = asStringValueMap(secondCall.function)
+      secondFunction.name == 'function_2'
+      secondFunction.arguments == 'args_2'
+    }
+  }
+
+  void 'test write tool output'() throws IOException {
+    given:
+    final message = AIGuard.Message.tool('call_1', 'output')
+
+    when:
+    writer.writeObject(message, encodingCache)
+
+    then:
+    try (final unpacker = MessagePack.newDefaultUnpacker(buffer.slice())) {
+      final value = asStringValueMap(unpacker.unpackValue())
+      value.size() == 3
+      value.role == 'tool'
+      value.tool_call_id == 'call_1'
+      value.content == 'output'
+    }
+  }
+
+  private static <K, V> Map<K, V> mapValue(
+  final Value values,
+  final Function<Value, K> keyMapper,
+  final Function<Value, V> valueMapper) {
+    return values.asMapValue().entrySet().collectEntries {
+      [(keyMapper.apply(it.key)):  valueMapper.apply(it.value)]
+    }
+  }
+
+  private static Map<String, Value> asStringKeyMap(final Value values) {
+    return mapValue(values, MessageWriterTest::asString, Function.identity())
+  }
+
+  private static Map<String, String> asStringValueMap(final Value values) {
+    return mapValue(values, MessageWriterTest::asString, MessageWriterTest::asString)
+  }
+
+  private static String asString(final Value value) {
+    return value.asStringValue().asString()
+  }
+}