Retrieve history from db when generating prompt

Author: sceuick

.dockerignore

@@ -0,0 +1,5 @@
+db/
+.parcel-cache
+node_modules
+.env
+.token_secret

Dockerfile

@@ -0,0 +1,20 @@
+FROM node:18-alpine
+
+WORKDIR /app
+VOLUME [ "/app/db" ]
+
+RUN npm install pnpm -g
+
+ADD package.json pnpm-lock.yaml ./
+RUN pnpm i --frozen-lockfile
+
+ADD public.ts requirements.txt tailwind.config.js tsconfig.json .babelrc .postcssrc .prettierrc ./
+ADD common/ ./common/
+ADD srv/ ./srv/
+ADD web/ ./web
+
+
+EXPOSE 3001
+EXPOSE 5001
+
+CMD ["pnpm", "start"]

package.json

@@ -1,5 +1,5 @@
 {
-  "name": "agn-ai",
+  "name": "agnaistic",
   "private": true,
   "version": "0.1.0",
   "description": "Agnostic AI Chat",

srv/api/adapter/chai.ts

@@ -2,7 +2,6 @@ import needle from 'needle'
 import { config } from '../../config'
 import { logger } from '../../logger'
 import { trimResponse } from '../chat/common'
-import { createPrompt } from './prompt'
 import { ModelAdapter } from './type'
 
 const base = {
@@ -13,14 +12,7 @@ const base = {
   top_p: 1,
 }
 
-export const handleChai: ModelAdapter = async function* ({
-  chat,
-  char,
-  history,
-  message,
-  sender,
-  members,
-}) {
+export const handleChai: ModelAdapter = async function* ({ char, members, prompt }) {
   if (!config.chai.url) {
     yield { error: 'Chai URL not set' }
     return
@@ -33,7 +25,7 @@ export const handleChai: ModelAdapter = async function* ({
 
   const body = {
     ...base,
-    text: createPrompt({ sender, chat, char, history, message, members }),
+    text: prompt,
   }
 
   const response = await needle('post', `${config.chai.url}/generate/gptj`, body, {

srv/api/adapter/generate.ts

@@ -7,14 +7,15 @@ import { errors, StatusError } from '../wrap'
 import { handleChai } from './chai'
 import { handleKobold } from './kobold'
 import { handleNovel } from './novel'
+import { createPrompt } from './prompt'
 import { ModelAdapter } from './type'
 
 export type GenerateOptions = {
   senderId: string
   chatId: string
-  history: AppSchema.ChatMessage[]
   message: string
   log: AppLog
+  retry?: AppSchema.ChatMessage
 }
 
 const handlers: { [key in ChatAdapter]: ModelAdapter } = {
@@ -41,7 +42,8 @@ export async function generateResponse(
   const adapter =
     (opts.chat.adapter === 'default' ? user.defaultAdapter : opts.chat.adapter) ||
     user.defaultAdapter
-  const adapterOpts = { ...opts, members, user, sender }
+  const prompt = await createPrompt({ ...opts, members, sender })
+  const adapterOpts = { ...opts, members, user, sender, prompt }
 
   const handler = handlers[adapter]
   return handler(adapterOpts)

srv/api/adapter/kobold.ts

@@ -2,7 +2,6 @@ import needle from 'needle'
 import { config } from '../../config'
 import { logger } from '../../logger'
 import { joinParts, trimResponse } from '../chat/common'
-import { createPrompt } from './prompt'
 import { ModelAdapter } from './type'
 
 const MAX_NEW_TOKENS = 196
@@ -31,19 +30,8 @@ const base = {
   sampler_order: [6, 0, 1, 2, 3, 4, 5],
 }
 
-export const handleKobold: ModelAdapter = async function* ({
-  chat,
-  char,
-  history,
-  message,
-  sender,
-  members,
-  user,
-}) {
-  const body = {
-    ...base,
-    prompt: createPrompt({ chat, char, history, message, sender, members }),
-  }
+export const handleKobold: ModelAdapter = async function* ({ char, members, user, prompt }) {
+  const body = { ...base, prompt }
 
   let attempts = 0
   let maxAttempts = body.max_length / MAX_NEW_TOKENS + 4

srv/api/adapter/novel.ts

@@ -2,7 +2,6 @@ import needle from 'needle'
 import { logger } from '../../logger'
 import { sanitise, trimResponse } from '../chat/common'
 import { badWordIds } from './novel-bad-words'
-import { createPrompt } from './prompt'
 import { ModelAdapter } from './type'
 
 const novelUrl = `https://api.novelai.net/ai/generate`
@@ -37,38 +36,20 @@ const base = {
   },
 }
 
-export const handleNovel: ModelAdapter = async function* ({
-  chat,
-  char,
-  history,
-  sender,
-  message,
-  members,
-  user,
-}) {
+export const handleNovel: ModelAdapter = async function* ({ char, members, user, prompt }) {
   if (!user.novelApiKey) {
     yield { error: 'Novel API key not set' }
     return
   }
 
-  const body = {
-    ...base,
-    input: createPrompt({
-      chat,
-      char,
-      history,
-      message,
-      sender,
-      members,
-    }),
-  }
+  const body = { ...base, input: prompt }
 
   const endTokens = ['***', 'Scenario:', '----']
 
   const response = await needle('post', novelUrl, body, {
     json: true,
     timeout: 2000,
-    response_timeout: 8000,
+    response_timeout: 10000,
     headers: { Authorization: `Bearer ${user.novelApiKey}` },
   }).catch((err) => ({ err }))
 

srv/api/adapter/prompt.ts

@@ -1,45 +1,88 @@
 import { AppSchema } from '../../db/schema'
 import gpt from 'gpt-3-encoder'
 import { logger } from '../../logger'
+import { store } from '../../db'
 
 type PromptOpts = {
   sender: AppSchema.Profile
   chat: AppSchema.Chat
   char: AppSchema.Character
-  history: AppSchema.ChatMessage[]
   message: string
   members: AppSchema.Profile[]
+  retry?: AppSchema.ChatMessage
 }
 
+const MAX_TOKENS = 2048
 const BOT_REPLACE = /\{\{char\}\}/g
 const SELF_REPLACE = /\{\{user\}\}/g
 
-export function createPrompt({ sender, chat, char, history, message, members }: PromptOpts) {
+export async function createPrompt({ sender, chat, char, message, members }: PromptOpts) {
   const username = sender.handle || 'You'
 
-  const lines: string[] = [`${char.name}'s Persona: ${formatCharacter(char.name, chat.overrides)}`]
+  const pre: string[] = [`${char.name}'s Persona: ${formatCharacter(char.name, chat.overrides)}`]
 
   if (chat.scenario) {
-    lines.push(`Scenario: ${chat.scenario}`)
+    pre.push(`Scenario: ${chat.scenario}`)
   }
 
-  lines.push(
-    `<START>`,
-    ...chat.sampleChat.split('\n'),
-    ...history.map((chat) => prefix(chat, char.name, members) + chat.msg),
-    `${username}: ${message}`,
-    `${char.name}:`
-  )
+  pre.push(`<START>`, ...chat.sampleChat.split('\n'))
+  const post = [`${username}: ${message}`, `${char.name}:`]
 
-  const prompt = lines
+  const prompt = await appendHistory(chat, members, char, pre, post)
+  return prompt
+}
+
+async function appendHistory(
+  chat: AppSchema.Chat,
+  members: AppSchema.Profile[],
+  char: AppSchema.Character,
+  pre: string[],
+  post: string[],
+  retry?: AppSchema.ChatMessage
+) {
+  const owner = members.find((mem) => mem.userId === chat.userId)
+  if (!owner) {
+    throw new Error(`Cannot produce prompt: Owner profile not found`)
+  }
+
+  // We need to do this early for accurate token counts
+  const preamble = pre
+    .filter(removeEmpty)
+    .join('\n')
+    .replace(BOT_REPLACE, char.name)
+    .replace(SELF_REPLACE, owner.handle)
+  const postamble = post
     .filter(removeEmpty)
     .join('\n')
     .replace(BOT_REPLACE, char.name)
-    .replace(SELF_REPLACE, username)
+    .replace(SELF_REPLACE, owner.handle)
 
-  const tokens = gpt.encode(prompt).length
-  logger.debug({ tokens, prompt }, 'Tokens')
+  let tokens = gpt.encode(preamble + '\n' + postamble).length
+  const lines: string[] = []
+  let before = retry?.updatedAt
+
+  do {
+    const messages = await store.chats.getRecentMessages(chat._id, before)
+    const history = messages.map((chat) => prefix(chat, char.name, members) + chat.msg)
+
+    for (const hist of history) {
+      const nextTokens = gpt.encode(hist).length
+      if (nextTokens + tokens > MAX_TOKENS) break
+      tokens += nextTokens
+      lines.unshift(hist)
+    }
+
+    if (tokens >= MAX_TOKENS || messages.length < 50) break
+    before = messages.slice(-1)[0].createdAt
+  } while (true)
+
+  const middle = lines
+    .join('\n')
+    .replace(BOT_REPLACE, char.name)
+    .replace(SELF_REPLACE, owner.handle)
 
+  const prompt = [preamble, middle, postamble].join('\n')
+  logger.info({ tokens, prompt }, 'Tokens used')
   return prompt
 }
 

srv/api/adapter/type.ts

@@ -5,7 +5,7 @@ export type ModelAdapter = (opts: {
   char: AppSchema.Character
   user: AppSchema.User
   members: AppSchema.Profile[]
-  history: AppSchema.ChatMessage[]
   message: string
   sender: AppSchema.Profile
+  prompt: string
 }) => AsyncGenerator<string | { error: any }>

srv/api/chat/message.ts

@@ -7,7 +7,7 @@ import { obtainLock, releaseLock, verifyLock } from './lock'
 
 export const generateMessage = handle(async ({ userId, params, body, log }, res) => {
   const id = params.id
-  assertValid({ message: 'string', history: 'any', ephemeral: 'boolean?', retry: 'boolean?' }, body)
+  assertValid({ message: 'string', ephemeral: 'boolean?', retry: 'boolean?' }, body)
 
   const lockId = await obtainLock(id)
 
@@ -41,7 +41,6 @@ export const generateMessage = handle(async ({ userId, params, body, log }, res)
     senderId: userId!,
     chatId: id,
     message: body.message,
-    history: body.history,
     log,
   })
 
@@ -73,19 +72,12 @@ export const generateMessage = handle(async ({ userId, params, body, log }, res)
 export const retryMessage = handle(async ({ body, params, userId, log }, res) => {
   const { id, messageId } = params
 
-  assertValid(
-    {
-      history: 'any',
-      message: 'string',
-      ephemeral: 'boolean?',
-    },
-    body
-  )
+  assertValid({ message: 'string', ephemeral: 'boolean?' }, body)
 
   const lockId = await obtainLock(id)
 
   const prev = await store.chats.getMessageAndChat(messageId)
-  if (!prev || !prev.chat) throw errors.NotFound
+  if (!prev || !prev.chat || !prev.msg) throw errors.NotFound
 
   const members = prev.chat.memberIds.concat(prev.chat.userId)
   if (!members.includes(userId!)) throw errors.Forbidden
@@ -99,10 +91,10 @@ export const retryMessage = handle(async ({ body, params, userId, log }, res) =>
 
   const { stream } = await createResponseStream({
     chatId: params.id,
-    history: body.history,
     message: body.message,
     senderId: userId!,
     log,
+    retry: prev.msg,
   })
 
   let generated = ''

srv/api/ws/handle.ts

@@ -91,14 +91,12 @@ export function publishMany<T extends { type: string }>(userIds: string[], data:
   for (const userId of unique) {
     count += publishOne(userId, data)
   }
-
-  logger.debug({ count }, 'Messages sent')
 }
 
 export function publishOne<T extends { type: string }>(userId: string, data: T) {
   let count = 0
   const sockets = userSockets.get(userId)
-  logger.info({ count: sockets?.length, type: data.type }, 'Publishing')
+
   if (!sockets) return count
 
   for (const socket of sockets) {

srv/db/chats.ts

@@ -128,6 +128,11 @@ export async function editMessage(id: string, content: string) {
   return doc
 }
 
+export async function getMessage(messageId: string) {
+  const msg = await msgs.findOne({ _id: messageId, kind: 'chat-message' })
+  return msg
+}
+
 export async function deleteMessages(messageIds: string[]) {
   await chats.deleteMany({ _id: { $in: messageIds } }, { multi: true })
 }
@@ -170,3 +175,18 @@ export async function getAllChats(userId: string) {
 
   return list
 }
+
+/**
+ *
+ * @param chatId
+ * @param before Date ISO string
+ */
+export async function getRecentMessages(chatId: string, before?: string) {
+  const query: any = { kind: 'chat-message', chatId }
+  if (before) {
+    query.createdAt = { $lt: before }
+  }
+
+  const messages = await msgs.find(query).sort({ createdAt: -1 }).limit(50)
+  return messages
+}