Update app.py

app.py

@@ -48,58 +48,25 @@ templates = Jinja2Templates(directory="templates")
 def read_root(request: Request):
     return templates.TemplateResponse("index.html", {"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)
-    
-#     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.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)
-        
-        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)
+    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)
+    
+    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]
         
-        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 = {
+        results.append({
             "PuzzleId": row['PuzzleId'],
             "FEN": matched_board.fen(),
             "Moves": row['Moves'],
@@ -107,7 +74,7 @@ async def search(data: dict):
             "Popularity": row['Popularity'],
             "Themes": row['Themes'],
             "MatchedMove": move_idx
-        }
-            yield json.dumps(result) + "\n"
+        })
     
-    return StreamingResponse(generate(), media_type="application/x-ndjson")
+    elapsed_ms = (time.time() - start) * 1000
+    return {"count": len(results), "results": results, "time_ms": elapsed_ms}