Coverage for src / mcp_server_langgraph / database / models.py: 100%

1""" 

2SQLAlchemy models for cost tracking persistence. 

3 

4This module defines database models for storing LLM token usage and cost metrics 

5with PostgreSQL persistence and automatic retention policies. 

6""" 

7 

8from datetime import datetime, UTC 

9from decimal import Decimal 

10from typing import Any 

11 

12from sqlalchemy import JSON, DateTime, Index, Integer, Numeric, String 

13from sqlalchemy.orm import Mapped, declarative_base, mapped_column 

14 

15Base = declarative_base() 

16 

17 

18class TokenUsageRecord(Base): # type: ignore[misc,valid-type] 

19 """ 

20 Persistent storage for LLM token usage and cost data. 

21 

22 This model stores detailed token usage metrics for cost tracking, 

23 budgeting, and analytics with automatic retention policies. 

24 

25 Indexes: 

26 - timestamp: For time-range queries and retention cleanup 

27 - user_id: For per-user cost analysis 

28 - session_id: For per-session cost tracking 

29 - model: For per-model cost analysis 

30 - composite (user_id, timestamp): For efficient user cost queries 

31 

32 Retention: 

33 Records are automatically purged after CONTEXT_RETENTION_DAYS (default: 90) 

34 via background cleanup job. 

35 """ 

36 

37 __tablename__ = "token_usage_records" 

38 

39 # Primary key 

40 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) 

41 

42 # Timestamps 

43 timestamp: Mapped[datetime] = mapped_column( 

44 DateTime(timezone=True), 

45 nullable=False, 

46 index=True, 

47 doc="When the LLM call was made (UTC)", 

48 ) 

49 created_at: Mapped[datetime] = mapped_column( 

50 DateTime(timezone=True), 

51 nullable=False, 

52 default=lambda: datetime.now(UTC), 

53 doc="When the record was created (UTC)", 

54 ) 

55 

56 # Identifiers 

57 user_id: Mapped[str] = mapped_column( 

58 String(255), 

59 nullable=False, 

60 index=True, 

61 doc="User identifier", 

62 ) 

63 session_id: Mapped[str] = mapped_column( 

64 String(255), 

65 nullable=False, 

66 index=True, 

67 doc="Session identifier", 

68 ) 

69 

70 # Model information 

71 model: Mapped[str] = mapped_column( 

72 String(255), 

73 nullable=False, 

74 index=True, 

75 doc="Model name (e.g., claude-sonnet-4-5-20250929)", 

76 ) 

77 provider: Mapped[str] = mapped_column( 

78 String(100), 

79 nullable=False, 

80 doc="Provider name (e.g., anthropic, openai, google)", 

81 ) 

82 

83 # Token counts 

84 prompt_tokens: Mapped[int] = mapped_column( 

85 Integer, 

86 nullable=False, 

87 doc="Number of input tokens", 

88 ) 

89 completion_tokens: Mapped[int] = mapped_column( 

90 Integer, 

91 nullable=False, 

92 doc="Number of output tokens", 

93 ) 

94 total_tokens: Mapped[int] = mapped_column( 

95 Integer, 

96 nullable=False, 

97 doc="Total tokens (prompt + completion)", 

98 ) 

99 

100 # Cost (stored as Decimal for precision) 

101 estimated_cost_usd: Mapped[Decimal] = mapped_column( 

102 Numeric(precision=10, scale=6), 

103 nullable=False, 

104 doc="Estimated cost in USD (6 decimal places for sub-cent precision)", 

105 ) 

106 

107 # Optional categorization 

108 feature: Mapped[str | None] = mapped_column( 

109 String(255), 

110 nullable=True, 

111 doc="Feature name (e.g., 'chat', 'summarization', 'context_compaction')", 

112 ) 

113 

114 # Metadata (stored as JSON) 

115 metadata_: Mapped[dict[str, Any] | None] = mapped_column( 

116 "metadata", # Column name in DB 

117 JSON, 

118 nullable=True, 

119 doc="Additional metadata (JSON)", 

120 ) 

121 

122 # Composite indexes for common queries 

123 __table_args__ = ( 

124 Index("ix_user_timestamp", "user_id", "timestamp"), 

125 Index("ix_provider_model", "provider", "model"), 

126 Index("ix_timestamp_desc", "timestamp", postgresql_ops={"timestamp": "DESC"}), 

127 ) 

128 

129 def __repr__(self) -> str: 

130 return ( 

131 f"<TokenUsageRecord(id={self.id}, " 

132 f"user_id={self.user_id}, " 

133 f"model={self.model}, " 

134 f"tokens={self.total_tokens}, " 

135 f"cost=${self.estimated_cost_usd})>" 

136 ) 

137 

138 def to_dict(self) -> dict[str, Any]: 

139 """Convert record to dictionary.""" 

140 return { 

141 "id": self.id, 

142 "timestamp": self.timestamp.isoformat(), 

143 "created_at": self.created_at.isoformat(), 

144 "user_id": self.user_id, 

145 "session_id": self.session_id, 

146 "model": self.model, 

147 "provider": self.provider, 

148 "prompt_tokens": self.prompt_tokens, 

149 "completion_tokens": self.completion_tokens, 

150 "total_tokens": self.total_tokens, 

151 "estimated_cost_usd": str(self.estimated_cost_usd), 

152 "feature": self.feature, 

153 "metadata": self.metadata_, 

154 }