Add guidance regions for prompts (#14)

* Add guidance regions for prompts

* update doc

* lint

* Add tools.extractCodeblockFromMarkdown to extract codeblocks from md

* tools: fix prompt space and markdown handling

* fix test

docs/MarkdownProcessing.md

@@ -34,4 +34,52 @@ This would result in:
 ```md
 Hello!
 You are running via API.
+```
+
+### Guidance Region
+
+LLMs don't always give you the output you desire for your prompt. One approach to fixing
+this is by modifying your prompt to be more clear and adding more examples.
+Another standard approach is to provide some initial guidance to the model for what the response 
+should look like, to guarantee that the model will give you output similar to that you desire.
+For example, instead of just asking your model to output JSON or YAML and then
+ending the prompt with a question for the model to answer, you might end it with
+a markdown code block (like <code>```yml</code>), that the LLM would then complete the
+body for.
+
+However, this can lead to messy code for you as you then have to prefix that guidance
+to the response you get from LXL. To make this easier, LXL provides a way to mark
+regions of your prompt as guidance. The guidance will then be automatically prepended
+to the model output you get from LXL. 
+
+As for how we send this data to the LLM, if the LLM is in chat mode, we split the user prompt
+by the base prompt and the guidance region, and mark the guidance region as being under
+the role `model` (Google Gemini API) or `assistant` (OpenAI GPT API) and the former as `system` or `user`.
+If the LLM is in completion mode, we simply append the guidance region to the prompt.
+
+To notate a region like this, you can use the following marker <code>%%%$GUIDANCE_START$%%%</code> in the prompt:
+
+User Prompt (prompt.md):
+<pre>
+Please convert this YAML to JSON:
+```yml
+hello: world
+```
+%%%$GUIDANCE_START$%%%
+```json
+</pre>
+
+This will result in:
+- role: system, message: "Please convert this YAML to JSON:\n```yml\nhello: world\n```\n"
+- role: model, message: "```json\n"
+
+And LXL's output will include the <code>```json</code> and the rest of the output as if they were both part of the model's output (this includes streaming).
+
+The usage in JS would look like:
+```js
+const { ChatSession, importPromptSync } = require('langxlang')
+const session = new ChatSession(service, 'gpt-3.5-turbo', '', {})
+// importPromptSync returns an object with the prompt and the guidance, that can be passed to sendMessage
+const prompt = importPromptSync('prompt.md', {})
+session.sendMessage(prompt).then(console.log)
 ```

src/ChatSession.js

@@ -149,10 +149,29 @@ class ChatSession {
 
   async sendMessage (message, chunkCb) {
     await this.loading
-    const content = message
+    const content = message.basePrompt ?? message.valueOf()
     this.messages.push({ role: 'user', content })
+    let guidance
+    if (message.guidanceText) {
+      guidance = { role: 'assistant', content: message.guidanceText }
+      this.messages.push(guidance)
+      chunkCb?.({ done: false, delta: message.guidanceText, content: message.guidanceText })
+    }
     this._calledFunctionsForRound = []
     const response = await this._submitRequest(chunkCb)
+    if (message.guidanceText) {
+      // update the current model output with the guidance message
+      const last = this.messages[this.messages.length - 1]
+      last.content = message.guidanceText + last.content
+      // remove the guidance message from the response
+      for (let i = 0; i < this.messages.length; i++) {
+        if (this.messages[i] === guidance) {
+          this.messages.splice(i, 1)
+          break
+        }
+      }
+      response.completeMessage = message.guidanceText + response.completeMessage
+    }
     return { text: response.completeMessage, calledFunctions: this._calledFunctionsForRound }
   }
 }

src/CompletionService.js

@@ -35,14 +35,20 @@ class CompletionService {
 
   async _requestCompletionOpenAI (model, system, user) {
     if (!this.openaiApiKey) throw new Error('OpenAI API key not set')
-    const result = await openai.generateCompletion(model, system, user, { apiKey: this.openaiApiKey })
-    return { text: result.message.content }
+    const guidance = system?.guidanceText || user?.guidanceText || ''
+    const result = await openai.generateCompletion(model, system.basePrompt || system, user.basePrompt || user, {
+      apiKey: this.openaiApiKey,
+      guidanceMessage: guidance
+    })
+    return { text: guidance + result.message.content }
   }
 
   async _requestCompletionGemini (model, system, user) {
     if (!this.geminiApiKey) throw new Error('Gemini API key not set')
-    const result = await gemini.generateCompletion(model, this.geminiApiKey, system, user)
-    return { text: result.text() }
+    const guidance = system?.guidanceText || user?.guidanceText || ''
+    const mergedPrompt = [system, user].join('\n')
+    const result = await gemini.generateCompletion(model, this.geminiApiKey, mergedPrompt)
+    return { text: guidance + result.text() }
   }
 
   async requestCompletion (model, system, user) {

src/GoogleAIStudioCompletionService.js

@@ -16,9 +16,11 @@ class GoogleAIStudioCompletionService {
     if (!supportedModels.includes(model)) {
       throw new Error(`Model ${model} is not supported`)
     }
+    const guidance = system?.guidanceText || user?.guidanceText || ''
+    if (guidance) chunkCb?.({ done: false, delta: guidance })
     const mergedPrompt = [system, user].join('\n')
     const result = await studio.generateCompletion(model, mergedPrompt, chunkCb)
-    return { text: result.text }
+    return { text: guidance + result.text }
   }
 
   async requestStreamingChat (model, { messages, maxTokens, functions }, chunkCb) {

src/gemini.js

@@ -1,10 +1,9 @@
 const { GoogleGenerativeAI } = require('@google/generative-ai')
 const debug = require('debug')('lxl')
 
-async function generateCompletion (model, apiKey, system, user) {
+async function generateCompletion (model, apiKey, prompt) {
   const google = new GoogleGenerativeAI(apiKey)
   const generator = google.getGenerativeModel({ model })
-  const prompt = system + '\n' + user
   const result = await generator.generateContent(prompt)
   const response = await result.response
   return response

src/index.d.ts

@@ -85,8 +85,14 @@ declare module 'langxlang' {
     importPrompt(filePath: string, vars: Record<string, string>): Promise<string>
     // Various string manipulation tools to minify/strip down strings
     stripping: {
-      stripMarkdown(input: string, options?: { stripEmailQuotes?: boolean, replacements?: Map<string | RegExp, string> }): string
+      stripMarkdown(input: string, options?: { 
+        stripEmailQuotes?: boolean,
+        replacements?: Map<string | RegExp, string>,
+        allowMalformed?: boolean
+      }): string
     }
+    // Extracts code blocks from markdown
+    extractCodeblockFromMarkdown(markdownInput: string): { raw: string, lang: string, code: string }[]
   }
 
   const tools: Tools

src/openai.js

@@ -4,11 +4,11 @@ const debug = require('debug')('lxl')
 
 async function generateCompletion (model, system, user, options = {}) {
   const openai = new OpenAI(options)
+  const messages = [{ role: 'user', content: user }]
+  if (system) messages.unshift({ role: 'system', content: system })
+  if (options.guidanceMessage) messages.push({ role: 'assistant', content: options.guidanceMessage })
   const completion = await openai.chat.completions.create({
-    messages: [
-      { role: 'system', content: system },
-      { role: 'user', content: user }
-    ],
+    messages,
     model
   })
   const choice = completion.choices[0]

src/tools.js

@@ -2,7 +2,7 @@ const codebase = require('./tools/codebase')
 const viz = require('./tools/viz')
 const yaml = require('./tools/yaml')
 const stripping = require('./tools/stripping')
-const { importPromptSync, importPrompt, loadPrompt, preMarkdown } = require('./tools/mdp')
+const { importPromptRaw, importPromptSync, importPrompt, loadPrompt, preMarkdown } = require('./tools/mdp')
 
 function createTypeWriterEffectStream (to = process.stdout) {
   // Instead of writing everything at once, we want a typewriter effect
@@ -26,15 +26,31 @@ function createTypeWriterEffectStream (to = process.stdout) {
   }
 }
 
+function extractCodeblockFromMarkdown (md) {
+  const tokens = stripping.tokenizeMarkdown(stripping.normalizeLineEndings(md), {})
+  return tokens.reduce((acc, token) => {
+    if (token[1] === 'code') {
+      acc.push({
+        raw: token[0],
+        lang: token[2],
+        code: token[3]
+      })
+    }
+    return acc
+  }, [])
+}
+
 module.exports = {
   makeVizForPrompt: viz.makeVizForPrompt,
   stripping,
   collectFolderFiles: codebase.collectFolderFiles,
   collectGithubRepoFiles: codebase.collectGithubRepoFiles,
   concatFilesToMarkdown: codebase.concatFilesToMarkdown,
   createTypeWriterEffectStream,
+  extractCodeblockFromMarkdown,
   preMarkdown,
   loadPrompt,
+  importPromptRaw,
   importPromptSync,
   importPrompt,
   encodeYAML: yaml.encodeYaml

src/tools/mdp.js

@@ -4,6 +4,14 @@ const getCaller = require('caller')
 
 // See doc/MarkdownPreprocessing.md for more information
 
+class PromptString extends String {
+  constructor (str, basePrompt, guidanceText) {
+    super(str)
+    this.basePrompt = basePrompt
+    this.guidanceText = guidanceText
+  }
+}
+
 function preMarkdown (text, vars = {}) {
   // Notes:
   // %%%()%%% refers to variable insertion
@@ -195,7 +203,16 @@ function preMarkdown (text, vars = {}) {
 }
 
 function loadPrompt (text, vars) {
-  return preMarkdown(text.replaceAll('\r\n', '\n'), vars).trim()
+  const str = preMarkdown(text.replaceAll('\r\n', '\n'), vars)
+  const TOKEN_GUIDANCE_START = '%%%$GUIDANCE_START$%%%'
+  const guidanceText = str.indexOf(TOKEN_GUIDANCE_START)
+  if (guidanceText !== -1) {
+    const [basePrompt, guidanceText] = str.split(TOKEN_GUIDANCE_START)
+    const newStr = str.replace(TOKEN_GUIDANCE_START, '')
+    return new PromptString(newStr, basePrompt, guidanceText)
+  } else {
+    return new PromptString(str)
+  }
 }
 
 function readSync (path, caller) {
@@ -220,7 +237,7 @@ function importPromptRaw (path) {
   const data = path.startsWith('.')
     ? readSync(path, getCaller(1))
     : readSync(path)
-  return data.replaceAll('\r\n', '\n').trim()
+  return data.replaceAll('\r\n', '\n')
 }
 
 function importPromptSync (path, vars) {

src/tools/stripping.js

@@ -411,44 +411,65 @@ function strOnlyContainsCharExcludingWhitespace (str, char) {
   return found
 }
 
-// This mainly erases extraneous new lines outside of code blocks, including ones with empty block quotes
-function stripMarkdown (comment, options = {}) {
-  const isACodeToken = (token) => token.startsWith('```') && token.endsWith('```')
-  if (!comment) return ''
-  comment = normalizeLineEndings(comment)
-  comment = removeSpecialUnicode(comment)
-  // First, split by any codeblocks
+function tokenizeMarkdown (comment, options) {
+  // console.log('Tokenize', comment)
   const tokens = []
+  let tokenSoFar = ''
   let inCodeBlock = false
-  let temp = ''
+  let inCodeLang
   for (let i = 0; i < comment.length; i++) {
     const currentChar = comment[i]
-    const nextChar = comment[i + 1]
-    const nextNextChar = comment[i + 2]
-    if (currentChar === '`' && nextChar === '`' && nextNextChar === '`') {
-      inCodeBlock = !inCodeBlock
-      if (inCodeBlock) {
-        tokens.push(temp)
-        temp = ''
+    const slice = comment.slice(i)
+    if (inCodeBlock) {
+      if (slice.startsWith(inCodeBlock)) {
+        const code = tokenSoFar.slice(inCodeBlock.length + inCodeLang.length + 1) // +1 for the newline
+        tokens.push([tokenSoFar + inCodeBlock, 'code', inCodeLang, code])
+        i += inCodeBlock.length - 1
+        inCodeBlock = false
+        tokenSoFar = ''
       } else {
-        tokens.push('```' + temp + '```')
-        temp = ''
+        tokenSoFar += currentChar
       }
-      i += 2
     } else {
-      temp += currentChar
+      const codeMatch = slice.match(/^([`]{3,})([a-zA-Z]*)\n/)
+      if (codeMatch) {
+        tokens.push([tokenSoFar, 'text'])
+        inCodeBlock = codeMatch[1]
+        inCodeLang = codeMatch[2]
+        tokenSoFar = codeMatch[0]
+        i += tokenSoFar.length - 1
+      } else {
+        tokenSoFar += currentChar
+      }
     }
   }
-  tokens.push(temp)
+  if (inCodeBlock) {
+    if (options.allowMalformed) {
+      tokens.push([tokenSoFar, 'text'])
+    } else {
+      throw new Error('Unmatched code block')
+    }
+  }
+  tokens.push([tokenSoFar, 'text'])
+  return tokens
+}
+
+// This mainly erases extraneous new lines outside of code blocks, including ones with empty block quotes
+function stripMarkdown (comment, options = {}) {
+  if (!comment) return ''
+  comment = normalizeLineEndings(comment)
+  comment = removeSpecialUnicode(comment)
+  // First, split by any codeblocks
+  const tokens = tokenizeMarkdown(comment, options)
   // Now go through the tokens
   const updated = []
   for (const token of tokens) {
-    if (isACodeToken(token)) {
+    if (token[1] === 'code') {
       // Don't update code
-      updated.push(token)
+      updated.push(token[0])
     } else {
       // Replace \n\n or any extra \n's with one \n
-      let update = removeExtraLines(token)
+      let update = removeExtraLines(token[0])
       if (options.replacements) {
         for (const replacement of options.replacements) {
           update = replacement[0] instanceof RegExp
@@ -478,4 +499,4 @@ function stripMarkdown (comment, options = {}) {
   return result.trim()
 }
 
-module.exports = { stripJava, stripPHP, stripGo, stripMarkdown, removeNonAscii, normalizeLineEndings }
+module.exports = { stripJava, stripPHP, stripGo, stripMarkdown, removeNonAscii, normalizeLineEndings, tokenizeMarkdown }