run_training.py

#!/usr/bin/env python3
"""
Training script for HealthFlow system.
Processes training data from JSONL files and saves experiences.
"""

import asyncio
import json
import sys
from pathlib import Path
from typing import List, Dict, Any
import argparse
from dataclasses import dataclass
import shutil

import typer
from rich.console import Console
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeElapsedColumn
from rich.table import Table
from loguru import logger

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent))

from healthflow.system import HealthFlowSystem
from healthflow.core.config import get_config, setup_logging

console = Console()

@dataclass
class TrainingExample:
    """Represents a single training example."""
    qid: str
    task: str
    answer: str

class TrainingRunner:
    """Handles the training process for HealthFlow."""

    def __init__(self, system: HealthFlowSystem, experience_path: Path):
        self.system = system
        self.experience_path = experience_path
        self.results: List[Dict[str, Any]] = []

    async def run_training(self, training_file: Path) -> Dict[str, Any]:
        """Run training on all examples in the training file."""
        training_examples = self._load_training_data(training_file)

        console.print(Panel(
            f"[bold cyan]Starting HealthFlow Training[/bold cyan]\n\n"
            f"[dim]Training file:[/dim] {training_file}\n"
            f"[dim]Total examples:[/dim] {len(training_examples)}\n"
            f"[dim]Experience path:[/dim] {self.experience_path}",
            border_style="cyan"
        ))

        successful_tasks = 0
        total_tasks = len(training_examples)

        with Progress(
            SpinnerColumn(),
            TextColumn("[bold blue]{task.description}"),
            BarColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            TimeElapsedColumn(),
            console=console
        ) as progress:
            main_task = progress.add_task("Training Progress", total=total_tasks)

            for i, example in enumerate(training_examples):
                progress.update(main_task, description=f"Processing example {i+1}/{total_tasks} (QID: {example.qid})")

                try:
                    result = await self.system.run_task(
                        user_request=example.task,
                    )

                    if result.get("success", False):
                        successful_tasks += 1

                    # Store result with training metadata
                    training_result = {
                        "qid": example.qid,
                        "task": example.task,
                        "reference_answer": example.answer,
                        "result": result,
                        "success": result.get("success", False),
                        "score": self._extract_score_from_result(result),
                        "workspace_path": result.get("workspace_path"),
                        "backend": result.get("backend"),
                        "planner_model": result.get("planner_model"),
                        "runtime_llm_keys": result.get("runtime_llm_keys"),
                        "memory_write_policy": result.get("memory_write_policy"),
                        "evaluation_status": result.get("evaluation_status"),
                        "evaluation_score": result.get("evaluation_score"),
                        "execution_time": result.get("execution_time", 0.0),
                        "sandbox_path": result.get("sandbox_path"),
                        "runtime_path": result.get("runtime_path"),
                        "task_state_path": result.get("task_state_path"),
                        "runtime_index_path": result.get("runtime_index_path"),
                        "run_summary_path": result.get("run_summary_path"),
                        "run_trajectory_path": result.get("run_trajectory_path"),
                        "run_costs_path": result.get("run_costs_path"),
                        "final_evaluation_path": result.get("final_evaluation_path"),
                        "report_requested": result.get("report_requested", False),
                        "report_generated": result.get("report_generated", False),
                        "report_path": result.get("report_path"),
                        "report_error": result.get("report_error"),
                    }
                    self.results.append(training_result)

                    logger.info(f"Training example {example.qid} completed. Success: {result.get('success', False)}")

                except Exception as e:
                    logger.error(f"Error processing training example {example.qid}: {e}")
                    error_result = {
                        "qid": example.qid,
                        "task": example.task,
                        "reference_answer": example.answer,
                        "result": {
                            "success": False,
                            "error": str(e),
                            "report_requested": False,
                            "report_generated": False,
                            "report_path": None,
                            "report_error": None,
                        },
                        "success": False,
                        "score": 0.0,
                        "workspace_path": None,
                        "backend": None,
                        "planner_model": None,
                        "runtime_llm_keys": None,
                        "memory_write_policy": None,
                        "evaluation_status": "failed",
                        "evaluation_score": 0.0,
                        "execution_time": 0.0,
                        "sandbox_path": None,
                        "runtime_path": None,
                        "task_state_path": None,
                        "runtime_index_path": None,
                        "run_summary_path": None,
                        "run_trajectory_path": None,
                        "run_costs_path": None,
                        "final_evaluation_path": None,
                        "report_requested": False,
                        "report_generated": False,
                        "report_path": None,
                        "report_error": None,
                    }
                    self.results.append(error_result)

                progress.advance(main_task)

        # Calculate summary statistics
        summary = self._calculate_summary(successful_tasks, total_tasks)
        self._display_summary(summary)

        return summary

    def _load_training_data(self, training_file: Path) -> List[TrainingExample]:
        """Load training examples from JSONL file."""
        examples = []

        if not training_file.exists():
            raise FileNotFoundError(f"Training file not found: {training_file}")

        with open(training_file, 'r', encoding='utf-8') as f:
            for line_num, line in enumerate(f, 1):
                line = line.strip()
                if not line:
                    continue

                try:
                    data = json.loads(line)
                    # Validate required keys
                    required_keys = ["qid", "task"]
                    missing_keys = [key for key in required_keys if key not in data]
                    if missing_keys:
                        logger.warning(f"Skipping line {line_num}: missing keys {missing_keys}")
                        continue
                    reference_answer = data.get("reference_answer", data.get("answer"))
                    if reference_answer is None:
                        logger.warning(f"Skipping line {line_num}: missing answer/reference_answer field")
                        continue

                    examples.append(TrainingExample(
                        qid=str(data["qid"]),
                        task=data["task"],
                        answer=reference_answer
                    ))

                except json.JSONDecodeError as e:
                    logger.warning(f"Skipping invalid JSON on line {line_num}: {e}")
                    continue
                except Exception as e:
                    logger.warning(f"Error processing line {line_num}: {e}")
                    continue

        if not examples:
            raise ValueError(f"No valid training examples found in {training_file}")

        logger.info(f"Loaded {len(examples)} training examples from {training_file}")
        return examples

    def _extract_score_from_result(self, result: Dict[str, Any]) -> float:
        """Extract evaluation score from result."""
        try:
            trajectory_path = result.get("run_trajectory_path")
            if trajectory_path:
                history_file = Path(trajectory_path)
            else:
                workspace_path = result.get("workspace_path")
                if not workspace_path:
                    return 0.0
                history_file = Path(workspace_path) / "runtime" / "run" / "trajectory.json"
            if history_file.exists():
                with open(history_file, 'r') as f:
                    history = json.loads(f.read())
                    attempts = history.get("attempts", [])
                    if attempts and "evaluation" in attempts[-1]:
                        return attempts[-1]["evaluation"].get("score", 0.0)
            return 0.0
        except (KeyError, IndexError, TypeError, json.JSONDecodeError, FileNotFoundError):
            return 0.0

    def _calculate_summary(self, successful_tasks: int, total_tasks: int) -> Dict[str, Any]:
        """Calculate training summary statistics."""
        success_rate = (successful_tasks / total_tasks) * 100 if total_tasks > 0 else 0

        # Calculate average score
        scores = [r["score"] for r in self.results if isinstance(r["score"], (int, float))]
        avg_score = sum(scores) / len(scores) if scores else 0.0
        avg_execution_time = (
            sum(r.get("execution_time", 0.0) for r in self.results) / len(self.results)
            if self.results
            else 0.0
        )
        evaluator_success_rate = (
            sum(1 for r in self.results if r.get("evaluation_status") == "success") / total_tasks * 100
            if total_tasks > 0
            else 0.0
        )

        summary = {
            "total_examples": total_tasks,
            "successful_examples": successful_tasks,
            "failed_examples": total_tasks - successful_tasks,
            "success_rate": success_rate,
            "average_score": avg_score,
            "average_execution_time": avg_execution_time,
            "evaluator_success_rate": evaluator_success_rate,
            "results": self.results
        }

        return summary

    def _display_summary(self, summary: Dict[str, Any]):
        """Display training summary."""
        table = Table(title="Training Results Summary")
        table.add_column("Metric", style="cyan")
        table.add_column("Value", style="green")

        table.add_row("Total Examples", str(summary["total_examples"]))
        table.add_row("Successful", str(summary["successful_examples"]))
        table.add_row("Failed", str(summary["failed_examples"]))
        table.add_row("Success Rate", f"{summary['success_rate']:.1f}%")
        table.add_row("Average Score", f"{summary['average_score']:.2f}/10.0")
        table.add_row("Evaluator Success Rate", f"{summary['evaluator_success_rate']:.1f}%")
        table.add_row("Average Execution Time", f"{summary['average_execution_time']:.2f}s")

        console.print("\n")
        console.print(table)

    def save_results(self, dataset_name: str, runtime_label: str):
        """Save training results using the benchmark-style directory structure."""
        # Create main results directory using the same structure as benchmark
        results_dir = Path("benchmark_results") / dataset_name / runtime_label
        results_dir.mkdir(parents=True, exist_ok=True)

        # Save individual results in qid directories (same as benchmark)
        for result in self.results:
            qid = result["qid"]
            qid_dir = results_dir / str(qid)
            qid_dir.mkdir(parents=True, exist_ok=True)

            # Save individual result JSON
            with open(qid_dir / "training_result.json", 'w', encoding='utf-8') as f:
                json.dump(result, f, indent=2)

            # Copy workspace files if available
            workspace_path = result.get("workspace_path")
            if workspace_path and Path(workspace_path).exists():
                self._copy_workspace_files(Path(workspace_path), qid_dir)

        # Save aggregated results
        results_file = results_dir / "training_results.jsonl"
        with open(results_file, 'w', encoding='utf-8') as f:
            for result in self.results:
                f.write(json.dumps(result, ensure_ascii=False) + '\n')

        console.print(f"\n[green]Training results saved to: {results_file}[/green]")
        logger.info(f"Training results saved to {results_file}")

        return results_dir

    def _copy_workspace_files(self, workspace_path: Path, output_dir: Path):
        """Copy generated files from the workspace to the output directory."""
        if not workspace_path.exists():
            return

        try:
            for item in workspace_path.iterdir():
                if item.is_file():
                    shutil.copy2(item, output_dir / item.name)
                elif item.is_dir():
                    shutil.copytree(item, output_dir / item.name, dirs_exist_ok=True)
        except Exception as e:
            logger.warning(f"Failed to copy workspace files: {e}")

def _runtime_output_label(config) -> str:
    runtime_llm_keys = config.runtime_llm_keys
    if len(set(runtime_llm_keys.values())) == 1:
        return next(iter(runtime_llm_keys.values())).replace("/", "__")
    return "__".join(
        f"{role}={llm_key.replace('/', '__')}"
        for role, llm_key in runtime_llm_keys.items()
    )


def _initialize_system(
    config_path: Path,
    experience_path: Path,
    planner_llm: str | None,
    evaluator_llm: str | None,
    reflector_llm: str | None,
    executor_llm: str | None,
    active_executor: str | None,
) -> HealthFlowSystem:
    """Initialize the HealthFlow system."""
    try:
        config = get_config(
            config_path,
            planner_llm=planner_llm,
            evaluator_llm=evaluator_llm,
            reflector_llm=reflector_llm,
            executor_llm=executor_llm,
            active_executor=active_executor,
        )
        setup_logging(config)
        return HealthFlowSystem(
            config=config,
            experience_path=experience_path
        )
    except (ValueError, FileNotFoundError) as e:
        console.print(Panel(f"[bold red]Initialization Error:[/bold red] {e}", title="Error", border_style="red"))
        raise typer.Exit(code=1)

async def main_async(
    training_file: Path,
    dataset_name: str,
    config_path: Path,
    experience_path: Path,
    planner_llm: str | None = None,
    evaluator_llm: str | None = None,
    reflector_llm: str | None = None,
    executor_llm: str | None = None,
    active_executor: str | None = None,
):
    """Main async function to run training."""
    system = _initialize_system(
        config_path,
        experience_path,
        planner_llm,
        evaluator_llm,
        reflector_llm,
        executor_llm,
        active_executor,
    )

    trainer = TrainingRunner(system, experience_path)
    summary = await trainer.run_training(training_file)
    results_dir = trainer.save_results(dataset_name, _runtime_output_label(system.config))

    # Create summary file (similar to benchmark)
    summary_data = {
        "dataset_name": dataset_name,
        "total_examples": summary["total_examples"],
        "successful_examples": summary["successful_examples"],
        "failed_examples": summary["failed_examples"],
        "success_rate": summary["success_rate"],
        "average_score": summary["average_score"],
        "average_execution_time": summary["average_execution_time"],
        "evaluator_success_rate": summary["evaluator_success_rate"],
        "backend": system.config.active_executor_name,
        "planner_model": system.config.planner_llm.model_name,
        "runtime_llm_keys": system.config.runtime_llm_keys,
        "memory_write_policy": system.config.memory.write_policy,
        "results_directory": str(results_dir)
    }

    with open(results_dir / "training_summary.json", "w", encoding="utf-8") as f:
        json.dump(summary_data, f, indent=2)

    # Exit with non-zero code if success rate is too low
    if summary["success_rate"] < 50:
        console.print(f"\n[bold red]Training completed with low success rate: {summary['success_rate']:.1f}%[/bold red]")
        raise typer.Exit(code=1)
    else:
        console.print(f"\n[bold green]Training completed successfully! Success rate: {summary['success_rate']:.1f}%[/bold green]")

def main():
    """Main entry point for the training script."""
    parser = argparse.ArgumentParser(description="Run HealthFlow training on a dataset")
    parser.add_argument("training_file", type=Path, help="Path to the training JSONL file")
    parser.add_argument("dataset_name", help="Name of the dataset (used for output directory structure)")
    parser.add_argument("--config", "-c", type=Path, default="config.toml", help="Path to the configuration file")
    parser.add_argument("--experience-path", type=Path, default="workspace/memory/experience.jsonl", help="Path to the experience knowledge base file")
    parser.add_argument("--planner-llm", default=None, help="Override runtime.planner_llm from config.toml")
    parser.add_argument("--evaluator-llm", default=None, help="Override runtime.evaluator_llm from config.toml")
    parser.add_argument("--reflector-llm", default=None, help="Override runtime.reflector_llm from config.toml")
    parser.add_argument("--executor-llm", default=None, help="Override runtime.executor_llm from config.toml")
    parser.add_argument("--active-executor", default=None, help="The executor backend to use (e.g., claude_code, opencode, pi)")

    args = parser.parse_args()

    try:
        asyncio.run(main_async(
            training_file=args.training_file,
            dataset_name=args.dataset_name,
            config_path=args.config,
            experience_path=args.experience_path,
            planner_llm=args.planner_llm,
            evaluator_llm=args.evaluator_llm,
            reflector_llm=args.reflector_llm,
            executor_llm=args.executor_llm,
            active_executor=args.active_executor,
        ))
    except KeyboardInterrupt:
        console.print("\n[yellow]Training interrupted by user.[/yellow]")
        sys.exit(1)
    except Exception as e:
        console.print(f"\n[bold red]Training failed: {e}[/bold red]")
        sys.exit(1)

if __name__ == "__main__":
    main()