Update app.py

app.py

@@ -3,6 +3,27 @@ from fastapi.responses import HTMLResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 import chess
+from datasets import load_dataset
+
+dset = load_dataset("Lichess/chess-puzzles", split="train")
+index, metadata = load_index('chess_index.pkl')
+
+def board_to_tokens(board):
+    return [(board.piece_at(sq).symbol(), chess.square_name(sq)) for sq in chess.SQUARES if board.piece_at(sq)]
+    
+def get_puzzle_positions(fen, moves_uci):
+    positions = []
+    
+    board = chess.Board(fen)
+    
+    board.push_uci(moves_uci.split()[0])
+    positions.append(board.copy())
+    
+    for move_uci in moves_uci.split()[1:]:
+        board.push_uci(move_uci)
+        positions.append(board.copy())
+    
+    return positions
 
 app = FastAPI()
 app.mount("/static", StaticFiles(directory="static"), name="static")
@@ -12,19 +33,29 @@ templates = Jinja2Templates(directory="templates")
 def read_root(request: Request):
     return templates.TemplateResponse("index.html", {"request": request})
 
+
 @app.post("/search")
 async def search(data: dict):
     board = chess.Board(data['fen'])
-    # TODO: Replace with your actual search logic
-    mock_puzzle = {
-        "PuzzleId": "000rZ",
-        "FEN": "2kr1b1r/p1p2pp1/2pqb3/7p/3N2n1/2NPB3/PPP2PPP/R2Q1RK1 w - - 2 13",
-        "Moves": "d4e6 d6h2",
-        "Rating": 923,
-        "Popularity": 100,
-        "Themes": ["kingsideAttack", "mate", "mateIn1"],
-        "PuzzleUrl": "https://lichess.org/training/000rZ",
-        "GameUrl": "https://lichess.org/seIMDWkD#25",
-        "MatchedMove": 0
-    }
-    return {"count": 100, "results": [mock_puzzle] * 100}
+    query_tokens = board_to_tokens(board)
+    matches = query_positions(index, metadata, query_tokens)
+    
+    results = []
+    for pos_id, (puzzle_row, move_idx) in matches[:100]:
+        row = dset[puzzle_row]
+        positions = get_puzzle_positions(row['FEN'], row['Moves'])
+        matched_board = positions[move_idx]
+        
+        results.append({
+            "PuzzleId": row['PuzzleId'],
+            "FEN": matched_board.fen(),
+            "Moves": row['Moves'],
+            "Rating": row['Rating'],
+            "Popularity": row['Popularity'],
+            "Themes": row['Themes'].split(),
+            "PuzzleUrl": f"https://lichess.org/training/{row['PuzzleId']}",
+            "GameUrl": row['GameUrl'],
+            "MatchedMove": move_idx
+        })
+    
+    return {"count": len(matches), "results": results}