From e39c813c15dbdab1d916828c5835ad085ce06f12 Mon Sep 17 00:00:00 2001
From: shaohuzhang1 <shaohu.zhang@fit2cloud.com>
Date: Thu, 14 Mar 2024 12:26:42 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=AF=8F=E6=97=A5=E5=87=8C=E6=99=A8?=
 =?UTF-8?q?=E9=87=8D=E7=BD=AE=E5=AE=A2=E6=88=B7=E8=AE=BF=E9=97=AE=E6=95=B0?=
 =?UTF-8?q?=E9=87=8F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/common/job/__init__.py              | 13 ++++++++++
 apps/common/job/client_access_num_job.py | 33 ++++++++++++++++++++++++
 apps/smartdoc/settings/base.py           |  1 +
 apps/smartdoc/wsgi.py                    |  2 ++
 pyproject.toml                           |  1 +
 5 files changed, 50 insertions(+)
 create mode 100644 apps/common/job/__init__.py
 create mode 100644 apps/common/job/client_access_num_job.py

diff --git a/apps/common/job/__init__.py b/apps/common/job/__init__.py
new file mode 100644
index 000000000..895bf7f5d
--- /dev/null
+++ b/apps/common/job/__init__.py
@@ -0,0 +1,13 @@
+# coding=utf-8
+"""
+    @project: maxkb
+    @Author：虎
+    @file： __init__.py
+    @date：2024/3/14 11:54
+    @desc:
+"""
+from .client_access_num_job import *
+
+
+def run():
+    client_access_num_job.run()
diff --git a/apps/common/job/client_access_num_job.py b/apps/common/job/client_access_num_job.py
new file mode 100644
index 000000000..4c03fd210
--- /dev/null
+++ b/apps/common/job/client_access_num_job.py
@@ -0,0 +1,33 @@
+# coding=utf-8
+"""
+    @project: maxkb
+    @Author：虎
+    @file： client_access_num_job.py
+    @date：2024/3/14 11:56
+    @desc:
+"""
+import logging
+
+from apscheduler.schedulers.background import BackgroundScheduler
+from django.db.models import QuerySet
+from django_apscheduler.jobstores import DjangoJobStore
+
+from application.models.api_key_model import ApplicationPublicAccessClient
+
+scheduler = BackgroundScheduler()
+scheduler.add_jobstore(DjangoJobStore(), "default")
+
+
+def client_access_num_reset_job():
+    logging.getLogger("max_kb").info('开始重置access_num')
+    QuerySet(ApplicationPublicAccessClient).update(intraday_access_num=0)
+    logging.getLogger("max_kb").info('结束重置access_num')
+
+
+def run():
+    scheduler.start()
+    access_num_reset = scheduler.get_job(job_id='access_num_reset')
+    if access_num_reset is not None:
+        access_num_reset.remove()
+    scheduler.add_job(client_access_num_reset_job, 'cron', hour='0', minute='0', second='0',
+                      id='access_num_reset')
diff --git a/apps/smartdoc/settings/base.py b/apps/smartdoc/settings/base.py
index 4b22d0702..621b0135e 100644
--- a/apps/smartdoc/settings/base.py
+++ b/apps/smartdoc/settings/base.py
@@ -38,6 +38,7 @@ INSTALLED_APPS = [
     'rest_framework',
     "drf_yasg",  # swagger 接口
     'django_filters',  # 条件过滤
+    'django_apscheduler'
 
 ]
 
diff --git a/apps/smartdoc/wsgi.py b/apps/smartdoc/wsgi.py
index df250f58e..e04ec5274 100644
--- a/apps/smartdoc/wsgi.py
+++ b/apps/smartdoc/wsgi.py
@@ -18,8 +18,10 @@ application = get_wsgi_application()
 
 def post_handler():
     from common import event
+    from common import job
     event.run()
     event.ListenerManagement.init_embedding_model_signal.send()
+    job.run()
 
 
 post_handler()
diff --git a/pyproject.toml b/pyproject.toml
index ca72630cc..cf26ea8fc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -29,6 +29,7 @@ beautifulsoup4 = "^4.12.2"
 html2text = "^2024.2.26"
 langchain-openai = "^0.0.8"
 django-ipware = "^6.0.4"
+django-apscheduler = "^0.6.2"
 
 [build-system]
 requires = ["poetry-core"]