Coverage for src/local_deep_research/metrics/search_tracker.py: 97%
153 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Search call tracking system for metrics collection.
3Similar to token_counter.py but tracks search engine usage.
4"""
6from typing import Any, Dict, List, Optional
8from loguru import logger
9from sqlalchemy import case, func
11from ..utilities.thread_context import get_search_context
12from ..utilities.metric_write_health import (
13 MetricWriteReason,
14 record_metric_write_drop,
15)
16from ..database.models import SearchCall
17from .database import MetricsDatabase
18from .query_utils import get_research_mode_condition, get_time_filter_condition
20# ``success_status`` of a search that was not sent because its engine was
21# disabled or cooling down for the research (see engine_availability).
22# Recorded so outages stay visible, but not counted as a search attempt:
23# it has no response time, and it is excluded from call counts, averages
24# and success rates.
25SEARCH_STATUS_SKIPPED = "skipped_unavailable"
28def _session_has_pending_orm_state(session: Any) -> bool:
29 """True if *session* has unflushed ORM work ``begin_nested()`` would flush.
31 ``Session.begin_nested()`` flushes ``new``/``dirty``/``deleted`` before
32 CREATE SAVEPOINT. A flush failure (SQLite writer lock held elsewhere)
33 invalidates the caller's transaction; catching it cannot restore
34 ``session.is_active``. Best-effort telemetry must skip in that case.
35 """
36 try:
37 for attr in ("new", "dirty", "deleted"):
38 collection = getattr(session, attr, None)
39 if collection and len(collection) > 0:
40 return True
41 except Exception:
42 return False
43 return False
46class SearchTracker:
47 """Track search engine calls and performance metrics."""
49 def __init__(self, db: Optional[MetricsDatabase] = None):
50 """Initialize the search tracker."""
51 self.db = db or MetricsDatabase()
53 @staticmethod
54 def record_search(
55 engine_name: str,
56 query: str,
57 results_count: int = 0,
58 response_time_ms: int = 0,
59 success: bool = True,
60 error_message: Optional[str] = None,
61 status: Optional[str] = None,
62 error_type: Optional[str] = None,
63 ) -> None:
64 """Record a completed search operation directly to database.
66 ``status`` overrides the success/error status, e.g.
67 ``SEARCH_STATUS_SKIPPED`` for a search skipped because its engine
68 was unavailable; ``error_message`` then carries the reason.
69 Pass ``error_type`` separately when the exception message has been
70 sanitized. Plain text alone cannot identify an exception class.
71 """
72 password = None
74 # Extract research context (thread-safe)
75 context = get_search_context()
77 # Skip metrics recording in programmatic mode or when no context is set
78 if context is None:
79 logger.warning(
80 "Skipping search metrics recording - no research context available "
81 "(likely in programmatic mode)"
82 )
83 return
85 research_id = context.get("research_id")
87 # Convert research_id to string if it's an integer (for backward compatibility)
88 if isinstance(research_id, int):
89 research_id = str(research_id)
90 research_query = context.get("research_query")
91 research_mode = context.get("research_mode", "unknown")
92 research_phase = context.get("research_phase", "search")
93 search_iteration = context.get("search_iteration", 0)
95 # Determine success status
96 success_status = status or ("success" if success else "error")
97 if status == SEARCH_STATUS_SKIPPED:
98 error_type = "engine_unavailable"
99 elif error_message and not error_type:
100 error_type = (
101 type(error_message).__name__
102 if isinstance(error_message, Exception)
103 else "unknown_error"
104 )
106 # Record search call in database - only from background threads
107 username = None
108 dropped_reason: MetricWriteReason | None = None
109 try:
110 # Get username from context for thread-safe database
111 username = context.get("username")
112 if not username:
113 logger.warning(
114 "Cannot save search metrics - no username in research "
115 f"context. Search engine: {engine_name}"
116 )
117 return
119 # Get password from context
120 password = context.get("user_password")
121 if not password:
122 dropped_reason = "missing_credentials"
123 logger.warning(
124 "Cannot save search metrics - no password in research "
125 f"context. Search engine: {engine_name}, "
126 f"username: {username}"
127 )
128 return
130 try:
132 def _write_search(session: Any) -> None:
133 search_call = SearchCall(
134 research_id=research_id,
135 research_query=research_query,
136 research_mode=research_mode,
137 research_phase=research_phase,
138 search_iteration=search_iteration,
139 search_engine=engine_name,
140 query=query,
141 results_count=results_count,
142 response_time_ms=response_time_ms,
143 success=int(success_status == "success"),
144 success_status=success_status,
145 error_type=error_type,
146 error_message=str(error_message)
147 if error_message
148 else None,
149 )
150 session.add(search_call)
152 logger.debug(
153 f"Search call recorded to encrypted DB: {engine_name} - "
154 f"{results_count} results in {response_time_ms}ms"
155 )
157 # Two paths, because get_user_db_session does not commit
158 # on exit and metrics_writer.get_session() is a second
159 # QueuePool connection:
160 #
161 # * in_scope(): reuse the thread-local session and write
162 # inside a SAVEPOINT (telemetry_savepoint()). A nested
163 # transaction isolates metrics failures from the caller's
164 # session (preventing the caller's staged work from being
165 # rolled back) and never calls session.commit(). The row
166 # joins the caller's open write transaction, or commits
167 # at once on its own when the scope has not written.
168 # Session.begin_nested() flushes ALL pending ORM state
169 # before CREATE SAVEPOINT. If that flush fails (e.g.
170 # another connection holds SQLite's writer lock),
171 # SQLAlchemy invalidates the caller's transaction and
172 # catching the exception cannot restore it. Best-effort
173 # telemetry therefore skips when the caller already has
174 # unflushed work — the write cannot be isolated.
175 # * otherwise: metrics_writer commits+closes its own
176 # session. A thread-local write here would be rolled
177 # back by the next get_session() at scope depth 0.
178 from ..database.thread_local_session import (
179 thread_session_manager,
180 )
182 if thread_session_manager.in_scope():
183 from ..database.session_context import (
184 telemetry_savepoint,
185 )
187 # Borrow directly: session provisioning can replace a
188 # different owner's session or recover a failed one by
189 # rolling back the caller's transaction.
190 session = thread_session_manager.get_scoped_session(
191 username
192 )
193 if session is None:
194 dropped_reason = "database_unavailable"
195 # Unlike rollback-coupled loss, an owner mismatch is
196 # an anomaly worth surfacing (as rate telemetry does).
197 logger.warning(
198 "Skipping in-scope search metrics without a "
199 "matching owner session"
200 )
201 return
202 if _session_has_pending_orm_state(session):
203 dropped_reason = "pending_transaction"
204 logger.debug(
205 "Skipping in-scope search metrics write; "
206 "caller session has pending ORM state "
207 "that Session.begin_nested() would flush "
208 "before creating a SAVEPOINT"
209 )
210 return
211 with telemetry_savepoint(session):
212 _write_search(session)
213 else:
214 from ..database.thread_metrics import metrics_writer
216 # Set password for this thread
217 metrics_writer.set_user_password(username, password)
219 with metrics_writer.get_session(username) as session:
220 _write_search(session)
221 except Exception as e:
222 dropped_reason = "write_error"
223 # Deferred import to avoid potential circular imports during module initialization
224 from ..security.log_sanitizer import scrub_error
226 safe_msg = scrub_error(e, password)
227 logger.warning(f"Failed to write search metrics: {safe_msg}")
229 except Exception as e:
230 dropped_reason = "write_error"
231 # Deferred import to avoid potential circular imports during module initialization
232 from ..security.log_sanitizer import scrub_error
234 safe_msg = scrub_error(e, password)
235 logger.warning(f"Failed to record search call: {safe_msg}")
236 finally:
237 # A write and its context-manager cleanup can both fail. Count
238 # this recording attempt once after all nested handlers finish.
239 if dropped_reason is not None:
240 record_metric_write_drop(username, "search", dropped_reason)
242 def get_search_metrics(
243 self,
244 period: str = "30d",
245 research_mode: str = "all",
246 username: Optional[str] = None,
247 password: Optional[str] = None,
248 ) -> Dict[str, Any]:
249 """Get search engine usage metrics."""
250 with self.db.get_session(
251 username=username, password=password
252 ) as session:
253 try:
254 # Build base query with filters
255 query = session.query(SearchCall).filter(
256 SearchCall.search_engine.isnot(None)
257 )
259 # Apply time filter
260 time_condition = get_time_filter_condition(
261 period, SearchCall.timestamp
262 )
263 if time_condition is not None:
264 query = query.filter(time_condition)
266 # Apply research mode filter
267 mode_condition = get_research_mode_condition(
268 research_mode, SearchCall.research_mode
269 )
270 if mode_condition is not None:
271 query = query.filter(mode_condition)
273 # Get search engine statistics using ORM aggregation
274 skipped = SearchCall.success_status == SEARCH_STATUS_SKIPPED
275 search_stats = session.query(
276 SearchCall.search_engine,
277 func.sum(case((skipped, 0), else_=1)).label("call_count"),
278 func.avg(
279 case((skipped, None), else_=SearchCall.response_time_ms)
280 ).label("avg_response_time"),
281 func.sum(SearchCall.results_count).label("total_results"),
282 func.avg(
283 case((skipped, None), else_=SearchCall.results_count)
284 ).label("avg_results_per_call"),
285 func.sum(
286 case(
287 (SearchCall.success_status == "success", 1), else_=0
288 )
289 ).label("success_count"),
290 func.sum(
291 case((SearchCall.success_status == "error", 1), else_=0)
292 ).label("error_count"),
293 func.sum(case((skipped, 1), else_=0)).label(
294 "skipped_count"
295 ),
296 ).filter(SearchCall.search_engine.isnot(None))
298 # Apply same filters to stats query
299 if time_condition is not None:
300 search_stats = search_stats.filter(time_condition)
301 if mode_condition is not None:
302 search_stats = search_stats.filter(mode_condition)
304 search_stats = (
305 search_stats.group_by(SearchCall.search_engine)
306 .order_by(func.count().desc())
307 .all()
308 )
310 # Get recent search calls
311 recent_calls_query = session.query(SearchCall)
312 if time_condition is not None:
313 recent_calls_query = recent_calls_query.filter(
314 time_condition
315 )
316 if mode_condition is not None:
317 recent_calls_query = recent_calls_query.filter(
318 mode_condition
319 )
321 recent_calls = (
322 recent_calls_query.order_by(SearchCall.timestamp.desc())
323 .limit(20)
324 .all()
325 )
327 return {
328 "search_engine_stats": [
329 {
330 "engine": stat.search_engine,
331 "call_count": stat.call_count,
332 "avg_response_time": stat.avg_response_time or 0,
333 "total_results": stat.total_results or 0,
334 "avg_results_per_call": stat.avg_results_per_call
335 or 0,
336 "success_rate": (
337 (stat.success_count / stat.call_count * 100)
338 if stat.call_count
339 else 0
340 ),
341 "error_count": stat.error_count or 0,
342 "skipped_count": stat.skipped_count or 0,
343 }
344 for stat in search_stats
345 ],
346 "recent_calls": [
347 {
348 "engine": call.search_engine,
349 "query": (
350 call.query[:100] + "..."
351 if len(call.query or "") > 100
352 else call.query
353 ),
354 "results_count": call.results_count,
355 "response_time_ms": call.response_time_ms,
356 "success_status": call.success_status,
357 "timestamp": str(call.timestamp),
358 }
359 for call in recent_calls
360 ],
361 }
363 except Exception:
364 logger.warning("Error getting search metrics")
365 return {"search_engine_stats": [], "recent_calls": []}
367 def get_research_search_metrics(
368 self,
369 research_id: str,
370 username: Optional[str] = None,
371 password: Optional[str] = None,
372 ) -> Dict[str, Any]:
373 """Get search metrics for a specific research session."""
374 with self.db.get_session(
375 username=username, password=password
376 ) as session:
377 try:
378 # Get all search calls for this research
379 search_calls = (
380 session.query(SearchCall)
381 .filter(SearchCall.research_id == research_id)
382 .order_by(SearchCall.timestamp.asc())
383 .all()
384 )
386 # Get search engine stats for this research
387 skipped = SearchCall.success_status == SEARCH_STATUS_SKIPPED
388 engine_stats = (
389 session.query(
390 SearchCall.search_engine,
391 func.sum(case((skipped, 0), else_=1)).label(
392 "call_count"
393 ),
394 func.avg(
395 case(
396 (skipped, None),
397 else_=SearchCall.response_time_ms,
398 )
399 ).label("avg_response_time"),
400 func.sum(SearchCall.results_count).label(
401 "total_results"
402 ),
403 func.sum(
404 case(
405 (SearchCall.success_status == "success", 1),
406 else_=0,
407 )
408 ).label("success_count"),
409 func.sum(case((skipped, 1), else_=0)).label(
410 "skipped_count"
411 ),
412 )
413 .filter(SearchCall.research_id == research_id)
414 .group_by(SearchCall.search_engine)
415 .order_by(func.count().desc())
416 .all()
417 )
419 # Calculate totals; skipped searches were never sent.
420 attempted = [
421 call
422 for call in search_calls
423 if call.success_status != SEARCH_STATUS_SKIPPED
424 ]
425 total_searches = len(attempted)
426 skipped_searches = len(search_calls) - total_searches
427 total_results = sum(
428 call.results_count or 0 for call in search_calls
429 )
430 avg_response_time = (
431 sum(call.response_time_ms or 0 for call in attempted)
432 / total_searches
433 if total_searches > 0
434 else 0
435 )
436 successful_searches = sum(
437 1 for call in attempted if call.success_status == "success"
438 )
439 success_rate = (
440 (successful_searches / total_searches * 100)
441 if total_searches > 0
442 else 0
443 )
445 return {
446 "total_searches": total_searches,
447 "skipped_searches": skipped_searches,
448 "total_results": total_results,
449 "avg_response_time": round(avg_response_time),
450 "success_rate": round(success_rate, 1),
451 "search_calls": [
452 {
453 "engine": call.search_engine,
454 "query": call.query,
455 "results_count": call.results_count,
456 "response_time_ms": call.response_time_ms,
457 "success_status": call.success_status,
458 "timestamp": str(call.timestamp),
459 }
460 for call in search_calls
461 ],
462 "engine_stats": [
463 {
464 "engine": stat.search_engine,
465 "call_count": stat.call_count,
466 "avg_response_time": stat.avg_response_time or 0,
467 "total_results": stat.total_results or 0,
468 "success_rate": (
469 (stat.success_count / stat.call_count * 100)
470 if stat.call_count
471 else 0
472 ),
473 "skipped_count": stat.skipped_count or 0,
474 }
475 for stat in engine_stats
476 ],
477 }
479 except Exception:
480 logger.warning("Error getting research search metrics")
481 return {
482 "total_searches": 0,
483 "skipped_searches": 0,
484 "total_results": 0,
485 "avg_response_time": 0,
486 "success_rate": 0,
487 "search_calls": [],
488 "engine_stats": [],
489 }
491 def get_search_time_series(
492 self,
493 period: str = "30d",
494 research_mode: str = "all",
495 username: Optional[str] = None,
496 password: Optional[str] = None,
497 ) -> List[Dict[str, Any]]:
498 """Get search activity time series data for charting.
500 Args:
501 period: Time period to filter by ('7d', '30d', '3m', '1y', 'all')
502 research_mode: Research mode to filter by ('quick', 'detailed', 'all')
503 username: Username for database access
504 password: Password for database access
506 Returns:
507 List of time series data points with search engine activity
508 """
509 with self.db.get_session(
510 username=username, password=password
511 ) as session:
512 try:
513 # Build base query
514 query = session.query(SearchCall).filter(
515 SearchCall.search_engine.isnot(None),
516 SearchCall.timestamp.isnot(None),
517 )
519 # Apply time filter
520 time_condition = get_time_filter_condition(
521 period, SearchCall.timestamp
522 )
523 if time_condition is not None: 523 ↛ 527line 523 didn't jump to line 527 because the condition on line 523 was always true
524 query = query.filter(time_condition)
526 # Apply research mode filter
527 mode_condition = get_research_mode_condition(
528 research_mode, SearchCall.research_mode
529 )
530 if mode_condition is not None:
531 query = query.filter(mode_condition)
533 # Get all search calls ordered by time
534 search_calls = query.order_by(SearchCall.timestamp.asc()).all()
536 # Create time series data
537 time_series = []
538 for call in search_calls:
539 if call.success_status == SEARCH_STATUS_SKIPPED: 539 ↛ 541line 539 didn't jump to line 541 because the condition on line 539 was never true
540 # Not search activity: nothing was sent.
541 continue
542 time_series.append(
543 {
544 "timestamp": (
545 str(call.timestamp) if call.timestamp else None
546 ),
547 "search_engine": call.search_engine,
548 "results_count": call.results_count or 0,
549 "response_time_ms": call.response_time_ms or 0,
550 "success_status": call.success_status,
551 "query": (
552 call.query[:50] + "..."
553 if call.query and len(call.query) > 50
554 else call.query
555 ),
556 }
557 )
559 return time_series
561 except Exception:
562 logger.warning("Error getting search time series")
563 return []
566def get_search_tracker() -> SearchTracker:
567 """Create a SearchTracker instance.
569 Returns a fresh instance each time. Callers should pass username/password
570 to the query methods so the correct per-user database is accessed.
571 """
572 return SearchTracker()