Update app.py

app.py

@@ -48,33 +48,60 @@ 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)
+# @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
-        })
+#         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):
+    def generate():
+        board = chess.Board(data['fen'])
+        query_tokens = board_to_tokens(board)
+        matches = query_positions(index, metadata, query_tokens)
+        
+        seen_puzzles = set()
+        for pos_id, (puzzle_row, move_idx) in matches:
+            if puzzle_row in seen_puzzles: continue
+            seen_puzzles.add(puzzle_row)
+            
+            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(),
+                      "Moves": row['Moves'],
+                      "Rating": row['Rating'],
+                      "Popularity": row['Popularity'],
+                      "Themes": row['Themes'],
+                      "MatchedMove": move_idx}
+            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")