Coverage for src/local_deep_research/metrics/token_counter.py: 96%
500 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-20 01:24 +0000
1"""Token counting functionality for LLM usage tracking."""
3import inspect
4import json
5import time
6from datetime import datetime
7from pathlib import Path
8from typing import Any, Dict, List, Optional
10from langchain_core.callbacks import BaseCallbackHandler
11from langchain_core.outputs import LLMResult
12from loguru import logger
13from sqlalchemy import func, text
15from ..database.models import ModelUsage, TokenUsage
16from .query_utils import (
17 get_period_cutoff,
18 get_research_mode_condition,
19 get_time_filter_condition,
20)
23class TokenCountingCallback(BaseCallbackHandler):
24 """Callback handler for counting tokens across different models."""
26 def __init__(
27 self,
28 research_id: Optional[str] = None,
29 research_context: Optional[Dict[str, Any]] = None,
30 ):
31 """Initialize the token counting callback.
33 Args:
34 research_id: The ID of the research to track tokens for
35 research_context: Additional research context for enhanced tracking
36 """
37 super().__init__()
38 self.research_id = research_id
39 self.research_context = research_context or {}
40 self.current_model = None
41 self.current_provider = None
42 self.preset_model = None # Model name set during callback creation
43 self.preset_provider = None # Provider set during callback creation
45 # Phase 1 Enhancement: Track timing and context
46 self.start_time = None
47 self.response_time_ms = None
48 self.success_status = "success"
49 self.error_type = None
51 # Call stack tracking
52 self.calling_file = None
53 self.calling_function = None
54 self.call_stack = None
56 # Context overflow tracking
57 self.context_limit = None
58 self.context_truncated = False
59 self.tokens_truncated = 0
60 self.truncation_ratio = 0.0
61 self.original_prompt_estimate = 0
63 # Raw Ollama response metrics
64 self.ollama_metrics = {}
66 # Whether we've already logged the "provider reports no usage"
67 # warning. The callback is shared across every LLM call in a
68 # research session (see llm_config.py wrap_llm), so without this a
69 # 50-call run against a non-reporting provider logs 50 identical
70 # warnings. Calls are still recorded every time — only the warning
71 # is deduplicated.
72 self._warned_no_usage = False
74 # Track token counts in memory
75 self.counts = {
76 "total_tokens": 0,
77 "total_prompt_tokens": 0,
78 "total_completion_tokens": 0,
79 "by_model": {},
80 }
82 def on_llm_start(
83 self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
84 ) -> None:
85 """Called when LLM starts running."""
86 # Phase 1 Enhancement: Start timing
87 self.start_time = time.time()
89 # Reset per-call truncation state. The callback instance is shared
90 # across every LLM call in a research session (see llm_config.py
91 # wrap_llm), so without this reset the post-loop estimation block's
92 # `if not self.context_truncated` guard would silently disable
93 # [estimated] / [estimated-total-context] detection on every call
94 # after the first one that truncates.
95 self.context_truncated = False
96 self.tokens_truncated = 0
97 self.truncation_ratio = 0.0
99 # Estimate original prompt size (rough estimate: ~4 chars per token)
100 if prompts:
101 total_chars = sum(len(prompt) for prompt in prompts)
102 self.original_prompt_estimate = total_chars // 4
103 logger.debug(
104 f"Estimated prompt tokens: {self.original_prompt_estimate} (from {total_chars} chars)"
105 )
107 # Get context limit from research context (will be set from settings)
108 self.context_limit = self.research_context.get("context_limit")
110 # Phase 1 Enhancement: Capture call stack information
111 try:
112 stack = inspect.stack()
114 # Skip the first few frames (this method, langchain internals)
115 # Look for the first frame that's in our project directory
116 for frame_info in stack[1:]:
117 file_path = frame_info.filename
118 # Look for any frame containing local_deep_research project
119 if (
120 "local_deep_research" in file_path
121 and "site-packages" not in file_path
122 and "venv" not in file_path
123 ):
124 # Extract relative path from local_deep_research
125 if "src/local_deep_research" in file_path:
126 relative_path = file_path.split(
127 "src/local_deep_research"
128 )[-1].lstrip("/")
129 elif "local_deep_research/src" in file_path: 129 ↛ 133line 129 didn't jump to line 133 because the condition on line 129 was always true
130 relative_path = file_path.split(
131 "local_deep_research/src"
132 )[-1].lstrip("/")
133 elif "local_deep_research" in file_path:
134 # Get everything after local_deep_research
135 relative_path = file_path.split("local_deep_research")[
136 -1
137 ].lstrip("/")
138 else:
139 relative_path = Path(file_path).name
141 self.calling_file = relative_path
142 self.calling_function = frame_info.function
144 # Capture a simplified call stack (just the relevant frames)
145 call_stack_frames = []
146 for frame in stack[1:6]: # Limit to 5 frames
147 if (
148 "local_deep_research" in frame.filename
149 and "site-packages" not in frame.filename
150 and "venv" not in frame.filename
151 ):
152 frame_name = f"{Path(frame.filename).name}:{frame.function}:{frame.lineno}"
153 call_stack_frames.append(frame_name)
155 self.call_stack = (
156 " -> ".join(call_stack_frames)
157 if call_stack_frames
158 else None
159 )
160 break
161 except Exception:
162 logger.warning("Error capturing call stack")
163 # Continue without call stack info if there's an error
165 # First, use preset values if available
166 if self.preset_model:
167 self.current_model = self.preset_model
168 else:
169 # Try multiple locations for model name
170 model_name = None
172 # First check invocation_params
173 invocation_params = kwargs.get("invocation_params", {})
174 model_name = invocation_params.get(
175 "model"
176 ) or invocation_params.get("model_name")
178 # Check kwargs directly
179 if not model_name:
180 model_name = kwargs.get("model") or kwargs.get("model_name")
182 # Check serialized data
183 if not model_name and "kwargs" in serialized:
184 model_name = serialized["kwargs"].get("model") or serialized[
185 "kwargs"
186 ].get("model_name")
188 # Check for name in serialized data
189 if not model_name and "name" in serialized:
190 model_name = serialized["name"]
192 # If still not found and we have Ollama, try to extract from the instance
193 if (
194 not model_name
195 and "_type" in serialized
196 and "ChatOllama" in serialized["_type"]
197 ):
198 # For Ollama, the model name might be in the serialized kwargs
199 if "kwargs" in serialized and "model" in serialized["kwargs"]: 199 ↛ 200line 199 didn't jump to line 200 because the condition on line 199 was never true
200 model_name = serialized["kwargs"]["model"]
201 else:
202 # Default to the type if we can't find the actual model
203 model_name = "ollama"
205 # Final fallback
206 if not model_name:
207 if "_type" in serialized:
208 model_name = serialized["_type"]
209 else:
210 model_name = "unknown"
212 self.current_model = model_name
214 # Use preset provider if available
215 if self.preset_provider:
216 self.current_provider = self.preset_provider
217 else:
218 # Extract provider from serialized type or kwargs
219 if "_type" in serialized:
220 type_str = serialized["_type"]
221 if "ChatOllama" in type_str:
222 self.current_provider = "ollama"
223 elif "ChatOpenAI" in type_str:
224 self.current_provider = "openai"
225 elif "ChatAnthropic" in type_str:
226 self.current_provider = "anthropic"
227 else:
228 self.current_provider = kwargs.get("provider", "unknown")
229 else:
230 self.current_provider = kwargs.get("provider", "unknown")
232 # Initialize model tracking if needed
233 if self.current_model not in self.counts["by_model"]:
234 self.counts["by_model"][self.current_model] = {
235 "prompt_tokens": 0,
236 "completion_tokens": 0,
237 "total_tokens": 0,
238 "calls": 0,
239 "provider": self.current_provider,
240 }
242 # Increment call count
243 self.counts["by_model"][self.current_model]["calls"] += 1
245 def _check_context_overflow(
246 self,
247 prompt_eval_count: int,
248 completion_tokens: int = 0,
249 source: str = "",
250 ) -> None:
251 """Check for context overflow based on prompt and total token usage.
253 Args:
254 prompt_eval_count: Number of tokens the model actually processed.
255 completion_tokens: Number of tokens generated (for total-context check).
256 source: Which branch provided the data (for logging).
257 """
258 logger.debug(
259 f"Context overflow check [{source}]: "
260 f"prompt_eval_count={prompt_eval_count}, "
261 f"completion_tokens={completion_tokens}, "
262 f"context_limit={self.context_limit}"
263 )
265 if not self.context_limit or prompt_eval_count <= 0:
266 return
268 # Input-only overflow: prompt at >= 80% of context limit. Matches the
269 # chart-warning threshold and PR #3840's deliberate choice (PR #3792
270 # lowered from 95% → 80%). The total-context branch below uses 95%
271 # because it's a stricter condition (input+output combined).
272 if prompt_eval_count >= self.context_limit * 0.80:
273 self.context_truncated = True
275 if self.original_prompt_estimate > prompt_eval_count:
276 self.tokens_truncated = max(
277 0,
278 self.original_prompt_estimate - prompt_eval_count,
279 )
280 if ( 280 ↛ 294line 280 didn't jump to line 294 because the condition on line 280 was always true
281 self.tokens_truncated > 0
282 and self.original_prompt_estimate > 0
283 ):
284 self.truncation_ratio = (
285 self.tokens_truncated / self.original_prompt_estimate
286 )
287 elif prompt_eval_count > self.context_limit: 287 ↛ 288line 287 didn't jump to line 288 because the condition on line 287 was never true
288 self.tokens_truncated = prompt_eval_count - self.context_limit
289 if self.tokens_truncated > 0 and prompt_eval_count > 0:
290 self.truncation_ratio = (
291 self.tokens_truncated / prompt_eval_count
292 )
294 logger.warning(
295 f"Context overflow detected [provider-confirmed] "
296 f"research_id={self.research_id} "
297 f"model={self.current_model} "
298 f"provider={self.current_provider} "
299 f"source={source} "
300 f"prompt_tokens={prompt_eval_count} "
301 f"context_limit={self.context_limit} "
302 f"tokens_truncated={self.tokens_truncated} "
303 f"truncation_ratio={self.truncation_ratio:.1%}"
304 )
306 # Total-context overflow: input + output exceeds 95% of context limit
307 elif (
308 completion_tokens > 0
309 and prompt_eval_count + completion_tokens
310 >= self.context_limit * 0.95
311 ):
312 total = prompt_eval_count + completion_tokens
313 self.context_truncated = True
314 self.tokens_truncated = max(0, total - self.context_limit)
315 self.truncation_ratio = (
316 self.tokens_truncated / total if total > 0 else 0
317 )
318 logger.warning(
319 f"Context overflow detected [total-context] "
320 f"research_id={self.research_id} "
321 f"model={self.current_model} "
322 f"provider={self.current_provider} "
323 f"source={source} "
324 f"prompt_tokens={prompt_eval_count} "
325 f"completion_tokens={completion_tokens} "
326 f"total_tokens={total} "
327 f"context_limit={self.context_limit} "
328 f"tokens_truncated={self.tokens_truncated} "
329 f"truncation_ratio={self.truncation_ratio:.1%}"
330 )
331 else:
332 logger.debug(
333 f"Context OK [{source}]: "
334 f"{prompt_eval_count}/{self.context_limit} "
335 f"({prompt_eval_count / self.context_limit:.1%})"
336 )
338 def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
339 """Called when LLM ends running."""
340 # Phase 1 Enhancement: Calculate response time
341 if self.start_time:
342 self.response_time_ms = int((time.time() - self.start_time) * 1000)
344 # Extract token usage from response
345 token_usage = None
347 # Check multiple locations for token usage
348 if hasattr(response, "llm_output") and response.llm_output:
349 token_usage = response.llm_output.get(
350 "token_usage"
351 ) or response.llm_output.get("usage", {})
353 # Check for usage metadata in generations (Ollama specific)
354 if not token_usage and hasattr(response, "generations"):
355 for generation_list in response.generations:
356 for generation in generation_list:
357 if hasattr(generation, "message") and hasattr(
358 generation.message, "usage_metadata"
359 ):
360 usage_meta = generation.message.usage_metadata
361 if usage_meta: # Check if usage_metadata is not None
362 token_usage = {
363 "prompt_tokens": usage_meta.get(
364 "input_tokens", 0
365 ),
366 "completion_tokens": usage_meta.get(
367 "output_tokens", 0
368 ),
369 "total_tokens": usage_meta.get(
370 "total_tokens", 0
371 ),
372 }
373 # Check context overflow before breaking
374 # (input_tokens == prompt_eval_count for Ollama)
375 self._check_context_overflow(
376 usage_meta.get("input_tokens", 0),
377 completion_tokens=usage_meta.get(
378 "output_tokens", 0
379 ),
380 source="usage_metadata",
381 )
382 break
383 # Also check response_metadata
384 if hasattr(generation, "message") and hasattr(
385 generation.message, "response_metadata"
386 ):
387 resp_meta = generation.message.response_metadata
388 if resp_meta.get("prompt_eval_count") or resp_meta.get(
389 "eval_count"
390 ):
391 # Capture raw Ollama metrics
392 self.ollama_metrics = {
393 "prompt_eval_count": resp_meta.get(
394 "prompt_eval_count"
395 ),
396 "eval_count": resp_meta.get("eval_count"),
397 "total_duration": resp_meta.get(
398 "total_duration"
399 ),
400 "load_duration": resp_meta.get("load_duration"),
401 "prompt_eval_duration": resp_meta.get(
402 "prompt_eval_duration"
403 ),
404 "eval_duration": resp_meta.get("eval_duration"),
405 }
407 # Check for context overflow (input only)
408 prompt_eval_count = resp_meta.get(
409 "prompt_eval_count", 0
410 )
411 self._check_context_overflow(
412 prompt_eval_count,
413 completion_tokens=resp_meta.get(
414 "eval_count", 0
415 ),
416 source="response_metadata",
417 )
419 token_usage = {
420 "prompt_tokens": resp_meta.get(
421 "prompt_eval_count", 0
422 ),
423 "completion_tokens": resp_meta.get(
424 "eval_count", 0
425 ),
426 "total_tokens": resp_meta.get(
427 "prompt_eval_count", 0
428 )
429 + resp_meta.get("eval_count", 0),
430 }
431 break
432 if token_usage:
433 break
435 # Estimation-based overflow detection for providers that don't echo
436 # prompt_eval_count (OpenAI, Anthropic, OpenRouter, etc.). The Ollama
437 # path above sets context_truncated=True if it detected provider-
438 # confirmed truncation; we only fire here if it didn't. For hosted
439 # providers the API typically rejects oversize prompts rather than
440 # silently truncating, so this signal flags "would-overflow per our
441 # estimate" — not the same as actual truncation, hence the [estimated]
442 # tag in the log message.
443 if not self.context_truncated and self.context_limit:
444 # Input-only overflow: prompt estimate exceeds context limit
445 if self.original_prompt_estimate > self.context_limit:
446 self.context_truncated = True
447 self.tokens_truncated = max(
448 0, self.original_prompt_estimate - self.context_limit
449 )
450 self.truncation_ratio = (
451 self.tokens_truncated / self.original_prompt_estimate
452 if self.original_prompt_estimate > 0
453 else 0
454 )
455 logger.warning(
456 "Context overflow detected [estimated] "
457 f"research_id={self.research_id} "
458 f"model={self.current_model} "
459 f"provider={self.current_provider} "
460 f"estimated_prompt_tokens={self.original_prompt_estimate} "
461 f"context_limit={self.context_limit} "
462 f"tokens_truncated={self.tokens_truncated} "
463 f"truncation_ratio={self.truncation_ratio:.1%}"
464 )
465 # Total-context overflow: input + output exceeds context limit.
466 # The "[estimated-total-context]" tag below refers to the
467 # *detection method* (post-loop fallback path that runs when
468 # _check_context_overflow couldn't fire), not the data source —
469 # the prompt_tokens / completion_tokens here come from the
470 # provider via response.llm_output.token_usage and are actual
471 # counts, not character-based estimates.
472 elif token_usage and isinstance(token_usage, dict): 472 ↛ 495line 472 didn't jump to line 495 because the condition on line 472 was always true
473 prompt_tokens = token_usage.get("prompt_tokens", 0)
474 completion_tokens = token_usage.get("completion_tokens", 0)
475 total = prompt_tokens + completion_tokens
476 if total >= self.context_limit * 0.95:
477 self.context_truncated = True
478 self.tokens_truncated = max(0, total - self.context_limit)
479 self.truncation_ratio = (
480 self.tokens_truncated / total if total > 0 else 0
481 )
482 logger.warning(
483 "Context overflow detected [estimated-total-context] "
484 f"research_id={self.research_id} "
485 f"model={self.current_model} "
486 f"provider={self.current_provider} "
487 f"prompt_tokens={prompt_tokens} "
488 f"completion_tokens={completion_tokens} "
489 f"total_tokens={total} "
490 f"context_limit={self.context_limit} "
491 f"tokens_truncated={self.tokens_truncated} "
492 f"truncation_ratio={self.truncation_ratio:.1%}"
493 )
495 if token_usage and isinstance(token_usage, dict):
496 prompt_tokens = token_usage.get("prompt_tokens", 0)
497 completion_tokens = token_usage.get("completion_tokens", 0)
498 total_tokens = token_usage.get(
499 "total_tokens", prompt_tokens + completion_tokens
500 )
502 # Update in-memory counts
503 self.counts["total_prompt_tokens"] += prompt_tokens
504 self.counts["total_completion_tokens"] += completion_tokens
505 self.counts["total_tokens"] += total_tokens
507 if self.current_model:
508 self.counts["by_model"][self.current_model][
509 "prompt_tokens"
510 ] += prompt_tokens
511 self.counts["by_model"][self.current_model][
512 "completion_tokens"
513 ] += completion_tokens
514 self.counts["by_model"][self.current_model]["total_tokens"] += (
515 total_tokens
516 )
518 # Save to database if we have a research_id
519 if self.research_id:
520 self._save_to_db(prompt_tokens, completion_tokens)
521 elif self.research_id:
522 # The provider returned no usage data at all (e.g. OpenAI-
523 # compatible servers that omit `usage` on streamed responses,
524 # proxies that strip it, Ollama omitting prompt_eval_count for
525 # fully-cached prompts). Record the call anyway with zero token
526 # counts — model, provider, phase, response time and context-
527 # overflow estimates are still real data. Skipping the row
528 # entirely is what made every metrics page render as if LDR had
529 # never been used (#4457).
530 # Warn once per research session, not once per call.
531 if not self._warned_no_usage:
532 self._warned_no_usage = True
533 hint = ""
534 if self.current_provider == "openai_endpoint": 534 ↛ 535line 534 didn't jump to line 535 because the condition on line 534 was never true
535 hint = (
536 " If your server supports stream_options.include_usage "
537 "(LM Studio 0.3.18+, llama.cpp, vLLM, OpenRouter), "
538 "enable the 'llm.openai_endpoint.stream_usage' setting "
539 "to get real token counts on streamed calls."
540 )
541 logger.warning(
542 f"LLM provider returned no token usage data - recording "
543 f"calls with zero counts. model={self.current_model} "
544 f"provider={self.current_provider} "
545 f"research_id={self.research_id}.{hint} "
546 f"(further occurrences this session are suppressed)"
547 )
548 self._save_to_db(0, 0)
550 def on_llm_error(self, error, **kwargs: Any) -> None:
551 """Called when LLM encounters an error."""
552 # Phase 1 Enhancement: Track errors
553 if self.start_time:
554 self.response_time_ms = int((time.time() - self.start_time) * 1000)
556 self.success_status = "error"
557 self.error_type = str(type(error).__name__)
559 # Still save to database to track failed calls
560 if self.research_id:
561 self._save_to_db(0, 0)
563 def _get_context_overflow_fields(self) -> Dict[str, Any]:
564 """Get context overflow detection fields for database saving."""
565 return {
566 "context_limit": self.context_limit,
567 "context_truncated": self.context_truncated, # Now Boolean
568 "tokens_truncated": self.tokens_truncated
569 if self.context_truncated
570 else None,
571 "truncation_ratio": self.truncation_ratio
572 if self.context_truncated
573 else None,
574 # Raw Ollama metrics
575 "ollama_prompt_eval_count": self.ollama_metrics.get(
576 "prompt_eval_count"
577 ),
578 "ollama_eval_count": self.ollama_metrics.get("eval_count"),
579 "ollama_total_duration": self.ollama_metrics.get("total_duration"),
580 "ollama_load_duration": self.ollama_metrics.get("load_duration"),
581 "ollama_prompt_eval_duration": self.ollama_metrics.get(
582 "prompt_eval_duration"
583 ),
584 "ollama_eval_duration": self.ollama_metrics.get("eval_duration"),
585 }
587 def _save_to_db(self, prompt_tokens: int, completion_tokens: int):
588 """Save token usage to the database."""
589 # Check if we're in a thread - if so, queue the save for later
590 import threading
592 if threading.current_thread().name != "MainThread":
593 # Use thread-safe metrics database for background threads
594 username = (
595 self.research_context.get("username")
596 if self.research_context
597 else None
598 )
600 if not username:
601 logger.warning(
602 f"Cannot save token metrics - no username in research context. "
603 f"Token usage: prompt={prompt_tokens}, completion={completion_tokens}, "
604 f"Research context: {self.research_context}"
605 )
606 return
608 # Import the thread-safe metrics database
610 # Prepare token data
611 token_data = {
612 "model_name": self.current_model,
613 "provider": self.current_provider,
614 "prompt_tokens": prompt_tokens,
615 "completion_tokens": completion_tokens,
616 "research_query": self.research_context.get("research_query"),
617 "research_mode": self.research_context.get("research_mode"),
618 "research_phase": self.research_context.get("research_phase"),
619 "search_iteration": self.research_context.get(
620 "search_iteration"
621 ),
622 "response_time_ms": self.response_time_ms,
623 "success_status": self.success_status,
624 "error_type": self.error_type,
625 "search_engines_planned": self.research_context.get(
626 "search_engines_planned"
627 ),
628 "search_engine_selected": self.research_context.get(
629 "search_engine_selected"
630 ),
631 "calling_file": self.calling_file,
632 "calling_function": self.calling_function,
633 "call_stack": self.call_stack,
634 # Add context overflow fields using helper method
635 **self._get_context_overflow_fields(),
636 }
638 # Convert list to JSON string if needed
639 if isinstance(token_data.get("search_engines_planned"), list):
640 token_data["search_engines_planned"] = json.dumps(
641 token_data["search_engines_planned"]
642 )
644 # Get password from research context
645 password = self.research_context.get("user_password")
646 if not password:
647 logger.warning(
648 f"Cannot save token metrics - no password in research context. "
649 f"Username: {username}, Token usage: prompt={prompt_tokens}, completion={completion_tokens}"
650 )
651 return
653 # Write metrics directly using thread-safe database
654 try:
655 from ..database.thread_metrics import metrics_writer
657 # Set password for this thread
658 metrics_writer.set_user_password(username, password)
660 # Write metrics to encrypted database
661 metrics_writer.write_token_metrics(
662 username, self.research_id, token_data
663 )
664 except Exception:
665 logger.warning("Failed to write metrics from thread")
666 return
668 # In MainThread, save directly
669 try:
670 from flask import session as flask_session
671 from ..database.session_context import get_user_db_session
673 username = flask_session.get("username")
674 if not username:
675 logger.debug("No user session, skipping token metrics save")
676 return
678 with get_user_db_session(username) as session:
679 # Phase 1 Enhancement: Prepare additional context
680 research_query = self.research_context.get("research_query")
681 research_mode = self.research_context.get("research_mode")
682 research_phase = self.research_context.get("research_phase")
683 search_iteration = self.research_context.get("search_iteration")
684 search_engines_planned = self.research_context.get(
685 "search_engines_planned"
686 )
687 search_engine_selected = self.research_context.get(
688 "search_engine_selected"
689 )
691 # Debug logging for search engine context
692 if search_engines_planned or search_engine_selected:
693 logger.info(
694 f"Token tracking - Search context: planned={search_engines_planned}, selected={search_engine_selected}, phase={research_phase}"
695 )
696 else:
697 logger.debug(
698 f"Token tracking - No search engine context yet, phase={research_phase}"
699 )
701 # Convert list to JSON string if needed
702 if isinstance(search_engines_planned, list):
703 search_engines_planned = json.dumps(search_engines_planned)
705 # Log context overflow detection values before saving
706 logger.debug(
707 f"Saving TokenUsage - context_limit: {self.context_limit}, "
708 f"context_truncated: {self.context_truncated}, "
709 f"tokens_truncated: {self.tokens_truncated}, "
710 f"ollama_prompt_eval_count: {self.ollama_metrics.get('prompt_eval_count')}, "
711 f"prompt_tokens: {prompt_tokens}, "
712 f"completion_tokens: {completion_tokens}"
713 )
715 # Add token usage record with enhanced fields
716 token_usage = TokenUsage(
717 research_id=self.research_id,
718 model_name=self.current_model,
719 model_provider=self.current_provider, # Added provider
720 # for accurate cost tracking
721 prompt_tokens=prompt_tokens,
722 completion_tokens=completion_tokens,
723 total_tokens=prompt_tokens + completion_tokens,
724 # Phase 1 Enhancement: Research context
725 research_query=research_query,
726 research_mode=research_mode,
727 research_phase=research_phase,
728 search_iteration=search_iteration,
729 # Phase 1 Enhancement: Performance metrics
730 response_time_ms=self.response_time_ms,
731 success_status=self.success_status,
732 error_type=self.error_type,
733 # Phase 1 Enhancement: Search engine context
734 search_engines_planned=search_engines_planned,
735 search_engine_selected=search_engine_selected,
736 # Phase 1 Enhancement: Call stack tracking
737 calling_file=self.calling_file,
738 calling_function=self.calling_function,
739 call_stack=self.call_stack,
740 # Add context overflow fields using helper method
741 **self._get_context_overflow_fields(),
742 )
743 session.add(token_usage)
745 # Update or create model usage statistics
746 model_usage = (
747 session.query(ModelUsage)
748 .filter_by(
749 model_name=self.current_model,
750 )
751 .first()
752 )
754 if model_usage:
755 model_usage.total_tokens += (
756 prompt_tokens + completion_tokens
757 )
758 model_usage.total_calls += 1
759 else:
760 model_usage = ModelUsage(
761 model_name=self.current_model,
762 model_provider=self.current_provider,
763 total_tokens=prompt_tokens + completion_tokens,
764 total_calls=1,
765 )
766 session.add(model_usage)
768 # Commit the transaction
769 session.commit()
771 except Exception:
772 logger.warning("Error saving token usage to database")
774 def get_counts(self) -> Dict[str, Any]:
775 """Get the current token counts."""
776 return self.counts
779class TokenCounter:
780 """Manager class for token counting across the application."""
782 def __init__(self):
783 """Initialize the token counter."""
785 def create_callback(
786 self,
787 research_id: Optional[str] = None,
788 research_context: Optional[Dict[str, Any]] = None,
789 ) -> TokenCountingCallback:
790 """Create a new token counting callback.
792 Args:
793 research_id: The ID of the research to track tokens for
794 research_context: Additional research context for enhanced tracking
796 Returns:
797 A new TokenCountingCallback instance
798 """
799 return TokenCountingCallback(
800 research_id=research_id, research_context=research_context
801 )
803 def get_research_metrics(self, research_id: str) -> Dict[str, Any]:
804 """Get token metrics for a specific research.
806 Args:
807 research_id: The ID of the research
809 Returns:
810 Dictionary containing token usage metrics
811 """
812 from flask import session as flask_session
814 from ..database.session_context import get_user_db_session
816 username = flask_session.get("username")
817 if not username:
818 return {
819 "research_id": research_id,
820 "total_tokens": 0,
821 "total_calls": 0,
822 "model_usage": [],
823 }
825 with get_user_db_session(username) as session:
826 # Get token usage for this research from TokenUsage table
827 from sqlalchemy import func
829 token_usages = (
830 session.query(
831 TokenUsage.model_name,
832 TokenUsage.model_provider,
833 func.sum(TokenUsage.prompt_tokens).label("prompt_tokens"),
834 func.sum(TokenUsage.completion_tokens).label(
835 "completion_tokens"
836 ),
837 func.sum(TokenUsage.total_tokens).label("total_tokens"),
838 func.count().label("calls"),
839 )
840 .filter_by(research_id=research_id)
841 .group_by(TokenUsage.model_name, TokenUsage.model_provider)
842 .order_by(func.sum(TokenUsage.total_tokens).desc())
843 .all()
844 )
846 model_usage = []
847 total_tokens = 0
848 total_calls = 0
850 for usage in token_usages:
851 model_usage.append(
852 {
853 "model": usage.model_name,
854 "provider": usage.model_provider,
855 "tokens": usage.total_tokens or 0,
856 "calls": usage.calls or 0,
857 "prompt_tokens": usage.prompt_tokens or 0,
858 "completion_tokens": usage.completion_tokens or 0,
859 }
860 )
861 total_tokens += usage.total_tokens or 0
862 total_calls += usage.calls or 0
864 return {
865 "research_id": research_id,
866 "total_tokens": total_tokens,
867 "total_calls": total_calls,
868 "model_usage": model_usage,
869 }
871 def get_overall_metrics(
872 self, period: str = "30d", research_mode: str = "all"
873 ) -> Dict[str, Any]:
874 """Get overall token metrics across all researches.
876 Args:
877 period: Time period to filter by ('7d', '30d', '3m', '1y', 'all')
878 research_mode: Research mode to filter by ('quick', 'detailed', 'all')
880 Returns:
881 Dictionary containing overall metrics
882 """
883 return self._get_metrics_from_encrypted_db(period, research_mode)
885 def _get_metrics_from_encrypted_db(
886 self, period: str, research_mode: str
887 ) -> Dict[str, Any]:
888 """Get metrics from user's encrypted database."""
889 from flask import session as flask_session
891 from ..database.session_context import get_user_db_session
893 username = flask_session.get("username")
894 if not username:
895 return self._get_empty_metrics()
897 try:
898 with get_user_db_session(username) as session:
899 # Build base query with filters
900 query = session.query(TokenUsage)
902 # Apply time filter
903 time_condition = get_time_filter_condition(
904 period, TokenUsage.timestamp
905 )
906 if time_condition is not None:
907 query = query.filter(time_condition)
909 # Apply research mode filter
910 mode_condition = get_research_mode_condition(
911 research_mode, TokenUsage.research_mode
912 )
913 if mode_condition is not None:
914 query = query.filter(mode_condition)
916 # Total tokens from TokenUsage
917 total_tokens = (
918 query.with_entities(
919 func.sum(TokenUsage.total_tokens)
920 ).scalar()
921 or 0
922 )
924 # Import ResearchHistory model
925 from ..database.models.research import ResearchHistory
927 # Count researches from ResearchHistory table
928 research_query = session.query(func.count(ResearchHistory.id))
930 # Debug: Check if any research history records exist at all
931 all_research_count = (
932 session.query(func.count(ResearchHistory.id)).scalar() or 0
933 )
934 logger.debug(
935 f"Total ResearchHistory records in database: {all_research_count}"
936 )
938 # Debug: List first few research IDs and their timestamps
939 sample_researches = (
940 session.query(
941 ResearchHistory.id,
942 ResearchHistory.created_at,
943 ResearchHistory.mode,
944 )
945 .limit(5)
946 .all()
947 )
948 if sample_researches:
949 logger.debug("Sample ResearchHistory records:")
950 for r_id, r_created, r_mode in sample_researches:
951 logger.debug(
952 f" - ID: {r_id}, Created: {r_created}, Mode: {r_mode}"
953 )
954 else:
955 logger.debug("No ResearchHistory records found in database")
957 # Time filter for the ResearchHistory count. created_at is
958 # an ISO-8601 TEXT column, so compare isoformat strings.
959 # Must use the API's period vocabulary ('7d'/'30d'/'3m'/
960 # '1y'/'all') — the previous 'today'/'week'/'month' branches
961 # never matched it, so the research count silently ignored
962 # the selected period and always showed all-time.
963 start_time = get_period_cutoff(period)
964 if start_time is not None:
965 research_query = research_query.filter(
966 ResearchHistory.created_at >= start_time.isoformat()
967 )
969 # Apply mode filter if specified
970 mode_filter = research_mode if research_mode != "all" else None
971 if mode_filter:
972 logger.debug(f"Applying mode filter: {mode_filter}")
973 research_query = research_query.filter(
974 ResearchHistory.mode == mode_filter
975 )
977 total_researches = research_query.scalar() or 0
978 logger.debug(
979 f"Final filtered research count: {total_researches}"
980 )
982 # Also check distinct research_ids in TokenUsage for comparison
983 token_research_count = (
984 session.query(
985 func.count(func.distinct(TokenUsage.research_id))
986 ).scalar()
987 or 0
988 )
989 logger.debug(
990 f"Distinct research_ids in TokenUsage: {token_research_count}"
991 )
993 # Model statistics using ORM aggregation
994 model_stats_query = session.query(
995 TokenUsage.model_name,
996 func.sum(TokenUsage.total_tokens).label("tokens"),
997 func.count().label("calls"),
998 func.sum(TokenUsage.prompt_tokens).label("prompt_tokens"),
999 func.sum(TokenUsage.completion_tokens).label(
1000 "completion_tokens"
1001 ),
1002 ).filter(TokenUsage.model_name.isnot(None))
1004 # Apply same filters to model stats
1005 if time_condition is not None:
1006 model_stats_query = model_stats_query.filter(time_condition)
1007 if mode_condition is not None:
1008 model_stats_query = model_stats_query.filter(mode_condition)
1010 model_stats = (
1011 model_stats_query.group_by(TokenUsage.model_name)
1012 .order_by(func.sum(TokenUsage.total_tokens).desc())
1013 .all()
1014 )
1016 # Batch load provider info from ModelUsage table (fix N+1)
1017 model_names = [stat.model_name for stat in model_stats]
1018 provider_map = {}
1019 if model_names:
1020 provider_results = (
1021 session.query(
1022 ModelUsage.model_name, ModelUsage.model_provider
1023 )
1024 .filter(ModelUsage.model_name.in_(model_names))
1025 .order_by(ModelUsage.id)
1026 .all()
1027 )
1028 for model_name, model_provider in provider_results:
1029 provider_map.setdefault(model_name, model_provider)
1031 by_model = []
1032 for stat in model_stats:
1033 provider = provider_map.get(stat.model_name, "unknown")
1035 by_model.append(
1036 {
1037 "model": stat.model_name,
1038 "provider": provider,
1039 "tokens": stat.tokens,
1040 "calls": stat.calls,
1041 "prompt_tokens": stat.prompt_tokens,
1042 "completion_tokens": stat.completion_tokens,
1043 }
1044 )
1046 # Get recent researches with token usage
1047 # Note: This requires research_history table - for now we'll use available data
1048 recent_research_query = session.query(
1049 TokenUsage.research_id,
1050 func.sum(TokenUsage.total_tokens).label("token_count"),
1051 func.max(TokenUsage.timestamp).label("latest_timestamp"),
1052 ).filter(TokenUsage.research_id.isnot(None))
1054 if time_condition is not None:
1055 recent_research_query = recent_research_query.filter(
1056 time_condition
1057 )
1058 if mode_condition is not None:
1059 recent_research_query = recent_research_query.filter(
1060 mode_condition
1061 )
1063 recent_research_data = (
1064 recent_research_query.group_by(TokenUsage.research_id)
1065 .order_by(func.max(TokenUsage.timestamp).desc())
1066 .limit(10)
1067 .all()
1068 )
1070 # Batch load research queries for recent researches (fix N+1)
1071 recent_research_ids = [
1072 r.research_id for r in recent_research_data
1073 ]
1074 research_query_map = {}
1075 if recent_research_ids:
1076 # Get first non-null research_query for each research_id
1077 query_results = (
1078 session.query(
1079 TokenUsage.research_id, TokenUsage.research_query
1080 )
1081 .filter(
1082 TokenUsage.research_id.in_(recent_research_ids),
1083 TokenUsage.research_query.isnot(None),
1084 )
1085 .order_by(TokenUsage.id)
1086 .all()
1087 )
1088 for research_id, research_query in query_results:
1089 if research_id not in research_query_map: 1089 ↛ 1088line 1089 didn't jump to line 1088 because the condition on line 1089 was always true
1090 research_query_map[research_id] = research_query
1092 recent_researches = []
1093 for research_data in recent_research_data:
1094 query_text = research_query_map.get(
1095 research_data.research_id,
1096 f"Research {research_data.research_id}",
1097 )
1099 recent_researches.append(
1100 {
1101 "id": research_data.research_id,
1102 "query": query_text,
1103 "tokens": research_data.token_count or 0,
1104 "created_at": research_data.latest_timestamp,
1105 }
1106 )
1108 # Token breakdown statistics
1109 breakdown_query = query.with_entities(
1110 func.sum(TokenUsage.prompt_tokens).label(
1111 "total_input_tokens"
1112 ),
1113 func.sum(TokenUsage.completion_tokens).label(
1114 "total_output_tokens"
1115 ),
1116 func.avg(TokenUsage.prompt_tokens).label(
1117 "avg_input_tokens"
1118 ),
1119 func.avg(TokenUsage.completion_tokens).label(
1120 "avg_output_tokens"
1121 ),
1122 func.avg(TokenUsage.total_tokens).label("avg_total_tokens"),
1123 )
1124 token_breakdown = breakdown_query.first()
1126 # Get rate limiting metrics
1127 from ..database.models import (
1128 RateLimitAttempt,
1129 RateLimitEstimate,
1130 )
1132 # Get rate limit attempts
1133 rate_limit_query = session.query(RateLimitAttempt)
1135 # Apply time filter
1136 if time_condition is not None:
1137 # RateLimitAttempt uses timestamp as float, not datetime
1138 if period == "7d":
1139 cutoff_time = time.time() - (7 * 24 * 3600)
1140 elif period == "30d":
1141 cutoff_time = time.time() - (30 * 24 * 3600)
1142 elif period == "3m":
1143 cutoff_time = time.time() - (90 * 24 * 3600)
1144 elif period == "1y":
1145 cutoff_time = time.time() - (365 * 24 * 3600)
1146 else: # all
1147 cutoff_time = 0
1149 if cutoff_time > 0:
1150 rate_limit_query = rate_limit_query.filter(
1151 RateLimitAttempt.timestamp >= cutoff_time
1152 )
1154 # Get rate limit statistics
1155 total_attempts = rate_limit_query.count()
1156 successful_attempts = rate_limit_query.filter(
1157 RateLimitAttempt.success
1158 ).count()
1159 failed_attempts = total_attempts - successful_attempts
1161 # Count rate limiting events (failures with RateLimitError)
1162 rate_limit_events = rate_limit_query.filter(
1163 ~RateLimitAttempt.success,
1164 RateLimitAttempt.error_type == "RateLimitError",
1165 ).count()
1167 logger.debug(
1168 f"Rate limit attempts in database: total={total_attempts}, successful={successful_attempts}"
1169 )
1171 # Get all attempts for detailed calculations
1172 attempts = rate_limit_query.all()
1174 # Calculate average wait times
1175 if attempts:
1176 avg_wait_time = sum(a.wait_time for a in attempts) / len(
1177 attempts
1178 )
1179 successful_wait_times = [
1180 a.wait_time for a in attempts if a.success
1181 ]
1182 avg_successful_wait = (
1183 sum(successful_wait_times) / len(successful_wait_times)
1184 if successful_wait_times
1185 else 0
1186 )
1187 else:
1188 avg_wait_time = 0
1189 avg_successful_wait = 0
1191 # Get tracked engines - count distinct engine types from attempts
1192 tracked_engines_query = session.query(
1193 func.count(func.distinct(RateLimitAttempt.engine_type))
1194 )
1195 if cutoff_time > 0:
1196 tracked_engines_query = tracked_engines_query.filter(
1197 RateLimitAttempt.timestamp >= cutoff_time
1198 )
1199 tracked_engines = tracked_engines_query.scalar() or 0
1201 # Get engine-specific stats from attempts
1202 engine_stats = []
1204 # Get distinct engine types from attempts
1205 engine_types_query = session.query(
1206 RateLimitAttempt.engine_type
1207 ).distinct()
1208 if cutoff_time > 0:
1209 engine_types_query = engine_types_query.filter(
1210 RateLimitAttempt.timestamp >= cutoff_time
1211 )
1212 engine_types = [
1213 row.engine_type for row in engine_types_query.all()
1214 ]
1216 # Batch-load all estimates (fix N+1 query)
1217 estimates_map = {}
1218 if engine_types:
1219 all_estimates = (
1220 session.query(RateLimitEstimate)
1221 .filter(RateLimitEstimate.engine_type.in_(engine_types))
1222 .all()
1223 )
1224 estimates_map = {e.engine_type: e for e in all_estimates}
1226 for engine_type in engine_types:
1227 engine_attempts_list = [
1228 a for a in attempts if a.engine_type == engine_type
1229 ]
1230 engine_attempts = len(engine_attempts_list)
1231 engine_success = len(
1232 [a for a in engine_attempts_list if a.success]
1233 )
1235 # Get estimate if exists
1236 estimate = estimates_map.get(engine_type)
1238 # Calculate recent success rate
1239 recent_success_rate = (
1240 (engine_success / engine_attempts * 100)
1241 if engine_attempts > 0
1242 else 0
1243 )
1245 # Determine status based on success rate
1246 if estimate:
1247 status = (
1248 "healthy"
1249 if estimate.success_rate > 0.8
1250 else "degraded"
1251 if estimate.success_rate > 0.5
1252 else "poor"
1253 )
1254 else:
1255 status = (
1256 "healthy"
1257 if recent_success_rate > 80
1258 else "degraded"
1259 if recent_success_rate > 50
1260 else "poor"
1261 )
1263 engine_stat = {
1264 "engine": engine_type,
1265 "base_wait": estimate.base_wait_seconds
1266 if estimate
1267 else 0.0,
1268 "base_wait_seconds": round(
1269 estimate.base_wait_seconds if estimate else 0.0, 2
1270 ),
1271 "min_wait_seconds": round(
1272 estimate.min_wait_seconds if estimate else 0.0, 2
1273 ),
1274 "max_wait_seconds": round(
1275 estimate.max_wait_seconds if estimate else 0.0, 2
1276 ),
1277 "success_rate": round(estimate.success_rate * 100, 1)
1278 if estimate
1279 else recent_success_rate,
1280 "total_attempts": estimate.total_attempts
1281 if estimate
1282 else engine_attempts,
1283 "recent_attempts": engine_attempts,
1284 "recent_success_rate": round(recent_success_rate, 1),
1285 "attempts": engine_attempts,
1286 "status": status,
1287 }
1289 if estimate:
1290 engine_stat["last_updated"] = datetime.fromtimestamp(
1291 estimate.last_updated
1292 ).strftime("%Y-%m-%d %H:%M:%S")
1293 else:
1294 engine_stat["last_updated"] = "Never"
1296 engine_stats.append(engine_stat)
1298 logger.debug(
1299 f"Tracked engines: {tracked_engines}, engine_stats: {engine_stats}"
1300 )
1302 result = {
1303 "total_tokens": total_tokens,
1304 "total_researches": total_researches,
1305 "by_model": by_model,
1306 "recent_researches": recent_researches,
1307 "token_breakdown": {
1308 "total_input_tokens": int(
1309 token_breakdown.total_input_tokens or 0
1310 ),
1311 "total_output_tokens": int(
1312 token_breakdown.total_output_tokens or 0
1313 ),
1314 "avg_input_tokens": int(
1315 token_breakdown.avg_input_tokens or 0
1316 ),
1317 "avg_output_tokens": int(
1318 token_breakdown.avg_output_tokens or 0
1319 ),
1320 "avg_total_tokens": int(
1321 token_breakdown.avg_total_tokens or 0
1322 ),
1323 },
1324 "rate_limiting": {
1325 "total_attempts": total_attempts,
1326 "successful_attempts": successful_attempts,
1327 "failed_attempts": failed_attempts,
1328 "success_rate": (
1329 successful_attempts / total_attempts * 100
1330 )
1331 if total_attempts > 0
1332 else 0,
1333 "rate_limit_events": rate_limit_events,
1334 "avg_wait_time": round(float(avg_wait_time), 2),
1335 "avg_successful_wait": round(
1336 float(avg_successful_wait), 2
1337 ),
1338 "tracked_engines": tracked_engines,
1339 "engine_stats": engine_stats,
1340 "total_engines_tracked": tracked_engines,
1341 "healthy_engines": len(
1342 [
1343 s
1344 for s in engine_stats
1345 if s["status"] == "healthy"
1346 ]
1347 ),
1348 "degraded_engines": len(
1349 [
1350 s
1351 for s in engine_stats
1352 if s["status"] == "degraded"
1353 ]
1354 ),
1355 "poor_engines": len(
1356 [s for s in engine_stats if s["status"] == "poor"]
1357 ),
1358 },
1359 }
1361 logger.debug(
1362 f"Returning from _get_metrics_from_encrypted_db - total_researches: {result['total_researches']}"
1363 )
1364 return result
1365 except Exception:
1366 logger.exception(
1367 "CRITICAL ERROR accessing encrypted database for metrics"
1368 )
1369 return self._get_empty_metrics()
1371 def _get_empty_metrics(self) -> Dict[str, Any]:
1372 """Return empty metrics structure when no data is available."""
1373 return {
1374 "total_tokens": 0,
1375 "total_researches": 0,
1376 "by_model": [],
1377 "recent_researches": [],
1378 "token_breakdown": {
1379 "prompt_tokens": 0,
1380 "completion_tokens": 0,
1381 "avg_prompt_tokens": 0,
1382 "avg_completion_tokens": 0,
1383 "avg_total_tokens": 0,
1384 },
1385 }
1387 def get_enhanced_metrics(
1388 self, period: str = "30d", research_mode: str = "all"
1389 ) -> Dict[str, Any]:
1390 """Get enhanced Phase 1 tracking metrics.
1392 Args:
1393 period: Time period to filter by ('7d', '30d', '3m', '1y', 'all')
1394 research_mode: Research mode to filter by ('quick', 'detailed', 'all')
1396 Returns:
1397 Dictionary containing enhanced metrics data including time series
1398 """
1399 from flask import session as flask_session
1401 from ..database.session_context import get_user_db_session
1403 username = flask_session.get("username")
1404 if not username:
1405 # Return empty metrics structure when no user session
1406 return {
1407 "recent_enhanced_data": [],
1408 "performance_stats": {
1409 "avg_response_time": 0,
1410 "min_response_time": 0,
1411 "max_response_time": 0,
1412 "success_rate": 0,
1413 "error_rate": 0,
1414 "total_enhanced_calls": 0,
1415 },
1416 "mode_breakdown": [],
1417 "search_engine_stats": [],
1418 "phase_breakdown": [],
1419 "time_series_data": [],
1420 "call_stack_analysis": {
1421 "by_file": [],
1422 "by_function": [],
1423 },
1424 }
1426 try:
1427 with get_user_db_session(username) as session:
1428 # Build base query with filters
1429 query = session.query(TokenUsage)
1431 # Apply time filter
1432 time_condition = get_time_filter_condition(
1433 period, TokenUsage.timestamp
1434 )
1435 if time_condition is not None: 1435 ↛ 1439line 1435 didn't jump to line 1439 because the condition on line 1435 was always true
1436 query = query.filter(time_condition)
1438 # Apply research mode filter
1439 mode_condition = get_research_mode_condition(
1440 research_mode, TokenUsage.research_mode
1441 )
1442 if mode_condition is not None: 1442 ↛ 1443line 1442 didn't jump to line 1443 because the condition on line 1442 was never true
1443 query = query.filter(mode_condition)
1445 # Get time series data for the chart - most important for "Token Consumption Over Time"
1446 time_series_query = query.filter(
1447 TokenUsage.timestamp.isnot(None),
1448 TokenUsage.total_tokens > 0,
1449 ).order_by(TokenUsage.timestamp.asc())
1451 # Limit to recent data for performance
1452 if period != "all": 1452 ↛ 1455line 1452 didn't jump to line 1455 because the condition on line 1452 was always true
1453 time_series_query = time_series_query.limit(200)
1455 time_series_data = time_series_query.all()
1457 # Format time series data with cumulative calculations
1458 time_series = []
1459 cumulative_tokens = 0
1460 cumulative_prompt_tokens = 0
1461 cumulative_completion_tokens = 0
1463 for usage in time_series_data: 1463 ↛ 1464line 1463 didn't jump to line 1464 because the loop on line 1463 never started
1464 cumulative_tokens += usage.total_tokens or 0
1465 cumulative_prompt_tokens += usage.prompt_tokens or 0
1466 cumulative_completion_tokens += usage.completion_tokens or 0
1468 time_series.append(
1469 {
1470 "timestamp": str(usage.timestamp)
1471 if usage.timestamp
1472 else None,
1473 "tokens": usage.total_tokens or 0,
1474 "prompt_tokens": usage.prompt_tokens or 0,
1475 "completion_tokens": usage.completion_tokens or 0,
1476 "cumulative_tokens": cumulative_tokens,
1477 "cumulative_prompt_tokens": cumulative_prompt_tokens,
1478 "cumulative_completion_tokens": cumulative_completion_tokens,
1479 "research_id": usage.research_id,
1480 }
1481 )
1483 # Basic performance stats using ORM
1484 performance_query = query.filter(
1485 TokenUsage.response_time_ms.isnot(None)
1486 )
1487 total_calls = performance_query.count()
1489 if total_calls > 0:
1490 avg_response_time = (
1491 performance_query.with_entities(
1492 func.avg(TokenUsage.response_time_ms)
1493 ).scalar()
1494 or 0
1495 )
1496 min_response_time = (
1497 performance_query.with_entities(
1498 func.min(TokenUsage.response_time_ms)
1499 ).scalar()
1500 or 0
1501 )
1502 max_response_time = (
1503 performance_query.with_entities(
1504 func.max(TokenUsage.response_time_ms)
1505 ).scalar()
1506 or 0
1507 )
1508 success_count = performance_query.filter(
1509 TokenUsage.success_status == "success"
1510 ).count()
1511 error_count = performance_query.filter(
1512 TokenUsage.success_status == "error"
1513 ).count()
1515 perf_stats = {
1516 "avg_response_time": round(avg_response_time),
1517 "min_response_time": min_response_time,
1518 "max_response_time": max_response_time,
1519 "success_rate": (
1520 round((success_count / total_calls * 100), 1)
1521 if total_calls > 0
1522 else 0
1523 ),
1524 "error_rate": (
1525 round((error_count / total_calls * 100), 1)
1526 if total_calls > 0
1527 else 0
1528 ),
1529 "total_enhanced_calls": total_calls,
1530 }
1531 else:
1532 perf_stats = {
1533 "avg_response_time": 0,
1534 "min_response_time": 0,
1535 "max_response_time": 0,
1536 "success_rate": 0,
1537 "error_rate": 0,
1538 "total_enhanced_calls": 0,
1539 }
1541 # Research mode breakdown using ORM
1542 mode_stats = (
1543 query.filter(TokenUsage.research_mode.isnot(None))
1544 .with_entities(
1545 TokenUsage.research_mode,
1546 func.count().label("count"),
1547 func.avg(TokenUsage.total_tokens).label("avg_tokens"),
1548 func.avg(TokenUsage.response_time_ms).label(
1549 "avg_response_time"
1550 ),
1551 )
1552 .group_by(TokenUsage.research_mode)
1553 .all()
1554 )
1556 modes = [
1557 {
1558 "mode": stat.research_mode,
1559 "count": stat.count,
1560 "avg_tokens": round(stat.avg_tokens or 0),
1561 "avg_response_time": round(stat.avg_response_time or 0),
1562 }
1563 for stat in mode_stats
1564 ]
1566 # Recent enhanced data (simplified)
1567 recent_enhanced_query = (
1568 query.filter(TokenUsage.research_query.isnot(None))
1569 .order_by(TokenUsage.timestamp.desc())
1570 .limit(50)
1571 )
1573 recent_enhanced_data = recent_enhanced_query.all()
1574 recent_enhanced = [
1575 {
1576 "research_query": usage.research_query,
1577 "research_mode": usage.research_mode,
1578 "research_phase": usage.research_phase,
1579 "search_iteration": usage.search_iteration,
1580 "response_time_ms": usage.response_time_ms,
1581 "success_status": usage.success_status,
1582 "error_type": usage.error_type,
1583 "search_engines_planned": usage.search_engines_planned,
1584 "search_engine_selected": usage.search_engine_selected,
1585 "total_tokens": usage.total_tokens,
1586 "prompt_tokens": usage.prompt_tokens,
1587 "completion_tokens": usage.completion_tokens,
1588 "timestamp": str(usage.timestamp)
1589 if usage.timestamp
1590 else None,
1591 "research_id": usage.research_id,
1592 "calling_file": usage.calling_file,
1593 "calling_function": usage.calling_function,
1594 "call_stack": usage.call_stack,
1595 }
1596 for usage in recent_enhanced_data
1597 ]
1599 # Search engine breakdown using ORM
1600 search_engine_stats = (
1601 query.filter(TokenUsage.search_engine_selected.isnot(None))
1602 .with_entities(
1603 TokenUsage.search_engine_selected,
1604 func.count().label("count"),
1605 func.avg(TokenUsage.total_tokens).label("avg_tokens"),
1606 func.avg(TokenUsage.response_time_ms).label(
1607 "avg_response_time"
1608 ),
1609 )
1610 .group_by(TokenUsage.search_engine_selected)
1611 .all()
1612 )
1614 search_engines = [
1615 {
1616 "search_engine": stat.search_engine_selected,
1617 "count": stat.count,
1618 "avg_tokens": round(stat.avg_tokens or 0),
1619 "avg_response_time": round(stat.avg_response_time or 0),
1620 }
1621 for stat in search_engine_stats
1622 ]
1624 # Research phase breakdown using ORM
1625 phase_stats = (
1626 query.filter(TokenUsage.research_phase.isnot(None))
1627 .with_entities(
1628 TokenUsage.research_phase,
1629 func.count().label("count"),
1630 func.avg(TokenUsage.total_tokens).label("avg_tokens"),
1631 func.avg(TokenUsage.response_time_ms).label(
1632 "avg_response_time"
1633 ),
1634 )
1635 .group_by(TokenUsage.research_phase)
1636 .all()
1637 )
1639 phases = [
1640 {
1641 "phase": stat.research_phase,
1642 "count": stat.count,
1643 "avg_tokens": round(stat.avg_tokens or 0),
1644 "avg_response_time": round(stat.avg_response_time or 0),
1645 }
1646 for stat in phase_stats
1647 ]
1649 # Call stack analysis using ORM
1650 file_stats = (
1651 query.filter(TokenUsage.calling_file.isnot(None))
1652 .with_entities(
1653 TokenUsage.calling_file,
1654 func.count().label("count"),
1655 func.avg(TokenUsage.total_tokens).label("avg_tokens"),
1656 )
1657 .group_by(TokenUsage.calling_file)
1658 .order_by(func.count().desc())
1659 .limit(10)
1660 .all()
1661 )
1663 files = [
1664 {
1665 "file": stat.calling_file,
1666 "count": stat.count,
1667 "avg_tokens": round(stat.avg_tokens or 0),
1668 }
1669 for stat in file_stats
1670 ]
1672 function_stats = (
1673 query.filter(TokenUsage.calling_function.isnot(None))
1674 .with_entities(
1675 TokenUsage.calling_function,
1676 func.count().label("count"),
1677 func.avg(TokenUsage.total_tokens).label("avg_tokens"),
1678 )
1679 .group_by(TokenUsage.calling_function)
1680 .order_by(func.count().desc())
1681 .limit(10)
1682 .all()
1683 )
1685 functions = [
1686 {
1687 "function": stat.calling_function,
1688 "count": stat.count,
1689 "avg_tokens": round(stat.avg_tokens or 0),
1690 }
1691 for stat in function_stats
1692 ]
1694 return {
1695 "recent_enhanced_data": recent_enhanced,
1696 "performance_stats": perf_stats,
1697 "mode_breakdown": modes,
1698 "search_engine_stats": search_engines,
1699 "phase_breakdown": phases,
1700 "time_series_data": time_series,
1701 "call_stack_analysis": {
1702 "by_file": files,
1703 "by_function": functions,
1704 },
1705 }
1706 except Exception:
1707 logger.exception("Error in get_enhanced_metrics")
1708 # Return simplified response without non-existent columns
1709 return {
1710 "recent_enhanced_data": [],
1711 "performance_stats": {
1712 "avg_response_time": 0,
1713 "min_response_time": 0,
1714 "max_response_time": 0,
1715 "success_rate": 0,
1716 "error_rate": 0,
1717 "total_enhanced_calls": 0,
1718 },
1719 "mode_breakdown": [],
1720 "search_engine_stats": [],
1721 "phase_breakdown": [],
1722 "time_series_data": [],
1723 "call_stack_analysis": {
1724 "by_file": [],
1725 "by_function": [],
1726 },
1727 }
1729 def get_research_timeline_metrics(self, research_id: str) -> Dict[str, Any]:
1730 """Get timeline metrics for a specific research.
1732 Args:
1733 research_id: The ID of the research
1735 Returns:
1736 Dictionary containing timeline metrics for the research
1737 """
1738 from flask import session as flask_session
1740 from ..database.session_context import get_user_db_session
1742 username = flask_session.get("username")
1743 if not username:
1744 return {
1745 "research_id": research_id,
1746 "research_details": {},
1747 "timeline": [],
1748 "summary": {
1749 "total_calls": 0,
1750 "total_tokens": 0,
1751 "total_prompt_tokens": 0,
1752 "total_completion_tokens": 0,
1753 "avg_response_time": 0,
1754 "success_rate": 0,
1755 },
1756 "phase_stats": {},
1757 }
1759 with get_user_db_session(username) as session:
1760 # Get all token usage for this research ordered by time including call stack
1761 timeline_data = session.execute(
1762 text(
1763 """
1764 SELECT
1765 timestamp,
1766 total_tokens,
1767 prompt_tokens,
1768 completion_tokens,
1769 response_time_ms,
1770 success_status,
1771 error_type,
1772 research_phase,
1773 search_iteration,
1774 search_engine_selected,
1775 model_name,
1776 calling_file,
1777 calling_function,
1778 call_stack
1779 FROM token_usage
1780 WHERE research_id = :research_id
1781 ORDER BY timestamp ASC
1782 """
1783 ),
1784 {"research_id": research_id},
1785 ).fetchall()
1787 # Format timeline data with cumulative tokens
1788 timeline = []
1789 cumulative_tokens = 0
1790 cumulative_prompt_tokens = 0
1791 cumulative_completion_tokens = 0
1793 for row in timeline_data:
1794 cumulative_tokens += row[1] or 0
1795 cumulative_prompt_tokens += row[2] or 0
1796 cumulative_completion_tokens += row[3] or 0
1798 timeline.append(
1799 {
1800 "timestamp": str(row[0]) if row[0] else None,
1801 "tokens": row[1] or 0,
1802 "prompt_tokens": row[2] or 0,
1803 "completion_tokens": row[3] or 0,
1804 "cumulative_tokens": cumulative_tokens,
1805 "cumulative_prompt_tokens": cumulative_prompt_tokens,
1806 "cumulative_completion_tokens": cumulative_completion_tokens,
1807 "response_time_ms": row[4],
1808 "success_status": row[5],
1809 "error_type": row[6],
1810 "research_phase": row[7],
1811 "search_iteration": row[8],
1812 "search_engine_selected": row[9],
1813 "model_name": row[10],
1814 "calling_file": row[11],
1815 "calling_function": row[12],
1816 "call_stack": row[13],
1817 }
1818 )
1820 # Get research basic info
1821 research_info = session.execute(
1822 text(
1823 """
1824 SELECT query, mode, status, created_at, completed_at
1825 FROM research_history
1826 WHERE id = :research_id
1827 """
1828 ),
1829 {"research_id": research_id},
1830 ).fetchone()
1832 research_details = {}
1833 if research_info:
1834 research_details = {
1835 "query": research_info[0],
1836 "mode": research_info[1],
1837 "status": research_info[2],
1838 "created_at": str(research_info[3])
1839 if research_info[3]
1840 else None,
1841 "completed_at": str(research_info[4])
1842 if research_info[4]
1843 else None,
1844 }
1846 # Calculate summary stats
1847 total_calls = len(timeline_data)
1848 total_tokens = cumulative_tokens
1849 avg_response_time = sum(row[4] or 0 for row in timeline_data) / max(
1850 total_calls, 1
1851 )
1852 success_rate = (
1853 sum(1 for row in timeline_data if row[5] == "success")
1854 / max(total_calls, 1)
1855 * 100
1856 )
1858 # Phase breakdown for this research
1859 phase_stats = {}
1860 for row in timeline_data:
1861 phase = row[7] or "unknown"
1862 if phase not in phase_stats:
1863 phase_stats[phase] = {
1864 "count": 0,
1865 "tokens": 0,
1866 "avg_response_time": 0,
1867 }
1868 phase_stats[phase]["count"] += 1
1869 phase_stats[phase]["tokens"] += row[1] or 0
1870 if row[4]: 1870 ↛ 1860line 1870 didn't jump to line 1860 because the condition on line 1870 was always true
1871 phase_stats[phase]["avg_response_time"] += row[4]
1873 # Calculate averages for phases
1874 for phase in phase_stats:
1875 if phase_stats[phase]["count"] > 0: 1875 ↛ 1874line 1875 didn't jump to line 1874 because the condition on line 1875 was always true
1876 phase_stats[phase]["avg_response_time"] = round(
1877 phase_stats[phase]["avg_response_time"]
1878 / phase_stats[phase]["count"]
1879 )
1881 return {
1882 "research_id": research_id,
1883 "research_details": research_details,
1884 "timeline": timeline,
1885 "summary": {
1886 "total_calls": total_calls,
1887 "total_tokens": total_tokens,
1888 "total_prompt_tokens": cumulative_prompt_tokens,
1889 "total_completion_tokens": cumulative_completion_tokens,
1890 "avg_response_time": round(avg_response_time),
1891 "success_rate": round(success_rate, 1),
1892 },
1893 "phase_stats": phase_stats,
1894 }