Handle "message_creation" ChatGPT response type

Sometimes, ChatGPT issues a "message_creation" response instead of
"tool_calls" when responding to run-steps requests. This update ensures
that the newly created messages by the ChatGPT assistant are properly
retrieved and processed in those cases as well.
Additionally, all functionalities related to thread messages have been
consolidated into the ChatGptThreadMessage class to enhance readability.

Change-Id: I717de414f91249dc1ea581820681b56648c6c423
Signed-off-by: Patrizio <patrizio.gelosi@amarulasolutions.com>
diff --git a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/PatchSetReviewer.java b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/PatchSetReviewer.java
index 440401f..c1a50ec 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/PatchSetReviewer.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/PatchSetReviewer.java
@@ -109,6 +109,10 @@
     }
 
     private void retrieveReviewBatches(ChatGptResponseContent reviewReply, GerritChange change) {
+        if (reviewReply.getMessageContent() != null && !reviewReply.getMessageContent().isEmpty()) {
+            reviewBatches.add(new ReviewBatch(reviewReply.getMessageContent()));
+            return;
+        }
         for (ChatGptReplyItem replyItem : reviewReply.getReplies()) {
             String reply = replyItem.getReply();
             Integer score = replyItem.getScore();
diff --git a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/common/model/api/chatgpt/ChatGptResponseContent.java b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/common/model/api/chatgpt/ChatGptResponseContent.java
index 2662155..7529395 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/common/model/api/chatgpt/ChatGptResponseContent.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/common/model/api/chatgpt/ChatGptResponseContent.java
@@ -12,5 +12,5 @@
     private List<ChatGptReplyItem> replies;
     private String changeId;
     @NonNull
-    private String errorMessage;
+    private String messageContent;
 }
diff --git a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/common/model/api/chatgpt/ChatGptResponseMessage.java b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/common/model/api/chatgpt/ChatGptResponseMessage.java
index 37c1e63..b9e169c 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/common/model/api/chatgpt/ChatGptResponseMessage.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/common/model/api/chatgpt/ChatGptResponseMessage.java
@@ -8,6 +8,15 @@
 @Data
 public class ChatGptResponseMessage {
     private String role;
+    private String type;
     @SerializedName("tool_calls")
     private List<ChatGptToolCall> toolCalls;
+    @SerializedName("message_creation")
+    private MessageCreation messageCreation;
+
+    @Data
+    public static class MessageCreation {
+        @SerializedName("message_id")
+        private String messageId;
+    }
 }
diff --git a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/UriResourceLocatorStateful.java b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/UriResourceLocatorStateful.java
index 6be772e..a313d5a 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/UriResourceLocatorStateful.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/UriResourceLocatorStateful.java
@@ -21,6 +21,10 @@
         return threadRetrieveUri(threadId) + "/messages";
     }
 
+    public static String threadMessageRetrieveUri(String threadId, String messageId) {
+        return threadMessagesUri(threadId) + "/" + messageId;
+    }
+
     public static String runsUri(String threadId) {
         return threadRetrieveUri(threadId) + "/runs";
     }
diff --git a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptClientStateful.java b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptClientStateful.java
index 833c0ca..02d50a6 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptClientStateful.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptClientStateful.java
@@ -10,11 +10,15 @@
 import com.googlesource.gerrit.plugins.chatgpt.mode.common.model.api.chatgpt.ChatGptResponseContent;
 import com.googlesource.gerrit.plugins.chatgpt.mode.common.model.data.ChangeSetData;
 import com.googlesource.gerrit.plugins.chatgpt.mode.interfaces.client.api.chatgpt.IChatGptClient;
+import com.googlesource.gerrit.plugins.chatgpt.mode.stateful.model.api.chatgpt.ChatGptThreadMessageResponse;
 import lombok.extern.slf4j.Slf4j;
 
 @Slf4j
 @Singleton
 public class ChatGptClientStateful extends ChatGptClient implements IChatGptClient {
+    private static final String TYPE_MESSAGE_CREATION = "message_creation";
+    private static final String TYPE_TOOL_CALLS = "tool_calls";
+
     private final PluginDataHandlerProvider pluginDataHandlerProvider;
 
     @VisibleForTesting
@@ -29,23 +33,42 @@
         String changeId = change.getFullChangeId();
         log.info("Processing STATEFUL ChatGPT Request with changeId: {}, Patch Set: {}", changeId, patchSet);
 
-        ChatGptThread chatGptThread = new ChatGptThread(
+        ChatGptThread chatGptThread = new ChatGptThread(config, pluginDataHandlerProvider);
+        String threadId = chatGptThread.createThread();
+
+        ChatGptThreadMessage chatGptThreadMessage = new ChatGptThreadMessage(
+                threadId,
                 config,
                 changeSetData,
                 change,
-                patchSet,
-                pluginDataHandlerProvider
+                patchSet
         );
-        String threadId = chatGptThread.createThread();
-        chatGptThread.addMessage();
+        chatGptThreadMessage.addMessage();
 
         ChatGptRun chatGptRun = new ChatGptRun(threadId, config, pluginDataHandlerProvider);
         chatGptRun.createRun();
         chatGptRun.pollRun();
         // Attribute `requestBody` is valued for testing purposes
-        requestBody = chatGptThread.getAddMessageRequestBody();
+        requestBody = chatGptThreadMessage.getAddMessageRequestBody();
         log.debug("ChatGPT request body: {}", requestBody);
 
-        return getResponseContent(chatGptRun.getFirstStep());
+        return getResponseContentStateful(threadId, chatGptRun);
+    }
+
+    private ChatGptResponseContent getResponseContentStateful(String threadId, ChatGptRun chatGptRun) {
+        return switch (chatGptRun.getFirstStepDetails().getType()) {
+            case TYPE_MESSAGE_CREATION -> retrieveThreadMessage(threadId, chatGptRun);
+            case TYPE_TOOL_CALLS -> getResponseContent(chatGptRun.getFirstStep());
+            default -> throw new IllegalStateException("Unexpected Step Type in stateful ChatGpt response: " +
+                    chatGptRun);
+        };
+    }
+
+    private ChatGptResponseContent retrieveThreadMessage(String threadId, ChatGptRun chatGptRun) {
+        ChatGptThreadMessage chatGptThreadMessage = new ChatGptThreadMessage(threadId, config);
+        ChatGptThreadMessageResponse threadMessageResponse = chatGptThreadMessage.retrieveMessage(
+                chatGptRun.getFirstStepDetails().getMessageCreation().getMessageId()
+        );
+        return new ChatGptResponseContent(threadMessageResponse.getContent().get(0).getText().getValue());
     }
 }
diff --git a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptRun.java b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptRun.java
index d078e67..bd734fa 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptRun.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptRun.java
@@ -4,6 +4,7 @@
 import com.googlesource.gerrit.plugins.chatgpt.data.PluginDataHandler;
 import com.googlesource.gerrit.plugins.chatgpt.data.PluginDataHandlerProvider;
 import com.googlesource.gerrit.plugins.chatgpt.mode.common.client.ClientBase;
+import com.googlesource.gerrit.plugins.chatgpt.mode.common.model.api.chatgpt.ChatGptResponseMessage;
 import com.googlesource.gerrit.plugins.chatgpt.mode.common.model.api.chatgpt.ChatGptToolCall;
 import com.googlesource.gerrit.plugins.chatgpt.mode.stateful.client.api.UriResourceLocatorStateful;
 import com.googlesource.gerrit.plugins.chatgpt.mode.stateful.model.api.chatgpt.*;
@@ -72,8 +73,12 @@
         log.info("Run executed after {} polling requests: {}", pollingCount, stepResponse);
     }
 
+    public ChatGptResponseMessage getFirstStepDetails() {
+        return stepResponse.getData().get(0).getStepDetails();
+    }
+
     public List<ChatGptToolCall> getFirstStep() {
-        return stepResponse.getData().get(0).getStepDetails().getToolCalls();
+        return getFirstStepDetails().getToolCalls();
     }
 
     private Request runCreateRequest() {
diff --git a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptThread.java b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptThread.java
index 6f9dd90..413bf98 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptThread.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptThread.java
@@ -3,12 +3,8 @@
 import com.googlesource.gerrit.plugins.chatgpt.config.Configuration;
 import com.googlesource.gerrit.plugins.chatgpt.data.PluginDataHandler;
 import com.googlesource.gerrit.plugins.chatgpt.data.PluginDataHandlerProvider;
-import com.googlesource.gerrit.plugins.chatgpt.mode.common.client.api.gerrit.GerritChange;
-import com.googlesource.gerrit.plugins.chatgpt.mode.common.model.api.chatgpt.ChatGptRequestMessage;
-import com.googlesource.gerrit.plugins.chatgpt.mode.common.model.data.ChangeSetData;
 import com.googlesource.gerrit.plugins.chatgpt.mode.stateful.client.api.UriResourceLocatorStateful;
-import com.googlesource.gerrit.plugins.chatgpt.mode.stateful.client.prompt.ChatGptPromptStateful;
-import com.googlesource.gerrit.plugins.chatgpt.mode.stateful.model.api.chatgpt.*;
+import com.googlesource.gerrit.plugins.chatgpt.mode.stateful.model.api.chatgpt.ChatGptResponse;
 import lombok.extern.slf4j.Slf4j;
 import okhttp3.Request;
 
@@ -22,30 +18,18 @@
 
     private final ChatGptHttpClient httpClient = new ChatGptHttpClient();
     private final Configuration config;
-    private final ChangeSetData changeSetData;
-    private final GerritChange change;
-    private final String patchSet;
     private final PluginDataHandler changeDataHandler;
 
-    private String threadId;
-    private ChatGptRequestMessage addMessageRequestBody;
-
     public ChatGptThread(
             Configuration config,
-            ChangeSetData changeSetData,
-            GerritChange change,
-            String patchSet,
             PluginDataHandlerProvider pluginDataHandlerProvider
     ) {
         this.config = config;
-        this.changeSetData = changeSetData;
-        this.change = change;
-        this.patchSet = patchSet;
         this.changeDataHandler = pluginDataHandlerProvider.getChangeScope();
     }
 
     public String createThread() {
-        threadId = changeDataHandler.getValue(KEY_THREAD_ID);
+        String threadId = changeDataHandler.getValue(KEY_THREAD_ID);
         if (threadId == null) {
             Request request = createThreadRequest();
             log.debug("ChatGPT Create Thread request: {}", request);
@@ -61,35 +45,10 @@
         return threadId;
     }
 
-    public void addMessage() {
-        Request request = addMessageRequest();
-        log.debug("ChatGPT Add Message request: {}", request);
-
-        ChatGptResponse addMessageResponse = getGson().fromJson(httpClient.execute(request), ChatGptResponse.class);
-        log.info("Message added: {}", addMessageResponse);
-    }
-
-    public String getAddMessageRequestBody() {
-        return getGson().toJson(addMessageRequestBody);
-    }
-
     private Request createThreadRequest() {
         URI uri = URI.create(config.getGptDomain() + UriResourceLocatorStateful.threadsUri());
         log.debug("ChatGPT Create Thread request URI: {}", uri);
 
         return httpClient.createRequestFromJson(uri.toString(), config.getGptToken(), new Object());
     }
-
-    private Request addMessageRequest() {
-        URI uri = URI.create(config.getGptDomain() + UriResourceLocatorStateful.threadMessagesUri(threadId));
-        log.debug("ChatGPT Add Message request URI: {}", uri);
-        ChatGptPromptStateful chatGptPromptStateful = new ChatGptPromptStateful(config, changeSetData, change);
-        addMessageRequestBody = ChatGptRequestMessage.builder()
-                .role("user")
-                .content(chatGptPromptStateful.getDefaultGptThreadReviewMessage(patchSet))
-                .build();
-        log.debug("ChatGPT Add Message request body: {}", addMessageRequestBody);
-
-        return httpClient.createRequestFromJson(uri.toString(), config.getGptToken(), addMessageRequestBody);
-    }
 }
diff --git a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptThreadMessage.java b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptThreadMessage.java
new file mode 100644
index 0000000..06c9171
--- /dev/null
+++ b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptThreadMessage.java
@@ -0,0 +1,90 @@
+package com.googlesource.gerrit.plugins.chatgpt.mode.stateful.client.api.chatgpt;
+
+import com.googlesource.gerrit.plugins.chatgpt.config.Configuration;
+import com.googlesource.gerrit.plugins.chatgpt.mode.common.client.ClientBase;
+import com.googlesource.gerrit.plugins.chatgpt.mode.common.client.api.gerrit.GerritChange;
+import com.googlesource.gerrit.plugins.chatgpt.mode.common.model.api.chatgpt.ChatGptRequestMessage;
+import com.googlesource.gerrit.plugins.chatgpt.mode.common.model.data.ChangeSetData;
+import com.googlesource.gerrit.plugins.chatgpt.mode.stateful.client.api.UriResourceLocatorStateful;
+import com.googlesource.gerrit.plugins.chatgpt.mode.stateful.client.prompt.ChatGptPromptStateful;
+import com.googlesource.gerrit.plugins.chatgpt.mode.stateful.model.api.chatgpt.ChatGptResponse;
+import com.googlesource.gerrit.plugins.chatgpt.mode.stateful.model.api.chatgpt.ChatGptThreadMessageResponse;
+import lombok.extern.slf4j.Slf4j;
+import okhttp3.Request;
+
+import java.net.URI;
+
+import static com.googlesource.gerrit.plugins.chatgpt.utils.GsonUtils.getGson;
+
+@Slf4j
+public class ChatGptThreadMessage extends ClientBase {
+    private final ChatGptHttpClient httpClient = new ChatGptHttpClient();
+    private final String threadId;
+
+    private ChangeSetData changeSetData;
+    private GerritChange change;
+    private String patchSet;
+    private ChatGptRequestMessage addMessageRequestBody;
+
+    public ChatGptThreadMessage(String threadId, Configuration config) {
+        super(config);
+        this.threadId = threadId;
+    }
+
+    public ChatGptThreadMessage(
+            String threadId,
+            Configuration config,
+            ChangeSetData changeSetData,
+            GerritChange change,
+            String patchSet
+    ) {
+        this(threadId, config);
+        this.changeSetData = changeSetData;
+        this.change = change;
+        this.patchSet = patchSet;
+    }
+
+    public ChatGptThreadMessageResponse retrieveMessage(String messageId) {
+        Request request = createRetrieveMessageRequest(messageId);
+        log.debug("ChatGPT Retrieve Thread Message request: {}", request);
+        ChatGptThreadMessageResponse threadMessageResponse = getGson().fromJson(
+                httpClient.execute(request), ChatGptThreadMessageResponse.class
+        );
+        log.info("Thread Message retrieved: {}", threadMessageResponse);
+
+        return threadMessageResponse;
+    }
+
+    public void addMessage() {
+        Request request = addMessageRequest();
+        log.debug("ChatGPT Add Message request: {}", request);
+
+        ChatGptResponse addMessageResponse = getGson().fromJson(httpClient.execute(request), ChatGptResponse.class);
+        log.info("Message added: {}", addMessageResponse);
+    }
+
+    public String getAddMessageRequestBody() {
+        return getGson().toJson(addMessageRequestBody);
+    }
+
+    private Request createRetrieveMessageRequest(String messageId) {
+        URI uri = URI.create(config.getGptDomain() +
+                UriResourceLocatorStateful.threadMessageRetrieveUri(threadId, messageId));
+        log.debug("ChatGPT Retrieve Thread Message request URI: {}", uri);
+
+        return httpClient.createRequestFromJson(uri.toString(), config.getGptToken(), null);
+    }
+
+    private Request addMessageRequest() {
+        URI uri = URI.create(config.getGptDomain() + UriResourceLocatorStateful.threadMessagesUri(threadId));
+        log.debug("ChatGPT Add Message request URI: {}", uri);
+        ChatGptPromptStateful chatGptPromptStateful = new ChatGptPromptStateful(config, changeSetData, change);
+        addMessageRequestBody = ChatGptRequestMessage.builder()
+                .role("user")
+                .content(chatGptPromptStateful.getDefaultGptThreadReviewMessage(patchSet))
+                .build();
+        log.debug("ChatGPT Add Message request body: {}", addMessageRequestBody);
+
+        return httpClient.createRequestFromJson(uri.toString(), config.getGptToken(), addMessageRequestBody);
+    }
+}
diff --git a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/model/api/chatgpt/ChatGptThreadMessageResponse.java b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/model/api/chatgpt/ChatGptThreadMessageResponse.java
new file mode 100644
index 0000000..5f063ea
--- /dev/null
+++ b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/model/api/chatgpt/ChatGptThreadMessageResponse.java
@@ -0,0 +1,23 @@
+package com.googlesource.gerrit.plugins.chatgpt.mode.stateful.model.api.chatgpt;
+
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+
+import java.util.List;
+
+@EqualsAndHashCode(callSuper = true)
+@Data
+public class ChatGptThreadMessageResponse extends ChatGptResponse {
+    private List<Content> content;
+
+    @Data
+    public static class Content {
+        private String type;
+        private Text text;
+
+        @Data
+        public static class Text {
+            private String value;
+        }
+    }
+}
diff --git a/src/test/java/com/googlesource/gerrit/plugins/chatgpt/ChatGptReviewStatefulTest.java b/src/test/java/com/googlesource/gerrit/plugins/chatgpt/ChatGptReviewStatefulTest.java
index 73539e7..d5ecc17 100644
--- a/src/test/java/com/googlesource/gerrit/plugins/chatgpt/ChatGptReviewStatefulTest.java
+++ b/src/test/java/com/googlesource/gerrit/plugins/chatgpt/ChatGptReviewStatefulTest.java
@@ -202,4 +202,30 @@
         Assert.assertEquals(promptTagComments, requestContent);
         Assert.assertEquals(reviewMessageCommitMessage, getCapturedMessage(captor, GERRIT_PATCH_SET_FILENAME));
     }
+
+    @Test
+    public void gptMentionedInCommentMessageResponse() throws RestApiException {
+        String reviewMessageCommitMessage = getReviewMessage("__files/chatGptResponseRequestStateful.json", 0);
+
+        chatGptPromptStateful.setCommentEvent(true);
+        // Mock the behavior of the ChatGPT retrieve-run-steps request
+        WireMock.stubFor(WireMock.get(WireMock.urlEqualTo(URI.create(config.getGptDomain()
+                        + UriResourceLocatorStateful.runStepsUri(CHAT_GPT_THREAD_ID, CHAT_GPT_RUN_ID)).getPath()))
+                .willReturn(WireMock.aResponse()
+                        .withStatus(HTTP_OK)
+                        .withHeader(HttpHeaders.CONTENT_TYPE, ContentType.APPLICATION_JSON.toString())
+                        .withBodyFile("chatGptResponseRequestMessageStateful.json")));
+        WireMock.stubFor(WireMock.get(WireMock.urlEqualTo(URI.create(config.getGptDomain()
+                        + UriResourceLocatorStateful.threadMessageRetrieveUri(CHAT_GPT_THREAD_ID, CHAT_GPT_MESSAGE_ID)).getPath()))
+                .willReturn(WireMock.aResponse()
+                        .withStatus(HTTP_OK)
+                        .withHeader(HttpHeaders.CONTENT_TYPE, ContentType.APPLICATION_JSON.toString())
+                        .withBodyFile("chatGptResponseThreadMessage.json")));
+
+        handleEventBasedOnType(true);
+
+        ArgumentCaptor<ReviewInput> captor = testRequestSent();
+        Assert.assertEquals(promptTagComments, requestContent);
+        Assert.assertEquals(reviewMessageCommitMessage, getCapturedMessage(captor, GERRIT_PATCH_SET_FILENAME));
+    }
 }
diff --git a/src/test/resources/__files/chatGptResponseRequestMessageStateful.json b/src/test/resources/__files/chatGptResponseRequestMessageStateful.json
new file mode 100644
index 0000000..812caa4
--- /dev/null
+++ b/src/test/resources/__files/chatGptResponseRequestMessageStateful.json
@@ -0,0 +1,15 @@
+{
+  "object": "list",
+  "data": [
+    {
+      "id": "step_UKMPnSinQy6XjSWO7SFleu1v",
+      "object": "thread.run.step",
+      "step_details": {
+        "type": "message_creation",
+        "message_creation": {
+          "message_id": "msg_TEST_MESSAGE_ID"
+        }
+      }
+    }
+  ]
+}
diff --git a/src/test/resources/__files/chatGptResponseThreadMessage.json b/src/test/resources/__files/chatGptResponseThreadMessage.json
new file mode 100644
index 0000000..a001e40
--- /dev/null
+++ b/src/test/resources/__files/chatGptResponseThreadMessage.json
@@ -0,0 +1,12 @@
+{
+  "object": "thread.message",
+  "role": "assistant",
+  "content": [
+    {
+      "type": "text",
+      "text": {
+        "value": "The commit message 'Corrected Indentation in Module-Class Retrieval Line' accurately represents the change made in the code."
+      }
+    }
+  ]
+}