feat:新增投顾agent和nl2sqlagent

nl2sql/limits.py

@@ -0,0 +1,58 @@
+"""NL2SQL 查询配额、并发和频率限制。"""
+from __future__ import annotations
+
+import time
+
+
+class QueryLimiter:
+    """基于 Redis 计数器的请求限制器，Redis 故障时允许降级放行。"""
+
+    def __init__(self, redis, *, rate_window_seconds: int = 60):
+        self.redis = redis
+        self.rate_window_seconds = rate_window_seconds
+
+    async def acquire(
+        self,
+        user_id: int,
+        *,
+        daily_quota: int = 0,
+        max_concurrent: int = 1,
+        rate_limit: int = 0,
+    ) -> bool:
+        """检查并占用一次配额、并发和频率额度，0 表示不限制。"""
+        try:
+            if daily_quota:
+                daily_key = f"nl2sql:daily:{user_id}:{time.strftime('%Y%m%d')}"
+                daily_count = await self.redis.incr(daily_key)
+                await self.redis.expire(daily_key, 86400)
+                if daily_count > daily_quota:
+                    await self.redis.decr(daily_key)
+                    return False
+            concurrent_key = f"nl2sql:concurrent:{user_id}"
+            concurrent_count = await self.redis.incr(concurrent_key)
+            await self.redis.expire(concurrent_key, 3600)
+            if max_concurrent and concurrent_count > max_concurrent:
+                await self.redis.decr(concurrent_key)
+                if daily_quota:
+                    await self.redis.decr(daily_key)
+                return False
+            if rate_limit:
+                rate_key = f"nl2sql:rate:{user_id}:{int(time.time()) // self.rate_window_seconds}"
+                rate_count = await self.redis.incr(rate_key)
+                await self.redis.expire(rate_key, self.rate_window_seconds + 1)
+                if rate_count > rate_limit:
+                    await self.redis.decr(rate_key)
+                    await self.redis.decr(concurrent_key)
+                    if daily_quota:
+                        await self.redis.decr(daily_key)
+                    return False
+            return True
+        except Exception:  # noqa: BLE001 Redis 故障时降级为不限制
+            return True
+
+    async def release(self, user_id: int) -> None:
+        """释放一次并发额度，失败时静默降级。"""
+        try:
+            await self.redis.decr(f"nl2sql:concurrent:{user_id}")
+        except Exception:  # noqa: BLE001 Redis 故障不能影响查询结果
+            return