Update app.py

app.py

@@ -1,5 +1,5 @@
 from fastapi import FastAPI, Request
-from fastapi.responses import HTMLResponse
+from fastapi.responses import HTMLResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 import chess
@@ -7,6 +7,7 @@ from datasets import load_dataset
 import pickle
 import time
 from pyroaring import BitMap
+import json
 
 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)]
@@ -49,33 +50,56 @@ def read_root(request: Request):
 
 
 
-@app.post("/search")
-async def search(data: dict):
-    start = time.time()
-    board = chess.Board(data['fen'])
-    query_tokens = board_to_tokens(board)
-    matches = query_positions(index, metadata, query_tokens)
+# @app.post("/search")
+# async def search(data: dict):
+#     start = time.time()
+#     board = chess.Board(data['fen'])
+#     query_tokens = board_to_tokens(board)
+#     matches = query_positions(index, metadata, query_tokens)
     
-    seen_puzzles = {}
-    for pos_id, (puzzle_row, move_idx) in matches:
-        if puzzle_row not in seen_puzzles:
-            seen_puzzles[puzzle_row] = (pos_id, move_idx)
+#     seen_puzzles = {}
+#     for pos_id, (puzzle_row, move_idx) in matches:
+#         if puzzle_row not in seen_puzzles:
+#             seen_puzzles[puzzle_row] = (pos_id, move_idx)
     
-    results = []
-    for puzzle_row, (pos_id, move_idx) in seen_puzzles.items():
-        row = dset[puzzle_row]
-        positions = get_puzzle_positions(row['FEN'], row['Moves'])
-        matched_board = positions[move_idx]
+#     results = []
+#     for puzzle_row, (pos_id, move_idx) in seen_puzzles.items():
+#         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'],
+#             "MatchedMove": move_idx
+#         })
+    
+#     elapsed_ms = (time.time() - start) * 1000
+#     return {"count": len(results), "results": results, "time_ms": elapsed_ms}
+
+
+@app.post("/search")
+async def search(data: dict):
+    async def generate():
+        board = chess.Board(data['fen'])
+        query_tokens = board_to_tokens(board)
+        matches = query_positions(index, metadata, query_tokens)
         
-        results.append({
-            "PuzzleId": row['PuzzleId'],
-            "FEN": matched_board.fen(),
-            "Moves": row['Moves'],
-            "Rating": row['Rating'],
-            "Popularity": row['Popularity'],
-            "Themes": row['Themes'],
-            "MatchedMove": move_idx
-        })
+        seen_puzzles = {}
+        for pos_id, (puzzle_row, move_idx) in matches:
+            if puzzle_row not in seen_puzzles:
+                seen_puzzles[puzzle_row] = (pos_id, move_idx)
+        
+        for puzzle_row, (pos_id, move_idx) in seen_puzzles.items():
+            row = dset[puzzle_row]
+            positions = get_puzzle_positions(row['FEN'], row['Moves'])
+            matched_board = positions[move_idx]
+            
+            result = {"PuzzleId": row['PuzzleId'], "FEN": matched_board.fen(), ...}
+            yield json.dumps(result) + "\n"
     
-    elapsed_ms = (time.time() - start) * 1000
-    return {"count": len(results), "results": results, "time_ms": elapsed_ms}
+    return StreamingResponse(generate(), media_type="application/x-ndjson")