Coverage for src / local_deep_research / database / models / metrics.py: 96%

1""" 

2Metrics and usage tracking models. 

3""" 

4 

5from sqlalchemy import ( 

6 JSON, 

7 Boolean, 

8 Column, 

9 Float, 

10 Index, 

11 Integer, 

12 String, 

13 Text, 

14) 

15from sqlalchemy_utc import UtcDateTime, utcnow 

16 

17from .base import Base 

18 

19 

20class TokenUsage(Base): 

21 """ 

22 Track token usage for LLM calls. 

23 """ 

24 

25 __tablename__ = "token_usage" 

26 

27 id = Column(Integer, primary_key=True, autoincrement=True) 

28 research_id = Column(String(36), nullable=False, index=True) 

29 timestamp = Column(UtcDateTime, nullable=False, default=utcnow()) 

30 

31 # Model information 

32 model_provider = Column(String(100), nullable=False) 

33 model_name = Column(String(255), nullable=False) 

34 

35 # Token counts 

36 prompt_tokens = Column(Integer, nullable=False) 

37 completion_tokens = Column(Integer, nullable=False) 

38 total_tokens = Column(Integer, nullable=False) 

39 

40 # Cost tracking (in USD) 

41 prompt_cost = Column(Float, default=0.0) 

42 completion_cost = Column(Float, default=0.0) 

43 total_cost = Column(Float, default=0.0) 

44 

45 # Context 

46 operation_type = Column(String(100)) # search, summarize, report, etc. 

47 operation_details = Column(JSON) 

48 research_mode = Column(String(50)) # standard, deep, expert, etc. 

49 

50 # Enhanced metrics columns 

51 response_time_ms = Column(Integer) 

52 success_status = Column(String(50)) 

53 error_type = Column(String(100)) 

54 research_query = Column(Text) 

55 research_phase = Column(String(100)) 

56 search_iteration = Column(Integer) 

57 search_engines_planned = Column(JSON) 

58 search_engine_selected = Column(String(100)) 

59 calling_file = Column(String(255)) 

60 calling_function = Column(String(255)) 

61 call_stack = Column(JSON) 

62 

63 # Context overflow detection columns 

64 context_limit = Column(Integer) # The configured num_ctx or max tokens 

65 context_truncated = Column( 

66 Boolean, default=False 

67 ) # True if request was truncated due to context limit 

68 tokens_truncated = Column(Integer) # Estimated tokens lost to truncation 

69 truncation_ratio = Column(Float) # Percentage of prompt that was truncated 

70 

71 # Raw Ollama response values for debugging 

72 ollama_prompt_eval_count = Column( 

73 Integer 

74 ) # Raw prompt_eval_count from Ollama 

75 ollama_eval_count = Column(Integer) # Raw eval_count from Ollama 

76 ollama_total_duration = Column( 

77 Integer 

78 ) # Total time in nanoseconds (raw from Ollama API) 

79 ollama_load_duration = Column( 

80 Integer 

81 ) # Model load time in nanoseconds (raw from Ollama API) 

82 ollama_prompt_eval_duration = Column( 

83 Integer 

84 ) # Prompt eval time in nanoseconds (raw from Ollama API) 

85 ollama_eval_duration = Column( 

86 Integer 

87 ) # Generation time in nanoseconds (raw from Ollama API) 

88 

89 # Compound indexes for query performance optimization 

90 __table_args__ = ( 

91 Index("idx_token_research_timestamp", "research_id", "timestamp"), 

92 Index( 

93 "idx_token_model_timestamp", 

94 "model_provider", 

95 "model_name", 

96 "timestamp", 

97 ), 

98 Index("idx_token_operation_timestamp", "operation_type", "timestamp"), 

99 Index("idx_token_research_phase", "research_id", "research_phase"), 

100 ) 

101 

102 def __repr__(self): 

103 return f"<TokenUsage(model={self.model_name}, total_tokens={self.total_tokens}, cost=${self.total_cost:.4f})>" 

104 

105 

106class ModelUsage(Base): 

107 """ 

108 Aggregate model usage statistics. 

109 """ 

110 

111 __tablename__ = "model_usage" 

112 

113 id = Column(Integer, primary_key=True) 

114 model_provider = Column(String(100), nullable=False) 

115 model_name = Column(String(255), nullable=False) 

116 

117 # Aggregate stats 

118 total_calls = Column(Integer, default=0) 

119 total_tokens = Column(Integer, default=0) 

120 total_cost = Column(Float, default=0.0) 

121 

122 # Performance metrics 

123 avg_response_time_ms = Column(Float) 

124 error_count = Column(Integer, default=0) 

125 success_rate = Column(Float, default=100.0) 

126 

127 # Time tracking 

128 first_used_at = Column(UtcDateTime, default=utcnow()) 

129 last_used_at = Column(UtcDateTime, default=utcnow(), onupdate=utcnow()) 

130 

131 def __repr__(self): 

132 return f"<ModelUsage(model={self.model_name}, calls={self.total_calls}, cost=${self.total_cost:.2f})>" 

133 

134 

135class ResearchRating(Base): 

136 """ 

137 User ratings for research results. 

138 """ 

139 

140 __tablename__ = "research_ratings" 

141 

142 id = Column(Integer, primary_key=True) 

143 research_id = Column(String(36), nullable=False, unique=True, index=True) 

144 

145 # Star rating (1-5) 

146 rating = Column(Integer, nullable=False) 

147 

148 # Feedback categories 

149 accuracy = Column(Integer) # 1-5 

150 completeness = Column(Integer) # 1-5 

151 relevance = Column(Integer) # 1-5 

152 readability = Column(Integer) # 1-5 

153 

154 # Written feedback 

155 feedback = Column(Text) 

156 

157 # Timestamps 

158 created_at = Column(UtcDateTime, default=utcnow()) 

159 updated_at = Column(UtcDateTime, default=utcnow(), onupdate=utcnow()) 

160 

161 def __repr__(self): 

162 return f"<ResearchRating(research_id={self.research_id}, rating={self.rating})>" 

163 

164 

165class SearchCall(Base): 

166 """ 

167 Track individual search engine calls. 

168 """ 

169 

170 __tablename__ = "search_calls" 

171 

172 id = Column(Integer, primary_key=True, autoincrement=True) 

173 research_id = Column(String(36), nullable=False, index=True) 

174 timestamp = Column(UtcDateTime, nullable=False, default=utcnow()) 

175 

176 # Search details 

177 search_engine = Column(String(100), nullable=False) 

178 query = Column(Text, nullable=False) 

179 num_results_requested = Column(Integer) 

180 num_results_returned = Column(Integer) 

181 

182 # Performance 

183 response_time_ms = Column(Float) 

184 success = Column(Integer, default=1) # 1 for success, 0 for failure 

185 error_message = Column(Text) 

186 

187 # Rate limiting 

188 rate_limited = Column(Integer, default=0) # 1 if rate limited 

189 wait_time_ms = Column(Float) 

190 

191 # Research context 

192 research_mode = Column(String(50)) # standard, deep, expert, etc. 

193 research_query = Column(Text) 

194 research_phase = Column(String(100)) 

195 search_iteration = Column(Integer) 

196 success_status = Column(String(50)) 

197 error_type = Column(String(100)) 

198 results_count = Column(Integer) 

199 

200 # Compound indexes for query performance optimization 

201 __table_args__ = ( 

202 Index("idx_search_research_timestamp", "research_id", "timestamp"), 

203 Index("idx_search_engine_timestamp", "search_engine", "timestamp"), 

204 Index("idx_search_success_timestamp", "success", "timestamp"), 

205 Index("idx_search_research_engine", "research_id", "search_engine"), 

206 ) 

207 

208 def __repr__(self): 

209 return f"<SearchCall(engine={self.search_engine}, query='{self.query[:50]}...', success={self.success})>"