dreadnode 1.0.0rc0__py3-none-any.whl

dreadnode/__init__.py

@@ -0,0 +1,51 @@
+from dreadnode.main import DEFAULT_INSTANCE, Dreadnode
+from dreadnode.metric import Metric, MetricDict, Scorer
+from dreadnode.object import Object
+from dreadnode.task import Task
+from dreadnode.tracing.span import RunSpan, Span, TaskSpan
+from dreadnode.version import VERSION
+
+configure = DEFAULT_INSTANCE.configure
+shutdown = DEFAULT_INSTANCE.shutdown
+
+api = DEFAULT_INSTANCE.api
+span = DEFAULT_INSTANCE.span
+task = DEFAULT_INSTANCE.task
+task_span = DEFAULT_INSTANCE.task_span
+run = DEFAULT_INSTANCE.run
+scorer = DEFAULT_INSTANCE.scorer
+task_span = DEFAULT_INSTANCE.task_span
+push_update = DEFAULT_INSTANCE.push_update
+
+log_metric = DEFAULT_INSTANCE.log_metric
+log_param = DEFAULT_INSTANCE.log_param
+log_params = DEFAULT_INSTANCE.log_params
+log_input = DEFAULT_INSTANCE.log_input
+log_inputs = DEFAULT_INSTANCE.log_inputs
+log_output = DEFAULT_INSTANCE.log_output
+link_objects = DEFAULT_INSTANCE.link_objects
+log_artifact = DEFAULT_INSTANCE.log_artifact
+
+__version__ = VERSION
+
+__all__ = [
+    "Dreadnode",
+    "Metric",
+    "MetricDict",
+    "Object",
+    "Run",
+    "RunSpan",
+    "Score",
+    "Scorer",
+    "Span",
+    "Task",
+    "TaskSpan",
+    "__version__",
+    "configure",
+    "log_metric",
+    "log_param",
+    "run",
+    "shutdown",
+    "span",
+    "task",
+]

dreadnode/api/client.py

@@ -0,0 +1,249 @@
+import io
+import json
+import typing as t
+
+import httpx
+import pandas as pd
+from pydantic import BaseModel
+from ulid import ULID
+
+from dreadnode.util import logger
+from dreadnode.version import VERSION
+
+from .models import (
+    MetricAggregationType,
+    Project,
+    Run,
+    StatusFilter,
+    Task,
+    TimeAggregationType,
+    TimeAxisType,
+    TraceSpan,
+    UserDataCredentials,
+)
+
+ModelT = t.TypeVar("ModelT", bound=BaseModel)
+
+
+class ApiClient:
+    """Client for the Dreadnode API."""
+
+    def __init__(
+        self,
+        base_url: str,
+        api_key: str,
+        *,
+        debug: bool = False,
+    ):
+        self._base_url = base_url.rstrip("/")
+        if not self._base_url.endswith("/api"):
+            self._base_url += "/api"
+
+        self._client = httpx.Client(
+            headers={
+                "User-Agent": f"dreadnode-sdk/{VERSION}",
+                "Accept": "application/json",
+                "X-API-Key": api_key,
+            },
+            base_url=self._base_url,
+            timeout=30,
+        )
+
+        if debug:
+            self._client.event_hooks["request"].append(self._log_request)
+            self._client.event_hooks["response"].append(self._log_response)
+
+    def _log_request(self, request: httpx.Request) -> None:
+        """Log every request to the console if debug is enabled."""
+
+        logger.debug("-------------------------------------------")
+        logger.debug("%s %s", request.method, request.url)
+        logger.debug("Headers: %s", request.headers)
+        logger.debug("Content: %s", request.content)
+        logger.debug("-------------------------------------------")
+
+    def _log_response(self, response: httpx.Response) -> None:
+        """Log every response to the console if debug is enabled."""
+
+        logger.debug("-------------------------------------------")
+        logger.debug("Response: %s", response.status_code)
+        logger.debug("Headers: %s", response.headers)
+        logger.debug("Content: %s", response.read())
+        logger.debug("--------------------------------------------")
+
+    def _get_error_message(self, response: httpx.Response) -> str:
+        """Get the error message from the response."""
+
+        try:
+            obj = response.json()
+            return f"{response.status_code}: {obj.get('detail', json.dumps(obj))}"
+        except Exception:  # noqa: BLE001
+            return str(response.content)
+
+    def _request(
+        self,
+        method: str,
+        path: str,
+        params: dict[str, t.Any] | None = None,
+        json_data: dict[str, t.Any] | None = None,
+    ) -> httpx.Response:
+        """Make a raw request to the API."""
+
+        return self._client.request(method, path, json=json_data, params=params)
+
+    def request(
+        self,
+        method: str,
+        path: str,
+        params: dict[str, t.Any] | None = None,
+        json_data: dict[str, t.Any] | None = None,
+    ) -> httpx.Response:
+        """Make a request to the API. Raise an exception for non-200 status codes."""
+
+        response = self._request(method, path, params, json_data)
+        if response.status_code == 401:  # noqa: PLR2004
+            raise RuntimeError("Authentication failed, please check your API token.")
+
+        try:
+            response.raise_for_status()
+        except httpx.HTTPStatusError as e:
+            raise RuntimeError(self._get_error_message(response)) from e
+
+        return response
+
+    # This currently won't work with API keys
+    # def get_user(self) -> UserResponse:
+    #     response = self.request("GET", "/user")
+    #     return UserResponse(**response.json())
+
+    def list_projects(self) -> list[Project]:
+        response = self.request("GET", "/strikes/projects")
+        return [Project(**project) for project in response.json()]
+
+    def get_project(self, project: str) -> Project:
+        response = self.request("GET", f"/strikes/projects/{project!s}")
+        return Project(**response.json())
+
+    def list_runs(self, project: str) -> list[Run]:
+        response = self.request("GET", f"/strikes/projects/{project!s}/runs")
+        return [Run(**run) for run in response.json()]
+
+    def get_run(self, run: str | ULID) -> Run:
+        response = self.request("GET", f"/strikes/projects/runs/{run!s}")
+        return Run(**response.json())
+
+    def get_run_tasks(self, run: str | ULID) -> list[Task]:
+        response = self.request("GET", f"/strikes/projects/runs/{run!s}/tasks")
+        return [Task(**task) for task in response.json()]
+
+    def get_run_trace(self, run: str | ULID) -> list[Task | TraceSpan]:
+        response = self.request("GET", f"/strikes/projects/runs/{run!s}/spans")
+        spans: list[Task | TraceSpan] = []
+        for item in response.json():
+            if "parent_task_span_id" in item:
+                spans.append(Task(**item))
+            else:
+                spans.append(TraceSpan(**item))
+        return spans
+
+    # Data exports
+
+    def export_runs(
+        self,
+        project: str,
+        *,
+        filter: str | None = None,
+        # format: ExportFormat = "parquet",
+        status: StatusFilter = "completed",
+        aggregations: list[MetricAggregationType] | None = None,
+    ) -> pd.DataFrame:
+        response = self.request(
+            "GET",
+            f"/strikes/projects/{project!s}/export",
+            params={
+                "format": "parquet",
+                "status": status,
+                **({"filter": filter} if filter else {}),
+                **({"aggregations": aggregations} if aggregations else {}),
+            },
+        )
+        return pd.read_parquet(io.BytesIO(response.content))
+
+    def export_metrics(
+        self,
+        project: str,
+        *,
+        filter: str | None = None,
+        # format: ExportFormat = "parquet",
+        status: StatusFilter = "completed",
+        metrics: list[str] | None = None,
+        aggregations: list[MetricAggregationType] | None = None,
+    ) -> pd.DataFrame:
+        response = self.request(
+            "GET",
+            f"/strikes/projects/{project!s}/export/metrics",
+            params={
+                "format": "parquet",
+                "status": status,
+                "filter": filter,
+                **({"metrics": metrics} if metrics else {}),
+                **({"aggregations": aggregations} if aggregations else {}),
+            },
+        )
+        return pd.read_parquet(io.BytesIO(response.content))
+
+    def export_parameters(
+        self,
+        project: str,
+        *,
+        filter: str | None = None,
+        # format: ExportFormat = "parquet",
+        status: StatusFilter = "completed",
+        parameters: list[str] | None = None,
+        metrics: list[str] | None = None,
+        aggregations: list[MetricAggregationType] | None = None,
+    ) -> pd.DataFrame:
+        response = self.request(
+            "GET",
+            f"/strikes/projects/{project!s}/export/parameters",
+            params={
+                "format": "parquet",
+                "status": status,
+                "filter": filter,
+                **({"parameters": parameters} if parameters else {}),
+                **({"metrics": metrics} if metrics else {}),
+                **({"aggregations": aggregations} if aggregations else {}),
+            },
+        )
+        return pd.read_parquet(io.BytesIO(response.content))
+
+    def export_timeseries(
+        self,
+        project: str,
+        *,
+        filter: str | None = None,
+        # format: ExportFormat = "parquet",
+        status: StatusFilter = "completed",
+        metrics: list[str] | None = None,
+        time_axis: TimeAxisType = "relative",
+        aggregations: list[TimeAggregationType] | None = None,
+    ) -> pd.DataFrame:
+        response = self.request(
+            "GET",
+            f"/strikes/projects/{project!s}/export/timeseries",
+            params={
+                "format": "parquet",
+                "status": status,
+                "filter": filter,
+                "time_axis": time_axis,
+                **({"metrics": metrics} if metrics else {}),
+                **({"aggregation": aggregations} if aggregations else {}),
+            },
+        )
+        return pd.read_parquet(io.BytesIO(response.content))
+
+    # User data access
+
+    def get_user_data_credentials(self) -> UserDataCredentials:
+        response = self.request("GET", "/user-data/credentials")
+        return UserDataCredentials(**response.json())

dreadnode/api/models.py

@@ -0,0 +1,210 @@
+import typing as t
+from datetime import datetime
+from uuid import UUID
+
+from pydantic import BaseModel, Field
+from ulid import ULID
+
+AnyDict = dict[str, t.Any]
+
+# User
+
+
+class UserAPIKey(BaseModel):
+    key: str
+
+
+class UserResponse(BaseModel):
+    id: UUID
+    email_address: str
+    username: str
+    api_key: UserAPIKey
+
+
+# Strikes
+
+SpanStatus = t.Literal[
+    "pending",  # A pending span has been created
+    "completed",  # The span has been finished
+    "failed",  # The raised an exception
+]
+
+ExportFormat = t.Literal["csv", "json", "jsonl", "parquet"]
+StatusFilter = t.Literal["all", "completed", "failed"]
+TimeAxisType = t.Literal["wall", "relative", "step"]
+TimeAggregationType = t.Literal["max", "min", "sum", "count"]
+MetricAggregationType = t.Literal[
+    "avg",
+    "median",
+    "min",
+    "max",
+    "sum",
+    "first",
+    "last",
+    "count",
+    "std",
+    "var",
+]
+
+
+class SpanException(BaseModel):
+    type: str
+    message: str
+    stacktrace: str
+
+
+class SpanEvent(BaseModel):
+    timestamp: datetime
+    name: str
+    attributes: AnyDict
+
+
+class SpanLink(BaseModel):
+    trace_id: str
+    span_id: str
+    attributes: AnyDict
+
+
+class TraceLog(BaseModel):
+    timestamp: datetime
+    body: str
+    severity: str
+    service: str | None
+    trace_id: str | None
+    span_id: str | None
+    attributes: AnyDict
+    container: str | None
+
+
+class TraceSpan(BaseModel):
+    timestamp: datetime
+    duration: int
+    trace_id: str
+    span_id: str
+    parent_span_id: str | None
+    service_name: str | None
+    status: SpanStatus
+    exception: SpanException | None
+    name: str
+    attributes: AnyDict
+    resource_attributes: AnyDict
+    events: list[SpanEvent]
+    links: list[SpanLink]
+
+
+class Metric(BaseModel):
+    value: float
+    step: int
+    timestamp: datetime
+    attributes: AnyDict
+
+
+class ObjectRef(BaseModel):
+    name: str
+    label: str
+    hash: str
+
+
+class ObjectUri(BaseModel):
+    hash: str
+    schema_hash: str
+    uri: str
+    size: int
+    type: t.Literal["uri"]
+
+
+class ObjectVal(BaseModel):
+    hash: str
+    schema_hash: str
+    value: t.Any
+    type: t.Literal["val"]
+
+
+Object = ObjectUri | ObjectVal
+
+
+class V0Object(BaseModel):
+    name: str
+    label: str
+    value: t.Any
+
+
+class Run(BaseModel):
+    id: ULID
+    name: str
+    span_id: str
+    trace_id: str
+    timestamp: datetime
+    duration: int
+    status: SpanStatus
+    exception: SpanException | None
+    tags: set[str]
+    params: AnyDict
+    metrics: dict[str, list[Metric]]
+    inputs: list[ObjectRef]
+    outputs: list[ObjectRef]
+    objects: dict[str, Object]
+    object_schemas: AnyDict
+    schema_: AnyDict = Field(alias="schema")
+
+
+class Task(BaseModel):
+    name: str
+    span_id: str
+    trace_id: str
+    parent_span_id: str | None
+    parent_task_span_id: str | None
+    timestamp: datetime
+    duration: int
+    status: SpanStatus
+    exception: SpanException | None
+    tags: set[str]
+    params: AnyDict
+    metrics: dict[str, list[Metric]]
+    inputs: list[ObjectRef] | list[V0Object]  # v0 compat
+    outputs: list[ObjectRef] | list[V0Object]  # v0 compat
+    schema_: AnyDict = Field(alias="schema")
+    attributes: AnyDict
+    resource_attributes: AnyDict
+    events: list[SpanEvent]
+    links: list[SpanLink]
+
+
+class Project(BaseModel):
+    id: UUID
+    key: str
+    name: str
+    description: str | None
+    created_at: datetime
+    updated_at: datetime
+    run_count: int
+    last_run: Run | None
+
+
+# Derived types
+
+
+class TaskTree(BaseModel):
+    task: Task
+    children: list["TaskTree"] = []
+
+
+class SpanTree(BaseModel):
+    """Tree representation of a trace span with its children"""
+
+    span: Task | TraceSpan
+    children: list["SpanTree"] = []
+
+
+# User data credentials
+
+
+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