afnio.tellurio

def configure_logging(verbosity: str = "info"):
    """
    Configure logging for the `afnio` library.

    Sets up logging format and levels for CLI, scripts, and Jupyter notebooks.
    In a notebook, adds a custom handler for INFO-level logs to display
    user-facing messages with color and formatting.

    Args:
        verbosity: Logging level as a string (`"info"`, `"debug"`, etc.).
    """
    if not isinstance(verbosity, str):
        raise TypeError("verbosity must be a string like 'info', 'debug', etc.")
    level = getattr(logging, verbosity.upper(), logging.INFO)
    if level <= logging.DEBUG:
        fmt = "%(asctime)s - %(name)-35s - %(levelname)-9s - %(message)s"
    else:
        fmt = "%(levelname)-9s: %(message)s"

    # Set root logger to WARNING (default for all libraries)
    logging.basicConfig(
        level=logging.WARNING,
        format=fmt,
        force=True,
    )

    # Set only afnio logs to the desired level
    logging.getLogger("afnio").setLevel(level)

    # Add notebook handler for INFO logs if in notebook
    if _in_notebook() and level == logging.INFO:
        handler = _NotebookInfoHandler()
        handler.setLevel(logging.INFO)
        handler.addFilter(lambda record: record.levelno == logging.INFO)
        logging.getLogger().addHandler(handler)

        class NoInfoFilter(logging.Filter):
            """
            Logging filter that suppresses INFO-level log records.
            Used to prevent duplicate INFO messages in notebook environments.
            """

            def filter(self, record):
                """
                Filter method to suppress INFO-level log records.

                Args:
                    record (logging.LogRecord): The log record to filter.

                Returns:
                    bool: True if the record should be logged, False otherwise.
                """
                return record.levelno != logging.INFO

        for h in logging.getLogger().handlers:
            if not isinstance(h, _NotebookInfoHandler):
                h.addFilter(NoInfoFilter())

@contextmanager
def suppress_variable_notifications():
    """
    Context manager to temporarily suppress variable change notifications.

    When this context manager is active, any attribute changes to afnio.Variable
    instances will not trigger `_on_variable_change` notifications. This is useful
    for internal/client-initiated updates where you do not want to broadcast changes
    back to the server.
    """
    global _SUPPRESS_NOTIFICATIONS
    token = _SUPPRESS_NOTIFICATIONS
    _SUPPRESS_NOTIFICATIONS = True
    try:
        yield
    finally:
        _SUPPRESS_NOTIFICATIONS = token

def init(
    namespace_slug: str,
    project_display_name: str,
    name: Optional[str] = None,
    description: Optional[str] = None,
    status: Optional[RunStatus] = RunStatus.RUNNING,
    client: Optional[TellurioClient] = None,
) -> Run:
    """Initializes a new Tellurio [`Run`][afnio.tellurio.run.Run].

    Args:
        namespace_slug: The namespace slug where the project resides. It can be
          either an organization slug or a user slug.
        project_display_name: The display name of the project. This will be used
            to retrive or create the project through its slugified version.
        name: The name of the run. If not provided, a random name is
            generated (e.g., "brave_pasta_123"). If the name is provided but already
            exists, an incremental number is appended to the name (e.g., "test_run_1",
            "test_run_2").
        description: A description of the run.
        status: The status of the run (default: "RUNNING").
        client: An instance of TellurioClient. If not provided,
            the default client will be used.

    Returns:
        A Run object representing the created run.

    Raises:
        Exception: If there is an error during the API requests to create the run or
            retrieve/create the [`Project`][afnio.tellurio.project.Project].
    """
    client = client or get_default_clients()[0]

    # Generate the project's slug from its name
    project_slug = slugify(project_display_name)

    # Ensure the project exists
    try:
        project_obj = get_project(
            namespace_slug=namespace_slug,
            project_slug=project_slug,
            client=client,
        )
        if project_obj is not None:
            logger.info(
                f"Project with slug {project_slug!r} already exists "
                f"in namespace {namespace_slug!r}."
            )
        else:
            logger.info(
                f"Project with slug {project_slug!r} does not exist "
                f"in namespace {namespace_slug!r}. "
                f"Creating it now with RESTRICTED visibility."
            )
            project_obj = create_project(
                namespace_slug=namespace_slug,
                display_name=project_display_name,
                visibility="RESTRICTED",
                client=client,
            )
    except Exception as e:
        logger.error(f"An error occurred while retrieving or creating the project: {e}")
        raise

    # Dynamically construct the payload to exclude None values
    payload = {}
    if name is not None:
        payload["name"] = name
    if description is not None:
        payload["description"] = description
    if status is not None:
        payload["status"] = status.value

    # Create the run
    endpoint = f"/api/v0/{namespace_slug}/projects/{project_slug}/runs/"

    try:
        response = client.post(endpoint, json=payload)

        if response.status_code == 201:
            data = response.json()
            base_url = os.getenv(
                "TELLURIO_BACKEND_HTTP_BASE_URL", "https://platform.tellurio.ai"
            )
            run_slug = slugify(data["name"])
            logger.info(
                f"Run {data['name']!r} created successfully at: "
                f"{base_url}/{namespace_slug}/projects/{project_slug}/runs/{run_slug}/"
            )

            # Parse date fields
            date_created = datetime.fromisoformat(
                data["date_created"].replace("Z", "+00:00")
            )
            date_updated = datetime.fromisoformat(
                data["date_updated"].replace("Z", "+00:00")
            )

            # Parse project and user fields
            org_obj = RunOrg(
                slug=namespace_slug,
            )
            project_obj = RunProject(
                uuid=data["project"]["uuid"],
                display_name=data["project"]["display_name"],
                slug=data["project"]["slug"],
            )
            user_obj = RunUser(
                uuid=data["user"]["uuid"],
                username=data["user"]["username"],
                slug=data["user"]["slug"],
            )

            # Create and return the Run object
            run = Run(
                uuid=data["uuid"],
                name=data["name"],
                description=data["description"],
                status=RunStatus(data["status"]),
                date_created=date_created,
                date_updated=date_updated,
                organization=org_obj,
                project=project_obj,
                user=user_obj,
            )
            set_active_run(run)
            _register_safeguard(run)
            return run
        else:
            logger.error(
                f"Failed to create run: {response.status_code} - {response.text}"
            )
            response.raise_for_status()
    except Exception as e:
        logger.error(f"An error occurred while creating the run: {e}")
        raise

def login(api_key: str = None, relogin=False) -> dict:
    """Logs in the user and establishes both HTTP and WebSocket connections to the
    [Tellurio Studio](https://platform.tellurio.ai/) backend.

    This function authenticates the user using an API key, which can be provided
    directly, via the `TELLURIO_API_KEY` environment variable, or retrieved from the
    local keyring. On successful authentication, it establishes a WebSocket connection
    for real-time communication and stores the API key securely for future use.

    Args:
        api_key: The user's API key. If not provided, the function will
            attempt to use the `TELLURIO_API_KEY` environment variable. If that is not
            set, it will look for a stored API key in the local keyring.
        relogin: If `True`, forces a re-login and requires the user to provide a
            new API key (either directly or via the environment variable).

    Returns:
        A dictionary containing the user's email, username, and session ID for the \
        WebSocket connection. Example:

            `{
                "email": "user@example.com",
                "username": "user123",
                "session_id": "abc123xyz"
            }`

    Raises:
        ValueError: If the API key is not provided during first login or re-login.
        afnio.tellurio.client.InvalidAPIKeyError: If the backend rejects the API key.
        RuntimeError: If the WebSocket connection fails.
        Exception: For any other unexpected errors during login.

    Notes:
        - On first login, the API key is stored securely in the local keyring for
            future use.
        - If relogin is True, a new API key must be provided (directly or via
            environment variable).
        - This function is synchronous and can be called from both scripts and
            interactive environments.
    """
    # Get the default HTTP and WebSocket clients
    # Login in a separate step below to pass parameters
    client, ws_client = get_default_clients(login=False)

    return run_in_background_loop(
        _login(client=client, ws_client=ws_client, api_key=api_key, relogin=relogin)
    )  # Handle both sync and async contexts

def log(
    name: str,
    value: Any,
    step: Optional[int] = None,
    client: Optional[TellurioClient] = None,
):
    """Log a metric to the active [`Run`][afnio.tellurio.run.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.
    """
    run = get_active_run()
    run.log(name=name, value=value, step=step, client=client)