Merge pull request #1837 from newrelic/aws-bedrock-llm-support

Aws bedrock llm support

agent-bridge/src/main/java/com/newrelic/agent/bridge/NoOpAgent.java

@@ -7,6 +7,7 @@
 
 package com.newrelic.agent.bridge;
 
+import com.newrelic.api.agent.AiMonitoring;
 import com.newrelic.api.agent.Config;
 import com.newrelic.api.agent.ErrorApi;
 import com.newrelic.api.agent.Insights;
@@ -18,6 +19,8 @@
 import java.util.Collections;
 import java.util.Map;
 
+import static com.newrelic.agent.bridge.NoOpAiMonitoring.INSTANCE;
+
 class NoOpAgent implements Agent {
 
     static final Agent INSTANCE = new NoOpAgent();
@@ -65,6 +68,11 @@ public Insights getInsights() {
         return NoOpInsights.INSTANCE;
     }
 
+    @Override
+    public AiMonitoring getAiMonitoring() {
+        return NoOpAiMonitoring.INSTANCE;
+    }
+
     @Override
     public ErrorApi getErrorApi() {
         return NoOpErrorApi.INSTANCE;

agent-bridge/src/main/java/com/newrelic/agent/bridge/NoOpAiMonitoring.java

@@ -0,0 +1,21 @@
+package com.newrelic.agent.bridge;
+
+import com.newrelic.api.agent.AiMonitoring;
+import com.newrelic.api.agent.LlmTokenCountCallback;
+
+import java.util.Map;
+
+public class NoOpAiMonitoring implements AiMonitoring {
+
+    static final AiMonitoring INSTANCE = new NoOpAiMonitoring();
+
+    private NoOpAiMonitoring() {}
+
+    @Override
+    public void recordLlmFeedbackEvent(Map<String, Object> llmFeedbackEventAttributes) {
+    }
+
+    @Override
+    public void setLlmTokenCountCallback(LlmTokenCountCallback llmTokenCountCallback) {
+    }
+}

agent-bridge/src/main/java/com/newrelic/agent/bridge/NoOpTransaction.java

@@ -27,6 +27,7 @@
 public class NoOpTransaction implements Transaction {
     public static final Transaction INSTANCE = new NoOpTransaction();
     public static final NoOpMap<String, Object> AGENT_ATTRIBUTES = new NoOpMap<>();
+    public static final NoOpMap<String, Object> USER_ATTRIBUTES = new NoOpMap<>();
 
     @Override
     public void beforeSendResponseHeaders() {
@@ -153,6 +154,11 @@ public Map<String, Object> getAgentAttributes() {
         return AGENT_ATTRIBUTES;
     }
 
+    @Override
+    public Map<String, Object> getUserAttributes() {
+        return USER_ATTRIBUTES;
+    }
+
     @Override
     public boolean registerAsyncActivity(Object activityContext) {
         return false;

agent-bridge/src/main/java/com/newrelic/agent/bridge/Transaction.java

@@ -23,6 +23,7 @@
 public interface Transaction extends com.newrelic.api.agent.Transaction {
 
     Map<String, Object> getAgentAttributes();
+    Map<String, Object> getUserAttributes();
 
     /**
      * Sets the current transaction's name.

agent-bridge/src/main/java/com/newrelic/agent/bridge/aimonitoring/AiMonitoringUtils.java

@@ -0,0 +1,78 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.newrelic.agent.bridge.aimonitoring;
+
+import com.newrelic.api.agent.Config;
+import com.newrelic.api.agent.NewRelic;
+
+import java.util.logging.Level;
+
+public class AiMonitoringUtils {
+    // Enabled defaults
+    private static final boolean AI_MONITORING_ENABLED_DEFAULT = false;
+    private static final boolean AI_MONITORING_STREAMING_ENABLED_DEFAULT = true;
+    private static final boolean AI_MONITORING_RECORD_CONTENT_ENABLED_DEFAULT = true;
+    private static final boolean HIGH_SECURITY_ENABLED_DEFAULT = false;
+
+    /**
+     * Check if ai_monitoring features are enabled.
+     * Indicates whether LLM instrumentation will be registered. If this is set to False, no metrics, events, or spans are to be sent.
+     *
+     * @return true if AI monitoring is enabled, else false
+     */
+    public static boolean isAiMonitoringEnabled() {
+        Config config = NewRelic.getAgent().getConfig();
+        Boolean aimEnabled = config.getValue("ai_monitoring.enabled", AI_MONITORING_ENABLED_DEFAULT);
+        Boolean highSecurity = config.getValue("high_security", HIGH_SECURITY_ENABLED_DEFAULT);
+
+        if (highSecurity || !aimEnabled) {
+            aimEnabled = false;
+            String disabledReason = highSecurity ? "High Security Mode." : "agent config.";
+            NewRelic.getAgent().getLogger().log(Level.FINE, "AIM: AI Monitoring is disabled due to " + disabledReason);
+            NewRelic.incrementCounter("Supportability/Java/ML/Disabled");
+        } else {
+            NewRelic.incrementCounter("Supportability/Java/ML/Enabled");
+        }
+
+        return aimEnabled;
+    }
+
+    /**
+     * Check if ai_monitoring.streaming features are enabled.
+     *
+     * @return true if streaming is enabled, else false
+     */
+    public static boolean isAiMonitoringStreamingEnabled() {
+        Boolean enabled = NewRelic.getAgent().getConfig().getValue("ai_monitoring.streaming.enabled", AI_MONITORING_STREAMING_ENABLED_DEFAULT);
+
+        if (enabled) {
+            NewRelic.incrementCounter("Supportability/Java/ML/Streaming/Enabled");
+        } else {
+            NewRelic.incrementCounter("Supportability/Java/ML/Streaming/Disabled");
+        }
+
+        return enabled;
+    }
+
+    /**
+     * Check if the input and output content should be added to LLM events.
+     *
+     * @return true if adding content is enabled, else false
+     */
+    public static boolean isAiMonitoringRecordContentEnabled() {
+        Boolean enabled = NewRelic.getAgent().getConfig().getValue("ai_monitoring.record_content.enabled", AI_MONITORING_RECORD_CONTENT_ENABLED_DEFAULT);
+
+        if (enabled) {
+            NewRelic.incrementCounter("Supportability/Java/ML/RecordContent/Enabled");
+        } else {
+            NewRelic.incrementCounter("Supportability/Java/ML/RecordContent/Disabled");
+        }
+
+        return enabled;
+    }
+}

agent-bridge/src/main/java/com/newrelic/agent/bridge/aimonitoring/LlmTokenCountCallbackHolder.java

@@ -0,0 +1,30 @@
+package com.newrelic.agent.bridge.aimonitoring;
+
+import com.newrelic.api.agent.LlmTokenCountCallback;
+
+/**
+ * A thread-safe holder for an instance of {@link LlmTokenCountCallback}.
+ * This class provides methods for setting and retrieving the callback instance.
+ */
+public class LlmTokenCountCallbackHolder {
+
+    private static volatile LlmTokenCountCallback llmTokenCountCallback = null;
+
+    /**
+     * Sets the {@link LlmTokenCountCallback} instance to be stored.
+     *
+     * @param newLlmTokenCountCallback the callback instance
+     */
+    public static void setLlmTokenCountCallback(LlmTokenCountCallback newLlmTokenCountCallback) {
+        llmTokenCountCallback = newLlmTokenCountCallback;
+    }
+
+    /**
+     * Retrieves the stored {@link LlmTokenCountCallback} instance.
+     *
+     * @return stored callback instance
+     */
+    public static LlmTokenCountCallback getLlmTokenCountCallback() {
+        return llmTokenCountCallback;
+    }
+}

agent-model/src/main/java/com/newrelic/agent/model/LlmCustomInsightsEvent.java

@@ -0,0 +1,30 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.newrelic.agent.model;
+
+/**
+ * Represents an internal subtype of a CustomInsightsEvent that is sent to the
+ * custom_event_data collector endpoint but potentially subject to different
+ * validation rules and agent configuration.
+ */
+public class LlmCustomInsightsEvent {
+    // LLM event types
+    private static final String LLM_EMBEDDING = "LlmEmbedding";
+    private static final String LLM_CHAT_COMPLETION_SUMMARY = "LlmChatCompletionSummary";
+    private static final String LLM_CHAT_COMPLETION_MESSAGE = "LlmChatCompletionMessage";
+
+    /**
+     * Determines if a CustomInsightsEvent should be treated as a LlmEvent
+     *
+     * @param eventType type of the current event
+     * @return true if eventType is an LlmEvent, else false
+     */
+    public static boolean isLlmEvent(String eventType) {
+        return eventType.equals(LLM_EMBEDDING) || eventType.equals(LLM_CHAT_COMPLETION_MESSAGE) || eventType.equals(LLM_CHAT_COMPLETION_SUMMARY);
+    }
+}

functional_test/src/test/java/com/newrelic/agent/extension/FakeExtensionAgent.java

@@ -10,6 +10,7 @@
 import com.newrelic.agent.bridge.Agent;
 import com.newrelic.agent.bridge.TracedMethod;
 import com.newrelic.agent.bridge.Transaction;
+import com.newrelic.api.agent.AiMonitoring;
 import com.newrelic.api.agent.Config;
 import com.newrelic.api.agent.ErrorApi;
 import com.newrelic.api.agent.Insights;
@@ -38,6 +39,11 @@ public Logger getLogger() {
     @Override
     public Insights getInsights() { throw new RuntimeException(); }
 
+    @Override
+    public AiMonitoring getAiMonitoring() {
+        return null;
+    }
+
     @Override
     public ErrorApi getErrorApi() { throw new RuntimeException(); }