And the rest

aiserver.py

@@ -9433,41 +9433,65 @@ def UI_2_generate_wi(data):
     field = data["field"]
     existing = data["existing"]
     gen_amount = data["genAmount"]
-    print(uid, field)
 
-    with open("data/wi_fewshot.txt", "r") as file:
-        fewshot_template = file.read()
-    
+    # The template to coax what we want from the model
+    extractor_string = ""
+
     if field == "title":
-        prompt = fewshot_template
         for thing in ["type", "desc"]:
             if not existing[thing]:
                 continue
             pretty = {"type": "Type", "desc": "Description"}[thing]
-            prompt += f"{pretty}: {existing[thing]}\n"
+            extractor_string += f"{pretty}: {existing[thing]}\n"
         
         pretty = "Title"
         if existing["desc"]:
             # Don't let the model think we're starting a new entry
             pretty = "Alternate Title"
 
-        prompt += pretty + ":"
+        extractor_string += pretty + ":"
     elif field == "desc":
         # MUST be title and type
         assert existing["title"]
         assert existing["type"]
-        prompt = f"{fewshot_template}Title: {existing['title']}\nType: {existing['type']}\nDescription:"
+        extractor_string = f"Title: {existing['title']}\nType: {existing['type']}\nDescription:"
     else:
         assert False, "What"
+
+    with open("data/wi_fewshot.txt", "r") as file:
+        fewshot_entries = [x.strip() for x in file.read().split("\n\n") if x]
+
+    # Use user's own WI entries in prompt
+    if koboldai_vars.wigen_use_own_wi:
+        fewshot_entries += koboldai_vars.worldinfo_v2.to_wi_fewshot_format(excluding_uid=uid)
     
-    logger.info(prompt)
+    # We must have this amount or less in our context.
+    target = koboldai_vars.max_length - gen_amount - len(tokenizer.encode(extractor_string))
+
+    used = []
+    # Walk the entries backwards until we can't cram anymore in
+    for entry in reversed(fewshot_entries):
+        maybe = [entry] + used
+        maybe_str = "\n\n".join(maybe)
+        possible_encoded = tokenizer.encode(maybe_str)
+        if len(possible_encoded) > target:
+            break
+        yes_str = maybe_str
+        used = maybe
+    
+    prompt = f"{yes_str}\n\n{extractor_string}"
+    
+    # logger.info(prompt)
+    # TODO: Make single_line mode that stops on newline rather than bans it (for title)
     out_text = tpool.execute(
         raw_generate,
         prompt,
         max_new=gen_amount,
         single_line=True,
     ).decoded[0]
-    print(f'{out_text}')
+    out_text = utils.trimincompletesentence(out_text.strip())
+
+    socketio.emit("generated_wi", {"uid": uid, "field": field, "out": out_text}, room="UI_2")
 
 @app.route("/generate_raw", methods=["GET"])
 def UI_2_generate_raw():

gensettings.py

@@ -880,6 +880,20 @@ gensettingstf = [
     "extra_classes": "simple_ui_only var_sync_alt_user_ui_level",
     "ui_level": 0
 	},
+    {
+    "UI_V2_Only": True,
+	"uitype": "toggle",
+	"unit": "bool",
+	"label": "Native WI Gen",
+	"id": "wigen_use_own_wi", 
+	"default": False,
+    "tooltip": "Uses your existing applicable (has title, type, content) World Info entries as inspiration for generated ones.",
+    "menu_path": "World Info",
+    "sub_path": "",
+    "classname": "user",
+    "name": "wigen_use_own_wi",
+    "ui_level": 2
+	},
 ]
 
 gensettingsik =[{

koboldai_settings.py

@@ -1152,6 +1152,7 @@ class user_settings(settings):
         self.img_gen_steps = 30
         self.img_gen_cfg_scale = 7.0
         self.cluster_requested_models = [] # The models which we allow to generate during cluster mode
+        self.wigen_use_own_wi = False
         
         
     def __setattr__(self, name, value):
@@ -2471,6 +2472,38 @@ class KoboldWorldInfo(object):
     
     def get_used_wi(self):
         return [x['content'] for x in self.world_info if x['used_in_game']]
+    
+    def to_wi_fewshot_format(self, excluding_uid: int) -> List[str]:
+        """
+        Returns a list of strings representing applicable (has title, text, and
+        object type) World Info entries. Intended for feeding into the fewshot
+        generator.
+        """
+        the_collection = []
+        for entry in self.world_info.values():
+            if entry["uid"] == excluding_uid:
+                continue
+
+            if not (
+                entry["title"]
+                and entry["manual_text"]
+                and entry["object_type"]
+            ):
+                continue
+                
+            processed_desc = entry["manual_text"].replace("\n", " ")
+            while "  " in processed_desc:
+                processed_desc = processed_desc.replace("  ", " ")
+            processed_desc = processed_desc.strip()
+
+            the_collection.append(
+                f"Title: {entry['title']}\n" \
+                f"Type: {entry['object_type']}\n"
+                f"Description: {processed_desc}"
+            )
+        
+        return the_collection
+
 
 @dataclass
 class SavePaths:

static/koboldai.css

@@ -1100,6 +1100,25 @@ td.server_vars {
 .world_info_title {
 	user-select: text;
 	-moz-user-select:text;
+	margin-bottom: 0px;
+}
+
+.wi-type-leadup {
+	opacity: 0.4;
+}
+
+.world_info_item_type {
+	/* Workaround for Firefox bug: https://stackoverflow.com/a/42283543 */
+	padding: 1px;
+
+	outline: none;
+	border-bottom: 1px solid gray;
+	transition: color 200ms, border-bottom-color 200ms;
+}
+
+.world_info_item_type.bad-input {
+	color: rgb(255, 103, 103);
+	border-bottom-color: red;
 }
 
 .world_info_delete {
@@ -1140,6 +1159,24 @@ td.server_vars {
 	filter: brightness(90%);
 }
 
+.world_info_label_container {
+	display: flex;
+	justify-content: space-between;
+	margin: 1px 5px;
+}
+
+.world_info_label_container > .wi-ui-label {
+	position: relative;
+	top: 2px;
+}
+
+.world_info_label_container > .generate-button {
+	opacity: 0.8;
+	cursor: pointer;
+}
+
+.world_info_label_container > .generate-button:hover { opacity: 1.0; }
+
 .tag {
 	background-color: var(--wi_tag_color);
 	color: var(--wi_tag_text_color);

static/koboldai.js

@@ -34,6 +34,7 @@ socket.on("log_message", function(data){process_log_message(data);});
 socket.on("debug_message", function(data){console.log(data);});
 socket.on("scratchpad_response", recieveScratchpadResponse);
 socket.on("show_error_notification", function(data) { reportError(data.title, data.text) });
+socket.on("generated_wi", showGeneratedWIData);
 //socket.onAny(function(event_name, data) {console.log({"event": event_name, "class": data.classname, "data": data});});
 
 // Must be done before any elements are made; we track their changes.
@@ -69,6 +70,7 @@ var setup_wi_toggles = [];
 var scroll_trigger_element = undefined; //undefined means not currently set. If set to null, it's disabled.
 var drag_id = null;
 var story_commentary_characters = {};
+var generating_summary = false;
 const on_colab = $el("#on_colab").textContent == "true";
 
 // Each entry into this array should be an object that looks like:
@@ -2248,6 +2250,48 @@ function world_info_entry(data) {
 		document.getElementById("world_info_tags_"+data.uid).classList.remove("hidden");
 		document.getElementById("world_info_secondtags_"+data.uid).classList.remove("hidden");
 	}
+
+	const genTypeInput = world_info_card.querySelector(".world_info_item_type");
+	const generateDescButton = world_info_card.querySelector(".wi-lc-text > .generate-button");
+	generateDescButton.addEventListener("click", function() {
+		if (generating_summary) return;
+		let type = genTypeInput.innerText;
+
+		if (!type) {
+			genTypeInput.classList.add("bad-input");
+			return;
+		} else {
+			genTypeInput.classList.remove("bad-input");
+		}
+
+		// TODO: Make type input element
+		generateWIData(data.uid, "desc", title.innerText, type, null, 80);
+		this.innerText = "autorenew";
+		this.classList.add("spinner");
+	});
+
+	genTypeInput.addEventListener("focus", function() {
+		this.classList.remove("bad-input");
+	});
+
+	genTypeInput.addEventListener("keydown", function(event) {
+		if (event.key === "Enter") {
+			event.preventDefault();
+			this.blur();
+		}
+	});
+
+	genTypeInput.addEventListener("blur", function() {
+		this.innerText = this.innerText.trim();
+
+		if (this.innerText == this.getAttribute("old-text")) return;
+		this.setAttribute("old-text", this.innerText);
+
+		world_info_data[data.uid].object_type = this.innerText;
+		send_world_info(data.uid);
+	});
+
+	genTypeInput.innerText = data.object_type;
 	
 	//$('#world_info_constant_'+data.uid).bootstrapToggle();
 	//$('#world_info_wpp_toggle_'+data.uid).bootstrapToggle();
@@ -3957,6 +4001,7 @@ function create_new_wi_entry(folder) {
                                     "selective": false,
 									"wpp": {'name': "", 'type': "", 'format': 'W++', 'attributes': {}},
 									'use_wpp': false,
+									"object_type": null,
                                     };
 	var card = world_info_entry(data);
 	//card.scrollIntoView(false);
@@ -6894,10 +6939,30 @@ for (const proxy of document.querySelectorAll("[sync-proxy-host]")) {
 }
 
 function generateWIData(uid, field, title=null, type=null, desc=null, genAmount=80) {
+	if (generating_summary) return;
+	generating_summary = true;
+
 	socket.emit("generate_wi", {
 		uid: uid,
 		field: field,
 		genAmount: genAmount || 80,
 		existing: {title: title, type: type, desc: desc}
 	});
+}
+
+function showGeneratedWIData(data) {
+	generating_summary = false;
+	console.warn(data)
+	const card = $el(`.world_info_card[uid="${data.uid}"]`);
+
+	// Stop spinning!
+	for (const littleRobotFriend of card.querySelectorAll(".generate-button.spinner")) {
+		littleRobotFriend.classList.remove("spinner");
+		littleRobotFriend.innerText = "smart_toy";
+	}
+
+	if (data.field === "desc") {
+		world_info_data[data.uid].manual_text = data.out;
+		send_world_info(data.uid);
+	}
 }

templates/templates.html

@@ -10,8 +10,22 @@
 				>add</span>
 				<img class="world_info_image">
 			</div>
-			<h4 draggable="true" ondragstart="event.preventDefault();event.stopPropagation();" id="world_info_title_"
-				class="world_info_title" onfocus="this.parentElement.parentElement.setAttribute('draggable', 'false');this.setAttribute('draggable', 'false');"></h4>
+			<div>
+				<h4
+					draggable="true"
+					ondragstart="event.preventDefault();event.stopPropagation();"
+					id="world_info_title_"
+					class="world_info_title"
+					onfocus="this.parentElement.parentElement.setAttribute('draggable', 'false');this.setAttribute('draggable', 'false');"
+				></h4>
+				<span class="wi-type-leadup">is a</span>
+				<span
+					class="world_info_item_type"
+					contenteditable="true"
+					data-placeholder="Person"
+					spellcheck="false"
+				></span>
+			</div>
 		</div>
 		<span id="world_info_delete_" class="world_info_delete">X</span>
 	</div>
@@ -61,7 +75,12 @@
 		<div class="wpp_attributes_area"></div>
 	</div>
 	<div id="world_info_basic_text_" class="world_info_basic_text_area">
-		Text:
+		<div class="world_info_label_container wi-lc-text">
+			<span class="wi-ui-label">Text</span>
+			<span class="material-icons-outlined generate-button" tooltip="Generate Content">
+				smart_toy
+			</span>
+		</div>
 		<textarea draggable="true" ondragstart="event.preventDefault();event.stopPropagation();" id="world_info_entry_text_" class="world_info_text world_info_entry_text fullwidth" oninput="autoResize(this, 60)" 
 		onfocus="this.parentElement.parentElement.setAttribute('draggable', 'false');this.setAttribute('draggable', 'false');autoResize(this, 60)" 
 		onblur="this.parentElement.parentElement.setAttribute('draggable', 'true');this.setAttribute('draggable', 'true');this.style.height='60px';"></textarea>
@@ -81,6 +100,10 @@
 		</div>
 	</div>
 	<div>
+		<div class="world_info_label_container wi-lc-comment">
+			<span class="wi-ui-label">Comment</span>
+			<padding></padding>
+		</div>
 		<textarea draggable="true" ondragstartondragstart="event.preventDefault();event.stopPropagation();" rows=1 id="world_info_comment_" 
 				  class="world_info_text world_info_comment fullwidth" oninput="autoResize(this, 60)" onfocus="autoResize(this, 60)" onblur="this.style.height='60px';"
 				  onfocus="this.parentElement.parentElement.setAttribute('draggable', 'false');this.setAttribute('draggable', 'false');;"