Merge pull request #23 from AssassinWS/main

add the smoke test for chess environment

Author: AssassinWS

src/agentfly/rewards/__init__.py

@@ -1,4 +1,5 @@
 from .alfworld_reward import alfworld_episode_reward
+from .chess_reward import chess_puzzle_reward, chess_puzzle_reward_simple
 from .code_reward import code_reward_test
 from .gui_reward import gui_reward
 from .math_reward import (
@@ -37,4 +38,6 @@
     "scienceworld_reward",
     "gui_reward",
     "code_reward_test",
+    "chess_puzzle_reward",
+    "chess_puzzle_reward_simple",
 ]

src/agentfly/rewards/chess_reward.py

@@ -14,7 +14,7 @@
 
 
 @reward(name="chess_puzzle_reward", env_cls=ChessPuzzleEnv, pool_size=8)
-async def chess_puzzle_reward(prediction: str, env: ChessPuzzleEnv) -> Dict[str, Any]:
+async def chess_puzzle_reward(final_response: str, env: ChessPuzzleEnv) -> Dict[str, Any]:
     """
     Calculate reward for chess puzzle solving based on Stockfish evaluation.
 
@@ -29,7 +29,7 @@ async def chess_puzzle_reward(prediction: str, env: ChessPuzzleEnv) -> Dict[str,
     - Making progress even with imperfect moves
 
     Args:
-        prediction (str): The agent's final response/output (not used directly).
+        final_response (str): The agent's final response/output (not used directly).
         env (ChessPuzzleEnv): The chess puzzle environment instance.
 
     Returns:
@@ -139,7 +139,7 @@ async def chess_puzzle_reward(prediction: str, env: ChessPuzzleEnv) -> Dict[str,
 
 @reward(name="chess_puzzle_reward_simple", env_cls=ChessPuzzleEnv, pool_size=8)
 async def chess_puzzle_reward_simple(
-    prediction: str, env: ChessPuzzleEnv
+    final_response: str, env: ChessPuzzleEnv
 ) -> Dict[str, Any]:
     """
     Simple binary reward for chess puzzle solving.
@@ -149,7 +149,7 @@ async def chess_puzzle_reward_simple(
     where you only care about correct solutions.
 
     Args:
-        prediction (str): The agent's final response/output (not used).
+        final_response (str): The agent's final response/output (not used).
         env (ChessPuzzleEnv): The chess puzzle environment instance.
 
     Returns:

tests/unit/envs/test_chess_smoke.py

@@ -0,0 +1,119 @@
+# test_chess_smoke.py
+"""
+End-to-end smoke test for the chess puzzle agent pipeline.
+
+Verifies that ReactAgent + chess tools + chess reward work together:
+- Agent receives a puzzle
+- Tools dispatch to ChessPuzzleEnv
+- Trajectories are populated after run()
+- Reward function produces a valid result
+
+Requires: Stockfish installed (brew install stockfish / apt-get install stockfish)
+"""
+
+import shutil
+from unittest.mock import AsyncMock
+
+import pytest
+
+from agentfly.agents import ReactAgent
+from agentfly.envs.chess_env import ChessPuzzleEnv
+from agentfly.rewards import chess_puzzle_reward
+from agentfly.tools import chess_get_legal_moves, chess_get_state, chess_move
+
+# Skip if Stockfish is not available
+pytestmark = pytest.mark.skipif(
+    shutil.which("stockfish") is None,
+    reason="Stockfish not installed",
+)
+
+# A simple mate-in-1 puzzle: White plays Qxf7#
+MATE_IN_1_PUZZLE = {
+    "puzzle_id": "smoke_mate1",
+    "fen": "r1bqkb1r/pppp1ppp/2n2n2/4p2Q/2B1P3/8/PPPP1PPP/RNB1K1NR w KQkq - 4 4",
+    "moves": "h5f7",
+}
+
+
+def _make_react_responses():
+    """Return a sequence of canned ReAct responses for the smoke test.
+
+    Turn 0: agent checks the board state
+    Turn 1: agent makes the winning move Qxf7#
+    """
+    return [
+        # Turn 0 – get state
+        (
+            'Thought: Let me look at the current board position first.\n'
+            'Action: chess_get_state\n'
+            'Input: {}'
+        ),
+        # Turn 1 – make the winning move
+        (
+            'Thought: I see White can play Qxf7# for checkmate.\n'
+            'Action: chess_move\n'
+            'Input: {"move": "h5f7"}'
+        ),
+    ]
+
+
+@pytest.mark.asyncio
+async def test_chess_smoke_e2e():
+    """Smoke test: ReactAgent solves a mate-in-1 puzzle with chess tools."""
+
+    canned = _make_react_responses()
+    call_idx = 0
+
+    async def fake_generate(messages_list, **kwargs):
+        nonlocal call_idx
+        idx = min(call_idx, len(canned) - 1)
+        call_idx += 1
+        return [canned[idx]]
+
+    tools = [chess_move, chess_get_state, chess_get_legal_moves]
+
+    agent = ReactAgent(
+        model_name_or_path="Qwen/Qwen2.5-3B-Instruct",
+        tools=tools,
+        backend="client",
+        reward_fn=chess_puzzle_reward,
+        monitors=[],
+        debug=True,
+    )
+
+    # Replace LLM engine methods with mocks
+    agent.llm_engine = AsyncMock()
+    agent.llm_engine.generate_async = fake_generate
+    agent.llm_engine.preprocess = lambda: None
+    agent.llm_engine.postprocess = lambda: None
+
+    messages = [
+        {
+            "messages": [
+                {
+                    "role": "user",
+                    "content": (
+                        "Solve this chess puzzle. The position is a mate-in-1. "
+                        "Find the winning move for White."
+                    ),
+                }
+            ],
+            **MATE_IN_1_PUZZLE,
+        }
+    ]
+
+    await agent.run(
+        messages=messages,
+        max_turns=2,
+        num_chains=1,
+        enable_streaming=False,
+    )
+
+    # Trajectories should be populated
+    trajectories = agent.trajectories
+    assert len(trajectories) > 0, "Expected at least one trajectory"
+
+    # The trajectory should contain messages
+    traj = trajectories[0]
+    assert "messages" in traj
+    assert len(traj["messages"]) > 1, "Expected multiple messages in trajectory"