Coverage for src/local_deep_research/benchmarks/runners.py: 100%
122 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-20 01:24 +0000
1"""
2Benchmark runners for Local Deep Research.
4This module provides the main functions for running benchmarks using LDR.
5"""
7import json
8from loguru import logger
9import os
10from pathlib import Path
11import time
12from typing import Any, Callable, Dict, Optional
14from ..api import quick_summary
15from .datasets import DEFAULT_DATASET_URLS, load_dataset
16from .datasets.base import DatasetRegistry
17from .graders import extract_answer_from_response, grade_results
18from .metrics import calculate_metrics, generate_report
19from .templates import BROWSECOMP_QUERY_TEMPLATE
22def format_query(question: str, dataset_type: str = "simpleqa") -> str:
23 """
24 Format query based on dataset type.
26 Args:
27 question: Original question
28 dataset_type: Type of dataset
30 Returns:
31 Formatted query for LDR
32 """
33 if dataset_type.lower() == "browsecomp":
34 # BrowseComp requires specific formatting
35 return BROWSECOMP_QUERY_TEMPLATE.format(question=question)
37 # Simple format for SimpleQA
38 return question
41def run_benchmark(
42 dataset_type: str,
43 dataset_path: Optional[str] = None,
44 num_examples: Optional[int] = None,
45 output_dir: str = "benchmark_results",
46 run_evaluation: bool = True,
47 evaluation_config: Optional[Dict[str, Any]] = None,
48 search_config: Optional[Dict[str, Any]] = None,
49 human_evaluation: bool = False,
50 progress_callback: Optional[Callable[[str, int, Dict], None]] = None,
51 seed: int = 42,
52 settings_snapshot: Optional[Dict[str, Any]] = None,
53) -> Dict[str, Any]:
54 """
55 Run a benchmark on the specified dataset.
57 Args:
58 dataset_type: Type of dataset ("simpleqa" or "browsecomp")
59 dataset_path: Optional custom dataset path
60 num_examples: Number of examples to use
61 output_dir: Directory to save results
62 run_evaluation: Whether to evaluate results
63 evaluation_config: Custom LLM config for evaluation
64 search_config: Custom search parameters
65 human_evaluation: Whether to use human evaluation
66 progress_callback: Optional callback for progress updates
67 seed: Random seed for reproducibility
69 Returns:
70 Dictionary with benchmark results and metrics
71 """
72 # Ensure output directory exists
73 os.makedirs(output_dir, exist_ok=True)
75 # Default search configuration
76 if not search_config:
77 search_config = {
78 "iterations": 3,
79 "questions_per_iteration": 3,
80 "search_tool": "searxng",
81 }
83 # Load dataset using the class-based approach
84 try:
85 # Create the dataset instance from registry
86 dataset_instance = DatasetRegistry.create_dataset(
87 dataset_id=dataset_type.lower(),
88 dataset_path=dataset_path,
89 num_examples=num_examples,
90 seed=seed,
91 )
92 # Load the examples
93 dataset = dataset_instance.load()
95 logger.info(
96 f"Loaded {len(dataset)} examples using dataset class {type(dataset_instance).__name__}"
97 )
98 except Exception:
99 # Fallback to legacy function if there's any issue
100 logger.warning(
101 "Error using dataset class. Falling back to legacy function."
102 )
103 dataset = load_dataset(
104 dataset_type=dataset_type,
105 dataset_path=dataset_path,
106 num_examples=num_examples,
107 seed=seed,
108 )
110 # Set up output files
111 timestamp = time.strftime("%Y%m%d_%H%M%S")
112 results_file = str(
113 Path(output_dir) / f"{dataset_type}_{timestamp}_results.jsonl"
114 )
115 evaluation_file = str(
116 Path(output_dir) / f"{dataset_type}_{timestamp}_evaluation.jsonl"
117 )
118 report_file = str(
119 Path(output_dir) / f"{dataset_type}_{timestamp}_report.md"
120 )
122 # Make sure output files don't exist
123 for file in [results_file, evaluation_file, report_file]:
124 file_path = Path(file)
125 if file_path.exists():
126 file_path.unlink()
128 # Progress tracking
129 total_examples = len(dataset)
131 if progress_callback:
132 progress_callback(
133 "Starting benchmark",
134 0,
135 {
136 "status": "started",
137 "dataset_type": dataset_type,
138 "total_examples": total_examples,
139 },
140 )
142 # Process each example
143 results = []
145 for i, example in enumerate(dataset):
146 # Extract question and answer in a way that uses the dataset class when available
147 if "dataset_instance" in locals() and isinstance(
148 dataset_instance,
149 DatasetRegistry.get_dataset_class(dataset_type.lower()),
150 ):
151 # Use the dataset class methods to extract question and answer
152 question = dataset_instance.get_question(example)
153 correct_answer = dataset_instance.get_answer(example)
154 logger.debug(
155 "Using dataset class methods to extract question and answer"
156 )
157 else:
158 # Fallback to the legacy approach
159 if dataset_type.lower() == "simpleqa":
160 question = example.get("problem", "")
161 correct_answer = example.get("answer", "")
162 else: # browsecomp
163 question = example.get("problem", "")
164 # For BrowseComp, the answer should be in "correct_answer" after decryption
165 correct_answer = example.get("correct_answer", "")
166 if not correct_answer and "answer" in example:
167 # Fallback to "answer" field if "correct_answer" is not available
168 correct_answer = example.get("answer", "")
170 # Update progress
171 if progress_callback:
172 progress_callback(
173 f"Processing example {i + 1}/{total_examples}",
174 int(i / total_examples * 50),
175 {
176 "status": "processing",
177 "current": i + 1,
178 "total": total_examples,
179 "question": (
180 question[:50] + "..."
181 if len(question) > 50
182 else question
183 ),
184 },
185 )
187 logger.info(f"Processing {i + 1}/{total_examples}: {question[:50]}...")
189 try:
190 # Format query based on dataset type
191 formatted_query = format_query(question, dataset_type)
193 # Time the search
194 start_time = time.time()
196 # Get response from LDR. Pass settings_snapshot through so
197 # the egress policy PEP fires on the benchmark's inner
198 # search-engine / LLM construction calls. Without it, the
199 # PEP's ``if settings_snapshot is not None`` guard skips and
200 # benchmarks silently bypass the user's saved scope.
201 search_result = quick_summary(
202 query=formatted_query,
203 iterations=search_config.get("iterations", 3),
204 questions_per_iteration=search_config.get(
205 "questions_per_iteration", 3
206 ),
207 search_tool=search_config.get("search_tool", "searxng"),
208 settings_snapshot=settings_snapshot,
209 )
211 end_time = time.time()
212 processing_time = end_time - start_time
214 # Extract response and search info
215 response = search_result.get("summary", "")
217 # Extract structured information
218 extracted = extract_answer_from_response(response, dataset_type)
220 # Format result
221 result = {
222 "id": example.get("id", f"example_{i}"),
223 "problem": question,
224 "correct_answer": correct_answer,
225 "response": response,
226 "extracted_answer": extracted["extracted_answer"],
227 "confidence": extracted["confidence"],
228 "processing_time": processing_time,
229 "sources": search_result.get("sources", []),
230 "search_config": search_config,
231 }
233 # Add to results list
234 results.append(result)
236 # Write result to file
237 with open(results_file, "a", encoding="utf-8") as f:
238 f.write(json.dumps(result) + "\n")
240 # Update progress
241 if progress_callback:
242 progress_callback(
243 f"Completed example {i + 1}/{total_examples}",
244 int((i + 0.5) / total_examples * 50),
245 {
246 "status": "completed_example",
247 "current": i + 1,
248 "total": total_examples,
249 "result": result,
250 },
251 )
253 except Exception as e:
254 logger.exception(f"Error processing example {i + 1}")
256 # Create error result
257 error_result = {
258 "id": example.get("id", f"example_{i}"),
259 "problem": question,
260 "correct_answer": correct_answer,
261 "error": str(e),
262 "processing_time": (
263 time.time() - start_time if "start_time" in locals() else 0
264 ),
265 }
267 # Add to results list
268 results.append(error_result)
270 # Write error result to file
271 with open(results_file, "a", encoding="utf-8") as f:
272 f.write(json.dumps(error_result) + "\n")
274 # Update progress
275 if progress_callback:
276 progress_callback(
277 f"Error processing example {i + 1}/{total_examples}",
278 int((i + 0.5) / total_examples * 50),
279 {
280 "status": "error",
281 "current": i + 1,
282 "total": total_examples,
283 "error": str(e),
284 "result": error_result,
285 },
286 )
288 logger.info(f"Completed processing {total_examples} examples")
290 # Run evaluation if requested
291 if run_evaluation:
292 if progress_callback:
293 progress_callback(
294 "Starting evaluation",
295 50,
296 {"status": "evaluating", "results_file": results_file},
297 )
299 if human_evaluation:
300 from .graders import human_evaluation as evaluate
302 logger.info("Running human evaluation...")
303 evaluation_results = evaluate(
304 results_file=results_file,
305 output_file=evaluation_file,
306 interactive=True,
307 )
308 else:
309 logger.info("Running automated evaluation...")
310 try:
311 evaluation_results = grade_results(
312 results_file=results_file,
313 output_file=evaluation_file,
314 dataset_type=dataset_type,
315 evaluation_config=evaluation_config,
316 settings_snapshot=settings_snapshot,
317 progress_callback=lambda current, total, meta: (
318 progress_callback(
319 f"Evaluating {current + 1}/{total}",
320 50 + int((current + 0.5) / total * 40),
321 {**meta, "status": "evaluating"},
322 )
323 if progress_callback
324 else None
325 ),
326 )
327 except Exception as e:
328 logger.exception("Automated evaluation failed")
330 if progress_callback:
331 progress_callback(
332 "Automated evaluation failed. Falling back to human evaluation.",
333 60,
334 {"status": "evaluation_fallback", "error": str(e)},
335 )
337 # Ask if user wants to fall back to human evaluation
338 fallback_to_human = False
339 print("\nAutomated evaluation failed with error:", str(e))
340 response = input(
341 "Do you want to fall back to human evaluation? (y/n): "
342 )
343 fallback_to_human = response.strip().lower() == "y"
345 if fallback_to_human:
346 logger.info("Falling back to human evaluation...")
347 from .graders import human_evaluation as evaluate
349 evaluation_results = evaluate(
350 results_file=results_file,
351 output_file=evaluation_file,
352 interactive=True,
353 )
354 else:
355 from ..security.file_write_verifier import (
356 write_file_verified,
357 )
359 logger.info("Skipping evaluation due to error.")
360 # Create an empty evaluation file to prevent issues
361 write_file_verified(
362 evaluation_file,
363 "",
364 "benchmark.allow_file_output",
365 context="empty evaluation placeholder",
366 )
368 return {
369 "status": "evaluation_error",
370 "dataset_type": dataset_type,
371 "results_path": results_file,
372 "evaluation_error": str(e),
373 "total_examples": total_examples,
374 }
376 # Calculate metrics
377 if progress_callback:
378 progress_callback(
379 "Calculating metrics", 90, {"status": "calculating_metrics"}
380 )
382 metrics = calculate_metrics(evaluation_file)
384 # Generate report
385 if progress_callback:
386 progress_callback(
387 "Generating report", 95, {"status": "generating_report"}
388 )
390 dataset_name = dataset_type.capitalize()
391 report_path = generate_report(
392 metrics=metrics,
393 results_file=evaluation_file,
394 output_file=report_file,
395 dataset_name=dataset_name,
396 config_info={
397 "Dataset": dataset_path
398 or DEFAULT_DATASET_URLS.get(dataset_type, "Unknown"),
399 "Examples": total_examples,
400 "Iterations": search_config.get("iterations", 3),
401 "Questions per iteration": search_config.get(
402 "questions_per_iteration", 3
403 ),
404 "Search tool": search_config.get("search_tool", "searxng"),
405 "Evaluation method": "Human"
406 if human_evaluation
407 else "Automated",
408 },
409 )
411 # Mark as complete
412 if progress_callback:
413 progress_callback(
414 "Benchmark complete",
415 100,
416 {
417 "status": "complete",
418 "metrics": metrics,
419 "report_path": report_path,
420 },
421 )
423 return {
424 "status": "complete",
425 "dataset_type": dataset_type,
426 "results_path": results_file,
427 "evaluation_path": evaluation_file,
428 "report_path": report_path,
429 "metrics": metrics,
430 "total_examples": total_examples,
431 "accuracy": metrics.get("accuracy", 0),
432 }
434 # No evaluation, just return results
435 if progress_callback:
436 progress_callback(
437 "Benchmark complete (no evaluation)",
438 100,
439 {"status": "complete_no_eval", "results_path": results_file},
440 )
442 return {
443 "status": "complete_no_eval",
444 "dataset_type": dataset_type,
445 "results_path": results_file,
446 "total_examples": total_examples,
447 }
450def run_simpleqa_benchmark(num_examples: int = 100, **kwargs) -> Dict[str, Any]:
451 """
452 Run SimpleQA benchmark with default settings.
454 Args:
455 num_examples: Number of examples to process
456 **kwargs: Additional arguments to pass to run_benchmark
458 Returns:
459 Dictionary with benchmark results
460 """
461 return run_benchmark(
462 dataset_type="simpleqa", num_examples=num_examples, **kwargs
463 )
466def run_browsecomp_benchmark(
467 num_examples: int = 100, **kwargs
468) -> Dict[str, Any]:
469 """
470 Run BrowseComp benchmark with default settings.
472 Args:
473 num_examples: Number of examples to process
474 **kwargs: Additional arguments to pass to run_benchmark
476 Returns:
477 Dictionary with benchmark results
478 """
479 return run_benchmark(
480 dataset_type="browsecomp", num_examples=num_examples, **kwargs
481 )
484def run_xbench_deepsearch_benchmark(
485 num_examples: int = 100, **kwargs
486) -> Dict[str, Any]:
487 """
488 Run xbench-DeepSearch benchmark with default settings.
490 Args:
491 num_examples: Number of examples to process
492 **kwargs: Additional arguments to pass to run_benchmark
494 Returns:
495 Dictionary with benchmark results
496 """
497 return run_benchmark(
498 dataset_type="xbench_deepsearch", num_examples=num_examples, **kwargs
499 )