diff --git a/communication/src/main/java/datadog/communication/serialization/custom/aiguard/MessageWriter.java b/communication/src/main/java/datadog/communication/serialization/custom/aiguard/MessageWriter.java index a1b089fbb1c..b057409e726 100644 --- a/communication/src/main/java/datadog/communication/serialization/custom/aiguard/MessageWriter.java +++ b/communication/src/main/java/datadog/communication/serialization/custom/aiguard/MessageWriter.java @@ -14,16 +14,51 @@ public void write( final AIGuard.Message value, final Writable writable, final EncodingCache encodingCache) { final int[] size = {0}; final boolean hasRole = isNotBlank(value.getRole(), size); - final boolean hasContent = isNotBlank(value.getContent(), size); final boolean hasToolCallId = isNotBlank(value.getToolCallId(), size); final boolean hasToolCalls = isNotEmpty(value.getToolCalls(), size); + + final boolean hasContentParts = isNotEmpty(value.getContentParts(), size); + final boolean hasContent = !hasContentParts && isNotBlank(value.getContent(), size); + writable.startMap(size[0]); writeString(hasRole, "role", value.getRole(), writable, encodingCache); - writeString(hasContent, "content", value.getContent(), writable, encodingCache); + + if (hasContentParts) { + writeContentParts("content", value.getContentParts(), writable, encodingCache); + } else { + writeString(hasContent, "content", value.getContent(), writable, encodingCache); + } + writeString(hasToolCallId, "tool_call_id", value.getToolCallId(), writable, encodingCache); writeToolCallArray(hasToolCalls, "tool_calls", value.getToolCalls(), writable, encodingCache); } + private static void writeContentParts( + final String key, + final List contentParts, + final Writable writable, + final EncodingCache encodingCache) { + writable.writeString(key, encodingCache); + writable.startArray(contentParts.size()); + + for (final AIGuard.ContentPart part : contentParts) { + writable.startMap(2); + + writable.writeString("type", encodingCache); + writable.writeString(part.getType().toString(), encodingCache); + + if (part.getType() == AIGuard.ContentPart.Type.TEXT) { + writable.writeString("text", encodingCache); + writable.writeString(part.getText(), encodingCache); + } else if (part.getType() == AIGuard.ContentPart.Type.IMAGE_URL) { + writable.writeString("image_url", encodingCache); + writable.startMap(1); + writable.writeString("url", encodingCache); + writable.writeString(part.getImageUrl().getUrl(), encodingCache); + } + } + } + private static void writeString( final boolean present, final String key, diff --git a/communication/src/test/groovy/datadog/communication/serialization/aiguard/MessageWriterTest.groovy b/communication/src/test/groovy/datadog/communication/serialization/aiguard/MessageWriterTest.groovy index 3328a2330ad..45191a2fc8f 100644 --- a/communication/src/test/groovy/datadog/communication/serialization/aiguard/MessageWriterTest.groovy +++ b/communication/src/test/groovy/datadog/communication/serialization/aiguard/MessageWriterTest.groovy @@ -116,4 +116,127 @@ class MessageWriterTest extends DDSpecification { private static String asString(final Value value) { return value.asStringValue().asString() } + + void 'test write message with text content parts'() { + given: + final message = AIGuard.Message.message('user', [ + AIGuard.ContentPart.text('Hello world') + ]) + + when: + writer.writeObject(message, encodingCache) + + then: + try (final unpacker = MessagePack.newDefaultUnpacker(buffer.slice())) { + final value = asStringKeyMap(unpacker.unpackValue()) + value.size() == 2 + asString(value.role) == 'user' + + final contentParts = value.content.asArrayValue().list() + contentParts.size() == 1 + + final part = asStringKeyMap(contentParts[0]) + asString(part.type) == 'text' + asString(part.text) == 'Hello world' + } + } + + void 'test write message with image_url content parts'() { + given: + final message = AIGuard.Message.message('user', [ + AIGuard.ContentPart.imageUrl('https://example.com/image.jpg') + ]) + + when: + writer.writeObject(message, encodingCache) + + then: + try (final unpacker = MessagePack.newDefaultUnpacker(buffer.slice())) { + final value = asStringKeyMap(unpacker.unpackValue()) + value.size() == 2 + asString(value.role) == 'user' + + final contentParts = value.content.asArrayValue().list() + contentParts.size() == 1 + + final part = asStringKeyMap(contentParts[0]) + asString(part.type) == 'image_url' + + final imageUrl = asStringKeyMap(part.image_url) + asString(imageUrl.url) == 'https://example.com/image.jpg' + } + } + + void 'test write message with mixed content parts'() { + given: + final message = AIGuard.Message.message('user', [ + AIGuard.ContentPart.text('Describe this:'), + AIGuard.ContentPart.imageUrl('https://example.com/image.jpg'), + AIGuard.ContentPart.text('What is it?') + ]) + + when: + writer.writeObject(message, encodingCache) + + then: + try (final unpacker = MessagePack.newDefaultUnpacker(buffer.slice())) { + final value = asStringKeyMap(unpacker.unpackValue()) + value.size() == 2 + asString(value.role) == 'user' + + final contentParts = value.content.asArrayValue().list() + contentParts.size() == 3 + + final part1 = asStringKeyMap(contentParts[0]) + asString(part1.type) == 'text' + asString(part1.text) == 'Describe this:' + + final part2 = asStringKeyMap(contentParts[1]) + asString(part2.type) == 'image_url' + final imageUrl = asStringKeyMap(part2.image_url) + asString(imageUrl.url) == 'https://example.com/image.jpg' + + final part3 = asStringKeyMap(contentParts[2]) + asString(part3.type) == 'text' + asString(part3.text) == 'What is it?' + } + } + + void 'test content parts type serializes as string not integer'() { + given: + final message = AIGuard.Message.message('user', [ + AIGuard.ContentPart.text('Test') + ]) + + when: + writer.writeObject(message, encodingCache) + + then: + try (final unpacker = MessagePack.newDefaultUnpacker(buffer.slice())) { + final value = asStringKeyMap(unpacker.unpackValue()) + final contentParts = value.content.asArrayValue().list() + final part = asStringKeyMap(contentParts[0]) + + // Verify type is a string value, not an integer + part.type.isStringValue() + !part.type.isIntegerValue() + asString(part.type) == 'text' + } + } + + void 'test backward compatibility with string content'() { + given: + final message = AIGuard.Message.message('user', 'Plain text message') + + 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 == 'Plain text message' + } + } } diff --git a/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java b/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java index cae66ce7e10..2ee18ba2b8a 100644 --- a/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java +++ b/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java @@ -16,6 +16,7 @@ import datadog.trace.api.aiguard.AIGuard.AIGuardAbortError; import datadog.trace.api.aiguard.AIGuard.AIGuardClientError; import datadog.trace.api.aiguard.AIGuard.Action; +import datadog.trace.api.aiguard.AIGuard.ContentPart; import datadog.trace.api.aiguard.AIGuard.Evaluation; import datadog.trace.api.aiguard.AIGuard.Message; import datadog.trace.api.aiguard.AIGuard.Options; @@ -136,16 +137,44 @@ private static List messagesForMetaStruct(List messages) { boolean contentTruncated = false; for (int i = messages.size() - size; i < messages.size(); i++) { final Message source = messages.get(i); - String content = source.getContent(); - if (content != null && content.length() > maxContent) { - contentTruncated = true; - content = content.substring(0, maxContent); - } - List toolCalls = source.getToolCalls(); - if (toolCalls != null) { - toolCalls = new ArrayList<>(toolCalls); + + List contentParts = source.getContentParts(); + if (contentParts != null) { + final List truncatedParts = new ArrayList<>(contentParts.size()); + for (final ContentPart part : contentParts) { + if (part.getType() == ContentPart.Type.TEXT) { + String text = part.getText(); + if (text != null && text.length() > maxContent) { + contentTruncated = true; + text = text.substring(0, maxContent); + truncatedParts.add(ContentPart.text(text)); + } else { + truncatedParts.add(part); + } + } else { + truncatedParts.add(part); + } + } + + List toolCalls = source.getToolCalls(); + if (toolCalls != null) { + toolCalls = new ArrayList<>(toolCalls); + } + result.add( + new Message(source.getRole(), truncatedParts, toolCalls, source.getToolCallId())); + } else { + // Handle plain text content (backward compatibility) + String content = source.getContent(); + if (content != null && content.length() > maxContent) { + contentTruncated = true; + content = content.substring(0, maxContent); + } + List toolCalls = source.getToolCalls(); + if (toolCalls != null) { + toolCalls = new ArrayList<>(toolCalls); + } + result.add(new Message(source.getRole(), content, toolCalls, source.getToolCallId())); } - result.add(new Message(source.getRole(), content, toolCalls, source.getToolCallId())); } if (contentTruncated) { WafMetricCollector.get().aiGuardTruncated(CONTENT); @@ -333,12 +362,45 @@ public Message fromJson(JsonReader reader) throws IOException { public void toJson(final JsonWriter writer, final Message value) throws IOException { writer.beginObject(); writeValue(writer, "role", value.getRole()); - writeValue(writer, "content", value.getContent()); + + if (value.getContentParts() != null) { + writeContentParts(writer, "content", value.getContentParts()); + } else { + writeValue(writer, "content", value.getContent()); + } + writeArray(writer, "tool_calls", value.getToolCalls()); writeValue(writer, "tool_call_id", value.getToolCallId()); writer.endObject(); } + private void writeContentParts( + final JsonWriter writer, final String name, final List contentParts) + throws IOException { + writer.name(name); + writer.beginArray(); + for (final ContentPart part : contentParts) { + writer.beginObject(); + + writer.name("type"); + writer.value(part.getType().toString()); + + if (part.getType() == ContentPart.Type.TEXT) { + writer.name("text"); + writer.value(part.getText()); + } else if (part.getType() == ContentPart.Type.IMAGE_URL) { + writer.name("image_url"); + writer.beginObject(); + writer.name("url"); + writer.value(part.getImageUrl().getUrl()); + writer.endObject(); + } + + writer.endObject(); + } + writer.endArray(); + } + private void writeValue(final JsonWriter writer, final String name, final Object value) throws IOException { if (value != null) { diff --git a/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy b/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy index f7c9de745ca..e3a5ffd6671 100644 --- a/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy +++ b/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy @@ -365,7 +365,7 @@ class AIGuardInternalTests extends DDSpecification { final messages = [ new AIGuard.Message( "assistant", - null, + (String) null, [AIGuard.ToolCall.toolCall('call_1', 'execute_shell', '{"cmd": "ls -lah"}')], null ) @@ -470,6 +470,180 @@ class AIGuardInternalTests extends DDSpecification { .build() } + void 'test JSON serialization with text content parts'() { + given: + final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Good']]]) + final messages = [AIGuard.Message.message('user', [AIGuard.ContentPart.text('Hello world')])] + + when: + aiguard.evaluate(messages, AIGuard.Options.DEFAULT) + + then: + 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { + final meta = it[1] as Map + final receivedMessages = meta.messages as List + assert receivedMessages.size() == 1 + assert receivedMessages[0].contentParts.size() == 1 + assert receivedMessages[0].contentParts[0].type == AIGuard.ContentPart.Type.TEXT + assert receivedMessages[0].contentParts[0].text == 'Hello world' + return span + } + } + + void 'test JSON serialization with image_url content parts'() { + given: + final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Good']]]) + final messages = [ + AIGuard.Message.message('user', [AIGuard.ContentPart.imageUrl('https://example.com/image.jpg')]) + ] + + when: + aiguard.evaluate(messages, AIGuard.Options.DEFAULT) + + then: + 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { + final meta = it[1] as Map + final receivedMessages = meta.messages as List + assert receivedMessages.size() == 1 + assert receivedMessages[0].contentParts.size() == 1 + assert receivedMessages[0].contentParts[0].type == AIGuard.ContentPart.Type.IMAGE_URL + assert receivedMessages[0].contentParts[0].imageUrl.url == 'https://example.com/image.jpg' + return span + } + } + + void 'test JSON serialization with mixed content parts'() { + given: + final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Good']]]) + final messages = [ + AIGuard.Message.message('user', [ + AIGuard.ContentPart.text('Describe this image:'), + AIGuard.ContentPart.imageUrl('https://example.com/image.jpg'), + AIGuard.ContentPart.text('What do you see?') + ]) + ] + + when: + aiguard.evaluate(messages, AIGuard.Options.DEFAULT) + + then: + 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { + final meta = it[1] as Map + final receivedMessages = meta.messages as List + assert receivedMessages.size() == 1 + assert receivedMessages[0].contentParts.size() == 3 + assert receivedMessages[0].contentParts[0].type == AIGuard.ContentPart.Type.TEXT + assert receivedMessages[0].contentParts[0].text == 'Describe this image:' + assert receivedMessages[0].contentParts[1].type == AIGuard.ContentPart.Type.IMAGE_URL + assert receivedMessages[0].contentParts[1].imageUrl.url == 'https://example.com/image.jpg' + assert receivedMessages[0].contentParts[2].type == AIGuard.ContentPart.Type.TEXT + assert receivedMessages[0].contentParts[2].text == 'What do you see?' + return span + } + } + + void 'test content parts order is preserved'() { + given: + final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Good']]]) + final parts = (0..4).collect { + it % 2 == 0 ? AIGuard.ContentPart.text("Text $it") : AIGuard.ContentPart.imageUrl("https://example.com/image${it}.jpg") + } + final messages = [AIGuard.Message.message('user', parts)] + + when: + aiguard.evaluate(messages, AIGuard.Options.DEFAULT) + + then: + 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { + final meta = it[1] as Map + final receivedMessages = meta.messages as List + assert receivedMessages[0].contentParts.size() == 5 + (0..4).each { i -> + if (i % 2 == 0) { + assert receivedMessages[0].contentParts[i].type == AIGuard.ContentPart.Type.TEXT + assert receivedMessages[0].contentParts[i].text == "Text $i" + } else { + assert receivedMessages[0].contentParts[i].type == AIGuard.ContentPart.Type.IMAGE_URL + assert receivedMessages[0].contentParts[i].imageUrl.url == "https://example.com/image${i}.jpg" + } + } + return span + } + } + + void 'test content part text truncation'() { + given: + final maxContent = Config.get().getAiGuardMaxContentSize() + final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Good']]]) + final longText = (0..maxContent).collect { 'A' }.join() + final messages = [ + AIGuard.Message.message('user', [AIGuard.ContentPart.text(longText), AIGuard.ContentPart.text('Short text')]) + ] + + when: + aiguard.evaluate(messages, AIGuard.Options.DEFAULT) + + then: + 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { + final meta = it[1] as Map + final receivedMessages = meta.messages as List + assert receivedMessages[0].contentParts.size() == 2 + assert receivedMessages[0].contentParts[0].text.length() == maxContent + assert receivedMessages[0].contentParts[0].text.length() < longText.length() + assert receivedMessages[0].contentParts[1].text == 'Short text' + return span + } + assertTelemetry('ai_guard.truncated', 'type:content') + } + + void 'test content part image_url not truncated even with long data URI'() { + given: + final maxContent = Config.get().getAiGuardMaxContentSize() + final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Good']]]) + // Create a very long data URI (longer than max content size) + final longDataUri = 'data:image/png;base64,' + (0..(maxContent + 1000)).collect { 'A' }.join() + final messages = [ + AIGuard.Message.message('user', [ + AIGuard.ContentPart.text('Image:'), + AIGuard.ContentPart.imageUrl(longDataUri) + ]) + ] + + when: + aiguard.evaluate(messages, AIGuard.Options.DEFAULT) + + then: + 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { + final meta = it[1] as Map + final receivedMessages = meta.messages as List + assert receivedMessages[0].contentParts.size() == 2 + assert receivedMessages[0].contentParts[1].type == AIGuard.ContentPart.Type.IMAGE_URL + // Image URL should NOT be truncated + assert receivedMessages[0].contentParts[1].imageUrl.url == longDataUri + assert receivedMessages[0].contentParts[1].imageUrl.url.length() > maxContent + return span + } + } + + void 'test backward compatibility with string content'() { + given: + final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Good']]]) + final messages = [AIGuard.Message.message('user', 'Hello world')] + + when: + aiguard.evaluate(messages, AIGuard.Options.DEFAULT) + + then: + 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { + final meta = it[1] as Map + final receivedMessages = meta.messages as List + assert receivedMessages.size() == 1 + assert receivedMessages[0].content == 'Hello world' + assert receivedMessages[0].contentParts == null + return span + } + } + private static class TestSuite { private final AIGuard.Action action private final String reason diff --git a/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/AIGuardController.java b/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/AIGuardController.java index fc3cb3942be..45048aba231 100644 --- a/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/AIGuardController.java +++ b/dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/AIGuardController.java @@ -65,6 +65,21 @@ public ResponseEntity abort(final @RequestHeader("X-Blocking-Enabled") boolea } } + @GetMapping(value = "/multimodal") + public ResponseEntity multimodal() { + final Evaluation result = + AIGuard.evaluate( + asList( + Message.message("system", "You are a beautiful AI"), + Message.message( + "user", + asList( + AIGuard.ContentPart.text("Describe this image:"), + AIGuard.ContentPart.imageUrl("https://example.com/image.jpg"), + AIGuard.ContentPart.text("What do you see?"))))); + return ResponseEntity.ok(result); + } + /** Mocking endpoint for the AI Guard REST API */ @SuppressWarnings("unchecked") @PostMapping( @@ -80,18 +95,40 @@ public ResponseEntity> evaluate( final Map last = messages.get(messages.size() - 1); String action = "ALLOW"; String reason = "The prompt looks harmless"; - String content = (String) last.get("content"); - if (content.startsWith("You should not trust me")) { - action = "DENY"; - reason = "I am feeling suspicious today"; - } else if (content.startsWith("Nuke yourself")) { - action = "ABORT"; - reason = "The user is trying to destroy me"; + + // Handle both string content and content parts + Object contentObj = last.get("content"); + String content = null; + boolean hasContentParts = false; + + if (contentObj instanceof String) { + content = (String) contentObj; + } else if (contentObj instanceof List) { + // Content parts - check if it contains an image + List> contentParts = (List>) contentObj; + hasContentParts = true; + for (Map part : contentParts) { + if ("image_url".equals(part.get("type"))) { + reason = "Multimodal prompt detected"; + break; + } + } } + + if (content != null) { + if (content.startsWith("You should not trust me")) { + action = "DENY"; + reason = "I am feeling suspicious today"; + } else if (content.startsWith("Nuke yourself")) { + action = "ABORT"; + reason = "The user is trying to destroy me"; + } + } + final Map evaluation = new HashMap<>(3); evaluation.put("action", action); evaluation.put("reason", reason); - evaluation.put("is_blocking_enabled", content.endsWith("[block]")); + evaluation.put("is_blocking_enabled", content != null && content.endsWith("[block]")); return ResponseEntity.ok() .body(Collections.singletonMap("data", Collections.singletonMap("attributes", evaluation))); } diff --git a/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/AIGuardSmokeTest.groovy b/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/AIGuardSmokeTest.groovy index 5c83d554269..973c96324ec 100644 --- a/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/AIGuardSmokeTest.groovy +++ b/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/AIGuardSmokeTest.groovy @@ -102,4 +102,62 @@ class AIGuardSmokeTest extends AbstractAppSecServerSmokeTest { } } } + + void 'test multimodal content parts evaluation'() { + given: + def request = new Request.Builder() + .url("http://localhost:${httpPort}/aiguard/multimodal") + .get() + .build() + + when: + final response = client.newCall(request).execute() + + then: + assert response.code() == 200 + final body = new JsonSlurper().parse(response.body().bytes()) + assert body.reason == 'Multimodal prompt detected' + assert body.action == 'ALLOW' + + and: + waitForTraceCount(2) // default call + internal API mock + final span = traces*.spans + ?.flatten() + ?.find { + it.resource == 'ai_guard' + } as DecodedSpan + assert span.meta.get('ai_guard.action') == 'ALLOW' + assert span.meta.get('ai_guard.reason') == 'Multimodal prompt detected' + assert span.meta.get('ai_guard.target') == 'prompt' + + // Verify content parts in metaStruct + final messages = span.metaStruct.get('ai_guard').messages as List> + assert messages.size() == 2 + with(messages[0]) { + role == 'system' + content == 'You are a beautiful AI' + } + with(messages[1]) { + role == 'user' + def contentParts = it.content as List> + assert contentParts != null + assert contentParts.size() == 3 + + with(contentParts[0]) { + type == 'text' + text == 'Describe this image:' + } + + with(contentParts[1]) { + type == 'image_url' + def imageUrl = it.image_url as Map + assert imageUrl.url == 'https://example.com/image.jpg' + } + + with(contentParts[2]) { + type == 'text' + text == 'What do you see?' + } + } + } } diff --git a/dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java b/dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java index 1a8c0a89dd7..ac76d349213 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java @@ -2,7 +2,12 @@ import datadog.trace.api.aiguard.noop.NoOpEvaluator; import java.util.Arrays; +import java.util.Collections; import java.util.List; +import java.util.Locale; +import java.util.Objects; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; /** * SDK for calling the AIGuard REST API to evaluate AI prompts and tool calls for security threats. @@ -186,16 +191,176 @@ public List getTags() { } } + /** + * Represents an image URL in a content part. Images can be provided as URLs or data URIs. + * + *

Example usage: + * + *

{@code
+   * // Image from URL
+   * var imageUrl = new AIGuard.ImageURL("https://example.com/image.jpg");
+   *
+   * // Image from data URI
+   * var dataUri = new AIGuard.ImageURL("data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAUA...");
+   * }
+ */ + public static class ImageURL { + + private final String url; + + /** + * Creates a new ImageURL with the specified URL. + * + * @param url the image URL or data URI + * @throws NullPointerException if url is null + */ + public ImageURL(@Nonnull final String url) { + this.url = Objects.requireNonNull(url, "url cannot be null"); + } + + /** + * Returns the image URL. + * + * @return the image URL or data URI + */ + @Nonnull + public String getUrl() { + return url; + } + } + + /** + * Represents a content part within a message. Content parts can be text or images, enabling + * multimodal AI prompts. + * + *

Example usage: + * + *

{@code
+   * // Text content
+   * var textPart = AIGuard.ContentPart.text("Describe this image:");
+   *
+   * // Image content from URL
+   * var imagePart = AIGuard.ContentPart.imageUrl("https://example.com/image.jpg");
+   *
+   * // Multimodal message with text and image
+   * var message = AIGuard.Message.message("user", List.of(textPart, imagePart));
+   * }
+ */ + public static class ContentPart { + + /** Type of content part. */ + public enum Type { + /** Text content */ + TEXT, + /** Image URL content */ + IMAGE_URL; + + @Override + public String toString() { + return name().toLowerCase(Locale.ROOT); + } + } + + private final Type type; + @Nullable private final String text; + @Nullable private final ImageURL imageUrl; + + /** + * Package-private constructor to enforce use of factory methods. + * + * @param type the content part type + * @param text the text content (required for TEXT type) + * @param imageUrl the image URL (required for IMAGE_URL type) + */ + ContentPart( + @Nonnull final Type type, @Nullable final String text, @Nullable final ImageURL imageUrl) { + this.type = type; + this.text = text; + this.imageUrl = imageUrl; + + if (type == Type.TEXT && text == null) { + throw new IllegalArgumentException("text content part requires text field"); + } + if (type == Type.IMAGE_URL && imageUrl == null) { + throw new IllegalArgumentException("image_url content part requires imageUrl field"); + } + } + + /** + * Returns the type of this content part. + * + * @return the content part type (TEXT or IMAGE_URL) + */ + @Nonnull + public Type getType() { + return type; + } + + /** + * Returns the text content of this part. + * + * @return the text content, or null if this is an IMAGE_URL part + */ + @Nullable + public String getText() { + return text; + } + + /** + * Returns the image URL of this part. + * + * @return the image URL, or null if this is a TEXT part + */ + @Nullable + public ImageURL getImageUrl() { + return imageUrl; + } + + /** + * Creates a text content part. + * + * @param text the text content + * @return a new ContentPart with TEXT type + * @throws NullPointerException if text is null + */ + @Nonnull + public static ContentPart text(@Nonnull final String text) { + Objects.requireNonNull(text, "text cannot be null"); + return new ContentPart(Type.TEXT, text, null); + } + + /** + * Creates an image content part from a URL string. + * + * @param url the image URL or data URI + * @return a new ContentPart with IMAGE_URL type + * @throws NullPointerException if url is null + */ + @Nonnull + public static ContentPart imageUrl(@Nonnull final String url) { + return new ContentPart(Type.IMAGE_URL, null, new ImageURL(url)); + } + } + /** * Represents a message in an AI conversation. Messages can represent user prompts, assistant * responses, system messages, or tool outputs. * + *

Messages can contain either plain text content or structured content parts (text and + * images): + * *

Example usage: * *

{@code
-   * // User prompt
+   * // User prompt with text
    * var userPrompt = AIGuard.Message.message("user", "What's the weather like?");
    *
+   * // User prompt with text and image
+   * var multimodalPrompt = AIGuard.Message.message("user", List.of(
+   *     AIGuard.ContentPart.text("Describe this image:"),
+   *     AIGuard.ContentPart.imageUrl("https://example.com/image.jpg")
+   * ));
+   *
    * // Assistant response with tool calls
    * var assistantWithTools = AIGuard.Message.assistant(
    *     AIGuard.ToolCall.toolCall("call_123", "get_weather", "{\"location\": \"New York\"}")
@@ -208,7 +373,8 @@ public List getTags() {
   public static class Message {
 
     private final String role;
-    private final String content;
+    @Nullable private final String content;
+    @Nullable private final List contentParts;
     private final List toolCalls;
     private final String toolCallId;
 
@@ -223,12 +389,41 @@ public static class Message {
      *     response
      */
     public Message(
-        final String role,
-        final String content,
-        final List toolCalls,
-        final String toolCallId) {
+        @Nonnull final String role,
+        @Nullable final String content,
+        @Nullable final List toolCalls,
+        @Nullable final String toolCallId) {
       this.role = role;
       this.content = content;
+      this.contentParts = null;
+      this.toolCalls = toolCalls;
+      this.toolCallId = toolCallId;
+    }
+
+    /**
+     * Creates a new message with content parts (text and/or images).
+     *
+     * @param role the role of the message sender
+     * @param contentParts list of content parts
+     * @param toolCalls list of tool calls, or null
+     * @param toolCallId the tool call ID this message responds to, or null
+     * @throws IllegalArgumentException if contentParts contains null elements
+     */
+    public Message(
+        @Nonnull final String role,
+        @Nonnull final List contentParts,
+        @Nullable final List toolCalls,
+        @Nullable final String toolCallId) {
+      this.role = role;
+      this.content = null;
+
+      for (final ContentPart part : contentParts) {
+        if (part == null) {
+          throw new IllegalArgumentException("contentParts must not contain null elements");
+        }
+      }
+
+      this.contentParts = Collections.unmodifiableList(contentParts);
       this.toolCalls = toolCalls;
       this.toolCallId = toolCallId;
     }
@@ -245,12 +440,24 @@ public String getRole() {
     /**
      * Returns the text content of the message.
      *
-     * @return the message content, or null for assistant messages with only tool calls
+     * @return the message content, or null if using content parts or for assistant messages with
+     *     only tool calls
      */
+    @Nullable
     public String getContent() {
       return content;
     }
 
+    /**
+     * Returns the content parts of the message.
+     *
+     * @return the content parts (text and images), or null if using plain text content
+     */
+    @Nullable
+    public List getContentParts() {
+      return contentParts;
+    }
+
     /**
      * Returns the list of tool calls associated with this message.
      *
@@ -276,10 +483,24 @@ public String getToolCallId() {
      * @param content the text content of the message
      * @return a new Message instance
      */
-    public static Message message(final String role, final String content) {
+    @Nonnull
+    public static Message message(@Nonnull final String role, @Nonnull final String content) {
       return new Message(role, content, null, null);
     }
 
+    /**
+     * Creates a message with specified role and content parts (text and/or images).
+     *
+     * @param role the role of the message sender (e.g., "user", "system")
+     * @param contentParts list of content parts (text and/or images)
+     * @return a new Message instance
+     */
+    @Nonnull
+    public static Message message(
+        @Nonnull final String role, @Nonnull final List contentParts) {
+      return new Message(role, contentParts, null, null);
+    }
+
     /**
      * Creates a tool response message.
      *
@@ -297,8 +518,9 @@ public static Message tool(final String toolCallId, final String content) {
      * @param toolCalls the tool calls the assistant wants to make
      * @return a new Message instance with role "assistant" and no text content
      */
-    public static Message assistant(final ToolCall... toolCalls) {
-      return new Message("assistant", null, Arrays.asList(toolCalls), null);
+    @Nonnull
+    public static Message assistant(@Nonnull final ToolCall... toolCalls) {
+      return new Message("assistant", (String) null, Arrays.asList(toolCalls), null);
     }
   }
 
diff --git a/dd-trace-api/src/test/groovy/datadog/trace/api/aiguard/AIGuardTest.groovy b/dd-trace-api/src/test/groovy/datadog/trace/api/aiguard/AIGuardTest.groovy
index f987cb307b0..6e9673c074b 100644
--- a/dd-trace-api/src/test/groovy/datadog/trace/api/aiguard/AIGuardTest.groovy
+++ b/dd-trace-api/src/test/groovy/datadog/trace/api/aiguard/AIGuardTest.groovy
@@ -69,4 +69,109 @@ class AIGuardTest extends Specification {
     eval.action == ALLOW
     eval.reason == 'AI Guard is not enabled'
   }
+
+  void 'test ImageURL creation'() {
+    when:
+    final imageUrl = new AIGuard.ImageURL('https://example.com/image.jpg')
+
+    then:
+    imageUrl.url == 'https://example.com/image.jpg'
+  }
+
+  void 'test ContentPart.text factory'() {
+    when:
+    final part = AIGuard.ContentPart.text('Hello world')
+
+    then:
+    part.type == AIGuard.ContentPart.Type.TEXT
+    part.text == 'Hello world'
+    part.imageUrl == null
+  }
+
+  void 'test ContentPart.imageUrl from String factory'() {
+    when:
+    final part = AIGuard.ContentPart.imageUrl('https://example.com/image.jpg')
+
+    then:
+    part.type == AIGuard.ContentPart.Type.IMAGE_URL
+    part.text == null
+    part.imageUrl != null
+    part.imageUrl.url == 'https://example.com/image.jpg'
+  }
+
+  void 'test Message with contentParts'() {
+    when:
+    final message = AIGuard.Message.message('user', [
+      AIGuard.ContentPart.text('Describe this image:'),
+      AIGuard.ContentPart.imageUrl('https://example.com/image.jpg')
+    ])
+
+    then:
+    message.role == 'user'
+    message.content == null
+    message.contentParts != null
+    message.contentParts.size() == 2
+    message.contentParts[0].type == AIGuard.ContentPart.Type.TEXT
+    message.contentParts[0].text == 'Describe this image:'
+    message.contentParts[1].type == AIGuard.ContentPart.Type.IMAGE_URL
+    message.contentParts[1].imageUrl.url == 'https://example.com/image.jpg'
+  }
+
+  void 'test Message with plain content returns null contentParts'() {
+    when:
+    final message = AIGuard.Message.message('user', 'Hello')
+
+    then:
+    message.content == 'Hello'
+    message.contentParts == null
+  }
+
+  void 'test Message with contentParts returns null content'() {
+    when:
+    final message = AIGuard.Message.message('user', [AIGuard.ContentPart.text('Hello')])
+
+    then:
+    message.content == null
+    message.contentParts != null
+  }
+
+
+
+  void 'test Message validation allows null content for assistant with tool calls'() {
+    when:
+    final message = AIGuard.Message.assistant(
+      AIGuard.ToolCall.toolCall('1', 'test', '{}')
+      )
+
+    then:
+    message.role == 'assistant'
+    message.content == null
+    message.contentParts == null
+    message.toolCalls != null
+  }
+
+  void 'test Message allows empty contentParts list'() {
+    when:
+    def message = new AIGuard.Message('user', [], null, null)
+
+    then:
+    message.contentParts != null
+    message.contentParts.isEmpty()
+  }
+
+  void 'test ContentPart validation fails for TEXT without text'() {
+    when:
+    new AIGuard.ContentPart(AIGuard.ContentPart.Type.TEXT, null, null)
+
+    then:
+    thrown(IllegalArgumentException)
+  }
+
+  void 'test ContentPart validation fails for IMAGE_URL without imageUrl'() {
+    when:
+    new AIGuard.ContentPart(AIGuard.ContentPart.Type.IMAGE_URL, null, null)
+
+    then:
+    thrown(IllegalArgumentException)
+  }
 }