From b61822131d64ab6313ac0ad1f70a3f97e449abb2 Mon Sep 17 00:00:00 2001
From: Gnome Ann <>
Date: Sat, 1 Jan 2022 12:46:37 -0500
Subject: [PATCH] Upload location scanner script
---
userscripts/kaipreset_location_scanner.lua | 208 +++++++++++++++++++++
1 file changed, 208 insertions(+)
create mode 100644 userscripts/kaipreset_location_scanner.lua
diff --git a/userscripts/kaipreset_location_scanner.lua b/userscripts/kaipreset_location_scanner.lua
new file mode 100644
index 00000000..dc9525fc
--- /dev/null
+++ b/userscripts/kaipreset_location_scanner.lua
@@ -0,0 +1,208 @@
+-- Location scanner
+-- Activates world info entries based on what the AI thinks the current location
+-- is.
+
+-- This file is part of KoboldAI.
+--
+-- KoboldAI is free software: you can redistribute it and/or modify
+-- it under the terms of the GNU Affero General Public License as published by
+-- the Free Software Foundation, either version 3 of the License, or
+-- (at your option) any later version.
+--
+-- This program is distributed in the hope that it will be useful,
+-- but WITHOUT ANY WARRANTY; without even the implied warranty of
+-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+-- GNU Affero General Public License for more details.
+--
+-- You should have received a copy of the GNU Affero General Public License
+-- along with this program. If not, see .
+
+kobold = require("bridge")() -- This line is optional and is only for EmmyLua type annotations
+local userscript = {} ---@class KoboldUserScript
+
+
+local example_config = [[;-- Location scanner
+;--
+;-- Usage instructions:
+;--
+;-- 1. Create a world info folder with name containing the string
+;-- "<||ls||>" (without the double quotes). The comment can be anything as
+;-- long as it contains that inside it somewhere -- for example, you could
+;-- set the name to "Locations <||ls||>".
+;--
+;-- 2. Create a non-selective, constant world info key _in that folder_ with key
+;-- "<||lslocation||>" (without the double quotes). Every once in a while,
+;-- this script will generate 20 tokens using "The current location"
+;-- as the submission and save the output into the <||lslocation||> entry.
+;--
+;-- 3. Put some other world info entries into the world info folder. These
+;-- entries will _only_ be triggered by the contents of the <||lslocation||>
+;-- entry and not by your story itself, or if it has constant key turned on.
+;--
+;-- You can edit some of the configuration values below to modify some of this
+;-- behaviour:
+;--
+return {
+ location_folder = "<||ls||>",
+ location_key = "<||lslocation||>",
+ submission = "\n\nThe current location:",
+ n_wait = 12, -- The script will run its extra generation every time your story grows by this many chunks.
+ n_tokens = 20, -- Number of tokens to generate in extra generation
+ singleline = true, -- true or false; true will result in the extra generation's output being cut off after the first line.
+ trim = true, -- true or false; true will result in the extra generation's output being cut off after the end of its last sentence.
+}
+]]
+
+local cfg ---@type table
+do
+ -- If config file is empty, write example config
+ local f = kobold.get_config_file()
+ f:seek("set")
+ if f:read(1) == nil then
+ f:write(example_config)
+ end
+ f:seek("set")
+ example_config = nil
+
+ -- Read config
+ local err
+ cfg, err = load(f:read("a"))
+ if err ~= nil then
+ error(err)
+ end
+ cfg = cfg()
+end
+
+
+local folder ---@type KoboldWorldInfoFolder|nil
+local entry ---@type KoboldWorldInfoEntry|nil
+local location = ""
+local orig_entry_map = {} ---@type table
+local repeated = false
+local last_quotient = math.huge
+
+local genamt = 0
+
+function userscript.inmod()
+ if repeated then
+ kobold.submission = cfg.submission
+ genamt = kobold.settings.genamt
+ kobold.settings.genamt = cfg.n_tokens
+ end
+
+ if entry == nil or folder == nil or not entry:is_valid() or not folder:is_valid() then
+ folder = nil
+ entry = nil
+ for i, f in ipairs(kobold.worldinfo.folders) do
+ if f.name:find(cfg.location_folder, 1, true) ~= nil then
+ folder = f
+ break
+ end
+ end
+ if folder ~= nil then
+ for i, e in ipairs(folder) do
+ if e.key:find(cfg.location_key, 1, true) ~= nil then
+ entry = e
+ break
+ end
+ end
+ end
+ end
+
+ orig_entry_map = {}
+
+ if entry ~= nil then
+ location = entry.content
+ entry.constant = false
+ end
+
+ if folder ~= nil then
+ for i, e in ipairs(folder) do
+ if entry == nil or e.uid ~= entry.uid then
+ orig_entry_map[e.uid] = {
+ constant = e.constant,
+ key = e.key,
+ keysecondary = e.keysecondary,
+ }
+ e.constant = not repeated and e:compute_context("") ~= e:compute_context(location)
+ e.key = ""
+ e.keysecondary = ""
+ end
+ end
+ end
+end
+
+function userscript.outmod()
+ if entry ~= nil and entry:is_valid() then
+ entry.constant = true
+ end
+
+ if repeated then
+ local output = kobold.outputs[1]
+ kobold.outputs[1] = ""
+
+ output = output:match("^%s*(.*)%s*$")
+
+ print("Extra generation result (prior to formatting): " .. output)
+
+ if cfg.singleline then
+ output = output:match("^[^\n]*")
+ end
+
+ if cfg.trim then
+ local i = 0
+ while true do
+ local j = output:find("[.?!)]", i + 1)
+ if j == nil then
+ break
+ end
+ i = j
+ end
+ if i > 0 then
+ if output:sub(i+1, i+1) == '"' then
+ i = i + 1
+ end
+ output = output:sub(1, i)
+ end
+ end
+
+ location = output
+ if entry ~= nil and entry:is_valid() then
+ entry.content = output
+ end
+
+ kobold.settings.genamt = genamt
+
+ for chunk in kobold.story:reverse_iter() do
+ if chunk.content ~= "" then
+ chunk.content = ""
+ break
+ end
+ end
+
+ print("Extra generation result (after formatting): " .. output)
+ end
+
+ local size = 0
+ for _ in kobold.story:forward_iter() do
+ size = size + 1
+ end
+
+ for uid, orig in pairs(orig_entry_map) do
+ for k, v in pairs(orig) do
+ kobold.worldinfo:finduid(uid)[k] = v
+ end
+ end
+
+ local quotient = math.floor(size / cfg.n_wait)
+ if repeated then
+ repeated = false
+ elseif quotient > last_quotient then
+ print("Running extra generation")
+ kobold.restart_generation()
+ repeated = true
+ end
+ last_quotient = quotient
+end
+
+return userscript