Coverage for llm_dataset_engine/cli/main.py: 32%

1"""

2 Main CLI entry point for LLM Dataset Engine.

3

4 Provides command-line interface for processing datasets, estimating costs,

5 and managing pipeline execution.

6"""

7

8 import sys

9 from pathlib import Path

10 from typing import Optional

11 from uuid import UUID

12

13 import click

14 from rich.console import Console

15 from rich.table import Table

16

17 from llm_dataset_engine import __version__

18 from llm_dataset_engine.api import Pipeline

19 from llm_dataset_engine.config import ConfigLoader

20 from llm_dataset_engine.core.specifications import (

21 DataSourceType,

22 LLMProvider,

23 PipelineSpecifications,

24)

25

26 console = Console()

27

28

29@click.group()

30@click.version_option(version=__version__, prog_name="llm-dataset")

31 def cli():

32 """

33 LLM Dataset Engine - Process tabular datasets using LLMs.

34

35 A production-grade SDK for processing CSV, Excel, and Parquet files

36 with Large Language Models, featuring reliability, cost control, and

37 observability.

38

39 Examples:

40

41 # Process a dataset

42 llm-dataset process --config config.yaml --input data.csv --output result.csv

43

44 # Estimate cost before processing

45 llm-dataset estimate --config config.yaml --input data.csv

46

47 # Resume from checkpoint

48 llm-dataset resume --session-id abc-123 --checkpoint-dir .checkpoints

49

50 # Validate configuration

51 llm-dataset validate --config config.yaml

52 """

53 pass

54

55

56@cli.command()

57@click.option(

58 "--config",

59 "-c",

60 required=True,

61 type=click.Path(exists=True, path_type=Path),

62 help="Path to YAML/JSON configuration file",

63)

64@click.option(

65 "--input",

66 "-i",

67 required=True,

68 type=click.Path(exists=True, path_type=Path),

69 help="Path to input data file (CSV, Excel, Parquet)",

70)

71@click.option(

72 "--output",

73 "-o",

74 required=True,

75 type=click.Path(path_type=Path),

76 help="Path to output file",

77)

78@click.option(

79 "--provider",

80 type=click.Choice(["openai", "azure_openai", "anthropic", "groq"]),

81 help="Override LLM provider from config",

82)

83@click.option(

84 "--model",

85 help="Override model name from config",

86)

87@click.option(

88 "--max-budget",

89 type=float,

90 help="Override maximum budget (USD) from config",

91)

92@click.option(

93 "--batch-size",

94 type=int,

95 help="Override batch size from config",

96)

97@click.option(

98 "--concurrency",

99 type=int,

100 help="Override concurrency from config",

101)

102@click.option(

103 "--checkpoint-dir",

104 type=click.Path(path_type=Path),

105 help="Override checkpoint directory from config",

106)

107@click.option(

108 "--dry-run",

109 is_flag=True,

110 help="Validate and estimate only, don't execute",

111)

112@click.option(

113 "--verbose",

114 "-v",

115 is_flag=True,

116 help="Enable verbose logging",

117)

118 def process(

119 config: Path,

120 input: Path,

121 output: Path,

122 provider: Optional[str],

123 model: Optional[str],

124 max_budget: Optional[float],

125 batch_size: Optional[int],

126 concurrency: Optional[int],

127 checkpoint_dir: Optional[Path],

128 dry_run: bool,

129 verbose: bool,

130):

131 """

132 Process a dataset using LLM transformations.

133

134 Reads data from INPUT file, applies LLM transformations according to CONFIG,

135 and writes results to OUTPUT file.

136

137 Examples:

138

139 # Basic usage

140 llm-dataset process -c config.yaml -i data.csv -o result.csv

141

142 # Override provider and model

143 llm-dataset process -c config.yaml -i data.csv -o result.csv \\

144 --provider groq --model openai/gpt-oss-120b

145

146 # Set budget limit

147 llm-dataset process -c config.yaml -i data.csv -o result.csv \\

148 --max-budget 10.0

149

150 # Dry run (estimate only)

151 llm-dataset process -c config.yaml -i data.csv -o result.csv --dry-run

152 """

153 try:

154 # Load configuration

155 console.print(f"[cyan]Loading configuration from {config}...[/cyan]")

156 specs = ConfigLoader.from_yaml(str(config))

157

158 # Override with CLI arguments

159 specs.dataset.source_path = input

160

161 # Set output configuration

162 if specs.output:

163 specs.output.destination_path = output

164 else:

165 # Create output spec if not in config

166 from llm_dataset_engine.core.specifications import OutputSpec, MergeStrategy

167

168 # Detect output type from extension

169 output_suffix = output.suffix.lower()

170 if output_suffix == ".csv":

171 output_type = DataSourceType.CSV

172 elif output_suffix in [".xlsx", ".xls"]:

173 output_type = DataSourceType.EXCEL

174 elif output_suffix == ".parquet":

175 output_type = DataSourceType.PARQUET

176 else:

177 output_type = DataSourceType.CSV # Default

178

179 specs.output = OutputSpec(

180 destination_type=output_type,

181 destination_path=output,

182 merge_strategy=MergeStrategy.REPLACE,

183 )

184

185 if provider:

186 specs.llm.provider = LLMProvider(provider)

187

188 if model:

189 specs.llm.model = model

190

191 if max_budget is not None:

192 from decimal import Decimal

193 specs.processing.max_budget = Decimal(str(max_budget))

194

195 if batch_size is not None:

196 specs.processing.batch_size = batch_size

197

198 if concurrency is not None:

199 specs.processing.concurrency = concurrency

200

201 if checkpoint_dir is not None:

202 specs.processing.checkpoint_dir = checkpoint_dir

203

204 # Create pipeline

205 console.print("[cyan]Creating pipeline...[/cyan]")

206 pipeline = Pipeline(specs)

207

208 # Validate

209 console.print("[cyan]Validating pipeline...[/cyan]")

210 validation = pipeline.validate()

211

212 if not validation.is_valid:

213 console.print("[red]❌ Validation failed:[/red]")

214 for error in validation.errors:

215 console.print(f" [red]• {error}[/red]")

216 sys.exit(1)

217

218 console.print("[green]✅ Validation passed[/green]")

219

220 # Estimate cost

221 console.print("\n[cyan]Estimating cost...[/cyan]")

222 estimate = pipeline.estimate_cost()

223

224 table = Table(title="Cost Estimate")

225 table.add_column("Metric", style="cyan")

226 table.add_column("Value", style="green")

227

228 table.add_row("Total Cost", f"${estimate.total_cost}")

229 table.add_row("Total Tokens", f"{estimate.total_tokens:,}")

230 table.add_row("Input Tokens", f"{estimate.input_tokens:,}")

231 table.add_row("Output Tokens", f"{estimate.output_tokens:,}")

232 table.add_row("Rows", f"{estimate.rows:,}")

233

234 console.print(table)

235

236 if dry_run:

237 console.print("\n[yellow]Dry run mode - skipping execution[/yellow]")

238 return

239

240 # Execute

241 console.print("\n[cyan]Processing dataset...[/cyan]")

242 result = pipeline.execute()

243

244 # Display results

245 console.print("\n[green]✅ Processing complete![/green]")

246

247 results_table = Table(title="Execution Results")

248 results_table.add_column("Metric", style="cyan")

249 results_table.add_column("Value", style="green")

250

251 results_table.add_row("Total Rows", str(result.metrics.total_rows))

252 results_table.add_row("Processed", str(result.metrics.processed_rows))

253 results_table.add_row("Failed", str(result.metrics.failed_rows))

254 results_table.add_row("Skipped", str(result.metrics.skipped_rows))

255 results_table.add_row("Duration", f"{result.duration:.2f}s")

256 results_table.add_row("Total Cost", f"${result.costs.total_cost}")

257 results_table.add_row("Cost per Row", f"${result.costs.total_cost / result.metrics.total_rows:.6f}")

258

259 console.print(results_table)

260

261 console.print(f"\n[green]Output written to: {output}[/green]")

262

263 except Exception as e:

264 console.print(f"[red]❌ Error: {e}[/red]")

265 if verbose:

266 console.print_exception()

267 sys.exit(1)

268

269

270@cli.command()

271@click.option(

272 "--config",

273 "-c",

274 required=True,

275 type=click.Path(exists=True, path_type=Path),

276 help="Path to YAML/JSON configuration file",

277)

278@click.option(

279 "--input",

280 "-i",

281 required=True,

282 type=click.Path(exists=True, path_type=Path),

283 help="Path to input data file",

284)

285@click.option(

286 "--provider",

287 type=click.Choice(["openai", "azure_openai", "anthropic", "groq"]),

288 help="Override LLM provider from config",

289)

290@click.option(

291 "--model",

292 help="Override model name from config",

293)

294 def estimate(

295 config: Path,

296 input: Path,

297 provider: Optional[str],

298 model: Optional[str],

299):

300 """

301 Estimate processing cost without executing.

302

303 Useful for budget planning and cost validation before running

304 expensive operations.

305

306 Examples:

307

308 # Estimate cost

309 llm-dataset estimate -c config.yaml -i data.csv

310

311 # Estimate with different model

312 llm-dataset estimate -c config.yaml -i data.csv --model gpt-4o

313 """

314 try:

315 # Load configuration

316 console.print(f"[cyan]Loading configuration from {config}...[/cyan]")

317 specs = ConfigLoader.from_yaml(str(config))

318

319 # Override

320 specs.dataset.source_path = input

321

322 if provider:

323 specs.llm.provider = LLMProvider(provider)

324

325 if model:

326 specs.llm.model = model

327

328 # Create pipeline

329 pipeline = Pipeline(specs)

330

331 # Validate

332 validation = pipeline.validate()

333 if not validation.is_valid:

334 console.print("[red]❌ Validation failed:[/red]")

335 for error in validation.errors:

336 console.print(f" [red]• {error}[/red]")

337 sys.exit(1)

338

339 # Estimate

340 console.print("[cyan]Estimating cost...[/cyan]")

341 estimate = pipeline.estimate_cost()

342

343 # Display results

344 table = Table(title="Cost Estimate", show_header=True)

345 table.add_column("Metric", style="cyan", width=20)

346 table.add_column("Value", style="green", width=20)

347

348 table.add_row("Total Cost", f"${estimate.total_cost}")

349 table.add_row("Total Tokens", f"{estimate.total_tokens:,}")

350 table.add_row("Input Tokens", f"{estimate.input_tokens:,}")

351 table.add_row("Output Tokens", f"{estimate.output_tokens:,}")

352 table.add_row("Rows to Process", f"{estimate.rows:,}")

353 table.add_row("Confidence", estimate.confidence)

354

355 console.print("\n")

356 console.print(table)

357

358 # Cost per row

359 if estimate.rows > 0:

360 cost_per_row = estimate.total_cost / estimate.rows

361 console.print(f"\n[cyan]Cost per row: ${cost_per_row:.6f}[/cyan]")

362

363 # Warning if expensive

364 if estimate.total_cost > 10.0:

365 console.print(f"\n[yellow]⚠️ Warning: Estimated cost (${estimate.total_cost}) exceeds $10[/yellow]")

366

367 except Exception as e:

368 console.print(f"[red]❌ Error: {e}[/red]")

369 sys.exit(1)

370

371

372@cli.command()

373@click.option(

374 "--session-id",

375 "-s",

376 required=True,

377 help="Session ID to resume (UUID)",

378)

379@click.option(

380 "--checkpoint-dir",

381 type=click.Path(exists=True, path_type=Path),

382 default=".checkpoints",

383 help="Checkpoint directory (default: .checkpoints)",

384)

385@click.option(

386 "--output",

387 "-o",

388 type=click.Path(path_type=Path),

389 help="Override output path",

390)

391 def resume(

392 session_id: str,

393 checkpoint_dir: Path,

394 output: Optional[Path],

395):

396 """

397 Resume pipeline execution from checkpoint.

398

399 Useful for recovering from failures or continuing interrupted processing.

400

401 Examples:

402

403 # Resume from checkpoint

404 llm-dataset resume --session-id abc-123-def

405

406 # Resume with custom checkpoint directory

407 llm-dataset resume -s abc-123 --checkpoint-dir /path/to/checkpoints

408 """

409 try:

410 from llm_dataset_engine.adapters import LocalFileCheckpointStorage

411 from llm_dataset_engine.orchestration import StateManager

412

413 # Load checkpoint

414 console.print(f"[cyan]Looking for checkpoint in {checkpoint_dir}...[/cyan]")

415

416 storage = LocalFileCheckpointStorage(str(checkpoint_dir))

417 state_manager = StateManager(storage)

418

419 session_uuid = UUID(session_id)

420

421 if not state_manager.can_resume(session_uuid):

422 console.print(f"[red]❌ No checkpoint found for session {session_id}[/red]")

423 console.print(f"[yellow]Check checkpoint directory: {checkpoint_dir}[/yellow]")

424 sys.exit(1)

425

426 # Load checkpoint

427 checkpoint_info = state_manager.get_latest_checkpoint(session_uuid)

428 console.print(f"[green]✅ Found checkpoint at row {checkpoint_info.row_index}[/green]")

429

430 # Resume execution

431 console.print("[cyan]Resuming execution...[/cyan]")

432

433 # Note: Full resume implementation would load the original pipeline

434 # and continue from checkpoint. For now, we show the checkpoint info.

435 console.print("\n[yellow]⚠️ Full resume functionality requires the original pipeline configuration[/yellow]")

436 console.print("[yellow]Please use Pipeline.execute(resume_from=session_id) in Python code[/yellow]")

437

438 # Display checkpoint info

439 table = Table(title="Checkpoint Information")

440 table.add_column("Property", style="cyan")

441 table.add_column("Value", style="green")

442

443 table.add_row("Session ID", str(checkpoint_info.session_id))

444 table.add_row("Checkpoint Path", checkpoint_info.checkpoint_path)

445 table.add_row("Last Row", str(checkpoint_info.row_index))

446 table.add_row("Last Stage", str(checkpoint_info.stage_index))

447 table.add_row("Timestamp", str(checkpoint_info.timestamp))

448 table.add_row("Size", f"{checkpoint_info.size_bytes:,} bytes")

449

450 console.print(table)

451

452 except ValueError:

453 console.print(f"[red]❌ Invalid session ID format: {session_id}[/red]")

454 console.print("[yellow]Session ID should be a UUID (e.g., abc-123-def-456)[/yellow]")

455 sys.exit(1)

456 except Exception as e:

457 console.print(f"[red]❌ Error: {e}[/red]")

458 sys.exit(1)

459

460

461@cli.command()

462@click.option(

463 "--config",

464 "-c",

465 required=True,

466 type=click.Path(exists=True, path_type=Path),

467 help="Path to YAML/JSON configuration file",

468)

469@click.option(

470 "--verbose",

471 "-v",

472 is_flag=True,

473 help="Show detailed validation results",

474)

475 def validate(config: Path, verbose: bool):

476 """

477 Validate pipeline configuration.

478

479 Checks configuration file for errors and warnings without executing

480 the pipeline.

481

482 Examples:

483

484 # Validate configuration

485 llm-dataset validate -c config.yaml

486

487 # Verbose validation

488 llm-dataset validate -c config.yaml --verbose

489 """

490 try:

491 # Load configuration

492 console.print(f"[cyan]Loading configuration from {config}...[/cyan]")

493 specs = ConfigLoader.from_yaml(str(config))

494

495 console.print("[green]✅ Configuration loaded successfully[/green]")

496

497 # Display configuration summary

498 if verbose:

499 table = Table(title="Configuration Summary")

500 table.add_column("Component", style="cyan")

501 table.add_column("Details", style="green")

502

503 table.add_row("Dataset", f"{specs.dataset.source_type.value}")

504 table.add_row("Input Columns", ", ".join(specs.dataset.input_columns))

505 table.add_row("Output Columns", ", ".join(specs.dataset.output_columns))

506 table.add_row("LLM Provider", specs.llm.provider.value)

507 table.add_row("Model", specs.llm.model)

508 table.add_row("Batch Size", str(specs.processing.batch_size))

509 table.add_row("Concurrency", str(specs.processing.concurrency))

510

511 if specs.processing.max_budget:

512 table.add_row("Max Budget", f"${specs.processing.max_budget}")

513

514 console.print("\n")

515 console.print(table)

516

517 # Create pipeline for validation

518 console.print("\n[cyan]Validating pipeline...[/cyan]")

519 pipeline = Pipeline(specs)

520 validation = pipeline.validate()

521

522 if validation.is_valid:

523 console.print("[green]✅ Pipeline configuration is valid[/green]")

524

525 if validation.warnings:

526 console.print("\n[yellow]Warnings:[/yellow]")

527 for warning in validation.warnings:

528 console.print(f" [yellow]• {warning}[/yellow]")

529 else:

530 console.print("[red]❌ Pipeline configuration is invalid[/red]")

531 console.print("\n[red]Errors:[/red]")

532 for error in validation.errors:

533 console.print(f" [red]• {error}[/red]")

534

535 if validation.warnings:

536 console.print("\n[yellow]Warnings:[/yellow]")

537 for warning in validation.warnings:

538 console.print(f" [yellow]• {warning}[/yellow]")

539

540 sys.exit(1)

541

542 except Exception as e:

543 console.print(f"[red]❌ Error: {e}[/red]")

544 sys.exit(1)

545

546

547@cli.command()

548@click.option(

549 "--checkpoint-dir",

550 type=click.Path(exists=True, path_type=Path),

551 default=".checkpoints",

552 help="Checkpoint directory to list (default: .checkpoints)",

553)

554 def list_checkpoints(checkpoint_dir: Path):

555 """

556 List available checkpoints.

557

558 Shows all saved checkpoints in the specified directory.

559

560 Examples:

561

562 # List checkpoints

563 llm-dataset list-checkpoints

564

565 # List from custom directory

566 llm-dataset list-checkpoints --checkpoint-dir /path/to/checkpoints

567 """

568 try:

569 from llm_dataset_engine.adapters import LocalFileCheckpointStorage

570

571 console.print(f"[cyan]Scanning {checkpoint_dir} for checkpoints...[/cyan]")

572

573 storage = LocalFileCheckpointStorage(checkpoint_dir)

574 checkpoints = storage.list_checkpoints()

575

576 if not checkpoints:

577 console.print("[yellow]No checkpoints found[/yellow]")

578 return

579

580 # Display checkpoints

581 table = Table(title=f"Checkpoints in {checkpoint_dir}")

582 table.add_column("Session ID", style="cyan")

583 table.add_column("Row", style="green")

584 table.add_column("Stage", style="green")

585 table.add_column("Timestamp", style="yellow")

586 table.add_column("Size", style="magenta")

587

588 for cp in checkpoints:

589 table.add_row(

590 str(cp.session_id)[:8] + "...",

591 str(cp.row_index),

592 str(cp.stage_index),

593 cp.timestamp.strftime("%Y-%m-%d %H:%M:%S"),

594 f"{cp.size_bytes:,} bytes",

595 )

596

597 console.print("\n")

598 console.print(table)

599 console.print(f"\n[cyan]Total checkpoints: {len(checkpoints)}[/cyan]")

600

601 except Exception as e:

602 console.print(f"[red]❌ Error: {e}[/red]")

603 sys.exit(1)

604

605

606@cli.command()

607@click.option(

608 "--input",

609 "-i",

610 required=True,

611 type=click.Path(exists=True, path_type=Path),

612 help="Path to input file to inspect",

613)

614@click.option(

615 "--head",

616 type=int,

617 default=5,

618 help="Number of rows to show (default: 5)",

619)

620 def inspect(input: Path, head: int):

621 """

622 Inspect input data file.

623

624 Shows file info and preview of first N rows.

625

626 Examples:

627

628 # Inspect CSV file

629 llm-dataset inspect -i data.csv

630

631 # Show first 10 rows

632 llm-dataset inspect -i data.csv --head 10

633 """

634 try:

635 import pandas as pd

636

637 console.print(f"[cyan]Inspecting {input}...[/cyan]")

638

639 # Detect file type

640 suffix = input.suffix.lower()

641

642 if suffix == ".csv":

643 df = pd.read_csv(input)

644 elif suffix in [".xlsx", ".xls"]:

645 df = pd.read_excel(input)

646 elif suffix == ".parquet":

647 df = pd.read_parquet(input)

648 else:

649 console.print(f"[red]❌ Unsupported file type: {suffix}[/red]")

650 sys.exit(1)

651

652 # File info

653 info_table = Table(title="File Information")

654 info_table.add_column("Property", style="cyan")

655 info_table.add_column("Value", style="green")

656

657 info_table.add_row("File Path", str(input))

658 info_table.add_row("File Type", suffix[1:].upper())

659 info_table.add_row("Total Rows", f"{len(df):,}")

660 info_table.add_row("Total Columns", str(len(df.columns)))

661 info_table.add_row("Memory Usage", f"{df.memory_usage(deep=True).sum() / 1024 / 1024:.2f} MB")

662

663 console.print("\n")

664 console.print(info_table)

665

666 # Columns

667 console.print("\n[cyan]Columns:[/cyan]")

668 for col in df.columns:

669 dtype = df[col].dtype

670 null_count = df[col].isnull().sum()

671 console.print(f" • {col} ({dtype}) - {null_count} nulls")

672

673 # Preview

674 console.print(f"\n[cyan]First {head} rows:[/cyan]")

675 console.print(df.head(head).to_string())

676

677 except Exception as e:

678 console.print(f"[red]❌ Error: {e}[/red]")

679 sys.exit(1)

680

681

682 if __name__ == "__main__":

683 cli()

684