afnio.tellurio.run

class RunStatus(Enum):
    """
    Represents the status of a Tellurio [`Run`][..Run].

    It can be one of the following values:

    - `"RUNNING"`: The Run is active and accepting logs, metrics, and artifacts.
    - `"COMPLETED"`: The run has completed successfully without any errors.
    - `"CRASHED"`: The run has crashed due to an unhandled exception or error.
    - `"FAILED"`: A safeguard/intermediate failure state used when a Run
        was left unfinished or superseded.
    """

    RUNNING = "RUNNING"
    CRASHED = "CRASHED"
    COMPLETED = "COMPLETED"
    FAILED = "FAILED"

class RunOrg:
    """Represents a Tellurio [`Run`][..Run]'s Organization.

    Each [`Run`][..Run] belongs to a [`Project`][afnio.tellurio.project.Project],
    and each [`Project`][afnio.tellurio.project.Project] belongs to a namespace,
    which can be either an Organization or a User (personal Organization).
    """

    def __init__(self, slug: str):
        self.slug = slug

    def __repr__(self):
        return f"<Organization slug={self.slug}>"

class RunProject:
    """
    Represents a Tellurio [`Run`][..Run]'s Project.

    Each [`Run`][..Run] belongs to a [`Project`][afnio.tellurio.project.Project],
    and each [`Project`][afnio.tellurio.project.Project] belongs to a namespace,
    which can be either an Organization or a User (personal Organization).
    """

    def __init__(self, uuid: str, display_name: str, slug: str):
        self.uuid = uuid
        self.display_name = display_name
        self.slug = slug

    def __repr__(self):
        return f"<Project uuid={self.uuid} display_name={self.display_name}>"

class RunUser:
    """
    Represents a Tellurio [`Run`][..Run]'s User.

    Each [`Run`][..Run] is associated with a user who created it.
    """

    def __init__(self, uuid: str, username: str, slug: str):
        self.uuid = uuid
        self.username = username
        self.slug = slug

    def __repr__(self):
        return f"<User uuid={self.uuid} username={self.username}>"

class Run:
    """
    Represents a Tellurio `Run`.

    Each [`Run`][.] belongs to a [`Project`][afnio.tellurio.project.Project],
    and each [`Project`][afnio.tellurio.project.Project] belongs to a namespace,
    which can be either an Organization or a User (personal Organization).

    A [`Run`][.] represents a single execution or experiment of an AI agent or
    workflow, that you want to track and log metrics for. It provides methods to log
    metrics and artifacts, and to mark the `Run` as completed or crashed. It can also
    be used as a context manager to automatically handle completion
    and crashing based on exceptions.
    """

    def __init__(
        self,
        uuid: str,
        name: str,
        description: str,
        status: RunStatus,
        date_created: Optional[datetime] = None,
        date_updated: Optional[datetime] = None,
        organization: Optional[RunOrg] = None,
        project: Optional[RunProject] = None,
        user: Optional[RunUser] = None,
    ):
        self.uuid = uuid
        self.name = name
        self.description = description
        self.status = RunStatus(status)
        self.date_created = date_created
        self.date_updated = date_updated
        self.organization = organization
        self.project = project
        self.user = user

    def __repr__(self):
        return (
            f"<Run uuid={self.uuid} name={self.name} status={self.status} "
            f"project={self.project.display_name if self.project else None}>"
        )

    def finish(
        self,
        client: Optional[TellurioClient] = None,
        status: Optional[RunStatus] = RunStatus.COMPLETED,
    ):
        """Marks the run as `"COMPLETED"` on the server by sending a PATCH request,
        and clears the active run UUID.

        Args:
            client: The client to use for the request. If not provided,
                the default client will be used.
            status: The [`RunStatus`][...RunStatus] to set for the run.

        Raises:
            Exception: If the PATCH request fails.
        """
        client = client or get_default_clients()[0]

        namespace_slug = self.organization.slug if self.organization else None
        project_slug = self.project.slug if self.project else None
        run_uuid = self.uuid

        if not (namespace_slug and project_slug and run_uuid):
            raise ValueError("Run object is missing required identifiers.")

        endpoint = f"/api/v0/{namespace_slug}/projects/{project_slug}/runs/{run_uuid}/"
        payload = {"status": status.value}

        try:
            response = client.patch(endpoint, json=payload)
            if response.status_code == 200:
                self.status = status
                if not _IN_ATEXIT:
                    logger.info(f"Run {self.name!r} marked as COMPLETED.")
            else:
                if not _IN_ATEXIT:
                    logger.error(
                        f"Failed to update run status: {response.status_code} - {response.text}"  # noqa: E501
                    )
                response.raise_for_status()
        except Exception as e:
            if not _IN_ATEXIT:
                logger.error(f"An error occurred while updating the run status: {e}")
            raise

        # Clear the active run UUID after finishing
        try:
            set_active_run(None)
        except Exception:
            pass

        # Mark safeguard as finished
        _unregister_safeguard(self)

    def log(
        self,
        name: str,
        value: Any,
        step: Optional[int] = None,
        client: Optional[TellurioClient] = None,
    ):
        """Log a metric for this run.

        Args:
            name: Name of the metric.
            value: Value of the metric. Can be any type that is JSON serializable.
            step: Step number. If not provided, the backend will auto-compute it.
            client: The client to use for the request. If not provided,
                the default client will be used.
        """
        client = client or get_default_clients()[0]

        namespace_slug = self.organization.slug if self.organization else None
        project_slug = self.project.slug if self.project else None
        run_uuid = self.uuid

        if not (namespace_slug and project_slug and run_uuid):
            raise ValueError("Run object is missing required identifiers.")

        endpoint = (
            f"/api/v0/{namespace_slug}/projects/{project_slug}/runs/{run_uuid}/metrics/"
        )
        payload = {
            "name": name,
            "value": value,
        }
        if step is not None:
            payload["step"] = step

        try:
            response = client.post(endpoint, json=payload)
            if response.status_code == 201:
                logger.info(f"Logged metric '{name}'={value} for run '{self.name}'.")
            else:
                logger.error(
                    f"Failed to log metric: {response.status_code} - {response.text}"
                )
                response.raise_for_status()
        except Exception as e:
            logger.error(f"An error occurred while logging the metric: {e}")
            raise

    def __enter__(self):
        """Enter the runtime context related to this object.

        Returns self so it can be used as a context manager.
        """
        _register_safeguard(self)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Exit the runtime context and finish the run.

        If an exception occurred, you may want to set status to `"CRASHED"`
        in the future.
        """
        status = RunStatus.CRASHED if exc_type else RunStatus.COMPLETED
        self.finish(status=status)