Coverage for src/local_deep_research/benchmarks/web_api/benchmark_service.py: 97%
501 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +0000
1"""Benchmark service for handling web-based benchmark execution."""
3import hashlib
4import json
5import threading
6import time
7from datetime import UTC, datetime
8from enum import Enum
9from typing import Any, Dict, List, Optional, Tuple
11from loguru import logger
12from sqlalchemy.exc import SQLAlchemyError
14from ...api.research_functions import quick_summary
15from ...settings.manager import SnapshotSettingsContext
16from ...web.services.research_service import _global_research_semaphore
17from ...database.models.benchmark import (
18 BenchmarkResult,
19 BenchmarkRun,
20 BenchmarkStatus,
21 DatasetType,
22)
23from ..datasets import load_dataset
24from ..graders import extract_answer_from_response, grade_single_result
25from ..runners import format_query
26from ...database.thread_local_session import thread_cleanup
28# Generic failure message surfaced to API clients in place of the raw benchmark
29# exception. ``benchmark_run.error_message`` is set from ``str(e)`` of the
30# benchmark run (LLM calls, search engines, grading — the same stack as
31# research) and returned by get_benchmark_status() via
32# GET /benchmark/api/status/<id>. The raw text can carry server-level LLM
33# endpoints / hosts / filesystem paths (settings_snapshot includes LDR_* env
34# overrides), so it must not reach the client (CWE-209). Full detail stays in
35# the logger.exception calls at the failure sites.
36_GENERIC_BENCHMARK_ERROR = (
37 "Benchmark run failed due to an internal error. "
38 "Check the server logs for details."
39)
40_RESULT_PERSISTENCE_ERROR = {
41 "code": "database_write_failed",
42 "message": (
43 "Benchmark results could not be saved. Check server logs and "
44 "database storage, then retry."
45 ),
46}
49class SocketIOService:
50 """Adapter over socketio_asgi's module-level emit functions.
52 Keeps the Flask-era ``SocketIOService`` interface (which
53 BenchmarkService and its tests are written against) while delegating
54 to the FastAPI-native emit path. An earlier port made this a no-op
55 stub, silently dropping every benchmark progress / completion /
56 rate-limit WebSocket event.
58 ``owner`` is keyword-only and REQUIRED here, where the upstream
59 ``SocketIOService.emit_to_subscribers`` takes an optional
60 ``username=None``. Same fix, stricter contract: a benchmark id is a
61 per-user integer, so an omitted owner silently files the event under
62 the wrong subscriber set. Making it required turns that into a
63 TypeError at the call site instead of a misrouted event.
64 """
66 def emit_to_subscribers(
67 self, event_base, research_id, data, *, owner, enable_logging=True
68 ):
69 from ...web.services.socketio_asgi import emit_to_subscribers
71 return emit_to_subscribers(
72 event_base,
73 research_id,
74 data,
75 owner=owner,
76 enable_logging=enable_logging,
77 )
79 def emit_to_room(self, event, data, room=None):
80 from ...web.services.socketio_asgi import emit_socket_event
82 return emit_socket_event(event, data, room=room)
85def _required_persistence_fields(
86 result: Any,
87) -> tuple[Any, str, DatasetType, Any, Any]:
88 """Extract required fields or raise a data-shape error."""
89 if not isinstance(result, dict): 89 ↛ 90line 89 didn't jump to line 90 because the condition on line 89 was never true
90 raise TypeError("result must be a mapping")
91 query_hash = result["query_hash"]
92 if not isinstance(query_hash, str) or not query_hash: 92 ↛ 93line 92 didn't jump to line 93 because the condition on line 92 was never true
93 raise ValueError("query_hash must be a non-empty string")
94 return (
95 result["example_id"],
96 query_hash,
97 DatasetType(result["dataset_type"]),
98 result["question"],
99 result["correct_answer"],
100 )
103class BenchmarkTaskStatus(Enum):
104 """Status values for benchmark tasks in the queue tracker."""
106 QUEUED = "queued"
107 PROCESSING = "processing"
108 COMPLETED = "completed"
109 FAILED = "failed"
110 CANCELLED = "cancelled"
113class BenchmarkQueueTracker:
114 """Simple in-memory tracker for benchmark queue status.
116 This replaces the removed memory_queue functionality for benchmarks.
117 Since benchmarks are temporary and don't need persistence,
118 this simple in-memory solution is sufficient.
120 Thread-safe for concurrent access from multiple benchmark threads.
121 """
123 def __init__(self):
124 self.tasks = {}
125 self._lock = threading.Lock()
127 def add_task(
128 self, task_id: str, username: str, task_type: str = "benchmark"
129 ):
130 """Add a new task to tracking.
132 Also performs opportunistic cleanup of old completed tasks.
133 """
134 # Cleanup old tasks before adding new one (outside lock for better performance)
135 self.cleanup_completed_tasks()
137 with self._lock:
138 self.tasks[task_id] = {
139 "username": username,
140 "task_type": task_type,
141 "status": BenchmarkTaskStatus.QUEUED.value,
142 "created_at": datetime.now(UTC),
143 }
145 def update_task_status(self, task_id: str, status: BenchmarkTaskStatus):
146 """Update the status of a task."""
147 with self._lock:
148 if task_id in self.tasks:
149 self.tasks[task_id]["status"] = status.value
150 self.tasks[task_id]["updated_at"] = datetime.now(UTC)
151 else:
152 logger.warning(
153 f"Attempted to update status for non-existent task: {task_id}"
154 )
156 def get_task_status(self, task_id: str) -> Optional[Dict]:
157 """Get the current status of a task."""
158 with self._lock:
159 return self.tasks.get(task_id)
161 def remove_task(self, task_id: str):
162 """Remove a task from tracking."""
163 with self._lock:
164 self.tasks.pop(task_id, None)
166 def cleanup_completed_tasks(self, max_age_seconds: int = 3600):
167 """Remove completed tasks older than max_age_seconds.
169 Args:
170 max_age_seconds: Maximum age in seconds for completed tasks (default 1 hour)
171 """
172 with self._lock:
173 now = datetime.now(UTC)
174 to_remove = []
175 for task_id, task_data in self.tasks.items():
176 # Only cleanup completed, failed, or cancelled tasks
177 if task_data["status"] in [
178 BenchmarkTaskStatus.COMPLETED.value,
179 BenchmarkTaskStatus.FAILED.value,
180 BenchmarkTaskStatus.CANCELLED.value,
181 ]:
182 # Check if task has updated_at timestamp
183 updated_at = task_data.get(
184 "updated_at", task_data.get("created_at")
185 )
186 if updated_at: 186 ↛ 175line 186 didn't jump to line 175 because the condition on line 186 was always true
187 age = (now - updated_at).total_seconds()
188 if age > max_age_seconds:
189 to_remove.append(task_id)
191 for task_id in to_remove:
192 self.tasks.pop(task_id, None)
193 logger.debug(f"Cleaned up old task: {task_id}")
195 if to_remove:
196 logger.info(f"Cleaned up {len(to_remove)} old benchmark tasks")
199class BenchmarkService:
200 """Service for managing benchmark runs through the web interface."""
202 def __init__(self, socket_service=None):
203 # Keyed by (username, BenchmarkRun.id). BenchmarkRun.id is a PER-USER
204 # autoincrement, so two users can each own run id 1; keying by the
205 # owner as well prevents one user's start_benchmark from clobbering
206 # another's in-memory run (and the cross-user reads that enabled).
207 self.active_runs: Dict[Tuple[str, int], Dict] = {}
208 self.socket_service = socket_service or SocketIOService()
209 # Track rate limiting per user benchmark run.
210 self.rate_limit_detected: Dict[Tuple[str, int], bool] = {}
211 self.queue_tracker = BenchmarkQueueTracker() # Initialize queue tracker
212 # Serializes benchmark-result persistence. _sync_results_to_database
213 # runs on the worker thread while sync_pending_results runs on the
214 # request thread; without this they can both INSERT the same
215 # (benchmark_run_id, query_hash) row and trip the uix_run_query
216 # unique constraint, which rolls back the whole pending batch.
217 self._results_sync_lock = threading.Lock()
219 def generate_config_hash(self, search_config: Dict[str, Any]) -> str:
220 """Generate a hash for search configuration compatibility checking."""
221 relevant_params = {
222 "iterations": search_config.get("iterations"),
223 "questions_per_iteration": search_config.get(
224 "questions_per_iteration"
225 ),
226 "search_tool": search_config.get("search_tool"),
227 "search_strategy": search_config.get("search_strategy"),
228 "model_name": search_config.get("model_name"),
229 "provider": search_config.get("provider"),
230 }
231 # Remove None values
232 relevant_params = {
233 k: v for k, v in relevant_params.items() if v is not None
234 }
235 config_str = json.dumps(relevant_params, sort_keys=True)
236 return hashlib.md5( # DevSkim: ignore DS126858
237 config_str.encode(), usedforsecurity=False
238 ).hexdigest()[:8]
240 def generate_query_hash(
241 self, question: str, dataset_type: str, example_id: str
242 ) -> str:
243 """Generate a hash for a query to enable deduplication.
245 The example id is part of the key so two distinct examples that share
246 the same question text (and dataset) get distinct hashes and each
247 persist a row, instead of colliding on the ``uix_run_query`` unique
248 constraint and dropping one. Re-hashing the same example is stable, so
249 the sync dedup still collapses a genuine re-sync of that example.
250 """
251 query_content = (
252 f"{example_id}|{question.strip()}|{dataset_type.lower()}"
253 )
254 return hashlib.md5( # DevSkim: ignore DS126858
255 query_content.encode(), usedforsecurity=False
256 ).hexdigest()
258 def create_benchmark_run(
259 self,
260 run_name: Optional[str],
261 search_config: Dict[str, Any],
262 evaluation_config: Dict[str, Any],
263 datasets_config: Dict[str, Dict],
264 username: Optional[str] = None,
265 user_password: Optional[str] = None,
266 ) -> int:
267 """Create a new benchmark run in the database."""
268 from ...database.session_context import get_user_db_session
270 with get_user_db_session(username, user_password) as session:
271 try:
272 config_hash = self.generate_config_hash(search_config)
274 # Calculate total examples
275 total_examples = sum(
276 config.get("count", 0)
277 for config in datasets_config.values()
278 )
280 benchmark_run = BenchmarkRun(
281 run_name=run_name,
282 config_hash=config_hash,
283 query_hash_list=[], # Will be populated as we process
284 search_config=search_config,
285 evaluation_config=evaluation_config,
286 datasets_config=datasets_config,
287 total_examples=total_examples,
288 status=BenchmarkStatus.PENDING,
289 )
291 session.add(benchmark_run)
292 session.commit()
294 logger.info(
295 f"Created benchmark run {benchmark_run.id} with config hash {config_hash}"
296 )
297 return int(benchmark_run.id)
299 except Exception:
300 session.rollback()
301 logger.exception("Error creating benchmark run")
302 raise
304 def start_benchmark(
305 self,
306 benchmark_run_id: int,
307 username: Optional[str] = None,
308 user_password: Optional[str] = None,
309 ) -> bool:
310 """Start a benchmark run in a background thread."""
311 from ...database.session_context import get_user_db_session
313 try:
314 # Get all data from the database in the main thread
315 # This avoids database access from the background thread
316 with get_user_db_session(username, user_password) as session:
317 # Get benchmark run details
318 benchmark_run = (
319 session.query(BenchmarkRun)
320 .filter(BenchmarkRun.id == benchmark_run_id)
321 .first()
322 )
323 if not benchmark_run:
324 raise ValueError( # noqa: TRY301 — caught by except, sets FAILED status in DB
325 f"Benchmark run {benchmark_run_id} not found"
326 )
328 # Create settings snapshot for thread safety + provenance.
329 # The settings snapshot is BEST-EFFORT provenance, not a
330 # critical path: a recoverable read failure must not break
331 # the user's ability to start a benchmark. We catch only the
332 # failures that can legitimately escape get_all_settings:
333 # SQLAlchemyError (transient DB issues) and LookupError /
334 # ValueError (a stale enum-typed settings row, e.g. a legacy
335 # 'CHAT' value that fails enum coercion). Those degrade to an
336 # empty snapshot — ldr_version is still recorded. Anything
337 # else (a genuine programming error) is deliberately NOT
338 # swallowed here: it propagates to the outer handler, which
339 # logs a full traceback and marks the run FAILED — the right
340 # outcome, since a settings subsystem that breaks unexpectedly
341 # would also corrupt the benchmark's own config reads.
342 from local_deep_research.settings import SettingsManager
344 settings_manager = SettingsManager(session)
345 try:
346 settings_snapshot = settings_manager.get_all_settings()
347 except (SQLAlchemyError, LookupError, ValueError):
348 # Use logger.exception so the traceback goes to stderr
349 # without splatting `repr(exc)` into the message body —
350 # exception details may include user data.
351 logger.exception(
352 "Failed to capture settings snapshot for benchmark "
353 f"{benchmark_run_id}; running with empty snapshot."
354 )
355 settings_snapshot = {}
357 # Get user password for metrics tracking in background thread
358 from ...utilities.request_context import (
359 get_current_session_id,
360 )
361 from ...database.session_passwords import session_password_store
363 _user_password = None
364 session_id = get_current_session_id()
365 if session_id and username:
366 _user_password = (
367 session_password_store.get_session_password(
368 username, session_id
369 )
370 )
371 if not _user_password: 371 ↛ 372line 371 didn't jump to line 372 because the condition on line 371 was never true
372 logger.warning(
373 f"No password found for user {username} in current session"
374 )
376 # Extract all data we need
377 benchmark_data = {
378 "benchmark_run_id": benchmark_run_id,
379 "username": username or "benchmark_user",
380 "user_password": _user_password, # Add password for metrics tracking
381 "config_hash": benchmark_run.config_hash,
382 "datasets_config": benchmark_run.datasets_config,
383 "search_config": benchmark_run.search_config,
384 "evaluation_config": benchmark_run.evaluation_config,
385 "settings_snapshot": settings_snapshot, # Add settings snapshot
386 }
388 # Update status in database. Persist provenance (ldr_version
389 # + redacted settings_snapshot) BEFORE the commit so the
390 # YAML download can later reproduce the run; the snapshot
391 # already exists in memory from line 306, we just store the
392 # redacted form so secrets don't sit in the DB even though
393 # SQLCipher protects at rest.
394 from local_deep_research import __version__
395 from local_deep_research.security.data_sanitizer import (
396 DataSanitizer,
397 )
399 benchmark_run.ldr_version = __version__
400 benchmark_run.settings_snapshot = (
401 DataSanitizer.redact_settings_snapshot(settings_snapshot)
402 )
403 benchmark_run.status = BenchmarkStatus.IN_PROGRESS
404 benchmark_run.start_time = datetime.now(UTC)
405 session.commit()
407 # Store data in memory for the thread. Key by (username, id) so a
408 # run another user owns under the same per-user id can't overwrite
409 # this entry (see _run_key).
410 run_key = self._run_key(username, benchmark_run_id)
411 self.active_runs[run_key] = {
412 "data": benchmark_data,
413 "start_time": datetime.now(UTC),
414 "status": "running",
415 "results": [],
416 }
418 # Start background thread. Pass the owner so the thread rebuilds the
419 # same composite key for its own active_runs accesses.
420 thread = threading.Thread(
421 target=self._run_benchmark_thread,
422 args=(username, benchmark_run_id),
423 daemon=True,
424 )
425 thread.start()
427 self.active_runs[run_key]["thread"] = thread
429 logger.info(f"Started benchmark run {benchmark_run_id}")
430 return True
432 except Exception as e:
433 logger.exception(f"Error starting benchmark {benchmark_run_id}")
434 # If we populated active_runs before the spawn failed, drop the
435 # stale entry — it has no thread and would mislead subsequent
436 # cancel_benchmark / get_run_status calls. Recompute the key here
437 # (the local run_key may not exist yet if we failed before it).
438 self.active_runs.pop(
439 self._run_key(username, benchmark_run_id), None
440 )
441 # Update status using user database
442 with get_user_db_session(username, user_password) as session:
443 benchmark_run = (
444 session.query(BenchmarkRun)
445 .filter(BenchmarkRun.id == benchmark_run_id)
446 .first()
447 )
448 if benchmark_run:
449 benchmark_run.status = BenchmarkStatus.FAILED
450 benchmark_run.error_message = str(e)
451 session.commit()
452 return False
454 @thread_cleanup
455 def _run_benchmark_thread(
456 self, username: Optional[str], benchmark_run_id: int
457 ):
458 """Main benchmark execution thread."""
459 # IMPORTANT: This runs in a background thread, so we cannot access the user database
460 # Using in-memory queue tracker for benchmark status tracking
462 task_id = None
463 # Rebuild the same composite key start_benchmark stored this run under.
464 run_key = self._run_key(username, benchmark_run_id)
466 # Get the benchmark data that was passed to us
467 # We need to retrieve this from the service database or from memory
468 benchmark_data = self.active_runs.get(run_key, {}).get("data")
470 try:
471 if not benchmark_data:
472 raise ValueError( # noqa: TRY301
473 f"Benchmark data for run {benchmark_run_id} not found"
474 )
475 # Set up settings context for thread-local access
476 settings_snapshot = benchmark_data.get("settings_snapshot", {})
477 username = benchmark_data.get("username", "benchmark_user")
479 # Create a settings context that threads can use
480 settings_context = SnapshotSettingsContext(
481 settings_snapshot,
482 username=username,
483 missing_key_log_level="WARNING",
484 )
486 # Set the context in thread-local storage
487 from ...config.thread_settings import set_settings_context
489 set_settings_context(settings_context)
491 # Extract all the data we need
492 datasets_config = benchmark_data["datasets_config"]
493 search_config = benchmark_data["search_config"]
494 evaluation_config = benchmark_data["evaluation_config"]
496 # Create task queue
497 task_queue = self._create_task_queue(
498 datasets_config,
499 benchmark_run_id,
500 )
502 # Calculate totals
503 total_examples = len(task_queue)
504 completed_examples = 0
506 # Initialize task tracking
507 task_id = f"benchmark_{benchmark_run_id}_{int(datetime.now(UTC).timestamp())}"
508 username = benchmark_data.get("username", "benchmark_user")
509 self.queue_tracker.add_task(task_id, username, "benchmark")
510 self.queue_tracker.update_task_status(
511 task_id, BenchmarkTaskStatus.PROCESSING
512 )
514 # Track progress in memory
515 progress_info = {
516 "total_examples": total_examples,
517 "completed_examples": completed_examples,
518 "failed_examples": 0,
519 "start_time": datetime.now(UTC),
520 }
522 # Process tasks
523 logger.info(
524 f"Benchmark {benchmark_run_id} starting to process {len(task_queue)} tasks"
525 )
526 for i, task in enumerate(task_queue):
527 # Check if benchmark has been cancelled
528 if (
529 run_key in self.active_runs
530 and self.active_runs[run_key].get("status") == "cancelled"
531 ):
532 logger.info(
533 f"Benchmark {benchmark_run_id} was cancelled, stopping processing"
534 )
535 break
537 logger.info(
538 f"Benchmark {benchmark_run_id} processing task {i + 1}/{len(task_queue)}"
539 )
540 try:
541 # Add username and password to task for metrics tracking
542 task["username"] = benchmark_data.get("username")
543 task["user_password"] = benchmark_data.get("user_password")
545 # Acquire the global research semaphore so benchmark
546 # tasks count against the server-wide concurrency limit
547 _global_research_semaphore.acquire()
548 try:
549 # Process single task
550 result = self._process_benchmark_task(
551 task,
552 search_config,
553 evaluation_config,
554 )
555 finally:
556 _global_research_semaphore.release()
558 # Store result in memory for now (will be saved later)
559 if "results" not in self.active_runs[run_key]: 559 ↛ 560line 559 didn't jump to line 560 because the condition on line 559 was never true
560 self.active_runs[run_key]["results"] = []
561 self.active_runs[run_key]["results"].append(result)
563 # Update progress
564 progress_info["completed_examples"] += 1
566 logger.info(
567 f"Benchmark {benchmark_run_id} task {i + 1}/{len(task_queue)} completed successfully. "
568 f"Progress: {progress_info['completed_examples']}/{progress_info['total_examples']} total examples"
569 )
571 # Send real-time update
572 self._send_progress_update(
573 benchmark_run_id,
574 progress_info["completed_examples"],
575 progress_info["total_examples"],
576 username=username,
577 )
579 except Exception as e:
580 logger.exception(f"Error processing task {i}")
581 progress_info["failed_examples"] += 1
582 logger.info(
583 f"Benchmark {benchmark_run_id} task {i + 1}/{len(task_queue)} failed. "
584 f"Total failed: {progress_info['failed_examples']}"
585 )
587 # Check if this is a rate limiting error
588 error_str = str(e).lower()
589 if (
590 "403" in error_str
591 or "rate limit" in error_str
592 or "forbidden" in error_str
593 ):
594 self.rate_limit_detected[run_key] = True
595 # Send rate limit warning via WebSocket. Pass the
596 # run's owner so the socket service files this under
597 # the right per-user subscriber set -- benchmark_run_id
598 # is a per-user integer, so two users can share the
599 # same id (see SocketIOService.__subscription_key).
600 self.socket_service.emit_to_subscribers(
601 "research_progress",
602 benchmark_run_id,
603 {
604 "rate_limit_detected": True,
605 "message": "SearXNG rate limiting detected",
606 },
607 owner=username,
608 )
610 # Mark as completed in memory tracker
611 progress_info["end_time"] = datetime.now(UTC)
613 # Check if benchmark was cancelled
614 was_cancelled = (
615 run_key in self.active_runs
616 and self.active_runs[run_key].get("status") == "cancelled"
617 )
619 if was_cancelled:
620 status = BenchmarkStatus.CANCELLED
621 message = "Benchmark cancelled by user"
622 if task_id: 622 ↛ 635line 622 didn't jump to line 635 because the condition on line 622 was always true
623 self.queue_tracker.update_task_status(
624 task_id, BenchmarkTaskStatus.CANCELLED
625 )
626 else:
627 status = BenchmarkStatus.COMPLETED
628 message = "Benchmark completed successfully"
629 if task_id: 629 ↛ 635line 629 didn't jump to line 635 because the condition on line 629 was always true
630 self.queue_tracker.update_task_status(
631 task_id, BenchmarkTaskStatus.COMPLETED
632 )
634 # Store completion info for later database update
635 self.active_runs[run_key]["completion_info"] = {
636 "status": status,
637 "end_time": progress_info["end_time"],
638 "completed_examples": progress_info["completed_examples"],
639 "failed_examples": progress_info["failed_examples"],
640 }
642 # Send completion notification. Pass the owner so this reaches
643 # only their subscriber set (see rate-limit emit above for why).
644 self.socket_service.emit_to_subscribers(
645 "research_progress",
646 benchmark_run_id,
647 {
648 "status": "cancelled" if was_cancelled else "completed",
649 "message": message,
650 "progress": (
651 progress_info["completed_examples"]
652 / progress_info["total_examples"]
653 * 100
654 )
655 if progress_info["total_examples"] > 0
656 else 0,
657 "benchmark_run_id": benchmark_run_id,
658 },
659 owner=username,
660 )
662 except Exception as e:
663 logger.exception(f"Benchmark run {benchmark_run_id} failed")
664 # Update task status if we have a task_id
665 if task_id: 665 ↛ 666line 665 didn't jump to line 666 because the condition on line 665 was never true
666 self.queue_tracker.update_task_status(
667 task_id, BenchmarkTaskStatus.FAILED
668 )
669 # Store failure info for later database update
670 if run_key in self.active_runs: 670 ↛ 677line 670 didn't jump to line 677 because the condition on line 670 was always true
671 self.active_runs[run_key]["completion_info"] = {
672 "status": BenchmarkStatus.FAILED,
673 "error_message": str(e),
674 }
675 finally:
676 # Clean up active run tracking
677 if run_key in self.active_runs: 677 ↛ exitline 677 didn't return from function '_run_benchmark_thread' because the condition on line 677 was always true
678 # Mark that thread is done but keep data for database update
679 self.active_runs[run_key]["thread_complete"] = True
681 # Try to save results to database immediately if possible
682 self._sync_results_to_database(username, benchmark_run_id)
684 def _create_task_queue(
685 self,
686 datasets_config: Dict,
687 benchmark_run_id: int,
688 ) -> List[Dict]:
689 """Create list of tasks to process.
691 Each dataset config may carry a "seed" for reproducible sampling:
692 the same seed and count select the same questions on every run,
693 keeping accuracy comparable across runs. Without a seed the sample
694 is random each time.
695 """
696 tasks: List[Dict[str, Any]] = []
698 for dataset_name, config in datasets_config.items():
699 if config.get("count", 0) > 0:
700 dataset = load_dataset(
701 dataset_type=dataset_name,
702 num_examples=config["count"],
703 seed=config.get("seed"),
704 )
706 for i, example in enumerate(dataset):
707 # All registered datasets expose "problem"/"answer"
708 question = example.get("problem", "")
709 correct_answer = example.get("answer", "")
710 example_id = example.get("id", f"example_{i}")
712 # Generate query hash (keyed on the example so distinct
713 # examples with identical question text don't collide)
714 query_hash = self.generate_query_hash(
715 question, dataset_name, example_id
716 )
718 tasks.append(
719 {
720 "benchmark_run_id": benchmark_run_id,
721 "example_id": example_id,
722 "dataset_type": dataset_name,
723 "question": question,
724 "correct_answer": correct_answer,
725 "query_hash": query_hash,
726 "task_index": len(tasks),
727 }
728 )
730 return tasks
732 def _process_benchmark_task(
733 self, task: Dict, search_config: Dict, evaluation_config: Dict
734 ) -> Dict:
735 """Process a single benchmark task."""
736 try:
737 logger.info(
738 f"Starting benchmark task {task['task_index'] + 1}: "
739 f"example_id={task['example_id']}, dataset={task['dataset_type']}, "
740 f"question_preview='{task['question'][:100]}...'"
741 )
743 # Get settings context from thread-local storage
744 from ...config.thread_settings import get_settings_context
746 settings_context = get_settings_context()
748 # Generate a unique tracking ID for this benchmark task
749 import uuid
751 tracking_id = str(uuid.uuid4())
752 logger.info(
753 f"Task {task['example_id']} assigned tracking_id: {tracking_id}"
754 )
756 # Format query
757 formatted_query = format_query(
758 task["question"], task["dataset_type"]
759 )
760 logger.info(
761 f"Task {task['example_id']} formatted query: '{formatted_query[:150]}...'"
762 )
764 # Run research with progress callback for WebSocket updates
765 start_time = time.time()
766 logger.info(f"Task {task['example_id']} starting research phase...")
768 def benchmark_progress_callback(
769 status: str, progress: int, data: dict
770 ):
771 """Progress callback to emit detailed research progress via WebSocket"""
772 try:
773 timestamp = datetime.now(UTC).isoformat()
775 # Create research-compatible log entry
776 log_entry = {
777 "time": timestamp,
778 "message": f"Example {task['example_id']}: {status}",
779 "progress": progress,
780 "metadata": {
781 "phase": data.get("phase", "benchmark_processing"),
782 "type": data.get("type", "info"),
783 "example_id": task["example_id"],
784 "benchmark_run_id": task["benchmark_run_id"],
785 **data, # Include all other data
786 },
787 }
789 # Determine log type based on status/message content
790 if (
791 "complete" in status.lower()
792 or "finished" in status.lower()
793 ):
794 log_entry["metadata"]["type"] = "milestone"
795 elif (
796 "error" in status.lower() or "failed" in status.lower()
797 ):
798 log_entry["metadata"]["type"] = "error"
799 elif (
800 "starting" in status.lower()
801 or "begin" in status.lower()
802 ):
803 log_entry["metadata"]["type"] = "milestone"
805 # Create progress data in research format
806 progress_data = {
807 "progress": progress,
808 "message": status,
809 "status": "in_progress",
810 "log_entry": log_entry,
811 "progress_log": json.dumps(
812 [log_entry]
813 ), # Array format expected by socket.js
814 }
816 # Emit using research_progress format that the UI expects.
817 # Pass the owner (set on task before this callback's
818 # enclosing call, see _run_benchmark_thread) so this
819 # reaches only their subscriber set -- benchmark_run_id
820 # is a per-user integer and can collide across users.
821 self.socket_service.emit_to_subscribers(
822 "research_progress",
823 task["benchmark_run_id"],
824 progress_data,
825 owner=task.get("username"),
826 )
828 except Exception:
829 logger.exception("Error sending benchmark progress update")
831 # Get user password from task data
832 user_password = task.get("user_password")
834 search_result = quick_summary(
835 query=formatted_query,
836 research_id=tracking_id, # Pass the tracking ID
837 iterations=search_config.get("iterations", 8),
838 questions_per_iteration=search_config.get(
839 "questions_per_iteration", 5
840 ),
841 search_tool=search_config.get("search_tool", "searxng"),
842 search_strategy=search_config.get(
843 "search_strategy", "focused_iteration"
844 ),
845 progress_callback=benchmark_progress_callback,
846 model_name=search_config.get("model_name"),
847 provider=search_config.get("provider"),
848 temperature=search_config.get("temperature", 0.7),
849 openai_endpoint_url=search_config.get("openai_endpoint_url"),
850 settings_snapshot=settings_context.snapshot, # Pass settings snapshot for thread safety
851 username=task.get("username"), # Pass username
852 user_password=user_password, # Pass password for metrics tracking
853 # The web benchmark runs against the user's encrypted DB
854 # (it has username/password and wants search metrics
855 # persisted). quick_summary's default is programmatic_mode=True
856 # for true library callers; override here so the engine's
857 # metrics path stays active.
858 programmatic_mode=False,
859 )
860 processing_time = time.time() - start_time
861 logger.info(
862 f"Task {task['example_id']} research completed in {processing_time:.2f}s, "
863 f"model={search_config.get('model_name')}, provider={search_config.get('provider')}"
864 )
866 # Extract answer
867 response = search_result.get("summary", "")
868 logger.info(
869 f"Task {task['example_id']} response length: {len(response)} chars"
870 )
872 extracted_data = extract_answer_from_response(
873 response, task["dataset_type"]
874 )
875 extracted_answer = (
876 extracted_data.get("extracted_answer", "")
877 if isinstance(extracted_data, dict)
878 else str(extracted_data)
879 )
880 logger.info(
881 f"Task {task['example_id']} extracted answer: '{extracted_answer[:100]}...'"
882 )
884 # Extract sources - handle both direct sources and all_links_of_system
885 sources = search_result.get("sources", [])
886 if not sources and "all_links_of_system" in search_result:
887 sources = search_result.get("all_links_of_system", [])
889 # Log for debugging
890 logger.debug(f"Search result keys: {list(search_result.keys())}")
891 logger.debug(f"Sources found: {len(sources)} items")
893 # Prepare result
894 result = {
895 **task,
896 "response": response,
897 "extracted_answer": extracted_answer,
898 "confidence": str(
899 extracted_data.get("confidence", "100")
900 if isinstance(extracted_data, dict)
901 else "100"
902 ),
903 "processing_time": processing_time,
904 "sources": json.dumps(sources), # Convert to JSON string
905 "completed_at": datetime.now(UTC),
906 "research_id": tracking_id, # Store the UUID in the research_id field
907 }
909 # Evaluate result - requires proper grading model
910 try:
911 logger.info(f"Task {task['example_id']} starting evaluation...")
912 eval_start_time = time.time()
914 # Always attempt evaluation, regardless of provider
915 # Modern local models like Ollama are capable of grading
916 # Try to evaluate with proper model
917 result_data = {
918 "id": task["example_id"],
919 "problem": task["question"],
920 "correct_answer": task["correct_answer"],
921 "response": response,
922 "extracted_answer": extracted_answer,
923 }
925 eval_result = grade_single_result(
926 result_data,
927 task["dataset_type"],
928 evaluation_config,
929 settings_context.snapshot,
930 )
931 eval_time = time.time() - eval_start_time
932 logger.info(
933 f"Task {task['example_id']} evaluation completed in {eval_time:.2f}s"
934 )
935 if eval_result and not eval_result.get("grading_error"):
936 result.update(
937 {
938 "is_correct": eval_result.get("is_correct", False),
939 "graded_confidence": eval_result.get(
940 "graded_confidence", "0"
941 ),
942 "grader_response": eval_result.get(
943 "grader_response", ""
944 ),
945 }
946 )
947 else:
948 error_msg = (
949 eval_result.get(
950 "grading_error", "Unknown evaluation error"
951 )
952 if eval_result
953 else "No evaluation results returned"
954 )
955 result.update(
956 {
957 "is_correct": None,
958 "graded_confidence": "0",
959 "grader_response": f"Evaluation failed: {error_msg}",
960 "evaluation_error": error_msg,
961 }
962 )
964 except Exception as e:
965 logger.exception("Evaluation error")
966 result.update(
967 {
968 "is_correct": None,
969 "graded_confidence": "0",
970 "grader_response": f"Evaluation failed: {e!s}",
971 "evaluation_error": str(e),
972 }
973 )
975 return result
977 except Exception as e:
978 logger.exception("Research error")
979 return {
980 **task,
981 "research_error": str(e),
982 "completed_at": datetime.now(UTC),
983 }
985 def _persist_unsaved_results(
986 self, session, benchmark_run_id: int, run_data: dict
987 ) -> list[int]:
988 """Stage INSERTs for this run's not-yet-persisted results; return the
989 result indices that were staged.
991 Idempotent and rollback-safe. A result is skipped when its index is
992 already known-saved (``saved_indices``), its ``query_hash`` is already
993 committed for this run, or it repeats a ``query_hash`` staged earlier
994 in this same batch. Malformed result entries are also skipped so one
995 bad payload cannot make every later valid result roll back (#4861).
996 Since ``query_hash`` covers ``example_id`` as well as the question text
997 (#5080), a repeat no longer means "a dataset legitimately repeats a
998 question": it is the same result entry reaching this method twice, so
999 the skip is logged at WARNING rather than passed over in silence
1000 (#4860).
1001 Dedup correctness rests on the DB-backed ``seen_hashes``, which
1002 survives a rollback (the next sync re-reads it), NOT on any in-memory
1003 flag.
1005 This method deliberately does NOT mark ``saved_indices``: the caller
1006 must do that ONLY after its commit succeeds. Marking before the commit
1007 would, on a commit failure (disk full, lock timeout, ...), leave rows
1008 flagged saved that were actually rolled back — silently dropped
1009 forever. Combined with ``self._results_sync_lock`` held by the caller
1010 ACROSS the commit, this also stops the worker thread and the request
1011 thread from both inserting the same ``(benchmark_run_id, query_hash)``
1012 row and tripping the ``uix_run_query`` unique constraint.
1013 """
1014 results = run_data.get("results", [])
1015 # Read-only here — marking indices saved is the caller's job, and only
1016 # after the commit lands (see the method docstring).
1017 saved_indices = run_data.get("saved_indices", set())
1019 # Every query_hash already committed for this run, fetched once. This
1020 # is the rollback-safe source of truth for dedup; it is extended below
1021 # so a question repeated within this batch is staged only once.
1022 seen_hashes = {
1023 row[0]
1024 for row in session.query(BenchmarkResult.query_hash)
1025 .filter(BenchmarkResult.benchmark_run_id == benchmark_run_id)
1026 .all()
1027 }
1029 staged_indices = []
1030 for idx, result in enumerate(results):
1031 if idx in saved_indices:
1032 continue
1033 try:
1034 (
1035 example_id,
1036 query_hash,
1037 dataset_type,
1038 question,
1039 correct_answer,
1040 ) = _required_persistence_fields(result)
1041 except (KeyError, TypeError, ValueError) as exc:
1042 example_id = (
1043 result.get("example_id")
1044 if isinstance(result, dict)
1045 else None
1046 )
1047 task_index = (
1048 result.get("task_index")
1049 if isinstance(result, dict)
1050 else None
1051 )
1052 logger.warning(
1053 f"Benchmark run {benchmark_run_id}: skipping malformed "
1054 f"result index {idx} (example_id={example_id!r}, "
1055 f"task_index={task_index!r}): {type(exc).__name__}"
1056 )
1057 continue
1058 if query_hash in seen_hashes:
1059 logger.warning(
1060 f"Benchmark run {benchmark_run_id}: skipping result index "
1061 f"{idx} (example_id={example_id!r}), its "
1062 f"query_hash is already persisted or staged for this run"
1063 )
1064 continue
1065 session.add(
1066 BenchmarkResult(
1067 benchmark_run_id=benchmark_run_id,
1068 example_id=example_id,
1069 query_hash=query_hash,
1070 dataset_type=dataset_type,
1071 research_id=result.get("research_id"),
1072 question=question,
1073 correct_answer=correct_answer,
1074 response=result.get("response"),
1075 extracted_answer=result.get("extracted_answer"),
1076 confidence=result.get("confidence"),
1077 processing_time=result.get("processing_time"),
1078 sources=result.get("sources"),
1079 is_correct=result.get("is_correct"),
1080 graded_confidence=result.get("graded_confidence"),
1081 grader_response=result.get("grader_response"),
1082 completed_at=result.get("completed_at"),
1083 research_error=result.get("research_error"),
1084 evaluation_error=result.get("evaluation_error"),
1085 task_index=result.get("task_index"),
1086 )
1087 )
1088 seen_hashes.add(query_hash)
1089 staged_indices.append(idx)
1091 return staged_indices
1093 @staticmethod
1094 def _run_key(
1095 username: Optional[str], benchmark_run_id: int
1096 ) -> Tuple[str, int]:
1097 """Composite ``active_runs`` key: ``(username, BenchmarkRun.id)``.
1099 ``active_runs`` is a process-global map, but ``BenchmarkRun.id`` is a
1100 PER-USER autoincrement integer -- so two different users can each own a
1101 run with the same id (e.g. both have id == 1). Keying by the owner too
1102 means a request can only ever reach the caller's own in-flight run:
1103 one user's ``start_benchmark`` cannot clobber another's entry, and no
1104 user-facing lookup can name another user's run. ``username`` is
1105 normalized the same way ``start_benchmark`` records it so the key a run
1106 is stored under matches the key every later lookup computes.
1107 """
1108 return (username or "benchmark_user", benchmark_run_id)
1110 @staticmethod
1111 def _run_owner(run_data: Optional[Dict]) -> Optional[str]:
1112 """Username that started the given in-memory run entry, if any."""
1113 if not run_data: 1113 ↛ 1114line 1113 didn't jump to line 1114 because the condition on line 1113 was never true
1114 return None
1115 return run_data.get("data", {}).get("username")
1117 def sync_pending_results(
1118 self, benchmark_run_id: int, username: Optional[str] = None
1119 ):
1120 """Sync any pending results to database. Can be called from main thread."""
1121 # Read outside _results_sync_lock, so the worker thread's
1122 # ``del self.active_runs[run_key]`` (in _sync_results_to_database,
1123 # also unlocked) can land between the membership check and the subscript.
1124 # A plain ``active_runs[key]`` would then raise KeyError, which escapes
1125 # before the try below and surfaces as a spurious 500 on /api/results.
1126 # Fetch once so the check and the read see the same state. See #4859.
1127 # active_runs is keyed by (username, id) (see _run_key), so this lookup
1128 # can only ever return the caller's own run -- a caller naming another
1129 # user's per-user id simply misses. That keeps one user's results (and,
1130 # via the run's stored password, their credentials) unreachable here.
1131 run_data = self.active_runs.get(
1132 self._run_key(username, benchmark_run_id)
1133 )
1134 if run_data is None:
1135 return 0
1137 # For an anonymous/internal caller (username None) fall back to the
1138 # owner recorded on the run for the DB session below.
1139 if not username:
1140 username = self._run_owner(run_data)
1141 user_password = run_data.get("data", {}).get("user_password")
1143 saved_count = 0
1144 session = None
1145 from ...database.session_context import (
1146 get_user_db_session,
1147 safe_rollback,
1148 )
1150 try:
1151 # Serialize with the worker thread's _sync_results_to_database so
1152 # the two can't insert the same (benchmark_run_id, query_hash) row.
1153 with self._results_sync_lock:
1154 with get_user_db_session(username, user_password) as session:
1155 staged = self._persist_unsaved_results(
1156 session, benchmark_run_id, run_data
1157 )
1158 if staged:
1159 session.commit()
1160 # Mark saved ONLY after the commit succeeds, so a
1161 # failed commit leaves these for the next sync to retry
1162 # instead of dropping them silently.
1163 run_data.setdefault("saved_indices", set()).update(
1164 staged
1165 )
1166 saved_count = len(staged)
1167 logger.info(
1168 f"Saved {saved_count} new results for benchmark "
1169 f"{benchmark_run_id}"
1170 )
1171 # Reaching the end of the transaction proves persistence
1172 # is healthy again. Clear a prior transient failure even
1173 # when another thread already saved the pending rows.
1174 run_data.pop("result_persistence_failed", None)
1175 except Exception:
1176 logger.exception(
1177 f"Error syncing pending results for benchmark {benchmark_run_id}"
1178 )
1179 # Thread-local sessions are reused, so a failed flush/commit must be
1180 # rolled back explicitly or the next use raises PendingRollbackError.
1181 if session is not None: 1181 ↛ 1183line 1181 didn't jump to line 1183 because the condition on line 1181 was always true
1182 safe_rollback(session, "sync_pending_results")
1183 run_data["result_persistence_failed"] = True
1185 return saved_count
1187 def get_result_persistence_error(
1188 self, benchmark_run_id: int, username: Optional[str] = None
1189 ) -> Optional[Dict[str, str]]:
1190 """Return a safe client-facing result persistence error, if active."""
1191 # Keyed by (username, id) (see _run_key): a caller can only ever read
1192 # the state of their own run, never one that belongs to someone else.
1193 run_data = self.active_runs.get(
1194 self._run_key(username, benchmark_run_id)
1195 )
1196 if not run_data or not run_data.get("result_persistence_failed"):
1197 return None
1198 # Return a copy so callers cannot mutate the shared constant.
1199 return dict(_RESULT_PERSISTENCE_ERROR)
1201 def _sync_results_to_database(
1202 self, username: Optional[str], benchmark_run_id: int
1203 ):
1204 """Sync benchmark results from memory to database after thread completes."""
1205 run_key = self._run_key(username, benchmark_run_id)
1206 if run_key not in self.active_runs:
1207 return
1209 run_data = self.active_runs[run_key]
1210 if not run_data.get("thread_complete"):
1211 return
1213 username = run_data.get("data", {}).get("username")
1214 user_password = run_data.get("data", {}).get("user_password")
1215 session = None
1216 from ...database.session_context import (
1217 get_user_db_session,
1218 safe_rollback,
1219 )
1221 try:
1222 # Serialize with sync_pending_results (request thread) so the two
1223 # can't insert the same (benchmark_run_id, query_hash) row.
1224 with (
1225 self._results_sync_lock,
1226 get_user_db_session(username, user_password) as session,
1227 ):
1228 # Update benchmark run status
1229 benchmark_run = (
1230 session.query(BenchmarkRun)
1231 .filter(BenchmarkRun.id == benchmark_run_id)
1232 .first()
1233 )
1235 if benchmark_run and "completion_info" in run_data:
1236 info = run_data["completion_info"]
1237 benchmark_run.status = info["status"]
1238 benchmark_run.end_time = info.get(
1239 "end_time", datetime.now(UTC)
1240 )
1241 benchmark_run.completed_examples = info.get(
1242 "completed_examples", 0
1243 )
1244 benchmark_run.failed_examples = info.get(
1245 "failed_examples", 0
1246 )
1247 benchmark_run.error_message = info.get("error_message")
1249 # Stage any results not yet persisted (idempotent; safe
1250 # against a concurrent sync_pending_results on the request
1251 # thread — see _persist_unsaved_results).
1252 staged = self._persist_unsaved_results(
1253 session, benchmark_run_id, run_data
1254 )
1256 # Calculate final accuracy
1257 if benchmark_run.status == BenchmarkStatus.COMPLETED:
1258 correct_results = [
1259 r
1260 for r in run_data.get("results", [])
1261 if r.get("is_correct")
1262 ]
1263 evaluated_results = [
1264 r
1265 for r in run_data.get("results", [])
1266 if r.get("is_correct") is not None
1267 ]
1269 if evaluated_results:
1270 benchmark_run.overall_accuracy = (
1271 len(correct_results) / len(evaluated_results)
1272 ) * 100
1274 # Calculate processing rate
1275 total_time = sum(
1276 r.get("processing_time", 0)
1277 for r in evaluated_results
1278 )
1279 if total_time > 0: 1279 ↛ 1284line 1279 didn't jump to line 1284 because the condition on line 1279 was always true
1280 benchmark_run.processing_rate = len(
1281 evaluated_results
1282 ) / (total_time / 60)
1284 session.commit()
1285 # Mark saved only after the commit lands (see
1286 # _persist_unsaved_results) so a failed commit retries.
1287 if staged:
1288 run_data.setdefault("saved_indices", set()).update(
1289 staged
1290 )
1291 run_data.pop("result_persistence_failed", None)
1292 logger.info(
1293 f"Successfully synced results for benchmark {benchmark_run_id}"
1294 )
1296 # Clean up memory
1297 del self.active_runs[run_key]
1299 except Exception:
1300 logger.exception("Error syncing benchmark results to database")
1301 # Thread-local session is reused — roll back a failed flush/commit
1302 # so the next use doesn't raise PendingRollbackError.
1303 if session is not None:
1304 safe_rollback(session, "_sync_results_to_database")
1305 run_data["result_persistence_failed"] = True
1307 def _send_progress_update(
1308 self,
1309 benchmark_run_id: int,
1310 completed: int,
1311 total: int,
1312 username: Optional[str] = None,
1313 ):
1314 """Send real-time progress update via websocket.
1316 ``username`` is the run's owner and is passed through to
1317 ``emit_to_subscribers`` so the event reaches only that user's
1318 socket subscriber set -- ``benchmark_run_id`` is a per-user integer
1319 and can collide across users (see
1320 ``SocketIOService.__subscription_key``).
1321 """
1322 try:
1323 percentage = (completed / total * 100) if total > 0 else 0
1325 # Create log entry for milestone progress
1326 log_entry = {
1327 "time": datetime.now(UTC).isoformat(),
1328 "message": f"Completed {completed}/{total} examples ({percentage:.1f}%)",
1329 "progress": percentage,
1330 "metadata": {
1331 "phase": "benchmark_progress",
1332 "type": "milestone",
1333 "completed": completed,
1334 "total": total,
1335 "benchmark_run_id": benchmark_run_id,
1336 },
1337 }
1339 progress_data = {
1340 "status": "in_progress",
1341 "message": f"Processing examples: {completed}/{total}",
1342 "progress": percentage,
1343 "completed": completed,
1344 "total": total,
1345 "benchmark_run_id": benchmark_run_id,
1346 "log_entry": log_entry,
1347 "progress_log": json.dumps([log_entry]),
1348 }
1350 self.socket_service.emit_to_subscribers(
1351 "research_progress",
1352 benchmark_run_id,
1353 progress_data,
1354 owner=username,
1355 )
1357 except Exception:
1358 logger.exception("Error sending progress update")
1360 def _calculate_final_accuracy(
1361 self, benchmark_run_id: int, username: Optional[str] = None
1362 ):
1363 """Calculate and save final accuracy metrics."""
1364 from ...database.session_context import get_user_db_session
1366 with get_user_db_session(username) as session:
1367 try:
1368 # Get all results for this run
1369 results = (
1370 session.query(BenchmarkResult)
1371 .filter(
1372 BenchmarkResult.benchmark_run_id == benchmark_run_id
1373 )
1374 .filter(BenchmarkResult.is_correct.isnot(None))
1375 .all()
1376 )
1378 if results:
1379 correct_count = sum(1 for r in results if r.is_correct)
1380 overall_accuracy = (correct_count / len(results)) * 100
1382 # Calculate processing rate
1383 total_time = sum(r.processing_time or 0 for r in results)
1384 processing_rate = (
1385 (len(results) / (total_time / 60))
1386 if total_time > 0
1387 else 0
1388 )
1390 # Update benchmark run
1391 benchmark_run = (
1392 session.query(BenchmarkRun)
1393 .filter(BenchmarkRun.id == benchmark_run_id)
1394 .first()
1395 )
1396 if benchmark_run: 1396 ↛ exitline 1396 didn't jump to the function exit
1397 benchmark_run.overall_accuracy = overall_accuracy
1398 benchmark_run.processing_rate = processing_rate
1399 session.commit()
1401 except Exception:
1402 logger.exception("Error calculating final accuracy")
1404 def update_benchmark_status(
1405 self,
1406 benchmark_run_id: int,
1407 status: BenchmarkStatus,
1408 error_message: Optional[str] = None,
1409 username: Optional[str] = None,
1410 ):
1411 """Update benchmark run status."""
1412 from ...database.session_context import get_user_db_session
1414 with get_user_db_session(username) as session:
1415 try:
1416 benchmark_run = (
1417 session.query(BenchmarkRun)
1418 .filter(BenchmarkRun.id == benchmark_run_id)
1419 .first()
1420 )
1421 if benchmark_run:
1422 benchmark_run.status = status
1423 benchmark_run.updated_at = datetime.now(UTC)
1425 if error_message:
1426 benchmark_run.error_message = error_message
1428 if (
1429 status == BenchmarkStatus.IN_PROGRESS
1430 and not benchmark_run.start_time
1431 ):
1432 benchmark_run.start_time = datetime.now(UTC)
1433 elif (
1434 status
1435 in [BenchmarkStatus.COMPLETED, BenchmarkStatus.FAILED]
1436 and not benchmark_run.end_time
1437 ):
1438 benchmark_run.end_time = datetime.now(UTC)
1440 session.commit()
1442 except Exception:
1443 session.rollback()
1444 logger.exception("Error updating benchmark status")
1446 def get_benchmark_status(
1447 self, benchmark_run_id: int, username: str = None
1448 ) -> Optional[Dict]:
1449 """Get current status of a benchmark run."""
1450 from ...database.session_context import get_user_db_session
1452 with get_user_db_session(username) as session:
1453 try:
1454 benchmark_run = (
1455 session.query(BenchmarkRun)
1456 .filter(BenchmarkRun.id == benchmark_run_id)
1457 .first()
1458 )
1459 if not benchmark_run:
1460 return None
1462 # Calculate running accuracy from this run's evaluated results
1463 results = (
1464 session.query(BenchmarkResult)
1465 .filter(
1466 BenchmarkResult.benchmark_run_id == benchmark_run_id
1467 )
1468 .filter(BenchmarkResult.is_correct.isnot(None))
1469 .all()
1470 )
1472 running_accuracy = None
1473 # Dynamic per-dataset accuracy tracking
1474 dataset_accuracies = {}
1476 if results:
1477 # Overall running accuracy
1478 correct_count = sum(1 for r in results if r.is_correct)
1479 running_accuracy = (correct_count / len(results)) * 100
1481 # Calculate accuracy for each dataset type dynamically
1482 from collections import defaultdict
1484 dataset_results = defaultdict(list)
1486 # Group results by dataset type
1487 for r in results:
1488 dataset_results[r.dataset_type.value].append(r)
1490 # Calculate accuracy for each dataset
1491 for (
1492 dataset_type,
1493 dataset_result_list,
1494 ) in dataset_results.items():
1495 if dataset_result_list: 1495 ↛ 1491line 1495 didn't jump to line 1491 because the condition on line 1495 was always true
1496 correct = sum(
1497 1 for r in dataset_result_list if r.is_correct
1498 )
1499 accuracy = (
1500 correct / len(dataset_result_list)
1501 ) * 100
1502 # Store with _accuracy suffix for consistency
1503 dataset_accuracies[f"{dataset_type}_accuracy"] = (
1504 accuracy
1505 )
1507 # Calculate time estimates and reliability metrics
1508 estimated_time_remaining = None
1509 total_elapsed_time = None
1510 avg_time_per_example = None
1511 accuracy_confidence = None
1513 # Get ALL results for timing calculation (including those pending evaluation)
1514 all_results_for_timing = (
1515 session.query(BenchmarkResult)
1516 .filter(
1517 BenchmarkResult.benchmark_run_id == benchmark_run_id
1518 )
1519 .all()
1520 )
1522 if benchmark_run.start_time and all_results_for_timing:
1523 # Calculate elapsed time
1524 current_time = datetime.now(UTC)
1525 total_elapsed_time = (
1526 current_time - benchmark_run.start_time
1527 ).total_seconds()
1529 # Calculate average processing time per example using actual count
1530 avg_time_per_example = total_elapsed_time / len(
1531 all_results_for_timing
1532 )
1534 logger.info(
1535 f"Time calculation - elapsed: {total_elapsed_time:.2f}s, "
1536 f"results_count: {len(all_results_for_timing)}, "
1537 f"avg_per_example: {avg_time_per_example:.2f}s"
1538 )
1540 # Estimate remaining time
1541 remaining_examples = benchmark_run.total_examples - len(
1542 all_results_for_timing
1543 )
1544 if remaining_examples > 0:
1545 estimated_time_remaining = (
1546 avg_time_per_example * remaining_examples
1547 )
1548 logger.info(
1549 f"Time estimation - total: {benchmark_run.total_examples}, "
1550 f"completed: {len(all_results_for_timing)}, remaining: {remaining_examples}, "
1551 f"avg_time: {avg_time_per_example:.2f}s, "
1552 f"estimated_remaining: {estimated_time_remaining:.2f}s"
1553 )
1555 # Calculate accuracy confidence interval (Wilson score, 95%)
1556 if results and len(results) >= 3:
1557 from local_deep_research.benchmarks.metrics.statistics import (
1558 wilson_score_interval,
1559 )
1561 n = len(results)
1562 correct = sum(1 for r in results if r.is_correct)
1563 ci = wilson_score_interval(correct, n)
1564 accuracy_confidence = {
1565 "lower_bound": ci["lower"] * 100,
1566 "upper_bound": ci["upper"] * 100,
1567 "margin_of_error": ci["margin_of_error"] * 100,
1568 "sample_size": n,
1569 }
1571 status_data = {
1572 "id": benchmark_run.id,
1573 "run_name": benchmark_run.run_name,
1574 "status": benchmark_run.status.value,
1575 "completed_examples": len(
1576 all_results_for_timing
1577 ), # Use actual count from DB
1578 "total_examples": benchmark_run.total_examples,
1579 "failed_examples": benchmark_run.failed_examples,
1580 "overall_accuracy": benchmark_run.overall_accuracy
1581 or running_accuracy, # Use running accuracy if final not calculated
1582 "running_accuracy": running_accuracy, # Current running accuracy
1583 "processing_rate": benchmark_run.processing_rate,
1584 "estimated_time_remaining": estimated_time_remaining, # seconds
1585 "total_elapsed_time": total_elapsed_time, # seconds
1586 "avg_time_per_example": avg_time_per_example, # seconds
1587 "accuracy_confidence": accuracy_confidence, # confidence interval
1588 "created_at": benchmark_run.created_at.isoformat()
1589 if benchmark_run.created_at
1590 else None,
1591 "start_time": benchmark_run.start_time.isoformat()
1592 if benchmark_run.start_time
1593 else None,
1594 "end_time": benchmark_run.end_time.isoformat()
1595 if benchmark_run.end_time
1596 else None,
1597 # Genericized at this boundary (CWE-209) — see
1598 # _GENERIC_BENCHMARK_ERROR. Returns None when there is no
1599 # failure so non-error states are unaffected.
1600 "error_message": (
1601 _GENERIC_BENCHMARK_ERROR
1602 if benchmark_run.error_message
1603 else None
1604 ),
1605 # Add all per-dataset accuracies dynamically
1606 **dataset_accuracies,
1607 }
1609 logger.info(
1610 f"Benchmark {benchmark_run_id} status - completed: {benchmark_run.completed_examples}, "
1611 f"running_acc: {running_accuracy}, dataset_accuracies: {dataset_accuracies}, "
1612 f"avg_time: {avg_time_per_example}"
1613 )
1615 return status_data
1617 except Exception:
1618 logger.exception("Error getting benchmark status")
1619 return None
1621 def cancel_benchmark(
1622 self, benchmark_run_id: int, username: Optional[str] = None
1623 ) -> bool:
1624 """Cancel a running benchmark."""
1625 try:
1626 # active_runs is keyed by (username, id) (see _run_key), so this
1627 # lookup only ever finds the requester's own run -- one user can no
1628 # longer set another user's in-flight run to "cancelled" and stop
1629 # it (cross-user DoS). The DB update below is already scoped to the
1630 # caller's own database.
1631 run = self.active_runs.get(
1632 self._run_key(username, benchmark_run_id)
1633 )
1634 if run is not None:
1635 run["status"] = "cancelled"
1637 self.update_benchmark_status(
1638 benchmark_run_id, BenchmarkStatus.CANCELLED, username=username
1639 )
1640 logger.info(f"Cancelled benchmark run {benchmark_run_id}")
1641 return True
1643 except Exception:
1644 logger.exception(f"Error cancelling benchmark {benchmark_run_id}")
1645 return False
1648# Global service instance
1649benchmark_service = BenchmarkService()