feature(workflows): added a agentic workflow for sudoku

tests/common/sudoku_test.py

@@ -0,0 +1,116 @@
+from trinity.common.workflows.envs.sudoku.sudoku_generator import SudokuGenerator
+from trinity.common.workflows.envs.sudoku.sudoku_judge import SudokuJudge
+
+# ---------- Generator Tests (9x9) ----------
+
+
+def test_9x9_generator_produces_valid_solution():
+    gen = SudokuGenerator()
+    puzzle, solution = gen.generate()
+
+    assert len(puzzle) == 9
+    assert len(solution) == 9
+    assert SudokuJudge.is_valid(solution)
+
+
+def test_9x9_generator_creates_holes():
+    gen = SudokuGenerator()
+    puzzle, _ = gen.generate()
+
+    zero_count = sum(row.count(0) for row in puzzle)
+    assert zero_count > 0
+
+
+def test_9x9_solution_is_fully_filled():
+    gen = SudokuGenerator()
+    _, solution = gen.generate()
+
+    for row in solution:
+        assert 0 not in row
+
+
+# ---------- Judge Tests (9x9) ----------
+
+
+def test_judge_allows_incomplete_board():
+    board = [
+        [5, 3, 0, 0, 7, 0, 0, 0, 0],
+        [6, 0, 0, 1, 9, 5, 0, 0, 0],
+        [0, 9, 8, 0, 0, 0, 0, 6, 0],
+        [8, 0, 0, 0, 6, 0, 0, 0, 3],
+        [4, 0, 0, 8, 0, 3, 0, 0, 1],
+        [7, 0, 0, 0, 2, 0, 0, 0, 6],
+        [0, 6, 0, 0, 0, 0, 2, 8, 0],
+        [0, 0, 0, 4, 1, 9, 0, 0, 5],
+        [0, 0, 0, 0, 8, 0, 0, 7, 9],
+    ]
+
+    assert SudokuJudge.is_valid(board)
+
+
+def test_judge_detects_row_violation():
+    board = [
+        [1, 1, 0, 0, 0, 0, 0, 0, 0],
+    ] + [[0] * 9 for _ in range(8)]
+
+    assert not SudokuJudge.is_valid(board)
+
+
+def test_judge_detects_column_violation():
+    board = [
+        [5, 0, 0, 0, 0, 0, 0, 0, 0],
+        [5, 0, 0, 0, 0, 0, 0, 0, 0],
+    ] + [[0] * 9 for _ in range(7)]
+
+    assert not SudokuJudge.is_valid(board)
+
+
+def test_judge_detects_block_violation():
+    board = [
+        [1, 2, 3, 0, 0, 0, 0, 0, 0],
+        [4, 1, 0, 0, 0, 0, 0, 0, 0],
+    ] + [[0] * 9 for _ in range(7)]
+
+    assert not SudokuJudge.is_valid(board)
+
+
+# ---------- Generator & Judge Tests (4x4) ----------
+
+
+def test_4x4_generator_produces_valid_solution():
+    gen = SudokuGenerator(size=4)
+    puzzle, solution = gen.generate()
+
+    assert len(puzzle) == 4
+    assert len(solution) == 4
+    assert SudokuJudge.is_valid(solution)
+
+
+def test_4x4_solution_is_fully_filled():
+    gen = SudokuGenerator(size=4)
+    _, solution = gen.generate()
+
+    for row in solution:
+        assert 0 not in row
+
+
+def test_4x4_judge_detects_row_violation():
+    board = [
+        [1, 1, 0, 0],
+        [0, 0, 0, 0],
+        [0, 0, 0, 0],
+        [0, 0, 0, 0],
+    ]
+
+    assert not SudokuJudge.is_valid(board)
+
+
+def test_4x4_judge_detects_block_violation():
+    board = [
+        [1, 2, 0, 0],
+        [3, 1, 0, 0],  # duplicate "1" in top-left 2x2 block
+        [0, 0, 0, 0],
+        [0, 0, 0, 0],
+    ]
+
+    assert not SudokuJudge.is_valid(board)

trinity/common/workflows/__init__.py

@@ -48,6 +48,8 @@
         # on-policy distillation workflows
         "on_policy_distill_workflow": "trinity.common.workflows.on_policy_distill_workflow.OnPolicyDistillWorkflow",
         "on_policy_distill_math_workflow": "trinity.common.workflows.on_policy_distill_workflow.OnPolicyDistillMathWorkflow",
+        # custom workflows
+        "sudoku_workflow": "trinity.common.workflows.envs.sudoku.sudoku_workflow.SudokuWorkflow",
     },
 )
 

trinity/common/workflows/envs/sudoku/sudoku_generator.py

@@ -0,0 +1,138 @@
+import math
+import random
+
+
+class SudokuGenerator:
+    """
+    Sudoku puzzle generator using randomized backtracking.
+
+    Features:
+    - Supports arbitrary square sizes (e.g., 9x9, 4x4)
+    - Generates a fully solved board first
+    - Removes cells based on difficulty to create a puzzle
+    - Avoids relying on a single canonical solution
+    """
+
+    def __init__(self, size: int = 9):
+        """
+        Initialize the generator.
+
+        Args:
+            size (int): Size of the Sudoku board (must be a perfect square).
+                        Examples: 9 for 9x9, 4 for 4x4.
+        """
+        self.size = size
+        self.block = int(math.sqrt(size))
+        assert self.block * self.block == size, "Size must be a perfect square"
+
+    def generate(self, difficulty: str = "medium"):
+        """
+        Generate a Sudoku puzzle and its solution.
+
+        Args:
+            difficulty (str): Difficulty level ("easy", "medium", "hard").
+
+        Returns:
+            tuple: (puzzle, solution), where puzzle contains zeros for empty cells.
+        """
+        holes_map = {
+            "easy": self.size * self.size // 3,
+            "medium": self.size * self.size // 2,
+            "hard": self.size * self.size * 2 // 3,
+        }
+        holes = holes_map.get(difficulty, holes_map["medium"])
+
+        board = [[0 for _ in range(self.size)] for _ in range(self.size)]
+        self._fill_board(board)
+
+        solution = [row[:] for row in board]
+        self._remove_cells(board, holes)
+
+        return board, solution
+
+    def _fill_board(self, board):
+        """
+        Recursively fill the board using backtracking.
+
+        Args:
+            board (list[list[int]]): Current board state.
+
+        Returns:
+            bool: True if the board is completely filled.
+        """
+        empty = self._find_empty(board)
+        if not empty:
+            return True
+
+        r, c = empty
+        nums = list(range(1, self.size + 1))
+        random.shuffle(nums)
+
+        for v in nums:
+            if self._is_valid(board, r, c, v):
+                board[r][c] = v
+                if self._fill_board(board):
+                    return True
+                board[r][c] = 0
+
+        return False
+
+    def _find_empty(self, board):
+        """
+        Find the next empty cell in the board.
+
+        Args:
+            board (list[list[int]]): Current board state.
+
+        Returns:
+            tuple | None: (row, col) of empty cell, or None if full.
+        """
+        for i in range(self.size):
+            for j in range(self.size):
+                if board[i][j] == 0:
+                    return i, j
+        return None
+
+    def _is_valid(self, board, r, c, v):
+        """
+        Check whether placing value v at (r, c) is valid.
+
+        Args:
+            board (list[list[int]]): Current board state.
+            r (int): Row index.
+            c (int): Column index.
+            v (int): Value to place.
+
+        Returns:
+            bool: True if valid, False otherwise.
+        """
+        if v in board[r]:
+            return False
+
+        for i in range(self.size):
+            if board[i][c] == v:
+                return False
+
+        br = (r // self.block) * self.block
+        bc = (c // self.block) * self.block
+        for i in range(br, br + self.block):
+            for j in range(bc, bc + self.block):
+                if board[i][j] == v:
+                    return False
+
+        return True
+
+    def _remove_cells(self, board, holes):
+        """
+        Remove cells from a solved board to create a puzzle.
+
+        Args:
+            board (list[list[int]]): Solved board.
+            holes (int): Number of cells to clear.
+        """
+        cells = [(i, j) for i in range(self.size) for j in range(self.size)]
+        random.shuffle(cells)
+
+        for i in range(min(holes, self.size * self.size)):
+            r, c = cells[i]
+            board[r][c] = 0

trinity/common/workflows/envs/sudoku/sudoku_judge.py

@@ -0,0 +1,53 @@
+import math
+
+
+class SudokuJudge:
+    """
+    Judge Sudoku board state.
+
+    - Supports both 9x9 and 4x4 Sudoku boards
+    - Allows incomplete boards (zeros are treated as empty cells)
+    - Checks:
+        * Row validity
+        * Column validity
+        * Sub-grid validity (3x3 for 9x9, 2x2 for 4x4)
+    """
+
+    @staticmethod
+    def is_valid(board):
+        size = len(board)
+        block = int(math.sqrt(size))
+
+        # Check rows
+        for row in board:
+            nums = [v for v in row if v != 0]
+            if len(nums) != len(set(nums)):
+                return False
+
+        # Check columns
+        for c in range(size):
+            nums = []
+            for r in range(size):
+                v = board[r][c]
+                if v != 0:
+                    nums.append(v)
+            if len(nums) != len(set(nums)):
+                return False
+
+        # Check sub-grids
+        for br in range(0, size, block):
+            for bc in range(0, size, block):
+                nums = []
+                for r in range(br, br + block):
+                    for c in range(bc, bc + block):
+                        v = board[r][c]
+                        if v != 0:
+                            nums.append(v)
+                if len(nums) != len(set(nums)):
+                    return False
+
+        return True
+
+    @staticmethod
+    def is_solved(board, solution):
+        return board == solution