mincomk/risuai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: prevent lua engine getting killed and pool engines

Browse Source
This commit is contained in:
Sunho Kim
2024-07-15 21:11:21 -07:00
parent 7a542c14d6
commit cb9514c508
2 changed files with 493 additions and 384 deletions
Download Patch File Download Diff File Expand all files Collapse all files

88
src/ts/mutex.ts Normal file
View File

@@ -0,0 +1,88 @@
/**
* A lock for synchronizing async operations.
* Use this to protect a critical section
* from getting modified by multiple async operations
* at the same time.
*/
export class Mutex {
/**
* When multiple operations attempt to acquire the lock,
* this queue remembers the order of operations.
*/
private _queue: {
resolve: (release: ReleaseFunction) => void
}[] = []
private _isLocked = false
/**
* Wait until the lock is acquired.
* @returns A function that releases the acquired lock.
*/
acquire() {
return new Promise<ReleaseFunction>((resolve) => {
this._queue.push({resolve})
this._dispatch()
});
}
/**
* Enqueue a function to be run serially.
*
* This ensures no other functions will start running
* until `callback` finishes running.
* @param callback Function to be run exclusively.
* @returns The return value of `callback`.
*/
async runExclusive<T>(callback: () => Promise<T>) {
const release = await this.acquire()
try {
return await callback()
} finally {
release()
}
}
/**
* Check the availability of the resource
* and provide access to the next operation in the queue.
*
* _dispatch is called whenever availability changes,
* such as after lock acquire request or lock release.
*/
private _dispatch() {
if (this._isLocked) {
// The resource is still locked.
// Wait until next time.
return
}
const nextEntry = this._queue.shift()
if (!nextEntry) {
// There is nothing in the queue.
// Do nothing until next dispatch.
return
}
// The resource is available.
this._isLocked = true
// and give access to the next operation
// in the queue.
nextEntry.resolve(this._buildRelease())
}
/**
* Build a release function for each operation
* so that it can release the lock after
* the operation is complete.
*/
private _buildRelease(): ReleaseFunction {
return () => {
// Each release function make
// the resource available again
this._isLocked = false
// and call dispatch.
this._dispatch()
}
}
}
type ReleaseFunction = () => void

789
src/ts/process/lua.ts
View File

@@ -11,14 +11,21 @@ import type { OpenAIChat } from ".";
import { requestChatData } from "./request"; import { requestChatData } from "./request";
import { v4 } from "uuid"; import { v4 } from "uuid";
import { getModuleTriggers } from "./modules"; import { getModuleTriggers } from "./modules";
import { Mutex } from "../mutex";
let luaFactory:LuaFactory let luaFactory:LuaFactory
let luaEngine:LuaEngine
let lastCode = ''
let LuaSafeIds = new Set<string>() let LuaSafeIds = new Set<string>()
let LuaEditDisplayIds = new Set<string>() let LuaEditDisplayIds = new Set<string>()
let LuaLowLevelIds = new Set<string>() let LuaLowLevelIds = new Set<string>()
interface LuaEngineState {
code: string;
engine: LuaEngine;
mutex: Mutex;
}
let LuaEngines = new Map<string, LuaEngineState>()
export async function runLua(code:string, arg:{ export async function runLua(code:string, arg:{
char?:character|groupChat|simpleCharacterArgument, char?:character|groupChat|simpleCharacterArgument,
chat?:Chat chat?:Chat
@@ -37,412 +44,426 @@ export async function runLua(code:string, arg:{
let stopSending = false let stopSending = false
let lowLevelAccess = arg.lowLevelAccess ?? false let lowLevelAccess = arg.lowLevelAccess ?? false
if(!luaEngine || lastCode !== code){ if(!luaFactory){
if(luaEngine){ await makeLuaFactory()
luaEngine.global.close() }
let luaEngineState = LuaEngines.get(mode)
let wasEmpty = false
if (!luaEngineState) {
luaEngineState = {
code,
engine: await luaFactory.createEngine({injectObjects: true}),
mutex: new Mutex()
} }
if(!luaFactory){ LuaEngines.set(mode, luaEngineState)
makeLuaFactory() wasEmpty = true
} }
luaEngine = await luaFactory.createEngine({injectObjects: true}) return await luaEngineState.mutex.runExclusive(async () => {
luaEngine.global.set('setChatVar', (id:string,key:string, value:string) => { if (wasEmpty || code !== luaEngineState.code) {
if(!LuaSafeIds.has(id) && !LuaEditDisplayIds.has(id)){ if (!wasEmpty)
return luaEngineState.engine.global.close()
} luaEngineState.engine = await luaFactory.createEngine({injectObjects: true})
setVar(key, value) const luaEngine = luaEngineState.engine
}) luaEngine.global.set('setChatVar', (id:string,key:string, value:string) => {
luaEngine.global.set('getChatVar', (id:string,key:string) => { if(!LuaSafeIds.has(id) && !LuaEditDisplayIds.has(id)){
if(!LuaSafeIds.has(id) && !LuaEditDisplayIds.has(id)){ return
return
}
return getVar(key)
})
luaEngine.global.set('stopChat', (id:string) => {
if(!LuaSafeIds.has(id)){
return
}
stopSending = true
})
luaEngine.global.set('alertError', (id:string, value:string) => {
if(!LuaSafeIds.has(id)){
return
}
alertError(value)
})
luaEngine.global.set('alertNormal', (id:string, value:string) => {
if(!LuaSafeIds.has(id)){
return
}
alertNormal(value)
})
luaEngine.global.set('alertInput', (id:string, value:string) => {
if(!LuaSafeIds.has(id)){
return
}
return alertInput(value)
})
luaEngine.global.set('setChat', (id:string, index:number, value:string) => {
if(!LuaSafeIds.has(id)){
return
}
const message = chat.message?.at(index)
if(message){
message.data = value
}
CurrentChat.set(chat)
})
luaEngine.global.set('setChatRole', (id:string, index:number, value:string) => {
if(!LuaSafeIds.has(id)){
return
}
const message = chat.message?.at(index)
if(message){
message.role = value === 'user' ? 'user' : 'char'
}
CurrentChat.set(chat)
})
luaEngine.global.set('cutChat', (id:string, start:number, end:number) => {
if(!LuaSafeIds.has(id)){
return
}
chat.message = chat.message.slice(start,end)
CurrentChat.set(chat)
})
luaEngine.global.set('removeChat', (id:string, index:number) => {
if(!LuaSafeIds.has(id)){
return
}
chat.message.splice(index, 1)
CurrentChat.set(chat)
})
luaEngine.global.set('addChat', (id:string, role:string, value:string) => {
if(!LuaSafeIds.has(id)){
return
}
let roleData:'user'|'char' = role === 'user' ? 'user' : 'char'
chat.message.push({role: roleData, data: value})
CurrentChat.set(chat)
})
luaEngine.global.set('insertChat', (id:string, index:number, role:string, value:string) => {
if(!LuaSafeIds.has(id)){
return
}
let roleData:'user'|'char' = role === 'user' ? 'user' : 'char'
chat.message.splice(index, 0, {role: roleData, data: value})
CurrentChat.set(chat)
})
luaEngine.global.set('removeChat', (id:string, index:number) => {
if(!LuaSafeIds.has(id)){
return
}
chat.message.splice(index, 1)
CurrentChat.set(chat)
})
luaEngine.global.set('getChatLength', (id:string) => {
if(!LuaSafeIds.has(id)){
return
}
return chat.message.length
})
luaEngine.global.set('getFullChatMain', (id:string) => {
const data = JSON.stringify(chat.message.map((v) => {
return {
role: v.role,
data: v.data
}
}))
return data
})
luaEngine.global.set('setFullChatMain', (id:string, value:string) => {
const realValue = JSON.parse(value)
if(!LuaSafeIds.has(id)){
return
}
chat.message = realValue.map((v) => {
return {
role: v.role,
data: v.data
} }
setVar(key, value)
})
luaEngine.global.set('getChatVar', (id:string,key:string) => {
if(!LuaSafeIds.has(id) && !LuaEditDisplayIds.has(id)){
return
}
return getVar(key)
})
luaEngine.global.set('stopChat', (id:string) => {
if(!LuaSafeIds.has(id)){
return
}
stopSending = true
})
luaEngine.global.set('alertError', (id:string, value:string) => {
if(!LuaSafeIds.has(id)){
return
}
alertError(value)
})
luaEngine.global.set('alertNormal', (id:string, value:string) => {
if(!LuaSafeIds.has(id)){
return
}
alertNormal(value)
})
luaEngine.global.set('alertInput', (id:string, value:string) => {
if(!LuaSafeIds.has(id)){
return
}
return alertInput(value)
})
luaEngine.global.set('setChat', (id:string, index:number, value:string) => {
if(!LuaSafeIds.has(id)){
return
}
const message = chat.message?.at(index)
if(message){
message.data = value
}
CurrentChat.set(chat)
})
luaEngine.global.set('setChatRole', (id:string, index:number, value:string) => {
if(!LuaSafeIds.has(id)){
return
}
const message = chat.message?.at(index)
if(message){
message.role = value === 'user' ? 'user' : 'char'
}
CurrentChat.set(chat)
})
luaEngine.global.set('cutChat', (id:string, start:number, end:number) => {
if(!LuaSafeIds.has(id)){
return
}
chat.message = chat.message.slice(start,end)
CurrentChat.set(chat)
})
luaEngine.global.set('removeChat', (id:string, index:number) => {
if(!LuaSafeIds.has(id)){
return
}
chat.message.splice(index, 1)
CurrentChat.set(chat)
})
luaEngine.global.set('addChat', (id:string, role:string, value:string) => {
if(!LuaSafeIds.has(id)){
return
}
let roleData:'user'|'char' = role === 'user' ? 'user' : 'char'
chat.message.push({role: roleData, data: value})
CurrentChat.set(chat)
})
luaEngine.global.set('insertChat', (id:string, index:number, role:string, value:string) => {
if(!LuaSafeIds.has(id)){
return
}
let roleData:'user'|'char' = role === 'user' ? 'user' : 'char'
chat.message.splice(index, 0, {role: roleData, data: value})
CurrentChat.set(chat)
})
luaEngine.global.set('removeChat', (id:string, index:number) => {
if(!LuaSafeIds.has(id)){
return
}
chat.message.splice(index, 1)
CurrentChat.set(chat)
})
luaEngine.global.set('getChatLength', (id:string) => {
if(!LuaSafeIds.has(id)){
return
}
return chat.message.length
})
luaEngine.global.set('getFullChatMain', (id:string) => {
const data = JSON.stringify(chat.message.map((v) => {
return {
role: v.role,
data: v.data
}
}))
return data
}) })
CurrentChat.set(chat)
})
luaEngine.global.set('logMain', (value:string) => { luaEngine.global.set('setFullChatMain', (id:string, value:string) => {
console.log(JSON.parse(value)) const realValue = JSON.parse(value)
}) if(!LuaSafeIds.has(id)){
return
}
chat.message = realValue.map((v) => {
return {
role: v.role,
data: v.data
}
})
CurrentChat.set(chat)
})
//Low Level Access luaEngine.global.set('logMain', (value:string) => {
luaEngine.global.set('similarity', async (id:string, source:string, value:string[]) => { console.log(JSON.parse(value))
if(!LuaLowLevelIds.has(id)){ })
return
}
const processer = new HypaProcesser('MiniLM')
await processer.addText(value)
return await processer.similaritySearch(source)
})
luaEngine.global.set('generateImage', async (id:string, value:string, negValue:string = '') => { //Low Level Access
if(!LuaLowLevelIds.has(id)){ luaEngine.global.set('similarity', async (id:string, source:string, value:string[]) => {
return if(!LuaLowLevelIds.has(id)){
} return
const gen = await generateAIImage(value, char as character, negValue, 'inlay') }
if(!gen){ const processer = new HypaProcesser('MiniLM')
return 'Error: Image generation failed' await processer.addText(value)
} return await processer.similaritySearch(source)
const imgHTML = new Image() })
imgHTML.src = gen
const inlay = await writeInlayImage(imgHTML)
return `{{inlay::${inlay}}}`
})
luaEngine.global.set('LLMMain', async (id:string, promptStr:string) => { luaEngine.global.set('generateImage', async (id:string, value:string, negValue:string = '') => {
let prompt:{ if(!LuaLowLevelIds.has(id)){
role: string, return
content: string }
}[] = JSON.parse(promptStr) const gen = await generateAIImage(value, char as character, negValue, 'inlay')
if(!LuaLowLevelIds.has(id)){ if(!gen){
return return 'Error: Image generation failed'
} }
let promptbody:OpenAIChat[] = prompt.map((dict) => { const imgHTML = new Image()
let role:'system'|'user'|'assistant' = 'assistant' imgHTML.src = gen
switch(dict['role']){ const inlay = await writeInlayImage(imgHTML)
case 'system': return `{{inlay::${inlay}}}`
case 'sys': })
role = 'system'
break luaEngine.global.set('LLMMain', async (id:string, promptStr:string) => {
case 'user': let prompt:{
role = 'user' role: string,
break content: string
case 'assistant': }[] = JSON.parse(promptStr)
case 'bot': if(!LuaLowLevelIds.has(id)){
case 'char':{ return
role = 'assistant' }
break let promptbody:OpenAIChat[] = prompt.map((dict) => {
let role:'system'|'user'|'assistant' = 'assistant'
switch(dict['role']){
case 'system':
case 'sys':
role = 'system'
break
case 'user':
role = 'user'
break
case 'assistant':
case 'bot':
case 'char':{
role = 'assistant'
break
}
}
return {
content: dict['content'] ?? '',
role: role,
}
})
const result = await requestChatData({
formated: promptbody,
bias: {},
useStreaming: false,
noMultiGen: true,
}, 'model')
if(result.type === 'fail'){
return JSON.stringify({
success: false,
result: 'Error: ' + result.result
})
}
if(result.type === 'streaming' || result.type === 'multiline'){
return JSON.stringify({
success: false,
result: result.result
})
}
return JSON.stringify({
success: true,
result: result.result
})
})
luaEngine.global.set('simpleLLM', async (id:string, prompt:string) => {
if(!LuaLowLevelIds.has(id)){
return
}
const result = await requestChatData({
formated: [{
role: 'user',
content: prompt
}],
bias: {},
useStreaming: false,
noMultiGen: true,
}, 'model')
if(result.type === 'fail'){
return {
success: false,
result: 'Error: ' + result.result
}
}
if(result.type === 'streaming' || result.type === 'multiline'){
return {
success: false,
result: result.result
} }
} }
return { return {
content: dict['content'] ?? '', success: true,
role: role,
}
})
const result = await requestChatData({
formated: promptbody,
bias: {},
useStreaming: false,
noMultiGen: true,
}, 'model')
if(result.type === 'fail'){
return JSON.stringify({
success: false,
result: 'Error: ' + result.result
})
}
if(result.type === 'streaming' || result.type === 'multiline'){
return JSON.stringify({
success: false,
result: result.result
})
}
return JSON.stringify({
success: true,
result: result.result
})
})
luaEngine.global.set('simpleLLM', async (id:string, prompt:string) => {
if(!LuaLowLevelIds.has(id)){
return
}
const result = await requestChatData({
formated: [{
role: 'user',
content: prompt
}],
bias: {},
useStreaming: false,
noMultiGen: true,
}, 'model')
if(result.type === 'fail'){
return {
success: false,
result: 'Error: ' + result.result
}
}
if(result.type === 'streaming' || result.type === 'multiline'){
return {
success: false,
result: result.result result: result.result
} }
} })
luaEngine.global.set('getName', async (id:string) => {
if(!LuaSafeIds.has(id)){
return
}
const db = get(DataBase)
const selectedChar = get(selectedCharID)
const char = db.characters[selectedChar]
return char.name
})
return { luaEngine.global.set('setName', async (id:string, name:string) => {
success: true, if(!LuaSafeIds.has(id)){
result: result.result return
} }
}) const db = get(DataBase)
const selectedChar = get(selectedCharID)
luaEngine.global.set('getName', async (id:string) => { if(typeof name !== 'string'){
if(!LuaSafeIds.has(id)){ throw('Invalid data type')
return }
} db.characters[selectedChar].name = name
const db = get(DataBase) setDatabase(db)
const selectedChar = get(selectedCharID) })
const char = db.characters[selectedChar]
return char.name
})
luaEngine.global.set('setName', async (id:string, name:string) => { luaEngine.global.set('setDescription', async (id:string, desc:string) => {
if(!LuaSafeIds.has(id)){ if(!LuaSafeIds.has(id)){
return return
} }
const db = get(DataBase) const db = get(DataBase)
const selectedChar = get(selectedCharID) const selectedChar = get(selectedCharID)
if(typeof name !== 'string'){ const char =db.characters[selectedChar]
throw('Invalid data type') if(typeof data !== 'string'){
} throw('Invalid data type')
db.characters[selectedChar].name = name }
setDatabase(db) if(char.type === 'group'){
}) throw('Character is a group')
}
char.desc = desc
db.characters[selectedChar] = char
setDatabase(db)
})
luaEngine.global.set('setDescription', async (id:string, desc:string) => { luaEngine.global.set('setCharacterFirstMessage', async (id:string, data:string) => {
if(!LuaSafeIds.has(id)){ if(!LuaSafeIds.has(id)){
return return
} }
const db = get(DataBase) const db = get(DataBase)
const selectedChar = get(selectedCharID) const selectedChar = get(selectedCharID)
const char =db.characters[selectedChar] const char = db.characters[selectedChar]
if(typeof data !== 'string'){ if(typeof data !== 'string'){
throw('Invalid data type') return false
} }
if(char.type === 'group'){ char.firstMessage = data
throw('Character is a group') db.characters[selectedChar] = char
} setDatabase(db)
char.desc = desc return true
db.characters[selectedChar] = char })
setDatabase(db)
})
luaEngine.global.set('setCharacterFirstMessage', async (id:string, data:string) => { luaEngine.global.set('getCharacterFirstMessage', async (id:string) => {
if(!LuaSafeIds.has(id)){ if(!LuaSafeIds.has(id)){
return return
} }
const db = get(DataBase) const db = get(DataBase)
const selectedChar = get(selectedCharID) const selectedChar = get(selectedCharID)
const char = db.characters[selectedChar] const char = db.characters[selectedChar]
if(typeof data !== 'string'){ return char.firstMessage
return false })
}
char.firstMessage = data
db.characters[selectedChar] = char
setDatabase(db)
return true
})
luaEngine.global.set('getCharacterFirstMessage', async (id:string) => { luaEngine.global.set('getBackgroundEmbedding', async (id:string) => {
if(!LuaSafeIds.has(id)){ if(!LuaSafeIds.has(id)){
return return
} }
const db = get(DataBase) const db = get(DataBase)
const selectedChar = get(selectedCharID) const selectedChar = get(selectedCharID)
const char = db.characters[selectedChar] const char = db.characters[selectedChar]
return char.firstMessage return char.backgroundHTML
}) })
luaEngine.global.set('getBackgroundEmbedding', async (id:string) => { luaEngine.global.set('setBackgroundEmbedding', async (id:string, data:string) => {
if(!LuaSafeIds.has(id)){ if(!LuaSafeIds.has(id)){
return return
} }
const db = get(DataBase) const db = get(DataBase)
const selectedChar = get(selectedCharID) const selectedChar = get(selectedCharID)
const char = db.characters[selectedChar] if(typeof data !== 'string'){
return char.backgroundHTML return false
}) }
db.characters[selectedChar].backgroundHTML = data
setDatabase(db)
return true
})
luaEngine.global.set('setBackgroundEmbedding', async (id:string, data:string) => { await luaEngine.doString(luaCodeWarper(code))
if(!LuaSafeIds.has(id)){ luaEngineState.code = code
return
}
const db = get(DataBase)
const selectedChar = get(selectedCharID)
if(typeof data !== 'string'){
return false
}
db.characters[selectedChar].backgroundHTML = data
setDatabase(db)
return true
})
await luaEngine.doString(luaCodeWarper(code))
lastCode = code
}
let accessKey = v4()
if(mode === 'editDisplay'){
LuaEditDisplayIds.add(accessKey)
}
else{
LuaSafeIds.add(accessKey)
if(lowLevelAccess){
LuaLowLevelIds.add(accessKey)
} }
} let accessKey = v4()
let res:any if(mode === 'editDisplay'){
try { LuaEditDisplayIds.add(accessKey)
switch(mode){ }
case 'input':{ else{
const func = luaEngine.global.get('onInput') LuaSafeIds.add(accessKey)
if(func){ if(lowLevelAccess){
res = await func(accessKey) LuaLowLevelIds.add(accessKey)
} }
} }
case 'output':{ let res:any
const func = luaEngine.global.get('onOutput') const luaEngine = luaEngineState.engine
if(func){ try {
res = await func(accessKey) switch(mode){
} case 'input':{
} const func = luaEngine.global.get('onInput')
case 'start':{ if(func){
const func = luaEngine.global.get('onStart') res = await func(accessKey)
if(func){ }
res = await func(accessKey) }
} case 'output':{
} const func = luaEngine.global.get('onOutput')
case 'editRequest': if(func){
case 'editDisplay': res = await func(accessKey)
case 'editInput': }
case 'editOutput':{ }
const func = luaEngine.global.get('callListenMain') case 'start':{
if(func){ const func = luaEngine.global.get('onStart')
res = await func(mode, accessKey, JSON.stringify(data)) if(func){
res = JSON.parse(res) res = await func(accessKey)
} }
} }
default:{ case 'editRequest':
const func = luaEngine.global.get(mode) case 'editDisplay':
if(func){ case 'editInput':
res = await func(accessKey) case 'editOutput':{
} const func = luaEngine.global.get('callListenMain')
} if(func){
} res = await func(mode, accessKey, JSON.stringify(data))
if(res === false){ res = JSON.parse(res)
stopSending = true }
}
default:{
const func = luaEngine.global.get(mode)
if(func){
res = await func(accessKey)
}
}
}
if(res === false){
stopSending = true
}
} catch (error) {
console.error(error)
} }
} catch (error) {
console.error(error)
}
LuaSafeIds.delete(accessKey) LuaSafeIds.delete(accessKey)
LuaLowLevelIds.delete(accessKey) LuaLowLevelIds.delete(accessKey)
return { return {
stopSending, chat, res stopSending, chat, res
} }
})
} }
async function makeLuaFactory(){ async function makeLuaFactory(){