Added priority to Role enum.

facebookincubator · Jan 18, 2024 · 2e1adfc · 2e1adfc
1 parent ea905cb
commit 2e1adfc
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/src/main/java/com/meta/cp4m/message/Message.java b/src/main/java/com/meta/cp4m/message/Message.java
@@ -44,8 +44,15 @@ default Identifier threadId() {
   }
 
   enum Role {
-    ASSISTANT,
-    USER,
-    SYSTEM
+    ASSISTANT(0),
+    USER(1),
+    SYSTEM(2);
+
+    public final Integer priority;
+
+    private Role(Integer priority){
+      this.priority = priority;
+    }
+
   }
 }
diff --git a/src/main/java/com/meta/cp4m/message/ThreadState.java b/src/main/java/com/meta/cp4m/message/ThreadState.java
@@ -42,7 +42,7 @@ private ThreadState(ThreadState<T> old, T newMessage) {
     T mWithParentMessage = newMessage.role() == Role.USER ? (T) newMessage.addParentMessage(old.tail()): newMessage;
     this.messages =
           Stream.concat(messages.stream(), Stream.of(mWithParentMessage))
-              .sorted((m1,m2) -> m1.parentMessage() == m2.parentMessage() ? (m1.role().compareTo(m2.role())) : (m1.timestamp().compareTo(m2.timestamp())))
+              .sorted((m1,m2) -> m1.parentMessage() == m2.parentMessage() ? (m1.role().priority.compareTo(m2.role().priority)) : (m1.timestamp().compareTo(m2.timestamp())))
               .collect(Collectors.toUnmodifiableList());
 
     Preconditions.checkArgument(