snowglobe 0.4.0__py3-none-any.whl

snowglobe/client/src/cli_utils.py

@@ -0,0 +1,361 @@
+import contextlib
+import math
+import os
+import sys
+import time
+from typing import Any, Dict, List, Optional, Tuple
+
+import requests
+from rich.console import Console
+from rich.prompt import Confirm, Prompt
+from rich.table import Table
+
+from .config import get_rc_file_path
+from .stats import get_shutdown_stats
+
+console = Console()
+
+
+class CliState:
+    """Global CLI state management"""
+
+    def __init__(self):
+        self.verbose = False
+        self.quiet = False
+        self.json_output = False
+
+
+cli_state = CliState()
+
+
+def get_api_key() -> Optional[str]:
+    """Get API key from environment or config file"""
+    api_key = os.getenv("SNOWGLOBE_API_KEY") or os.getenv("GUARDRAILS_API_KEY")
+    if not api_key:
+        rc_path = get_rc_file_path()
+        if os.path.exists(rc_path):
+            with open(rc_path, "r") as rc_file:
+                for line in rc_file:
+                    if line.startswith("SNOWGLOBE_API_KEY="):
+                        api_key = line.strip().split("=", 1)[1]
+                        break
+    return api_key
+
+
+def get_control_plane_url() -> str:
+    """Get control plane URL from environment or config file"""
+    control_plane_url = os.getenv("CONTROL_PLANE_URL")
+    if not control_plane_url:
+        control_plane_url = "https://api.snowglobe.guardrailsai.com"
+        rc_path = get_rc_file_path()
+        if os.path.exists(rc_path):
+            with open(rc_path, "r") as rc_file:
+                for line in rc_file:
+                    if line.startswith("CONTROL_PLANE_URL="):
+                        control_plane_url = line.strip().split("=", 1)[1]
+                        break
+    return control_plane_url
+
+
+def success(message: str) -> None:
+    """Print success message with formatting"""
+    if cli_state.json_output:
+        return
+    if not cli_state.quiet:
+        console.print(f"✅ {message}", style="green")
+
+
+def warning(message: str) -> None:
+    """Print warning message with formatting"""
+    if cli_state.json_output:
+        return
+    if not cli_state.quiet:
+        console.print(f"⚠️  {message}", style="yellow")
+
+
+def error(message: str) -> None:
+    """Print error message with formatting"""
+    if cli_state.json_output:
+        return
+    console.print(f"❌ {message}", style="red")
+
+
+def info(message: str) -> None:
+    """Print info message with formatting"""
+    if cli_state.json_output:
+        return
+    if not cli_state.quiet:
+        console.print(f"💡 {message}", style="blue")
+
+
+def debug(message: str) -> None:
+    """Print debug message if verbose mode is enabled"""
+    if cli_state.json_output:
+        return
+    if cli_state.verbose:
+        console.print(f"🔍 {message}", style="dim")
+
+
+def docs_link(message: str, url: str = "https://www.snowglobe.so/docs") -> None:
+    """Print documentation link"""
+    if cli_state.json_output:
+        return
+    if not cli_state.quiet:
+        console.print(f"📖 {message}: {url}", style="cyan")
+
+
+def graceful_shutdown():
+    """Handle graceful shutdown with session summary"""
+    console.print("\n")
+    warning("🛑 Shutting down gracefully...")
+    success("Completing current scenarios")
+    success("Connection closed")
+
+    stats = get_shutdown_stats()
+
+    if stats and stats["total_messages"] > 0:
+        success("Session summary:")
+        if len(stats["experiment_totals"]) > 1:
+            # Multiple experiments - show breakdown
+            for exp_name, count in stats["experiment_totals"].items():
+                console.print(f"   • {exp_name}: {count} scenarios processed")
+            console.print(
+                f"   • Total: {stats['total_messages']} scenarios in {stats['uptime']}"
+            )
+        else:
+            # Single experiment or total only
+            console.print(
+                f"   • {stats['total_messages']} scenarios processed in {stats['uptime']}"
+            )
+    else:
+        success("No scenarios processed during this session")
+
+    console.print()
+    success("Agent disconnected successfully")
+    sys.exit(0)
+
+
+@contextlib.contextmanager
+def spinner(text: str):
+    """Context manager for showing a spinner during operations"""
+    if cli_state.json_output or cli_state.quiet:
+        yield
+        return
+
+    with console.status(f"[bold blue]{text}..."):
+        yield
+
+
+def check_auth_status() -> Tuple[bool, str, Dict[str, Any]]:
+    """Check authentication status"""
+    api_key = get_api_key()
+    if not api_key:
+        return False, "No API key found", {}
+
+    control_plane_url = get_control_plane_url()
+    try:
+        response = requests.get(
+            f"{control_plane_url}/api/applications",
+            headers={"x-api-key": api_key},
+            timeout=10,
+        )
+        if response.status_code == 200:
+            return True, "Authenticated", response.json()
+        else:
+            return False, f"Authentication failed: {response.status_code}", {}
+    except requests.RequestException as e:
+        return False, f"Connection error: {str(e)}", {}
+
+
+def select_stateful_interactive(
+    stateful: bool = False,
+) -> bool:
+    """Interactive prompt to confirm if the agent is stateful"""
+    if cli_state.json_output:
+        # For JSON mode, just return the default stateful value
+        return stateful
+    info("Some stateful agents such as ones that maintain communication over a websocket or convo specific completion endpoint require stateful integration.")
+    info("If your agent takes messages and completions on a single completion endpoint regardless of context, you can answer no to the following question.")
+    if Confirm.ask("Would you like to create a new application?"):
+        return True
+    return False
+
+def select_application_interactive(
+    applications: List[Dict[str, Any]],
+) -> Optional[Dict[str, Any]]:
+    """Clean, readable application selection interface"""
+    if cli_state.json_output:
+        # For JSON mode, just return the first app or None
+        return applications[0] if applications else None
+
+    if not applications:
+        info("No applications found")
+        if Confirm.ask("Would you like to create a new application?"):
+            return "new"
+        return None
+
+    # Sort applications by updated_at (most recent first)
+    sorted_applications = sort_applications_by_date(applications)
+
+    return display_applications_clean(sorted_applications)
+
+
+def sort_applications_by_date(
+    applications: List[Dict[str, Any]],
+) -> List[Dict[str, Any]]:
+    """Sort applications by updated_at date, most recent first"""
+
+    def get_sort_key(app):
+        updated_at = app.get("updated_at", "")
+        if not updated_at:
+            return ""  # Apps without dates go to the end
+        return updated_at
+
+    # Sort in reverse order (most recent first)
+    return sorted(applications, key=get_sort_key, reverse=True)
+
+
+def display_applications_clean(
+    applications: List[Dict[str, Any]], page_size: int = 15
+) -> Optional[Dict[str, Any]]:
+    """Display applications in a clean table format"""
+    total_apps = len(applications)
+    total_pages = math.ceil(total_apps / page_size) if total_apps > 0 else 1
+    current_page = 0
+
+    # Check if any apps have updated_at date information
+    has_dates = any(app.get("updated_at") for app in applications)
+
+    while True:
+        # Calculate page boundaries
+        start_idx = current_page * page_size
+        end_idx = min(start_idx + page_size, total_apps)
+        page_apps = applications[start_idx:end_idx]
+
+        # Create table
+        table = Table(title=f"📱 Your Applications ({total_apps} total)")
+        if total_pages > 1:
+            table.title = f"📱 Your Applications ({total_apps} total) - Page {current_page + 1}/{total_pages}"
+
+        table.add_column("#", style="bold blue", width=4)
+        table.add_column("Name", style="bold", min_width=15)
+
+        # Add date column if date info is available
+        if has_dates:
+            table.add_column("Updated", style="dim", min_width=10)
+
+        table.add_column("Description", style="green", min_width=20)
+
+        # Add applications to table
+        for i, app in enumerate(page_apps):
+            app_idx = start_idx + i + 1
+            name = app.get("name", "Unknown")
+            description = app.get("description", "No description")
+
+            # Clean up description - remove newlines and extra spaces
+            description = " ".join(description.split())
+
+            # Truncate description to 20 characters
+            if len(description) > 20:
+                description = description[:17] + "..."
+
+            # Get the best available date
+            date_str = "-"
+            if has_dates:
+                date_str = get_best_date(app)
+
+            # Build row based on whether we have dates
+            if has_dates:
+                table.add_row(str(app_idx), name, date_str, description)
+            else:
+                table.add_row(str(app_idx), name, description)
+
+        # Add create new option
+        if has_dates:
+            table.add_row("new", "🆕 Create New App", "-", "Set up new application")
+        else:
+            table.add_row("new", "🆕 Create New App", "Set up new application")
+
+        console.print(table)
+
+        # Navigation instructions
+        nav_options = []
+        if current_page > 0:
+            nav_options.append("[bold cyan]p[/bold cyan] Previous")
+        if current_page < total_pages - 1:
+            nav_options.append("[bold cyan]n[/bold cyan] Next")
+        nav_options.extend(
+            [
+                f"[bold yellow]1-{total_apps}[/bold yellow] Select app",
+                "[bold green]new[/bold green] Create new",
+                "[bold red]q[/bold red] Quit",
+            ]
+        )
+
+        console.print("\nOptions: " + " | ".join(nav_options))
+
+        # Get user input
+        try:
+            choice = Prompt.ask("\n[bold]Your choice[/bold]").strip().lower()
+
+            if choice == "q":
+                return None
+            elif choice == "p" and current_page > 0:
+                current_page -= 1
+                continue
+            elif choice == "n" and current_page < total_pages - 1:
+                current_page += 1
+                continue
+            elif choice == "new":
+                return "new"
+            elif choice.isdigit():
+                idx = int(choice)
+                if 1 <= idx <= total_apps:
+                    return applications[idx - 1]
+                else:
+                    error(f"Please choose between 1 and {total_apps}")
+                    time.sleep(1)
+            else:
+                error("Invalid choice. Try again.")
+                time.sleep(1)
+
+        except (KeyboardInterrupt, EOFError):
+            warning("\nSelection cancelled")
+            return None
+
+
+def get_best_date(app: Dict[str, Any]) -> str:
+    """Get updated_at date formatted for display"""
+    date_value = app.get("updated_at")
+
+    if not date_value:
+        return "-"
+
+    # Format ISO date like "2025-07-29T04:35:22.093Z" to "2025-07-29"
+    date_str = str(date_value)
+    if "T" in date_str:
+        return date_str.split("T")[0]  # Take just the date part
+    elif len(date_str) > 10:
+        return date_str[:10]
+    return date_str
+
+
+def get_remote_applications() -> Tuple[bool, List[Dict[str, Any]], str]:
+    """Fetch applications from the remote API"""
+    api_key = get_api_key()
+    if not api_key:
+        return False, [], "No API key found"
+
+    control_plane_url = get_control_plane_url()
+    try:
+        response = requests.get(
+            f"{control_plane_url}/api/applications",
+            headers={"x-api-key": api_key},
+            timeout=10,
+        )
+        if response.status_code == 200:
+            return True, response.json(), "Success"
+        else:
+            return False, [], f"HTTP {response.status_code}: {response.text}"
+    except requests.RequestException as e:
+        return False, [], f"Connection error: {str(e)}"

snowglobe/client/src/config.py

@@ -0,0 +1,213 @@
+import logging
+import os
+
+from rich.console import Console
+
+LOGGER = logging.getLogger(__name__)
+console = Console()
+
+
+def get_rc_file_path() -> str:
+    """Get path to snowglobe rc file"""
+    return os.path.join(os.getcwd(), ".snowglobe", "config.rc")
+
+
+def get_legacy_rc_file_path() -> str:
+    """Get path to legacy .snowgloberc file for migration purposes"""
+    return os.path.join(os.getcwd(), ".snowgloberc")
+
+
+def migrate_rc_file_if_needed():
+    """Migrate .snowgloberc to .snowglobe/config.rc if needed"""
+    legacy_path = get_legacy_rc_file_path()
+    new_path = get_rc_file_path()
+
+    # If new file already exists, no migration needed
+    if os.path.exists(new_path):
+        return
+
+    # If legacy file exists, migrate it
+    if os.path.exists(legacy_path):
+        # Ensure .snowglobe directory exists
+        os.makedirs(os.path.dirname(new_path), exist_ok=True)
+
+        # Copy content from legacy to new location
+        with open(legacy_path, "r") as legacy_file:
+            content = legacy_file.read()
+
+        with open(new_path, "w") as new_file:
+            new_file.write(content)
+
+        LOGGER.info(f"Migrated {legacy_path} to {new_path}")
+
+
+class Config:
+    def __init__(self):
+        # Migrate legacy .snowgloberc if needed
+        migrate_rc_file_if_needed()
+
+        self.SNOWGLOBE_APPLICATION_HEARTBEAT_INTERVAL_MINUTES = int(
+            os.getenv("SNOWGLOBE_APPLICATION_HEARTBEAT_INTERVAL_MINUTES", "5")
+        )
+        self.CONCURRENT_HEARTBEATS_PER_INTERVAL = 120
+        self.CONCURRENT_HEARTBEATS_INTERVAL_SECONDS = 60
+
+        # Initialize API key - just raise exception if missing, let caller handle it
+        self.API_KEY = self.get_api_key()
+        self.APPLICATION_ID = self.get_application_id()
+        self.CONTROL_PLANE_URL = self.get_control_plane_url()
+        self.SNOWGLOBE_CLIENT_URL = self.get_snowglobe_client_url()
+        self.SNOWGLOBE_CLIENT_PORT = self.get_snowglobe_client_port()
+        self.CONCURRENT_COMPLETIONS_PER_INTERVAL = self.get_completions_per_interval()
+        self.CONCURRENT_COMPLETIONS_INTERVAL_SECONDS = (
+            self.get_completions_interval_seconds()
+        )
+        self.CONCURRENT_RISK_EVALUATIONS = self.get_concurrent_risk_evaluations()
+        self.CONCURRENT_RISK_EVALUATIONS_INTERVAL_SECONDS = (
+            self.get_concurrent_risk_evaluations_interval_seconds()
+        )
+        self.DEBUG = os.getenv("DEBUG", "false").lower() == "true"
+
+    def get_snowglobe_client_port(self) -> int:
+        snowglobe_client_port = os.getenv("SNOWGLOBE_CLIENT_PORT")
+        if not snowglobe_client_port:
+            snowglobe_client_port = "8000"
+            rc_path = get_rc_file_path()
+            if os.path.exists(rc_path):
+                with open(rc_path, "r") as rc_file:
+                    for line in rc_file:
+                        if line.startswith("SNOWGLOBE_CLIENT_PORT="):
+                            snowglobe_client_port = line.strip().split("=", 1)[1]
+                            break
+        LOGGER.debug(f"setting SNOWGLOBE_CLIENT_PORT: {snowglobe_client_port}")
+        return int(snowglobe_client_port)
+
+    def get_snowglobe_client_url(self) -> str:
+        snowglobe_client_url = os.getenv("SNOWGLOBE_CLIENT_URL")
+        if not snowglobe_client_url:
+            snowglobe_client_url = "http://localhost:8000"
+            rc_path = get_rc_file_path()
+            if os.path.exists(rc_path):
+                with open(rc_path, "r") as rc_file:
+                    for line in rc_file:
+                        if line.startswith("SNOWGLOBE_CLIENT_URL="):
+                            snowglobe_client_url = line.strip().split("=", 1)[1]
+                            break
+        LOGGER.debug(f"setting SNOWGLOBE_CLIENT_URL: {snowglobe_client_url}")
+        return snowglobe_client_url
+
+    def get_completions_interval_seconds(self) -> int:
+        completions_interval_seconds = os.getenv("COMPLETIONS_INTERVAL_SECONDS")
+        if not completions_interval_seconds:
+            completions_interval_seconds = "60"
+            LOGGER.debug(
+                "COMPLETIONS_INTERVAL_SECONDS not found in environment variables, using COMPLETIONS_INTERVAL as fallback"
+            )
+            rc_path = get_rc_file_path()
+            if os.path.exists(rc_path):
+                with open(rc_path, "r") as rc_file:
+                    for line in rc_file:
+                        if line.startswith("COMPLETIONS_INTERVAL_SECONDS="):
+                            completions_interval_seconds = line.strip().split("=", 1)[1]
+                            break
+        return int(completions_interval_seconds)
+
+    def get_completions_per_interval(self) -> int:
+        completions_per_second = os.getenv("COMPLETIONS_PER_SECOND")
+        if not completions_per_second:
+            completions_per_second = "120"
+            LOGGER.debug(
+                "COMPLETIONS_PER_SECOND not found in environment variables, using MAX_COMPLETIONS as fallback"
+            )
+            rc_path = get_rc_file_path()
+            if os.path.exists(rc_path):
+                with open(rc_path, "r") as rc_file:
+                    for line in rc_file:
+                        if line.startswith("COMPLETIONS_PER_SECOND="):
+                            completions_per_second = line.strip().split("=", 1)[1]
+                            break
+        return int(completions_per_second)
+
+    def get_concurrent_risk_evaluations(self) -> int:
+        concurrent_risk_evaluations = os.getenv("CONCURRENT_RISK_EVALUATIONS")
+        if not concurrent_risk_evaluations:
+            concurrent_risk_evaluations = "120"
+            rc_path = get_rc_file_path()
+            if os.path.exists(rc_path):
+                with open(rc_path, "r") as rc_file:
+                    for line in rc_file:
+                        if line.startswith("CONCURRENT_RISK_EVALUATIONS="):
+                            concurrent_risk_evaluations = line.strip().split("=", 1)[1]
+                            break
+        return int(concurrent_risk_evaluations)
+
+    def get_concurrent_risk_evaluations_interval_seconds(self) -> int:
+        concurrent_risk_evaluations_interval_seconds = os.getenv(
+            "CONCURRENT_RISK_EVALUATIONS_INTERVAL_SECONDS"
+        )
+        if not concurrent_risk_evaluations_interval_seconds:
+            concurrent_risk_evaluations_interval_seconds = "60"
+            rc_path = get_rc_file_path()
+            if os.path.exists(rc_path):
+                with open(rc_path, "r") as rc_file:
+                    for line in rc_file:
+                        if line.startswith(
+                            "CONCURRENT_RISK_EVALUATIONS_INTERVAL_SECONDS="
+                        ):
+                            concurrent_risk_evaluations_interval_seconds = (
+                                line.strip().split("=", 1)[1]
+                            )
+                            break
+        return int(concurrent_risk_evaluations_interval_seconds)
+
+    def get_control_plane_url(self) -> str:
+        control_plane_url = os.getenv("CONTROL_PLANE_URL")
+        if not control_plane_url:
+            control_plane_url = "https://api.snowglobe.guardrailsai.com"
+            rc_path = get_rc_file_path()
+            if os.path.exists(rc_path):
+                with open(rc_path, "r") as rc_file:
+                    for line in rc_file:
+                        LOGGER.debug(f"Checking line: {line.strip()}")
+                        if line.startswith("CONTROL_PLANE_URL="):
+                            control_plane_url = line.strip().split("=", 1)[1]
+                            break
+        LOGGER.debug(f"setting CONTROL_PLANE_URL: {control_plane_url}")
+        return control_plane_url
+
+    def get_api_key(self) -> str:
+        api_key = os.getenv("SNOWGLOBE_API_KEY") or os.getenv("GUARDRAILS_API_KEY")
+        if not api_key:
+            rc_path = get_rc_file_path()
+            if os.path.exists(rc_path):
+                with open(rc_path, "r") as rc_file:
+                    for line in rc_file:
+                        if line.startswith("SNOWGLOBE_API_KEY="):
+                            api_key = line.strip().split("=", 1)[1]
+                            break
+        if not api_key:
+            raise ValueError(
+                "API key is required either passed as an argument or set as an environment variable. \n"
+                "You can set the API key as an environment variable by running: \n"
+                "export SNOWGLOBE_API_KEY=<your_api_key> \n"
+                "or \n"
+                "export GUARDRAILS_API_KEY=<your_api_key> \n"
+                "Or you can create a .snowglobe/config.rc file in the current directory with the line: \n"
+                "SNOWGLOBE_API_KEY=<your_api_key> \n"
+            )
+        return api_key
+
+    def get_application_id(self) -> str:
+        application_id = os.getenv("SNOWGLOBE_APP_ID")
+        if not application_id:
+            rc_path = get_rc_file_path()
+            if os.path.exists(rc_path):
+                with open(rc_path, "r") as rc_file:
+                    for line in rc_file:
+                        if line.startswith("SNOWGLOBE_APP_ID="):
+                            application_id = line.strip().split("=", 1)[1]
+                            break
+        return application_id
+
+
+config = Config()

snowglobe/client/src/models.py

@@ -0,0 +1,37 @@
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel
+
+
+class SnowglobeData(BaseModel):
+    conversation_id: str
+    test_id: str
+
+
+class SnowglobeMessage(BaseModel):
+    role: str
+    content: str
+    snowglobe_data: Optional[SnowglobeData] = None
+
+
+class CompletionFunctionOutputs(BaseModel):
+    response: str
+
+
+class CompletionRequest(BaseModel):
+    messages: List[SnowglobeMessage]
+
+    def to_openai_messages(self) -> List[Dict]:
+        """Return a list of OpenAI messages from the Snowglobe messages"""
+        return [{"role": msg.role, "content": msg.content} for msg in self.messages]
+
+
+class RiskEvaluationRequest(BaseModel):
+    messages: List[SnowglobeMessage]
+
+
+class RiskEvaluationOutputs(BaseModel):
+    triggered: bool
+    tags: Optional[Dict[str, str]] = None
+    reason: Optional[str] = None
+    severity: Optional[int] = None