Minor fixes (#1)

1. Add tests for rollout entrypoint decorator
2. Add dependency constraints and resolutions

Author: luyuzhe111

README.md

@@ -60,7 +60,7 @@ uv venv --python 3.13
 source .venv/bin/activate
 
 # Install with development dependencies
-uv pip install -e ".[dev]"
+uv sync --frozen --extra dev
 
 # Install pre-commit hooks
 pre-commit install

examples/strands_math_agent/basic_app.py

@@ -31,6 +31,7 @@ async def invoke_agent(payload):
     print("User input:", user_input)
 
     response = await agent.invoke_async(user_input)
+
     return response.message["content"][0]["text"]
 
 

examples/strands_math_agent/pyproject.toml

@@ -6,7 +6,7 @@ readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
     "bedrock-agentcore>=1.0.3",
-    "bedrock-agentcore-starter-toolkit>=0.1.34",
+    "bedrock-agentcore-starter-toolkit<=0.2.0",
     "boto3>=1.40.55",
     "python-dotenv>=1.0.0",
     # TODO: replace the above dependencies with agentcore-rl-toolkit>=0.1.0 after PyPI indexing

examples/strands_math_agent/reward.py

@@ -4,6 +4,25 @@
 
 
 class GSM8KReward(RewardFunction):
+    def __call__(
+        self,
+        response_text="",
+        ground_truth="",
+        method="strict",
+        format_score=0.0,
+        score=1.0,
+        **kwargs,
+    ):
+        answer = self.extract_solution(solution_str=response_text, method=method)
+        if answer is None:
+            reward = 0
+        else:
+            if answer == ground_truth:
+                reward = score
+            else:
+                reward = format_score
+        return reward
+
     @staticmethod
     def extract_solution(solution_str, method="strict"):
         """
@@ -40,22 +59,3 @@ def extract_solution(solution_str, method="strict"):
                     if final_answer not in invalid_str:
                         break
         return final_answer
-
-    def __call__(
-        self,
-        response_text="",
-        ground_truth="",
-        method="strict",
-        format_score=0.0,
-        score=1.0,
-        **kwargs,
-    ):
-        answer = self.extract_solution(solution_str=response_text, method=method)
-        if answer is None:
-            reward = 0
-        else:
-            if answer == ground_truth:
-                reward = score
-            else:
-                reward = format_score
-        return reward

pyproject.toml

@@ -10,7 +10,7 @@ readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
     "bedrock-agentcore>=1.0.3",
-    "bedrock-agentcore-starter-toolkit>=0.1.34",
+    "bedrock-agentcore-starter-toolkit<=0.2.0",
     "boto3>=1.40.55",
     "python-dotenv>=1.0.0",
 ]

src/agentcore_rl_toolkit/app.py

@@ -243,5 +243,10 @@ async def rollout_entrypoint_wrapper(payload, context):
             asyncio.create_task(rollout_background_task(payload, context))
             return {"status": "processing"}
 
+        # Remove __wrapped__ so inspect.signature() sees the wrapper's actual signature
+        # (payload, context) instead of the user function's signature. This ensures
+        # BedrockAgentCoreApp._takes_context() correctly passes context to this wrapper.
+        del rollout_entrypoint_wrapper.__wrapped__
+
         # Register using existing BedrockAgentCoreApp entrypoint infrastructure
         return self.entrypoint(rollout_entrypoint_wrapper)

src/agentcore_rl_toolkit/frameworks/strands/app.py

@@ -2,13 +2,29 @@
 
 
 class StrandsAgentCoreRLApp(AgentCoreRLApp):
-    def create_openai_compatible_model(self, **kwargs):
-        """Create Strands OpenAI-compatible model for vLLM/SGLang server."""
+    def create_openai_compatible_model(self, provider_model_id=None, **kwargs):
+        """
+        Create Strands model that's compatible with the OpenAI format. When provider_model_id
+        is provided, LiteLLM model will be used. Otherwise, an OpenAI compatible model with
+        base_url and model_id will be used.
+
+        :param provider_model_id: Provide this parameter when using cloud providers (bedrock,
+        anthropic, openai, etc.) that does not use a base_url. Example: Otherwise, leave it to None.
+        """
         try:
             from strands.models.openai import OpenAIModel
         except ImportError:
-            raise ImportError("Strands not installed. Install with: uv pip install strands-agents[openai]") from None
+            raise ImportError("Strands not installed. Install with: " "uv pip install strands-agents[openai]") from None
+
+        if not provider_model_id:
+            base_url, model_id = self._get_model_config()
+            return OpenAIModel(client_args={"api_key": "dummy", "base_url": base_url}, model_id=model_id, **kwargs)
 
-        base_url, model_id = self._get_model_config()
+        try:
+            from strands.models.litellm import LiteLLMModel
+        except ImportError:
+            raise ImportError(
+                "Strands not installed. Install with: " "uv pip install strands-agents[litellm]"
+            ) from None
 
-        return OpenAIModel(client_args={"api_key": "dummy", "base_url": base_url}, model_id=model_id, **kwargs)
+        return LiteLLMModel(model_id=provider_model_id, **kwargs)

src/agentcore_rl_toolkit/frameworks/strands/rollout_collector.py

@@ -13,7 +13,7 @@ def register_hooks(self, registry):
             from strands.experimental.hooks import BeforeModelInvocationEvent
             from strands.hooks import AfterInvocationEvent
         except ImportError:
-            raise ImportError("Strands not installed. Install with: uv pip install strands-agents[openai]") from None
+            raise ImportError("Strands not installed. Install with: " "uv pip install strands-agents[openai]") from None
 
         registry.add_callback(BeforeModelInvocationEvent, self.collect_messages)
         registry.add_callback(AfterInvocationEvent, self.prepare_rollout)

tests/test_rollout_entrypoint.py

@@ -0,0 +1,91 @@
+"""Tests for the @rollout_entrypoint decorator."""
+
+import inspect
+
+from starlette.testclient import TestClient
+
+from agentcore_rl_toolkit import AgentCoreRLApp
+
+
+class MockAgentCoreRLApp(AgentCoreRLApp):
+    """Minimal concrete implementation for testing."""
+
+    def create_openai_compatible_model(self, **kwargs):
+        return None
+
+
+def test_wrapper_signature_has_context():
+    """Test that the wrapper's signature includes (payload, context) for BedrockAgentCoreApp."""
+    app = MockAgentCoreRLApp()
+
+    @app.rollout_entrypoint
+    async def my_handler(payload: dict):
+        return {"rollout_data": [], "rewards": [0]}
+
+    wrapper = app.handlers["main"]
+    params = list(inspect.signature(wrapper).parameters.keys())
+
+    assert len(params) == 2
+    assert params[0] == "payload"
+    assert params[1] == "context"
+
+
+def test_wrapper_preserves_function_name():
+    """Test that @wraps preserves the original function name."""
+    app = MockAgentCoreRLApp()
+
+    @app.rollout_entrypoint
+    async def my_custom_handler(payload: dict):
+        return {"rollout_data": [], "rewards": [0]}
+
+    wrapper = app.handlers["main"]
+    assert wrapper.__name__ == "my_custom_handler"
+
+
+def test_entrypoint_with_payload_only():
+    """Test that user function with signature (payload) works."""
+    app = MockAgentCoreRLApp()
+
+    @app.rollout_entrypoint
+    async def handler(payload: dict):
+        return {"rollout_data": [{"test": True}], "rewards": [1.0]}
+
+    client = TestClient(app)
+    response = client.post("/invocations", json={"prompt": "test"})
+
+    assert response.status_code == 200
+    assert response.json() == {"status": "processing"}
+
+
+def test_entrypoint_with_payload_and_context():
+    """Test that user function with signature (payload, context) works."""
+    app = MockAgentCoreRLApp()
+
+    @app.rollout_entrypoint
+    async def handler(payload: dict, context):
+        return {"rollout_data": [{"session": context.session_id}], "rewards": [1.0]}
+
+    client = TestClient(app)
+    response = client.post(
+        "/invocations",
+        json={"prompt": "test"},
+        headers={"X-Amz-Bedrock-AgentCore-Session-Id": "session-123"},
+    )
+
+    assert response.status_code == 200
+    assert response.json() == {"status": "processing"}
+
+
+def test_entrypoint_with_sync_handler():
+    """Test that sync user function works."""
+    app = MockAgentCoreRLApp()
+
+    @app.rollout_entrypoint
+    def handler(payload: dict):
+        return {"rollout_data": [{"sync": True}], "rewards": [1.0]}
+
+    client = TestClient(app)
+    response = client.post("/invocations", json={"prompt": "test"})
+
+    assert response.status_code == 200
+    assert response.json() == {"status": "processing"}