|
| 1 | +import logging |
| 2 | +import os |
| 3 | +import sys |
| 4 | +import time |
| 5 | +from functools import wraps |
| 6 | +from typing import Callable, Optional |
| 7 | + |
| 8 | +import click |
| 9 | +from dotenv import load_dotenv |
| 10 | + |
| 11 | +from humanloop import Humanloop |
| 12 | +from humanloop.sync.sync_client import SyncClient |
| 13 | + |
| 14 | +# Set up logging |
| 15 | +logger = logging.getLogger(__name__) |
| 16 | +logger.setLevel(logging.INFO) # Set back to INFO level |
| 17 | +console_handler = logging.StreamHandler() |
| 18 | +formatter = logging.Formatter("%(message)s") # Simplified formatter |
| 19 | +console_handler.setFormatter(formatter) |
| 20 | +if not logger.hasHandlers(): |
| 21 | + logger.addHandler(console_handler) |
| 22 | + |
| 23 | +# Color constants |
| 24 | +SUCCESS_COLOR = "green" |
| 25 | +ERROR_COLOR = "red" |
| 26 | +INFO_COLOR = "blue" |
| 27 | +WARNING_COLOR = "yellow" |
| 28 | + |
| 29 | + |
| 30 | +def load_api_key(env_file: Optional[str] = None) -> str: |
| 31 | + """Load API key from .env file or environment variable. |
| 32 | +
|
| 33 | + Args: |
| 34 | + env_file: Optional path to .env file |
| 35 | +
|
| 36 | + Returns: |
| 37 | + str: The loaded API key |
| 38 | +
|
| 39 | + Raises: |
| 40 | + click.ClickException: If no API key is found |
| 41 | + """ |
| 42 | + # Try specific .env file if provided, otherwise default to .env in current directory |
| 43 | + if env_file: |
| 44 | + if not load_dotenv(env_file): # load_dotenv returns False if file not found/invalid |
| 45 | + raise click.ClickException( |
| 46 | + click.style( |
| 47 | + f"Failed to load environment file: {env_file} (file not found or invalid format)", |
| 48 | + fg=ERROR_COLOR, |
| 49 | + ) |
| 50 | + ) |
| 51 | + else: |
| 52 | + load_dotenv() # Attempt to load from default .env in current directory |
| 53 | + |
| 54 | + # Get API key from environment |
| 55 | + api_key = os.getenv("HUMANLOOP_API_KEY") |
| 56 | + if not api_key: |
| 57 | + raise click.ClickException( |
| 58 | + click.style( |
| 59 | + "No API key found. Set HUMANLOOP_API_KEY in .env file or environment, or use --api-key", fg=ERROR_COLOR |
| 60 | + ) |
| 61 | + ) |
| 62 | + |
| 63 | + return api_key |
| 64 | + |
| 65 | + |
| 66 | +def get_client( |
| 67 | + api_key: Optional[str] = None, env_file: Optional[str] = None, base_url: Optional[str] = None |
| 68 | +) -> Humanloop: |
| 69 | + """Instantiate a Humanloop client for the CLI. |
| 70 | +
|
| 71 | + Args: |
| 72 | + api_key: Optional API key provided directly |
| 73 | + env_file: Optional path to .env file |
| 74 | + base_url: Optional base URL for the API |
| 75 | +
|
| 76 | + Returns: |
| 77 | + Humanloop: Configured client instance |
| 78 | +
|
| 79 | + Raises: |
| 80 | + click.ClickException: If no API key is found |
| 81 | + """ |
| 82 | + if not api_key: |
| 83 | + api_key = load_api_key(env_file) |
| 84 | + return Humanloop(api_key=api_key, base_url=base_url) |
| 85 | + |
| 86 | + |
| 87 | +def common_options(f: Callable) -> Callable: |
| 88 | + """Decorator for common CLI options.""" |
| 89 | + |
| 90 | + @click.option( |
| 91 | + "--api-key", |
| 92 | + help="Humanloop API key. If not provided, uses HUMANLOOP_API_KEY from .env or environment.", |
| 93 | + default=None, |
| 94 | + show_default=False, |
| 95 | + ) |
| 96 | + @click.option( |
| 97 | + "--env-file", |
| 98 | + help="Path to .env file. If not provided, looks for .env in current directory.", |
| 99 | + default=None, |
| 100 | + type=click.Path(exists=True), |
| 101 | + show_default=False, |
| 102 | + ) |
| 103 | + @click.option( |
| 104 | + "--local-files-directory", |
| 105 | + "--local-dir", |
| 106 | + help="Directory (relative to the current working directory) where Humanloop files are stored locally (default: humanloop/).", |
| 107 | + default="humanloop", |
| 108 | + type=click.Path(), |
| 109 | + ) |
| 110 | + @click.option( |
| 111 | + "--base-url", |
| 112 | + default=None, |
| 113 | + hidden=True, |
| 114 | + ) |
| 115 | + @wraps(f) |
| 116 | + def wrapper(*args, **kwargs): |
| 117 | + return f(*args, **kwargs) |
| 118 | + |
| 119 | + return wrapper |
| 120 | + |
| 121 | + |
| 122 | +def handle_sync_errors(f: Callable) -> Callable: |
| 123 | + """Decorator for handling sync operation errors. |
| 124 | +
|
| 125 | + If an error occurs in any operation that uses this decorator, it will be logged and the program will exit with a non-zero exit code. |
| 126 | + """ |
| 127 | + |
| 128 | + @wraps(f) |
| 129 | + def wrapper(*args, **kwargs): |
| 130 | + try: |
| 131 | + return f(*args, **kwargs) |
| 132 | + except Exception as e: |
| 133 | + click.echo(click.style(str(f"Error: {e}"), fg=ERROR_COLOR)) |
| 134 | + sys.exit(1) |
| 135 | + |
| 136 | + return wrapper |
| 137 | + |
| 138 | + |
| 139 | +@click.group( |
| 140 | + help="Humanloop CLI for managing sync operations.", |
| 141 | + context_settings={ |
| 142 | + "help_option_names": ["-h", "--help"], |
| 143 | + "max_content_width": 100, |
| 144 | + }, |
| 145 | +) |
| 146 | +def cli(): # Does nothing because used as a group for other subcommands (pull, push, etc.) |
| 147 | + """Humanloop CLI for managing sync operations.""" |
| 148 | + pass |
| 149 | + |
| 150 | + |
| 151 | +@cli.command() |
| 152 | +@click.option( |
| 153 | + "--path", |
| 154 | + "-p", |
| 155 | + help="Path in the Humanloop workspace to pull from (file or directory). You can pull an entire directory (e.g. 'my/directory') " |
| 156 | + "or a specific file (e.g. 'my/directory/my_prompt.prompt'). When pulling a directory, all files within that directory and its subdirectories will be included. " |
| 157 | + "If not specified, pulls from the root of the remote workspace.", |
| 158 | + default=None, |
| 159 | +) |
| 160 | +@click.option( |
| 161 | + "--environment", |
| 162 | + "-e", |
| 163 | + help="Environment to pull from (e.g. 'production', 'staging')", |
| 164 | + default=None, |
| 165 | +) |
| 166 | +@click.option( |
| 167 | + "--verbose", |
| 168 | + "-v", |
| 169 | + is_flag=True, |
| 170 | + help="Show detailed information about the operation", |
| 171 | +) |
| 172 | +@click.option( |
| 173 | + "--quiet", |
| 174 | + "-q", |
| 175 | + is_flag=True, |
| 176 | + help="Suppress output of successful files", |
| 177 | +) |
| 178 | +@handle_sync_errors |
| 179 | +@common_options |
| 180 | +def pull( |
| 181 | + path: Optional[str], |
| 182 | + environment: Optional[str], |
| 183 | + api_key: Optional[str], |
| 184 | + env_file: Optional[str], |
| 185 | + local_files_directory: str, |
| 186 | + base_url: Optional[str], |
| 187 | + verbose: bool, |
| 188 | + quiet: bool, |
| 189 | +): |
| 190 | + """Pull Prompt and Agent files from Humanloop to your local filesystem. |
| 191 | +
|
| 192 | + \b |
| 193 | + This command will: |
| 194 | + 1. Fetch Prompt and Agent files from your Humanloop workspace |
| 195 | + 2. Save them to your local filesystem (directory specified by --local-files-directory, default: humanloop/) |
| 196 | + 3. Maintain the same directory structure as in Humanloop |
| 197 | + 4. Add appropriate file extensions (.prompt or .agent) |
| 198 | +
|
| 199 | + \b |
| 200 | + For example, with the default --local-files-directory=humanloop, files will be saved as: |
| 201 | + ./humanloop/ |
| 202 | + ├── my_project/ |
| 203 | + │ ├── prompts/ |
| 204 | + │ │ ├── my_prompt.prompt |
| 205 | + │ │ └── nested/ |
| 206 | + │ │ └── another_prompt.prompt |
| 207 | + │ └── agents/ |
| 208 | + │ └── my_agent.agent |
| 209 | + └── another_project/ |
| 210 | + └── prompts/ |
| 211 | + └── other_prompt.prompt |
| 212 | +
|
| 213 | + \b |
| 214 | + If you specify --local-files-directory=data/humanloop, files will be saved in ./data/humanloop/ instead. |
| 215 | +
|
| 216 | + If a file exists both locally and in the Humanloop workspace, the local file will be overwritten |
| 217 | + with the version from Humanloop. Files that only exist locally will not be affected. |
| 218 | +
|
| 219 | + Currently only supports syncing Prompt and Agent files. Other file types will be skipped.""" |
| 220 | + client = get_client(api_key, env_file, base_url) |
| 221 | + sync_client = SyncClient( |
| 222 | + client, base_dir=local_files_directory, log_level=logging.DEBUG if verbose else logging.WARNING |
| 223 | + ) |
| 224 | + |
| 225 | + click.echo(click.style("Pulling files from Humanloop...", fg=INFO_COLOR)) |
| 226 | + click.echo(click.style(f"Path: {path or '(root)'}", fg=INFO_COLOR)) |
| 227 | + click.echo(click.style(f"Environment: {environment or '(default)'}", fg=INFO_COLOR)) |
| 228 | + |
| 229 | + start_time = time.time() |
| 230 | + successful_files, failed_files = sync_client.pull(path, environment) |
| 231 | + duration_ms = int((time.time() - start_time) * 1000) |
| 232 | + |
| 233 | + # Determine if the operation was successful based on failed_files |
| 234 | + is_successful = not failed_files |
| 235 | + duration_color = SUCCESS_COLOR if is_successful else ERROR_COLOR |
| 236 | + click.echo(click.style(f"Pull completed in {duration_ms}ms", fg=duration_color)) |
| 237 | + |
| 238 | + if successful_files and not quiet: |
| 239 | + click.echo(click.style(f"\nSuccessfully pulled {len(successful_files)} files:", fg=SUCCESS_COLOR)) |
| 240 | + for file in successful_files: |
| 241 | + click.echo(click.style(f" ✓ {file}", fg=SUCCESS_COLOR)) |
| 242 | + |
| 243 | + if failed_files: |
| 244 | + click.echo(click.style(f"\nFailed to pull {len(failed_files)} files:", fg=ERROR_COLOR)) |
| 245 | + for file in failed_files: |
| 246 | + click.echo(click.style(f" ✗ {file}", fg=ERROR_COLOR)) |
| 247 | + |
| 248 | + |
| 249 | +if __name__ == "__main__": |
| 250 | + cli() |
0 commit comments