use Blueprint and Snapshot objects directly in ScenarioBuilder

sid-rl · sid-rl · commit d3a53ce38e2c · 2025-12-09T13:27:18.000-08:00
diff --git a/src/runloop_api_client/sdk/async_scenario_builder.py b/src/runloop_api_client/sdk/async_scenario_builder.py
@@ -7,6 +7,8 @@
 
 from .._client import AsyncRunloop
 from .async_scenario import AsyncScenario
+from .async_snapshot import AsyncSnapshot
+from .async_blueprint import AsyncBlueprint
 from ..types.scoring_function_param import (
     Scorer,
     ScoringFunctionParam,
@@ -28,7 +30,7 @@ class AsyncScenarioBuilder:
 
     Example:
         >>> builder = sdk.scenario.builder("my-scenario")
-        >>> builder.from_blueprint_id("bp-xxx")
+        >>> builder.from_blueprint(blueprint)
         >>> builder.with_working_directory("/app")
         >>> builder.with_problem_statement("Fix the bug in main.py")
         >>> builder.add_test_scorer("tests", test_command="pytest")
@@ -47,8 +49,8 @@ def __init__(self, client: AsyncRunloop, name: str) -> None:
         self._name = name
 
         # Environment configuration
-        self._blueprint_id: Optional[str] = None
-        self._snapshot_id: Optional[str] = None
+        self._blueprint: Optional[AsyncBlueprint] = None
+        self._snapshot: Optional[AsyncSnapshot] = None
         self._working_directory: Optional[str] = None
 
         # Input context
@@ -78,28 +80,28 @@ def name(self) -> str:
         """
         return self._name
 
-    def from_blueprint_id(self, blueprint_id: str) -> Self:
-        """Set the blueprint ID for the scenario environment.
+    def from_blueprint(self, blueprint: AsyncBlueprint) -> Self:
+        """Set the blueprint for the scenario environment.
 
-        :param blueprint_id: Blueprint ID to use
-        :type blueprint_id: str
+        :param blueprint: Blueprint to use
+        :type blueprint: AsyncBlueprint
         :return: Self for method chaining
         :rtype: Self
         """
-        self._blueprint_id = blueprint_id
-        self._snapshot_id = None  # Clear snapshot if blueprint is set
+        self._blueprint = blueprint
+        self._snapshot = None  # Clear snapshot if blueprint is set
         return self
 
-    def from_snapshot_id(self, snapshot_id: str) -> Self:
-        """Set the snapshot ID for the scenario environment.
+    def from_snapshot(self, snapshot: AsyncSnapshot) -> Self:
+        """Set the snapshot for the scenario environment.
 
-        :param snapshot_id: Snapshot ID to use
-        :type snapshot_id: str
+        :param snapshot: Snapshot to use
+        :type snapshot: AsyncSnapshot
         :return: Self for method chaining
         :rtype: Self
         """
-        self._snapshot_id = snapshot_id
-        self._blueprint_id = None  # Clear blueprint if snapshot is set
+        self._snapshot = snapshot
+        self._blueprint = None  # Clear blueprint if snapshot is set
         return self
 
     def with_working_directory(self, directory: str) -> Self:
@@ -415,10 +417,10 @@ def _build_params(self) -> Dict[str, Any]:
 
         # Build environment parameters if any are set
         env_params: Dict[str, Any] = {}
-        if self._blueprint_id:
-            env_params["blueprint_id"] = self._blueprint_id
-        if self._snapshot_id:
-            env_params["snapshot_id"] = self._snapshot_id
+        if self._blueprint:
+            env_params["blueprint_id"] = self._blueprint.id
+        if self._snapshot:
+            env_params["snapshot_id"] = self._snapshot.id
         if self._working_directory:
             env_params["working_directory"] = self._working_directory
 
diff --git a/src/runloop_api_client/sdk/scenario_builder.py b/src/runloop_api_client/sdk/scenario_builder.py
@@ -7,6 +7,8 @@
 
 from .._client import Runloop
 from .scenario import Scenario
+from .snapshot import Snapshot
+from .blueprint import Blueprint
 from ..types.scoring_function_param import (
     Scorer,
     ScoringFunctionParam,
@@ -28,7 +30,7 @@ class ScenarioBuilder:
 
     Example:
         >>> builder = sdk.scenario.builder("my-scenario")
-        >>> builder.from_blueprint_id("bp-xxx")
+        >>> builder.from_blueprint(blueprint)
         >>> builder.with_working_directory("/app")
         >>> builder.with_problem_statement("Fix the bug in main.py")
         >>> builder.add_test_scorer("tests", test_command="pytest")
@@ -47,8 +49,8 @@ def __init__(self, client: Runloop, name: str) -> None:
         self._name = name
 
         # Environment configuration
-        self._blueprint_id: Optional[str] = None
-        self._snapshot_id: Optional[str] = None
+        self._blueprint: Optional[Blueprint] = None
+        self._snapshot: Optional[Snapshot] = None
         self._working_directory: Optional[str] = None
 
         # Input context
@@ -78,28 +80,28 @@ def name(self) -> str:
         """
         return self._name
 
-    def from_blueprint_id(self, blueprint_id: str) -> Self:
-        """Set the blueprint ID for the scenario environment.
+    def from_blueprint(self, blueprint: Blueprint) -> Self:
+        """Set the blueprint for the scenario environment.
 
-        :param blueprint_id: Blueprint ID to use
-        :type blueprint_id: str
+        :param blueprint: Blueprint to use
+        :type blueprint: Blueprint
         :return: Self for method chaining
         :rtype: Self
         """
-        self._blueprint_id = blueprint_id
-        self._snapshot_id = None  # Clear snapshot if blueprint is set
+        self._blueprint = blueprint
+        self._snapshot = None  # Clear snapshot if blueprint is set
         return self
 
-    def from_snapshot_id(self, snapshot_id: str) -> Self:
-        """Set the snapshot ID for the scenario environment.
+    def from_snapshot(self, snapshot: Snapshot) -> Self:
+        """Set the snapshot for the scenario environment.
 
-        :param snapshot_id: Snapshot ID to use
-        :type snapshot_id: str
+        :param snapshot: Snapshot to use
+        :type snapshot: Snapshot
         :return: Self for method chaining
         :rtype: Self
         """
-        self._snapshot_id = snapshot_id
-        self._blueprint_id = None  # Clear blueprint if snapshot is set
+        self._snapshot = snapshot
+        self._blueprint = None  # Clear blueprint if snapshot is set
         return self
 
     def with_working_directory(self, directory: str) -> Self:
@@ -415,10 +417,10 @@ def _build_params(self) -> Dict[str, Any]:
 
         # Build environment parameters if any are set
         env_params: Dict[str, Any] = {}
-        if self._blueprint_id:
-            env_params["blueprint_id"] = self._blueprint_id
-        if self._snapshot_id:
-            env_params["snapshot_id"] = self._snapshot_id
+        if self._blueprint:
+            env_params["blueprint_id"] = self._blueprint.id
+        if self._snapshot:
+            env_params["snapshot_id"] = self._snapshot.id
         if self._working_directory:
             env_params["working_directory"] = self._working_directory
 
diff --git a/tests/sdk/test_async_scenario_builder.py b/tests/sdk/test_async_scenario_builder.py
@@ -6,6 +6,8 @@
 
 import pytest
 
+from runloop_api_client.sdk.async_snapshot import AsyncSnapshot
+from runloop_api_client.sdk.async_blueprint import AsyncBlueprint
 from runloop_api_client.sdk.async_scenario_builder import AsyncScenarioBuilder
 
 
@@ -20,6 +22,16 @@ def mock_async_client(self) -> MagicMock:
         client.scenarios.create = AsyncMock()
         return client
 
+    @pytest.fixture
+    def mock_blueprint(self, mock_async_client: MagicMock) -> AsyncBlueprint:
+        """Create a mock AsyncBlueprint object."""
+        return AsyncBlueprint(mock_async_client, "bp-123")
+
+    @pytest.fixture
+    def mock_snapshot(self, mock_async_client: MagicMock) -> AsyncSnapshot:
+        """Create a mock AsyncSnapshot object."""
+        return AsyncSnapshot(mock_async_client, "snap-123")
+
     @pytest.fixture
     def builder(self, mock_async_client: MagicMock) -> AsyncScenarioBuilder:
         """Create an AsyncScenarioBuilder instance with mock client."""
@@ -37,21 +49,21 @@ def test_repr(self, builder: AsyncScenarioBuilder) -> None:
         """Test builder __repr__."""
         assert repr(builder) == "<AsyncScenarioBuilder name='test-scenario'>"
 
-    def test_from_blueprint_id_returns_self(self, builder: AsyncScenarioBuilder) -> None:
-        """Test from_blueprint_id returns self for chaining."""
-        result = builder.from_blueprint_id("bp-123")
+    def test_from_blueprint_returns_self(self, builder: AsyncScenarioBuilder, mock_blueprint: AsyncBlueprint) -> None:
+        """Test from_blueprint returns self for chaining."""
+        result = builder.from_blueprint(mock_blueprint)
 
         assert result is builder
-        assert builder._blueprint_id == "bp-123"
-        assert builder._snapshot_id is None
+        assert builder._blueprint is mock_blueprint
+        assert builder._snapshot is None
 
-    def test_from_snapshot_id_returns_self(self, builder: AsyncScenarioBuilder) -> None:
-        """Test from_snapshot_id returns self for chaining."""
-        result = builder.from_snapshot_id("snap-123")
+    def test_from_snapshot_returns_self(self, builder: AsyncScenarioBuilder, mock_snapshot: AsyncSnapshot) -> None:
+        """Test from_snapshot returns self for chaining."""
+        result = builder.from_snapshot(mock_snapshot)
 
         assert result is builder
-        assert builder._snapshot_id == "snap-123"
-        assert builder._blueprint_id is None
+        assert builder._snapshot is mock_snapshot
+        assert builder._blueprint is None
 
     def test_with_working_directory_returns_self(self, builder: AsyncScenarioBuilder) -> None:
         """Test with_working_directory returns self for chaining."""
@@ -125,11 +137,11 @@ def test_build_params_minimal(self, builder: AsyncScenarioBuilder) -> None:
         assert params["input_context"]["problem_statement"] == "Fix the bug"
         assert len(params["scoring_contract"]["scoring_function_parameters"]) == 1
 
-    def test_build_params_with_environment(self, builder: AsyncScenarioBuilder) -> None:
+    def test_build_params_with_environment(self, builder: AsyncScenarioBuilder, mock_blueprint: AsyncBlueprint) -> None:
         """Test _build_params includes environment parameters."""
         builder.with_problem_statement("Fix the bug")
         builder.add_test_scorer("tests", test_command="pytest")
-        builder.from_blueprint_id("bp-123")
+        builder.from_blueprint(mock_blueprint)
         builder.with_working_directory("/app")
 
         params = builder._build_params()
@@ -156,10 +168,10 @@ async def test_push_calls_api_and_returns_scenario(
 
         assert scenario.id == "scn-new-123"
 
-    def test_fluent_chaining(self, builder: AsyncScenarioBuilder) -> None:
+    def test_fluent_chaining(self, builder: AsyncScenarioBuilder, mock_blueprint: AsyncBlueprint) -> None:
         """Test that all builder methods can be chained fluently."""
         result = (
-            builder.from_blueprint_id("bp-123")
+            builder.from_blueprint(mock_blueprint)
             .with_working_directory("/app")
             .with_problem_statement("Fix the bug")
             .with_additional_context({"hint": "check main.py"})
@@ -172,7 +184,7 @@ def test_fluent_chaining(self, builder: AsyncScenarioBuilder) -> None:
         )
 
         assert result is builder
-        assert builder._blueprint_id == "bp-123"
+        assert builder._blueprint is mock_blueprint
         assert builder._working_directory == "/app"
         assert builder._problem_statement == "Fix the bug"
         assert len(builder._scorers) == 1
diff --git a/tests/sdk/test_scenario_builder.py b/tests/sdk/test_scenario_builder.py
