From 2e9c41347b0179cad5054a34001fa701e7c6e1fb Mon Sep 17 00:00:00 2001 From: EdgeNeko Date: Mon, 20 May 2024 19:14:30 +0800 Subject: [PATCH 1/3] Allow user to configure max index queue length --- app/Controllers/admin.py | 2 +- app/Services/upload_service.py | 3 ++- app/config.py | 1 + config/default.env | 2 ++ tests/api/integrate_test.py | 1 + 5 files changed, 7 insertions(+), 2 deletions(-) diff --git a/app/Controllers/admin.py b/app/Controllers/admin.py index 21cdf26..e0c5566 100644 --- a/app/Controllers/admin.py +++ b/app/Controllers/admin.py @@ -108,7 +108,7 @@ async def upload_image(image_file: Annotated[UploadFile, File(description="The i try: image = Image.open(BytesIO(img_bytes)) except UnidentifiedImageError as ex: - raise HTTPException(400, "Cannot open the image file.") from ex + raise HTTPException(422, "Cannot open the image file.") from ex image_data = ImageData(id=img_id, url=model.url, diff --git a/app/Services/upload_service.py b/app/Services/upload_service.py index 91accaf..baa025d 100644 --- a/app/Services/upload_service.py +++ b/app/Services/upload_service.py @@ -9,6 +9,7 @@ from app.Services.index_service import IndexService from app.Services.storage import StorageService from app.Services.vector_db_context import VectorDbContext +from app.config import config class UploadService: @@ -17,7 +18,7 @@ def __init__(self, storage_service: StorageService, db_context: VectorDbContext, self._db_context = db_context self._index_service = index_service - self._queue = asyncio.Queue(200) + self._queue = asyncio.Queue(config.admin_index_queue_max_length) self._upload_worker_task = asyncio.create_task(self._upload_worker()) self._processed_count = 0 diff --git a/app/config.py b/app/config.py index bbe71f3..4c8afdf 100644 --- a/app/config.py +++ b/app/config.py @@ -87,6 +87,7 @@ class Config(BaseSettings): cors_origins: set[str] = {'*'} admin_api_enable: bool = False admin_token: str = '' + admin_index_queue_max_length: int = 200 access_protected: bool = False access_token: str = '' diff --git a/config/default.env b/config/default.env index caaf218..ca3a940 100644 --- a/config/default.env +++ b/config/default.env @@ -71,6 +71,8 @@ # APP_ADMIN_API_ENABLE=False # Uncomment the line below if you enabled admin API. Use this token to access admin API. For security reasons, the admin token is always required if you want to use admin API. # APP_ADMIN_TOKEN="your-super-secret-admin-token" +# Max length of the upload queue for admin API, higher value means more indexing requests can be queued but also means more memory usage. Upload requests will be blocked when the queue is full. +# APP_ADMIN_INDEX_QUEUE_MAX_LENGTH=200 # ------ diff --git a/tests/api/integrate_test.py b/tests/api/integrate_test.py index 562b20d..aaa9e76 100644 --- a/tests/api/integrate_test.py +++ b/tests/api/integrate_test.py @@ -38,6 +38,7 @@ async def test_integrate(test_client): break await asyncio.sleep(1) + assert resp.json()['index_queue_length'] == 0 resp = test_client.get('/search/text/hatsune+miku', headers=credentials) assert resp.status_code == 200 From 0b07ef819f538731da892b14a47155722d66afc7 Mon Sep 17 00:00:00 2001 From: EdgeNeko Date: Mon, 20 May 2024 19:14:56 +0800 Subject: [PATCH 2/3] Add negative testcases for upload api --- tests/api/test_upload.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 tests/api/test_upload.py diff --git a/tests/api/test_upload.py b/tests/api/test_upload.py new file mode 100644 index 0000000..a6fae14 --- /dev/null +++ b/tests/api/test_upload.py @@ -0,0 +1,27 @@ +import io +import random + +from .conftest import TEST_ADMIN_TOKEN + + +def test_upload_bad_img_file(test_client): + bad_img_file = io.BytesIO(bytearray(random.getrandbits(8) for _ in range(1024 * 1024))) + bad_img_file.name = 'bad_image.jpg' + + resp = test_client.post('/admin/upload', + files={'image_file': bad_img_file}, + headers={'x-admin-token': TEST_ADMIN_TOKEN}, + params={'local': True}) + print(resp.content) + assert resp.status_code == 422 + + +def test_upload_unsupported_types(test_client): + bad_img_file = io.BytesIO(bytearray(random.getrandbits(8) for _ in range(1024 * 1024))) + bad_img_file.name = 'bad_image.tga' + + resp = test_client.post('/admin/upload', + files={'image_file': ('bad_img.tga', bad_img_file, 'image/tga')}, + headers={'x-admin-token': TEST_ADMIN_TOKEN}, + params={'local': True}) + assert resp.status_code == 415 From c5bf1da2cbd92f255eedf2bf13d1c8774dacc932 Mon Sep 17 00:00:00 2001 From: EdgeNeko Date: Mon, 20 May 2024 19:18:49 +0800 Subject: [PATCH 3/3] Update docs to match recent changes --- readme.md | 6 ++++-- readme_cn.md | 4 ++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/readme.md b/readme.md index 684b122..3126068 100644 --- a/readme.md +++ b/readme.md @@ -63,7 +63,7 @@ Local file storage does not require an additional database deployment process, b #### Deploy NekoImageGallery -1. Clone the project directory to your own PC or server. +1. Clone the project directory to your own PC or server, then checkout to a specific version tag (like `v1.0.0`). 2. It is highly recommended to install the dependencies required for this project in a Python venv virtual environment. Run the following command: ```shell @@ -80,7 +80,9 @@ Local file storage does not require an additional database deployment process, b ``` 5. Modify the project configuration file inside `config/`, you can edit `default.env` directly, but it's recommended to create a new file named `local.env` and override the configuration in `default.env`. -6. Initialize the Qdrant database by running the following command: +6. ~~Initialize the Qdrant database by running the following command:~~ + Since NekoImageGallery will now automatically create the collection if not present, this steps is no longer required. + However, if you want to explicitly create the collection, you can still run the following command: ```shell python main.py --init-database ``` diff --git a/readme_cn.md b/readme_cn.md index 99380ec..548f2db 100644 --- a/readme_cn.md +++ b/readme_cn.md @@ -55,7 +55,7 @@ NekoImageGallery支持两种元数据存储方式：Qdrant数据库存储与本 #### 部署NekoImageGallery -1. 将项目目录clone到你自己的PC或服务器中。 +1. 将项目目录clone到你自己的PC或服务器中，然后按需checkout到特定版本tag（如`v1.0.0`）。 2. 强烈建议在python venv虚拟环境中安装本项目所需依赖， 运行下面命令： ```shell python -m venv .venv @@ -70,7 +70,7 @@ NekoImageGallery支持两种元数据存储方式：Qdrant数据库存储与本 ``` 5. 按需修改位于`config`目录下的配置文件，您可以直接修改`default.env`，但是建议创建一个名为`local.env` 的文件，覆盖`default.env`中的配置。 -6. 初始化Qdrant数据库，运行下面命令： +6. ~~初始化Qdrant数据库，运行下面命令：~~ 由于NekoImageGallery会自动在第一次运行时初始化Qdrant数据库，因此此步骤可省略。但您仍可通过以下命令显式初始化Qdrant数据库： ```shell python main.py --init-database ```