feat: port to new endpoints + move towards adding modal before using maia2

Author: kevinjosethomas

src/api/analysis/analysis.ts

@@ -472,26 +472,14 @@ export const getClientAnalyzedTournamentGame = async (
   } as any as ClientAnalyzedGame
 }
 
-export const getClientAnalyzedLichessGame = async (
-  id: string,
-  pgn: string,
-  maia_model = 'maia_kdd_1500',
-) => {
-  const res = await fetch(
-    buildUrl(
-      'analysis/analyze_user_game?' +
-        new URLSearchParams({
-          maia_model,
-        }),
-    ),
-    {
-      method: 'POST',
-      body: pgn,
-      headers: {
-        'Content-Type': 'text/plain',
-      },
+export const getClientAnalyzedLichessGame = async (id: string, pgn: string) => {
+  const res = await fetch(buildUrl('analysis/analyze_user_game'), {
+    method: 'POST',
+    body: pgn,
+    headers: {
+      'Content-Type': 'text/plain',
     },
-  )
+  })
 
   if (res.status === 401) {
     throw new Error('Unauthorized')
@@ -580,14 +568,12 @@ export const getClientAnalyzedLichessGame = async (
 export const getClientAnalyzedUserGame = async (
   id: string,
   game_type: 'play' | 'hand' | 'brain',
-  maia_model = 'maia_kdd_1500',
 ) => {
   const res = await fetch(
     buildUrl(
       `analysis/user/analyze_user_maia_game/${id}?` +
         new URLSearchParams({
           game_type,
-          maia_model,
         }),
     ),
     {

src/hooks/useAnalysisController/useAnalysisController.ts

@@ -59,7 +59,8 @@ export const useAnalysisController = (
     const board = new Chess(game.moves[controller.currentIndex].board)
 
     ;(async () => {
-      if (!maia.ready || maiaEvaluations[controller.currentIndex]) return
+      if (maia?.status !== 'ready' || maiaEvaluations[controller.currentIndex])
+        return
 
       const { result } = await maia.batchEvaluate(
         Array(9).fill(board.fen()),
@@ -86,7 +87,7 @@ export const useAnalysisController = (
         return newEvaluations
       })
     })()
-  }, [controller.currentIndex, game.type, maia.ready])
+  }, [controller.currentIndex, game.type, maia?.status])
 
   useEffect(() => {
     if (game.type === 'tournament') return
@@ -398,8 +399,8 @@ export const useAnalysisController = (
   }, [controller.currentIndex, currentMove, game.availableMoves])
 
   return {
+    maia,
     move,
-    data: {},
     moves,
     controller,
     currentMaiaModel,

src/hooks/useMaiaEngine/useMaiaEngine.ts

@@ -1,12 +1,18 @@
-import { useMemo } from 'react'
+import { useState, useMemo, useEffect } from 'react'
 
 import Maia from 'src/utils/maia2'
 
 export const useMaiaEngine = () => {
-  const maia = useMemo(() => {
-    const model = new Maia({ model: '/maia2/maia_rapid.onnx' })
-    return model
+  const [maia, setMaia] = useState<Maia>()
+
+  useEffect(() => {
+    setMaia(new Maia({ model: '/maia2/maia_rapid.onnx', type: 'rapid' }))
   }, [])
 
+  // const maia = useMemo(() => {
+  //   const model = new Maia({ model: '/maia2/maia_rapid.onnx', type: 'rapid' })
+  //   return model
+  // }, [])
+
   return maia
 }

src/pages/analysis/[...id].tsx

@@ -237,6 +237,7 @@ const Analysis: React.FC<Props> = ({
     },
   ]
   const {
+    maia,
     move,
     moves,
     controller,
@@ -576,6 +577,15 @@ const Analysis: React.FC<Props> = ({
           content="Collection of chess training and analysis tools centered around Maia."
         />
       </Head>
+      {maia?.status !== 'ready' ? (
+        <>
+          <div className="absolute left-0 top-0 z-50 flex h-screen w-screen flex-col bg-black">
+            <p className="text-white">{maia?.status}</p>
+          </div>
+        </>
+      ) : (
+        <></>
+      )}
       <GameControllerContext.Provider value={{ ...controller }}>
         {analyzedGame && (isMobile ? mobileLayout : desktopLayout)}
       </GameControllerContext.Provider>

src/utils/maia2/model.ts

@@ -4,38 +4,58 @@ import { mirrorMove, preprocess, allPossibleMovesReversed } from './utils'
 
 class Maia {
   public model!: InferenceSession
-  public ready: boolean
+  public type: 'rapid' | 'blitz'
+  public status: 'loading' | 'no-cache' | 'downloading' | 'ready'
 
-  constructor(options: { model: string }) {
-    this.ready = false
+  constructor(options: { model: string; type: 'rapid' | 'blitz' }) {
+    this.status = 'loading'
+    this.type = options.type ?? 'rapid'
     ;(async () => {
       try {
-        const buffer = await this.getCachedModel(options.model)
-        this.model = await InferenceSession.create(buffer)
-
-        this.ready = true
+        console.log('Getting cached')
+        const buffer = await this.getCachedModel(options.model, options.type)
+        await this.initializeModel(buffer)
       } catch (e) {
-        this.ready = false
+        console.log('Missing cache')
+        this.status = 'no-cache'
       }
     })()
   }
 
-  private async getCachedModel(url: string): Promise<ArrayBuffer> {
-    const cache = await caches.open('maia2-model')
+  public getStatus() {
+    return this.status
+  }
+
+  public async getCachedModel(
+    url: string,
+    type: 'rapid' | 'blitz',
+  ): Promise<ArrayBuffer> {
+    const cache = await caches.open(`MAIA2-${type.toUpperCase()}-MODEL`)
     const response = await cache.match(url)
     if (response) {
       return response.arrayBuffer()
     } else {
-      const response = await fetch(url)
-      if (response.ok) {
-        await cache.put(url, response.clone())
-        return response.arrayBuffer()
-      } else {
-        throw new Error('Failed to fetch model')
-      }
+      throw new Error('Model not found in cache')
     }
   }
 
+  public async fetchModel(url: string, type: 'rapid' | 'blitz') {
+    const cache = await caches.open(`MAIA2-${type.toUpperCase()}-MODEL`)
+    const response = await fetch(url)
+    if (response.ok) {
+      await cache.put(url, response.clone())
+      return response.arrayBuffer()
+    } else {
+      throw new Error('Failed to fetch model')
+    }
+  }
+
+  public async initializeModel(buffer: ArrayBuffer) {
+    this.model = await InferenceSession.create(buffer)
+    this.status = 'ready'
+    console.log('initialized')
+  }
+
   /**
    * Evaluates a given chess position using the Maia model.
    *

src/utils/maia2/utils.ts

@@ -288,38 +288,4 @@ function swapColorsInRank(rank: string): string {
   return swappedRank
 }
 
-/**
- * Swap castling rights between white and black.
- *
- * @param castling - The castling rights string.
- * @returns The swapped castling rights string.
- */
-function swapCastlingRights(castling: string): string {
-  const castlingMap: { [key: string]: string } = {
-    K: 'k',
-    Q: 'q',
-    k: 'K',
-    q: 'Q',
-  }
-  return castling
-    .split('')
-    .map((char) => castlingMap[char] || char)
-    .join('')
-}
-
-/**
- * Mirrors the en passant target square by flipping the rank.
- * @param square The en passant target square to be mirrored.
- * @returns The mirrored en passant target square.
- */
-function mirrorEnPassant(square: string): string {
-  const file = square[0]
-  const rank = square[1]
-
-  // Flip the rank: '1' ↔ '8', '2' ↔ '7', etc.
-  const mirroredRank = (9 - parseInt(rank, 10)).toString()
-
-  return `${file}${mirroredRank}`
-}
-
 export { preprocess, mirrorMove, allPossibleMovesReversed }