Coverage for src/local_deep_research/metrics/token_counter.py: 97%
551 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +0000
1"""Token counting functionality for LLM usage tracking."""
3import inspect
4import json
5import threading
6import time
7import uuid
8from datetime import UTC, datetime
9from pathlib import Path
10from typing import Any, Dict, List, Optional
12from langchain_core.callbacks import BaseCallbackHandler
13from langchain_core.outputs import LLMResult
14from loguru import logger
15from sqlalchemy import func, text
17from ..database.models import ModelUsage, TokenUsage
18from .query_utils import (
19 get_period_cutoff,
20 get_research_mode_condition,
21 get_time_filter_condition,
22)
23from ..utilities.request_context import (
24 get_current_session_id, # noqa: F401
25 get_current_username,
26)
28# Maximum number of frames captured per LLM call when building the
29# simplified call stack shown in the metrics UI. Five frames is enough to
30# surface the LDR call site and the immediate LangChain/LangGraph frames
31# above it without flooding the persisted column.
32_CALL_STACK_FRAME_LIMIT = 5
35class TokenCountingCallback(BaseCallbackHandler):
36 """Callback handler for counting tokens across different models."""
38 def __init__(
39 self,
40 research_id: Optional[str] = None,
41 research_context: Optional[Dict[str, Any]] = None,
42 ):
43 """Initialize the token counting callback.
45 Args:
46 research_id: The ID of the research to track tokens for
47 research_context: Additional research context for enhanced tracking
48 """
49 super().__init__()
50 self.research_id = research_id
51 self.research_context = research_context or {}
52 self.current_model = None
53 self.current_provider = None
54 self.preset_model = None # Model name set during callback creation
55 self.preset_provider = None # Provider set during callback creation
57 # Phase 1 Enhancement: Track timing and context
58 self.start_time = None
59 self.response_time_ms = None
60 self.success_status = "success"
61 self.error_type = None
63 # Per-run call metadata. The callback is shared across every LLM
64 # call in a research session (see llm_config.py wrap_llm), and a
65 # single session can drive overlapping chains (e.g. agent + judge
66 # sub-strategies). Storing the captured calling_file /
67 # calling_function / call_stack on instance attributes caused
68 # overlapping or unmatched runs to inherit each other's stack
69 # (or to keep a previous run's stack when the new run had no
70 # matching LDR frame). Keying by langchain's run_id and removing
71 # the entry on end/error bounds the metadata to its call.
72 self._call_meta: Dict[str, Dict[str, Any]] = {}
73 self._lock = threading.Lock()
74 # Mirror the most recently captured values for code paths (and
75 # legacy tests) that read these as plain attributes. They are not
76 # authoritative; the per-run dict is.
77 self.calling_file = None
78 self.calling_function = None
79 self.call_stack = None
81 # Context overflow tracking
82 self.context_limit = None
83 self.context_truncated = False
84 self.tokens_truncated = 0
85 self.truncation_ratio = 0.0
86 self.original_prompt_estimate = 0
88 # Raw Ollama response metrics
89 self.ollama_metrics = {}
91 # Whether we've already logged the "provider reports no usage"
92 # warning. The callback is shared across every LLM call in a
93 # research session (see llm_config.py wrap_llm), so without this a
94 # 50-call run against a non-reporting provider logs 50 identical
95 # warnings. Calls are still recorded every time — only the warning
96 # is deduplicated.
97 self._warned_no_usage = False
99 # Track token counts in memory
100 self.counts = {
101 "total_tokens": 0,
102 "total_prompt_tokens": 0,
103 "total_completion_tokens": 0,
104 "by_model": {},
105 }
107 def on_llm_start(
108 self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
109 ) -> None:
110 """Called when LLM starts running."""
111 # Phase 1 Enhancement: Start timing
112 self.start_time = time.time()
114 # Reset per-call truncation state. The callback instance is shared
115 # across every LLM call in a research session (see llm_config.py
116 # wrap_llm), so without this reset the post-loop estimation block's
117 # `if not self.context_truncated` guard would silently disable
118 # [estimated] / [estimated-total-context] detection on every call
119 # after the first one that truncates.
120 self.context_truncated = False
121 self.tokens_truncated = 0
122 self.truncation_ratio = 0.0
124 # Estimate original prompt size (rough estimate: ~4 chars per token)
125 if prompts:
126 total_chars = sum(len(prompt) for prompt in prompts)
127 self.original_prompt_estimate = total_chars // 4
128 logger.debug(
129 f"Estimated prompt tokens: {self.original_prompt_estimate} (from {total_chars} chars)"
130 )
132 # Get context limit from research context (will be set from settings)
133 self.context_limit = self.research_context.get("context_limit")
135 # Per-run call metadata, keyed by langchain's run_id so overlapping
136 # runs (or a no-match run following a match) don't see each other's
137 # stack. The run_id comes in via kwargs from langchain; fall back
138 # to a synthetic key when invoked outside the normal chain so
139 # tests / direct callers still work.
140 run_id = kwargs.get("run_id") or f"_synthetic-{uuid.uuid4()}"
142 # Phase 1 Enhancement: Capture call stack information
143 calling_file = None
144 calling_function = None
145 call_stack = None
146 try:
147 stack = inspect.stack()
149 # Skip the first few frames (this method, langchain internals)
150 # Look for the first frame that's in our project directory.
151 #
152 # Note: the project is installed as a wheel inside
153 # .venv/.../site-packages in the Docker image, so the project
154 # frames DO contain "site-packages" and "venv" as substrings.
155 # The "local_deep_research" package-name check is sufficient to
156 # identify project files — third-party packages with the same
157 # name do not exist.
158 for matched_index, frame_info in enumerate(stack[1:], start=1):
159 file_path = frame_info.filename
160 if "local_deep_research" in file_path:
161 # Extract relative path from local_deep_research
162 if "src/local_deep_research" in file_path:
163 relative_path = file_path.split(
164 "src/local_deep_research"
165 )[-1].lstrip("/")
166 elif "local_deep_research/src" in file_path:
167 relative_path = file_path.split(
168 "local_deep_research/src"
169 )[-1].lstrip("/")
170 elif "local_deep_research" in file_path: 170 ↛ 176line 170 didn't jump to line 176 because the condition on line 170 was always true
171 # Get everything after local_deep_research
172 relative_path = file_path.split("local_deep_research")[
173 -1
174 ].lstrip("/")
175 else:
176 relative_path = Path(file_path).name
178 calling_file = relative_path
179 calling_function = frame_info.function
181 # Capture a simplified call stack starting from the
182 # matched frame onward (limit to _CALL_STACK_FRAME_LIMIT
183 # frames). Previously this sliced stack[1:6]
184 # unconditionally, so when the matched LDR frame sat
185 # deeper than index 5 (common with langgraph),
186 # call_stack_frames only inspected langchain/langgraph
187 # frames and stayed empty.
188 call_stack_frames = []
189 for frame in stack[
190 matched_index : matched_index + _CALL_STACK_FRAME_LIMIT
191 ]:
192 if "local_deep_research" in frame.filename:
193 frame_name = f"{Path(frame.filename).name}:{frame.function}:{frame.lineno}"
194 call_stack_frames.append(frame_name)
196 call_stack = (
197 " -> ".join(call_stack_frames)
198 if call_stack_frames
199 else None
200 )
201 break
202 except Exception:
203 logger.warning("Error capturing call stack")
204 # Continue without call stack info if there's an error
206 with self._lock:
207 self._call_meta[run_id] = {
208 "calling_file": calling_file,
209 "calling_function": calling_function,
210 "call_stack": call_stack,
211 }
212 # Mirror to the legacy attributes for callers that read them
213 # directly. The per-run dict is authoritative; these are a
214 # convenience that reflects the most recent on_llm_start.
215 self.calling_file = calling_file
216 self.calling_function = calling_function
217 self.call_stack = call_stack
219 # First, use preset values if available
220 if self.preset_model:
221 self.current_model = self.preset_model
222 else:
223 # Try multiple locations for model name
224 model_name = None
226 # First check invocation_params
227 invocation_params = kwargs.get("invocation_params", {})
228 model_name = invocation_params.get(
229 "model"
230 ) or invocation_params.get("model_name")
232 # Check kwargs directly
233 if not model_name:
234 model_name = kwargs.get("model") or kwargs.get("model_name")
236 # Check serialized data
237 if not model_name and "kwargs" in serialized:
238 model_name = serialized["kwargs"].get("model") or serialized[
239 "kwargs"
240 ].get("model_name")
242 # Check for name in serialized data
243 if not model_name and "name" in serialized:
244 model_name = serialized["name"]
246 # If still not found and we have Ollama, try to extract from the instance
247 if (
248 not model_name
249 and "_type" in serialized
250 and "ChatOllama" in serialized["_type"]
251 ):
252 # For Ollama, the model name might be in the serialized kwargs
253 if "kwargs" in serialized and "model" in serialized["kwargs"]: 253 ↛ 254line 253 didn't jump to line 254 because the condition on line 253 was never true
254 model_name = serialized["kwargs"]["model"]
255 else:
256 # Default to the type if we can't find the actual model
257 model_name = "ollama"
259 # Final fallback
260 if not model_name:
261 if "_type" in serialized:
262 model_name = serialized["_type"]
263 else:
264 model_name = "unknown"
266 self.current_model = model_name
268 # Use preset provider if available
269 if self.preset_provider:
270 self.current_provider = self.preset_provider
271 else:
272 # Extract provider from serialized type or kwargs
273 if "_type" in serialized:
274 type_str = serialized["_type"]
275 if "ChatOllama" in type_str:
276 self.current_provider = "ollama"
277 elif "ChatOpenAI" in type_str:
278 self.current_provider = "openai"
279 elif "ChatAnthropic" in type_str:
280 self.current_provider = "anthropic"
281 else:
282 self.current_provider = kwargs.get("provider", "unknown")
283 else:
284 self.current_provider = kwargs.get("provider", "unknown")
286 # Initialize model tracking if needed
287 if self.current_model not in self.counts["by_model"]:
288 self.counts["by_model"][self.current_model] = {
289 "prompt_tokens": 0,
290 "completion_tokens": 0,
291 "total_tokens": 0,
292 "calls": 0,
293 "provider": self.current_provider,
294 }
296 # Increment call count
297 self.counts["by_model"][self.current_model]["calls"] += 1
299 def _check_context_overflow(
300 self,
301 prompt_eval_count: int,
302 completion_tokens: int = 0,
303 source: str = "",
304 ) -> None:
305 """Check for context overflow based on prompt and total token usage.
307 Args:
308 prompt_eval_count: Number of tokens the model actually processed.
309 completion_tokens: Number of tokens generated (for total-context check).
310 source: Which branch provided the data (for logging).
311 """
312 logger.debug(
313 f"Context overflow check [{source}]: "
314 f"prompt_eval_count={prompt_eval_count}, "
315 f"completion_tokens={completion_tokens}, "
316 f"context_limit={self.context_limit}"
317 )
319 if not self.context_limit or prompt_eval_count <= 0:
320 return
322 # Input-only overflow: prompt at >= 80% of context limit. Matches the
323 # chart-warning threshold and PR #3840's deliberate choice (PR #3792
324 # lowered from 95% → 80%). The total-context branch below uses 95%
325 # because it's a stricter condition (input+output combined).
326 if prompt_eval_count >= self.context_limit * 0.80:
327 self.context_truncated = True
329 if self.original_prompt_estimate > prompt_eval_count:
330 self.tokens_truncated = max(
331 0,
332 self.original_prompt_estimate - prompt_eval_count,
333 )
334 if ( 334 ↛ 348line 334 didn't jump to line 348 because the condition on line 334 was always true
335 self.tokens_truncated > 0
336 and self.original_prompt_estimate > 0
337 ):
338 self.truncation_ratio = (
339 self.tokens_truncated / self.original_prompt_estimate
340 )
341 elif prompt_eval_count > self.context_limit: 341 ↛ 342line 341 didn't jump to line 342 because the condition on line 341 was never true
342 self.tokens_truncated = prompt_eval_count - self.context_limit
343 if self.tokens_truncated > 0 and prompt_eval_count > 0:
344 self.truncation_ratio = (
345 self.tokens_truncated / prompt_eval_count
346 )
348 logger.warning(
349 f"Context overflow detected [provider-confirmed] "
350 f"research_id={self.research_id} "
351 f"model={self.current_model} "
352 f"provider={self.current_provider} "
353 f"source={source} "
354 f"prompt_tokens={prompt_eval_count} "
355 f"context_limit={self.context_limit} "
356 f"tokens_truncated={self.tokens_truncated} "
357 f"truncation_ratio={self.truncation_ratio:.1%}"
358 )
360 # Total-context overflow: input + output exceeds 95% of context limit
361 elif (
362 completion_tokens > 0
363 and prompt_eval_count + completion_tokens
364 >= self.context_limit * 0.95
365 ):
366 total = prompt_eval_count + completion_tokens
367 self.context_truncated = True
368 self.tokens_truncated = max(0, total - self.context_limit)
369 self.truncation_ratio = (
370 self.tokens_truncated / total if total > 0 else 0
371 )
372 logger.warning(
373 f"Context overflow detected [total-context] "
374 f"research_id={self.research_id} "
375 f"model={self.current_model} "
376 f"provider={self.current_provider} "
377 f"source={source} "
378 f"prompt_tokens={prompt_eval_count} "
379 f"completion_tokens={completion_tokens} "
380 f"total_tokens={total} "
381 f"context_limit={self.context_limit} "
382 f"tokens_truncated={self.tokens_truncated} "
383 f"truncation_ratio={self.truncation_ratio:.1%}"
384 )
385 else:
386 logger.debug(
387 f"Context OK [{source}]: "
388 f"{prompt_eval_count}/{self.context_limit} "
389 f"({prompt_eval_count / self.context_limit:.1%})"
390 )
392 def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
393 """Called when LLM ends running."""
394 # Phase 1 Enhancement: Calculate response time
395 if self.start_time:
396 self.response_time_ms = int((time.time() - self.start_time) * 1000)
398 # Pop the per-run call metadata captured at on_llm_start so a
399 # later unmatched run cannot persist a previous run's stack.
400 # Missing-key case (e.g. direct test invocation that skipped
401 # on_llm_start) yields a None-filled dict, which the API/UI
402 # already treat as "no stack".
403 run_id = kwargs.get("run_id")
404 call_meta = (
405 self._pop_call_meta(run_id) if run_id else self._legacy_meta()
406 )
408 # Extract token usage from response
409 token_usage = None
411 # Check multiple locations for token usage
412 if hasattr(response, "llm_output") and response.llm_output:
413 token_usage = response.llm_output.get(
414 "token_usage"
415 ) or response.llm_output.get("usage", {})
417 # Check for usage metadata in generations (Ollama specific)
418 if not token_usage and hasattr(response, "generations"):
419 for generation_list in response.generations:
420 for generation in generation_list:
421 if hasattr(generation, "message") and hasattr(
422 generation.message, "usage_metadata"
423 ):
424 usage_meta = generation.message.usage_metadata
425 if usage_meta: # Check if usage_metadata is not None
426 token_usage = {
427 "prompt_tokens": usage_meta.get(
428 "input_tokens", 0
429 ),
430 "completion_tokens": usage_meta.get(
431 "output_tokens", 0
432 ),
433 # Fall back to the sum: a provider that omits
434 # the total otherwise contributes 0 here.
435 # `or` rather than `is None` is deliberate: a
436 # reported total of 0 alongside nonzero input
437 # and output is bad data, and the sum recovers
438 # it. The `is None` checks further down
439 # (_save_to_db, thread_metrics) persist a real
440 # 0, because the total has already been
441 # resolved by the time they run.
442 "total_tokens": usage_meta.get("total_tokens")
443 or (
444 usage_meta.get("input_tokens", 0)
445 + usage_meta.get("output_tokens", 0)
446 ),
447 }
448 # Check context overflow before breaking
449 # (input_tokens == prompt_eval_count for Ollama)
450 self._check_context_overflow(
451 usage_meta.get("input_tokens", 0),
452 completion_tokens=usage_meta.get(
453 "output_tokens", 0
454 ),
455 source="usage_metadata",
456 )
457 break
458 # Also check response_metadata
459 if hasattr(generation, "message") and hasattr(
460 generation.message, "response_metadata"
461 ):
462 resp_meta = generation.message.response_metadata
463 if resp_meta.get("prompt_eval_count") or resp_meta.get(
464 "eval_count"
465 ):
466 # Capture raw Ollama metrics
467 self.ollama_metrics = {
468 "prompt_eval_count": resp_meta.get(
469 "prompt_eval_count"
470 ),
471 "eval_count": resp_meta.get("eval_count"),
472 "total_duration": resp_meta.get(
473 "total_duration"
474 ),
475 "load_duration": resp_meta.get("load_duration"),
476 "prompt_eval_duration": resp_meta.get(
477 "prompt_eval_duration"
478 ),
479 "eval_duration": resp_meta.get("eval_duration"),
480 }
482 # Check for context overflow (input only)
483 prompt_eval_count = resp_meta.get(
484 "prompt_eval_count", 0
485 )
486 self._check_context_overflow(
487 prompt_eval_count,
488 completion_tokens=resp_meta.get(
489 "eval_count", 0
490 ),
491 source="response_metadata",
492 )
494 token_usage = {
495 "prompt_tokens": resp_meta.get(
496 "prompt_eval_count", 0
497 ),
498 "completion_tokens": resp_meta.get(
499 "eval_count", 0
500 ),
501 "total_tokens": resp_meta.get(
502 "prompt_eval_count", 0
503 )
504 + resp_meta.get("eval_count", 0),
505 }
506 break
507 if token_usage:
508 break
510 # Estimation-based overflow detection for providers that don't echo
511 # prompt_eval_count (OpenAI, Anthropic, OpenRouter, etc.). The Ollama
512 # path above sets context_truncated=True if it detected provider-
513 # confirmed truncation; we only fire here if it didn't. For hosted
514 # providers the API typically rejects oversize prompts rather than
515 # silently truncating, so this signal flags "would-overflow per our
516 # estimate" — not the same as actual truncation, hence the [estimated]
517 # tag in the log message.
518 if not self.context_truncated and self.context_limit:
519 # Input-only overflow: prompt estimate exceeds context limit
520 if self.original_prompt_estimate > self.context_limit:
521 self.context_truncated = True
522 self.tokens_truncated = max(
523 0, self.original_prompt_estimate - self.context_limit
524 )
525 self.truncation_ratio = (
526 self.tokens_truncated / self.original_prompt_estimate
527 if self.original_prompt_estimate > 0
528 else 0
529 )
530 logger.warning(
531 "Context overflow detected [estimated] "
532 f"research_id={self.research_id} "
533 f"model={self.current_model} "
534 f"provider={self.current_provider} "
535 f"estimated_prompt_tokens={self.original_prompt_estimate} "
536 f"context_limit={self.context_limit} "
537 f"tokens_truncated={self.tokens_truncated} "
538 f"truncation_ratio={self.truncation_ratio:.1%}"
539 )
540 # Total-context overflow: input + output exceeds context limit.
541 # The "[estimated-total-context]" tag below refers to the
542 # *detection method* (post-loop fallback path that runs when
543 # _check_context_overflow couldn't fire), not the data source —
544 # the prompt_tokens / completion_tokens here come from the
545 # provider via response.llm_output.token_usage and are actual
546 # counts, not character-based estimates.
547 elif token_usage and isinstance(token_usage, dict):
548 prompt_tokens = token_usage.get("prompt_tokens", 0)
549 completion_tokens = token_usage.get("completion_tokens", 0)
550 total = prompt_tokens + completion_tokens
551 if total >= self.context_limit * 0.95:
552 self.context_truncated = True
553 self.tokens_truncated = max(0, total - self.context_limit)
554 self.truncation_ratio = (
555 self.tokens_truncated / total if total > 0 else 0
556 )
557 logger.warning(
558 "Context overflow detected [estimated-total-context] "
559 f"research_id={self.research_id} "
560 f"model={self.current_model} "
561 f"provider={self.current_provider} "
562 f"prompt_tokens={prompt_tokens} "
563 f"completion_tokens={completion_tokens} "
564 f"total_tokens={total} "
565 f"context_limit={self.context_limit} "
566 f"tokens_truncated={self.tokens_truncated} "
567 f"truncation_ratio={self.truncation_ratio:.1%}"
568 )
570 if token_usage and isinstance(token_usage, dict):
571 prompt_tokens = token_usage.get("prompt_tokens", 0)
572 completion_tokens = token_usage.get("completion_tokens", 0)
573 total_tokens = token_usage.get(
574 "total_tokens", prompt_tokens + completion_tokens
575 )
577 # Update in-memory counts
578 self.counts["total_prompt_tokens"] += prompt_tokens
579 self.counts["total_completion_tokens"] += completion_tokens
580 self.counts["total_tokens"] += total_tokens
582 if self.current_model:
583 self.counts["by_model"][self.current_model][
584 "prompt_tokens"
585 ] += prompt_tokens
586 self.counts["by_model"][self.current_model][
587 "completion_tokens"
588 ] += completion_tokens
589 self.counts["by_model"][self.current_model]["total_tokens"] += (
590 total_tokens
591 )
593 # Save to database if we have a research_id
594 if self.research_id:
595 self._save_to_db(
596 prompt_tokens,
597 completion_tokens,
598 call_meta,
599 total_tokens=total_tokens,
600 )
601 elif self.research_id:
602 # The provider returned no usage data at all (e.g. OpenAI-
603 # compatible servers that omit `usage` on streamed responses,
604 # proxies that strip it, Ollama omitting prompt_eval_count for
605 # fully-cached prompts). Record the call anyway with zero token
606 # counts — model, provider, phase, response time and context-
607 # overflow estimates are still real data. Skipping the row
608 # entirely is what made every metrics page render as if LDR had
609 # never been used (#4457).
610 # Warn once per research session, not once per call.
611 if not self._warned_no_usage:
612 self._warned_no_usage = True
613 hint = ""
614 if self.current_provider == "openai_endpoint": 614 ↛ 615line 614 didn't jump to line 615 because the condition on line 614 was never true
615 hint = (
616 " If your server supports stream_options.include_usage "
617 "(LM Studio 0.3.18+, llama.cpp, vLLM, OpenRouter), "
618 "enable the 'llm.openai_endpoint.stream_usage' setting "
619 "to get real token counts on streamed calls."
620 )
621 logger.warning(
622 f"LLM provider returned no token usage data - recording "
623 f"calls with zero counts. model={self.current_model} "
624 f"provider={self.current_provider} "
625 f"research_id={self.research_id}.{hint} "
626 f"(further occurrences this session are suppressed)"
627 )
628 self._save_to_db(0, 0, call_meta)
630 def on_llm_error(self, error, **kwargs: Any) -> None:
631 """Called when LLM encounters an error."""
632 # Phase 1 Enhancement: Track errors
633 if self.start_time:
634 self.response_time_ms = int((time.time() - self.start_time) * 1000)
636 self.success_status = "error"
637 self.error_type = str(type(error).__name__)
639 # Pop the per-run call metadata captured at on_llm_start (if any)
640 # so a later unmatched run does not inherit a failed run's stack.
641 run_id = kwargs.get("run_id")
642 call_meta = (
643 self._pop_call_meta(run_id) if run_id else self._legacy_meta()
644 )
646 # Still save to database to track failed calls
647 if self.research_id:
648 self._save_to_db(0, 0, call_meta)
650 def _pop_call_meta(self, run_id: str) -> Dict[str, Any]:
651 """Remove and return the per-run call metadata for ``run_id``.
653 Falls back to ``_legacy_meta()`` (the most recently captured
654 values) if the run was never registered — e.g. ``on_llm_error``
655 firing without a corresponding ``on_llm_start`` (provider raised
656 during argument validation) or a test that drives the callbacks
657 out of order. After popping, refresh the legacy mirror
658 attributes from the next-newest remaining entry so direct
659 ``cb.calling_file`` reads in overlapping-run scenarios do not
660 show a deleted run's stack.
661 """
662 with self._lock:
663 meta = self._call_meta.pop(
664 run_id,
665 {
666 "calling_file": None,
667 "calling_function": None,
668 "call_stack": None,
669 },
670 )
671 if self._call_meta:
672 latest_id = next(reversed(self._call_meta))
673 latest = self._call_meta[latest_id]
674 else:
675 latest = meta
676 self.calling_file = latest.get("calling_file")
677 self.calling_function = latest.get("calling_function")
678 self.call_stack = latest.get("call_stack")
679 return meta
681 def _legacy_meta(self) -> Dict[str, Any]:
682 """Return the legacy mirror attributes as a meta dict.
684 Used by callers that did not register a run (no run_id kwarg),
685 e.g. direct test invocations of ``on_llm_end`` /
686 ``on_llm_error``. The values reflect the most recent
687 ``on_llm_start`` (or whatever the test set explicitly).
688 """
689 with self._lock:
690 return {
691 "calling_file": self.calling_file,
692 "calling_function": self.calling_function,
693 "call_stack": self.call_stack,
694 }
696 def _get_context_overflow_fields(self) -> Dict[str, Any]:
697 """Get context overflow detection fields for database saving."""
698 return {
699 "context_limit": self.context_limit,
700 "context_truncated": self.context_truncated, # Now Boolean
701 "tokens_truncated": self.tokens_truncated
702 if self.context_truncated
703 else None,
704 "truncation_ratio": self.truncation_ratio
705 if self.context_truncated
706 else None,
707 # Raw Ollama metrics
708 "ollama_prompt_eval_count": self.ollama_metrics.get(
709 "prompt_eval_count"
710 ),
711 "ollama_eval_count": self.ollama_metrics.get("eval_count"),
712 "ollama_total_duration": self.ollama_metrics.get("total_duration"),
713 "ollama_load_duration": self.ollama_metrics.get("load_duration"),
714 "ollama_prompt_eval_duration": self.ollama_metrics.get(
715 "prompt_eval_duration"
716 ),
717 "ollama_eval_duration": self.ollama_metrics.get("eval_duration"),
718 }
720 def _save_to_db(
721 self,
722 prompt_tokens: int,
723 completion_tokens: int,
724 call_meta: Optional[Dict[str, Any]] = None,
725 total_tokens: Optional[int] = None,
726 ):
727 """Save token usage to the database.
729 Args:
730 prompt_tokens: Prompt-side token count for this call.
731 completion_tokens: Completion-side token count for this call.
732 total_tokens: Total reported by the provider, as resolved by
733 the caller. Defaults to the sum of the two counts above
734 when the caller has no provider-reported total.
735 call_meta: Per-run call metadata captured by on_llm_start
736 (calling_file / calling_function / call_stack). When the
737 callback is invoked outside on_llm_end / on_llm_error
738 (e.g. legacy tests that set the attributes directly),
739 ``call_meta`` may be ``None`` and the instance attributes
740 are used as a fallback.
741 """
742 if total_tokens is None:
743 total_tokens = prompt_tokens + completion_tokens
744 if call_meta is None:
745 with self._lock:
746 call_meta = {
747 "calling_file": self.calling_file,
748 "calling_function": self.calling_function,
749 "call_stack": self.call_stack,
750 }
751 # Check if we're in a thread - if so, queue the save for later
752 import threading
754 if threading.current_thread().name != "MainThread":
755 # Use thread-safe metrics database for background threads
756 username = (
757 self.research_context.get("username")
758 if self.research_context
759 else None
760 )
762 if not username:
763 logger.warning(
764 f"Cannot save token metrics - no username in research context. "
765 f"Token usage: prompt={prompt_tokens}, completion={completion_tokens}, "
766 f"Research context: {self.research_context}"
767 )
768 return
770 # Import the thread-safe metrics database
772 # Prepare token data
773 token_data = {
774 "model_name": self.current_model,
775 "provider": self.current_provider,
776 "prompt_tokens": prompt_tokens,
777 "completion_tokens": completion_tokens,
778 "total_tokens": total_tokens,
779 "research_query": self.research_context.get("research_query"),
780 "research_mode": self.research_context.get("research_mode"),
781 "research_phase": self.research_context.get("research_phase"),
782 "search_iteration": self.research_context.get(
783 "search_iteration"
784 ),
785 "response_time_ms": self.response_time_ms,
786 "success_status": self.success_status,
787 "error_type": self.error_type,
788 "search_engines_planned": self.research_context.get(
789 "search_engines_planned"
790 ),
791 "search_engine_selected": self.research_context.get(
792 "search_engine_selected"
793 ),
794 "calling_file": call_meta.get("calling_file"),
795 "calling_function": call_meta.get("calling_function"),
796 "call_stack": call_meta.get("call_stack"),
797 # Add context overflow fields using helper method
798 **self._get_context_overflow_fields(),
799 }
801 # Convert list to JSON string if needed
802 if isinstance(token_data.get("search_engines_planned"), list):
803 token_data["search_engines_planned"] = json.dumps(
804 token_data["search_engines_planned"]
805 )
807 # Get password from research context
808 password = self.research_context.get("user_password")
809 if not password:
810 logger.warning(
811 f"Cannot save token metrics - no password in research context. "
812 f"Username: {username}, Token usage: prompt={prompt_tokens}, completion={completion_tokens}"
813 )
814 return
816 # Write metrics directly using thread-safe database
817 try:
818 from ..database.thread_metrics import metrics_writer
820 # Set password for this thread
821 metrics_writer.set_user_password(username, password)
823 # Write metrics to encrypted database
824 metrics_writer.write_token_metrics(
825 username, self.research_id, token_data
826 )
827 except Exception as e:
828 # Deferred import to avoid potential circular imports during module initialization
829 from ..security.log_sanitizer import scrub_error
831 safe_msg = scrub_error(e, password)
832 logger.warning(
833 f"Failed to write metrics from thread: {safe_msg}"
834 )
835 return
837 # In MainThread, save directly
838 try:
839 from ..database.session_context import get_user_db_session
841 username = get_current_username()
842 if not username:
843 logger.debug("No user session, skipping token metrics save")
844 return
846 with get_user_db_session(username) as session:
847 # Phase 1 Enhancement: Prepare additional context
848 research_query = self.research_context.get("research_query")
849 research_mode = self.research_context.get("research_mode")
850 research_phase = self.research_context.get("research_phase")
851 search_iteration = self.research_context.get("search_iteration")
852 search_engines_planned = self.research_context.get(
853 "search_engines_planned"
854 )
855 search_engine_selected = self.research_context.get(
856 "search_engine_selected"
857 )
859 # Debug logging for search engine context
860 if search_engines_planned or search_engine_selected:
861 logger.info(
862 f"Token tracking - Search context: planned={search_engines_planned}, selected={search_engine_selected}, phase={research_phase}"
863 )
864 else:
865 logger.debug(
866 f"Token tracking - No search engine context yet, phase={research_phase}"
867 )
869 # Convert list to JSON string if needed
870 if isinstance(search_engines_planned, list):
871 search_engines_planned = json.dumps(search_engines_planned)
873 # Log context overflow detection values before saving
874 logger.debug(
875 f"Saving TokenUsage - context_limit: {self.context_limit}, "
876 f"context_truncated: {self.context_truncated}, "
877 f"tokens_truncated: {self.tokens_truncated}, "
878 f"ollama_prompt_eval_count: {self.ollama_metrics.get('prompt_eval_count')}, "
879 f"prompt_tokens: {prompt_tokens}, "
880 f"completion_tokens: {completion_tokens}"
881 )
883 # Add token usage record with enhanced fields
884 token_usage = TokenUsage(
885 research_id=self.research_id,
886 model_name=self.current_model,
887 model_provider=self.current_provider, # Added provider
888 # for accurate cost tracking
889 prompt_tokens=prompt_tokens,
890 completion_tokens=completion_tokens,
891 total_tokens=total_tokens,
892 # Phase 1 Enhancement: Research context
893 research_query=research_query,
894 research_mode=research_mode,
895 research_phase=research_phase,
896 search_iteration=search_iteration,
897 # Phase 1 Enhancement: Performance metrics
898 response_time_ms=self.response_time_ms,
899 success_status=self.success_status,
900 error_type=self.error_type,
901 # Phase 1 Enhancement: Search engine context
902 search_engines_planned=search_engines_planned,
903 search_engine_selected=search_engine_selected,
904 # Phase 1 Enhancement: Call stack tracking
905 calling_file=call_meta.get("calling_file"),
906 calling_function=call_meta.get("calling_function"),
907 call_stack=call_meta.get("call_stack"),
908 # Add context overflow fields using helper method
909 **self._get_context_overflow_fields(),
910 )
911 session.add(token_usage)
913 # Update or create model usage statistics
914 model_usage = (
915 session.query(ModelUsage)
916 .filter_by(
917 model_name=self.current_model,
918 )
919 .first()
920 )
922 if model_usage:
923 model_usage.total_tokens += total_tokens
924 model_usage.total_calls += 1
925 else:
926 model_usage = ModelUsage(
927 model_name=self.current_model,
928 model_provider=self.current_provider,
929 total_tokens=total_tokens,
930 total_calls=1,
931 )
932 session.add(model_usage)
934 # Commit the transaction
935 session.commit()
937 except Exception as e:
938 # Deferred import to avoid potential circular imports during module initialization
939 from ..security.log_sanitizer import scrub_error
941 safe_msg = scrub_error(e)
942 logger.warning(f"Error saving token usage to database: {safe_msg}")
944 def get_counts(self) -> Dict[str, Any]:
945 """Get the current token counts."""
946 return self.counts
949class TokenCounter:
950 """Manager class for token counting across the application."""
952 def __init__(self):
953 """Initialize the token counter."""
955 def create_callback(
956 self,
957 research_id: Optional[str] = None,
958 research_context: Optional[Dict[str, Any]] = None,
959 ) -> TokenCountingCallback:
960 """Create a new token counting callback.
962 Args:
963 research_id: The ID of the research to track tokens for
964 research_context: Additional research context for enhanced tracking
966 Returns:
967 A new TokenCountingCallback instance
968 """
969 return TokenCountingCallback(
970 research_id=research_id, research_context=research_context
971 )
973 def get_research_metrics(
974 self, research_id: str, username: Optional[str] = None
975 ) -> Dict[str, Any]:
976 """Get token metrics for a specific research.
978 Args:
979 research_id: The ID of the research
980 username: Authenticated user whose encrypted DB to open. If
981 omitted, falls back to the request's contextvar — but
982 passing it explicitly from the route handler is
983 preferred so the lookup is obviously user-scoped even
984 if the contextvar path ever regresses.
986 Returns:
987 Dictionary containing token usage metrics
988 """
990 from ..database.session_context import get_user_db_session
992 if username is None:
993 username = get_current_username()
994 if not username:
995 return {
996 "research_id": research_id,
997 "total_tokens": 0,
998 "total_calls": 0,
999 "model_usage": [],
1000 }
1002 with get_user_db_session(username) as session:
1003 # Get token usage for this research from TokenUsage table
1004 from sqlalchemy import func
1006 token_usages = (
1007 session.query(
1008 TokenUsage.model_name,
1009 TokenUsage.model_provider,
1010 func.sum(TokenUsage.prompt_tokens).label("prompt_tokens"),
1011 func.sum(TokenUsage.completion_tokens).label(
1012 "completion_tokens"
1013 ),
1014 func.sum(TokenUsage.total_tokens).label("total_tokens"),
1015 func.count().label("calls"),
1016 )
1017 .filter_by(research_id=research_id)
1018 .group_by(TokenUsage.model_name, TokenUsage.model_provider)
1019 .order_by(func.sum(TokenUsage.total_tokens).desc())
1020 .all()
1021 )
1023 model_usage = []
1024 total_tokens = 0
1025 total_calls = 0
1027 for usage in token_usages:
1028 model_usage.append(
1029 {
1030 "model": usage.model_name,
1031 "provider": usage.model_provider,
1032 "tokens": usage.total_tokens or 0,
1033 "calls": usage.calls or 0,
1034 "prompt_tokens": usage.prompt_tokens or 0,
1035 "completion_tokens": usage.completion_tokens or 0,
1036 }
1037 )
1038 total_tokens += usage.total_tokens or 0
1039 total_calls += usage.calls or 0
1041 return {
1042 "research_id": research_id,
1043 "total_tokens": total_tokens,
1044 "total_calls": total_calls,
1045 "model_usage": model_usage,
1046 }
1048 def get_overall_metrics(
1049 self,
1050 period: str = "30d",
1051 research_mode: str = "all",
1052 username: Optional[str] = None,
1053 ) -> Dict[str, Any]:
1054 """Get overall token metrics across all researches.
1056 Args:
1057 period: Time period to filter by ('7d', '30d', '3m', '1y', 'all')
1058 research_mode: Research mode to filter by ('quick', 'detailed', 'all')
1059 username: Username to fetch metrics for. If omitted, falls back
1060 to the Flask session (Flask-only path).
1062 Returns:
1063 Dictionary containing overall metrics
1064 """
1065 return self._get_metrics_from_encrypted_db(
1066 period, research_mode, username=username
1067 )
1069 def _get_metrics_from_encrypted_db(
1070 self,
1071 period: str,
1072 research_mode: str,
1073 username: Optional[str] = None,
1074 ) -> Dict[str, Any]:
1075 """Get metrics from user's encrypted database."""
1076 from ..database.session_context import get_user_db_session
1078 if not username:
1079 username = get_current_username()
1081 if not username:
1082 return self._get_empty_metrics()
1084 try:
1085 with get_user_db_session(username) as session:
1086 # Build base query with filters
1087 query = session.query(TokenUsage)
1089 # Apply time filter
1090 time_condition = get_time_filter_condition(
1091 period, TokenUsage.timestamp
1092 )
1093 if time_condition is not None:
1094 query = query.filter(time_condition)
1096 # Apply research mode filter
1097 mode_condition = get_research_mode_condition(
1098 research_mode, TokenUsage.research_mode
1099 )
1100 if mode_condition is not None:
1101 query = query.filter(mode_condition)
1103 # Total tokens from TokenUsage
1104 total_tokens = (
1105 query.with_entities(
1106 func.sum(TokenUsage.total_tokens)
1107 ).scalar()
1108 or 0
1109 )
1111 # Import ResearchHistory model
1112 from ..database.models.research import ResearchHistory
1114 # Count researches from ResearchHistory table
1115 research_query = session.query(func.count(ResearchHistory.id))
1117 # Debug: Check if any research history records exist at all
1118 all_research_count = (
1119 session.query(func.count(ResearchHistory.id)).scalar() or 0
1120 )
1121 logger.debug(
1122 f"Total ResearchHistory records in database: {all_research_count}"
1123 )
1125 # Debug: List first few research IDs and their timestamps
1126 sample_researches = (
1127 session.query(
1128 ResearchHistory.id,
1129 ResearchHistory.created_at,
1130 ResearchHistory.mode,
1131 )
1132 .limit(5)
1133 .all()
1134 )
1135 if sample_researches:
1136 logger.debug("Sample ResearchHistory records:")
1137 for r_id, r_created, r_mode in sample_researches:
1138 logger.debug(
1139 f" - ID: {r_id}, Created: {r_created}, Mode: {r_mode}"
1140 )
1141 else:
1142 logger.debug("No ResearchHistory records found in database")
1144 # Time filter for the ResearchHistory count. created_at is
1145 # an ISO-8601 TEXT column, so compare isoformat strings.
1146 # Must use the API's period vocabulary ('7d'/'30d'/'3m'/
1147 # '1y'/'all') — the previous 'today'/'week'/'month' branches
1148 # never matched it, so the research count silently ignored
1149 # the selected period and always showed all-time.
1150 start_time = get_period_cutoff(period)
1151 if start_time is not None:
1152 research_query = research_query.filter(
1153 ResearchHistory.created_at >= start_time.isoformat()
1154 )
1156 # Apply mode filter if specified
1157 mode_filter = research_mode if research_mode != "all" else None
1158 if mode_filter:
1159 logger.debug(f"Applying mode filter: {mode_filter}")
1160 research_query = research_query.filter(
1161 ResearchHistory.mode == mode_filter
1162 )
1164 total_researches = research_query.scalar() or 0
1165 logger.debug(
1166 f"Final filtered research count: {total_researches}"
1167 )
1169 # Also check distinct research_ids in TokenUsage for comparison
1170 token_research_count = (
1171 session.query(
1172 func.count(func.distinct(TokenUsage.research_id))
1173 ).scalar()
1174 or 0
1175 )
1176 logger.debug(
1177 f"Distinct research_ids in TokenUsage: {token_research_count}"
1178 )
1180 # Model statistics using ORM aggregation
1181 model_stats_query = session.query(
1182 TokenUsage.model_name,
1183 func.sum(TokenUsage.total_tokens).label("tokens"),
1184 func.count().label("calls"),
1185 func.sum(TokenUsage.prompt_tokens).label("prompt_tokens"),
1186 func.sum(TokenUsage.completion_tokens).label(
1187 "completion_tokens"
1188 ),
1189 ).filter(TokenUsage.model_name.isnot(None))
1191 # Apply same filters to model stats
1192 if time_condition is not None:
1193 model_stats_query = model_stats_query.filter(time_condition)
1194 if mode_condition is not None:
1195 model_stats_query = model_stats_query.filter(mode_condition)
1197 model_stats = (
1198 model_stats_query.group_by(TokenUsage.model_name)
1199 .order_by(func.sum(TokenUsage.total_tokens).desc())
1200 .all()
1201 )
1203 # Batch load provider info from ModelUsage table (fix N+1)
1204 model_names = [stat.model_name for stat in model_stats]
1205 provider_map = {}
1206 if model_names:
1207 provider_results = (
1208 session.query(
1209 ModelUsage.model_name, ModelUsage.model_provider
1210 )
1211 .filter(ModelUsage.model_name.in_(model_names))
1212 .order_by(ModelUsage.id)
1213 .all()
1214 )
1215 for model_name, model_provider in provider_results:
1216 provider_map.setdefault(model_name, model_provider)
1218 by_model = []
1219 for stat in model_stats:
1220 provider = provider_map.get(stat.model_name, "unknown")
1222 by_model.append(
1223 {
1224 "model": stat.model_name,
1225 "provider": provider,
1226 "tokens": stat.tokens,
1227 "calls": stat.calls,
1228 "prompt_tokens": stat.prompt_tokens,
1229 "completion_tokens": stat.completion_tokens,
1230 }
1231 )
1233 # Get recent researches with token usage
1234 # Note: This requires research_history table - for now we'll use available data
1235 recent_research_query = session.query(
1236 TokenUsage.research_id,
1237 func.sum(TokenUsage.total_tokens).label("token_count"),
1238 func.max(TokenUsage.timestamp).label("latest_timestamp"),
1239 ).filter(TokenUsage.research_id.isnot(None))
1241 if time_condition is not None:
1242 recent_research_query = recent_research_query.filter(
1243 time_condition
1244 )
1245 if mode_condition is not None:
1246 recent_research_query = recent_research_query.filter(
1247 mode_condition
1248 )
1250 recent_research_data = (
1251 recent_research_query.group_by(TokenUsage.research_id)
1252 .order_by(func.max(TokenUsage.timestamp).desc())
1253 .limit(10)
1254 .all()
1255 )
1257 # Batch load research queries for recent researches (fix N+1)
1258 recent_research_ids = [
1259 r.research_id for r in recent_research_data
1260 ]
1261 research_query_map = {}
1262 if recent_research_ids:
1263 # Get first non-null research_query for each research_id
1264 query_results = (
1265 session.query(
1266 TokenUsage.research_id, TokenUsage.research_query
1267 )
1268 .filter(
1269 TokenUsage.research_id.in_(recent_research_ids),
1270 TokenUsage.research_query.isnot(None),
1271 )
1272 .order_by(TokenUsage.id)
1273 .all()
1274 )
1275 for research_id, research_query in query_results:
1276 if research_id not in research_query_map: 1276 ↛ 1275line 1276 didn't jump to line 1275 because the condition on line 1276 was always true
1277 research_query_map[research_id] = research_query
1279 recent_researches = []
1280 for research_data in recent_research_data:
1281 query_text = research_query_map.get(
1282 research_data.research_id,
1283 f"Research {research_data.research_id}",
1284 )
1286 recent_researches.append(
1287 {
1288 "id": research_data.research_id,
1289 "query": query_text,
1290 "tokens": research_data.token_count or 0,
1291 "created_at": research_data.latest_timestamp,
1292 }
1293 )
1295 # Token breakdown statistics
1296 breakdown_query = query.with_entities(
1297 func.sum(TokenUsage.prompt_tokens).label(
1298 "total_input_tokens"
1299 ),
1300 func.sum(TokenUsage.completion_tokens).label(
1301 "total_output_tokens"
1302 ),
1303 func.avg(TokenUsage.prompt_tokens).label(
1304 "avg_input_tokens"
1305 ),
1306 func.avg(TokenUsage.completion_tokens).label(
1307 "avg_output_tokens"
1308 ),
1309 func.avg(TokenUsage.total_tokens).label("avg_total_tokens"),
1310 )
1311 token_breakdown = breakdown_query.first()
1313 # Get rate limiting metrics
1314 from ..database.models import (
1315 RateLimitAttempt,
1316 RateLimitEstimate,
1317 )
1319 # Get rate limit attempts
1320 rate_limit_query = session.query(RateLimitAttempt)
1322 # Apply time filter
1323 if time_condition is not None:
1324 # RateLimitAttempt uses timestamp as float, not datetime
1325 if period == "7d":
1326 cutoff_time = time.time() - (7 * 24 * 3600)
1327 elif period == "30d":
1328 cutoff_time = time.time() - (30 * 24 * 3600)
1329 elif period == "3m":
1330 cutoff_time = time.time() - (90 * 24 * 3600)
1331 elif period == "1y":
1332 cutoff_time = time.time() - (365 * 24 * 3600)
1333 else: # all
1334 cutoff_time = 0
1336 if cutoff_time > 0:
1337 rate_limit_query = rate_limit_query.filter(
1338 RateLimitAttempt.timestamp >= cutoff_time
1339 )
1341 # Get rate limit statistics
1342 total_attempts = rate_limit_query.count()
1343 successful_attempts = rate_limit_query.filter(
1344 RateLimitAttempt.success
1345 ).count()
1346 failed_attempts = total_attempts - successful_attempts
1348 # Count rate limiting events (failures with RateLimitError)
1349 rate_limit_events = rate_limit_query.filter(
1350 ~RateLimitAttempt.success,
1351 RateLimitAttempt.error_type == "RateLimitError",
1352 ).count()
1354 logger.debug(
1355 f"Rate limit attempts in database: total={total_attempts}, successful={successful_attempts}"
1356 )
1358 # Get all attempts for detailed calculations
1359 attempts = rate_limit_query.all()
1361 # Calculate average wait times
1362 if attempts:
1363 avg_wait_time = sum(a.wait_time for a in attempts) / len(
1364 attempts
1365 )
1366 successful_wait_times = [
1367 a.wait_time for a in attempts if a.success
1368 ]
1369 avg_successful_wait = (
1370 sum(successful_wait_times) / len(successful_wait_times)
1371 if successful_wait_times
1372 else 0
1373 )
1374 else:
1375 avg_wait_time = 0
1376 avg_successful_wait = 0
1378 # Get tracked engines - count distinct engine types from attempts
1379 tracked_engines_query = session.query(
1380 func.count(func.distinct(RateLimitAttempt.engine_type))
1381 )
1382 if cutoff_time > 0:
1383 tracked_engines_query = tracked_engines_query.filter(
1384 RateLimitAttempt.timestamp >= cutoff_time
1385 )
1386 tracked_engines = tracked_engines_query.scalar() or 0
1388 # Get engine-specific stats from attempts
1389 engine_stats = []
1391 # Get distinct engine types from attempts
1392 engine_types_query = session.query(
1393 RateLimitAttempt.engine_type
1394 ).distinct()
1395 if cutoff_time > 0:
1396 engine_types_query = engine_types_query.filter(
1397 RateLimitAttempt.timestamp >= cutoff_time
1398 )
1399 engine_types = [
1400 row.engine_type for row in engine_types_query.all()
1401 ]
1403 # Batch-load all estimates (fix N+1 query)
1404 estimates_map = {}
1405 if engine_types:
1406 all_estimates = (
1407 session.query(RateLimitEstimate)
1408 .filter(RateLimitEstimate.engine_type.in_(engine_types))
1409 .all()
1410 )
1411 estimates_map = {e.engine_type: e for e in all_estimates}
1413 for engine_type in engine_types:
1414 engine_attempts_list = [
1415 a for a in attempts if a.engine_type == engine_type
1416 ]
1417 engine_attempts = len(engine_attempts_list)
1418 engine_success = len(
1419 [a for a in engine_attempts_list if a.success]
1420 )
1422 # Get estimate if exists
1423 estimate = estimates_map.get(engine_type)
1425 # Calculate recent success rate
1426 recent_success_rate = (
1427 (engine_success / engine_attempts * 100)
1428 if engine_attempts > 0
1429 else 0
1430 )
1432 # Determine status based on success rate
1433 if estimate:
1434 status = (
1435 "healthy"
1436 if estimate.success_rate > 0.8
1437 else "degraded"
1438 if estimate.success_rate > 0.5
1439 else "poor"
1440 )
1441 else:
1442 status = (
1443 "healthy"
1444 if recent_success_rate > 80
1445 else "degraded"
1446 if recent_success_rate > 50
1447 else "poor"
1448 )
1450 engine_stat = {
1451 "engine": engine_type,
1452 "base_wait": estimate.base_wait_seconds
1453 if estimate
1454 else 0.0,
1455 "base_wait_seconds": round(
1456 estimate.base_wait_seconds if estimate else 0.0, 2
1457 ),
1458 "min_wait_seconds": round(
1459 estimate.min_wait_seconds if estimate else 0.0, 2
1460 ),
1461 "max_wait_seconds": round(
1462 estimate.max_wait_seconds if estimate else 0.0, 2
1463 ),
1464 "success_rate": round(estimate.success_rate * 100, 1)
1465 if estimate
1466 else recent_success_rate,
1467 "total_attempts": estimate.total_attempts
1468 if estimate
1469 else engine_attempts,
1470 "recent_attempts": engine_attempts,
1471 "recent_success_rate": round(recent_success_rate, 1),
1472 "attempts": engine_attempts,
1473 "status": status,
1474 }
1476 if estimate:
1477 # Pass UTC explicitly: web/routers/metrics.py renders
1478 # this same RateLimitEstimate.last_updated field with
1479 # ``fromtimestamp(..., UTC)`` (lines ~700 and ~967).
1480 # Without the tz argument this is server-local, so the
1481 # two surfaces showed the same timestamp shifted by the
1482 # host's UTC offset.
1483 engine_stat["last_updated"] = datetime.fromtimestamp(
1484 estimate.last_updated, UTC
1485 ).strftime("%Y-%m-%d %H:%M:%S")
1486 else:
1487 engine_stat["last_updated"] = "Never"
1489 engine_stats.append(engine_stat)
1491 logger.debug(
1492 f"Tracked engines: {tracked_engines}, engine_stats: {engine_stats}"
1493 )
1495 result = {
1496 "total_tokens": total_tokens,
1497 "total_researches": total_researches,
1498 "by_model": by_model,
1499 "recent_researches": recent_researches,
1500 "token_breakdown": {
1501 "total_input_tokens": int(
1502 token_breakdown.total_input_tokens or 0
1503 ),
1504 "total_output_tokens": int(
1505 token_breakdown.total_output_tokens or 0
1506 ),
1507 "avg_input_tokens": int(
1508 token_breakdown.avg_input_tokens or 0
1509 ),
1510 "avg_output_tokens": int(
1511 token_breakdown.avg_output_tokens or 0
1512 ),
1513 "avg_total_tokens": int(
1514 token_breakdown.avg_total_tokens or 0
1515 ),
1516 },
1517 "rate_limiting": {
1518 "total_attempts": total_attempts,
1519 "successful_attempts": successful_attempts,
1520 "failed_attempts": failed_attempts,
1521 "success_rate": (
1522 successful_attempts / total_attempts * 100
1523 )
1524 if total_attempts > 0
1525 else 0,
1526 "rate_limit_events": rate_limit_events,
1527 "avg_wait_time": round(float(avg_wait_time), 2),
1528 "avg_successful_wait": round(
1529 float(avg_successful_wait), 2
1530 ),
1531 "tracked_engines": tracked_engines,
1532 "engine_stats": engine_stats,
1533 "total_engines_tracked": tracked_engines,
1534 "healthy_engines": len(
1535 [
1536 s
1537 for s in engine_stats
1538 if s["status"] == "healthy"
1539 ]
1540 ),
1541 "degraded_engines": len(
1542 [
1543 s
1544 for s in engine_stats
1545 if s["status"] == "degraded"
1546 ]
1547 ),
1548 "poor_engines": len(
1549 [s for s in engine_stats if s["status"] == "poor"]
1550 ),
1551 },
1552 }
1554 logger.debug(
1555 f"Returning from _get_metrics_from_encrypted_db - total_researches: {result['total_researches']}"
1556 )
1557 return result
1558 except Exception:
1559 logger.exception(
1560 "CRITICAL ERROR accessing encrypted database for metrics"
1561 )
1562 return self._get_empty_metrics()
1564 def _get_empty_metrics(self) -> Dict[str, Any]:
1565 """Return empty metrics structure when no data is available."""
1566 return {
1567 "total_tokens": 0,
1568 "total_researches": 0,
1569 "by_model": [],
1570 "recent_researches": [],
1571 "token_breakdown": {
1572 "prompt_tokens": 0,
1573 "completion_tokens": 0,
1574 "avg_prompt_tokens": 0,
1575 "avg_completion_tokens": 0,
1576 "avg_total_tokens": 0,
1577 },
1578 }
1580 def get_enhanced_metrics(
1581 self,
1582 period: str = "30d",
1583 research_mode: str = "all",
1584 username: Optional[str] = None,
1585 ) -> Dict[str, Any]:
1586 """Get enhanced Phase 1 tracking metrics.
1588 Args:
1589 period: Time period to filter by ('7d', '30d', '3m', '1y', 'all')
1590 research_mode: Research mode to filter by ('quick', 'detailed', 'all')
1591 username: Authenticated user whose encrypted DB to open. If
1592 omitted, falls back to the request's contextvar — route
1593 handlers should pass it explicitly so the lookup is
1594 obviously user-scoped.
1596 Returns:
1597 Dictionary containing enhanced metrics data including time series
1598 """
1600 from ..database.session_context import get_user_db_session
1602 if username is None:
1603 username = get_current_username()
1604 if not username:
1605 # Return empty metrics structure when no user session
1606 return {
1607 "recent_enhanced_data": [],
1608 "performance_stats": {
1609 "avg_response_time": 0,
1610 "min_response_time": 0,
1611 "max_response_time": 0,
1612 "success_rate": 0,
1613 "error_rate": 0,
1614 "total_enhanced_calls": 0,
1615 },
1616 "mode_breakdown": [],
1617 "search_engine_stats": [],
1618 "phase_breakdown": [],
1619 "time_series_data": [],
1620 "call_stack_analysis": {
1621 "by_file": [],
1622 "by_function": [],
1623 },
1624 }
1626 try:
1627 with get_user_db_session(username) as session:
1628 # Build base query with filters
1629 query = session.query(TokenUsage)
1631 # Apply time filter
1632 time_condition = get_time_filter_condition(
1633 period, TokenUsage.timestamp
1634 )
1635 if time_condition is not None: 1635 ↛ 1639line 1635 didn't jump to line 1639 because the condition on line 1635 was always true
1636 query = query.filter(time_condition)
1638 # Apply research mode filter
1639 mode_condition = get_research_mode_condition(
1640 research_mode, TokenUsage.research_mode
1641 )
1642 if mode_condition is not None: 1642 ↛ 1643line 1642 didn't jump to line 1643 because the condition on line 1642 was never true
1643 query = query.filter(mode_condition)
1645 # Get time series data for the chart - most important for "Token Consumption Over Time"
1646 time_series_query = query.filter(
1647 TokenUsage.timestamp.isnot(None),
1648 TokenUsage.total_tokens > 0,
1649 ).order_by(TokenUsage.timestamp.asc())
1651 # Limit to recent data for performance
1652 if period != "all": 1652 ↛ 1655line 1652 didn't jump to line 1655 because the condition on line 1652 was always true
1653 time_series_query = time_series_query.limit(200)
1655 time_series_data = time_series_query.all()
1657 # Format time series data with cumulative calculations
1658 time_series = []
1659 cumulative_tokens = 0
1660 cumulative_prompt_tokens = 0
1661 cumulative_completion_tokens = 0
1663 for usage in time_series_data: 1663 ↛ 1664line 1663 didn't jump to line 1664 because the loop on line 1663 never started
1664 cumulative_tokens += usage.total_tokens or 0
1665 cumulative_prompt_tokens += usage.prompt_tokens or 0
1666 cumulative_completion_tokens += usage.completion_tokens or 0
1668 time_series.append(
1669 {
1670 "timestamp": str(usage.timestamp)
1671 if usage.timestamp
1672 else None,
1673 "tokens": usage.total_tokens or 0,
1674 "prompt_tokens": usage.prompt_tokens or 0,
1675 "completion_tokens": usage.completion_tokens or 0,
1676 "cumulative_tokens": cumulative_tokens,
1677 "cumulative_prompt_tokens": cumulative_prompt_tokens,
1678 "cumulative_completion_tokens": cumulative_completion_tokens,
1679 "research_id": usage.research_id,
1680 }
1681 )
1683 # Basic performance stats using ORM
1684 performance_query = query.filter(
1685 TokenUsage.response_time_ms.isnot(None)
1686 )
1687 total_calls = performance_query.count()
1689 if total_calls > 0:
1690 avg_response_time = (
1691 performance_query.with_entities(
1692 func.avg(TokenUsage.response_time_ms)
1693 ).scalar()
1694 or 0
1695 )
1696 min_response_time = (
1697 performance_query.with_entities(
1698 func.min(TokenUsage.response_time_ms)
1699 ).scalar()
1700 or 0
1701 )
1702 max_response_time = (
1703 performance_query.with_entities(
1704 func.max(TokenUsage.response_time_ms)
1705 ).scalar()
1706 or 0
1707 )
1708 success_count = performance_query.filter(
1709 TokenUsage.success_status == "success"
1710 ).count()
1711 error_count = performance_query.filter(
1712 TokenUsage.success_status == "error"
1713 ).count()
1715 perf_stats = {
1716 "avg_response_time": round(avg_response_time),
1717 "min_response_time": min_response_time,
1718 "max_response_time": max_response_time,
1719 "success_rate": (
1720 round((success_count / total_calls * 100), 1)
1721 if total_calls > 0
1722 else 0
1723 ),
1724 "error_rate": (
1725 round((error_count / total_calls * 100), 1)
1726 if total_calls > 0
1727 else 0
1728 ),
1729 "total_enhanced_calls": total_calls,
1730 }
1731 else:
1732 perf_stats = {
1733 "avg_response_time": 0,
1734 "min_response_time": 0,
1735 "max_response_time": 0,
1736 "success_rate": 0,
1737 "error_rate": 0,
1738 "total_enhanced_calls": 0,
1739 }
1741 # Research mode breakdown using ORM
1742 mode_stats = (
1743 query.filter(TokenUsage.research_mode.isnot(None))
1744 .with_entities(
1745 TokenUsage.research_mode,
1746 func.count().label("count"),
1747 func.avg(TokenUsage.total_tokens).label("avg_tokens"),
1748 func.avg(TokenUsage.response_time_ms).label(
1749 "avg_response_time"
1750 ),
1751 )
1752 .group_by(TokenUsage.research_mode)
1753 .all()
1754 )
1756 modes = [
1757 {
1758 "mode": stat.research_mode,
1759 "count": stat.count,
1760 "avg_tokens": round(stat.avg_tokens or 0),
1761 "avg_response_time": round(stat.avg_response_time or 0),
1762 }
1763 for stat in mode_stats
1764 ]
1766 # Recent enhanced data (simplified)
1767 recent_enhanced_query = (
1768 query.filter(TokenUsage.research_query.isnot(None))
1769 .order_by(TokenUsage.timestamp.desc())
1770 .limit(50)
1771 )
1773 recent_enhanced_data = recent_enhanced_query.all()
1774 recent_enhanced = [
1775 {
1776 "research_query": usage.research_query,
1777 "research_mode": usage.research_mode,
1778 "research_phase": usage.research_phase,
1779 "search_iteration": usage.search_iteration,
1780 "response_time_ms": usage.response_time_ms,
1781 "success_status": usage.success_status,
1782 "error_type": usage.error_type,
1783 "search_engines_planned": usage.search_engines_planned,
1784 "search_engine_selected": usage.search_engine_selected,
1785 "total_tokens": usage.total_tokens,
1786 "prompt_tokens": usage.prompt_tokens,
1787 "completion_tokens": usage.completion_tokens,
1788 "timestamp": str(usage.timestamp)
1789 if usage.timestamp
1790 else None,
1791 "research_id": usage.research_id,
1792 "calling_file": usage.calling_file,
1793 "calling_function": usage.calling_function,
1794 "call_stack": usage.call_stack,
1795 }
1796 for usage in recent_enhanced_data
1797 ]
1799 # Search engine breakdown using ORM
1800 search_engine_stats = (
1801 query.filter(TokenUsage.search_engine_selected.isnot(None))
1802 .with_entities(
1803 TokenUsage.search_engine_selected,
1804 func.count().label("count"),
1805 func.avg(TokenUsage.total_tokens).label("avg_tokens"),
1806 func.avg(TokenUsage.response_time_ms).label(
1807 "avg_response_time"
1808 ),
1809 )
1810 .group_by(TokenUsage.search_engine_selected)
1811 .all()
1812 )
1814 search_engines = [
1815 {
1816 "search_engine": stat.search_engine_selected,
1817 "count": stat.count,
1818 "avg_tokens": round(stat.avg_tokens or 0),
1819 "avg_response_time": round(stat.avg_response_time or 0),
1820 }
1821 for stat in search_engine_stats
1822 ]
1824 # Research phase breakdown using ORM
1825 phase_stats = (
1826 query.filter(TokenUsage.research_phase.isnot(None))
1827 .with_entities(
1828 TokenUsage.research_phase,
1829 func.count().label("count"),
1830 func.avg(TokenUsage.total_tokens).label("avg_tokens"),
1831 func.avg(TokenUsage.response_time_ms).label(
1832 "avg_response_time"
1833 ),
1834 )
1835 .group_by(TokenUsage.research_phase)
1836 .all()
1837 )
1839 phases = [
1840 {
1841 "phase": stat.research_phase,
1842 "count": stat.count,
1843 "avg_tokens": round(stat.avg_tokens or 0),
1844 "avg_response_time": round(stat.avg_response_time or 0),
1845 }
1846 for stat in phase_stats
1847 ]
1849 # Call stack analysis using ORM
1850 file_stats = (
1851 query.filter(TokenUsage.calling_file.isnot(None))
1852 .with_entities(
1853 TokenUsage.calling_file,
1854 func.count().label("count"),
1855 func.avg(TokenUsage.total_tokens).label("avg_tokens"),
1856 )
1857 .group_by(TokenUsage.calling_file)
1858 .order_by(func.count().desc())
1859 .limit(10)
1860 .all()
1861 )
1863 files = [
1864 {
1865 "file": stat.calling_file,
1866 "count": stat.count,
1867 "avg_tokens": round(stat.avg_tokens or 0),
1868 }
1869 for stat in file_stats
1870 ]
1872 function_stats = (
1873 query.filter(TokenUsage.calling_function.isnot(None))
1874 .with_entities(
1875 TokenUsage.calling_function,
1876 func.count().label("count"),
1877 func.avg(TokenUsage.total_tokens).label("avg_tokens"),
1878 )
1879 .group_by(TokenUsage.calling_function)
1880 .order_by(func.count().desc())
1881 .limit(10)
1882 .all()
1883 )
1885 functions = [
1886 {
1887 "function": stat.calling_function,
1888 "count": stat.count,
1889 "avg_tokens": round(stat.avg_tokens or 0),
1890 }
1891 for stat in function_stats
1892 ]
1894 return {
1895 "recent_enhanced_data": recent_enhanced,
1896 "performance_stats": perf_stats,
1897 "mode_breakdown": modes,
1898 "search_engine_stats": search_engines,
1899 "phase_breakdown": phases,
1900 "time_series_data": time_series,
1901 "call_stack_analysis": {
1902 "by_file": files,
1903 "by_function": functions,
1904 },
1905 }
1906 except Exception:
1907 logger.exception("Error in get_enhanced_metrics")
1908 # Return simplified response without non-existent columns
1909 return {
1910 "recent_enhanced_data": [],
1911 "performance_stats": {
1912 "avg_response_time": 0,
1913 "min_response_time": 0,
1914 "max_response_time": 0,
1915 "success_rate": 0,
1916 "error_rate": 0,
1917 "total_enhanced_calls": 0,
1918 },
1919 "mode_breakdown": [],
1920 "search_engine_stats": [],
1921 "phase_breakdown": [],
1922 "time_series_data": [],
1923 "call_stack_analysis": {
1924 "by_file": [],
1925 "by_function": [],
1926 },
1927 }
1929 def get_research_timeline_metrics(
1930 self, research_id: str, username: Optional[str] = None
1931 ) -> Dict[str, Any]:
1932 """Get timeline metrics for a specific research.
1934 Args:
1935 research_id: The ID of the research
1936 username: Authenticated user whose encrypted DB to open. If
1937 omitted, falls back to the request's contextvar.
1939 Returns:
1940 Dictionary containing timeline metrics for the research
1941 """
1943 from ..database.session_context import get_user_db_session
1945 if username is None:
1946 username = get_current_username()
1947 if not username:
1948 return {
1949 "research_id": research_id,
1950 "research_details": {},
1951 "timeline": [],
1952 "summary": {
1953 "total_calls": 0,
1954 "total_tokens": 0,
1955 "total_prompt_tokens": 0,
1956 "total_completion_tokens": 0,
1957 "avg_response_time": 0,
1958 "success_rate": 0,
1959 },
1960 "phase_stats": {},
1961 }
1963 with get_user_db_session(username) as session:
1964 # Get all token usage for this research ordered by time including call stack
1965 timeline_data = session.execute(
1966 text(
1967 """
1968 SELECT
1969 timestamp,
1970 total_tokens,
1971 prompt_tokens,
1972 completion_tokens,
1973 response_time_ms,
1974 success_status,
1975 error_type,
1976 research_phase,
1977 search_iteration,
1978 search_engine_selected,
1979 model_name,
1980 calling_file,
1981 calling_function,
1982 call_stack
1983 FROM token_usage
1984 WHERE research_id = :research_id
1985 ORDER BY timestamp ASC
1986 """
1987 ),
1988 {"research_id": research_id},
1989 ).fetchall()
1991 # Format timeline data with cumulative tokens
1992 timeline = []
1993 cumulative_tokens = 0
1994 cumulative_prompt_tokens = 0
1995 cumulative_completion_tokens = 0
1997 for row in timeline_data:
1998 cumulative_tokens += row[1] or 0
1999 cumulative_prompt_tokens += row[2] or 0
2000 cumulative_completion_tokens += row[3] or 0
2002 # call_stack is a JSON column. Depending on the
2003 # SQLAlchemy/SQLite/driver combo the raw text() SELECT may
2004 # either hand back the raw JSON-encoded TEXT or already
2005 # the decoded Python value. Normalise to the API/UI
2006 # contract (`string | null`): pass through None and
2007 # strings; re-serialise any other stored JSON shape
2008 # (number / bool / list / object) so a Python int/bool/
2009 # list never leaks to the browser. Legacy unparseable
2010 # strings pass through as-is.
2011 raw_call_stack = row[13]
2012 if raw_call_stack is None or raw_call_stack == "":
2013 decoded_call_stack = None
2014 elif isinstance(raw_call_stack, str):
2015 try:
2016 decoded = json.loads(raw_call_stack)
2017 except (TypeError, ValueError):
2018 decoded_call_stack = raw_call_stack
2019 else:
2020 decoded_call_stack = (
2021 decoded
2022 if decoded is None or isinstance(decoded, str)
2023 else json.dumps(decoded)
2024 )
2025 else:
2026 decoded_call_stack = (
2027 raw_call_stack
2028 if isinstance(raw_call_stack, str)
2029 else json.dumps(raw_call_stack)
2030 )
2032 timeline.append(
2033 {
2034 "timestamp": str(row[0]) if row[0] else None,
2035 "tokens": row[1] or 0,
2036 "prompt_tokens": row[2] or 0,
2037 "completion_tokens": row[3] or 0,
2038 "cumulative_tokens": cumulative_tokens,
2039 "cumulative_prompt_tokens": cumulative_prompt_tokens,
2040 "cumulative_completion_tokens": cumulative_completion_tokens,
2041 "response_time_ms": row[4],
2042 "success_status": row[5],
2043 "error_type": row[6],
2044 "research_phase": row[7],
2045 "search_iteration": row[8],
2046 "search_engine_selected": row[9],
2047 "model_name": row[10],
2048 "calling_file": row[11],
2049 "calling_function": row[12],
2050 "call_stack": decoded_call_stack,
2051 }
2052 )
2054 # Get research basic info
2055 research_info = session.execute(
2056 text(
2057 """
2058 SELECT query, mode, status, created_at, completed_at
2059 FROM research_history
2060 WHERE id = :research_id
2061 """
2062 ),
2063 {"research_id": research_id},
2064 ).fetchone()
2066 research_details = {}
2067 if research_info:
2068 research_details = {
2069 "query": research_info[0],
2070 "mode": research_info[1],
2071 "status": research_info[2],
2072 "created_at": str(research_info[3])
2073 if research_info[3]
2074 else None,
2075 "completed_at": str(research_info[4])
2076 if research_info[4]
2077 else None,
2078 }
2080 # Calculate summary stats
2081 total_calls = len(timeline_data)
2082 total_tokens = cumulative_tokens
2083 avg_response_time = sum(row[4] or 0 for row in timeline_data) / max(
2084 total_calls, 1
2085 )
2086 success_rate = (
2087 sum(1 for row in timeline_data if row[5] == "success")
2088 / max(total_calls, 1)
2089 * 100
2090 )
2092 # Phase breakdown for this research
2093 phase_stats = {}
2094 for row in timeline_data:
2095 phase = row[7] or "unknown"
2096 if phase not in phase_stats:
2097 phase_stats[phase] = {
2098 "count": 0,
2099 "tokens": 0,
2100 "avg_response_time": 0,
2101 }
2102 phase_stats[phase]["count"] += 1
2103 phase_stats[phase]["tokens"] += row[1] or 0
2104 if row[4]:
2105 phase_stats[phase]["avg_response_time"] += row[4]
2107 # Calculate averages for phases
2108 for phase in phase_stats:
2109 if phase_stats[phase]["count"] > 0: 2109 ↛ 2108line 2109 didn't jump to line 2108 because the condition on line 2109 was always true
2110 phase_stats[phase]["avg_response_time"] = round(
2111 phase_stats[phase]["avg_response_time"]
2112 / phase_stats[phase]["count"]
2113 )
2115 return {
2116 "research_id": research_id,
2117 "research_details": research_details,
2118 "timeline": timeline,
2119 "summary": {
2120 "total_calls": total_calls,
2121 "total_tokens": total_tokens,
2122 "total_prompt_tokens": cumulative_prompt_tokens,
2123 "total_completion_tokens": cumulative_completion_tokens,
2124 "avg_response_time": round(avg_response_time),
2125 "success_rate": round(success_rate, 1),
2126 },
2127 "phase_stats": phase_stats,
2128 }