@@ -86,31 +86,31 @@ def _run(self) -> NodeRunResult: # type: ignore
86
86
return NodeRunResult (
87
87
status = WorkflowNodeExecutionStatus .FAILED , inputs = variables , error = "Query is required."
88
88
)
89
+ # TODO(-LAN-): Move this check outside.
89
90
# 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" ,
113
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" ,
113
+ )
114
114
115
115
# retrieve knowledge
116
116
try :
0 commit comments