refactoring for OpenAI and Groq

src/app/queries/getSystemPrompt.ts

@@ -0,0 +1,27 @@
+import { Preset } from "../engine/presets"
+
+export function getSystemPrompt({
+  preset,
+  // prompt,
+  // existingPanelsTemplate,
+  firstNextOrLast,
+  maxNbPanels,
+  nbPanelsToGenerate,
+  // nbMaxNewTokens,
+}: {
+  preset: Preset
+  // prompt: string
+  // existingPanelsTemplate: string
+  firstNextOrLast: string
+  maxNbPanels: number
+  nbPanelsToGenerate: number
+  // nbMaxNewTokens: number
+}) {
+  return [
+    `You are a writer specialized in ${preset.llmPrompt}`,
+    `Please write detailed drawing instructions and short (2-3 sentences long) speech captions for the ${firstNextOrLast} ${nbPanelsToGenerate} panels (out of ${maxNbPanels} in total) of a new story, but keep it open-ended (it will be continued and expanded later). Please make sure each of those ${nbPanelsToGenerate} panels include info about character gender, age, origin, clothes, colors, location, lights, etc. Only generate those ${nbPanelsToGenerate} panels, but take into account the fact the panels are part of a longer story (${maxNbPanels} panels long).`,
+    `Give your response as a VALID JSON array like this: \`Array<{ panel: number; instructions: string; caption: string; }>\`.`,
+    // `Give your response as Markdown bullet points.`,
+    `Be brief in the instructions and narrative captions of those ${nbPanelsToGenerate} panels, don't add your own comments. The captions must be captivating, smart, entertaining. Be straight to the point, and never reply things like "Sure, I can.." etc. Reply using valid JSON!! Important: Write valid JSON!`
+  ].filter(item => item).join("\n")
+}

src/app/queries/getUserPrompt.ts

@@ -0,0 +1,9 @@
+export function getUserPrompt({
+  prompt,
+  existingPanelsTemplate,
+}: {
+  prompt: string
+  existingPanelsTemplate: string
+}) {
+  return `The story is about: ${prompt}.${existingPanelsTemplate}`
+}

src/app/queries/predictNextPanels.ts

@@ -7,6 +7,8 @@ import { createZephyrPrompt } from "@/lib/createZephyrPrompt"
 import { dirtyGeneratedPanelCleaner } from "@/lib/dirtyGeneratedPanelCleaner"
 import { dirtyGeneratedPanelsParser } from "@/lib/dirtyGeneratedPanelsParser"
 import { sleep } from "@/lib/sleep"
+import { getSystemPrompt } from "./getSystemPrompt"
+import { getUserPrompt } from "./getUserPrompt"
 
 export const predictNextPanels = async ({
   preset,
@@ -31,7 +33,6 @@ export const predictNextPanels = async ({
     ? ` To help you, here are the previous panels and their captions (note: if you see an anomaly here eg. no caption or the same description repeated multiple times, do not hesitate to fix the story): ${JSON.stringify(existingPanels, null, 2)}`
     : ''
 
-
   const firstNextOrLast =
     existingPanels.length === 0
       ? "first"
@@ -39,24 +40,23 @@ export const predictNextPanels = async ({
       ? "last"
       : "next"
 
-  const query = createZephyrPrompt([
-    {
-      role: "system",
-      content: [
-        `You are a writer specialized in ${preset.llmPrompt}`,
-        `Please write detailed drawing instructions and short (2-3 sentences long) speech captions for the ${firstNextOrLast} ${nbPanelsToGenerate} panels (out of ${maxNbPanels} in total) of a new story, but keep it open-ended (it will be continued and expanded later). Please make sure each of those ${nbPanelsToGenerate} panels include info about character gender, age, origin, clothes, colors, location, lights, etc. Only generate those ${nbPanelsToGenerate} panels, but take into account the fact the panels are part of a longer story (${maxNbPanels} panels long).`,
-        `Give your response as a VALID JSON array like this: \`Array<{ panel: number; instructions: string; caption: string; }>\`.`,
-        // `Give your response as Markdown bullet points.`,
-        `Be brief in the instructions and narrative captions of those ${nbPanelsToGenerate} panels, don't add your own comments. The captions must be captivating, smart, entertaining. Be straight to the point, and never reply things like "Sure, I can.." etc. Reply using valid JSON!! Important: Write valid JSON!`
-      ].filter(item => item).join("\n")
-    },
-    {
-      role: "user",
-      content: `The story is about: ${prompt}.${existingPanelsTemplate}`,
-    }
+  const systemPrompt = getSystemPrompt({
+    preset,
+    firstNextOrLast,
+    maxNbPanels,
+    nbPanelsToGenerate,
+  })
+
+  const userPrompt = getUserPrompt({
+    prompt,
+    existingPanelsTemplate,
+  })
+  
+  const zephyPrompt = createZephyrPrompt([
+    { role: "system", content: systemPrompt },
+    { role: "user", content: userPrompt }
   ]) + "\n[{"
 
-
   let result = ""
 
   // we don't require a lot of token for our task
@@ -66,8 +66,8 @@ export const predictNextPanels = async ({
   const nbMaxNewTokens = nbPanelsToGenerate * nbTokensPerPanel
 
   try {
-    // console.log(`calling predict(${query}, ${nbTotalPanels})`)
-    result = `${await predict(query, nbMaxNewTokens)}`.trim()
+    // console.log(`calling predict:`, { systemPrompt, userPrompt, nbMaxNewTokens })
+    result = `${await predict({ systemPrompt, userPrompt, nbMaxNewTokens })}`.trim()
     console.log("LLM result (1st trial):", result)
     if (!result.length) {
       throw new Error("empty result on 1st trial!")
@@ -78,7 +78,7 @@ export const predictNextPanels = async ({
     await sleep(2000)
 
     try {
-      result = `${await predict(query + " \n ", nbMaxNewTokens)}`.trim()
+      result = `${await predict({ systemPrompt: systemPrompt + " \n ", userPrompt, nbMaxNewTokens })}`.trim()
       console.log("LLM result (2nd trial):", result)
       if (!result.length) {
         throw new Error("empty result on 2nd trial!")

src/app/queries/predictWithGroq.ts

@@ -2,7 +2,15 @@
 
 import Groq from "groq-sdk"
 
-export async function predict(inputs: string, nbMaxNewTokens: number): Promise<string> {
+export async function predict({
+  systemPrompt,
+  userPrompt,
+  nbMaxNewTokens,
+}: {
+  systemPrompt: string
+  userPrompt: string
+  nbMaxNewTokens: number
+}): Promise<string> {
   const groqApiKey = `${process.env.AUTH_GROQ_API_KEY || ""}`
   const groqApiModel = `${process.env.LLM_GROQ_API_MODEL || "mixtral-8x7b-32768"}`
   
@@ -11,7 +19,8 @@ export async function predict(inputs: string, nbMaxNewTokens: number): Promise<s
   })
 
   const messages: Groq.Chat.Completions.CompletionCreateParams.Message[] = [
-    { role: "assistant", content: inputs },
+    { role: "system", content: systemPrompt },
+    { role: "user", content: userPrompt },
   ]
 
   try {

src/app/queries/predictWithHuggingFace.ts

@@ -2,8 +2,17 @@
 
 import { HfInference, HfInferenceEndpoint } from "@huggingface/inference"
 import { LLMEngine } from "@/types"
+import { createZephyrPrompt } from "@/lib/createZephyrPrompt"
 
-export async function predict(inputs: string, nbMaxNewTokens: number): Promise<string> {
+export async function predict({
+  systemPrompt,
+  userPrompt,
+  nbMaxNewTokens,
+}: {
+  systemPrompt: string
+  userPrompt: string
+  nbMaxNewTokens: number
+}): Promise<string> {
   const hf = new HfInference(process.env.AUTH_HF_API_TOKEN)
 
   const llmEngine = `${process.env.LLM_ENGINE || ""}` as LLMEngine
@@ -46,7 +55,12 @@ export async function predict(inputs: string, nbMaxNewTokens: number): Promise<s
   try {
     for await (const output of api.textGenerationStream({
       model: llmEngine === "INFERENCE_ENDPOINT" ? undefined : (inferenceModel || undefined),
-      inputs,
+      
+      inputs: createZephyrPrompt([
+        { role: "system", content: systemPrompt },
+        { role: "user", content: userPrompt }
+      ]) + "\n[{", // <-- important: we force its hand
+
       parameters: {
         do_sample: true,
         max_new_tokens: nbMaxNewTokens,

src/app/queries/predictWithOpenAI.ts

@@ -1,9 +1,17 @@
 "use server"
 
-import type { ChatCompletionMessage } from "openai/resources/chat"
+import type { ChatCompletionMessageParam } from "openai/resources/chat"
 import OpenAI from "openai"
 
-export async function predict(inputs: string, nbMaxNewTokens: number): Promise<string> {
+export async function predict({
+  systemPrompt,
+  userPrompt,
+  nbMaxNewTokens,
+}: {
+  systemPrompt: string
+  userPrompt: string
+  nbMaxNewTokens: number
+}): Promise<string> {
   const openaiApiKey = `${process.env.AUTH_OPENAI_API_KEY || ""}`
   const openaiApiBaseUrl = `${process.env.LLM_OPENAI_API_BASE_URL || "https://api.openai.com/v1"}`
   const openaiApiModel = `${process.env.LLM_OPENAI_API_MODEL || "gpt-3.5-turbo"}`
@@ -13,8 +21,9 @@ export async function predict(inputs: string, nbMaxNewTokens: number): Promise<s
     baseURL: openaiApiBaseUrl,
   })
 
-  const messages: ChatCompletionMessage[] = [
-    { role: "assistant", content: inputs },
+  const messages: ChatCompletionMessageParam[] = [
+    { role: "system", content: systemPrompt },
+    { role: "user", content: userPrompt },
   ]
 
   try {