Skip to content

Commit

Permalink
Merge pull request #134 from OpeOginni/main
Browse files Browse the repository at this point in the history
feat(display tokens used and total cost)
  • Loading branch information
webdevcody authored Oct 31, 2024
2 parents b52a6b4 + f5e181a commit 866acde
Show file tree
Hide file tree
Showing 10 changed files with 383 additions and 81 deletions.
65 changes: 29 additions & 36 deletions app/play/[level]/test-mode.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -13,57 +13,38 @@ import {
} from "@/components/ui/select";
import { api } from "@/convex/_generated/api";
import { AI_MODELS } from "@/convex/constants";
import { errorMessage } from "@/lib/utils";
import { useAITesting } from "@/hooks/useAITesting";

interface TestModeProps {
level: number;
map: string[][];
}

export default function TestMode({ level, map }: TestModeProps) {
const [playerMap, setPlayerMap] = useState<string[][]>([]);
const [isSimulating, setIsSimulating] = useState(false);
const [gameResult, setGameResult] = useState<"WON" | "LOST" | null>(null);
const [selectedModel, setSelectedModel] = useState(AI_MODELS[0].model);
const testAIModel = useAction(api.maps.testAIModel);
const [aiError, setAiError] = useState<string | null>(null);
const [aiReasoning, setAiReasoning] = useState<string | null>(null);
const [showOriginalMap, setShowOriginalMap] = useState(true);
const {
isSimulating,
gameResult,
aiError,
aiReasoning,
aiPromptTokensUsed,
aiOutputTokensUsed,
aiTotalTokensUsed,
aiTotalRunCost,
resultMap,
runTest,
resetAITest,
} = useAITesting({ testingType: "MODEL", level });

const handleAITest = async () => {
setIsSimulating(true);
setGameResult(null);
setAiError(null);
setAiReasoning(null);
setShowOriginalMap(false);

try {
const result = await testAIModel({
level,
modelId: selectedModel,
});

if (!result.map) {
throw new Error("No map found");
}

setPlayerMap(result.map);
setGameResult(result.isWin ? "WON" : "LOST");
setAiReasoning(result.reasoning);
} catch (error) {
console.error("Error testing AI model:", error);
setAiError(errorMessage(error));
} finally {
setIsSimulating(false);
}
await runTest(selectedModel, map);
};

const handleReset = () => {
setShowOriginalMap(true);
setPlayerMap([]);
setGameResult(null);
setAiError(null);
setAiReasoning(null);
resetAITest();
};

return (
Expand Down Expand Up @@ -105,7 +86,7 @@ export default function TestMode({ level, map }: TestModeProps) {
{aiError && <div className="mt-4 text-red-500">{aiError}</div>}
{gameResult && (
<div className="mx-auto mt-4 flex max-w-4xl flex-col items-center gap-8">
<Visualizer map={playerMap} autoStart onReset={handleReset} />
<Visualizer map={resultMap} autoStart onReset={handleReset} />
<div className="flex flex-1 flex-col gap-4">
<div
className={`text-2xl font-bold ${
Expand All @@ -120,6 +101,18 @@ export default function TestMode({ level, map }: TestModeProps) {
<p>{aiReasoning}</p>
</div>
)}
<div className="p-4">
<h3 className="mb-2 font-semibold">Token Usage:</h3>
<ul className="space-y-1 text-sm">
<li>Prompt Tokens: {aiPromptTokensUsed ?? "N/A"}</li>
<li>Output Tokens: {aiOutputTokensUsed ?? "N/A"}</li>
<li>Total Tokens: {aiTotalTokensUsed ?? "N/A"}</li>
<li>
Total Cost:{" "}
{aiTotalRunCost ? `$${aiTotalRunCost.toFixed(6)}` : "N/A"}
</li>
</ul>
</div>
</div>
</div>
)}
Expand Down
79 changes: 46 additions & 33 deletions app/playground/page.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ import { useToast } from "@/components/ui/use-toast";
import { api } from "@/convex/_generated/api";
import { Id } from "@/convex/_generated/dataModel";
import { SIGN_IN_ERROR_MESSAGE } from "@/convex/users";
import { useAITesting } from "@/hooks/useAITesting";
import { errorMessage } from "@/lib/utils";
import { ZombieSurvival } from "@/simulators/zombie-survival";

Expand Down Expand Up @@ -56,15 +57,26 @@ export default function PlaygroundPage() {
]);
const [model, setModel] = useState("");
const [error, setError] = useState<string | null>(null);
const [solution, setSolution] = useState<string[][] | null>(null);
const [reasoning, setReasoning] = useState<string | null>(null);
const [publishing, setPublishing] = useState(false);
const [simulating, setSimulating] = useState(false);
const [userPlaying, setUserPlaying] = useState(false);
const [userSolution, setUserSolution] = useState<string[][]>([]);
const [visualizingUserSolution, setVisualizingUserSolution] = useState(false);
const [openSignInModal, setOpenSignInModal] = useState(false);

const {
isSimulating,
gameResult: solution,
aiError,
aiReasoning,
aiPromptTokensUsed,
aiOutputTokensUsed,
aiTotalTokensUsed,
aiTotalRunCost,
resultMap,
runTest,
resetAITest,
} = useAITesting({ testingType: "MAP" });

async function handlePublish() {
if (!ZombieSurvival.mapHasZombies(map)) {
alert("Add some zombies to the map first");
Expand Down Expand Up @@ -98,30 +110,12 @@ export default function PlaygroundPage() {
}

async function handleSimulate() {
setError(null);
setSolution(null);
setReasoning(null);

if (!ZombieSurvival.mapHasZombies(map)) {
alert("Add some zombies to the map first");
return;
}

setSimulating(true);

const { error, solution, reasoning } = await testMap({
modelId: model,
map: map,
});

if (typeof error !== "undefined") {
setError(error);
} else {
setSolution(solution!);
setReasoning(reasoning!);
}

setSimulating(false);
await runTest(model, map);
}

function handleChangeMap(value: string[][]) {
Expand All @@ -139,8 +133,7 @@ export default function PlaygroundPage() {
}

function handleEdit() {
setSolution(null);
setReasoning(null);
resetAITest();
setUserPlaying(false);
setVisualizingUserSolution(false);

Expand Down Expand Up @@ -223,7 +216,7 @@ export default function PlaygroundPage() {
}
}

// function handleReset() {
// function handleresetAITest() {
// handleChangeMap([]);
// setSolution(null);
// setReasoning(null);
Expand Down Expand Up @@ -304,7 +297,7 @@ export default function PlaygroundPage() {
autoReplay
autoStart
controls={false}
map={visualizingUserSolution ? userSolution : solution!}
map={visualizingUserSolution ? userSolution : resultMap}
/>
)}
{!visualizing && (
Expand Down Expand Up @@ -485,7 +478,7 @@ export default function PlaygroundPage() {

<Button
disabled={
!(solution === null && !simulating && !userPlaying)
!(solution === null && !isSimulating && !userPlaying)
}
onClick={handleUserPlay}
type="button"
Expand All @@ -499,10 +492,30 @@ export default function PlaygroundPage() {
<p className="text-sm text-red-500">{error}</p>
)}
{visualizingUserSolution && <MapStatus map={userSolution} />}
{reasoning !== null && (
<div className="flex flex-col gap-0">
<MapStatus map={solution!} />
<p className="text-sm">{reasoning}</p>
{aiReasoning !== null && (
<div className="flex flex-col gap-2">
<MapStatus map={resultMap} />
<p className="text-sm">{aiReasoning}</p>

<div className="mt-2 space-y-1 text-xs text-gray-500">
<p>
Prompt Tokens:{" "}
{aiPromptTokensUsed?.toLocaleString() ?? "N/A"}
</p>

<p>
Output Tokens:{" "}
{aiOutputTokensUsed?.toLocaleString() ?? "N/A"}
</p>
<p>
Total Tokens:{" "}
{aiTotalTokensUsed?.toLocaleString() ?? "N/A"}
</p>
<p>
Total Cost:{" "}
{aiTotalRunCost ? `$${aiTotalRunCost.toFixed(6)}` : "N/A"}
</p>
</div>
</div>
)}

Expand All @@ -514,11 +527,11 @@ export default function PlaygroundPage() {
<ModelSelector onChange={handleChangeModel} value={model} />
<Button
className="w-full"
disabled={model === "" || simulating}
disabled={model === "" || isSimulating}
onClick={handleSimulate}
type="button"
>
{simulating ? "Simulating..." : "Play With AI"}
{isSimulating ? "Simulating..." : "Play With AI"}
</Button>
</div>
</>
Expand Down
18 changes: 13 additions & 5 deletions convex/maps.ts
Original file line number Diff line number Diff line change
Expand Up @@ -404,16 +404,24 @@ export const testAIModel = action({
throw new Error("Active prompt not found");
}

const { solution, reasoning, error } = await runModel(
args.modelId,
map.grid,
activePrompt.prompt,
);
const {
solution,
reasoning,
error,
promptTokens,
outputTokens,
totalTokensUsed,
totalRunCost,
} = await runModel(args.modelId, map.grid, activePrompt.prompt);

return {
map: solution,
isWin: error ? false : ZombieSurvival.isWin(solution!),
reasoning,
promptTokens,
outputTokens,
totalTokensUsed,
totalRunCost,
error,
};
},
Expand Down
Loading

0 comments on commit 866acde

Please sign in to comment.