local-deep-research/examples/optimization/run_optimization.py

#!/usr/bin/env python
"""
Parameter Optimization Runner for Local Deep Research.

This script provides a convenient way to run hyperparameter optimization.

Usage:
    # Install dependencies with PDM
    cd /path/to/local-deep-research
    pdm install

    # Run the script with PDM
    pdm run python examples/optimization/run_optimization.py --help
"""

import argparse
import json
import os
import sys
from datetime import datetime

# Add the src directory to the Python path
project_root = os.path.abspath(
    os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
)
sys.path.insert(0, os.path.join(project_root, "src"))

# Import the optimization functionality
from local_deep_research.benchmarks.optimization import (
    optimize_for_efficiency,
    optimize_for_quality,
    optimize_for_speed,
    optimize_parameters,
)


def main():
    """Run parameter optimization with command-line arguments."""
    parser = argparse.ArgumentParser(
        description="Run parameter optimization for Local Deep Research"
    )
    parser.add_argument("query", help="Research query to optimize for")
    parser.add_argument(
        "--output-dir",
        default=os.path.join("examples", "optimization", "results"),
        help="Directory to save results",
    )
    parser.add_argument("--search-tool", default="searxng", help="Search tool to use")

    # LLM configuration options
    parser.add_argument(
        "--model", help="Model name for the LLM (e.g., 'claude-3-sonnet-20240229')"
    )
    parser.add_argument(
        "--provider",
        help="Provider for the LLM (e.g., 'anthropic', 'openai', 'openai_endpoint')",
    )
    parser.add_argument(
        "--endpoint-url",
        help="Custom endpoint URL (e.g., 'https://openrouter.ai/api/v1')",
    )
    parser.add_argument("--api-key", help="API key for the LLM provider")
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.7,
        help="Temperature for the LLM (default: 0.7)",
    )

    parser.add_argument(
        "--trials", type=int, default=30, help="Number of parameter combinations to try"
    )
    parser.add_argument(
        "--mode",
        choices=["balanced", "speed", "quality", "efficiency"],
        default="balanced",
        help="Optimization mode",
    )
    parser.add_argument(
        "--weights",
        help='Custom weights as JSON string, e.g., \'{"quality": 0.7, "speed": 0.3}\'',
    )

    args = parser.parse_args()

    # Create timestamp for unique output directory
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir = os.path.join(args.output_dir, f"opt_{timestamp}")
    os.makedirs(output_dir, exist_ok=True)

    print(
        f"Starting optimization ({args.mode} mode) - results will be saved to {output_dir}"
    )

    # Parse custom weights if provided
    custom_weights = None
    if args.weights:
        try:
            custom_weights = json.loads(args.weights)
        except json.JSONDecodeError:
            print("Error parsing weights JSON. Using default weights.")

    # Set environment variables for the API key and endpoint URL if provided
    if args.api_key:
        os.environ["OPENAI_ENDPOINT_API_KEY"] = args.api_key
        os.environ["LDR_LLM__OPENAI_ENDPOINT_API_KEY"] = args.api_key

    if args.endpoint_url:
        os.environ["OPENAI_ENDPOINT_URL"] = args.endpoint_url
        os.environ["LDR_LLM__OPENAI_ENDPOINT_URL"] = args.endpoint_url

    if args.model:
        os.environ["LDR_LLM__MODEL"] = args.model

    if args.provider:
        os.environ["LDR_LLM__PROVIDER"] = args.provider

    # Run optimization based on mode
    if args.mode == "speed":
        best_params, best_score = optimize_for_speed(
            query=args.query,
            search_tool=args.search_tool,
            n_trials=args.trials,
            model_name=args.model,
            provider=args.provider,
            openai_endpoint_url=args.endpoint_url,
            temperature=args.temperature,
            api_key=args.api_key,
            output_dir=output_dir,
        )
    elif args.mode == "quality":
        best_params, best_score = optimize_for_quality(
            query=args.query,
            search_tool=args.search_tool,
            n_trials=args.trials,
            model_name=args.model,
            provider=args.provider,
            openai_endpoint_url=args.endpoint_url,
            temperature=args.temperature,
            api_key=args.api_key,
            output_dir=output_dir,
        )
    elif args.mode == "efficiency":
        best_params, best_score = optimize_for_efficiency(
            query=args.query,
            search_tool=args.search_tool,
            n_trials=args.trials,
            model_name=args.model,
            provider=args.provider,
            openai_endpoint_url=args.endpoint_url,
            temperature=args.temperature,
            api_key=args.api_key,
            output_dir=output_dir,
        )
    else:  # balanced
        best_params, best_score = optimize_parameters(
            query=args.query,
            search_tool=args.search_tool,
            n_trials=args.trials,
            model_name=args.model,
            provider=args.provider,
            openai_endpoint_url=args.endpoint_url,
            temperature=args.temperature,
            api_key=args.api_key,
            output_dir=output_dir,
            metric_weights=custom_weights,
        )

    print(f"\nOptimization complete! Results saved to {output_dir}")
    print(f"Best parameters: {best_params}")
    print(f"Best score: {best_score:.4f}")

    # Save summary to a JSON file
    summary = {
        "timestamp": timestamp,
        "query": args.query,
        "mode": args.mode,
        "trials": args.trials,
        "search_tool": args.search_tool,
        "model": args.model,
        "provider": args.provider,
        "temperature": args.temperature,
        "best_parameters": best_params,
        "best_score": best_score,
        "custom_weights": custom_weights,
    }

    with open(os.path.join(output_dir, "optimization_summary.json"), "w") as f:
        json.dump(summary, f, indent=2)

    return 0


if __name__ == "__main__":
    sys.exit(main())