feat: read quota from settings instead of db

Author: lwj2015

config/settings.py

@@ -282,6 +282,11 @@
 FEISHU_APP_SECRET = env.str("FEISHU_APP_SECRET", default="")
 FEISHU_ENCRYPT_KEY = env.str("FEISHU_ENCRYPT_KEY", default="")
 
+MAX_BOT_COUNT = env.str("MAX_BOT_COUNT", default=10)
+MAX_COLLECTION_COUNT = env.str("MAX_COLLECTION_COUNT", default=50)
+MAX_DOCUMENT_COUNT = env.str("MAX_DOCUMENT_COUNT", default=1000)
+MAX_CONVERSATION_COUNT = env.str("MAX_CONVERSATION_COUNT", default=100)
+
 ENABLE_QA_GENERATOR = env.bool("ENABLE_QA_GENERATOR", default=False)
 
 CHAT_CONSUMER_IMPLEMENTATION = env.str("CHAT_CONSUMER_IMPLEMENTATION", default="document-qa")

deploy/kubechat/templates/django-deployment.yaml

@@ -220,6 +220,26 @@ spec:
                 configMapKeyRef:
                   key: FEISHU_ENCRYPT_KEY
                   name: kubechat-config
+            - name: MAX_BOT_COUNT
+              valueFrom:
+                configMapKeyRef:
+                  key: MAX_BOT_COUNT
+                  name: kubechat-config
+            - name: MAX_COLLECTION_COUNT
+              valueFrom:
+                configMapKeyRef:
+                  key: MAX_COLLECTION_COUNT
+                  name: kubechat-config
+            - name: MAX_DOCUMENT_COUNT
+              valueFrom:
+                configMapKeyRef:
+                  key: MAX_DOCUMENT_COUNT
+                  name: kubechat-config
+            - name: MAX_CONVERSATION_COUNT
+              valueFrom:
+                configMapKeyRef:
+                  key: MAX_CONVERSATION_COUNT
+                  name: kubechat-config
             - name: EMBEDDING_MODEL
               valueFrom:
                 configMapKeyRef:

deploy/kubechat/templates/kubechat-config-configmap.yaml

@@ -23,6 +23,10 @@ data:
   FEISHU_APP_ID: "{{ .Values.django.feishu.appID }}"
   FEISHU_APP_SECRET: "{{ .Values.django.feishu.appSecret }}"
   FEISHU_ENCRYPT_KEY: "{{ .Values.django.feishu.encryptKey }}"
+  MAX_BOT_COUNT: "{{ .Values.django.quota.maxBotCount }}"
+  MAX_COLLECTION_COUNT: "{{ .Values.django.quota.maxCollectionCount }}"
+  MAX_DOCUMENT_COUNT: "{{ .Values.django.quota.maxDocumentCount }}"
+  MAX_CONVERSATION_COUNT: "{{ .Values.django.quota.maxConversationCount }}"
   EMBEDDING_MODEL: "{{ .Values.django.embeddingModel }}"
   EMBEDDING_DEVICE: "{{ .Values.django.embeddingDevice }}"
   ENABLE_QA_GENERATOR: "{{ .Values.django.enableQAGenerator }}"

deploy/kubechat/values.yaml

@@ -90,6 +90,11 @@ django:
     appID: ""
     appSecret: ""
     encryptKey: ""
+  quota:
+    maxBotCount: 10
+    maxCollectionCount: 50
+    maxDocumentCount: 1000
+    maxConversationCount: 100
   resources:
     # We usually recommend not to specify default resources and to leave this as a conscious
     # choice for the user. This also increases chances charts run on environments with little

kubechat/apps.py

@@ -1,35 +1,10 @@
-import sys
-
 from django.apps import AppConfig
-import asyncio
-
-from django.db import transaction
-from asgiref.sync import sync_to_async
 
 
 class KubechatConfig(AppConfig):
     default_auto_field = "django.db.models.BigAutoField"
     name = "kubechat"
 
-    def ready(self):
-        # if manage.py is running, this might be a migration, so we don't need to load quotas
-        if "manage.py" in sys.argv:
-            return
-
-        if asyncio.get_event_loop().is_running():
-            # call load_quotas() asynchronously in uvicorn
-            asyncio.create_task(sync_to_async(load_quotas)())
-        else:
-            # call load_quotas() after the commit transaction in celery
-            transaction.on_commit(load_quotas)
-
-
-class DefaultQuota:
-    MAX_BOT_COUNT = None
-    MAX_COLLECTION_COUNT = None
-    MAX_DOCUMENT_COUNT = None
-    MAX_CONVERSATION_COUNT = None
-
 
 class QuotaType:
     MAX_BOT_COUNT = 'max_bot_count'
@@ -38,10 +13,3 @@ class QuotaType:
     MAX_CONVERSATION_COUNT = 'max_conversation_count'
 
 
-def load_quotas():
-    from kubechat.db.ops import query_quota
-
-    DefaultQuota.MAX_BOT_COUNT = query_quota(QuotaType.MAX_BOT_COUNT)
-    DefaultQuota.MAX_COLLECTION_COUNT = query_quota(QuotaType.MAX_COLLECTION_COUNT)
-    DefaultQuota.MAX_DOCUMENT_COUNT = query_quota(QuotaType.MAX_DOCUMENT_COUNT)
-    DefaultQuota.MAX_CONVERSATION_COUNT = query_quota(QuotaType.MAX_CONVERSATION_COUNT)

kubechat/chat/websocket/base_consumer.py

@@ -12,7 +12,7 @@
 import redis.asyncio as redis
 
 import config.settings as settings
-from kubechat.apps import DefaultQuota, QuotaType
+from kubechat.apps import QuotaType
 from kubechat.auth.validator import DEFAULT_USER
 from kubechat.chat.utils import start_response, success_response, stop_response, fail_response, welcome_response
 from kubechat.pipeline.pipeline import KUBE_CHAT_DOC_QA_REFERENCES, KeywordPipeline,KUBE_CHAT_RELATED_QUESTIONS
@@ -58,7 +58,7 @@ async def connect(self):
         self.redis_client = redis.Redis.from_url(settings.MEMORY_REDIS_URL)
         self.conversation_limit = await query_user_quota(self.user, QuotaType.MAX_CONVERSATION_COUNT)
         if self.conversation_limit is None:
-            self.conversation_limit = DefaultQuota.MAX_CONVERSATION_COUNT
+            self.conversation_limit = settings.MAX_CONVERSATION_COUNT
 
         # If using daphne, we need to put the token in the headers and include the headers in the subprotocol.
         # headers = {}

kubechat/tasks/sync_documents_task.py

@@ -13,7 +13,7 @@
 from pydantic import BaseModel
 
 from config.celery import app
-from kubechat.apps import DefaultQuota
+from config.settings import MAX_DOCUMENT_COUNT
 from kubechat.db.models import Document, DocumentStatus, Collection, CollectionStatus, CollectionSyncHistory, \
     CollectionSyncStatus
 from kubechat.source.base import get_source
@@ -63,7 +63,7 @@ def sync_documents(self, **kwargs):
 
     document_limit = kwargs["document_user_quota"]
     if document_limit is None:
-        document_limit = DefaultQuota.MAX_DOCUMENT_COUNT
+        document_limit = MAX_DOCUMENT_COUNT
 
     src_docs = {}
     tasks = []

kubechat/views/main.py

@@ -16,7 +16,7 @@
 import kubechat.chat.message
 from config import settings
 from config.celery import app
-from kubechat.apps import DefaultQuota, QuotaType
+from kubechat.apps import QuotaType
 from kubechat.chat.history.redis import RedisChatMessageHistory
 from kubechat.llm.base import Predictor
 from kubechat.llm.prompts import DEFAULT_MODEL_PROMPT_TEMPLATES, DEFAULT_CHINESE_PROMPT_TEMPLATE_V2, \
@@ -225,10 +225,10 @@ async def create_collection(request, collection: CollectionIn):
             return fail(HTTPStatus.BAD_REQUEST, error_msg)
 
     # there is quota limit on collection
-    if DefaultQuota.MAX_COLLECTION_COUNT:
+    if settings.MAX_COLLECTION_COUNT:
         collection_limit = await query_user_quota(user, QuotaType.MAX_COLLECTION_COUNT)
         if collection_limit is None:
-            collection_limit = DefaultQuota.MAX_COLLECTION_COUNT
+            collection_limit = settings.MAX_COLLECTION_COUNT
         if collection_limit and await query_collections_count(user) >= collection_limit:
             return fail(HTTPStatus.FORBIDDEN, f"collection number has reached the limit of {collection_limit}")
 
@@ -353,10 +353,10 @@ async def create_document(request, collection_id, file: List[UploadedFile] = Fil
         return fail(HTTPStatus.NOT_FOUND, "Collection not found")
 
     # there is quota limit on document
-    if DefaultQuota.MAX_DOCUMENT_COUNT:
+    if settings.MAX_DOCUMENT_COUNT:
         document_limit = await query_user_quota(user, QuotaType.MAX_DOCUMENT_COUNT)
         if document_limit is None:
-            document_limit = DefaultQuota.MAX_DOCUMENT_COUNT
+            document_limit = settings.MAX_DOCUMENT_COUNT
         if await query_documents_count(user, collection_id) >= document_limit:
             return fail(HTTPStatus.FORBIDDEN, f"document number has reached the limit of {document_limit}")
 
@@ -399,10 +399,10 @@ async def create_url_document(request, collection_id):
         return fail(HTTPStatus.NOT_FOUND, "Collection not found")
 
     # there is quota limit on document
-    if DefaultQuota.MAX_DOCUMENT_COUNT:
+    if settings.MAX_DOCUMENT_COUNT:
         document_limit = await query_user_quota(user, QuotaType.MAX_DOCUMENT_COUNT)
         if document_limit is None:
-            document_limit = DefaultQuota.MAX_DOCUMENT_COUNT
+            document_limit = settings.MAX_DOCUMENT_COUNT
         if await query_documents_count(user, collection_id) >= document_limit:
             return fail(HTTPStatus.FORBIDDEN, f"document number has reached the limit of {document_limit}")
 
@@ -603,10 +603,10 @@ async def create_bot(request, bot_in: BotIn):
     user = get_user(request)
 
     # there is quota limit on bot
-    if DefaultQuota.MAX_BOT_COUNT:
+    if settings.MAX_BOT_COUNT:
         bot_limit = await query_user_quota(user, QuotaType.MAX_BOT_COUNT)
         if bot_limit is None:
-            bot_limit = DefaultQuota.MAX_BOT_COUNT
+            bot_limit = settings.MAX_BOT_COUNT
         if await query_bots_count(user) >= bot_limit:
             return fail(HTTPStatus.FORBIDDEN, f"bot number has reached the limit of {bot_limit}")
 

kubechat/views/weixin.py

@@ -7,7 +7,8 @@
 from config import settings
 import kubechat.chat.message
 
-from kubechat.apps import QuotaType, DefaultQuota
+from kubechat.apps import QuotaType
+from config.settings import MAX_CONVERSATION_COUNT
 from kubechat.chat.history.redis import RedisChatMessageHistory
 from kubechat.pipeline.pipeline import KeywordPipeline
 from kubechat.db.ops import *
@@ -62,7 +63,7 @@ async def weixin_text_response(client, user, bot, query, msg_id):
 
     conversation_limit = await query_user_quota(user, QuotaType.MAX_CONVERSATION_COUNT)
     if conversation_limit is None:
-        conversation_limit = DefaultQuota.MAX_CONVERSATION_COUNT
+        conversation_limit = MAX_CONVERSATION_COUNT
 
     try:
         if use_default_token and conversation_limit:
@@ -103,7 +104,7 @@ async def weixin_card_response(client, user, bot, query, msg_id):
 
     conversation_limit = await query_user_quota(user, QuotaType.MAX_CONVERSATION_COUNT)
     if conversation_limit is None:
-        conversation_limit = DefaultQuota.MAX_CONVERSATION_COUNT
+        conversation_limit = MAX_CONVERSATION_COUNT
 
     try:
         if use_default_token and conversation_limit: