mirror of
https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git
synced 2025-08-17 05:45:53 +08:00
refactor: Removes tenant ID check from rate limit logic (#20585)
Signed-off-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
parent
077d627953
commit
275e86a26c
@ -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:
|
||||||
|
Loading…
x
Reference in New Issue
Block a user