diff --git a/models/gemini-1.5-pro.ts b/models/gemini-1.5-pro.ts index 7cc9fcc..0d6bba1 100644 --- a/models/gemini-1.5-pro.ts +++ b/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 { +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 { }); 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 { const parsedResponse = JSON.parse(result.response.text()) as GeminiResponse; return { - solution: parsedResponse.map, + boxCoordinates: parsedResponse.boxCoordinates, + playerCoordinates: parsedResponse.playerCoordinates, reasoning: parsedResponse.reasoning, }; -} +}; diff --git a/models/gpt-4o.ts b/models/gpt-4o.ts index 498c8eb..632ae5b 100644 --- a/models/gpt-4o.ts +++ b/models/gpt-4o.ts @@ -1,7 +1,7 @@ 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(), @@ -9,25 +9,9 @@ const ResponseSchema = z.object({ boxCoordinates: z.array(z.array(z.number())), }); -export async function gpt4o(map: string[][]): Promise { +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 { 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, }; -} +}; diff --git a/models/index.ts b/models/index.ts index 0b994f7..b8bd75d 100644 --- a/models/index.ts +++ b/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 { +): 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, + }; }