fix context pills (#1100)

Author: sceuick

common/presets.ts

@@ -102,6 +102,7 @@ export const presetValidator = {
   epsilonCutoff: 'number?',
   etaCutoff: 'number?',
   mirostatToggle: 'boolean?',
+  presetMode: ['simple', 'advanced', null],
 } as const
 
 const disabledValues: { [key in keyof GenMap]?: AppSchema.GenSettings[key] } = {

common/prompt.ts

@@ -283,7 +283,7 @@ export async function assemblePrompt(
   const template = getTemplate(opts)
 
   const history = { lines, order: 'asc' } as const
-  let { parsed, inserts, length, sections } = await injectPlaceholders(template, {
+  let { parsed, inserts, length, sections, linesAddedCount } = await injectPlaceholders(template, {
     opts,
     parts,
     history,
@@ -293,7 +293,16 @@ export async function assemblePrompt(
     jsonValues: opts.jsonValues,
   })
 
-  return { lines: history.lines, prompt: parsed, inserts, parts, post, length, sections }
+  return {
+    lines: history.lines,
+    prompt: parsed,
+    inserts,
+    parts,
+    post,
+    length,
+    sections,
+    linesAddedCount,
+  }
 }
 
 export function getTemplate(opts: Pick<GenerateRequestV2, 'settings' | 'chat'>) {

common/template-parser.ts

@@ -221,6 +221,13 @@ export async function parseTemplate(
   //   }
   // }
 
+  /**
+   * Some placeholders require re-parsing as they also contain placeholders
+   */
+  opts.isFinal = true
+  const result = render(output, opts).replace(/\r\n/g, '\n').replace(/\n\n+/g, '\n\n').trim()
+  opts.isFinal = false
+
   /** Replace iterators */
   let history: string[] = []
   if (opts.limit && opts.limit.output) {
@@ -236,7 +243,7 @@ export async function parseTemplate(
       })
       unusedTokens = filled.unusedTokens
       const trimmed = filled.adding.slice().reverse()
-      output = output.replace(id, trimmed.join('\n'))
+      output = output.replace(new RegExp(id, 'gi'), trimmed.join('\n'))
       linesAddedCount += filled.linesAddedCount
       history = trimmed
     }
@@ -256,10 +263,6 @@ export async function parseTemplate(
     }
   }
 
-  opts.isFinal = true
-  const result = render(output, opts).replace(/\r\n/g, '\n').replace(/\n\n+/g, '\n\n').trim()
-  opts.isFinal = false
-
   sections.sections.history = history
 
   // console.log(
@@ -273,8 +276,10 @@ export async function parseTemplate(
   //   sections.sections.post.join('')
   // )
 
+  output = output.replace(/\r\n/g, '\n').replace(/\n\n+/g, '\n\n').trim()
+
   return {
-    parsed: result,
+    parsed: output,
     inserts: opts.inserts ?? new Map(),
     length: await opts.limit?.encoder?.(result),
     linesAddedCount,

common/types/presets.ts

@@ -76,7 +76,7 @@ export interface UserGenPreset extends GenSettings {
 export interface GenSettings {
   name: string
   description?: string
-  presetMode?: 'simple' | 'advanced'
+  presetMode?: 'simple' | 'advanced' | undefined
 
   service?: AIAdapter
 

common/valid/validate.ts

@@ -61,21 +61,21 @@ export function isValidPartial<T extends Validator>(
 }
 
 export function validateBody<T extends Validator>(
-  type: T,
+  guard: T,
   compare: any,
   opts: { partial?: boolean; prefix?: string; notThrow?: boolean } = {}
-): { errors: string[]; actual: UnwrapBody<T> } {
+): { errors: string[]; actual: UnwrapBody<T>; original: UnwrapBody<T> } {
   const prefix = opts.prefix ? `${opts.prefix}.` : ''
   const errors: string[] = []
   const actual: any = {}
 
-  if (!compare && '?' in type && (type as any)['?'] === '?') {
-    return { errors, actual }
+  if (!compare && '?' in guard && (guard as any)['?'] === '?') {
+    return { errors, actual, original: compare }
   }
 
-  start: for (const key in type) {
+  start: for (const key in guard) {
     const prop = `${prefix}${key}`
-    const bodyType = type[key]
+    const bodyType = guard[key]
     let value
     try {
       value = compare?.[key]
@@ -246,5 +246,5 @@ export function validateBody<T extends Validator>(
     throw new Error(`Object does not match type: ${errors.join(', ')}`)
   }
 
-  return { errors, actual }
+  return { errors, actual, original: compare }
 }

srv/adapter/gemini.ts

@@ -1,12 +1,13 @@
 import needle from 'needle'
 import { decryptText } from '../db/util'
-import { getEncoderByName } from '../tokenize'
+import { getEncoder, getEncoderByName } from '../tokenize'
 import { toChatCompletionPayload } from './chat-completion'
 import { getStoppingStrings } from './prompt'
-import { ModelAdapter } from './type'
+import { AdapterProps, ModelAdapter } from './type'
 import { AppLog } from '../middleware'
 import { sanitise, sanitiseAndTrim, trimResponseV2 } from '/common/requests/util'
 import { requestStream } from './stream'
+import { injectPlaceholders } from '/common/prompt'
 
 const BASE_URL = `https://generativelanguage.googleapis.com/v1beta/models/`
 
@@ -34,12 +35,13 @@ export const handleGemini: ModelAdapter = async function* (opts) {
     },
   }
 
-  const systems: string[] = []
+  const fallback = await fallbackSystemMessage(opts)
+  const systems: string[] = [opts.parts.systemPrompt || fallback.parsed]
   const contents: any[] = []
 
   for (const msg of messages) {
     if (msg.role === 'system') {
-      systems.push(msg.content)
+      contents.push({ role: 'user', parts: [{ text: msg.content }] })
       continue
     }
 
@@ -51,11 +53,7 @@ export const handleGemini: ModelAdapter = async function* (opts) {
   if (systems.length) {
     if (!SYSTEM_INCAPABLE[opts.gen.googleModel]) {
       payload.system_instruction = {
-        parts: [
-          {
-            text: systems.join('\n'),
-          },
-        ],
+        parts: [{ text: systems.join('\n') }],
       }
     } else {
       contents.unshift({ role: 'user', parts: [{ text: systems.join('\n') }] })
@@ -224,3 +222,18 @@ const safetySettings = [
     threshold: 'BLOCK_NONE',
   },
 ]
+
+function fallbackSystemMessage(opts: AdapterProps) {
+  const message = injectPlaceholders(
+    `Write "{{char}}'s" next reply in a fictional roleplay chat between "{{user}}" and "{{char}}"`,
+    {
+      characters: opts.characters,
+      encoder: getEncoder('main').count,
+      jsonValues: {},
+      parts: opts.parts,
+      opts,
+    }
+  )
+
+  return message
+}

srv/adapter/generate.ts

@@ -64,6 +64,7 @@ export type InferenceRequest = {
   guest?: string
   user: AppSchema.User
   settings?: Partial<AppSchema.UserGenPreset>
+  maxKnownLines?: number
 
   guidance?: boolean
   placeholders?: any
@@ -394,6 +395,13 @@ export async function createChatStream(
    */
 
   const prompt = await assemblePrompt(opts, opts.parts, opts.lines, encoder)
+  if (prompt.linesAddedCount === 0 && opts.linesCount) {
+    throw new StatusError(
+      `Could not fit any messages in prompt. Check your character definition, context size, and template`,
+      400
+    )
+  }
+
   const messages = await toChatMessages(opts, prompt, encoder)
 
   const size = encoder(

srv/adapter/payloads.ts

@@ -102,6 +102,7 @@ function getBasePayload(opts: AdapterProps, stops: string[] = []) {
       skip_special_tokens: gen.skipSpecialTokens ?? true,
       stopping_strings: getStoppingStrings(opts, stops),
       dynamic_temperature: gen.dynatemp_range ? true : false,
+      smoothing_curve: gen.smoothingCurve,
       smoothing_factor: gen.smoothingFactor,
       token_healing: gen.tokenHealing,
       temp_last: gen.tempLast,

srv/adapter/type.ts

@@ -61,6 +61,7 @@ export type GenerateRequestV2 = {
 
   parts: PromptParts
   lines: string[]
+  linesCount?: number
   text?: string
   settings?: Partial<AppSchema.GenSettings>
   replacing?: AppSchema.ChatMessage

srv/api/chat/message.ts

@@ -66,6 +66,7 @@ const genValidator = {
     userEmbeds: 'any?',
   },
   lines: ['string'],
+  linesCount: 'number?',
   text: 'string?',
   settings: 'any?',
   lastMessage: 'string?',
@@ -247,15 +248,6 @@ export const generateMessageV2 = handle(async (req, res) => {
     })
   }
 
-  res.json({
-    requestId,
-    success: true,
-    generating: true,
-    message: 'Generating message',
-    messageId,
-    created: userMsg,
-  })
-
   const entities = await getResponseEntities(chat, body.sender.userId, body.settings)
   const schema = entities.gen.jsonSource === 'character' ? replyAs.json : entities.gen.json
   const hydrator = entities.gen.jsonEnabled && schema ? jsonHydrator(schema) : undefined
@@ -271,9 +263,10 @@ export const generateMessageV2 = handle(async (req, res) => {
   let probs: any
 
   if (body.response === undefined) {
-    const { stream, ...metadata } = await createChatStream(
+    const chatStream = await createChatStream(
       {
         ...body,
+        linesCount: body.linesCount,
         chat,
         replyAs,
         impersonate,
@@ -282,7 +275,22 @@ export const generateMessageV2 = handle(async (req, res) => {
         chatSchema: schema,
       },
       log
-    )
+    ).catch((err) => ({ err }))
+
+    if ('err' in chatStream) {
+      throw chatStream.err
+    }
+
+    res.json({
+      requestId,
+      success: true,
+      generating: true,
+      message: 'Generating message',
+      messageId,
+      created: userMsg,
+    })
+
+    const { stream, ...metadata } = chatStream
 
     adapter = metadata.adapter
 
@@ -584,8 +592,6 @@ async function handleGuestGenerate(body: GenRequest, req: AppRequest, res: Respo
     return { success: true }
   }
 
-  res.json({ success: true, generating: true, message: 'Generating message', requestId })
-
   const schema = body.settings.jsonSource === 'character' ? body.char.json : body.settings.json
   const hydrator = body.settings.jsonEnabled && schema ? jsonHydrator(schema) : undefined
   let generated = body.response || ''
@@ -597,12 +603,20 @@ async function handleGuestGenerate(body: GenRequest, req: AppRequest, res: Respo
   let jsonPartial: any
 
   if (body.response === undefined) {
-    const { stream, ...entities } = await createChatStream(
-      { ...body, chat, replyAs, requestId, chatSchema: schema },
+    const chatStream = await createChatStream(
+      { ...body, chat, replyAs, requestId, chatSchema: schema, linesCount: body.linesCount },
       log,
       guest
     )
 
+    if ('err' in chatStream) {
+      throw chatStream.err
+    }
+
+    const { stream, ...entities } = chatStream
+
+    res.json({ success: true, generating: true, message: 'Generating message', requestId })
+
     log.setBindings({ adapter })
 
     adapter = entities.adapter

srv/api/user/presets.ts

@@ -1,7 +1,7 @@
 import { assertValid } from '/common/valid'
 import { defaultPresets, presetValidator } from '../../../common/presets'
 import { store } from '../../db'
-import { StatusError, errors, handle } from '../wrap'
+import { StatusError, handle } from '../wrap'
 import { AIAdapter } from '../../../common/adapters'
 import { AppSchema } from '/common/types'
 import { toSamplerOrder } from '/common/sampler-order'
@@ -43,10 +43,7 @@ export const createUserPreset = handle(async ({ userId, body, authed }) => {
   }
 
   if (body.chatId) {
-    const res = await store.chats.getChat(body.chatId)
-    if (res?.chat.userId !== userId) {
-      throw errors.Forbidden
-    }
+    delete body.chatId
   }
 
   const samplers = toSamplerOrder(body.service, body.order, body.disabledSamplers)

tests/__snapshots__/prompt.spec.js.snap

@@ -220,7 +220,6 @@ Object {
   },
   "template": Object {
     "inserts": Map {},
-    "length": 38,
     "linesAddedCount": 2,
     "parsed": "GASLIGHT TEMPLATE
 ChatOwner

tests/prompt.spec.ts

@@ -171,6 +171,7 @@ This is how {{char}} should talk: {{example_dialogue}}`,
       },
     })
     delete (actual.template as any).sections
+    delete actual.template.length
     expect(actual).to.matchSnapshot()
   })