Move prompt/checks to models/index.ts

models/gemini-1.5-pro.ts

@@ -1,5 +1,5 @@
 import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai";
-import { type ModelResult } from ".";
+import { type ModelHandler } from ".";
 
 const schema = {
   description: "Game Round Results",
@@ -41,13 +41,13 @@ const schema = {
 };
 
 interface GeminiResponse {
+  boxCoordinates: number[][];
   map: string[][];
-  reasoning: string;
   playerCoordinates: number[];
-  boxCoordinates: number[][];
+  reasoning: string;
 }
 
-export async function gemini15pro(map: string[][]): Promise<ModelResult> {
+export const gemini15pro: ModelHandler = async (prompt, map) => {
   const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY!);
 
   const model = genAI.getGenerativeModel({
@@ -59,20 +59,7 @@ export async function gemini15pro(map: string[][]): Promise<ModelResult> {
   });
 
   const result = await model.generateContent(
-    `You're given a 2d grid of nums such that.
-    " " represents an empty space.
-    "Z" represents a zombie. Zombies move one Manhattan step every turn and aim to reach the player.
-    "R" represents rocks, which players can shoot over but zombies cannot pass through or break.
-    "P" represents the player, who cannot move. The player's goal is to shoot and kill zombies before they reach them.
-    "B" represents blocks that can be placed before the round begins to hinder the zombies. You can place up to two blocks on the map.
-
-    Your goal is to place the player ("P") and two blocks ("B") in locations that maximize the player's survival by delaying the zombies' approach.
-    You can shoot any zombie regardless of where it is on the grid.
-    Returning a 2d grid with the player and blocks placed in the optimal locations, with the coordinates player ("P") and the blocks ("B"), also provide reasoning for the choices.
-
-    You can't replace rocks R or zombies Z with blocks.  If there is no room to place a block, do not place any.
-
-    Grid: ${JSON.stringify(map)}`,
+    `${prompt}\n\nGrid: ${JSON.stringify(map)}`,
   );
 
   // todo: check if the response is valid acc to types and the player and box coordinates are valid,
@@ -81,7 +68,8 @@ export async function gemini15pro(map: string[][]): Promise<ModelResult> {
   const parsedResponse = JSON.parse(result.response.text()) as GeminiResponse;
 
   return {
-    solution: parsedResponse.map,
+    boxCoordinates: parsedResponse.boxCoordinates,
+    playerCoordinates: parsedResponse.playerCoordinates,
     reasoning: parsedResponse.reasoning,
   };
-}
+};

models/gpt-4o.ts

@@ -1,33 +1,17 @@
 import OpenAI from "openai";
 import { z } from "zod";
 import { zodResponseFormat } from "openai/helpers/zod";
-import { type ModelResult } from ".";
+import { type ModelHandler } from ".";
 
 const ResponseSchema = z.object({
   reasoning: z.string(),
   playerCoordinates: z.array(z.number()),
   boxCoordinates: z.array(z.array(z.number())),
 });
 
-export async function gpt4o(map: string[][]): Promise<ModelResult> {
+export const gpt4o: ModelHandler = async (prompt, map) => {
   const openai = new OpenAI();
 
-  const prompt = `You're given a 2d grid of nums such that.
-    " " represents an empty space.
-    "Z" represents a zombie. Zombies move one Manhattan step every turn and aim to reach the player.
-    "R" represents rocks, which players can shoot over but zombies cannot pass through or break.
-    "P" represents the player, who cannot move. The player's goal is to shoot and kill zombies before they reach them.
-    "B" represents blocks that can be placed before the round begins to hinder the zombies.
-
-    Your goal is to place the player ("P") in a location which maximize the player's survival.
-    You must place two blocks ("B") in locations which maximize the player's survival. 
-    You can shoot any zombie regardless of where it is on the grid.
-    Returning a 2d grid with the player and blocks placed in the optimal locations, with the coordinates player ("P") and the blocks ("B"), also provide reasoning for the choices.
-    Zombies can only move horizontally or vertically, not diagonally.
-    You can't replace rocks R or zombies Z with blocks.  
-    Players will always shoot at the closest zombie each turn.
-    If there is no room to place a block, do not place any.`;
-
   const completion = await openai.beta.chat.completions.parse({
     model: "gpt-4o-2024-08-06",
     messages: [
@@ -52,25 +36,9 @@ export async function gpt4o(map: string[][]): Promise<ModelResult> {
     throw new Error("Failed to run model GPT-4o");
   }
 
-  const originalMap = JSON.parse(JSON.stringify(map));
-
-  const [playerRow, playerCol] = response.parsed.playerCoordinates;
-  if (originalMap[playerRow][playerCol] !== " ") {
-    throw new Error("Cannot place player in a non-empty space");
-  }
-  originalMap[playerRow][playerCol] = "P";
-
-  for (const block of response.parsed.boxCoordinates) {
-    const [blockRow, blockCol] = block;
-    if (originalMap[blockRow][blockCol] !== " ") {
-      throw new Error("Cannot place block in a non-empty space");
-    }
-
-    originalMap[blockRow][blockCol] = "B";
-  }
-
   return {
-    solution: originalMap,
+    boxCoordinates: response.parsed.boxCoordinates,
+    playerCoordinates: response.parsed.playerCoordinates,
     reasoning: response.parsed.reasoning,
   };
-}
+};

models/index.ts

@@ -1,24 +1,75 @@
 import { gemini15pro } from "./gemini-1.5-pro";
 import { gpt4o } from "./gpt-4o";
 
-export interface ModelResult {
-  solution: string[][];
+export type ModelHandler = (
+  prompt: string,
+  map: string[][],
+) => Promise<{
+  boxCoordinates: number[][];
+  playerCoordinates: number[];
   reasoning: string;
-}
+}>;
+
+const prompt = `You're given a 2d grid of nums such that.
+" " represents an empty space.
+"Z" represents a zombie. Zombies move one Manhattan step every turn and aim to reach the player.
+"R" represents rocks, which players can shoot over but zombies cannot pass through or break.
+"P" represents the player, who cannot move. The player's goal is to shoot and kill zombies before they reach them.
+"B" represents blocks that can be placed before the round begins to hinder the zombies.
+
+Your goal is to place the player ("P") in a location which maximize the player's survival.
+You must place two blocks ("B") in locations which maximize the player's survival.
+You can shoot any zombie regardless of where it is on the grid.
+Returning a 2d grid with the player and blocks placed in the optimal locations, with the coordinates player ("P") and the blocks ("B"), also provide reasoning for the choices.
+Zombies can only move horizontally or vertically, not diagonally.
+You can't replace rocks R or zombies Z with blocks.
+Players will always shoot at the closest zombie each turn.
+If there is no room to place a block, do not place any.`;
 
 export async function runModel(
   modelId: string,
   map: string[][],
-): Promise<ModelResult> {
+): Promise<{
+  solution: string[][];
+  reasoning: string;
+}> {
+  let result;
+
   switch (modelId) {
     case "gemini-1.5-pro": {
-      return gemini15pro(map);
+      result = await gemini15pro(prompt, map);
+      break;
     }
     case "gpt-4o": {
-      return gpt4o(map);
+      result = await gpt4o(prompt, map);
+      break;
     }
     default: {
       throw new Error(`Tried running unknown model '${modelId}'`);
     }
   }
+
+  const originalMap = JSON.parse(JSON.stringify(map));
+  const [playerRow, playerCol] = result.playerCoordinates;
+
+  if (originalMap[playerRow][playerCol] !== " ") {
+    throw new Error("Cannot place player in a non-empty space");
+  }
+
+  originalMap[playerRow][playerCol] = "P";
+
+  for (const block of result.boxCoordinates) {
+    const [blockRow, blockCol] = block;
+
+    if (originalMap[blockRow][blockCol] !== " ") {
+      throw new Error("Cannot place block in a non-empty space");
+    }
+
+    originalMap[blockRow][blockCol] = "B";
+  }
+
+  return {
+    solution: originalMap,
+    reasoning: result.reasoning,
+  };
 }