feat: chat session response prefix

Author: giladgd

src/evaluator/LlamaChatSession/LlamaChatSession.ts

@@ -148,6 +148,14 @@ export type LLamaChatPromptOptions<Functions extends ChatSessionModelFunctions |
      */
     trimWhitespaceSuffix?: boolean,
 
+    /**
+     * Force a given text prefix to be the start of the model response, to make the model follow a specific path.
+     *
+     * May cause some models to not use the given functions in some scenarios where they would have been used otherwise,
+     * so avoid using it together with function calling if you notice unexpected behavior.
+     */
+    responsePrefix?: string,
+
     /**
      * See the parameter `evaluationPriority` on the `LlamaContextSequence.evaluate()` function for more information.
      */
@@ -403,6 +411,7 @@ export class LlamaChatSession {
             seed,
             grammar,
             trimWhitespaceSuffix = false,
+            responsePrefix,
             repeatPenalty,
             tokenBias,
             customStopTriggers
@@ -415,7 +424,7 @@ export class LlamaChatSession {
             maxParallelFunctionCalls: maxParallelFunctionCalls as undefined,
 
             onTextChunk, onToken, signal, stopOnAbortSignal, maxTokens, temperature, minP, topK, topP, seed, grammar, trimWhitespaceSuffix,
-            repeatPenalty, tokenBias, customStopTriggers
+            responsePrefix, repeatPenalty, tokenBias, customStopTriggers
         });
 
         return responseText;
@@ -441,6 +450,7 @@ export class LlamaChatSession {
         seed,
         grammar,
         trimWhitespaceSuffix = false,
+        responsePrefix,
         repeatPenalty,
         tokenBias,
         customStopTriggers,
@@ -467,17 +477,30 @@ export class LlamaChatSession {
                 ? undefined
                 : appendUserMessageToChatHistory(lastEvaluation?.contextWindow, prompt);
 
+            const resolvedResponsePrefix = (responsePrefix != null && responsePrefix !== "")
+                ? responsePrefix
+                : undefined;
+
             newChatHistory.push({
                 type: "model",
-                response: []
+                response: resolvedResponsePrefix != null
+                    ? [resolvedResponsePrefix]
+                    : []
             });
 
             if (newContextWindowChatHistory != null)
                 newContextWindowChatHistory.push({
                     type: "model",
-                    response: []
+                    response: resolvedResponsePrefix != null
+                        ? [resolvedResponsePrefix]
+                        : []
                 });
 
+            if (resolvedResponsePrefix != null) {
+                safeEventCallback(onToken)?.(this.model.tokenize(resolvedResponsePrefix));
+                safeEventCallback(onTextChunk)?.(resolvedResponsePrefix);
+            }
+
             // eslint-disable-next-line no-constant-condition
             while (true) {
                 const functionCallsAndResults: Array<Promise<null | {

test/modelDependent/llama3/chatSession.test.ts

@@ -107,6 +107,32 @@ describe("llama 3", () => {
             });
             expect(completion).to.eql(" it is.");
         });
+        test("using response prefix", {timeout: 1000 * 60 * 60 * 2}, async () => {
+            const modelPath = await getModelFile("Meta-Llama-3-8B-Instruct-Q4_K_M.gguf");
+            const llama = await getTestLlama();
+
+            const model = await llama.loadModel({
+                modelPath
+            });
+            const context = await model.createContext({
+                contextSize: 2048
+            });
+            const chatSession = new LlamaChatSession({
+                contextSequence: context.getSequence()
+            });
+
+            expect(chatSession.chatWrapper).to.be.an.instanceof(Llama3ChatWrapper);
+
+            const prompt = "Describe the appearance of a llama";
+            const responsePrefix = "Of course! A llama is";
+            const res = await chatSession.prompt(prompt, {
+                responsePrefix,
+                maxTokens: 10
+            });
+
+            expect(res.startsWith(responsePrefix)).to.eql(true);
+            expect(res).toMatchInlineSnapshot('"Of course! A llama is a domesticated mammal that belongs to the camel"');
+        });
 
         // disabled due to getting timeout in the CI due to taking too long
         test.skip("context shift works correctly", {timeout: 1000 * 60 * 60 * 2}, async () => {