Spaces:
Running
Running
File size: 4,924 Bytes
dc3879e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 |
"""Rate limiting service for chat API.
[Task]: T021
[From]: specs/004-ai-chatbot/tasks.md
This service enforces the 100 messages/day limit per user (NFR-011).
"""
import uuid
from datetime import datetime, timedelta
from typing import Optional
from sqlmodel import Session, select
from sqlalchemy import func
from models.message import Message
# Rate limit constants
DAILY_MESSAGE_LIMIT = 100 # NFR-011: Maximum messages per user per day
def check_rate_limit(
db: Session,
user_id: uuid.UUID
) -> tuple[bool, int, Optional[datetime]]:
"""Check if user has exceeded their daily message limit.
[From]: specs/004-ai-chatbot/spec.md - NFR-011
Args:
db: Database session (synchronous)
user_id: User ID to check
Returns:
Tuple of (allowed, remaining_count, reset_time)
- allowed: True if user can send message, False if limit exceeded
- remaining_count: Number of messages remaining today
- reset_time: When the limit resets (midnight UTC), or None if allowed
Example:
>>> allowed, remaining, reset = await check_rate_limit(db, user_id)
>>> if not allowed:
... print(f"Rate limited. Resets at {reset}")
"""
# Calculate today's date range (UTC)
now = datetime.utcnow()
today_start = now.replace(hour=0, minute=0, second=0, microsecond=0)
today_end = today_start + timedelta(days=1)
# Count messages sent by user today
# [From]: specs/004-ai-chatbot/spec.md - NFR-011
# Count both user and assistant messages (all messages in conversation)
statement = select(func.count(Message.id)).where(
Message.user_id == user_id,
Message.created_at >= today_start,
Message.created_at < today_end
)
message_count = db.exec(statement).one() or 0
# Calculate remaining messages
remaining = DAILY_MESSAGE_LIMIT - message_count
if remaining <= 0:
# Rate limit exceeded
return False, 0, today_end
else:
# User can send message
return True, remaining - 1, None
def record_message(
db: Session,
user_id: uuid.UUID,
conversation_id: uuid.UUID,
role: str,
content: str
) -> Message:
"""Record a message in the database (for rate limit tracking).
[From]: specs/004-ai-chatbot/plan.md - Message Persistence
Note: This function is primarily for rate limit tracking.
The actual message persistence should happen in the chat API endpoint
before AI processing (T017) and after AI response (T018).
Args:
db: Database session
user_id: User ID who sent the message
conversation_id: Conversation ID
role: Message role ("user" or "assistant")
content: Message content
Returns:
Created message object
"""
message = Message(
id=uuid.uuid4(),
conversation_id=conversation_id,
user_id=user_id,
role=role,
content=content,
created_at=datetime.utcnow()
)
db.add(message)
db.commit()
db.refresh(message)
return message
def get_message_count_today(
db: Session,
user_id: uuid.UUID
) -> int:
"""Get the number of messages sent by user today.
[From]: specs/004-ai-chatbot/spec.md - NFR-011
Args:
db: Database session
user_id: User ID to check
Returns:
Number of messages sent today (both user and assistant)
"""
now = datetime.utcnow()
today_start = now.replace(hour=0, minute=0, second=0, microsecond=0)
today_end = today_start + timedelta(days=1)
statement = select(func.count(Message.id)).where(
Message.user_id == user_id,
Message.created_at >= today_start,
Message.created_at < today_end
)
return db.exec(statement).one() or 0
def get_rate_limit_status(
db: Session,
user_id: uuid.UUID
) -> dict:
"""Get comprehensive rate limit status for a user.
[From]: specs/004-ai-chatbot/spec.md - NFR-011
Args:
db: Database session
user_id: User ID to check
Returns:
Dictionary with rate limit information:
{
"limit": 100,
"used": 45,
"remaining": 55,
"resets_at": "2025-01-16T00:00:00Z"
}
"""
now = datetime.utcnow()
today_start = now.replace(hour=0, minute=0, second=0, microsecond=0)
today_end = today_start + timedelta(days=1)
# Count messages sent today
statement = select(func.count(Message.id)).where(
Message.user_id == user_id,
Message.created_at >= today_start,
Message.created_at < today_end
)
message_count = db.exec(statement).one() or 0
remaining = max(0, DAILY_MESSAGE_LIMIT - message_count)
return {
"limit": DAILY_MESSAGE_LIMIT,
"used": message_count,
"remaining": remaining,
"resets_at": today_end.isoformat() + "Z"
}
|