Coverage for src/local_deep_research/benchmarks/web_api/benchmark_service.py: 95%
543 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +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 ...exceptions import ResearchTerminatedException
16from ...settings.exceptions import SettingsValueError
17from ...settings.manager import SnapshotSettingsContext
18from ...web.services.research_service import _global_research_semaphore
19from ...database.models.benchmark import (
20 BenchmarkResult,
21 BenchmarkRun,
22 BenchmarkStatus,
23 DatasetType,
24)
25from ..datasets import load_dataset
26from ..graders import extract_answer_from_response, grade_single_result
27from ..runners import format_query
28from ...database.thread_local_session import thread_cleanup
30# Generic failure message surfaced to API clients in place of the raw benchmark
31# exception. ``benchmark_run.error_message`` is set from ``str(e)`` of the
32# benchmark run (LLM calls, search engines, grading — the same stack as
33# research) and returned by get_benchmark_status() via
34# GET /benchmark/api/status/<id>. The raw text can carry server-level LLM
35# endpoints / hosts / filesystem paths (settings_snapshot includes LDR_* env
36# overrides), so it must not reach the client (CWE-209). Full detail stays in
37# the logger.exception calls at the failure sites.
38_GENERIC_BENCHMARK_ERROR = (
39 "Benchmark run failed due to an internal error. "
40 "Check the server logs for details."
41)
42_RESULT_PERSISTENCE_ERROR = {
43 "code": "database_write_failed",
44 "message": (
45 "Benchmark results could not be saved. Check server logs and "
46 "database storage, then retry."
47 ),
48}
51class SocketIOService:
52 """Adapter over socketio_asgi's module-level emit functions.
54 Keeps the Flask-era ``SocketIOService`` interface (which
55 BenchmarkService and its tests are written against) while delegating
56 to the FastAPI-native emit path. An earlier port made this a no-op
57 stub, silently dropping every benchmark progress / completion /
58 rate-limit WebSocket event.
60 ``owner`` is keyword-only and REQUIRED here, where the upstream
61 ``SocketIOService.emit_to_subscribers`` takes an optional
62 ``username=None``. Same fix, stricter contract: a benchmark id is a
63 per-user integer, so an omitted owner silently files the event under
64 the wrong subscriber set. Making it required turns that into a
65 TypeError at the call site instead of a misrouted event.
66 """
68 def emit_to_subscribers(
69 self, event_base, research_id, data, *, owner, enable_logging=True
70 ):
71 from ...web.services.socketio_asgi import emit_to_subscribers
73 return emit_to_subscribers(
74 event_base,
75 research_id,
76 data,
77 owner=owner,
78 enable_logging=enable_logging,
79 )
81 def emit_to_room(self, event, data, room=None):
82 from ...web.services.socketio_asgi import emit_socket_event
84 return emit_socket_event(event, data, room=room)
87def _required_persistence_fields(
88 result: Any,
89) -> tuple[Any, str, DatasetType, Any, Any]:
90 """Extract required fields or raise a data-shape error."""
91 if not isinstance(result, dict): 91 ↛ 92line 91 didn't jump to line 92 because the condition on line 91 was never true
92 raise TypeError("result must be a mapping")
93 query_hash = result["query_hash"]
94 if not isinstance(query_hash, str) or not query_hash: 94 ↛ 95line 94 didn't jump to line 95 because the condition on line 94 was never true
95 raise ValueError("query_hash must be a non-empty string")
96 return (
97 result["example_id"],
98 query_hash,
99 DatasetType(result["dataset_type"]),
100 result["question"],
101 result["correct_answer"],
102 )
105class BenchmarkTaskStatus(Enum):
106 """Status values for benchmark tasks in the queue tracker."""
108 QUEUED = "queued"
109 PROCESSING = "processing"
110 COMPLETED = "completed"
111 FAILED = "failed"
112 CANCELLED = "cancelled"
115class BenchmarkQueueTracker:
116 """Simple in-memory tracker for benchmark queue status.
118 This replaces the removed memory_queue functionality for benchmarks.
119 Since benchmarks are temporary and don't need persistence,
120 this simple in-memory solution is sufficient.
122 Thread-safe for concurrent access from multiple benchmark threads.
123 """
125 def __init__(self):
126 self.tasks = {}
127 self._lock = threading.Lock()
129 def add_task(
130 self, task_id: str, username: str, task_type: str = "benchmark"
131 ):
132 """Add a new task to tracking.
134 Also performs opportunistic cleanup of old completed tasks.
135 """
136 # Cleanup old tasks before adding new one (outside lock for better performance)
137 self.cleanup_completed_tasks()
139 with self._lock:
140 self.tasks[task_id] = {
141 "username": username,
142 "task_type": task_type,
143 "status": BenchmarkTaskStatus.QUEUED.value,
144 "created_at": datetime.now(UTC),
145 }
147 def update_task_status(self, task_id: str, status: BenchmarkTaskStatus):
148 """Update the status of a task."""
149 with self._lock:
150 if task_id in self.tasks:
151 self.tasks[task_id]["status"] = status.value
152 self.tasks[task_id]["updated_at"] = datetime.now(UTC)
153 else:
154 logger.warning(
155 f"Attempted to update status for non-existent task: {task_id}"
156 )
158 def get_task_status(self, task_id: str) -> Optional[Dict]:
159 """Get the current status of a task."""
160 with self._lock:
161 return self.tasks.get(task_id)
163 def remove_task(self, task_id: str):
164 """Remove a task from tracking."""
165 with self._lock:
166 self.tasks.pop(task_id, None)
168 def cleanup_completed_tasks(self, max_age_seconds: int = 3600):
169 """Remove completed tasks older than max_age_seconds.
171 Args:
172 max_age_seconds: Maximum age in seconds for completed tasks (default 1 hour)
173 """
174 with self._lock:
175 now = datetime.now(UTC)
176 to_remove = []
177 for task_id, task_data in self.tasks.items():
178 # Only cleanup completed, failed, or cancelled tasks
179 if task_data["status"] in [
180 BenchmarkTaskStatus.COMPLETED.value,
181 BenchmarkTaskStatus.FAILED.value,
182 BenchmarkTaskStatus.CANCELLED.value,
183 ]:
184 # Check if task has updated_at timestamp
185 updated_at = task_data.get(
186 "updated_at", task_data.get("created_at")
187 )
188 if updated_at: 188 ↛ 177line 188 didn't jump to line 177 because the condition on line 188 was always true
189 age = (now - updated_at).total_seconds()
190 if age > max_age_seconds:
191 to_remove.append(task_id)
193 for task_id in to_remove:
194 self.tasks.pop(task_id, None)
195 logger.debug(f"Cleaned up old task: {task_id}")
197 if to_remove:
198 logger.info(f"Cleaned up {len(to_remove)} old benchmark tasks")
201class BenchmarkService:
202 """Service for managing benchmark runs through the web interface."""
204 def __init__(self, socket_service=None):
205 # Keyed by (username, BenchmarkRun.id). BenchmarkRun.id is a PER-USER
206 # autoincrement, so two users can each own run id 1; keying by the
207 # owner as well prevents one user's start_benchmark from clobbering
208 # another's in-memory run (and the cross-user reads that enabled).
209 self.active_runs: Dict[Tuple[str, int], Dict] = {}
210 self.socket_service = socket_service or SocketIOService()
211 # Track rate limiting per user benchmark run.
212 self.rate_limit_detected: Dict[Tuple[str, int], bool] = {}
213 self.queue_tracker = BenchmarkQueueTracker() # Initialize queue tracker
214 # Serializes benchmark-result persistence. _sync_results_to_database
215 # runs on the worker thread while sync_pending_results runs on the
216 # request thread; without this they can both INSERT the same
217 # (benchmark_run_id, query_hash) row and trip the uix_run_query
218 # unique constraint, which rolls back the whole pending batch.
219 self._results_sync_lock = threading.Lock()
221 def generate_config_hash(self, search_config: Dict[str, Any]) -> str:
222 """Generate a hash for search configuration compatibility checking."""
223 relevant_params = {
224 "iterations": search_config.get("iterations"),
225 "questions_per_iteration": search_config.get(
226 "questions_per_iteration"
227 ),
228 "search_tool": search_config.get("search_tool"),
229 "search_strategy": search_config.get("search_strategy"),
230 "model_name": search_config.get("model_name"),
231 "provider": search_config.get("provider"),
232 }
233 # Remove None values
234 relevant_params = {
235 k: v for k, v in relevant_params.items() if v is not None
236 }
237 config_str = json.dumps(relevant_params, sort_keys=True)
238 return hashlib.md5( # DevSkim: ignore DS126858
239 config_str.encode(), usedforsecurity=False
240 ).hexdigest()[:8]
242 def generate_query_hash(
243 self, question: str, dataset_type: str, example_id: str
244 ) -> str:
245 """Generate a hash for a query to enable deduplication.
247 The example id is part of the key so two distinct examples that share
248 the same question text (and dataset) get distinct hashes and each
249 persist a row, instead of colliding on the ``uix_run_query`` unique
250 constraint and dropping one. Re-hashing the same example is stable, so
251 the sync dedup still collapses a genuine re-sync of that example.
252 """
253 query_content = (
254 f"{example_id}|{question.strip()}|{dataset_type.lower()}"
255 )
256 return hashlib.md5( # DevSkim: ignore DS126858
257 query_content.encode(), usedforsecurity=False
258 ).hexdigest()
260 def create_benchmark_run(
261 self,
262 run_name: Optional[str],
263 search_config: Dict[str, Any],
264 evaluation_config: Dict[str, Any],
265 datasets_config: Dict[str, Dict],
266 username: Optional[str] = None,
267 user_password: Optional[str] = None,
268 ) -> int:
269 """Create a new benchmark run in the database."""
270 from ...database.session_context import get_user_db_session
272 with get_user_db_session(username, user_password) as session:
273 try:
274 config_hash = self.generate_config_hash(search_config)
276 # Calculate total examples
277 total_examples = sum(
278 config.get("count", 0)
279 for config in datasets_config.values()
280 )
282 benchmark_run = BenchmarkRun(
283 run_name=run_name,
284 config_hash=config_hash,
285 query_hash_list=[], # Will be populated as we process
286 search_config=search_config,
287 evaluation_config=evaluation_config,
288 datasets_config=datasets_config,
289 total_examples=total_examples,
290 status=BenchmarkStatus.PENDING,
291 )
293 session.add(benchmark_run)
294 session.commit()
296 logger.info(
297 f"Created benchmark run {benchmark_run.id} with config hash {config_hash}"
298 )
299 return int(benchmark_run.id)
301 except Exception:
302 session.rollback()
303 logger.exception("Error creating benchmark run")
304 raise
306 def start_benchmark(
307 self,
308 benchmark_run_id: int,
309 username: Optional[str] = None,
310 user_password: Optional[str] = None,
311 ) -> bool:
312 """Start a benchmark run in a background thread."""
313 from ...database.session_context import get_user_db_session
315 try:
316 # Get all data from the database in the main thread
317 # This avoids database access from the background thread
318 with get_user_db_session(username, user_password) as session:
319 # Get benchmark run details
320 benchmark_run = (
321 session.query(BenchmarkRun)
322 .filter(BenchmarkRun.id == benchmark_run_id)
323 .first()
324 )
325 if not benchmark_run:
326 raise ValueError( # noqa: TRY301 — caught by except, sets FAILED status in DB
327 f"Benchmark run {benchmark_run_id} not found"
328 )
330 # Create settings snapshot for thread safety + provenance.
331 # The settings snapshot is BEST-EFFORT provenance, not a
332 # critical path: a recoverable read failure must not break
333 # the user's ability to start a benchmark. We catch only the
334 # failures that can legitimately escape get_all_settings:
335 # SQLAlchemyError (transient DB issues) and LookupError /
336 # ValueError (a stale enum-typed settings row, e.g. a legacy
337 # 'CHAT' value that fails enum coercion). Those degrade to an
338 # empty snapshot — ldr_version is still recorded. Anything
339 # else (a genuine programming error) is deliberately NOT
340 # swallowed here: it propagates to the outer handler, which
341 # logs a full traceback and marks the run FAILED — the right
342 # outcome, since a settings subsystem that breaks unexpectedly
343 # would also corrupt the benchmark's own config reads.
344 from local_deep_research.settings import SettingsManager
346 settings_manager = SettingsManager(session)
347 try:
348 settings_snapshot = settings_manager.get_all_settings()
349 except SettingsValueError:
350 # Invalid configuration cannot run with an empty snapshot.
351 raise
352 except (SQLAlchemyError, LookupError, ValueError):
353 # Use logger.exception so the traceback goes to stderr
354 # without splatting `repr(exc)` into the message body —
355 # exception details may include user data.
356 logger.exception(
357 "Failed to capture settings snapshot for benchmark "
358 f"{benchmark_run_id}; running with empty snapshot."
359 )
360 settings_snapshot = {}
362 # Get user password for metrics tracking in background thread
363 from ...utilities.request_context import (
364 get_current_session_id,
365 )
366 from ...database.session_passwords import session_password_store
368 _user_password = None
369 session_id = get_current_session_id()
370 if session_id and username:
371 _user_password = (
372 session_password_store.get_session_password(
373 username, session_id
374 )
375 )
376 if not _user_password: 376 ↛ 377line 376 didn't jump to line 377 because the condition on line 376 was never true
377 logger.warning(
378 f"No password found for user {username} in current session"
379 )
381 # Extract all data we need
382 benchmark_data = {
383 "benchmark_run_id": benchmark_run_id,
384 "username": username or "benchmark_user",
385 "user_password": _user_password, # Add password for metrics tracking
386 "config_hash": benchmark_run.config_hash,
387 "datasets_config": benchmark_run.datasets_config,
388 "search_config": benchmark_run.search_config,
389 "evaluation_config": benchmark_run.evaluation_config,
390 # The grader later calls get_llm using only this snapshot.
391 # Scope it to the authenticated owner so ADAPTIVE can see
392 # their private retriever primary and refuse cloud egress.
393 "settings_snapshot": (
394 {**settings_snapshot, "_username": username}
395 if username
396 else settings_snapshot
397 ),
398 }
400 # Update status in database. Persist provenance (ldr_version
401 # + redacted settings_snapshot) BEFORE the commit so the
402 # YAML download can later reproduce the run; the snapshot
403 # already exists in memory from line 306, we just store the
404 # redacted form so secrets don't sit in the DB even though
405 # SQLCipher protects at rest.
406 from local_deep_research import __version__
407 from local_deep_research.security.data_sanitizer import (
408 DataSanitizer,
409 )
411 benchmark_run.ldr_version = __version__
412 benchmark_run.settings_snapshot = (
413 DataSanitizer.redact_settings_snapshot(settings_snapshot)
414 )
415 benchmark_run.status = BenchmarkStatus.IN_PROGRESS
416 benchmark_run.start_time = datetime.now(UTC)
417 session.commit()
419 # Store data in memory for the thread. Key by (username, id) so a
420 # run another user owns under the same per-user id can't overwrite
421 # this entry (see _run_key).
422 run_key = self._run_key(username, benchmark_run_id)
423 self.active_runs[run_key] = {
424 "data": benchmark_data,
425 "start_time": datetime.now(UTC),
426 "status": "running",
427 "results": [],
428 }
430 # Start background thread. Pass the owner so the thread rebuilds the
431 # same composite key for its own active_runs accesses.
432 thread = threading.Thread(
433 target=self._run_benchmark_thread,
434 args=(username, benchmark_run_id),
435 daemon=True,
436 )
437 # Publish the handle before the worker can finish and remove its
438 # active_runs entry during result synchronization.
439 self.active_runs[run_key]["thread"] = thread
440 thread.start()
442 logger.info(f"Started benchmark run {benchmark_run_id}")
443 return True
445 except Exception as e:
446 logger.exception(f"Error starting benchmark {benchmark_run_id}")
447 # If we populated active_runs before the spawn failed, drop the
448 # stale entry — it has no thread and would mislead subsequent
449 # cancel_benchmark / get_run_status calls. Recompute the key here
450 # (the local run_key may not exist yet if we failed before it).
451 self.active_runs.pop(
452 self._run_key(username, benchmark_run_id), None
453 )
454 # Update status using user database
455 with get_user_db_session(username, user_password) as session:
456 benchmark_run = (
457 session.query(BenchmarkRun)
458 .filter(BenchmarkRun.id == benchmark_run_id)
459 .first()
460 )
461 if benchmark_run:
462 benchmark_run.status = BenchmarkStatus.FAILED
463 benchmark_run.error_message = str(e)
464 session.commit()
465 return False
467 @thread_cleanup
468 def _run_benchmark_thread(
469 self, username: Optional[str], benchmark_run_id: int
470 ):
471 """Main benchmark execution thread."""
472 # IMPORTANT: This runs in a background thread, so we cannot access the user database
473 # Using in-memory queue tracker for benchmark status tracking
475 task_id = None
476 # Rebuild the same composite key start_benchmark stored this run under.
477 run_key = self._run_key(username, benchmark_run_id)
479 # Get the benchmark data that was passed to us
480 # We need to retrieve this from the service database or from memory
481 benchmark_data = self.active_runs.get(run_key, {}).get("data")
483 try:
484 if not benchmark_data:
485 raise ValueError( # noqa: TRY301
486 f"Benchmark data for run {benchmark_run_id} not found"
487 )
488 # Set up settings context for thread-local access
489 settings_snapshot = benchmark_data.get("settings_snapshot", {})
490 # ``or`` (not ``.get(..., default)``) so an explicit None still
491 # falls back — symmetric with _run_key's normalization.
492 username = benchmark_data.get("username") or "benchmark_user"
494 # Create a settings context that threads can use
495 settings_context = SnapshotSettingsContext(
496 settings_snapshot,
497 username=username,
498 missing_key_log_level="WARNING",
499 )
501 # Set the context in thread-local storage
502 from ...config.thread_settings import set_settings_context
504 set_settings_context(settings_context)
506 # Extract all the data we need
507 datasets_config = benchmark_data["datasets_config"]
508 search_config = benchmark_data["search_config"]
509 evaluation_config = benchmark_data["evaluation_config"]
511 # Create task queue
512 task_queue = self._create_task_queue(
513 datasets_config,
514 benchmark_run_id,
515 )
517 # Calculate totals
518 total_examples = len(task_queue)
519 completed_examples = 0
521 # Initialize task tracking
522 task_id = f"benchmark_{benchmark_run_id}_{int(datetime.now(UTC).timestamp())}"
523 username = benchmark_data.get("username") or "benchmark_user"
524 self.queue_tracker.add_task(task_id, username, "benchmark")
525 self.queue_tracker.update_task_status(
526 task_id, BenchmarkTaskStatus.PROCESSING
527 )
529 # Track progress in memory
530 progress_info = {
531 "total_examples": total_examples,
532 "completed_examples": completed_examples,
533 "failed_examples": 0,
534 "start_time": datetime.now(UTC),
535 }
537 # Process tasks
538 logger.info(
539 f"Benchmark {benchmark_run_id} starting to process {len(task_queue)} tasks"
540 )
541 for i, task in enumerate(task_queue):
542 # Check if benchmark has been cancelled
543 if (
544 run_key in self.active_runs
545 and self.active_runs[run_key].get("status") == "cancelled"
546 ):
547 logger.info(
548 f"Benchmark {benchmark_run_id} was cancelled, stopping processing"
549 )
550 break
552 logger.info(
553 f"Benchmark {benchmark_run_id} processing task {i + 1}/{len(task_queue)}"
554 )
555 try:
556 # Add username and password to task for metrics tracking.
557 # ``or`` keeps the fallback symmetric with _run_key: a
558 # missing/None owner still maps to "benchmark_user" so
559 # the publish below hits the same entry cancel reads.
560 task["username"] = benchmark_data.get("username") or (
561 "benchmark_user"
562 )
563 task["user_password"] = benchmark_data.get("user_password")
565 # Acquire the global research semaphore so benchmark
566 # tasks count against the server-wide concurrency limit
567 _global_research_semaphore.acquire()
568 try:
569 # Re-check after the semaphore wait: Cancel may have
570 # landed while blocked, when there was no tracking ID
571 # to flag. Starting the task now would publish a fresh
572 # ID the finished handshake cannot see and run to
573 # completion after Cancel.
574 if (
575 run_key in self.active_runs
576 and self.active_runs[run_key].get("status")
577 == "cancelled"
578 ):
579 logger.info(
580 f"Benchmark {benchmark_run_id} was cancelled "
581 f"while waiting for capacity, skipping task "
582 f"{i + 1}/{len(task_queue)}"
583 )
584 break
585 # Process single task
586 result = self._process_benchmark_task(
587 task,
588 search_config,
589 evaluation_config,
590 )
591 except ResearchTerminatedException:
592 # Cancel arrived mid-task: the in-flight LLM call
593 # aborted via its termination flag (published as
594 # current_tracking_id, see _process_benchmark_task).
595 # Stop the run without counting the aborted task as
596 # failed — the was_cancelled check below still marks
597 # the run CANCELLED.
598 logger.info(
599 f"Benchmark {benchmark_run_id} task {i + 1} "
600 f"aborted mid-task on cancel; stopping run"
601 )
602 break
603 finally:
604 _global_research_semaphore.release()
605 # Drop the finished task's research scope and flag
606 # so one task's cancellation state can never leak
607 # into the next task's LLM calls.
608 finished_entry = self.active_runs.get(run_key)
609 finished_tid = (
610 finished_entry.pop("current_tracking_id", None)
611 if finished_entry is not None
612 else None
613 )
614 if finished_tid:
615 try:
616 from ...web.research_state import (
617 clear_termination_flag,
618 )
620 clear_termination_flag(finished_tid)
621 except Exception:
622 logger.debug(
623 "Failed to clear benchmark task flag "
624 "(non-critical)"
625 )
627 # Store result in memory for now (will be saved later)
628 if "results" not in self.active_runs[run_key]: 628 ↛ 629line 628 didn't jump to line 629 because the condition on line 628 was never true
629 self.active_runs[run_key]["results"] = []
630 self.active_runs[run_key]["results"].append(result)
632 # Update progress
633 progress_info["completed_examples"] += 1
635 logger.info(
636 f"Benchmark {benchmark_run_id} task {i + 1}/{len(task_queue)} completed successfully. "
637 f"Progress: {progress_info['completed_examples']}/{progress_info['total_examples']} total examples"
638 )
640 # Send real-time update
641 self._send_progress_update(
642 benchmark_run_id,
643 progress_info["completed_examples"],
644 progress_info["total_examples"],
645 username=username,
646 )
648 except Exception as e:
649 logger.exception(f"Error processing task {i}")
650 progress_info["failed_examples"] += 1
651 logger.info(
652 f"Benchmark {benchmark_run_id} task {i + 1}/{len(task_queue)} failed. "
653 f"Total failed: {progress_info['failed_examples']}"
654 )
656 # Check if this is a rate limiting error
657 error_str = str(e).lower()
658 if (
659 "403" in error_str
660 or "rate limit" in error_str
661 or "forbidden" in error_str
662 ):
663 self.rate_limit_detected[run_key] = True
664 # Send rate limit warning via WebSocket. Pass the
665 # run's owner so the socket service files this under
666 # the right per-user subscriber set -- benchmark_run_id
667 # is a per-user integer, so two users can share the
668 # same id (see SocketIOService.__subscription_key).
669 self.socket_service.emit_to_subscribers(
670 "research_progress",
671 benchmark_run_id,
672 {
673 "rate_limit_detected": True,
674 "message": "SearXNG rate limiting detected",
675 },
676 owner=username,
677 )
679 # Mark as completed in memory tracker
680 progress_info["end_time"] = datetime.now(UTC)
682 # Check if benchmark was cancelled
683 was_cancelled = (
684 run_key in self.active_runs
685 and self.active_runs[run_key].get("status") == "cancelled"
686 )
688 if was_cancelled:
689 status = BenchmarkStatus.CANCELLED
690 message = "Benchmark cancelled by user"
691 if task_id: 691 ↛ 704line 691 didn't jump to line 704 because the condition on line 691 was always true
692 self.queue_tracker.update_task_status(
693 task_id, BenchmarkTaskStatus.CANCELLED
694 )
695 else:
696 status = BenchmarkStatus.COMPLETED
697 message = "Benchmark completed successfully"
698 if task_id: 698 ↛ 704line 698 didn't jump to line 704 because the condition on line 698 was always true
699 self.queue_tracker.update_task_status(
700 task_id, BenchmarkTaskStatus.COMPLETED
701 )
703 # Store completion info for later database update
704 self.active_runs[run_key]["completion_info"] = {
705 "status": status,
706 "end_time": progress_info["end_time"],
707 "completed_examples": progress_info["completed_examples"],
708 "failed_examples": progress_info["failed_examples"],
709 }
711 # Send completion notification. Pass the owner so this reaches
712 # only their subscriber set (see rate-limit emit above for why).
713 self.socket_service.emit_to_subscribers(
714 "research_progress",
715 benchmark_run_id,
716 {
717 "status": "cancelled" if was_cancelled else "completed",
718 "message": message,
719 "progress": (
720 progress_info["completed_examples"]
721 / progress_info["total_examples"]
722 * 100
723 )
724 if progress_info["total_examples"] > 0
725 else 0,
726 "benchmark_run_id": benchmark_run_id,
727 },
728 owner=username,
729 )
731 except Exception as e:
732 logger.exception(f"Benchmark run {benchmark_run_id} failed")
733 # Update task status if we have a task_id
734 if task_id: 734 ↛ 735line 734 didn't jump to line 735 because the condition on line 734 was never true
735 self.queue_tracker.update_task_status(
736 task_id, BenchmarkTaskStatus.FAILED
737 )
738 # Store failure info for later database update
739 if run_key in self.active_runs: 739 ↛ 746line 739 didn't jump to line 746 because the condition on line 739 was always true
740 self.active_runs[run_key]["completion_info"] = {
741 "status": BenchmarkStatus.FAILED,
742 "error_message": str(e),
743 }
744 finally:
745 # Clean up active run tracking
746 if run_key in self.active_runs: 746 ↛ exitline 746 didn't return from function '_run_benchmark_thread' because the condition on line 746 was always true
747 # Mark that thread is done but keep data for database update
748 self.active_runs[run_key]["thread_complete"] = True
750 # Try to save results to database immediately if possible
751 self._sync_results_to_database(username, benchmark_run_id)
753 def _create_task_queue(
754 self,
755 datasets_config: Dict,
756 benchmark_run_id: int,
757 ) -> List[Dict]:
758 """Create list of tasks to process.
760 Each dataset config may carry a "seed" for reproducible sampling:
761 the same seed and count select the same questions on every run,
762 keeping accuracy comparable across runs. Without a seed the sample
763 is random each time.
764 """
765 tasks: List[Dict[str, Any]] = []
767 for dataset_name, config in datasets_config.items():
768 if config.get("count", 0) > 0:
769 # This is the local benchmark registry factory, not a Hugging Face download API.
770 dataset = load_dataset( # nosec B615
771 dataset_type=dataset_name,
772 num_examples=config["count"],
773 seed=config.get("seed"),
774 )
776 for i, example in enumerate(dataset):
777 # All registered datasets expose "problem"/"answer"
778 question = example.get("problem", "")
779 correct_answer = example.get("answer", "")
780 example_id = example.get("id", f"example_{i}")
782 # Generate query hash (keyed on the example so distinct
783 # examples with identical question text don't collide)
784 query_hash = self.generate_query_hash(
785 question, dataset_name, example_id
786 )
788 tasks.append(
789 {
790 "benchmark_run_id": benchmark_run_id,
791 "example_id": example_id,
792 "dataset_type": dataset_name,
793 "question": question,
794 "correct_answer": correct_answer,
795 "query_hash": query_hash,
796 "task_index": len(tasks),
797 }
798 )
800 return tasks
802 def _process_benchmark_task(
803 self, task: Dict, search_config: Dict, evaluation_config: Dict
804 ) -> Dict:
805 """Process a single benchmark task."""
806 try:
807 logger.info(
808 f"Starting benchmark task {task['task_index'] + 1}: "
809 f"example_id={task['example_id']}, dataset={task['dataset_type']}, "
810 f"question_preview='{task['question'][:100]}...'"
811 )
813 # Get settings context from thread-local storage
814 from ...config.thread_settings import get_settings_context
816 settings_context = get_settings_context()
818 # Generate a unique tracking ID for this benchmark task
819 import uuid
821 tracking_id = str(uuid.uuid4())
822 logger.info(
823 f"Task {task['example_id']} assigned tracking_id: {tracking_id}"
824 )
826 # Publish the in-flight research scope so cancel_benchmark can
827 # abort this task's LLM calls mid-generation (not just between
828 # tasks). quick_summary binds this id to its LLM (explicit
829 # research_id) and thread context, and the grading phase below
830 # re-arms the same scope — the termination flag then aborts the
831 # current call within one streamed chunk. Cleared by the worker
832 # loop after every task, so ids never outlive their task.
833 # The key must match the worker loop's run_key: task["username"]
834 # is set from benchmark_data (with the same "benchmark_user"
835 # fallback) before this call, so web runs always hit. A direct
836 # caller that bypasses the loop with a mismatched username
837 # misses here instead of publishing to another user's same-id
838 # run (BenchmarkRun.id is per-user) — log so the trap is visible.
839 tracker_entry = self.active_runs.get(
840 self._run_key(task.get("username"), task["benchmark_run_id"])
841 )
842 if tracker_entry is None:
843 logger.warning(
844 f"Benchmark {task['benchmark_run_id']} task "
845 f"{task['example_id']} has no live run entry for "
846 f"owner {task.get('username')!r}; skipping scope publish "
847 f"(cancel will only stop between tasks)"
848 )
849 else:
850 tracker_entry["current_tracking_id"] = tracking_id
851 # Cooperative handoff: Cancel may have landed between the
852 # worker's status checks and this publish, seeing no ID to
853 # flag. If the run is already cancelled, arm the flag for the
854 # just-published ID and abort before any LLM call starts, so
855 # neither ordering loses the stop request.
856 if tracker_entry.get("status") == "cancelled":
857 try:
858 from ...web.research_state import (
859 set_termination_flag,
860 )
862 set_termination_flag(tracking_id)
863 except Exception:
864 logger.debug(
865 "Failed to set benchmark termination flag "
866 "(non-critical)"
867 )
868 raise ResearchTerminatedException( # noqa: TRY301
869 "Benchmark cancelled before task started"
870 )
872 # Format query
873 formatted_query = format_query(
874 task["question"], task["dataset_type"]
875 )
876 logger.info(
877 f"Task {task['example_id']} formatted query: '{formatted_query[:150]}...'"
878 )
880 # Run research with progress callback for WebSocket updates
881 start_time = time.time()
882 logger.info(f"Task {task['example_id']} starting research phase...")
884 def benchmark_progress_callback(
885 status: str, progress: int, data: dict
886 ):
887 """Progress callback to emit detailed research progress via WebSocket"""
888 try:
889 timestamp = datetime.now(UTC).isoformat()
891 # Create research-compatible log entry
892 log_entry = {
893 "time": timestamp,
894 "message": f"Example {task['example_id']}: {status}",
895 "progress": progress,
896 "metadata": {
897 "phase": data.get("phase", "benchmark_processing"),
898 "type": data.get("type", "info"),
899 "example_id": task["example_id"],
900 "benchmark_run_id": task["benchmark_run_id"],
901 **data, # Include all other data
902 },
903 }
905 # Determine log type based on status/message content
906 if (
907 "complete" in status.lower()
908 or "finished" in status.lower()
909 ):
910 log_entry["metadata"]["type"] = "milestone"
911 elif (
912 "error" in status.lower() or "failed" in status.lower()
913 ):
914 log_entry["metadata"]["type"] = "error"
915 elif (
916 "starting" in status.lower()
917 or "begin" in status.lower()
918 ):
919 log_entry["metadata"]["type"] = "milestone"
921 # Create progress data in research format
922 progress_data = {
923 "progress": progress,
924 "message": status,
925 "status": "in_progress",
926 "log_entry": log_entry,
927 "progress_log": json.dumps(
928 [log_entry]
929 ), # Array format expected by socket.js
930 }
932 # Emit using research_progress format that the UI expects.
933 # Pass the owner (set on task before this callback's
934 # enclosing call, see _run_benchmark_thread) so this
935 # reaches only their subscriber set -- benchmark_run_id
936 # is a per-user integer and can collide across users.
937 self.socket_service.emit_to_subscribers(
938 "research_progress",
939 task["benchmark_run_id"],
940 progress_data,
941 owner=task.get("username"),
942 )
944 except Exception:
945 logger.exception("Error sending benchmark progress update")
947 # Get user password from task data
948 user_password = task.get("user_password")
950 search_result = quick_summary(
951 query=formatted_query,
952 research_id=tracking_id, # Pass the tracking ID
953 iterations=search_config.get("iterations", 8),
954 questions_per_iteration=search_config.get(
955 "questions_per_iteration", 5
956 ),
957 search_tool=search_config.get("search_tool", "searxng"),
958 search_strategy=search_config.get(
959 "search_strategy", "focused_iteration"
960 ),
961 progress_callback=benchmark_progress_callback,
962 model_name=search_config.get("model_name"),
963 provider=search_config.get("provider"),
964 temperature=search_config.get("temperature", 0.7),
965 openai_endpoint_url=search_config.get("openai_endpoint_url"),
966 settings_snapshot=settings_context.snapshot, # Pass settings snapshot for thread safety
967 username=task.get("username"), # Pass username
968 user_password=user_password, # Pass password for metrics tracking
969 # The web benchmark runs against the user's encrypted DB
970 # (it has username/password and wants search metrics
971 # persisted). quick_summary's default is programmatic_mode=True
972 # for true library callers; override here so the engine's
973 # metrics path stays active.
974 programmatic_mode=False,
975 )
976 processing_time = time.time() - start_time
977 logger.info(
978 f"Task {task['example_id']} research completed in {processing_time:.2f}s, "
979 f"model={search_config.get('model_name')}, provider={search_config.get('provider')}"
980 )
982 # Extract answer
983 response = search_result.get("summary", "")
984 logger.info(
985 f"Task {task['example_id']} response length: {len(response)} chars"
986 )
988 extracted_data = extract_answer_from_response(
989 response, task["dataset_type"]
990 )
991 extracted_answer = (
992 extracted_data.get("extracted_answer", "")
993 if isinstance(extracted_data, dict)
994 else str(extracted_data)
995 )
996 logger.info(
997 f"Task {task['example_id']} extracted answer: '{extracted_answer[:100]}...'"
998 )
1000 # Extract sources - handle both direct sources and all_links_of_system
1001 sources = search_result.get("sources", [])
1002 if not sources and "all_links_of_system" in search_result:
1003 sources = search_result.get("all_links_of_system", [])
1005 # Log for debugging
1006 logger.debug(f"Search result keys: {list(search_result.keys())}")
1007 logger.debug(f"Sources found: {len(sources)} items")
1009 # Prepare result
1010 result = {
1011 **task,
1012 "response": response,
1013 "extracted_answer": extracted_answer,
1014 "confidence": str(
1015 extracted_data.get("confidence", "100")
1016 if isinstance(extracted_data, dict)
1017 else "100"
1018 ),
1019 "processing_time": processing_time,
1020 "sources": json.dumps(sources), # Convert to JSON string
1021 "completed_at": datetime.now(UTC),
1022 "research_id": tracking_id, # Store the UUID in the research_id field
1023 }
1025 # Evaluate result - requires proper grading model
1026 try:
1027 logger.info(f"Task {task['example_id']} starting evaluation...")
1028 eval_start_time = time.time()
1030 # Always attempt evaluation, regardless of provider
1031 # Modern local models like Ollama are capable of grading
1032 # Try to evaluate with proper model
1033 result_data = {
1034 "id": task["example_id"],
1035 "problem": task["question"],
1036 "correct_answer": task["correct_answer"],
1037 "response": response,
1038 "extracted_answer": extracted_answer,
1039 }
1041 # Re-arm the task's research scope around grading:
1042 # quick_summary cleared its own context on return, and the
1043 # grader LLM resolves cancellation from the thread context
1044 # (graders call get_llm without a research_id). The same
1045 # termination flag therefore aborts a long grading call
1046 # mid-stream too.
1047 from ...utilities.thread_context import (
1048 clear_search_context,
1049 set_search_context,
1050 )
1052 set_search_context(
1053 {
1054 "research_id": tracking_id,
1055 "research_query": task["question"],
1056 "username": task.get("username"),
1057 }
1058 )
1059 try:
1060 eval_result = grade_single_result(
1061 result_data,
1062 task["dataset_type"],
1063 evaluation_config,
1064 settings_context.snapshot,
1065 )
1066 finally:
1067 clear_search_context()
1068 eval_time = time.time() - eval_start_time
1069 logger.info(
1070 f"Task {task['example_id']} evaluation completed in {eval_time:.2f}s"
1071 )
1072 if eval_result and not eval_result.get("grading_error"):
1073 result.update(
1074 {
1075 "is_correct": eval_result.get("is_correct", False),
1076 "graded_confidence": eval_result.get(
1077 "graded_confidence", "0"
1078 ),
1079 "grader_response": eval_result.get(
1080 "grader_response", ""
1081 ),
1082 }
1083 )
1084 else:
1085 error_msg = (
1086 eval_result.get(
1087 "grading_error", "Unknown evaluation error"
1088 )
1089 if eval_result
1090 else "No evaluation results returned"
1091 )
1092 result.update(
1093 {
1094 "is_correct": None,
1095 "graded_confidence": "0",
1096 "grader_response": f"Evaluation failed: {error_msg}",
1097 "evaluation_error": error_msg,
1098 }
1099 )
1101 except Exception as e:
1102 logger.exception("Evaluation error")
1103 result.update(
1104 {
1105 "is_correct": None,
1106 "graded_confidence": "0",
1107 "grader_response": f"Evaluation failed: {e!s}",
1108 "evaluation_error": str(e),
1109 }
1110 )
1112 return result
1114 except Exception as e:
1115 logger.exception("Research error")
1116 return {
1117 **task,
1118 "research_error": str(e),
1119 "completed_at": datetime.now(UTC),
1120 }
1122 def _persist_unsaved_results(
1123 self, session, benchmark_run_id: int, run_data: dict
1124 ) -> list[int]:
1125 """Stage INSERTs for this run's not-yet-persisted results; return the
1126 result indices that were staged.
1128 Idempotent and rollback-safe. A result is skipped when its index is
1129 already known-saved (``saved_indices``), its ``query_hash`` is already
1130 committed for this run, or it repeats a ``query_hash`` staged earlier
1131 in this same batch. Malformed result entries are also skipped so one
1132 bad payload cannot make every later valid result roll back (#4861).
1133 Since ``query_hash`` covers ``example_id`` as well as the question text
1134 (#5080), a repeat no longer means "a dataset legitimately repeats a
1135 question": it is the same result entry reaching this method twice, so
1136 the skip is logged at WARNING rather than passed over in silence
1137 (#4860).
1138 Dedup correctness rests on the DB-backed ``seen_hashes``, which
1139 survives a rollback (the next sync re-reads it), NOT on any in-memory
1140 flag.
1142 This method deliberately does NOT mark ``saved_indices``: the caller
1143 must do that ONLY after its commit succeeds. Marking before the commit
1144 would, on a commit failure (disk full, lock timeout, ...), leave rows
1145 flagged saved that were actually rolled back — silently dropped
1146 forever. Combined with ``self._results_sync_lock`` held by the caller
1147 ACROSS the commit, this also stops the worker thread and the request
1148 thread from both inserting the same ``(benchmark_run_id, query_hash)``
1149 row and tripping the ``uix_run_query`` unique constraint.
1150 """
1151 results = run_data.get("results", [])
1152 # Read-only here — marking indices saved is the caller's job, and only
1153 # after the commit lands (see the method docstring).
1154 saved_indices = run_data.get("saved_indices", set())
1156 # Every query_hash already committed for this run, fetched once. This
1157 # is the rollback-safe source of truth for dedup; it is extended below
1158 # so a question repeated within this batch is staged only once.
1159 seen_hashes = {
1160 row[0]
1161 for row in session.query(BenchmarkResult.query_hash)
1162 .filter(BenchmarkResult.benchmark_run_id == benchmark_run_id)
1163 .all()
1164 }
1166 staged_indices = []
1167 for idx, result in enumerate(results):
1168 if idx in saved_indices:
1169 continue
1170 try:
1171 (
1172 example_id,
1173 query_hash,
1174 dataset_type,
1175 question,
1176 correct_answer,
1177 ) = _required_persistence_fields(result)
1178 except (KeyError, TypeError, ValueError) as exc:
1179 example_id = (
1180 result.get("example_id")
1181 if isinstance(result, dict)
1182 else None
1183 )
1184 task_index = (
1185 result.get("task_index")
1186 if isinstance(result, dict)
1187 else None
1188 )
1189 logger.warning(
1190 f"Benchmark run {benchmark_run_id}: skipping malformed "
1191 f"result index {idx} (example_id={example_id!r}, "
1192 f"task_index={task_index!r}): {type(exc).__name__}"
1193 )
1194 continue
1195 if query_hash in seen_hashes:
1196 logger.warning(
1197 f"Benchmark run {benchmark_run_id}: skipping result index "
1198 f"{idx} (example_id={example_id!r}), its "
1199 f"query_hash is already persisted or staged for this run"
1200 )
1201 continue
1202 session.add(
1203 BenchmarkResult(
1204 benchmark_run_id=benchmark_run_id,
1205 example_id=example_id,
1206 query_hash=query_hash,
1207 dataset_type=dataset_type,
1208 research_id=result.get("research_id"),
1209 question=question,
1210 correct_answer=correct_answer,
1211 response=result.get("response"),
1212 extracted_answer=result.get("extracted_answer"),
1213 confidence=result.get("confidence"),
1214 processing_time=result.get("processing_time"),
1215 sources=result.get("sources"),
1216 is_correct=result.get("is_correct"),
1217 graded_confidence=result.get("graded_confidence"),
1218 grader_response=result.get("grader_response"),
1219 completed_at=result.get("completed_at"),
1220 research_error=result.get("research_error"),
1221 evaluation_error=result.get("evaluation_error"),
1222 task_index=result.get("task_index"),
1223 )
1224 )
1225 seen_hashes.add(query_hash)
1226 staged_indices.append(idx)
1228 return staged_indices
1230 @staticmethod
1231 def _run_key(
1232 username: Optional[str], benchmark_run_id: int
1233 ) -> Tuple[str, int]:
1234 """Composite ``active_runs`` key: ``(username, BenchmarkRun.id)``.
1236 ``active_runs`` is a process-global map, but ``BenchmarkRun.id`` is a
1237 PER-USER autoincrement integer -- so two different users can each own a
1238 run with the same id (e.g. both have id == 1). Keying by the owner too
1239 means a request can only ever reach the caller's own in-flight run:
1240 one user's ``start_benchmark`` cannot clobber another's entry, and no
1241 user-facing lookup can name another user's run. ``username`` is
1242 normalized the same way ``start_benchmark`` records it so the key a run
1243 is stored under matches the key every later lookup computes.
1244 """
1245 return (username or "benchmark_user", benchmark_run_id)
1247 @staticmethod
1248 def _run_owner(run_data: Optional[Dict]) -> Optional[str]:
1249 """Username that started the given in-memory run entry, if any."""
1250 if not run_data: 1250 ↛ 1251line 1250 didn't jump to line 1251 because the condition on line 1250 was never true
1251 return None
1252 return run_data.get("data", {}).get("username")
1254 def sync_pending_results(
1255 self, benchmark_run_id: int, username: Optional[str] = None
1256 ):
1257 """Sync any pending results to database. Can be called from main thread."""
1258 # Read outside _results_sync_lock, so the worker thread's
1259 # ``del self.active_runs[run_key]`` (in _sync_results_to_database,
1260 # also unlocked) can land between the membership check and the subscript.
1261 # A plain ``active_runs[key]`` would then raise KeyError, which escapes
1262 # before the try below and surfaces as a spurious 500 on /api/results.
1263 # Fetch once so the check and the read see the same state. See #4859.
1264 # active_runs is keyed by (username, id) (see _run_key), so this lookup
1265 # can only ever return the caller's own run -- a caller naming another
1266 # user's per-user id simply misses. That keeps one user's results (and,
1267 # via the run's stored password, their credentials) unreachable here.
1268 run_data = self.active_runs.get(
1269 self._run_key(username, benchmark_run_id)
1270 )
1271 if run_data is None:
1272 return 0
1274 # For an anonymous/internal caller (username None) fall back to the
1275 # owner recorded on the run for the DB session below.
1276 if not username:
1277 username = self._run_owner(run_data)
1278 user_password = run_data.get("data", {}).get("user_password")
1280 saved_count = 0
1281 session = None
1282 from ...database.session_context import (
1283 get_user_db_session,
1284 safe_rollback,
1285 )
1287 try:
1288 # Serialize with the worker thread's _sync_results_to_database so
1289 # the two can't insert the same (benchmark_run_id, query_hash) row.
1290 with self._results_sync_lock:
1291 with get_user_db_session(username, user_password) as session:
1292 staged = self._persist_unsaved_results(
1293 session, benchmark_run_id, run_data
1294 )
1295 if staged:
1296 session.commit()
1297 # Mark saved ONLY after the commit succeeds, so a
1298 # failed commit leaves these for the next sync to retry
1299 # instead of dropping them silently.
1300 run_data.setdefault("saved_indices", set()).update(
1301 staged
1302 )
1303 saved_count = len(staged)
1304 logger.info(
1305 f"Saved {saved_count} new results for benchmark "
1306 f"{benchmark_run_id}"
1307 )
1308 # Reaching the end of the transaction proves persistence
1309 # is healthy again. Clear a prior transient failure even
1310 # when another thread already saved the pending rows.
1311 run_data.pop("result_persistence_failed", None)
1312 except Exception:
1313 logger.exception(
1314 f"Error syncing pending results for benchmark {benchmark_run_id}"
1315 )
1316 # Thread-local sessions are reused, so a failed flush/commit must be
1317 # rolled back explicitly or the next use raises PendingRollbackError.
1318 if session is not None: 1318 ↛ 1320line 1318 didn't jump to line 1320 because the condition on line 1318 was always true
1319 safe_rollback(session, "sync_pending_results")
1320 run_data["result_persistence_failed"] = True
1322 return saved_count
1324 def get_result_persistence_error(
1325 self, benchmark_run_id: int, username: Optional[str] = None
1326 ) -> Optional[Dict[str, str]]:
1327 """Return a safe client-facing result persistence error, if active."""
1328 # Keyed by (username, id) (see _run_key): a caller can only ever read
1329 # the state of their own run, never one that belongs to someone else.
1330 run_data = self.active_runs.get(
1331 self._run_key(username, benchmark_run_id)
1332 )
1333 if not run_data or not run_data.get("result_persistence_failed"):
1334 return None
1335 # Return a copy so callers cannot mutate the shared constant.
1336 return dict(_RESULT_PERSISTENCE_ERROR)
1338 def _sync_results_to_database(
1339 self, username: Optional[str], benchmark_run_id: int
1340 ):
1341 """Sync benchmark results from memory to database after thread completes."""
1342 run_key = self._run_key(username, benchmark_run_id)
1343 if run_key not in self.active_runs:
1344 return
1346 run_data = self.active_runs[run_key]
1347 if not run_data.get("thread_complete"):
1348 return
1350 username = run_data.get("data", {}).get("username")
1351 user_password = run_data.get("data", {}).get("user_password")
1352 session = None
1353 from ...database.session_context import (
1354 get_user_db_session,
1355 safe_rollback,
1356 )
1358 try:
1359 # Serialize with sync_pending_results (request thread) so the two
1360 # can't insert the same (benchmark_run_id, query_hash) row.
1361 with (
1362 self._results_sync_lock,
1363 get_user_db_session(username, user_password) as session,
1364 ):
1365 # Update benchmark run status
1366 benchmark_run = (
1367 session.query(BenchmarkRun)
1368 .filter(BenchmarkRun.id == benchmark_run_id)
1369 .first()
1370 )
1372 if benchmark_run and "completion_info" in run_data:
1373 info = run_data["completion_info"]
1374 benchmark_run.status = info["status"]
1375 benchmark_run.end_time = info.get(
1376 "end_time", datetime.now(UTC)
1377 )
1378 benchmark_run.completed_examples = info.get(
1379 "completed_examples", 0
1380 )
1381 benchmark_run.failed_examples = info.get(
1382 "failed_examples", 0
1383 )
1384 benchmark_run.error_message = info.get("error_message")
1386 # Stage any results not yet persisted (idempotent; safe
1387 # against a concurrent sync_pending_results on the request
1388 # thread — see _persist_unsaved_results).
1389 staged = self._persist_unsaved_results(
1390 session, benchmark_run_id, run_data
1391 )
1393 # Calculate final accuracy
1394 if benchmark_run.status == BenchmarkStatus.COMPLETED:
1395 correct_results = [
1396 r
1397 for r in run_data.get("results", [])
1398 if r.get("is_correct")
1399 ]
1400 evaluated_results = [
1401 r
1402 for r in run_data.get("results", [])
1403 if r.get("is_correct") is not None
1404 ]
1406 if evaluated_results:
1407 benchmark_run.overall_accuracy = (
1408 len(correct_results) / len(evaluated_results)
1409 ) * 100
1411 # Calculate processing rate
1412 total_time = sum(
1413 r.get("processing_time", 0)
1414 for r in evaluated_results
1415 )
1416 if total_time > 0: 1416 ↛ 1421line 1416 didn't jump to line 1421 because the condition on line 1416 was always true
1417 benchmark_run.processing_rate = len(
1418 evaluated_results
1419 ) / (total_time / 60)
1421 session.commit()
1422 # Mark saved only after the commit lands (see
1423 # _persist_unsaved_results) so a failed commit retries.
1424 if staged:
1425 run_data.setdefault("saved_indices", set()).update(
1426 staged
1427 )
1428 run_data.pop("result_persistence_failed", None)
1429 logger.info(
1430 f"Successfully synced results for benchmark {benchmark_run_id}"
1431 )
1433 # Clean up memory
1434 del self.active_runs[run_key]
1436 except Exception:
1437 logger.exception("Error syncing benchmark results to database")
1438 # Thread-local session is reused — roll back a failed flush/commit
1439 # so the next use doesn't raise PendingRollbackError.
1440 if session is not None:
1441 safe_rollback(session, "_sync_results_to_database")
1442 run_data["result_persistence_failed"] = True
1444 def _send_progress_update(
1445 self,
1446 benchmark_run_id: int,
1447 completed: int,
1448 total: int,
1449 username: Optional[str] = None,
1450 ):
1451 """Send real-time progress update via websocket.
1453 ``username`` is the run's owner and is passed through to
1454 ``emit_to_subscribers`` so the event reaches only that user's
1455 socket subscriber set -- ``benchmark_run_id`` is a per-user integer
1456 and can collide across users (see
1457 ``SocketIOService.__subscription_key``).
1458 """
1459 try:
1460 percentage = (completed / total * 100) if total > 0 else 0
1462 # Create log entry for milestone progress
1463 log_entry = {
1464 "time": datetime.now(UTC).isoformat(),
1465 "message": f"Completed {completed}/{total} examples ({percentage:.1f}%)",
1466 "progress": percentage,
1467 "metadata": {
1468 "phase": "benchmark_progress",
1469 "type": "milestone",
1470 "completed": completed,
1471 "total": total,
1472 "benchmark_run_id": benchmark_run_id,
1473 },
1474 }
1476 progress_data = {
1477 "status": "in_progress",
1478 "message": f"Processing examples: {completed}/{total}",
1479 "progress": percentage,
1480 "completed": completed,
1481 "total": total,
1482 "benchmark_run_id": benchmark_run_id,
1483 "log_entry": log_entry,
1484 "progress_log": json.dumps([log_entry]),
1485 }
1487 self.socket_service.emit_to_subscribers(
1488 "research_progress",
1489 benchmark_run_id,
1490 progress_data,
1491 owner=username,
1492 )
1494 except Exception:
1495 logger.exception("Error sending progress update")
1497 def _calculate_final_accuracy(
1498 self, benchmark_run_id: int, username: Optional[str] = None
1499 ):
1500 """Calculate and save final accuracy metrics."""
1501 from ...database.session_context import get_user_db_session
1503 with get_user_db_session(username) as session:
1504 try:
1505 # Get all results for this run
1506 results = (
1507 session.query(BenchmarkResult)
1508 .filter(
1509 BenchmarkResult.benchmark_run_id == benchmark_run_id
1510 )
1511 .filter(BenchmarkResult.is_correct.isnot(None))
1512 .all()
1513 )
1515 if results:
1516 correct_count = sum(1 for r in results if r.is_correct)
1517 overall_accuracy = (correct_count / len(results)) * 100
1519 # Calculate processing rate
1520 total_time = sum(r.processing_time or 0 for r in results)
1521 processing_rate = (
1522 (len(results) / (total_time / 60))
1523 if total_time > 0
1524 else 0
1525 )
1527 # Update benchmark run
1528 benchmark_run = (
1529 session.query(BenchmarkRun)
1530 .filter(BenchmarkRun.id == benchmark_run_id)
1531 .first()
1532 )
1533 if benchmark_run: 1533 ↛ exitline 1533 didn't jump to the function exit
1534 benchmark_run.overall_accuracy = overall_accuracy
1535 benchmark_run.processing_rate = processing_rate
1536 session.commit()
1538 except Exception:
1539 logger.exception("Error calculating final accuracy")
1541 def update_benchmark_status(
1542 self,
1543 benchmark_run_id: int,
1544 status: BenchmarkStatus,
1545 error_message: Optional[str] = None,
1546 username: Optional[str] = None,
1547 ):
1548 """Update benchmark run status."""
1549 from ...database.session_context import get_user_db_session
1551 with get_user_db_session(username) as session:
1552 try:
1553 benchmark_run = (
1554 session.query(BenchmarkRun)
1555 .filter(BenchmarkRun.id == benchmark_run_id)
1556 .first()
1557 )
1558 if benchmark_run:
1559 benchmark_run.status = status
1560 benchmark_run.updated_at = datetime.now(UTC)
1562 if error_message:
1563 benchmark_run.error_message = error_message
1565 if (
1566 status == BenchmarkStatus.IN_PROGRESS
1567 and not benchmark_run.start_time
1568 ):
1569 benchmark_run.start_time = datetime.now(UTC)
1570 elif (
1571 status
1572 in [BenchmarkStatus.COMPLETED, BenchmarkStatus.FAILED]
1573 and not benchmark_run.end_time
1574 ):
1575 benchmark_run.end_time = datetime.now(UTC)
1577 session.commit()
1579 except Exception:
1580 session.rollback()
1581 logger.exception("Error updating benchmark status")
1583 def get_benchmark_status(
1584 self, benchmark_run_id: int, username: str = None
1585 ) -> Optional[Dict]:
1586 """Get current status of a benchmark run."""
1587 from ...database.session_context import get_user_db_session
1589 with get_user_db_session(username) as session:
1590 try:
1591 benchmark_run = (
1592 session.query(BenchmarkRun)
1593 .filter(BenchmarkRun.id == benchmark_run_id)
1594 .first()
1595 )
1596 if not benchmark_run:
1597 return None
1599 # Calculate running accuracy from this run's evaluated results
1600 results = (
1601 session.query(BenchmarkResult)
1602 .filter(
1603 BenchmarkResult.benchmark_run_id == benchmark_run_id
1604 )
1605 .filter(BenchmarkResult.is_correct.isnot(None))
1606 .all()
1607 )
1609 running_accuracy = None
1610 # Dynamic per-dataset accuracy tracking
1611 dataset_accuracies = {}
1613 if results:
1614 # Overall running accuracy
1615 correct_count = sum(1 for r in results if r.is_correct)
1616 running_accuracy = (correct_count / len(results)) * 100
1618 # Calculate accuracy for each dataset type dynamically
1619 from collections import defaultdict
1621 dataset_results = defaultdict(list)
1623 # Group results by dataset type
1624 for r in results:
1625 dataset_results[r.dataset_type.value].append(r)
1627 # Calculate accuracy for each dataset
1628 for (
1629 dataset_type,
1630 dataset_result_list,
1631 ) in dataset_results.items():
1632 if dataset_result_list: 1632 ↛ 1628line 1632 didn't jump to line 1628 because the condition on line 1632 was always true
1633 correct = sum(
1634 1 for r in dataset_result_list if r.is_correct
1635 )
1636 accuracy = (
1637 correct / len(dataset_result_list)
1638 ) * 100
1639 # Store with _accuracy suffix for consistency
1640 dataset_accuracies[f"{dataset_type}_accuracy"] = (
1641 accuracy
1642 )
1644 # Calculate time estimates and reliability metrics
1645 estimated_time_remaining = None
1646 total_elapsed_time = None
1647 avg_time_per_example = None
1648 accuracy_confidence = None
1650 # Get ALL results for timing calculation (including those pending evaluation)
1651 all_results_for_timing = (
1652 session.query(BenchmarkResult)
1653 .filter(
1654 BenchmarkResult.benchmark_run_id == benchmark_run_id
1655 )
1656 .all()
1657 )
1659 if benchmark_run.start_time and all_results_for_timing:
1660 # Calculate elapsed time
1661 current_time = datetime.now(UTC)
1662 total_elapsed_time = (
1663 current_time - benchmark_run.start_time
1664 ).total_seconds()
1666 # Calculate average processing time per example using actual count
1667 avg_time_per_example = total_elapsed_time / len(
1668 all_results_for_timing
1669 )
1671 logger.info(
1672 f"Time calculation - elapsed: {total_elapsed_time:.2f}s, "
1673 f"results_count: {len(all_results_for_timing)}, "
1674 f"avg_per_example: {avg_time_per_example:.2f}s"
1675 )
1677 # Estimate remaining time
1678 remaining_examples = benchmark_run.total_examples - len(
1679 all_results_for_timing
1680 )
1681 if remaining_examples > 0:
1682 estimated_time_remaining = (
1683 avg_time_per_example * remaining_examples
1684 )
1685 logger.info(
1686 f"Time estimation - total: {benchmark_run.total_examples}, "
1687 f"completed: {len(all_results_for_timing)}, remaining: {remaining_examples}, "
1688 f"avg_time: {avg_time_per_example:.2f}s, "
1689 f"estimated_remaining: {estimated_time_remaining:.2f}s"
1690 )
1692 # Calculate accuracy confidence interval (Wilson score, 95%)
1693 if results and len(results) >= 3:
1694 from local_deep_research.benchmarks.metrics.statistics import (
1695 wilson_score_interval,
1696 )
1698 n = len(results)
1699 correct = sum(1 for r in results if r.is_correct)
1700 ci = wilson_score_interval(correct, n)
1701 accuracy_confidence = {
1702 "lower_bound": ci["lower"] * 100,
1703 "upper_bound": ci["upper"] * 100,
1704 "margin_of_error": ci["margin_of_error"] * 100,
1705 "sample_size": n,
1706 }
1708 status_data = {
1709 "id": benchmark_run.id,
1710 "run_name": benchmark_run.run_name,
1711 "status": benchmark_run.status.value,
1712 "completed_examples": len(
1713 all_results_for_timing
1714 ), # Use actual count from DB
1715 "total_examples": benchmark_run.total_examples,
1716 "failed_examples": benchmark_run.failed_examples,
1717 "overall_accuracy": benchmark_run.overall_accuracy
1718 or running_accuracy, # Use running accuracy if final not calculated
1719 "running_accuracy": running_accuracy, # Current running accuracy
1720 "processing_rate": benchmark_run.processing_rate,
1721 "estimated_time_remaining": estimated_time_remaining, # seconds
1722 "total_elapsed_time": total_elapsed_time, # seconds
1723 "avg_time_per_example": avg_time_per_example, # seconds
1724 "accuracy_confidence": accuracy_confidence, # confidence interval
1725 "created_at": benchmark_run.created_at.isoformat()
1726 if benchmark_run.created_at
1727 else None,
1728 "start_time": benchmark_run.start_time.isoformat()
1729 if benchmark_run.start_time
1730 else None,
1731 "end_time": benchmark_run.end_time.isoformat()
1732 if benchmark_run.end_time
1733 else None,
1734 # Genericized at this boundary (CWE-209) — see
1735 # _GENERIC_BENCHMARK_ERROR. Returns None when there is no
1736 # failure so non-error states are unaffected.
1737 "error_message": (
1738 _GENERIC_BENCHMARK_ERROR
1739 if benchmark_run.error_message
1740 else None
1741 ),
1742 # Add all per-dataset accuracies dynamically
1743 **dataset_accuracies,
1744 }
1746 logger.info(
1747 f"Benchmark {benchmark_run_id} status - completed: {benchmark_run.completed_examples}, "
1748 f"running_acc: {running_accuracy}, dataset_accuracies: {dataset_accuracies}, "
1749 f"avg_time: {avg_time_per_example}"
1750 )
1752 return status_data
1754 except Exception:
1755 logger.exception("Error getting benchmark status")
1756 return None
1758 def cancel_benchmark(
1759 self, benchmark_run_id: int, username: Optional[str] = None
1760 ) -> bool:
1761 """Cancel a running benchmark."""
1762 try:
1763 # active_runs is keyed by (username, id) (see _run_key), so this
1764 # lookup only ever finds the requester's own run -- one user can no
1765 # longer set another user's in-flight run to "cancelled" and stop
1766 # it (cross-user DoS). The DB update below is already scoped to the
1767 # caller's own database.
1768 run = self.active_runs.get(
1769 self._run_key(username, benchmark_run_id)
1770 )
1771 if run is not None:
1772 run["status"] = "cancelled"
1773 # Abort the in-flight task's LLM calls mid-generation, not
1774 # just between tasks (see _process_benchmark_task): the
1775 # worker publishes each task's research scope here, and the
1776 # LLM wrapper aborts on the flag within one streamed chunk.
1777 # Scoped to the caller's own run via _run_key (above), and
1778 # the id is an unguessable per-task UUID besides.
1779 tracking_id = run.get("current_tracking_id")
1780 if tracking_id:
1781 try:
1782 from ...web.research_state import (
1783 clear_termination_flag,
1784 set_termination_flag,
1785 )
1787 set_termination_flag(tracking_id)
1788 # The worker pops the id then clears its flag in a
1789 # per-task finally. If it finished between our read
1790 # and set, we just created a stale process-global
1791 # flag for a fresh UUID that will never be reused
1792 # (cosmetic leak). Re-check and reclaim it.
1793 if run.get("current_tracking_id") != tracking_id: 1793 ↛ 1794line 1793 didn't jump to line 1794 because the condition on line 1793 was never true
1794 clear_termination_flag(tracking_id)
1795 except Exception:
1796 logger.debug(
1797 "Failed to set benchmark termination flag "
1798 "(non-critical)"
1799 )
1801 self.update_benchmark_status(
1802 benchmark_run_id, BenchmarkStatus.CANCELLED, username=username
1803 )
1804 logger.info(f"Cancelled benchmark run {benchmark_run_id}")
1805 return True
1807 except Exception:
1808 logger.exception(f"Error cancelling benchmark {benchmark_run_id}")
1809 return False
1812# Global service instance
1813benchmark_service = BenchmarkService()