dreadnode 1.12.2__tar.gz → 1.13.1__tar.gz

{dreadnode-1.12.2 → dreadnode-1.13.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: dreadnode
-Version: 1.12.2
+Version: 1.13.1
 Summary: Dreadnode SDK
 Author: Nick Landers
 Author-email: monoxgas@gmail.com
@@ -14,6 +14,7 @@ Provides-Extra: all
 Provides-Extra: multimodal
 Provides-Extra: training
 Requires-Dist: coolname (>=2.2.0,<3.0.0)
+Requires-Dist: cyclopts (>=3.22.2,<4.0.0)
 Requires-Dist: fsspec[s3] (>=2023.1.0,<=2025.3.0)
 Requires-Dist: httpx (>=0.28.0,<0.29.0)
 Requires-Dist: logfire (>=3.5.3,<=3.20.0)

{dreadnode-1.12.2 → dreadnode-1.13.1}/dreadnode/__init__.py

@@ -1,9 +1,10 @@
 from dreadnode import convert, data_types, scorers
 from dreadnode.data_types import Audio, Code, Image, Markdown, Object3D, Table, Text, Video
+from dreadnode.lookup import Lookup, lookup_input, lookup_output, lookup_param, resolve_lookup
 from dreadnode.main import DEFAULT_INSTANCE, Dreadnode
 from dreadnode.metric import Metric, MetricDict, Scorer
 from dreadnode.object import Object
-from dreadnode.task import Task, TaskInput
+from dreadnode.task import Task
 from dreadnode.tracing.span import RunSpan, Span, TaskSpan
 from dreadnode.version import VERSION
 
@@ -39,6 +40,7 @@ __all__ = [
     "Code",
     "Dreadnode",
     "Image",
+    "Lookup",
     "Markdown",
     "Metric",
     "MetricDict",
@@ -50,7 +52,6 @@ __all__ = [
     "Span",
     "Table",
     "Task",
-    "TaskInput",
     "TaskSpan",
     "Text",
     "Video",
@@ -69,7 +70,11 @@ __all__ = [
     "log_output",
     "log_param",
     "log_params",
+    "lookup_input",
+    "lookup_output",
+    "lookup_param",
     "push_update",
+    "resolve_lookup",
     "run",
     "scorer",
     "scorers",

dreadnode-1.13.1/dreadnode/__main__.py

@@ -0,0 +1,10 @@
+from dreadnode.cli import cli
+
+
+def run() -> None:
+    """Run the Dreadnode CLI."""
+    cli.meta()
+
+
+if __name__ == "__main__":
+    run()

{dreadnode-1.12.2 → dreadnode-1.13.1}/dreadnode/api/client.py

@@ -1,22 +1,19 @@
 import io
 import json
+import time
 import typing as t
+from datetime import datetime, timezone
+from urllib.parse import urlparse
 
 import httpx
 import pandas as pd
 from pydantic import BaseModel
 from ulid import ULID
 
-from dreadnode.api.util import (
-    convert_flat_tasks_to_tree,
-    convert_flat_trace_to_tree,
-    process_run,
-    process_task,
-)
-from dreadnode.util import logger
-from dreadnode.version import VERSION
-
-from .models import (
+from dreadnode.api.models import (
+    AccessRefreshTokenResponse,
+    DeviceCodeResponse,
+    GithubTokenResponse,
     MetricAggregationType,
     Project,
     RawRun,
@@ -31,7 +28,21 @@ from .models import (
     TraceSpan,
     TraceTree,
     UserDataCredentials,
+    UserResponse,
 )
+from dreadnode.api.util import (
+    convert_flat_tasks_to_tree,
+    convert_flat_trace_to_tree,
+    process_run,
+    process_task,
+)
+from dreadnode.constants import (
+    DEFAULT_FS_CREDENTIAL_DURATION,
+    DEFAULT_MAX_POLL_TIME,
+    DEFAULT_POLL_INTERVAL,
+)
+from dreadnode.util import logger
+from dreadnode.version import VERSION
 
 ModelT = t.TypeVar("ModelT", bound=BaseModel)
 
@@ -47,11 +58,13 @@ class ApiClient:
     def __init__(
         self,
         base_url: str,
-        api_key: str,
         *,
+        api_key: str | None = None,
+        cookies: dict[str, str] | None = None,
         debug: bool = False,
     ):
-        """Initializes the API client.
+        """
+        Initializes the API client.
 
         Args:
             base_url (str): The base URL of the Dreadnode API.
@@ -62,12 +75,28 @@ class ApiClient:
         if not self._base_url.endswith("/api"):
             self._base_url += "/api"
 
+        _cookies = httpx.Cookies()
+        cookie_domain = urlparse(base_url).hostname
+        if cookie_domain is None:
+            raise ValueError(f"Invalid URL: {base_url}")
+
+        if cookie_domain == "localhost":
+            cookie_domain = "localhost.local"
+
+        for key, value in (cookies or {}).items():
+            _cookies.set(key, value, domain=cookie_domain)
+
+        headers = {
+            "User-Agent": f"dreadnode-sdk/{VERSION}",
+            "Accept": "application/json",
+        }
+
+        if api_key:
+            headers["X-Api-Key"] = api_key
+
         self._client = httpx.Client(
-            headers={
-                "User-Agent": f"dreadnode-sdk/{VERSION}",
-                "Accept": "application/json",
-                "X-API-Key": api_key,
-            },
+            headers=headers,
+            cookies=_cookies,
             base_url=self._base_url,
             timeout=30,
         )
@@ -77,7 +106,8 @@ class ApiClient:
             self._client.event_hooks["response"].append(self._log_response)
 
     def _log_request(self, request: httpx.Request) -> None:
-        """Logs HTTP requests if debug mode is enabled.
+        """
+        Logs HTTP requests if debug mode is enabled.
 
         Args:
             request (httpx.Request): The HTTP request object.
@@ -90,7 +120,8 @@ class ApiClient:
         logger.debug("-------------------------------------------")
 
     def _log_response(self, response: httpx.Response) -> None:
-        """Logs HTTP responses if debug mode is enabled.
+        """
+        Logs HTTP responses if debug mode is enabled.
 
         Args:
             response (httpx.Response): The HTTP response object.
@@ -103,7 +134,8 @@ class ApiClient:
         logger.debug("--------------------------------------------")
 
     def _get_error_message(self, response: httpx.Response) -> str:
-        """Extracts the error message from an HTTP response.
+        """
+        Extracts the error message from an HTTP response.
 
         Args:
             response (httpx.Response): The HTTP response object.
@@ -125,7 +157,8 @@ class ApiClient:
         params: dict[str, t.Any] | None = None,
         json_data: dict[str, t.Any] | None = None,
     ) -> httpx.Response:
-        """Makes a raw HTTP request to the API.
+        """
+        Makes a raw HTTP request to the API.
 
         Args:
             method (str): The HTTP method (e.g., "GET", "POST").
@@ -146,7 +179,8 @@ class ApiClient:
         params: dict[str, t.Any] | None = None,
         json_data: dict[str, t.Any] | None = None,
     ) -> httpx.Response:
-        """Makes an HTTP request to the API and raises exceptions for errors.
+        """
+        Makes an HTTP request to the API and raises exceptions for errors.
 
         Args:
             method (str): The HTTP method (e.g., "GET", "POST").
@@ -170,6 +204,59 @@ class ApiClient:
 
         return response
 
+    # Auth
+
+    def url_for_user_code(self, user_code: str) -> str:
+        """Get the URL to verify the user code."""
+
+        return f"{self._base_url.removesuffix('/api')}/account/device?code={user_code}"
+
+    def get_device_codes(self) -> DeviceCodeResponse:
+        """Start the authentication flow by requesting user and device codes."""
+
+        response = self.request("POST", "/auth/device/code")
+        return DeviceCodeResponse(**response.json())
+
+    def poll_for_token(
+        self,
+        device_code: str,
+        interval: int = DEFAULT_POLL_INTERVAL,
+        max_poll_time: int = DEFAULT_MAX_POLL_TIME,
+    ) -> AccessRefreshTokenResponse:
+        """Poll for the access token with the given device code."""
+
+        start_time = datetime.now(timezone.utc)
+        while (datetime.now(timezone.utc) - start_time).total_seconds() < max_poll_time:
+            response = self._request(
+                "POST", "/auth/device/token", json_data={"device_code": device_code}
+            )
+
+            if response.status_code == 200:  # noqa: PLR2004
+                return AccessRefreshTokenResponse(**response.json())
+            if response.status_code != 401:  # noqa: PLR2004
+                raise RuntimeError(self._get_error_message(response))
+
+            time.sleep(interval)
+
+        raise RuntimeError("Polling for token timed out")
+
+    # User
+
+    def get_user(self) -> UserResponse:
+        """Get the user email and username."""
+
+        response = self.request("GET", "/user")
+        return UserResponse(**response.json())
+
+    # Github
+
+    def get_github_access_token(self, repos: list[str]) -> GithubTokenResponse:
+        """Try to get a GitHub access token for the given repositories."""
+        response = self.request("POST", "/github/token", json_data={"repos": repos})
+        return GithubTokenResponse(**response.json())
+
+    # Strikes
+
     def list_projects(self) -> list[Project]:
         """Retrieves a list of projects.
 
@@ -294,7 +381,8 @@ class ApiClient:
         status: StatusFilter = "completed",
         aggregations: list[MetricAggregationType] | None = None,
     ) -> pd.DataFrame:
-        """Exports run data for a specific project.
+        """
+        Exports run data for a specific project.
 
         Args:
             project: The project identifier.
@@ -327,7 +415,8 @@ class ApiClient:
         metrics: list[str] | None = None,
         aggregations: list[MetricAggregationType] | None = None,
     ) -> pd.DataFrame:
-        """Exports metric data for a specific project.
+        """
+        Exports metric data for a specific project.
 
         Args:
             project: The project identifier.
@@ -363,7 +452,8 @@ class ApiClient:
         metrics: list[str] | None = None,
         aggregations: list[MetricAggregationType] | None = None,
     ) -> pd.DataFrame:
-        """Exports parameter data for a specific project.
+        """
+        Exports parameter data for a specific project.
 
         Args:
             project: The project identifier.
@@ -401,7 +491,8 @@ class ApiClient:
         time_axis: TimeAxisType = "relative",
         aggregations: list[TimeAggregationType] | None = None,
     ) -> pd.DataFrame:
-        """Exports timeseries data for a specific project.
+        """
+        Exports timeseries data for a specific project.
 
         Args:
             project: The project identifier.
@@ -430,12 +521,17 @@ class ApiClient:
 
     # User data access
 
-    def get_user_data_credentials(self) -> UserDataCredentials:
+    def get_user_data_credentials(
+        self, duration: int = DEFAULT_FS_CREDENTIAL_DURATION
+    ) -> UserDataCredentials:
         """
         Retrieves user data credentials for secondary storage access.
 
+        Args:
+            duration: Credential lifetime in seconds (default: 4 hours)
+
         Returns:
             The user data credentials object.
         """
-        response = self.request("GET", "/user-data/credentials")
+        response = self._request("GET", "/user-data/credentials", params={"duration": duration})
         return UserDataCredentials(**response.json())

{dreadnode-1.12.2 → dreadnode-1.13.1}/dreadnode/api/models.py

@@ -32,6 +32,35 @@ class UserResponse(BaseModel):
     api_key: UserAPIKey
 
 
+class UserDataCredentials(BaseModel):
+    access_key_id: str
+    secret_access_key: str
+    session_token: str
+    expiration: datetime
+    region: str
+    bucket: str
+    prefix: str
+    endpoint: str | None
+
+
+# Auth
+
+
+class DeviceCodeResponse(BaseModel):
+    id: UUID
+    completed: bool
+    device_code: str
+    expires_at: datetime
+    expires_in: int
+    user_code: str
+    verification_url: str
+
+
+class AccessRefreshTokenResponse(BaseModel):
+    access_token: str
+    refresh_token: str
+
+
 # Strikes
 
 SpanStatus = t.Literal[
@@ -406,15 +435,10 @@ class TraceTree(BaseModel):
     """Children of this span, representing nested spans or tasks."""
 
 
-# User data credentials
+# Github
 
 
-class UserDataCredentials(BaseModel):
-    access_key_id: str
-    secret_access_key: str
-    session_token: str
-    expiration: datetime
-    region: str
-    bucket: str
-    prefix: str
-    endpoint: str | None
+class GithubTokenResponse(BaseModel):
+    token: str
+    expires_at: datetime
+    repos: list[str]

{dreadnode-1.12.2 → dreadnode-1.13.1}/dreadnode/api/util.py

@@ -1,6 +1,6 @@
 from logging import getLogger
 
-from .models import (
+from dreadnode.api.models import (
     Object,
     ObjectUri,
     ObjectVal,

{dreadnode-1.12.2 → dreadnode-1.13.1}/dreadnode/artifact/storage.py

@@ -4,10 +4,12 @@ Provides efficient uploading of files and directories with deduplication.
 """
 
 import hashlib
+import typing as t
 from pathlib import Path
 
 import fsspec  # type: ignore[import-untyped]
 
+from dreadnode.storage_utils import with_credential_refresh
 from dreadnode.util import logger
 
 CHUNK_SIZE = 8 * 1024 * 1024  # 8MB
@@ -22,15 +24,27 @@ class ArtifactStorage:
     - Batch uploads for directories handled by fsspec
     """
 
-    def __init__(self, file_system: fsspec.AbstractFileSystem):
+    def __init__(
+        self,
+        file_system: fsspec.AbstractFileSystem,
+        credential_refresher: t.Callable[[], bool] | None = None,
+    ):
         """
         Initialize artifact storage with a file system and prefix path.
 
         Args:
             file_system: FSSpec-compatible file system
+            credential_refresher: Optional function to refresh credentials when it's about to expire
         """
         self._file_system = file_system
+        self._credential_refresher = credential_refresher
 
+    def _refresh_credentials_if_needed(self) -> None:
+        """Refresh credentials if refresher is available."""
+        if self._credential_refresher:
+            self._credential_refresher()
+
+    @with_credential_refresh
     def store_file(self, file_path: Path, target_key: str) -> str:
         """
         Store a file in the storage system, using multipart upload for large files.

dreadnode-1.13.1/dreadnode/cli/__init__.py

@@ -0,0 +1,3 @@
+from dreadnode.cli.main import cli
+
+__all__ = ["cli"]

dreadnode-1.13.1/dreadnode/cli/api.py

@@ -0,0 +1,79 @@
+import atexit
+import base64
+import json
+from datetime import datetime, timezone
+
+from dreadnode.api.client import ApiClient
+from dreadnode.config import UserConfig
+from dreadnode.constants import (
+    DEFAULT_TOKEN_MAX_TTL,
+)
+
+
+class Token:
+    """A JWT token with an expiration time."""
+
+    data: str
+    expires_at: datetime
+
+    @staticmethod
+    def parse_jwt_token_expiration(token: str) -> datetime:
+        """Return the expiration date from a JWT token."""
+
+        _, b64payload, _ = token.split(".")
+        payload = base64.urlsafe_b64decode(b64payload + "==").decode("utf-8")
+        return datetime.fromtimestamp(json.loads(payload).get("exp"), tz=timezone.utc)
+
+    def __init__(self, token: str):
+        self.data = token
+        self.expires_at = Token.parse_jwt_token_expiration(token)
+
+    def ttl(self) -> int:
+        """Get number of seconds left until the token expires."""
+        return int((self.expires_at - datetime.now(tz=timezone.utc)).total_seconds())
+
+    def is_expired(self) -> bool:
+        """Return True if the token is expired."""
+        return self.ttl() <= 0
+
+    def is_close_to_expiry(self) -> bool:
+        """Return True if the token is close to expiry."""
+        return self.ttl() <= DEFAULT_TOKEN_MAX_TTL
+
+
+def create_api_client(*, profile: str | None = None) -> ApiClient:
+    """Create an authenticated API client using stored configuration data."""
+
+    user_config = UserConfig.read()
+    config = user_config.get_server_config(profile)
+
+    client = ApiClient(
+        config.url,
+        cookies={"access_token": config.access_token, "refresh_token": config.refresh_token},
+    )
+
+    # Preemptively check if the token is expired
+    if Token(config.refresh_token).is_expired():
+        raise RuntimeError("Authentication expired, use [bold]dreadnode login[/]")
+
+    def _flush_auth_changes() -> None:
+        """Flush the authentication data to disk if it has been updated."""
+
+        access_token = client._client.cookies.get("access_token")  # noqa: SLF001
+        refresh_token = client._client.cookies.get("refresh_token")  # noqa: SLF001
+
+        changed: bool = False
+        if access_token and access_token != config.access_token:
+            changed = True
+            config.access_token = access_token
+
+        if refresh_token and refresh_token != config.refresh_token:
+            changed = True
+            config.refresh_token = refresh_token
+
+        if changed:
+            user_config.set_server_config(config, profile).write()
+
+    atexit.register(_flush_auth_changes)
+
+    return client