import argparse import json from geo_bot import GeoBot from benchmark import MapGuesserBenchmark from data_collector import DataCollector from config import MODELS_CONFIG, get_data_paths, SUCCESS_THRESHOLD_KM, get_model_class def agent_mode( model_name: str, steps: int, headless: bool, samples: int, dataset_name: str = "default", temperature: float = 0.0, ): """ Runs the AI Agent in a benchmark loop over multiple samples, using multi-step exploration for each. """ print( f"Starting Agent Mode: model={model_name}, steps={steps}, samples={samples}, dataset={dataset_name}, temperature={temperature}" ) data_paths = get_data_paths(dataset_name) try: with open(data_paths["golden_labels"], "r", encoding="utf-8") as f: golden_labels = json.load(f).get("samples", []) except FileNotFoundError: print( f"Error: Dataset '{dataset_name}' not found at {data_paths['golden_labels']}." ) return if not golden_labels: print(f"Error: No samples found in dataset '{dataset_name}'.") return num_to_test = min(samples, len(golden_labels)) test_samples = golden_labels[:num_to_test] print(f"Will run on {len(test_samples)} samples from dataset '{dataset_name}'.") config = MODELS_CONFIG.get(model_name) model_class = get_model_class(config["class"]) model_instance_name = config["model_name"] benchmark_helper = MapGuesserBenchmark(dataset_name=dataset_name, headless=True) all_results = [] with GeoBot( model=model_class, model_name=model_instance_name, headless=headless, temperature=temperature, ) as bot: for i, sample in enumerate(test_samples): print( f"\n--- Running Sample {i + 1}/{len(test_samples)} (ID: {sample.get('id')}) ---" ) if not bot.controller.load_location_from_data(sample): print( f" ❌ Failed to load location for sample {sample.get('id')}. Skipping." ) continue bot.controller.setup_clean_environment() final_guess = bot.run_agent_loop(max_steps=steps) true_coords = {"lat": sample.get("lat"), "lng": sample.get("lng")} distance_km = None is_success = False if final_guess: distance_km = benchmark_helper.calculate_distance( true_coords, final_guess ) if distance_km is not None: is_success = distance_km <= SUCCESS_THRESHOLD_KM print(f"\nResult for Sample ID: {sample.get('id')}") print( f" Ground Truth: Lat={true_coords['lat']:.4f}, Lon={true_coords['lng']:.4f}" ) print( f" Final Guess: Lat={final_guess[0]:.4f}, Lon={final_guess[1]:.4f}" ) dist_str = f"{distance_km:.1f} km" if distance_km is not None else "N/A" print(f" Distance: {dist_str}, Success: {is_success}") else: print("Agent did not make a final guess for this sample.") all_results.append( { "sample_id": sample.get("id"), "model": bot.model_name, "true_coordinates": true_coords, "predicted_coordinates": final_guess, "distance_km": distance_km, "success": is_success, } ) summary = benchmark_helper.generate_summary(all_results) if summary: print( f"\n\n--- Agent Benchmark Complete for dataset '{dataset_name}'! Summary ---" ) for model, stats in summary.items(): print(f"Model: {model}") print(f" Success Rate: {stats['success_rate'] * 100:.1f}%") print(f" Avg Distance: {stats['average_distance_km']:.1f} km") print("Agent Mode finished.") def benchmark_mode( models: list, samples: int, headless: bool, dataset_name: str = "default", temperature: float = 0.0, ): """Runs the benchmark on pre-collected data.""" print( f"Starting Benchmark Mode: models={models}, samples={samples}, dataset={dataset_name}, temperature={temperature}" ) benchmark = MapGuesserBenchmark(dataset_name=dataset_name, headless=headless) summary = benchmark.run_benchmark( models=models, max_samples=samples, temperature=temperature ) if summary: print(f"\n--- Benchmark Complete for dataset '{dataset_name}'! Summary ---") for model, stats in summary.items(): print(f"Model: {model}") print(f" Success Rate: {stats['success_rate'] * 100:.1f}%") print(f" Avg Distance: {stats['average_distance_km']:.1f} km") def collect_mode(dataset_name: str, samples: int, headless: bool): """Collects data for a new dataset.""" print(f"Starting Data Collection: dataset={dataset_name}, samples={samples}") with DataCollector(dataset_name=dataset_name, headless=headless) as collector: collector.collect_samples(num_samples=samples) print(f"Data collection complete for dataset '{dataset_name}'.") def main(): parser = argparse.ArgumentParser(description="MapCrunch AI Agent & Benchmark") parser.add_argument( "--mode", choices=["agent", "benchmark", "collect"], default="agent", help="Operation mode.", ) parser.add_argument( "--dataset", default="default", help="Dataset name to use or create.", ) parser.add_argument( "--model", choices=list(MODELS_CONFIG.keys()), default="gpt-4o", help="Model to use.", ) parser.add_argument( "--steps", type=int, default=10, help="[Agent] Number of exploration steps." ) parser.add_argument( "--samples", type=int, default=50, help="Number of samples to process for the selected mode.", ) parser.add_argument( "--headless", action="store_true", help="Run browser in headless mode." ) parser.add_argument( "--models", nargs="+", choices=list(MODELS_CONFIG.keys()), help="[Benchmark] Models to benchmark.", ) parser.add_argument( "--temperature", type=float, default=0.0, help="Temperature parameter for LLM sampling (0.0 = deterministic, higher = more random). Default: 0.0", ) args = parser.parse_args() if args.mode == "collect": collect_mode( dataset_name=args.dataset, samples=args.samples, headless=args.headless, ) elif args.mode == "agent": agent_mode( model_name=args.model, steps=args.steps, headless=args.headless, samples=args.samples, dataset_name=args.dataset, temperature=args.temperature, ) elif args.mode == "benchmark": benchmark_mode( models=args.models or [args.model], samples=args.samples, headless=args.headless, dataset_name=args.dataset, temperature=args.temperature, ) if __name__ == "__main__": main()