Coverage for src / local_deep_research / database / thread_metrics.py: 92%

1""" 

2Thread-safe metrics database access. 

3 

4This module provides a way for background threads to write metrics 

5to the user's encrypted database by creating thread-local connections 

6with the provided password. 

7""" 

8 

9import threading 

10from contextlib import contextmanager 

11from typing import Optional 

12 

13from loguru import logger 

14from sqlalchemy.orm import Session 

15 

16from .encrypted_db import db_manager 

17 

18 

19class ThreadSafeMetricsWriter: 

20 """ 

21 Thread-safe writer for metrics to encrypted user databases. 

22 Creates encrypted connections per thread using provided passwords. 

23 """ 

24 

25 def __init__(self): 

26 self._thread_local = threading.local() 

27 

28 def set_user_password(self, username: str, password: str): 

29 """ 

30 Store user password for the current thread. 

31 This allows the thread to create its own encrypted connection. 

32 

33 IMPORTANT: This is safe because: 

34 1. Password is already in memory (user is logged in) 

35 2. It's only stored thread-locally 

36 3. It's cleared when the thread ends 

37 """ 

38 

39 if not hasattr(self._thread_local, "passwords"): 

40 self._thread_local.passwords = {} 

41 self._thread_local.passwords[username] = password 

42 

43 @contextmanager 

44 def get_session(self, username: str = None) -> Session: 

45 """ 

46 Get a database session for metrics in the current thread. 

47 Creates a new encrypted connection if needed. 

48 

49 Args: 

50 username: The username for database access. If not provided, 

51 will attempt to get it from Flask session. 

52 """ 

53 # If username not provided, try to get it from Flask session 

54 if username is None: 

55 try: 

56 from flask import session as flask_session 

57 from werkzeug.exceptions import Unauthorized 

58 

59 username = flask_session.get("username") 

60 if not username: 

61 raise Unauthorized("No username in Flask session") 

62 except (ImportError, RuntimeError) as e: 

63 # Flask context not available or no session 

64 raise ValueError(f"Cannot determine username: {e}") 

65 

66 # Get password for this user in this thread 

67 if not hasattr(self._thread_local, "passwords"): 

68 raise ValueError("No password set for thread metrics access") 

69 

70 password = self._thread_local.passwords.get(username) 

71 

72 if not password: 

73 raise ValueError( 

74 f"No password available for user {username} in this thread" 

75 ) 

76 

77 # Create a thread-safe session for this user 

78 session = None 

79 try: 

80 session = db_manager.create_thread_safe_session_for_metrics( 

81 username, password 

82 ) 

83 if not session: 

84 raise ValueError( # noqa: TRY301 — except does session rollback before re-raise 

85 f"Failed to create session for user {username}" 

86 ) 

87 yield session 

88 session.commit() 

89 except Exception: 

90 logger.exception(f"Session error for {username}") 

91 if session: 

92 session.rollback() 

93 raise 

94 finally: 

95 if session: 95 ↛ exitline 95 didn't return from function 'get_session' because the condition on line 95 was always true

96 from ..utilities.resource_utils import safe_close 

97 

98 safe_close(session, "thread metrics session") 

99 

100 def write_token_metrics( 

101 self, username: str, research_id: Optional[int], token_data: dict 

102 ): 

103 """ 

104 Write token metrics from any thread. 

105 

106 Args: 

107 username: The username (for database access) 

108 research_id: The research ID 

109 token_data: Dictionary with token metrics data 

110 """ 

111 with self.get_session(username) as session: 

112 # Import here to avoid circular imports 

113 from .models import TokenUsage 

114 

115 # Create TokenUsage record 

116 token_usage = TokenUsage( 

117 research_id=research_id, 

118 model_name=token_data.get("model_name"), 

119 model_provider=token_data.get("provider"), 

120 prompt_tokens=token_data.get("prompt_tokens", 0), 

121 completion_tokens=token_data.get("completion_tokens", 0), 

122 total_tokens=token_data.get("prompt_tokens", 0) 

123 + token_data.get("completion_tokens", 0), 

124 # Research context 

125 research_query=token_data.get("research_query"), 

126 research_mode=token_data.get("research_mode"), 

127 research_phase=token_data.get("research_phase"), 

128 search_iteration=token_data.get("search_iteration"), 

129 # Performance metrics 

130 response_time_ms=token_data.get("response_time_ms"), 

131 success_status=token_data.get("success_status", "success"), 

132 error_type=token_data.get("error_type"), 

133 # Search engine context 

134 search_engines_planned=token_data.get("search_engines_planned"), 

135 search_engine_selected=token_data.get("search_engine_selected"), 

136 # Call stack tracking 

137 calling_file=token_data.get("calling_file"), 

138 calling_function=token_data.get("calling_function"), 

139 call_stack=token_data.get("call_stack"), 

140 # Context overflow detection 

141 context_limit=token_data.get("context_limit"), 

142 context_truncated=token_data.get("context_truncated", False), 

143 tokens_truncated=token_data.get("tokens_truncated"), 

144 truncation_ratio=token_data.get("truncation_ratio"), 

145 # Raw Ollama metrics 

146 ollama_prompt_eval_count=token_data.get( 

147 "ollama_prompt_eval_count" 

148 ), 

149 ollama_eval_count=token_data.get("ollama_eval_count"), 

150 ollama_total_duration=token_data.get("ollama_total_duration"), 

151 ollama_load_duration=token_data.get("ollama_load_duration"), 

152 ollama_prompt_eval_duration=token_data.get( 

153 "ollama_prompt_eval_duration" 

154 ), 

155 ollama_eval_duration=token_data.get("ollama_eval_duration"), 

156 ) 

157 session.add(token_usage) 

158 

159 

160# Global instance for thread-safe metrics 

161metrics_writer = ThreadSafeMetricsWriter()