From 1998eb58d834dbe474c72cfd8d77c1ca15cac616 Mon Sep 17 00:00:00 2001
From: Ido Berkovich <ido@comet.com>
Date: Fri, 3 Jan 2025 13:45:03 +0200
Subject: [PATCH] OPIK-610 system message support

---
 .../opik/domain/llmproviders/Anthropic.java   | 23 +++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/llmproviders/Anthropic.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/llmproviders/Anthropic.java
index 7801410f3f..7e50dbdc02 100644
--- a/apps/opik-backend/src/main/java/com/comet/opik/domain/llmproviders/Anthropic.java
+++ b/apps/opik-backend/src/main/java/com/comet/opik/domain/llmproviders/Anthropic.java
@@ -8,6 +8,7 @@
 import dev.ai4j.openai4j.chat.Delta;
 import dev.ai4j.openai4j.chat.Message;
 import dev.ai4j.openai4j.chat.Role;
+import dev.ai4j.openai4j.chat.SystemMessage;
 import dev.ai4j.openai4j.chat.UserMessage;
 import dev.ai4j.openai4j.shared.Usage;
 import dev.langchain4j.data.message.AiMessage;
@@ -108,7 +109,12 @@ private AnthropicCreateMessageRequest toAnthropicCreateMessageRequest(ChatComple
         return builder
                 .stream(request.stream())
                 .model(request.model())
-                .messages(request.messages().stream().map(this::toMessage).toList())
+                .messages(request.messages().stream()
+                        .filter(message -> List.of(Role.ASSISTANT, Role.USER).contains(message.role()))
+                        .map(this::toMessage).toList())
+                .system(request.messages().stream()
+                        .filter(message -> message.role() == Role.SYSTEM)
+                        .map(this::toSystemMessage).toList())
                 .temperature(request.temperature())
                 .topP(request.topP())
                 .stopSequences(request.stop())
@@ -122,15 +128,24 @@ private AnthropicMessage toMessage(Message message) {
                     .role(AnthropicRole.ASSISTANT)
                     .content(List.of(new AnthropicTextContent(((AssistantMessage) message).content())))
                     .build();
-        } else if (message.role() == Role.USER) {
+        }
+
+        if (message.role() == Role.USER) {
             return AnthropicMessage.builder()
                     .role(AnthropicRole.USER)
                     .content(List.of(toAnthropicMessageContent(((UserMessage) message).content())))
                     .build();
         }
 
-        // Anthropic only supports assistant and user roles
-        throw new BadRequestException("not supported message role: " + message.role());
+        throw new BadRequestException("unexpected message role: " + message.role());
+    }
+
+    private AnthropicTextContent toSystemMessage(Message message) {
+        if (message.role() != Role.SYSTEM) {
+            throw new BadRequestException("expecting only system role, got: " + message.role());
+        }
+
+        return new AnthropicTextContent(((SystemMessage) message).content());
     }
 
     private AnthropicMessageContent toAnthropicMessageContent(Object rawContent) {