refactor: Removes tenant ID check from rate limit logic (#20585)

Signed-off-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
-LAN- 2025-06-03 18:56:38 +08:00 committed by GitHub
parent 077d627953
commit 275e86a26c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -86,31 +86,31 @@ class KnowledgeRetrievalNode(LLMNode):
return NodeRunResult( return NodeRunResult(
status=WorkflowNodeExecutionStatus.FAILED, inputs=variables, error="Query is required." status=WorkflowNodeExecutionStatus.FAILED, inputs=variables, error="Query is required."
) )
# TODO(-LAN-): Move this check outside.
# check rate limit # check rate limit
if self.tenant_id: knowledge_rate_limit = FeatureService.get_knowledge_rate_limit(self.tenant_id)
knowledge_rate_limit = FeatureService.get_knowledge_rate_limit(self.tenant_id) if knowledge_rate_limit.enabled:
if knowledge_rate_limit.enabled: current_time = int(time.time() * 1000)
current_time = int(time.time() * 1000) key = f"rate_limit_{self.tenant_id}"
key = f"rate_limit_{self.tenant_id}" redis_client.zadd(key, {current_time: current_time})
redis_client.zadd(key, {current_time: current_time}) redis_client.zremrangebyscore(key, 0, current_time - 60000)
redis_client.zremrangebyscore(key, 0, current_time - 60000) request_count = redis_client.zcard(key)
request_count = redis_client.zcard(key) if request_count > knowledge_rate_limit.limit:
if request_count > knowledge_rate_limit.limit: with Session(db.engine) as session:
with Session(db.engine) as session: # add ratelimit record
# add ratelimit record rate_limit_log = RateLimitLog(
rate_limit_log = RateLimitLog( tenant_id=self.tenant_id,
tenant_id=self.tenant_id, subscription_plan=knowledge_rate_limit.subscription_plan,
subscription_plan=knowledge_rate_limit.subscription_plan, operation="knowledge",
operation="knowledge",
)
session.add(rate_limit_log)
session.commit()
return NodeRunResult(
status=WorkflowNodeExecutionStatus.FAILED,
inputs=variables,
error="Sorry, you have reached the knowledge base request rate limit of your subscription.",
error_type="RateLimitExceeded",
) )
session.add(rate_limit_log)
session.commit()
return NodeRunResult(
status=WorkflowNodeExecutionStatus.FAILED,
inputs=variables,
error="Sorry, you have reached the knowledge base request rate limit of your subscription.",
error_type="RateLimitExceeded",
)
# retrieve knowledge # retrieve knowledge
try: try: