#ai 代码优化

zhou zhou

15 小时以前 82624affb0251b75b62b35567d3eb260c06efe78

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
package com.vincent.rsf.server.ai.store;
 
import com.vincent.rsf.server.ai.config.AiDefaults;
import com.vincent.rsf.server.ai.store.support.AiRedisExecutor;
import com.vincent.rsf.server.ai.store.support.AiRedisKeys;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Component;
 
import java.time.Instant;
import java.util.UUID;
 
@Component
@RequiredArgsConstructor
public class AiChatRateLimiter {
 
    private final AiRedisExecutor aiRedisExecutor;
    private final AiRedisKeys aiRedisKeys;
 
    public boolean allowChatRequest(Long tenantId, Long userId, String promptCode) {
        String key = aiRedisKeys.buildRateLimitKey(tenantId, userId, promptCode);
        long now = Instant.now().toEpochMilli();
        long windowStart = now - (AiDefaults.CHAT_RATE_LIMIT_WINDOW_SECONDS * 1000L);
        Boolean allowed = aiRedisExecutor.execute(jedis -> {
            jedis.zremrangeByScore(key, 0, windowStart);
            long count = jedis.zcard(key);
            if (count >= AiDefaults.CHAT_RATE_LIMIT_MAX_REQUESTS) {
                jedis.expire(key, AiDefaults.CHAT_RATE_LIMIT_WINDOW_SECONDS);
                return Boolean.FALSE;
            }
            jedis.zadd(key, now, now + ":" + UUID.randomUUID());
            jedis.expire(key, AiDefaults.CHAT_RATE_LIMIT_WINDOW_SECONDS);
            return Boolean.TRUE;
        });
        return Boolean.TRUE.equals(allowed);
    }
}