42a9a17133
- Add cache::ai module with Redis rate limiting for AI generations
- Add functions: check_ai_rate_limit, get_ai_usage, cache_ai_response,
get_cached_ai_response, invalidate_ai_cache, reset_daily_usage
- Update check_and_increment_usage to use Redis fast-path before DB
- Redis key pattern: ai:rate:{user_id} for 24hr sliding window counter
80 lines
2.4 KiB
Rust
80 lines
2.4 KiB
Rust
//! Redis caching for AI generation rate limiting and response caching.
|
|
//!
|
|
//! Key patterns:
|
|
//! - `ai:rate:{user_id}` - sliding window counter for rate limiting
|
|
//! - `ai:resp:{hash}` - cached AI response (by prompt hash)
|
|
|
|
use redis::AsyncCommands;
|
|
use crate::RedisPool;
|
|
|
|
const AI_RATE_WINDOW_SECS: i64 = 86_400; // 24 hours
|
|
const AI_CACHE_TTL_SECS: i64 = 3_600; // 1 hour
|
|
|
|
/// Check + increment AI generation rate limit counter.
|
|
/// Uses a simple counter with TTL reset on first write.
|
|
///
|
|
/// Returns `Ok(true)` if allowed, `Ok(false)` if rate limited.
|
|
pub async fn check_ai_rate_limit(
|
|
redis: &mut RedisPool,
|
|
user_id: &str,
|
|
max_generations: i64,
|
|
) -> Result<bool, redis::RedisError> {
|
|
let key = format!("ai:rate:{}", user_id);
|
|
let count: i64 = redis.incr(&key, 1i64).await?;
|
|
if count == 1 {
|
|
redis.expire::<_, ()>(&key, AI_RATE_WINDOW_SECS).await?;
|
|
}
|
|
Ok(count <= max_generations)
|
|
}
|
|
|
|
/// Get current AI generation count for a user.
|
|
pub async fn get_ai_usage(
|
|
redis: &mut RedisPool,
|
|
user_id: &str,
|
|
) -> Result<i64, redis::RedisError> {
|
|
let key = format!("ai:rate:{}", user_id);
|
|
let count: Option<i64> = redis.get(&key).await?;
|
|
Ok(count.unwrap_or(0))
|
|
}
|
|
|
|
/// Store AI-generated response in cache.
|
|
pub async fn cache_ai_response(
|
|
redis: &mut RedisPool,
|
|
prompt_hash: &str,
|
|
response: &str,
|
|
) -> Result<(), redis::RedisError> {
|
|
let key = format!("ai:resp:{}", prompt_hash);
|
|
let ttl: u64 = AI_CACHE_TTL_SECS.try_into().unwrap();
|
|
let _: () = redis.set_ex(&key, response, ttl).await?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Get cached AI response if available.
|
|
pub async fn get_cached_ai_response(
|
|
redis: &mut RedisPool,
|
|
prompt_hash: &str,
|
|
) -> Result<Option<String>, redis::RedisError> {
|
|
let key = format!("ai:resp:{}", prompt_hash);
|
|
let result: Option<String> = redis.get(&key).await?;
|
|
Ok(result)
|
|
}
|
|
|
|
/// Invalidate cached AI response.
|
|
pub async fn invalidate_ai_cache(
|
|
redis: &mut RedisPool,
|
|
prompt_hash: &str,
|
|
) -> Result<(), redis::RedisError> {
|
|
let key = format!("ai:resp:{}", prompt_hash);
|
|
let _: () = redis.del(&key).await?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Reset daily AI usage counter (called at start of new day or when daily limit changes).
|
|
pub async fn reset_daily_usage(
|
|
redis: &mut RedisPool,
|
|
user_id: &str,
|
|
) -> Result<(), redis::RedisError> {
|
|
let key = format!("ai:rate:{}", user_id);
|
|
let _: () = redis.del(&key).await?;
|
|
Ok(())
|
|
}
|