plato-sdk-v2 2.0.50__py3-none-any.whl → 2.2.4__py3-none-any.whl

plato/v1/cli/world.py

@@ -235,6 +235,8 @@ def world_publish(
     upload_url = f"{api_url}/v2/pypi/worlds/"
     console.print(f"\n[cyan]Uploading to {upload_url}...[/cyan]")
 
+    # api_key is guaranteed to be set (checked earlier when not dry_run)
+    assert api_key is not None, "api_key must be set when not in dry_run mode"
     try:
         result = subprocess.run(
             [
@@ -250,6 +252,7 @@ def world_publish(
             ],
             capture_output=True,
             text=True,
+            check=False,
         )
 
         if result.returncode == 0:

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,8 @@ from typing import TYPE_CHECKING
 
 from plato._generated.api.v2 import jobs
 from plato._generated.models import (
+    ConnectRoutingInfoResult,
+    CreateCheckpointRequest,
     CreateSnapshotResult,
     ExecuteCommandRequest,
     ExecuteCommandResult,
@@ -130,7 +132,7 @@ class Environment:
         return await jobs.snapshot.asyncio(
             client=self._http,
             job_id=self.job_id,
-            body=None,
+            body=CreateCheckpointRequest(),
             x_api_key=self._api_key,
         )
 
@@ -142,5 +144,22 @@ class Environment:
             x_api_key=self._api_key,
         )
 
+    async def get_connection_info(self) -> ConnectRoutingInfoResult:
+        """Get connection/routing info for this environment.
+
+        Returns connection details including:
+        - vm_gateway_ip: Gateway IP for VM to reach the host
+        - vm_private_ip: Private IP of the VM
+        - worker_private_ip: Private IP of the worker
+        - ready: Whether the job is ready
+
+        Returns:
+            ConnectRoutingInfoResult with routing information.
+        """
+        return await jobs.connect_routing_info.asyncio(
+            client=self._http,
+            job_id=self.job_id,
+        )
+
     def __repr__(self) -> str:
         return f"Environment(alias={self.alias!r}, job_id={self.job_id!r})"

plato/v2/async_/session.py

@@ -12,12 +12,15 @@ import uuid
 from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import httpx
 import tenacity
-from playwright.async_api import Browser, BrowserContext, Page
 from pydantic import BaseModel
 
+if TYPE_CHECKING:
+    from playwright.async_api import Browser, BrowserContext, Page
+
 from plato._generated.api.v2.jobs import get_flows as jobs_get_flows
 from plato._generated.api.v2.jobs import public_url as jobs_public_url
 from plato._generated.api.v2.sessions import close as sessions_close
@@ -29,6 +32,7 @@ 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 state as sessions_state
 from plato._generated.api.v2.sessions import wait_for_ready as sessions_wait_for_ready
@@ -37,6 +41,8 @@ from plato._generated.models import (
     AppApiV2SchemasSessionCreateSnapshotResponse,
     AppApiV2SchemasSessionEvaluateResponse,
     AppApiV2SchemasSessionHeartbeatResponse,
+    AppApiV2SchemasSessionSetupSandboxRequest,
+    AppApiV2SchemasSessionSetupSandboxResponse,
     CreateDiskSnapshotRequest,
     CreateDiskSnapshotResponse,
     CreateSessionFromEnvs,
@@ -67,7 +73,10 @@ logger = logging.getLogger(__name__)
 
 @dataclass
 class LoginResult:
-    """Result of login operation containing browser context and pages."""
+    """Result of login operation containing browser context and pages.
+
+    Requires playwright to be installed.
+    """
 
     context: BrowserContext
     pages: dict[str, Page]
@@ -291,9 +300,10 @@ class Session:
         )
 
         # Use generated API function
+        # Note: API supports both CreateSessionFromEnvs and CreateSessionFromTask via discriminator
         response = await sessions_make.asyncio(
             client=http_client,
-            body=request_body,
+            body=request_body,  # type: ignore[arg-type]
             x_api_key=api_key,
         )
 
@@ -550,6 +560,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.
 
@@ -738,6 +780,9 @@ class Session:
         Navigates each page to the environment's public URL and executes
         the login flow.
 
+        Requires playwright to be installed:
+            pip install playwright
+
         Args:
             browser: Playwright Browser instance.
             dataset: Dataset name for login flow (default: "base" uses "login" flow).
@@ -750,9 +795,15 @@ class Session:
 
         Raises:
             RuntimeError: If login fails.
+            ImportError: If playwright is not installed.
         """
         self._check_closed()
 
+        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 (
+    CreateCheckpointRequest,
     CreateSnapshotResult,
     ExecuteCommandRequest,
     ExecuteCommandResult,
@@ -130,7 +131,7 @@ class Environment:
         return jobs.snapshot.sync(
             client=self._http,
             job_id=self.job_id,
-            body=None,
+            body=CreateCheckpointRequest(),
             x_api_key=self._api_key,
         )