Apply PR feedback

Author: manuel-alvarez-alvarez

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

@@ -31,14 +31,6 @@ jar {
   archiveClassifier = 'unbundled'
 }
 
-ext {
-  minimumBranchCoverage = 0.6
-  minimumInstructionCoverage = 0.8
-  excludedClassesCoverage = []
-  excludedClassesBranchCoverage = []
-  excludedClassesInstructionCoverage = []
-}
-
 spotless {
   java {
     target 'src/**/*.java'

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

@@ -21,6 +21,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import java.io.IOException;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -37,6 +38,12 @@
 
 public class AIGuardInternal implements Evaluator {
 
+  public static class BadConfigurationException extends RuntimeException {
+    public BadConfigurationException(final String message) {
+      super(message);
+    }
+  }
+
   static final String SPAN_NAME = "ai_guard";
   static final String TARGET_TAG = "ai_guard.target";
   static final String TOOL_TAG = "ai_guard.tool";
@@ -51,7 +58,7 @@ public static void install() {
     final String apiKey = config.getApiKey();
     final String appKey = config.getApplicationKey();
     if (isEmpty(apiKey) || isEmpty(appKey)) {
-      throw new RuntimeException(
+      throw new BadConfigurationException(
           "AI Guard: Missing api and/or application key, use DD_API_KEY and DD_APP_KEY");
     }
     String endpoint = config.getAiGuardEndpoint();
@@ -98,11 +105,12 @@ private static List<Message> truncate(List<Message> messages) {
     final int maxContent = config.getAiGuardMaxContentSize();
     for (int i = 0; i < messages.size(); i++) {
       Message source = messages.get(i);
-      if (source.getContent() != null && source.getContent().length() > maxContent) {
+      final String content = source.getContent();
+      if (content != null && content.length() > maxContent) {
         source =
             new Message(
                 source.getRole(),
-                source.getContent().substring(0, maxContent),
+                content.substring(0, maxContent),
                 source.getToolCalls(),
                 source.getToolCallId());
         messages.set(i, source);
@@ -146,7 +154,7 @@ private boolean isBlockingEnabled(final Object isBlockingEnabled) {
   @Override
   public Evaluation evaluate(final List<Message> messages, final Options options) {
     if (messages == null || messages.isEmpty()) {
-      throw new IllegalArgumentException("messages must not be empty");
+      throw new IllegalArgumentException("Messages must not be empty");
     }
     final AgentTracer.TracerAPI tracer = AgentTracer.get();
     final AgentSpan span = tracer.buildSpan(SPAN_NAME, SPAN_NAME).start();
@@ -161,8 +169,8 @@ public Evaluation evaluate(final List<Message> messages, final Options options)
       } else {
         span.setTag(TARGET_TAG, "prompt");
       }
-      final Map<String, Object> metaStruct = new HashMap<>(1);
-      metaStruct.put(META_STRUCT_KEY, truncate(messages));
+      final Map<String, Object> metaStruct =
+          Collections.singletonMap(META_STRUCT_KEY, truncate(messages));
       span.setMetaStruct(META_STRUCT_TAG, metaStruct);
       final Request.Builder request =
           new Request.Builder()
@@ -173,7 +181,7 @@ public Evaluation evaluate(final List<Message> messages, final Options options)
         final Map<String, Object> result = parseResponseBody(response);
         final String actionStr = (String) result.get("action");
         if (actionStr == null) {
-          throw new IllegalArgumentException("action field is missing in the response");
+          throw new IllegalArgumentException("Action field is missing in the response");
         }
         final Action action = Action.valueOf(actionStr);
         final String reason = (String) result.get("reason");
@@ -214,7 +222,7 @@ private Map<String, Object> parseResponseBody(final Response response) throws IO
   }
 
   private AIGuardClientError fail(final int statusCode, final Object errors) {
-    return new AIGuardClientError("AI Guard service call failed, status" + statusCode, errors);
+    return new AIGuardClientError("AI Guard service call failed, status: " + statusCode, errors);
   }
 
   private static OkHttpClient buildClient(final HttpUrl url, final long timeout) {

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

@@ -77,11 +77,28 @@ class AIGuardInternalTests extends DDSpecification {
   }
 
   void 'test missing api/app keys'() {
+    given:
+    if (apiKey) {
+      injectEnvConfig('API_KEY', apiKey)
+    }
+    if (appKey) {
+      injectEnvConfig('APP_KEY', appKey)
+    }
+
     when:
     AIGuardInternal.install()
 
     then:
-    thrown(RuntimeException)
+    thrown(AIGuardInternal.BadConfigurationException)
+
+    where:
+    apiKey   | appKey
+    'apiKey' | null
+    'apiKey' | ''
+    null     | 'appKey'
+    ''       | 'appKey'
+    null     | null
+    ''       | ''
   }
 
   void 'test endpoint discovery'() {
@@ -269,6 +286,30 @@ class AIGuardInternalTests extends DDSpecification {
     1 * span.addThrowable(_ as AIGuard.AIGuardClientError)
   }
 
+  void 'test evaluate with empty response'() {
+    given:
+    Request request = null
+    final call = Mock(Call) {
+      execute() >> {
+        return mockResponse(request, 200, null)
+      }
+    }
+    final client = Mock(OkHttpClient) {
+      newCall(_ as Request) >> {
+        request = (Request) it[0]
+        return call
+      }
+    }
+    final aiguard = new AIGuardInternal(URL, HEADERS, client)
+
+    when:
+    aiguard.evaluate(TOOL_CALL, AIGuard.Options.DEFAULT)
+
+    then:
+    thrown(AIGuard.AIGuardClientError)
+    1 * span.addThrowable(_ as AIGuard.AIGuardClientError)
+  }
+
   void 'test message length truncation'() {
     given:
     final maxMessages = Config.get().getAiGuardMaxMessagesLength()
@@ -331,6 +372,49 @@ class AIGuardInternalTests extends DDSpecification {
     }
   }
 
+  void 'test no messages'() {
+    given:
+    final aiguard = new AIGuardInternal(URL, HEADERS, Stub(OkHttpClient))
+
+    when:
+    aiguard.evaluate(messages, AIGuard.Options.DEFAULT)
+
+    then:
+    thrown(IllegalArgumentException)
+
+
+    where:
+    messages << [[], null]
+  }
+
+  void 'test missing tool name'() {
+    given:
+    def request
+    final call = Mock(Call) {
+      execute() >> {
+        return mockResponse(
+          request,
+          200,
+          [data: [attributes: [action: 'ALLOW', reason: 'Just do it']]]
+          )
+      }
+    }
+    final client = Mock(OkHttpClient) {
+      newCall(_ as Request) >> {
+        request = (Request) it[0]
+        return call
+      }
+    }
+    final aiguard = new AIGuardInternal(URL, HEADERS, client)
+
+    when:
+    aiguard.evaluate([AIGuard.Message.tool('call_1', 'Content')], AIGuard.Options.DEFAULT)
+
+    then:
+    1 * span.setTag(AIGuardInternal.TARGET_TAG, 'tool')
+    0 * span.setTag(AIGuardInternal.TOOL_TAG, _)
+  }
+
   private static assertRequest(final Request request, final List<AIGuard.Message> messages) {
     assert request.url() == URL
     assert request.method() == 'POST'
@@ -358,7 +442,7 @@ class AIGuardInternalTests extends DDSpecification {
       .message('ok')
       .request(request)
       .code(status)
-      .body(ResponseBody.create(MediaType.parse('application/json'), MOSHI.adapter(Object).toJson(body)))
+      .body(body == null ? null : ResponseBody.create(MediaType.parse('application/json'), MOSHI.adapter(Object).toJson(body)))
       .build()
   }