Coverage for src/local_deep_research/metrics/search_tracker.py: 99%
112 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +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 ..database.models import SearchCall
13from .database import MetricsDatabase
14from .query_utils import get_research_mode_condition, get_time_filter_condition
17class SearchTracker:
18 """Track search engine calls and performance metrics."""
20 def __init__(self, db: Optional[MetricsDatabase] = None):
21 """Initialize the search tracker."""
22 self.db = db or MetricsDatabase()
24 @staticmethod
25 def record_search(
26 engine_name: str,
27 query: str,
28 results_count: int = 0,
29 response_time_ms: int = 0,
30 success: bool = True,
31 error_message: Optional[str] = None,
32 ) -> None:
33 """Record a completed search operation directly to database."""
34 password = None
36 # Extract research context (thread-safe)
37 context = get_search_context()
39 # Skip metrics recording in programmatic mode or when no context is set
40 if context is None:
41 logger.warning(
42 "Skipping search metrics recording - no research context available "
43 "(likely in programmatic mode)"
44 )
45 return
47 research_id = context.get("research_id")
49 # Convert research_id to string if it's an integer (for backward compatibility)
50 if isinstance(research_id, int):
51 research_id = str(research_id)
52 research_query = context.get("research_query")
53 research_mode = context.get("research_mode", "unknown")
54 research_phase = context.get("research_phase", "search")
55 search_iteration = context.get("search_iteration", 0)
57 # Determine success status
58 success_status = "success" if success else "error"
59 error_type = None
60 if error_message:
61 error_type = (
62 type(error_message).__name__
63 if isinstance(error_message, Exception)
64 else "unknown_error"
65 )
67 # Record search call in database - only from background threads
68 try:
69 # Get username from context for thread-safe database
70 username = context.get("username")
71 if not username:
72 logger.warning(
73 f"Cannot save search metrics - no username in research context. "
74 f"Search: {engine_name} for '{query}'"
75 )
76 return
78 # Get password from context
79 password = context.get("user_password")
80 if not password:
81 logger.warning(
82 f"Cannot save search metrics - no password in research context. "
83 f"Search: {engine_name} for '{query}', username: {username}"
84 )
85 return
87 # Use thread-safe metrics writer
88 from ..database.thread_metrics import metrics_writer
90 try:
91 # Set password for this thread
92 metrics_writer.set_user_password(username, password)
94 with metrics_writer.get_session(username) as session:
95 search_call = SearchCall(
96 research_id=research_id,
97 research_query=research_query,
98 research_mode=research_mode,
99 research_phase=research_phase,
100 search_iteration=search_iteration,
101 search_engine=engine_name,
102 query=query,
103 results_count=results_count,
104 response_time_ms=response_time_ms,
105 success_status=success_status,
106 error_type=error_type,
107 error_message=str(error_message)
108 if error_message
109 else None,
110 )
111 session.add(search_call)
113 logger.debug(
114 f"Search call recorded to encrypted DB: {engine_name} - "
115 f"{results_count} results in {response_time_ms}ms"
116 )
117 except Exception as e:
118 # Deferred import to avoid potential circular imports during module initialization
119 from ..security.log_sanitizer import scrub_error
121 safe_msg = scrub_error(e, password)
122 logger.warning(f"Failed to write search metrics: {safe_msg}")
124 except Exception as e:
125 # Deferred import to avoid potential circular imports during module initialization
126 from ..security.log_sanitizer import scrub_error
128 safe_msg = scrub_error(e, password)
129 logger.warning(f"Failed to record search call: {safe_msg}")
131 def get_search_metrics(
132 self,
133 period: str = "30d",
134 research_mode: str = "all",
135 username: Optional[str] = None,
136 password: Optional[str] = None,
137 ) -> Dict[str, Any]:
138 """Get search engine usage metrics."""
139 with self.db.get_session(
140 username=username, password=password
141 ) as session:
142 try:
143 # Build base query with filters
144 query = session.query(SearchCall).filter(
145 SearchCall.search_engine.isnot(None)
146 )
148 # Apply time filter
149 time_condition = get_time_filter_condition(
150 period, SearchCall.timestamp
151 )
152 if time_condition is not None:
153 query = query.filter(time_condition)
155 # Apply research mode filter
156 mode_condition = get_research_mode_condition(
157 research_mode, SearchCall.research_mode
158 )
159 if mode_condition is not None:
160 query = query.filter(mode_condition)
162 # Get search engine statistics using ORM aggregation
163 search_stats = session.query(
164 SearchCall.search_engine,
165 func.count().label("call_count"),
166 func.avg(SearchCall.response_time_ms).label(
167 "avg_response_time"
168 ),
169 func.sum(SearchCall.results_count).label("total_results"),
170 func.avg(SearchCall.results_count).label(
171 "avg_results_per_call"
172 ),
173 func.sum(
174 case(
175 (SearchCall.success_status == "success", 1), else_=0
176 )
177 ).label("success_count"),
178 func.sum(
179 case((SearchCall.success_status == "error", 1), else_=0)
180 ).label("error_count"),
181 ).filter(SearchCall.search_engine.isnot(None))
183 # Apply same filters to stats query
184 if time_condition is not None:
185 search_stats = search_stats.filter(time_condition)
186 if mode_condition is not None:
187 search_stats = search_stats.filter(mode_condition)
189 search_stats = (
190 search_stats.group_by(SearchCall.search_engine)
191 .order_by(func.count().desc())
192 .all()
193 )
195 # Get recent search calls
196 recent_calls_query = session.query(SearchCall)
197 if time_condition is not None:
198 recent_calls_query = recent_calls_query.filter(
199 time_condition
200 )
201 if mode_condition is not None:
202 recent_calls_query = recent_calls_query.filter(
203 mode_condition
204 )
206 recent_calls = (
207 recent_calls_query.order_by(SearchCall.timestamp.desc())
208 .limit(20)
209 .all()
210 )
212 return {
213 "search_engine_stats": [
214 {
215 "engine": stat.search_engine,
216 "call_count": stat.call_count,
217 "avg_response_time": stat.avg_response_time or 0,
218 "total_results": stat.total_results or 0,
219 "avg_results_per_call": stat.avg_results_per_call
220 or 0,
221 "success_rate": (
222 (stat.success_count / stat.call_count * 100)
223 if stat.call_count > 0
224 else 0
225 ),
226 "error_count": stat.error_count or 0,
227 }
228 for stat in search_stats
229 ],
230 "recent_calls": [
231 {
232 "engine": call.search_engine,
233 "query": (
234 call.query[:100] + "..."
235 if len(call.query or "") > 100
236 else call.query
237 ),
238 "results_count": call.results_count,
239 "response_time_ms": call.response_time_ms,
240 "success_status": call.success_status,
241 "timestamp": str(call.timestamp),
242 }
243 for call in recent_calls
244 ],
245 }
247 except Exception:
248 logger.warning("Error getting search metrics")
249 return {"search_engine_stats": [], "recent_calls": []}
251 def get_research_search_metrics(
252 self,
253 research_id: str,
254 username: Optional[str] = None,
255 password: Optional[str] = None,
256 ) -> Dict[str, Any]:
257 """Get search metrics for a specific research session."""
258 with self.db.get_session(
259 username=username, password=password
260 ) as session:
261 try:
262 # Get all search calls for this research
263 search_calls = (
264 session.query(SearchCall)
265 .filter(SearchCall.research_id == research_id)
266 .order_by(SearchCall.timestamp.asc())
267 .all()
268 )
270 # Get search engine stats for this research
271 engine_stats = (
272 session.query(
273 SearchCall.search_engine,
274 func.count().label("call_count"),
275 func.avg(SearchCall.response_time_ms).label(
276 "avg_response_time"
277 ),
278 func.sum(SearchCall.results_count).label(
279 "total_results"
280 ),
281 func.sum(
282 case(
283 (SearchCall.success_status == "success", 1),
284 else_=0,
285 )
286 ).label("success_count"),
287 )
288 .filter(SearchCall.research_id == research_id)
289 .group_by(SearchCall.search_engine)
290 .order_by(func.count().desc())
291 .all()
292 )
294 # Calculate totals
295 total_searches = len(search_calls)
296 total_results = sum(
297 call.results_count or 0 for call in search_calls
298 )
299 avg_response_time = (
300 sum(call.response_time_ms or 0 for call in search_calls)
301 / total_searches
302 if total_searches > 0
303 else 0
304 )
305 successful_searches = sum(
306 1
307 for call in search_calls
308 if call.success_status == "success"
309 )
310 success_rate = (
311 (successful_searches / total_searches * 100)
312 if total_searches > 0
313 else 0
314 )
316 return {
317 "total_searches": total_searches,
318 "total_results": total_results,
319 "avg_response_time": round(avg_response_time),
320 "success_rate": round(success_rate, 1),
321 "search_calls": [
322 {
323 "engine": call.search_engine,
324 "query": call.query,
325 "results_count": call.results_count,
326 "response_time_ms": call.response_time_ms,
327 "success_status": call.success_status,
328 "timestamp": str(call.timestamp),
329 }
330 for call in search_calls
331 ],
332 "engine_stats": [
333 {
334 "engine": stat.search_engine,
335 "call_count": stat.call_count,
336 "avg_response_time": stat.avg_response_time or 0,
337 "total_results": stat.total_results or 0,
338 "success_rate": (
339 (stat.success_count / stat.call_count * 100)
340 if stat.call_count > 0
341 else 0
342 ),
343 }
344 for stat in engine_stats
345 ],
346 }
348 except Exception:
349 logger.warning("Error getting research search metrics")
350 return {
351 "total_searches": 0,
352 "total_results": 0,
353 "avg_response_time": 0,
354 "success_rate": 0,
355 "search_calls": [],
356 "engine_stats": [],
357 }
359 def get_search_time_series(
360 self,
361 period: str = "30d",
362 research_mode: str = "all",
363 username: Optional[str] = None,
364 password: Optional[str] = None,
365 ) -> List[Dict[str, Any]]:
366 """Get search activity time series data for charting.
368 Args:
369 period: Time period to filter by ('7d', '30d', '3m', '1y', 'all')
370 research_mode: Research mode to filter by ('quick', 'detailed', 'all')
371 username: Username for database access
372 password: Password for database access
374 Returns:
375 List of time series data points with search engine activity
376 """
377 with self.db.get_session(
378 username=username, password=password
379 ) as session:
380 try:
381 # Build base query
382 query = session.query(SearchCall).filter(
383 SearchCall.search_engine.isnot(None),
384 SearchCall.timestamp.isnot(None),
385 )
387 # Apply time filter
388 time_condition = get_time_filter_condition(
389 period, SearchCall.timestamp
390 )
391 if time_condition is not None: 391 ↛ 395line 391 didn't jump to line 395 because the condition on line 391 was always true
392 query = query.filter(time_condition)
394 # Apply research mode filter
395 mode_condition = get_research_mode_condition(
396 research_mode, SearchCall.research_mode
397 )
398 if mode_condition is not None:
399 query = query.filter(mode_condition)
401 # Get all search calls ordered by time
402 search_calls = query.order_by(SearchCall.timestamp.asc()).all()
404 # Create time series data
405 time_series = []
406 for call in search_calls:
407 time_series.append(
408 {
409 "timestamp": (
410 str(call.timestamp) if call.timestamp else None
411 ),
412 "search_engine": call.search_engine,
413 "results_count": call.results_count or 0,
414 "response_time_ms": call.response_time_ms or 0,
415 "success_status": call.success_status,
416 "query": (
417 call.query[:50] + "..."
418 if call.query and len(call.query) > 50
419 else call.query
420 ),
421 }
422 )
424 return time_series
426 except Exception:
427 logger.warning("Error getting search time series")
428 return []
431def get_search_tracker() -> SearchTracker:
432 """Create a SearchTracker instance.
434 Returns a fresh instance each time. Callers should pass username/password
435 to the query methods so the correct per-user database is accessed.
436 """
437 return SearchTracker()