From e1af34a986ec13bfc0ca9206ff2429e08918caaa Mon Sep 17 00:00:00 2001
From: yuetsh <517252939@qq.com>
Date: Mon, 27 Apr 2026 08:57:18 -0600
Subject: [PATCH] add maxkb toggle

---
 database.py | 16 +++++++++++++---
 main.py     | 11 ++++++++++-
 models.py   | 11 +++++++++--
 3 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/database.py b/database.py
index e794895..0972a42 100644
--- a/database.py
+++ b/database.py
@@ -1,5 +1,5 @@
 from tortoise.contrib.fastapi import register_tortoise
-from models import PresetCode
+from models import PresetCode, Setting
 from schemas import PresetCodeCreate, PresetCodeResponse
 from typing import List, Optional
 
@@ -38,10 +38,20 @@ class DatabaseService:
     
     @staticmethod
     async def delete_code(code_id: int) -> bool:
-        """删除指定 ID 的代码"""
         code = await PresetCode.get_or_none(id=code_id)
         if not code:
             return False
-        
         await code.delete()
         return True
+
+    @staticmethod
+    async def get_maxkb_enabled() -> bool:
+        setting, _ = await Setting.get_or_create(key="maxkb_enabled", defaults={"value": True})
+        return setting.value
+
+    @staticmethod
+    async def toggle_maxkb_enabled() -> bool:
+        setting, _ = await Setting.get_or_create(key="maxkb_enabled", defaults={"value": True})
+        setting.value = not setting.value
+        await setting.save()
+        return setting.value
diff --git a/main.py b/main.py
index ec4e2b8..ee27b3a 100644
--- a/main.py
+++ b/main.py
@@ -1,4 +1,3 @@
-from io import StringIO
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse
@@ -36,6 +35,16 @@ DATABASE_URL = "sqlite://database.db"
 DatabaseService.init_database(app, DATABASE_URL)
 
 
+@app.get("/maxkb")
+async def get_maxkb():
+    return {"enabled": await DatabaseService.get_maxkb_enabled()}
+
+
+@app.post("/maxkb")
+async def toggle_maxkb():
+    return {"enabled": await DatabaseService.toggle_maxkb_enabled()}
+
+
 @app.get("/")
 async def get_all_codes() -> dict:
     """获取所有预设代码"""
diff --git a/models.py b/models.py
index 7b1fafd..bfbec96 100644
--- a/models.py
+++ b/models.py
@@ -2,10 +2,17 @@ from tortoise.models import Model
 from tortoise import fields
 
 class PresetCode(Model):
-    """预设代码数据库模型"""
     id = fields.IntField(pk=True)
     query = fields.CharField(max_length=255, unique=True)
     code = fields.TextField()
-    
+
     class Meta:
         table = "preset_codes"
+
+
+class Setting(Model):
+    key = fields.CharField(max_length=100, pk=True)
+    value = fields.BooleanField()
+
+    class Meta:
+        table = "settings"