|
|
|
|
@ -86,31 +86,31 @@ class KnowledgeRetrievalNode(LLMNode):
|
|
|
|
|
return NodeRunResult(
|
|
|
|
|
status=WorkflowNodeExecutionStatus.FAILED, inputs=variables, error="Query is required."
|
|
|
|
|
)
|
|
|
|
|
# TODO(-LAN-): Move this check outside.
|
|
|
|
|
# check rate limit
|
|
|
|
|
if self.tenant_id:
|
|
|
|
|
knowledge_rate_limit = FeatureService.get_knowledge_rate_limit(self.tenant_id)
|
|
|
|
|
if knowledge_rate_limit.enabled:
|
|
|
|
|
current_time = int(time.time() * 1000)
|
|
|
|
|
key = f"rate_limit_{self.tenant_id}"
|
|
|
|
|
redis_client.zadd(key, {current_time: current_time})
|
|
|
|
|
redis_client.zremrangebyscore(key, 0, current_time - 60000)
|
|
|
|
|
request_count = redis_client.zcard(key)
|
|
|
|
|
if request_count > knowledge_rate_limit.limit:
|
|
|
|
|
with Session(db.engine) as session:
|
|
|
|
|
# add ratelimit record
|
|
|
|
|
rate_limit_log = RateLimitLog(
|
|
|
|
|
tenant_id=self.tenant_id,
|
|
|
|
|
subscription_plan=knowledge_rate_limit.subscription_plan,
|
|
|
|
|
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",
|
|
|
|
|
knowledge_rate_limit = FeatureService.get_knowledge_rate_limit(self.tenant_id)
|
|
|
|
|
if knowledge_rate_limit.enabled:
|
|
|
|
|
current_time = int(time.time() * 1000)
|
|
|
|
|
key = f"rate_limit_{self.tenant_id}"
|
|
|
|
|
redis_client.zadd(key, {current_time: current_time})
|
|
|
|
|
redis_client.zremrangebyscore(key, 0, current_time - 60000)
|
|
|
|
|
request_count = redis_client.zcard(key)
|
|
|
|
|
if request_count > knowledge_rate_limit.limit:
|
|
|
|
|
with Session(db.engine) as session:
|
|
|
|
|
# add ratelimit record
|
|
|
|
|
rate_limit_log = RateLimitLog(
|
|
|
|
|
tenant_id=self.tenant_id,
|
|
|
|
|
subscription_plan=knowledge_rate_limit.subscription_plan,
|
|
|
|
|
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",
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
# retrieve knowledge
|
|
|
|
|
try:
|
|
|
|
|
|