diff --git a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/common/client/api/gerrit/GerritClientPatchSet.java b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/common/client/api/gerrit/GerritClientPatchSet.java
index 4e17e7e..ff54e0e 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/common/client/api/gerrit/GerritClientPatchSet.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/common/client/api/gerrit/GerritClientPatchSet.java
@@ -1,24 +1,38 @@
 package com.googlesource.gerrit.plugins.chatgpt.mode.common.client.api.gerrit;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
 
 import com.google.gerrit.extensions.client.ListChangesOption;
 import com.google.gerrit.extensions.common.ChangeInfo;
+import com.google.gerrit.extensions.common.DiffInfo;
 import com.google.gerrit.server.account.AccountCache;
 import com.google.gerrit.server.util.ManualRequestContext;
 import com.googlesource.gerrit.plugins.chatgpt.config.Configuration;
-import com.googlesource.gerrit.plugins.chatgpt.data.ChangeSetDataHandler;
+import com.googlesource.gerrit.plugins.chatgpt.mode.common.client.patch.diff.FileDiffProcessed;
+import com.googlesource.gerrit.plugins.chatgpt.mode.common.model.api.gerrit.GerritFileDiff;
+import com.googlesource.gerrit.plugins.chatgpt.mode.common.model.api.gerrit.GerritPatchSetFileDiff;
+import com.googlesource.gerrit.plugins.chatgpt.mode.common.model.api.gerrit.GerritReviewFileDiff;
 import com.googlesource.gerrit.plugins.chatgpt.mode.common.model.data.ChangeSetData;
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
 
+import static com.googlesource.gerrit.plugins.chatgpt.utils.GsonUtils.getNoEscapedGson;
+import static java.util.stream.Collectors.toList;
+
 @Slf4j
 public class GerritClientPatchSet extends GerritClientAccount {
+    protected final List<String> diffs;
+
     @Getter
     protected Integer revisionBase = 0;
 
+    private boolean isCommitMessage;
+
     public GerritClientPatchSet(Configuration config, AccountCache accountCache) {
         super(config, accountCache);
+        diffs = new ArrayList<>();
     }
 
     public void retrieveRevisionBase(GerritChange change) {
@@ -48,8 +62,74 @@
         return isChangeSetBased(changeSetData) ? 0 : revisionBase;
     }
 
+    protected void retrieveFileDiff(GerritChange change, List<String> files, int revisionBase) throws Exception {
+        List<String> enabledFileExtensions = config.getEnabledFileExtensions();
+        try (ManualRequestContext requestContext = config.openRequestContext()) {
+            for (String filename : files) {
+                isCommitMessage = filename.equals("/COMMIT_MSG");
+                if (!isCommitMessage && (filename.lastIndexOf(".") < 1 ||
+                        !enabledFileExtensions.contains(filename.substring(filename.lastIndexOf("."))))) {
+                    continue;
+                }
+                DiffInfo diff =
+                        config
+                                .getGerritApi()
+                                .changes()
+                                .id(
+                                        change.getProjectName(),
+                                        change.getBranchNameKey().shortName(),
+                                        change.getChangeKey().get())
+                                .current()
+                                .file(filename)
+                                .diff(revisionBase);
+                processFileDiff(filename, diff);
+            }
+        }
+    }
+
     private boolean isChangeSetBased(ChangeSetData changeSetData) {
         return !changeSetData.getForcedReviewLastPatchSet();
     }
 
+    private void processFileDiff(String filename, DiffInfo diff) {
+        log.debug("FileDiff content processed: {}", filename);
+
+        GerritPatchSetFileDiff gerritPatchSetFileDiff = new GerritPatchSetFileDiff();
+        Optional.ofNullable(diff.metaA)
+                .ifPresent(
+                        meta -> gerritPatchSetFileDiff.setMetaA(GerritClientPatchSet.toMeta(meta)));
+        Optional.ofNullable(diff.metaB)
+                .ifPresent(
+                        meta -> gerritPatchSetFileDiff.setMetaB(GerritClientPatchSet.toMeta(meta)));
+        Optional.ofNullable(diff.content)
+                .ifPresent(
+                        content ->
+                                gerritPatchSetFileDiff.setContent(
+                                        content.stream()
+                                                .map(GerritClientPatchSet::toContent)
+                                                .collect(toList())));
+
+        // Initialize the reduced file diff for the Gerrit review with fields `meta_a` and `meta_b`
+        GerritReviewFileDiff gerritReviewFileDiff = new GerritReviewFileDiff(gerritPatchSetFileDiff.getMetaA(),
+                gerritPatchSetFileDiff.getMetaB());
+        FileDiffProcessed fileDiffProcessed = new FileDiffProcessed(config, isCommitMessage, gerritPatchSetFileDiff);
+        fileDiffsProcessed.put(filename, fileDiffProcessed);
+        gerritReviewFileDiff.setContent(fileDiffProcessed.getReviewDiffContent());
+        diffs.add(getNoEscapedGson().toJson(gerritReviewFileDiff));
+    }
+
+    protected static GerritFileDiff.Meta toMeta(DiffInfo.FileMeta input) {
+        GerritFileDiff.Meta meta = new GerritFileDiff.Meta();
+        meta.setContentType(input.contentType);
+        meta.setName(input.name);
+        return meta;
+    }
+
+    protected static GerritPatchSetFileDiff.Content toContent(DiffInfo.ContentEntry input) {
+        GerritPatchSetFileDiff.Content content = new GerritPatchSetFileDiff.Content();
+        content.a = input.a;
+        content.b = input.b;
+        content.ab = input.ab;
+        return content;
+    }
 }
diff --git a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/common/client/patch/code/CodeFinder.java b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/common/client/patch/code/CodeFinder.java
index 2267261..0f736cf 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/common/client/patch/code/CodeFinder.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/common/client/patch/code/CodeFinder.java
@@ -15,6 +15,7 @@
 @Slf4j
 public class CodeFinder {
     private static final String PUNCTUATION_REGEX = "([()\\[\\]{}<>:;,?&+\\-*/%|=])";
+    private static final String BEGINNING_DIFF_REGEX = "(?:^|\n)[+\\-]";
     private static final String ENDING_ELLIPSIS_REGEX = "\\.\\.\\.\\W*$";
 
     private final String NON_PRINTING_REPLACEMENT;
@@ -58,7 +59,10 @@
     }
 
     private void updateCodePattern(ChatGptReplyItem replyItem) {
-        String commentedCode = replyItem.getCodeSnippet().replaceAll(ENDING_ELLIPSIS_REGEX, "").trim();
+        String commentedCode = replyItem.getCodeSnippet()
+                .replaceAll(BEGINNING_DIFF_REGEX, "")
+                .replaceAll(ENDING_ELLIPSIS_REGEX, "")
+                .trim();
         String commentedCodeRegex = Pattern.quote(commentedCode);
         // Generalize the regex to capture snippets where existing sequences of non-printing chars have been modified
         // from the original code
diff --git a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/gerrit/GerritClientPatchSetStateful.java b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/gerrit/GerritClientPatchSetStateful.java
index 74f0cd3..e7c92c9 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/gerrit/GerritClientPatchSetStateful.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/gerrit/GerritClientPatchSetStateful.java
@@ -14,12 +14,18 @@
 import com.googlesource.gerrit.plugins.chatgpt.mode.stateful.client.api.git.GitRepoFiles;
 import lombok.extern.slf4j.Slf4j;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
 import static com.googlesource.gerrit.plugins.chatgpt.settings.Settings.COMMIT_MESSAGE_FILTER_OUT_PREFIXES;
 
 @Slf4j
 public class GerritClientPatchSetStateful extends GerritClientPatchSet implements IGerritClientPatchSet {
+    private static final Pattern EXTRACT_B_FILENAMES_FROM_PATCH_SET = Pattern.compile("^diff --git .*? b/(.*)$",
+            Pattern.MULTILINE);
+
     private final GitRepoFiles gitRepoFiles;
     private final PluginDataHandler pluginDataHandler;
 
@@ -42,12 +48,16 @@
         ChatGptAssistant chatGptAssistant = new ChatGptAssistant(config, change, gitRepoFiles, pluginDataHandler);
         chatGptAssistant.setupAssistant();
 
-        return getPatchFromGerrit();
+        String formattedPatch = getPatchFromGerrit();
+        List<String> files = extractFilesFromPatch(formattedPatch);
+        retrieveFileDiff(change, files, revisionBase);
+
+        return formattedPatch;
     }
 
     private String getPatchFromGerrit() throws Exception {
         try (ManualRequestContext requestContext = config.openRequestContext()) {
-            String gerritPatch = config
+            String formattedPatch = config
                 .getGerritApi()
                 .changes()
                 .id(
@@ -57,19 +67,28 @@
                 .current()
                 .patch()
                 .asString();
-            log.debug("Gerrit Patch retrieved: {}", gerritPatch);
+            log.debug("Formatted Patch retrieved: {}", formattedPatch);
 
-            return filterPatch(gerritPatch);
+            return filterPatch(formattedPatch);
         }
     }
 
-    private String filterPatch(String gerritPatch) {
+    private String filterPatch(String formattedPatch) {
         // Remove Patch heading up to the Change-Id annotation
         Pattern CONFIG_ID_HEADING_PATTERN = Pattern.compile(
                 "^.*?" + COMMIT_MESSAGE_FILTER_OUT_PREFIXES.get("CHANGE_ID") + " " + change.getChangeKey().get(),
                 Pattern.DOTALL
         );
-        return CONFIG_ID_HEADING_PATTERN.matcher(gerritPatch).replaceAll("");
+        return CONFIG_ID_HEADING_PATTERN.matcher(formattedPatch).replaceAll("");
+    }
+
+    private List<String> extractFilesFromPatch(String formattedPatch) {
+        Matcher extractFilenameMatcher = EXTRACT_B_FILENAMES_FROM_PATCH_SET.matcher(formattedPatch);
+        List<String> files = new ArrayList<>();
+        while (extractFilenameMatcher.find()) {
+            files.add(extractFilenameMatcher.group(1));
+        }
+        return files;
     }
 
 }
diff --git a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateless/client/api/gerrit/GerritClientPatchSetStateless.java b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateless/client/api/gerrit/GerritClientPatchSetStateless.java
index bbba27f..a757858 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateless/client/api/gerrit/GerritClientPatchSetStateless.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateless/client/api/gerrit/GerritClientPatchSetStateless.java
@@ -2,40 +2,27 @@
 
 import com.google.inject.Inject;
 import com.google.common.annotations.VisibleForTesting;
-import com.google.gerrit.extensions.common.DiffInfo;
 import com.google.gerrit.extensions.common.FileInfo;
 import com.google.gerrit.server.account.AccountCache;
 import com.google.gerrit.server.util.ManualRequestContext;
 import com.googlesource.gerrit.plugins.chatgpt.config.Configuration;
 import com.googlesource.gerrit.plugins.chatgpt.mode.common.client.api.gerrit.GerritChange;
 import com.googlesource.gerrit.plugins.chatgpt.mode.common.client.api.gerrit.GerritClientPatchSet;
-import com.googlesource.gerrit.plugins.chatgpt.mode.common.client.patch.diff.FileDiffProcessed;
-import com.googlesource.gerrit.plugins.chatgpt.mode.common.model.api.gerrit.GerritFileDiff;
-import com.googlesource.gerrit.plugins.chatgpt.mode.common.model.api.gerrit.GerritPatchSetFileDiff;
-import com.googlesource.gerrit.plugins.chatgpt.mode.common.model.api.gerrit.GerritReviewFileDiff;
 import com.googlesource.gerrit.plugins.chatgpt.mode.common.model.data.ChangeSetData;
 import com.googlesource.gerrit.plugins.chatgpt.mode.interfaces.client.api.gerrit.IGerritClientPatchSet;
 import lombok.extern.slf4j.Slf4j;
 
 import static java.util.stream.Collectors.toList;
 
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
-
-import static com.googlesource.gerrit.plugins.chatgpt.utils.GsonUtils.getNoEscapedGson;
 
 @Slf4j
 public class GerritClientPatchSetStateless extends GerritClientPatchSet implements IGerritClientPatchSet {
-    private final List<String> diffs;
-    private boolean isCommitMessage;
-
     @VisibleForTesting
     @Inject
     public GerritClientPatchSetStateless(Configuration config, AccountCache accountCache) {
         super(config, accountCache);
-        diffs = new ArrayList<>();
     }
 
     public String getPatchSet(ChangeSetData changeSetData, GerritChange change) throws Exception {
@@ -83,72 +70,9 @@
         }
     }
 
-    private void processFileDiff(String filename, DiffInfo diff) {
-        log.debug("FileDiff content processed: {}", filename);
-
-        GerritPatchSetFileDiff gerritPatchSetFileDiff = new GerritPatchSetFileDiff();
-        Optional.ofNullable(diff.metaA)
-            .ifPresent(
-                meta -> gerritPatchSetFileDiff.setMetaA(GerritClientPatchSetStateless.toMeta(meta)));
-        Optional.ofNullable(diff.metaB)
-            .ifPresent(
-                meta -> gerritPatchSetFileDiff.setMetaB(GerritClientPatchSetStateless.toMeta(meta)));
-        Optional.ofNullable(diff.content)
-            .ifPresent(
-                content ->
-                    gerritPatchSetFileDiff.setContent(
-                        content.stream()
-                            .map(GerritClientPatchSetStateless::toContent)
-                            .collect(toList())));
-
-        // Initialize the reduced file diff for the Gerrit review with fields `meta_a` and `meta_b`
-        GerritReviewFileDiff gerritReviewFileDiff = new GerritReviewFileDiff(gerritPatchSetFileDiff.getMetaA(),
-                gerritPatchSetFileDiff.getMetaB());
-        FileDiffProcessed fileDiffProcessed = new FileDiffProcessed(config, isCommitMessage, gerritPatchSetFileDiff);
-        fileDiffsProcessed.put(filename, fileDiffProcessed);
-        gerritReviewFileDiff.setContent(fileDiffProcessed.getReviewDiffContent());
-        diffs.add(getNoEscapedGson().toJson(gerritReviewFileDiff));
-    }
-
     private String getFileDiffsJson(GerritChange change, List<String> files, int revisionBase) throws Exception {
-        List<String> enabledFileExtensions = config.getEnabledFileExtensions();
-        try (ManualRequestContext requestContext = config.openRequestContext()) {
-            for (String filename : files) {
-                isCommitMessage = filename.equals("/COMMIT_MSG");
-                if (!isCommitMessage && (filename.lastIndexOf(".") < 1 ||
-                        !enabledFileExtensions.contains(filename.substring(filename.lastIndexOf("."))))) {
-                    continue;
-                }
-                DiffInfo diff =
-                    config
-                        .getGerritApi()
-                        .changes()
-                        .id(
-                            change.getProjectName(),
-                            change.getBranchNameKey().shortName(),
-                            change.getChangeKey().get())
-                        .current()
-                        .file(filename)
-                        .diff(revisionBase);
-                processFileDiff(filename, diff);
-            }
-        }
+        retrieveFileDiff(change, files, revisionBase);
         diffs.add(String.format("{\"changeId\": \"%s\"}", change.getFullChangeId()));
         return "[" + String.join(",", diffs) + "]\n";
     }
-
-    private static GerritFileDiff.Meta toMeta(DiffInfo.FileMeta input) {
-        GerritFileDiff.Meta meta = new GerritFileDiff.Meta();
-        meta.setContentType(input.contentType);
-        meta.setName(input.name);
-        return meta;
-    }
-
-    private static GerritPatchSetFileDiff.Content toContent(DiffInfo.ContentEntry input) {
-        GerritPatchSetFileDiff.Content content = new GerritPatchSetFileDiff.Content();
-        content.a = input.a;
-        content.b = input.b;
-        content.ab = input.ab;
-        return content;
-    }
 }
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 55f0ac8..79e15ea 100644
--- a/src/test/java/com/googlesource/gerrit/plugins/chatgpt/ChatGptReviewStatefulTest.java
+++ b/src/test/java/com/googlesource/gerrit/plugins/chatgpt/ChatGptReviewStatefulTest.java
@@ -2,7 +2,9 @@
 
 import com.github.tomakehurst.wiremock.client.WireMock;
 import com.google.common.net.HttpHeaders;
+import com.google.gerrit.extensions.api.changes.FileApi;
 import com.google.gerrit.extensions.api.changes.ReviewInput;
+import com.google.gerrit.extensions.common.DiffInfo;
 import com.google.gerrit.extensions.restapi.BinaryResult;
 import com.google.gerrit.extensions.restapi.RestApiException;
 import com.google.gson.JsonObject;
@@ -135,6 +137,11 @@
                 .setContentType("text/plain")
                 .setContentLength(formattedPatchContent.length());
         when(revisionApiMock.patch()).thenReturn(binaryResult);
+
+        FileApi testFileMock = mock(FileApi.class);
+        when(revisionApiMock.file("test_file_1.py")).thenReturn(testFileMock);
+        DiffInfo testFileDiff = readTestFileToClass("__files/stateful/gerritPatchSetDiffTestFile.json", DiffInfo.class);
+        when(testFileMock.diff(0)).thenReturn(testFileDiff);
     }
 
     protected void initComparisonContent() {
diff --git a/src/test/resources/__files/stateful/gerritPatchSetDiffTestFile.json b/src/test/resources/__files/stateful/gerritPatchSetDiffTestFile.json
new file mode 100644
index 0000000..151f0a9
--- /dev/null
+++ b/src/test/resources/__files/stateful/gerritPatchSetDiffTestFile.json
@@ -0,0 +1,78 @@
+{
+  "meta_a": {
+    "name": "test_file_1.py",
+    "content_type": "text/x-python",
+    "lines": 43
+  },
+  "meta_b": {
+    "name": "test_file_1.py",
+    "content_type": "text/x-python",
+    "lines": 43
+  },
+  "change_type": "MODIFIED",
+  "diff_header": [
+    "diff --git a/test_file_1.py b/test_file_1.py",
+    "index 95ab5e7..e368bb0 100644",
+    "--- a/test_file_1.py",
+    "+++ b/test_file_1.py"
+  ],
+  "content": [
+    {
+      "ab": [
+        "from types import Any, Callable, Type, Union",
+        "",
+        "__all__ = [\"importclass\", \"preprocess_classes\", \"TypeClassOrPath\"]",
+        "",
+        "TypeClassOrPath = Union[Type, str]",
+        "",
+        "",
+        "def importclass(",
+        "    module_name: str,",
+        "    class_name: Union[str, None] = None",
+        ") -> Type:",
+        "    \"\"\"",
+        "    Dynamically import a class from a specified module.",
+        "",
+        "    :param module_name: The name of the module to import.",
+        "    :param class_name: The name of the class in the module to import. Defaults to None.",
+        "    :return: The dynamically imported class.",
+        "    \"\"\"",
+        "    if not class_name:"
+      ]
+    },
+    {
+      "a": [
+        "        module_name, class_name = module_name.rsplit('.', 1)"
+      ],
+      "b": [
+        "        module_name, class_name = module_name.rsplit('.', 2)"
+      ],
+      "common": true
+    },
+    {
+      "ab": [
+        "    loaded_module = importclass(module_name, fromlist=[class_name])",
+        "    return getattr(loaded_module, class_name)",
+        "",
+        "",
+        "def preprocess_classes(func: Callable) -> Callable:",
+        "    \"\"\"Decorator to convert dot-notated class paths into strings from positional arguments.\"\"\"",
+        "    def __preprocess_classes_wrapper(*all_classes: TypeClassOrPath, **kwargs: Any) -> Any:",
+        "        \"\"\"",
+        "        Dynamically import classes if they are passed as strings.",
+        "",
+        "        :param all_classes: A variable number of class paths (strings or actual types).",
+        "        :param kwargs: Any keyword arguments to pass to the decorated function.",
+        "        :return: The result of the decorated function.",
+        "        \"\"\"",
+        "        classes_processed = (",
+        "            class_id if isinstance(class_id, type)",
+        "            else importclass(class_id)",
+        "            for class_id in all_classes",
+        "        )",
+        "        return func(*classes_processed, *kwargs),",
+        "    return __preprocess_classes_wrapper"
+      ]
+    }
+  ]
+}
\ No newline at end of file
diff --git a/src/test/resources/__files/stateful/gerritPatchSetFiles.json b/src/test/resources/__files/stateful/gerritPatchSetFiles.json
new file mode 100644
index 0000000..75320cb
--- /dev/null
+++ b/src/test/resources/__files/stateful/gerritPatchSetFiles.json
@@ -0,0 +1,8 @@
+{
+  "test_file_1.py":{
+    "old_mode":33188,
+    "new_mode":33188,
+    "lines_inserted":1,
+    "lines_deleted":1
+  }
+}
\ No newline at end of file