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