Fix guest chat trees (#956)

- Fix ChatML end marker
- Fix empty chat load caused by deleted message edge case
- Add more recommendations to presets
- Add missing placeholders to prompt help modal
Co-authored-by: Sceuick <[email protected]>

common/chat.ts

@@ -42,7 +42,7 @@ export function toChatGraph(messages: AppSchema.ChatMessage[]): { tree: ChatTree
   return { tree, root: messages[0]?._id || '' }
 }
 
-export function updateChatTree(tree: ChatTree, msg: AppSchema.ChatMessage) {
+export function updateChatTreeNode(tree: ChatTree, msg: AppSchema.ChatMessage) {
   tree[msg._id] = {
     msg,
     children: new Set(),
@@ -56,7 +56,19 @@ export function updateChatTree(tree: ChatTree, msg: AppSchema.ChatMessage) {
     parent.children.add(msg._id)
   }
 
-  return tree
+  return Object.assign({}, tree)
+}
+
+export function removeChatTreeNodes(tree: ChatTree, ids: string[]) {
+  for (const id of ids) {
+    const parent = tree[id]
+    if (parent) {
+      parent.children.delete(id)
+    }
+    delete tree[id]
+  }
+
+  return Object.assign({}, tree)
 }
 
 export function getChatDepths(tree: ChatTree) {

common/presets/templates.ts

@@ -45,7 +45,7 @@ export const BUILTIN_FORMATS: { [key in ModelFormat]: FormatTags } = {
   },
   ChatML: {
     openUser: '<|im_start|>user\n',
-    closeUser: '<|im_end>',
+    closeUser: '<|im_end|>',
     openBot: '<|im_start|>assistant\n',
     closeBot: '<|im_end|>',
     openSystem: '<|im_start|>system\n',

common/prompt-order.ts

@@ -39,6 +39,7 @@ export const formatHolders: Record<string, Record<string, string>> = {
     {{/each}}`,
     post: `### Response:\n{{post}}`,
     system_prompt: `{{#if system_prompt}}### Instruction:\n{{system_prompt}}{{/if}}`,
+    ujb: `{{#if ujb}}({{ujb}}){{/if}}`,
   },
   Vicuna: {
     preamble: neat`Below is an instruction that describes a task. Write a response that appropriately completes the request.\n
@@ -49,6 +50,7 @@ export const formatHolders: Record<string, Record<string, string>> = {
     {{/each}}`,
     post: `ASSISTANT: {{post}}`,
     system_prompt: `{{#if system_prompt}}SYSTEM: {{system_prompt}}{{/if}}`,
+    ujb: `{{#if ujb}}({{ujb}}){{/if}}`,
   },
   Mistral: {
     preamble: neat`Below is an instruction that describes a task. Write a response that appropriately completes the request.\n

common/types/schema.ts

@@ -137,7 +137,7 @@ export namespace AppSchema {
     level: number
     service: AIAdapter
     guidance: boolean
-    preset: GenSettings
+    preset: GenSettings & Pick<SubscriptionPreset, 'allowGuestUsage'>
   }
 
   export interface AppConfig {

common/util.ts

@@ -7,6 +7,11 @@ export function replace<T extends { _id: string }>(id: string, list: T[], item:
   return list.map((li) => (li._id === id ? { ...li, ...item } : li))
 }
 
+export function exclude<T extends { _id: string }>(list: T[], ids: string[]) {
+  const set = new Set(ids)
+  return list.filter((item) => !set.has(item._id))
+}
+
 export function findOne<T extends { _id: string }>(id: string, list: T[]): T | void {
   for (const item of list) {
     if (item._id === id) return item

srv/api/chat/message.ts

@@ -326,6 +326,7 @@ export const generateMessageV2 = handle(async (req, res) => {
   }
 
   const responseText = body.kind === 'continue' ? `${body.continuing.msg} ${generated}` : generated
+  const parent = getNewMessageParent(body, userMsg)
   const actions: AppSchema.ChatAction[] = []
   let treeLeafId = ''
 
@@ -363,7 +364,7 @@ export const generateMessageV2 = handle(async (req, res) => {
         meta,
         retries,
         event: undefined,
-        parent: userMsg?._id,
+        parent,
       })
 
       sendMany(members, {
@@ -387,7 +388,7 @@ export const generateMessageV2 = handle(async (req, res) => {
 
         await store.msgs.editMessage(body.replacing._id, {
           msg: responseText,
-          parent: body.parent,
+          parent,
           actions,
           adapter,
           meta,
@@ -419,7 +420,7 @@ export const generateMessageV2 = handle(async (req, res) => {
           meta,
           retries,
           event: undefined,
-          parent: body.parent,
+          parent,
         })
         treeLeafId = requestId
         sendMany(members, {
@@ -491,12 +492,14 @@ async function handleGuestGenerate(body: GenRequest, req: AppRequest, res: Respo
       characterId: body.impersonate?._id,
       ooc: body.kind === 'ooc',
       event: undefined,
+      parent: body.parent,
     })
   } else if (body.kind.startsWith('send-event:')) {
     newMsg = newMessage(v4(), chatId, body.text!, {
       characterId: replyAs?._id,
       ooc: false,
       event: getScenarioEventType(body.kind),
+      parent: body.parent,
     })
   }
 
@@ -571,6 +574,7 @@ async function handleGuestGenerate(body: GenRequest, req: AppRequest, res: Respo
 
   const characterId = body.kind === 'self' ? undefined : body.replyAs?._id || body.char?._id
   const senderId = body.kind === 'self' ? 'anon' : undefined
+  const parent = getNewMessageParent(body, newMsg)
 
   if (body.kind === 'retry' && body.replacing) {
     retries = [body.replacing.msg].concat(retries).concat(body.replacing.retries || [])
@@ -583,6 +587,7 @@ async function handleGuestGenerate(body: GenRequest, req: AppRequest, res: Respo
     meta,
     event: undefined,
     retries,
+    parent,
   })
 
   switch (body.kind) {
@@ -623,6 +628,7 @@ function newMessage(
     meta?: any
     event: undefined | AppSchema.ScenarioEventType
     retries?: string[]
+    parent?: string
   }
 ) {
   const userMsg: AppSchema.ChatMessage = {
@@ -674,3 +680,25 @@ async function ensureBotMembership(
   update.characters = characters
   await store.chats.update(chat._id, update)
 }
+
+function getNewMessageParent(body: GenRequest, userMsg: AppSchema.ChatMessage | undefined): string {
+  switch (body.kind) {
+    case 'summary':
+    case 'chat-query':
+      return ''
+
+    case 'retry':
+    case 'continue':
+      return body.parent || ''
+
+    case 'request':
+    case 'ooc':
+    case 'self':
+    case 'send':
+    case 'send-event:character':
+    case 'send-event:hidden':
+    case 'send-event:ooc':
+    case 'send-event:world':
+      return userMsg?._id || ''
+  }
+}

srv/api/chat/remove.ts

@@ -5,7 +5,7 @@ import { sendMany } from '../ws'
 
 export const deleteMessages = handle(async ({ body, params, userId }) => {
   const chatId = params.id
-  assertValid({ ids: ['string'] }, body)
+  assertValid({ ids: ['string'], leafId: 'string?' }, body)
 
   const chat = await store.chats.getChatOnly(chatId)
   if (!chat) {
@@ -17,6 +17,11 @@ export const deleteMessages = handle(async ({ body, params, userId }) => {
   }
 
   await store.msgs.deleteMessages(body.ids)
+
+  if (body.leafId) {
+    await store.chats.update(chatId, { treeLeafId: body.leafId })
+  }
+
   sendMany(chat.memberIds.concat(chat.userId), { type: 'messages-deleted', ids: body.ids })
   return { success: true }
 })

srv/app.ts

@@ -9,6 +9,7 @@ import { setupSockets } from './api/ws'
 import { config } from './config'
 import { createServer } from './server'
 import pipeline from './api/pipeline'
+import { getDb } from './db/client'
 
 export function createApp() {
   const upload = multer({
@@ -41,6 +42,17 @@ export function createApp() {
   app.use('/v1', keyedRouter)
   app.use('/api', api)
 
+  if (config.db.host || config.db.uri) {
+    app.get('/healthcheck', (_, res) => {
+      try {
+        getDb()
+        res.status(200).json({ message: 'okay', status: true })
+      } catch (ex) {
+        res.status(503).json({ message: 'Database not ready', status: false })
+      }
+    })
+  }
+
   if (config.pipelineProxy) {
     app.use('/pipeline', pipeline)
   }

srv/db/subscriptions.ts

@@ -86,6 +86,7 @@ export function getCachedSubscriptions(user?: AppSchema.User | null) {
       guidance: !!sub.guidanceCapable,
       preset: {
         name: sub.name,
+
         maxContextLength: sub.maxContextLength,
         maxTokens: sub.maxTokens,
         frequencyPenalty: sub.frequencyPenalty,
@@ -110,6 +111,15 @@ export function getCachedSubscriptions(user?: AppSchema.User | null) {
         encoderRepitionPenalty: sub.encoderRepitionPenalty,
         mirostatLR: sub.mirostatLR,
         mirostatTau: sub.mirostatTau,
+        allowGuestUsage: sub.allowGuestUsage,
+        dynatemp_exponent: sub.dynatemp_exponent,
+        dynatemp_range: sub.dynatemp_range,
+        modelFormat: sub.modelFormat,
+        penaltyAlpha: sub.penaltyAlpha,
+        smoothingCurve: sub.smoothingCurve,
+        smoothingFactor: sub.smoothingFactor,
+        skipSpecialTokens: sub.skipSpecialTokens,
+        tempLast: sub.tempLast,
       },
     }))
     .sort((l, r) => (l.level === r.level ? l.name.localeCompare(r.name) : l.level - r.level))

web/pages/Chat/ChatDetail.tsx

@@ -428,7 +428,7 @@ const ChatDetail: Component = () => {
               </div>
             </Show>
             <Show when={chats.loaded && chatMsgs().length === 0 && !msgs.waiting}>
-              <div class="flex justify-center">
+              <div class="flex justify-center gap-2">
                 <Button onClick={generateFirst}>Generate Message</Button>
               </div>
             </Show>

web/pages/Chat/components/GraphModal.tsx

@@ -37,8 +37,8 @@ export const ChatGraphModal: Component<{
       footer={
         <>
           <Button onClick={() => setShort(short() === 'full' ? 'short' : 'full')}>
-            <Show when={short() === 'short'} fallback={'Detailed'}>
-              Summary
+            <Show when={short() === 'short'} fallback={'Collapse Graph'}>
+              Expand Graph
             </Show>
           </Button>
           <Button onClick={reset}>Reset View</Button>

web/pages/Settings/Agnaistic.tsx

@@ -20,6 +20,7 @@ export const AgnaisticSettings: Component<{
 
     return config.subs
       .filter((sub) => sub.level <= level)
+      .filter((sub) => (level === -1 ? !!sub.preset.allowGuestUsage : true))
       .map((sub) => ({ label: sub.name, value: sub._id, level: sub.level }))
       .sort((l, r) => r.level - l.level)
   })