snowglobe 0.4.0__tar.gz → 0.4.2__tar.gz

{snowglobe-0.4.0/src/snowglobe.egg-info → snowglobe-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowglobe
-Version: 0.4.0
+Version: 0.4.2
 Summary: client server for usage with snowglobe experiments
 Author-email: Guardrails AI <contact@guardrailsai.com>
 License: MIT License
@@ -49,21 +49,14 @@ The Snowglobe Connect SDK helps you connect your AI agents to Snowglobe. It send
 
 ## Installation
 
-```bash
-# Paste Guardrails Client or extract from guardrailsrc
-export GUARDRAILS_TOKEN=$(cat ~/.guardrailsrc| awk -F 'token=' '{print $2}' | awk '{print $1}' | tr -d '\n')
-```
-
 ```
 # Install client
-pip install -U --index-url="https://__token__:$GUARDRAILS_TOKEN@pypi.guardrailsai.com/simple" \
-    --extra-index-url="https://pypi.org/simple" snowglobe-connect
+pip install snowglobe-connect
 ```
 
 If using uv, set the `--prerelease=allow` flag
 ```
-pip install -U --index-url="https://__token__:$GUARDRAILS_TOKEN@pypi.guardrailsai.com/simple" \
-    --extra-index-url="https://pypi.org/simple" --prerelease=allow snowglobe-connect
+uv pip install --prerelease=allow snowglobe-connect
 ```
 
 

{snowglobe-0.4.0 → snowglobe-0.4.2}/README.md

@@ -4,21 +4,14 @@ The Snowglobe Connect SDK helps you connect your AI agents to Snowglobe. It send
 
 ## Installation
 
-```bash
-# Paste Guardrails Client or extract from guardrailsrc
-export GUARDRAILS_TOKEN=$(cat ~/.guardrailsrc| awk -F 'token=' '{print $2}' | awk '{print $1}' | tr -d '\n')
-```
-
 ```
 # Install client
-pip install -U --index-url="https://__token__:$GUARDRAILS_TOKEN@pypi.guardrailsai.com/simple" \
-    --extra-index-url="https://pypi.org/simple" snowglobe-connect
+pip install snowglobe-connect
 ```
 
 If using uv, set the `--prerelease=allow` flag
 ```
-pip install -U --index-url="https://__token__:$GUARDRAILS_TOKEN@pypi.guardrailsai.com/simple" \
-    --extra-index-url="https://pypi.org/simple" --prerelease=allow snowglobe-connect
+uv pip install --prerelease=allow snowglobe-connect
 ```
 
 

{snowglobe-0.4.0 → snowglobe-0.4.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "snowglobe"
-version = "0.4.0"
+version = "0.4.2"
 authors = [
     {name = "Guardrails AI", email = "contact@guardrailsai.com"}
 ]

{snowglobe-0.4.0 → snowglobe-0.4.2}/src/snowglobe/client/src/app.py

@@ -20,8 +20,8 @@ from apscheduler import AsyncScheduler
 from apscheduler.triggers.interval import IntervalTrigger
 from fastapi import FastAPI, HTTPException, Request
 
-from .cli_utils import info
-from .config import config
+from .cli_utils import info, shutdown_manager
+from .config import config, get_api_key_or_raise
 from .models import CompletionFunctionOutputs, CompletionRequest, RiskEvaluationRequest
 from .stats import initialize_stats, track_batch_completion
 from .utils import fetch_experiments, fetch_messages
@@ -190,7 +190,7 @@ async def process_application_heartbeat(app_id):
         connection_test_response = await client.post(
             connection_test_url,
             json=connection_test_payload,
-            headers={"x-api-key": config.API_KEY},
+            headers={"x-api-key": get_api_key_or_raise()},
         )
 
     if not connection_test_response.is_success:
@@ -240,7 +240,7 @@ async def process_risk_evaluation(test, risk_name):
                 "risk_type": risk_name,
                 "risk_triggered": risk_evaluation.triggered,
             },
-            headers={"x-api-key": config.API_KEY},
+            headers={"x-api-key": get_api_key_or_raise()},
         )
     LOGGER.debug(f"Time taken: {time.time() - start} seconds")
     if not risk_evaluation_response.is_success:
@@ -271,7 +271,7 @@ async def process_test(test, completion_fn, app_id):
                 "response": completionOutput.response,
                 "persona": test["persona"],
             },
-            headers={"x-api-key": config.API_KEY},
+            headers={"x-api-key": get_api_key_or_raise()},
         )
     LOGGER.debug(f"Time taken: {time.time() - start} seconds")
     # remove test['id'] from queued_tests
@@ -282,186 +282,382 @@ async def process_test(test, completion_fn, app_id):
 
 # The task to run
 async def poll_for_completions():
+    # Exit early if shutdown requested
+    if await shutdown_manager.is_async_shutdown_requested():
+        return
+
+    job_id = "poll_for_completions"
+
     if len(apps) == 0:
         LOGGER.warning("No applications found. Skipping completions check.")
         return
-    experiments = await fetch_experiments()
 
-    async with httpx.AsyncClient() as client:
-        LOGGER.info(f"Polling {len(experiments)} experiments for completions...")
-        for experiment in experiments:
-            LOGGER.debug(
-                f"Checking experiment: id={experiment.get('id')}, name={experiment.get('name', 'unknown')}"
-            )
-        for experiment in experiments:
-            app_id = experiment.get("app_id")
-            if app_id not in apps:
+    # Register this job as active
+    shutdown_manager.register_active_job(job_id)
+
+    try:
+        experiments = await fetch_experiments()
+    except Exception as e:
+        # Handle shutdown-related connection errors gracefully
+        if shutdown_manager.is_shutdown_requested():
+            LOGGER.debug(f"Connection error during shutdown (expected): {e}")
+            return
+        else:
+            LOGGER.error(f"Error fetching experiments: {e}")
+            raise
+    finally:
+        shutdown_manager.unregister_active_job(job_id)
+
+    try:
+        async with httpx.AsyncClient() as client:
+            LOGGER.info(f"Polling {len(experiments)} experiments for completions...")
+            for experiment in experiments:
                 LOGGER.debug(
-                    f"Skipping experiment as we do not have a completion function for app_id: {app_id}"
-                )
-                continue
-            experiment_id = experiment["id"]
-            experiment_request = await client.get(
-                f"{config.CONTROL_PLANE_URL}/api/experiments/{experiment_id}?appId={app_id}",
-                headers={"x-api-key": config.API_KEY},
-            )
-            if not experiment_request.is_success:
-                LOGGER.error(
-                    f"Error fetching experiment {experiment_id}: {experiment_request.text}"
+                    f"Checking experiment: id={experiment.get('id')}, name={experiment.get('name', 'unknown')}"
                 )
-                continue
-            experiment = experiment_request.json()
 
-            limit = 10
-            LOGGER.debug(f"Checking for tests for experiment {experiment_id}")
-            tests_response = await client.get(
-                f"{config.CONTROL_PLANE_URL}/api/experiments/{experiment_id}/tests?appId={app_id}&include-risk-evaluations=false&limit={limit}&unprocessed-only=true",
-                headers={"x-api-key": config.API_KEY},
-            )
+            for experiment in experiments:
+                # Check for shutdown between each experiment
+                if await shutdown_manager.is_async_shutdown_requested():
+                    LOGGER.debug("Shutdown requested, stopping completions polling")
+                    return
 
-            if not tests_response.is_success:
-                LOGGER.error(f"Error fetching tests: {tests_response.text}")
-                continue
+                app_id = experiment.get("app_id")
+                if app_id not in apps:
+                    LOGGER.debug(
+                        f"Skipping experiment as we do not have a completion function for app_id: {app_id}"
+                    )
+                    continue
+                experiment_id = experiment["id"]
 
-            tests = tests_response.json()
-            completion_count = 0
-            for test in tests:
-                LOGGER.debug(f"Found test {test['id']} for experiment {experiment_id}")
-                if not test["response"] and test["id"] not in queued_tests:
-                    completion_request = await httpx.AsyncClient().post(
-                        f"{config.SNOWGLOBE_CLIENT_URL}/completion",
-                        json={"test": test, "app_id": app_id},
-                        timeout=30,
+                try:
+                    experiment_request = await client.get(
+                        f"{config.CONTROL_PLANE_URL}/api/experiments/{experiment_id}?appId={app_id}",
+                        headers={"x-api-key": get_api_key_or_raise()},
                     )
-                    # if 429 raise and exception and stop this batch
-                    if (
-                        not completion_request.is_success
-                        and completion_request.status_code == 429
-                    ):
-                        LOGGER.warning(
-                            f"Completion Rate limit exceeded for test {test['id']}: {completion_request.text}"
-                        )
-                        raise ValueError(
-                            status_code=429,
-                            detail=f"Rate limit exceeded for test {test['id']}",
+                except (httpx.ConnectError, httpx.TimeoutException) as e:
+                    if shutdown_manager.is_shutdown_requested():
+                        LOGGER.debug(f"HTTP error during shutdown (expected): {e}")
+                        return
+                    else:
+                        LOGGER.error(
+                            f"Connection error fetching experiment {experiment_id}: {e}"
                         )
-                    completion_count += 1
-                    queued_tests[test["id"]] = True
-            if completion_count > 0:
-                experiment_name = experiment.get("name", "unknown")
-                if LOGGER.level <= logging.INFO:  # Verbose mode
-                    LOGGER.info(
-                        f"Processed {completion_count} completions for experiment {experiment_name} ({experiment_id})"
+                        continue
+
+                if not experiment_request.is_success:
+                    LOGGER.error(
+                        f"Error fetching experiment {experiment_id}: {experiment_request.text}"
                     )
-                else:  # Clean UI mode
-                    timestamp = datetime.datetime.now().strftime("%H:%M:%S")
-                    info(
-                        f"[{timestamp}] ✓ Batch complete: {completion_count} responses sent ({experiment_name})"
+                    continue
+                experiment = experiment_request.json()
+
+                limit = 10
+                LOGGER.debug(f"Checking for tests for experiment {experiment_id}")
+
+                try:
+                    tests_response = await client.get(
+                        f"{config.CONTROL_PLANE_URL}/api/experiments/{experiment_id}/tests?appId={app_id}&include-risk-evaluations=false&limit={limit}&unprocessed-only=true",
+                        headers={"x-api-key": get_api_key_or_raise()},
                     )
+                except (httpx.ConnectError, httpx.TimeoutException) as e:
+                    if shutdown_manager.is_shutdown_requested():
+                        LOGGER.debug(f"HTTP error during shutdown (expected): {e}")
+                        return
+                    else:
+                        LOGGER.error(f"Connection error fetching tests: {e}")
+                        continue
+
+                if not tests_response.is_success:
+                    LOGGER.error(f"Error fetching tests: {tests_response.text}")
+                    continue
 
-                    # Track batch completion
-                    track_batch_completion(experiment_name, completion_count)
+                tests = tests_response.json()
+                completion_count = 0
+                for test in tests:
+                    # Check for shutdown before processing each test
+                    if await shutdown_manager.is_async_shutdown_requested():
+                        LOGGER.debug("Shutdown requested, stopping test processing")
+                        return
+
+                    LOGGER.debug(
+                        f"Found test {test['id']} for experiment {experiment_id}"
+                    )
+                    if not test["response"] and test["id"] not in queued_tests:
+                        try:
+                            completion_request = await httpx.AsyncClient().post(
+                                f"{config.SNOWGLOBE_CLIENT_URL}/completion",
+                                json={"test": test, "app_id": app_id},
+                                timeout=30,
+                            )
+                        except (httpx.ConnectError, httpx.TimeoutException) as e:
+                            if shutdown_manager.is_shutdown_requested():
+                                LOGGER.debug(
+                                    f"HTTP error during shutdown (expected): {e}"
+                                )
+                                return
+                            else:
+                                LOGGER.error(
+                                    f"Connection error posting completion: {e}"
+                                )
+                                continue
+
+                        # if 429 raise and exception and stop this batch
+                        if (
+                            not completion_request.is_success
+                            and completion_request.status_code == 429
+                        ):
+                            LOGGER.warning(
+                                f"Completion Rate limit exceeded for test {test['id']}: {completion_request.text}"
+                            )
+                            raise ValueError(
+                                status_code=429,
+                                detail=f"Rate limit exceeded for test {test['id']}",
+                            )
+                        completion_count += 1
+                        queued_tests[test["id"]] = True
+
+                if completion_count > 0:
+                    experiment_name = experiment.get("name", "unknown")
+                    if LOGGER.level <= logging.INFO:  # Verbose mode
+                        LOGGER.info(
+                            f"Processed {completion_count} completions for experiment {experiment_name} ({experiment_id})"
+                        )
+                    else:  # Clean UI mode
+                        timestamp = datetime.datetime.now().strftime("%H:%M:%S")
+                        info(
+                            f"[{timestamp}] ✓ Batch complete: {completion_count} responses sent ({experiment_name})"
+                        )
+
+                        # Track batch completion
+                        track_batch_completion(experiment_name, completion_count)
+
+    except Exception as e:
+        # Handle any other errors
+        if shutdown_manager.is_shutdown_requested():
+            LOGGER.debug(f"Error during shutdown (expected): {e}")
+        else:
+            LOGGER.error(f"Unexpected error in poll_for_completions: {e}")
+            raise
 
 
 async def process_application_heartbeats():
-    connection_test_count = 0
-    LOGGER.info("Processing application heartbeats...")
-    for app_id, app_info in apps.items():
-        connection_test_request = await httpx.AsyncClient().post(
-            f"{config.SNOWGLOBE_CLIENT_URL}/heartbeat",
-            json={"app_id": app_id},
-            timeout=30,
-        )
-        if not connection_test_request.is_success:
-            LOGGER.error(
-                f"Error sending heartbeat for application {app_id}: {connection_test_request.text}"
-            )
-            continue
-        connection_test_count += 1
+    # Exit early if shutdown requested
+    if await shutdown_manager.is_async_shutdown_requested():
+        return
 
-    LOGGER.info(f"Processed {connection_test_count} heartbeats for applications.")
+    job_id = "process_application_heartbeats"
+    shutdown_manager.register_active_job(job_id)
 
+    try:
+        connection_test_count = 0
+        LOGGER.info("Processing application heartbeats...")
+        for app_id, app_info in apps.items():
+            # Check for shutdown between each heartbeat
+            if await shutdown_manager.is_async_shutdown_requested():
+                LOGGER.debug("Shutdown requested, stopping heartbeats")
+                return
 
-async def poll_for_risk_evaluations():
-    """Poll for risk evaluations and process them."""
-    experiments = await fetch_experiments()
-    LOGGER.info("Checking for pending risk evaluations...")
-    LOGGER.debug(f"Found {len(experiments)} experiments with validation in progress")
-    # experiments = [{"id": "123"}]
-    async with httpx.AsyncClient() as client:
-        for experiment in experiments:
-            experiment_request = await client.get(
-                f"{config.CONTROL_PLANE_URL}/api/experiments/{experiment['id']}",
-                headers={"x-api-key": config.API_KEY},
-            )
-            if not experiment_request.is_success:
+            try:
+                connection_test_request = await httpx.AsyncClient().post(
+                    f"{config.SNOWGLOBE_CLIENT_URL}/heartbeat",
+                    json={"app_id": app_id},
+                    timeout=30,
+                )
+            except (httpx.ConnectError, httpx.TimeoutException) as e:
+                if shutdown_manager.is_shutdown_requested():
+                    LOGGER.debug(f"HTTP error during shutdown (expected): {e}")
+                    return
+                else:
+                    LOGGER.error(
+                        f"Connection error sending heartbeat for {app_id}: {e}"
+                    )
+                    continue
+
+            if not connection_test_request.is_success:
                 LOGGER.error(
-                    f"Error fetching experiment {experiment['id']}: {experiment_request.text}"
+                    f"Error sending heartbeat for application {app_id}: {connection_test_request.text}"
                 )
                 continue
-            experiment = experiment_request.json()
-            risk_eval_count = 0
-            for risk_name in risks.keys():
+            connection_test_count += 1
+
+        LOGGER.info(f"Processed {connection_test_count} heartbeats for applications.")
+    finally:
+        shutdown_manager.unregister_active_job(job_id)
+
+
+async def poll_for_risk_evaluations():
+    """Poll for risk evaluations and process them."""
+    # Exit early if shutdown requested
+    if await shutdown_manager.is_async_shutdown_requested():
+        return
+
+    job_id = "poll_for_risk_evaluations"
+    shutdown_manager.register_active_job(job_id)
+
+    try:
+        experiments = await fetch_experiments()
+    except Exception as e:
+        # Handle shutdown-related connection errors gracefully
+        if shutdown_manager.is_shutdown_requested():
+            LOGGER.debug(f"Connection error during shutdown (expected): {e}")
+            return
+        else:
+            LOGGER.error(f"Error fetching experiments: {e}")
+            raise
+    finally:
+        shutdown_manager.unregister_active_job(job_id)
+
+    try:
+        LOGGER.info("Checking for pending risk evaluations...")
+        LOGGER.debug(
+            f"Found {len(experiments)} experiments with validation in progress"
+        )
+
+        async with httpx.AsyncClient() as client:
+            for experiment in experiments:
+                # Check for shutdown between experiments
+                if await shutdown_manager.is_async_shutdown_requested():
+                    LOGGER.debug("Shutdown requested, stopping risk evaluations")
+                    return
+
                 try:
-                    if (
-                        risk_name
-                        not in experiment.get("source_data", {})
-                        .get("evaluation_configuration", {})
-                        .keys()
-                    ):
-                        LOGGER.debug(
-                            f"Skipping experiment {experiment['id']} as it does not have risk {risk_name}"
+                    experiment_request = await client.get(
+                        f"{config.CONTROL_PLANE_URL}/api/experiments/{experiment['id']}",
+                        headers={"x-api-key": get_api_key_or_raise()},
+                    )
+                except (httpx.ConnectError, httpx.TimeoutException) as e:
+                    if shutdown_manager.is_shutdown_requested():
+                        LOGGER.debug(f"HTTP error during shutdown (expected): {e}")
+                        return
+                    else:
+                        LOGGER.error(
+                            f"Connection error fetching experiment {experiment['id']}: {e}"
                         )
                         continue
-                    LOGGER.debug(
-                        f"checking for tests for experiment {experiment['id']}"
-                    )
-                    tests_response = await client.get(
-                        f"{config.CONTROL_PLANE_URL}/api/experiments/{experiment['id']}/tests?unevaluated-risk={quote_plus(risk_name)}&include-risk-evaluations=true",
-                        headers={"x-api-key": config.API_KEY},
+
+                if not experiment_request.is_success:
+                    LOGGER.error(
+                        f"Error fetching experiment {experiment['id']}: {experiment_request.text}"
                     )
+                    continue
+                experiment = experiment_request.json()
+                risk_eval_count = 0
 
-                    if not tests_response.is_success:
-                        message = (
-                            tests_response.json().get("message") or tests_response.text
-                        )
-                        raise ValueError(
-                            status_code=tests_response.status_code,
-                            message=message,
-                        )
-                    tests = tests_response.json()
+                for risk_name in risks.keys():
+                    # Check for shutdown before processing each risk
+                    if await shutdown_manager.is_async_shutdown_requested():
+                        LOGGER.debug("Shutdown requested, stopping risk processing")
+                        return
 
-                    for test in tests:
-                        test_id = test["id"]
+                    try:
                         if (
-                            test_id not in queued_evaluations
-                            and test.get("response") is not None
+                            risk_name
+                            not in experiment.get("source_data", {})
+                            .get("evaluation_configuration", {})
+                            .keys()
                         ):
-                            risk_eval_response = await httpx.AsyncClient().post(
-                                f"{config.SNOWGLOBE_CLIENT_URL}/risk-evaluation",
-                                json={"test": test, "risk_name": risk_name},
-                                timeout=30,
+                            LOGGER.debug(
+                                f"Skipping experiment {experiment['id']} as it does not have risk {risk_name}"
                             )
-                            # if risk evaltion response is 429 raise and exception and bail on this batch
-                            if (
-                                not risk_eval_response.is_success
-                                and risk_eval_response.status_code == 429
-                            ):
-                                LOGGER.error(
-                                    f"Rate limit exceeded for risk evaluation {test['id']}: {risk_eval_response.text}"
+                            continue
+                        LOGGER.debug(
+                            f"checking for tests for experiment {experiment['id']}"
+                        )
+
+                        try:
+                            tests_response = await client.get(
+                                f"{config.CONTROL_PLANE_URL}/api/experiments/{experiment['id']}/tests?unevaluated-risk={quote_plus(risk_name)}&include-risk-evaluations=true",
+                                headers={"x-api-key": get_api_key_or_raise()},
+                            )
+                        except (httpx.ConnectError, httpx.TimeoutException) as e:
+                            if shutdown_manager.is_shutdown_requested():
+                                LOGGER.debug(
+                                    f"HTTP error during shutdown (expected): {e}"
                                 )
-                                raise ValueError(
-                                    status_code=429,
-                                    detail=f"Rate limit exceeded for risk evaluation {test['id']}",
+                                return
+                            else:
+                                LOGGER.error(f"Connection error fetching tests: {e}")
+                                continue
+
+                        if not tests_response.is_success:
+                            message = (
+                                tests_response.json().get("message")
+                                or tests_response.text
+                            )
+                            raise ValueError(
+                                status_code=tests_response.status_code,
+                                message=message,
+                            )
+                        tests = tests_response.json()
+
+                        for test in tests:
+                            # Check for shutdown before processing each test
+                            if await shutdown_manager.is_async_shutdown_requested():
+                                LOGGER.debug(
+                                    "Shutdown requested, stopping test evaluation"
                                 )
-                            queued_evaluations[test_id] = True
-                            risk_eval_count += 1
-                except Exception as e:
-                    LOGGER.error(f"Error fetching tests: {e}")
-            if risk_eval_count > 0:
-                LOGGER.info(
-                    f"Processed {risk_eval_count} risk evaluations for experiment {experiment.get('name', 'unknown')} ({experiment['id']})"
-                )
+                                return
+
+                            test_id = test["id"]
+                            if (
+                                test_id not in queued_evaluations
+                                and test.get("response") is not None
+                            ):
+                                try:
+                                    risk_eval_response = await httpx.AsyncClient().post(
+                                        f"{config.SNOWGLOBE_CLIENT_URL}/risk-evaluation",
+                                        json={"test": test, "risk_name": risk_name},
+                                        timeout=30,
+                                    )
+                                except (
+                                    httpx.ConnectError,
+                                    httpx.TimeoutException,
+                                ) as e:
+                                    if shutdown_manager.is_shutdown_requested():
+                                        LOGGER.debug(
+                                            f"HTTP error during shutdown (expected): {e}"
+                                        )
+                                        return
+                                    else:
+                                        LOGGER.error(
+                                            f"Connection error posting risk evaluation: {e}"
+                                        )
+                                        continue
+
+                                # if risk evaltion response is 429 raise and exception and bail on this batch
+                                if (
+                                    not risk_eval_response.is_success
+                                    and risk_eval_response.status_code == 429
+                                ):
+                                    LOGGER.error(
+                                        f"Rate limit exceeded for risk evaluation {test['id']}: {risk_eval_response.text}"
+                                    )
+                                    raise ValueError(
+                                        status_code=429,
+                                        detail=f"Rate limit exceeded for risk evaluation {test['id']}",
+                                    )
+                                queued_evaluations[test_id] = True
+                                risk_eval_count += 1
+                    except Exception as e:
+                        if shutdown_manager.is_shutdown_requested():
+                            LOGGER.debug(f"Error during shutdown (expected): {e}")
+                            return
+                        else:
+                            LOGGER.error(f"Error fetching tests: {e}")
+
+                if risk_eval_count > 0:
+                    LOGGER.info(
+                        f"Processed {risk_eval_count} risk evaluations for experiment {experiment.get('name', 'unknown')} ({experiment['id']})"
+                    )
+
+    except Exception as e:
+        # Handle any other errors
+        if shutdown_manager.is_shutdown_requested():
+            LOGGER.debug(f"Error during shutdown (expected): {e}")
+        else:
+            LOGGER.error(f"Unexpected error in poll_for_risk_evaluations: {e}")
+            raise
 
 
 # Ensure the scheduler shuts down properly on application exit.
@@ -599,7 +795,7 @@ async def lifespan(app: FastAPI):
     for app_id in apps.keys():
         try:
             connection_test_url = (
-            f"{config.CONTROL_PLANE_URL}/api/failed-code-connection-tests"
+                f"{config.CONTROL_PLANE_URL}/api/failed-code-connection-tests"
             )
             connection_test_payload = {
                 "appId": app_id,