Skip to content

Commit 275e86a

Browse files
authored
refactor: Removes tenant ID check from rate limit logic (#20585)
Signed-off-by: -LAN- <[email protected]>
1 parent 077d627 commit 275e86a

File tree

1 file changed

+23
-23
lines changed

1 file changed

+23
-23
lines changed

api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py

Lines changed: 23 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -86,31 +86,31 @@ def _run(self) -> NodeRunResult: # type: ignore
8686
return NodeRunResult(
8787
status=WorkflowNodeExecutionStatus.FAILED, inputs=variables, error="Query is required."
8888
)
89+
# TODO(-LAN-): Move this check outside.
8990
# check rate limit
90-
if self.tenant_id:
91-
knowledge_rate_limit = FeatureService.get_knowledge_rate_limit(self.tenant_id)
92-
if knowledge_rate_limit.enabled:
93-
current_time = int(time.time() * 1000)
94-
key = f"rate_limit_{self.tenant_id}"
95-
redis_client.zadd(key, {current_time: current_time})
96-
redis_client.zremrangebyscore(key, 0, current_time - 60000)
97-
request_count = redis_client.zcard(key)
98-
if request_count > knowledge_rate_limit.limit:
99-
with Session(db.engine) as session:
100-
# add ratelimit record
101-
rate_limit_log = RateLimitLog(
102-
tenant_id=self.tenant_id,
103-
subscription_plan=knowledge_rate_limit.subscription_plan,
104-
operation="knowledge",
105-
)
106-
session.add(rate_limit_log)
107-
session.commit()
108-
return NodeRunResult(
109-
status=WorkflowNodeExecutionStatus.FAILED,
110-
inputs=variables,
111-
error="Sorry, you have reached the knowledge base request rate limit of your subscription.",
112-
error_type="RateLimitExceeded",
91+
knowledge_rate_limit = FeatureService.get_knowledge_rate_limit(self.tenant_id)
92+
if knowledge_rate_limit.enabled:
93+
current_time = int(time.time() * 1000)
94+
key = f"rate_limit_{self.tenant_id}"
95+
redis_client.zadd(key, {current_time: current_time})
96+
redis_client.zremrangebyscore(key, 0, current_time - 60000)
97+
request_count = redis_client.zcard(key)
98+
if request_count > knowledge_rate_limit.limit:
99+
with Session(db.engine) as session:
100+
# add ratelimit record
101+
rate_limit_log = RateLimitLog(
102+
tenant_id=self.tenant_id,
103+
subscription_plan=knowledge_rate_limit.subscription_plan,
104+
operation="knowledge",
113105
)
106+
session.add(rate_limit_log)
107+
session.commit()
108+
return NodeRunResult(
109+
status=WorkflowNodeExecutionStatus.FAILED,
110+
inputs=variables,
111+
error="Sorry, you have reached the knowledge base request rate limit of your subscription.",
112+
error_type="RateLimitExceeded",
113+
)
114114

115115
# retrieve knowledge
116116
try:

0 commit comments

Comments
 (0)