plato-sdk-v2 2.0.64__py3-none-any.whl → 2.3.4__py3-none-any.whl

plato/v1/flow_executor.py

@@ -7,15 +7,13 @@ import logging
 import os
 import sys
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, cast
 from urllib.parse import urljoin
 
 sys.path.insert(0, os.path.dirname(__file__))
 
-try:
+if TYPE_CHECKING:
     from playwright.async_api import Page
-except ImportError:
-    Page = Any  # Fallback if playwright not installed
 
 from plato.v1.models.flow import (
     CheckElementStep,
@@ -42,7 +40,7 @@ class FlowExecutor:
 
     def __init__(
         self,
-        page: Page,
+        page: "Page",
         flow: Flow,
         screenshots_dir: Path | None = None,
         logger: logging.Logger = logging.getLogger(__name__),
@@ -108,29 +106,29 @@ class FlowExecutor:
     async def _execute_step(self, step: FlowStep) -> bool:
         """Execute a single step in a flow using type attribute."""
         if step.type == "wait_for_selector":
-            return await self._wait_for_selector(step)
+            return await self._wait_for_selector(cast(WaitForSelectorStep, step))
         elif step.type == "click":
-            return await self._click(step)
+            return await self._click(cast(ClickStep, step))
         elif step.type == "fill":
-            return await self._fill(step)
+            return await self._fill(cast(FillStep, step))
         elif step.type == "wait":
-            return await self._wait(step)
+            return await self._wait(cast(WaitStep, step))
         elif step.type == "navigate":
-            return await self._navigate(step)
+            return await self._navigate(cast(NavigateStep, step))
         elif step.type == "wait_for_url":
-            return await self._wait_for_url(step)
+            return await self._wait_for_url(cast(WaitForUrlStep, step))
         elif step.type == "check_element":
-            return await self._check_element(step)
+            return await self._check_element(cast(CheckElementStep, step))
         elif step.type == "verify":
-            return await self._verify(step)
+            return await self._verify(cast(VerifyStep, step))
         elif step.type == "screenshot":
-            return await self._screenshot(step)
+            return await self._screenshot(cast(ScreenshotStep, step))
         elif step.type == "verify_text":
-            return await self._verify_text(step)
+            return await self._verify_text(cast(VerifyTextStep, step))
         elif step.type == "verify_url":
-            return await self._verify_url(step)
+            return await self._verify_url(cast(VerifyUrlStep, step))
         elif step.type == "verify_no_errors":
-            return await self._verify_no_errors(step)
+            return await self._verify_no_errors(cast(VerifyNoErrorsStep, step))
         else:
             self.logger.error(f"❌ Unknown step type: {step.type}")
             return False
@@ -248,6 +246,7 @@ class FlowExecutor:
 
     async def _verify_element_exists(self, step: VerifyStep) -> bool:
         """Verify that an element exists in the DOM."""
+        assert step.selector is not None  # Guaranteed by VerifyStep model validator
         try:
             element = await self.page.query_selector(step.selector)
             if element:
@@ -262,6 +261,7 @@ class FlowExecutor:
 
     async def _verify_element_visible(self, step: VerifyStep) -> bool:
         """Verify that an element is visible on the page."""
+        assert step.selector is not None  # Guaranteed by VerifyStep model validator
         try:
             element = await self.page.query_selector(step.selector)
             if element:
@@ -281,6 +281,8 @@ class FlowExecutor:
 
     async def _verify_element_text(self, step: VerifyStep) -> bool:
         """Verify that an element contains specific text."""
+        assert step.selector is not None  # Guaranteed by VerifyStep model validator
+        assert step.text is not None  # Guaranteed by VerifyStep model validator
         try:
             element = await self.page.query_selector(step.selector)
             if element:
@@ -318,6 +320,7 @@ class FlowExecutor:
 
     async def _verify_element_count(self, step: VerifyStep) -> bool:
         """Verify the count of elements matching a selector."""
+        assert step.selector is not None  # Guaranteed by VerifyStep model validator
         try:
             elements = await self.page.query_selector_all(step.selector)
             actual_count = len(elements)
@@ -336,6 +339,7 @@ class FlowExecutor:
 
     async def _verify_page_title(self, step: VerifyStep) -> bool:
         """Verify the page title."""
+        assert step.title is not None  # Guaranteed by VerifyStep model validator
         try:
             actual_title = await self.page.title()
 

plato/v1/models/env.py

@@ -13,11 +13,6 @@ from urllib.parse import urlparse
 
 import yaml
 
-try:
-    from playwright.async_api import Page
-except ImportError:
-    Page = Any  # Fallback if playwright not installed
-
 from plato.v1.exceptions import PlatoClientError
 from plato.v1.flow_executor import FlowExecutor
 from plato.v1.models.flow import Flow
@@ -27,6 +22,8 @@ logger = logging.getLogger(__name__)
 
 # Using TYPE_CHECKING for proper type annotation without circular imports
 if TYPE_CHECKING:
+    from playwright.async_api import Page
+
     from plato.sdk import Plato
 
 
@@ -48,8 +45,8 @@ class PlatoEnvironment:
 
     _current_task: PlatoTask | None = None
     _client: "Plato" = None
-    id: str = None
-    env_id: str = None
+    id: str = None  # type: ignore
+    env_id: str = None  # type: ignore
     alias: str | None = None
     _run_session_id: str | None = None
     _heartbeat_task: asyncio.Task | None = None
@@ -66,7 +63,7 @@ class PlatoEnvironment:
     ):
         self._client = client
         self.id = id
-        self.env_id = env_id
+        self.env_id = env_id  # type: ignore[assignment]
         self.alias = alias
         self._run_session_id = None
         self._heartbeat_task = None
@@ -77,7 +74,7 @@ class PlatoEnvironment:
 
     async def login(
         self,
-        page: Page,
+        page: "Page",
         throw_on_login_error: bool = False,
         screenshots_dir: Path | None = None,
         dataset: str = "base",
@@ -97,7 +94,7 @@ class PlatoEnvironment:
                     f"{self._client.base_url}/env/{self.id}/flows",
                     headers=headers,
                 ) as resp:
-                    await self._client._handle_response_error(resp)  # type: ignore[attr-defined]
+                    await self._client._handle_response_error(resp)
                     body_text = await resp.text()
                     # Endpoint may return JSON with { data: { flows: "...yaml..." } } or raw YAML
                     try:
@@ -394,7 +391,10 @@ class PlatoEnvironment:
         if not self._run_session_id:
             raise PlatoClientError("No active run session. Call reset() first.")
 
-        if not self._current_task or not self._current_task.eval_config:
+        if not self._current_task:
+            logger.warning("No current task set")
+            raise PlatoClientError("No evaluation config found for task")
+        if not self._current_task.eval_config:
             logger.warning(f"No evaluation config found for task: {self._current_task.name}")
             raise PlatoClientError("No evaluation config found for task")
 

plato/v1/sdk.py

@@ -538,7 +538,7 @@ class Plato:
                     default_scoring_config=t.get("defaultScoringConfig", {}),
                     scoring_type=[ScoringType(st) for st in t.get("scoringTypes", [])]
                     if t.get("scoringTypes")
-                    else None,
+                    else [ScoringType.MUTATIONS],  # Use default when not provided
                     output_schema=t.get("outputSchema"),
                     is_sample=t.get("isSample", False),
                     simulator_artifact_id=(
@@ -685,7 +685,7 @@ class Plato:
             return await response.json()
 
     async def create_simulator(
-        self, name: str, description: str = None, sim_type: str = "docker_app"
+        self, name: str, description: str | None = None, sim_type: str = "docker_app"
     ) -> dict[str, Any]:
         """Create a new simulator.
 

plato/v1/sync_env.py

@@ -10,11 +10,6 @@ from urllib.parse import urlparse
 
 import yaml
 
-try:
-    from playwright.sync_api import Page
-except ImportError:
-    Page = Any  # Fallback if playwright not installed
-
 from plato.v1.exceptions import PlatoClientError
 from plato.v1.models.flow import Flow
 from plato.v1.models.task import CustomEvalConfig, EvaluationResult, PlatoTask
@@ -24,6 +19,8 @@ logger = logging.getLogger(__name__)
 
 # Using TYPE_CHECKING for proper type annotation without circular imports
 if TYPE_CHECKING:
+    from playwright.sync_api import Page
+
     from plato.sync_sdk import SyncPlato
 
 
@@ -45,8 +42,8 @@ class SyncPlatoEnvironment:
 
     _current_task: PlatoTask | None = None
     _client: "SyncPlato" = None
-    id: str = None
-    env_id: str = None
+    id: str = None  # type: ignore
+    env_id: str = None  # type: ignore
     alias: str | None = None
     _run_session_id: str | None = None
     _heartbeat_thread: threading.Thread | None = None
@@ -64,7 +61,7 @@ class SyncPlatoEnvironment:
     ):
         self._client = client
         self.id = id
-        self.env_id = env_id
+        self.env_id = env_id  # type: ignore[assignment]
         self.alias = alias
         self._run_session_id = active_session
         self._heartbeat_thread = None
@@ -75,7 +72,7 @@ class SyncPlatoEnvironment:
 
     def login(
         self,
-        page: Page,
+        page: "Page",
         throw_on_login_error: bool = False,
         screenshots_dir: Path | None = None,
         dataset: str = "base",
@@ -95,7 +92,7 @@ class SyncPlatoEnvironment:
             try:
                 headers = {"X-API-Key": self._client.api_key}
                 resp = self._client.http_session.get(f"{self._client.base_url}/env/{self.id}/flows", headers=headers)
-                self._client._handle_response_error(resp)  # type: ignore[attr-defined]
+                self._client._handle_response_error(resp)
                 body_text = resp.text
                 # Endpoint may return JSON with { data: { flows: "...yaml..." } } or raw YAML
                 try:
@@ -379,7 +376,10 @@ class SyncPlatoEnvironment:
         if not self._run_session_id:
             raise PlatoClientError("No active run session. Call reset() first.")
 
-        if not self._current_task or not self._current_task.eval_config:
+        if not self._current_task:
+            logger.warning("No current task set")
+            raise PlatoClientError("No evaluation config found for task")
+        if not self._current_task.eval_config:
             logger.warning(f"No evaluation config found for task: {self._current_task.name}")
             raise PlatoClientError("No evaluation config found for task")
 

plato/v1/sync_flow_executor.py

@@ -7,15 +7,13 @@ import logging
 import os
 import sys
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, cast
 from urllib.parse import urljoin
 
 sys.path.insert(0, os.path.dirname(__file__))
 
-try:
+if TYPE_CHECKING:
     from playwright.sync_api import Page
-except ImportError:
-    Page = Any  # Fallback if playwright not installed
 
 from plato.v1.models.flow import (
     CheckElementStep,
@@ -40,7 +38,7 @@ class SyncFlowExecutor:
 
     def __init__(
         self,
-        page: Page,
+        page: "Page",
         flow: Flow,
         screenshots_dir: Path | None = None,
         logger: logging.Logger = logging.getLogger(__name__),
@@ -104,29 +102,29 @@ class SyncFlowExecutor:
     def _execute_step(self, step: FlowStep) -> bool:
         """Execute a single step in a flow using action attribute."""
         if step.type == "wait_for_selector":
-            return self._wait_for_selector(step)
+            return self._wait_for_selector(cast(WaitForSelectorStep, step))
         elif step.type == "click":
-            return self._click(step)
+            return self._click(cast(ClickStep, step))
         elif step.type == "fill":
-            return self._fill(step)
+            return self._fill(cast(FillStep, step))
         elif step.type == "wait":
-            return self._wait(step)
+            return self._wait(cast(WaitStep, step))
         elif step.type == "navigate":
-            return self._navigate(step)
+            return self._navigate(cast(NavigateStep, step))
         elif step.type == "wait_for_url":
-            return self._wait_for_url(step)
+            return self._wait_for_url(cast(WaitForUrlStep, step))
         elif step.type == "check_element":
-            return self._check_element(step)
+            return self._check_element(cast(CheckElementStep, step))
         elif step.type == "verify":
-            return self._verify(step)
+            return self._verify(cast(VerifyStep, step))
         elif step.type == "screenshot":
-            return self._screenshot(step)
+            return self._screenshot(cast(ScreenshotStep, step))
         elif step.type == "verify_text":
-            return self._verify_text(step)
+            return self._verify_text(cast(VerifyTextStep, step))
         elif step.type == "verify_url":
-            return self._verify_url(step)
+            return self._verify_url(cast(VerifyUrlStep, step))
         elif step.type == "verify_no_errors":
-            return self._verify_no_errors(step)
+            return self._verify_no_errors(cast(VerifyNoErrorsStep, step))
         else:
             self.logger.error(f"❌ Unknown step action: {step.type}")
             return False
@@ -242,6 +240,7 @@ class SyncFlowExecutor:
 
     def _verify_element_exists(self, step: VerifyStep) -> bool:
         """Verify that an element exists in the DOM."""
+        assert step.selector is not None  # Guaranteed by VerifyStep model validator
         try:
             element = self.page.query_selector(step.selector)
             if element:
@@ -256,6 +255,7 @@ class SyncFlowExecutor:
 
     def _verify_element_visible(self, step: VerifyStep) -> bool:
         """Verify that an element is visible on the page."""
+        assert step.selector is not None  # Guaranteed by VerifyStep model validator
         try:
             element = self.page.query_selector(step.selector)
             if element:
@@ -275,6 +275,8 @@ class SyncFlowExecutor:
 
     def _verify_element_text(self, step: VerifyStep) -> bool:
         """Verify that an element contains specific text."""
+        assert step.selector is not None  # Guaranteed by VerifyStep model validator
+        assert step.text is not None  # Guaranteed by VerifyStep model validator
         try:
             element = self.page.query_selector(step.selector)
             if element:
@@ -312,6 +314,7 @@ class SyncFlowExecutor:
 
     def _verify_element_count(self, step: VerifyStep) -> bool:
         """Verify the count of elements matching a selector."""
+        assert step.selector is not None  # Guaranteed by VerifyStep model validator
         try:
             elements = self.page.query_selector_all(step.selector)
             actual_count = len(elements)
@@ -330,6 +333,7 @@ class SyncFlowExecutor:
 
     def _verify_page_title(self, step: VerifyStep) -> bool:
         """Verify the page title."""
+        assert step.title is not None  # Guaranteed by VerifyStep model validator
         try:
             actual_title = self.page.title()
 

plato/v1/sync_sdk.py

@@ -89,7 +89,7 @@ class SyncPlato:
         Raises:
             PlatoClientError: With the actual error message from the response
         """
-        if response.status_code >= 400:
+        if response.status_code >= 400:  # type: ignore[operator]
             try:
                 # Try to get the error message from the response body
                 error_data = response.json()
@@ -525,7 +525,9 @@ class SyncPlato:
                 average_steps=t.get("averageStepsTaken"),
                 num_validator_human_scores=t.get("defaultScoringConfig", {}).get("num_sessions_used", 0),
                 default_scoring_config=t.get("defaultScoringConfig", {}),
-                scoring_type=[ScoringType(st) for st in t.get("scoringTypes", [])] if t.get("scoringTypes") else None,
+                scoring_type=[ScoringType(st) for st in t.get("scoringTypes", [])]
+                if t.get("scoringTypes")
+                else [ScoringType.MUTATIONS],
                 output_schema=t.get("outputSchema"),
                 is_sample=t.get("isSample", False),
                 simulator_artifact_id=(

plato/v2/__init__.py

@@ -15,6 +15,7 @@ from plato.v2.async_.client import AsyncPlato
 from plato.v2.async_.environment import Environment as AsyncEnvironment
 from plato.v2.async_.flow_executor import FlowExecutionError as AsyncFlowExecutionError
 from plato.v2.async_.flow_executor import FlowExecutor as AsyncFlowExecutor
+from plato.v2.async_.session import SerializedSession
 from plato.v2.async_.session import Session as AsyncSession
 from plato.v2.sync.client import Plato
 from plato.v2.sync.environment import Environment
@@ -45,6 +46,7 @@ __all__ = [
     "AsyncEnvironment",
     "AsyncFlowExecutor",
     "AsyncFlowExecutionError",
+    "SerializedSession",
     # Models
     "Flow",
     # Helpers

plato/v2/async_/environment.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING
 
 from plato._generated.api.v2 import jobs
 from plato._generated.models import (
+    AppApiV2SchemasSessionCreateSnapshotRequest,
     ConnectRoutingInfoResult,
     CreateCheckpointRequest,
     CreateSnapshotResult,
@@ -136,6 +137,36 @@ class Environment:
             x_api_key=self._api_key,
         )
 
+    async def snapshot_store(
+        self,
+        override_service: str | None = None,
+        override_version: str | None = None,
+        override_dataset: str | None = None,
+    ) -> CreateSnapshotResult:
+        """Create a snapshot-store snapshot of this environment.
+
+        Uses the snapshot-store pipeline for chunk-based deduplication and
+        efficient storage. This is the preferred method for new base snapshots.
+
+        Args:
+            override_service: Override simulator/service name in artifact metadata.
+            override_version: Override version/git_hash in artifact metadata.
+            override_dataset: Override dataset name in artifact metadata.
+
+        Returns:
+            CreateSnapshotResult with artifact_id.
+        """
+        return await jobs.snapshot_store.asyncio(
+            client=self._http,
+            job_id=self.job_id,
+            body=AppApiV2SchemasSessionCreateSnapshotRequest(
+                override_service=override_service,
+                override_version=override_version,
+                override_dataset=override_dataset,
+            ),
+            x_api_key=self._api_key,
+        )
+
     async def close(self) -> None:
         """Close this environment."""
         await jobs.close.asyncio(

plato/v2/async_/session.py

@@ -32,7 +32,9 @@ from plato._generated.api.v2.sessions import heartbeat as sessions_heartbeat
 from plato._generated.api.v2.sessions import make as sessions_make
 from plato._generated.api.v2.sessions import reset as sessions_reset
 from plato._generated.api.v2.sessions import set_date as sessions_set_date
+from plato._generated.api.v2.sessions import setup_sandbox as sessions_setup_sandbox
 from plato._generated.api.v2.sessions import snapshot as sessions_snapshot
+from plato._generated.api.v2.sessions import snapshot_store as sessions_snapshot_store
 from plato._generated.api.v2.sessions import state as sessions_state
 from plato._generated.api.v2.sessions import wait_for_ready as sessions_wait_for_ready
 from plato._generated.models import (
@@ -40,6 +42,8 @@ from plato._generated.models import (
     AppApiV2SchemasSessionCreateSnapshotResponse,
     AppApiV2SchemasSessionEvaluateResponse,
     AppApiV2SchemasSessionHeartbeatResponse,
+    AppApiV2SchemasSessionSetupSandboxRequest,
+    AppApiV2SchemasSessionSetupSandboxResponse,
     CreateDiskSnapshotRequest,
     CreateDiskSnapshotResponse,
     CreateSessionFromEnvs,
@@ -557,6 +561,38 @@ class Session:
             x_api_key=self._api_key,
         )
 
+    async def setup_sandbox(
+        self,
+        timeout: int = 120,
+    ) -> AppApiV2SchemasSessionSetupSandboxResponse:
+        """Setup sandbox environment with Docker overlay on all environments.
+
+        This configures the VMs for Docker usage with overlay2 storage driver,
+        which is significantly faster than the default vfs driver. Should be called
+        after session creation and before pulling Docker images.
+
+        The setup includes:
+        - Mounting /dev/vdb to /mnt/docker for Docker storage
+        - Configuring Docker with overlay2 storage driver
+        - Setting up ECR and Docker Hub authentication
+        - Creating a docker-user service for non-root Docker access
+
+        Args:
+            timeout: Setup timeout in seconds (default: 120).
+
+        Returns:
+            SetupSandboxResponse with results per job_id.
+        """
+        self._check_closed()
+
+        request = AppApiV2SchemasSessionSetupSandboxRequest(timeout=timeout)
+        return await sessions_setup_sandbox.asyncio(
+            client=self._http,
+            session_id=self.session_id,
+            body=request,
+            x_api_key=self._api_key,
+        )
+
     async def evaluate(self, **kwargs) -> AppApiV2SchemasSessionEvaluateResponse:
         """Evaluate the session against task criteria.
 
@@ -587,6 +623,38 @@ class Session:
             x_api_key=self._api_key,
         )
 
+    async def snapshot_store(
+        self,
+        override_service: str | None = None,
+        override_version: str | None = None,
+        override_dataset: str | None = None,
+    ) -> AppApiV2SchemasSessionCreateSnapshotResponse:
+        """Create a snapshot-store snapshot of all environments in the session.
+
+        Uses the snapshot-store pipeline for chunk-based deduplication and
+        efficient storage. This is the preferred method for new base snapshots.
+
+        Args:
+            override_service: Override simulator/service name in artifact metadata.
+            override_version: Override version/git_hash in artifact metadata.
+            override_dataset: Override dataset name in artifact metadata.
+
+        Returns:
+            Snapshot response with info per job_id.
+        """
+        self._check_closed()
+
+        return await sessions_snapshot_store.asyncio(
+            client=self._http,
+            session_id=self.session_id,
+            body=AppApiV2SchemasSessionCreateSnapshotRequest(
+                override_service=override_service,
+                override_version=override_version,
+                override_dataset=override_dataset,
+            ),
+            x_api_key=self._api_key,
+        )
+
     async def disk_snapshot(
         self,
         override_service: str | None = None,
@@ -764,10 +832,10 @@ class Session:
         """
         self._check_closed()
 
-        try:
-            from playwright.async_api import Page
-        except ImportError as e:
-            raise ImportError("The login() method requires playwright. Install it with: pip install playwright") from e
+        import importlib.util
+
+        if importlib.util.find_spec("playwright") is None:
+            raise ImportError("The login() method requires playwright. Install it with: pip install playwright")
 
         context = await browser.new_context()
         pages: dict[str, Page] = {}

plato/v2/sync/environment.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING
 
 from plato._generated.api.v2 import jobs
 from plato._generated.models import (
+    AppApiV2SchemasSessionCreateSnapshotRequest,
     CreateCheckpointRequest,
     CreateSnapshotResult,
     ExecuteCommandRequest,
@@ -135,6 +136,36 @@ class Environment:
             x_api_key=self._api_key,
         )
 
+    def snapshot_store(
+        self,
+        override_service: str | None = None,
+        override_version: str | None = None,
+        override_dataset: str | None = None,
+    ) -> CreateSnapshotResult:
+        """Create a snapshot-store snapshot of this environment.
+
+        Uses the snapshot-store pipeline for chunk-based deduplication and
+        efficient storage. This is the preferred method for new base snapshots.
+
+        Args:
+            override_service: Override simulator/service name in artifact metadata.
+            override_version: Override version/git_hash in artifact metadata.
+            override_dataset: Override dataset name in artifact metadata.
+
+        Returns:
+            CreateSnapshotResult with artifact_id.
+        """
+        return jobs.snapshot_store.sync(
+            client=self._http,
+            job_id=self.job_id,
+            body=AppApiV2SchemasSessionCreateSnapshotRequest(
+                override_service=override_service,
+                override_version=override_version,
+                override_dataset=override_dataset,
+            ),
+            x_api_key=self._api_key,
+        )
+
     def close(self) -> None:
         """Close this environment."""
         jobs.close.sync(