From e772860d18cc1c68eb2718f9b06963ca58ce567a Mon Sep 17 00:00:00 2001
From: somebody <onesome01@protonmail.com>
Date: Tue, 27 Dec 2022 17:33:46 -0600
Subject: [PATCH] Add endpoint to get action composition

---
 aiserver.py          | 23 +++++++++++++++++++++++
 koboldai_settings.py | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/aiserver.py b/aiserver.py
index ce1f372c..d35b917b 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -9106,6 +9106,29 @@ def UI_2_get_image_db():
     except FileNotFoundError:
         return jsonify([])
 
+@app.route("/action_composition.json", methods=["GET"])
+@logger.catch
+def UI_2_get_action_composition():
+    try:
+        actions = request.args.get("actions").split(",")
+        if not actions:
+            raise ValueError()
+    except (ValueError, AttributeError):
+        return "No actions", 400
+
+    try:
+        actions = [int(action) for action in actions]
+    except TypeError:
+        return "Not all actions int", 400
+
+    ret = []
+    for action_id in actions:
+        try:
+            ret.append(koboldai_vars.actions.get_action_composition(action_id))
+        except KeyError:
+            ret.append([])
+    return jsonify(ret)
+
 @app.route("/generated_images/<path:path>")
 def UI_2_send_generated_images(path):
     return send_from_directory(koboldai_vars.save_paths.generated_images, path)
diff --git a/koboldai_settings.py b/koboldai_settings.py
index 9adf1059..4490a009 100644
--- a/koboldai_settings.py
+++ b/koboldai_settings.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 from dataclasses import dataclass
+import difflib
 import importlib
 import os, re, time, threading, json, pickle, base64, copy, tqdm, datetime, sys
 import shutil
@@ -2106,6 +2107,37 @@ class KoboldStoryRegister(object):
             return filename, prompt
         return None, None
     
+    def get_action_composition(self, action_id: int) -> List[dict]:
+        """
+        Returns a list of chunks that comprise an action in dictionaries
+        formatted as follows:
+            type: string identifying chunk type ("ai", "user", or "edit")
+            content: the actual content of the chunk
+        """
+        current_text = self.actions[action_id]["Selected Text"]
+        action_original_type = "ai"
+
+        matching_blocks = difflib.SequenceMatcher(
+            None,
+            self.actions[action_id]["Original Text"],
+            current_text
+        ).get_matching_blocks()
+
+        chunks = []
+        base = 0
+        for chunk_match in matching_blocks:
+            inserted = current_text[base:chunk_match.b]
+            content = current_text[chunk_match.b:chunk_match.b + chunk_match.size]
+
+            base = chunk_match.b + chunk_match.size
+
+            if inserted:
+                chunks.append({"type": "edit", "content": inserted})
+            if content:
+                chunks.append({"type": action_original_type, "content": content})
+
+        return chunks
+
     def __setattr__(self, name, value):
         new_variable = name not in self.__dict__
         old_value = getattr(self, name, None)