commit 050402f2b75336ea9ee54516f908294010565e79 Author: 0nikod Date: Sat Jan 17 18:47:07 2026 +0800 npm run build diff --git a/.agent/rules/llms.md b/.agent/rules/llms.md new file mode 100644 index 0000000..6f8a556 --- /dev/null +++ b/.agent/rules/llms.md @@ -0,0 +1,8 @@ +--- +trigger: always_on +--- + +Always read official documents before perform tasks. + +- +- \ No newline at end of file diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000..d4a2c44 --- /dev/null +++ b/.editorconfig @@ -0,0 +1,21 @@ +# http://editorconfig.org + +root = true + +[*] +indent_style = space +indent_size = 4 +trim_trailing_whitespace = true +insert_final_newline = true +charset = utf-8 +end_of_line = lf + +[*.bat] +indent_style = tab +end_of_line = crlf + +[LICENSE] +insert_final_newline = false + +[Makefile] +indent_style = tab diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..eb5d0d3 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +.gitignore merge=ours \ No newline at end of file diff --git a/.github/workflows/publish_action.yml b/.github/workflows/publish_action.yml new file mode 100644 index 0000000..122a1ae --- /dev/null +++ b/.github/workflows/publish_action.yml @@ -0,0 +1,20 @@ +name: Publish to Comfy registry +on: + workflow_dispatch: + push: + branches: + - main + paths: + - "pyproject.toml" + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + steps: + - name: Check out code + uses: actions/checkout@v4 + - name: Publish Custom Node + uses: Comfy-Org/publish-node-action@main + with: + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} ## Add your own personal access token to your Github Repository secrets and reference it here. \ No newline at end of file diff --git a/.github/workflows/vite-build.yml b/.github/workflows/vite-build.yml new file mode 100644 index 0000000..ff7c299 --- /dev/null +++ b/.github/workflows/vite-build.yml @@ -0,0 +1,35 @@ +name: NodeJS with vite build + Publish to Comfy registry + +on: + push: + branches: [ "main" ] + pull_request: + branches: [ "main" ] + +jobs: + build: + runs-on: ubuntu-latest + if: ${{ github.repository_owner == 'jtydhr88' }} + + strategy: + matrix: + node-version: [18.x] + + steps: + - uses: actions/checkout@v4 + + - name: Use Node.js ${{ matrix.node-version }} + uses: actions/setup-node@v4 + with: + node-version: ${{ matrix.node-version }} + + - name: Build + run: | + npm install + npm run build + + - name: Publish Custom Node + uses: Comfy-Org/publish-node-action@main + with: + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} + skip_checkout: 'true' diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..6732b3f --- /dev/null +++ b/.gitignore @@ -0,0 +1,121 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# OSX useful to ignore +*.DS_Store +.AppleDouble +.LSOverride + +# Thumbnails +._* + +# Files that might appear in the root of a volume +.DocumentRevisions-V100 +.fseventsd +.Spotlight-V100 +.TemporaryItems +.Trashes +.VolumeIcon.icns +.com.apple.timemachine.donotpresent + +# Directories potentially created on remote AFP share +.AppleDB +.AppleDesktop +Network Trash Folder +Temporary Items +.apdisk + +# C extensions +*.so + +# Distribution / packaging +.Python +env/ +venv/ +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +*.egg-info/ +.installed.cfg +*.egg + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log + +# Sphinx documentation +docs/_build/ + +# IntelliJ Idea +.idea +*.iml +*.ipr +*.iws + +# PyBuilder +target/ + +# Cookiecutter +output/ +python_boilerplate/ +cookiecutter-pypackage-env/ + +# vscode settings +.history/ +*.code-workspace + +# Frontend extension +node_modules/ +.env +.env.local +.env.development.local +.env.test.local +.env.production.local +npm-debug.log* +yarn-debug.log* +yarn-error.log* +node.zip +.vscode/ +.claude/ +.GEMINI.md + +# duckdb +data/* +!data/default* + +uv.lock diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..2539924 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,10 @@ +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + # Ruff version. + rev: v0.4.9 + hooks: + # Run the linter. + - id: ruff + args: [ --fix ] + # Run the formatter. + - id: ruff-format diff --git a/GEMINI.md b/GEMINI.md new file mode 100644 index 0000000..e69de29 diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..46db539 --- /dev/null +++ b/LICENSE @@ -0,0 +1,22 @@ +MIT License + +Copyright (c) 2026, ComfyUI-Simple-Prompt + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + diff --git a/README.md b/README.md new file mode 100644 index 0000000..44f677e --- /dev/null +++ b/README.md @@ -0,0 +1,83 @@ +# ComfyUI Simple Prompt + +[中文](readme_zh.md) | [English](README.md) + + +A powerful, visual prompt editing node for ComfyUI that blends the best of text-based and visual tag management. + +![Simple Prompt Interface](./doc/preview.png) + +## ✨ Key Features + +* **⚡ High-Performance Backend**: Built with **DuckDB** to handle 70k+ Danbooru tags with millisecond-level search and autocomplete latency. +* **🎨 Visual & Text Modes**: + * **Visual Mode**: Drag-and-drop tags, double-click to adjust weights, color-coded categories (Character, Artist, etc.). + * **Text Mode**: Classic text editing with syntax highlighting and smart autocomplete. +* **🔍 Advanced Search**: + * Search by tag name or **alias** (e.g., "miku" finds "hatsune_miku"). + * Filter by categories (Artist, Style, Character, etc.). + * Sort by post count to find the most popular tags. +* **❤️ Personalization**: + * **Like Tags**: Star your favorite tags for quick access. + * **Custom Tags**: Add your own tags that aren't in the database. + * **Data Updates**: One-click update of tag data from the cloud. +* **🛠️ Deep Integration**: Fully integrated into the ComfyUI node graph workflow. + +## 🚀 Installation + +### Method 1: ComfyUI Manager (Recommended) +1. Open ComfyUI Manager. +2. Search for `Simple Prompt`. +3. Click **Install**. + +### Method 2: Manual Installation +1. Navigate to your ComfyUI `custom_nodes` directory: + ```bash + cd ComfyUI/custom_nodes/ + ``` +2. Clone this repository: + ```bash + git clone https://github.com/0nikod/ComfyUI-Simple-Prompt.git + cd Comfyui-Simple-Prompt + ``` +3. Install dependencies (DuckDB is required): + ```bash + pip install -r requirements.txt + # OR directly + pip install duckdb + ``` +4. Restart ComfyUI. + +## 📖 Usage + +### Basic Controls +* **Adding Tags**: Type in the text box or use the Search button to find tags. +* **Autocomplete**: Press `Ctrl + Space` or simply type to see suggestions. +* **Weight Adjustment**: + * **Visual**: Double-click a tag capsule to edit weight. + * **Shortcuts**: Select a tag and press `Ctrl + Up` / `Ctrl + Down` to increase/decrease weight. + +### Data Management +* **Updating Tags**: Go to Settings -> Data Management -> Click **Update Now** to get the latest Danbooru tag dataset. +* **User Tags**: Add custom tags via the Search Modal -> "Add Custom Tag". + +## ⚙️ Settings + +Click the **Settings** (gear icon) to configure: +* **Smart Backspace**: Automatically delete associated commas/spaces when deleting a tag. +* **Underscore Conversion**: Auto-convert `_` to space when adding tags. +* **Language**: Switch between English and Chinese. + +## 🏗️ Architecture + +* **Frontend**: Vue 3 + Vite + PrimeVue +* **Backend**: Python + DuckDB + Parquet +* **Communication**: ComfyUI API Routes + +## 🤝 Contributing + +Contributions are welcome! Please check the [Design Document (EN)](./doc/DESIGN_DOC_EN.md) / [设计文档 (ZH)](./doc/DESIGN_DOC_ZH.md) for architectural details. + +## 📄 License + +MIT License diff --git a/README_ARCH.md b/README_ARCH.md new file mode 100644 index 0000000..3cb8496 --- /dev/null +++ b/README_ARCH.md @@ -0,0 +1,70 @@ +# ComfyUI-Simple-Prompt Architecture Documentation + +This document describes the architectural design of `ComfyUI-Simple-Prompt` after the v0.1.0 refactor. The main goal of this refactor was to decouple the core business logic from the ComfyUI framework and support hot reloading during development. + +## Directory Structure + +``` +ComfyUI-Simple-Prompt/ +├── simple_prompt/ +│ ├── __init__.py # Package initialization +│ ├── core/ # Core business logic (platform-independent) +│ │ ├── config.py # Path configurations +│ │ ├── database.py # DuckDB connection and initialization +│ │ ├── tags.py # Tag search and management logic +│ │ ├── categories.py # Category management logic +│ │ ├── presets.py # Preset management logic +│ │ └── utils.py # General utility functions +│ ├── api/ # API abstraction layer +│ │ └── handlers.py # Pure functional request handlers +│ └── comfyui/ # ComfyUI adapter layer +│ ├── nodes.py # Node definitions (SimplePrompt) +│ └── server.py # PromptServer route registration & hot reload logic +├── tests/ # Unit tests +├── js/ # Frontend build artifacts +├── src/ # Frontend source code +├── __init__.py # ComfyUI entry point +└── pyproject.toml # Python project configuration +``` + +## Core Design + +### 1. Layered Architecture +- **Core Layer (`simple_prompt/core`)**: Contains pure Python logic, independent of `aiohttp` or `ComfyUI`. All data operations (DuckDB, JSON read/write) are handled in this layer. +- **API Layer (`simple_prompt/api`)**: `handlers.py` acts as an intermediate layer, receiving parameters, calling Core layer functions, and returning standard Python dictionaries/lists. +- **Adapter Layer (`simple_prompt/comfyui`)**: Responsible for forwarding `PromptServer` requests to the API layer and wrapping the results in `web.json_response`. + +### 2. Data Management +- Uses **DuckDB** as the primary query engine, merging multiple Parquet data sources (Main repository, User custom, Liked, Default) through Views. +- Configuration files (categories, presets) are stored as JSON. +- All data paths are centrally managed in `simple_prompt/core/config.py`. + +## Development Guide + +### Hot Reload + +To improve development efficiency, the backend supports modifying code **without restarting ComfyUI**. + +**How to enable:** +Set the environment variable `SIMPLE_PROMPT_DEV=1` when starting ComfyUI. + +**Windows (PowerShell):** +```powershell +$env:SIMPLE_PROMPT_DEV="1" +.\run_nvidia_gpu.bat # Or your startup script +``` + +**How it works:** +- `simple_prompt/comfyui/server.py` contains a `_maybe_reload()` function. +- Every API request (e.g., tag search) triggers this function. If in development mode, it uses `importlib.reload` to reload all modules under `simple_prompt.core` and `simple_prompt.api`. +- **Note**: After modifying `nodes.py` (node definitions) or `server.py` (route registration), you **still need to restart** ComfyUI as these are registered at startup. + +### Running Tests + +Use `pytest` to run unit tests: + +```bash +uv run pytest tests/ -v +``` + +The test environment is independent and does not rely on the ComfyUI runtime. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..758dfdc --- /dev/null +++ b/__init__.py @@ -0,0 +1,35 @@ +""" +Simple Prompt - ComfyUI Custom Node +一个带有高级可视化标签编辑器的提示词构建节点 + +架构说明： +- simple_prompt/core/ - 核心业务逻辑（与平台无关） +- simple_prompt/api/ - API 处理器层 +- simple_prompt/comfyui/ - ComfyUI 特定适配 + +热重载：设置环境变量 SIMPLE_PROMPT_DEV=1 启用开发模式 +""" + +import os +import sys + +# 将当前目录添加到 sys.path，以便能导入 simple_prompt 包 +current_dir = os.path.dirname(os.path.abspath(__file__)) +if current_dir not in sys.path: + sys.path.append(current_dir) + +# 初始化核心模块 +from simple_prompt.core.database import init_duckdb + +init_duckdb() + +# 注册 API 路由（副作用导入） +from simple_prompt.comfyui import server # noqa: F401 + +# 导出节点定义 +from simple_prompt.comfyui.nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +# 前端扩展目录 +WEB_DIRECTORY = "./js" + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/data/default_category.json b/data/default_category.json new file mode 100644 index 0000000..de5010a --- /dev/null +++ b/data/default_category.json @@ -0,0 +1,37 @@ +[ + { + "id": 0, + "name": "General", + "color": "#0075db" + }, + { + "id": 1, + "name": "Artist", + "color": "#c00004" + }, + { + "id": 3, + "name": "Copyright", + "color": "#c000c0" + }, + { + "id": 4, + "name": "Character", + "color": "#00aa00" + }, + { + "id": 5, + "name": "Meta", + "color": "#fd9200" + }, + { + "id": 6, + "name": "Special", + "color": "#ffd700" + }, + { + "id": 7, + "name": "Rating", + "color": "#ff69b4" + } +] \ No newline at end of file diff --git a/data/default_presets.json b/data/default_presets.json new file mode 100644 index 0000000..0637a08 --- /dev/null +++ b/data/default_presets.json @@ -0,0 +1 @@ +[] \ No newline at end of file diff --git a/data/default_tags.parquet b/data/default_tags.parquet new file mode 100644 index 0000000..ffabcf9 Binary files /dev/null and b/data/default_tags.parquet differ diff --git a/doc/DESIGN_DOC_EN.md b/doc/DESIGN_DOC_EN.md new file mode 100644 index 0000000..e8ec210 --- /dev/null +++ b/doc/DESIGN_DOC_EN.md @@ -0,0 +1,97 @@ +# ComfyUI Simple Prompt Design Document + +## 1. Project Overview + +**ComfyUI Simple Prompt** is an advanced prompt editing node designed specifically for ComfyUI. It aims to solve the pain points of traditional text input boxes when dealing with complex prompts, providing a WebUI-like visual editing experience while maintaining ComfyUI's fluidity. + +Core Goals: +* **Visual Editing**: Provides capsule-style tag displays that are easy to drag, delete, and adjust weights. +* **High-Performance Data Retrieval**: Utilizes **DuckDB** and **Parquet** files to handle massive Danbooru tag data, achieving millisecond-level search and completion. +* **Seamless Integration**: Functions as a custom ComfyUI node, integrating perfectly into existing workflows. +* **Customization & Expansion**: Supports user-defined tags, favorite tags, and allows for data updates via the network. + +## 2. System Architecture + +The project adopts a decoupled frontend-backend architecture: + +* **Backend (Python)**: Responsible for data storage, retrieval APIs, and node logic. +* **Frontend (Vue 3 + TypeScript)**: Responsible for UI interaction, state management, and communicating with the backend via API. + +### 2.1 Backend Architecture (`simple_prompt/`) + +The backend core has been refactored into the `simple_prompt` package, separating core business logic from the ComfyUI adaptation layer. + +* **Data Storage (DuckDB)**: + Uses DuckDB's In-Memory database to query Parquet files directly from disk via `read_parquet`. + * **Layered Data Sources** (Priority from high to low): + 1. `liked_tags.parquet` (Custom Data): User-bookmarked/liked tags. + 2. `user_tags.parquet` (Custom Data): User-added custom tags. + 3. `default_tags.parquet`: Pre-configured default tags. + 4. `tags.parquet`: Main dataset containing 70k+ Danbooru tags. + +* **Key APIs** (Defined in `simple_prompt/comfyui/server.py` & `simple_prompt/api/handlers.py`): + * `GET /simple-prompt/tags/list`: Get tag list (supports pagination and search). + * `GET /simple-prompt/search-tags`: High-performance search (supports alias matching and category filtering). + * `POST /simple-prompt/add-custom-tag`: Add or modify a tag. + * `POST /simple-prompt/toggle-like-tag`: Favorite or unfavorite a tag. + * `POST /simple-prompt/update-tags`: Update data files from GitHub Releases. + +### 2.2 Frontend Architecture (`src/`) + +The frontend is built with **Vue 3** (Composition API) and bundled using **Vite**. + +* **Component Structure**: + * `App.vue`: Root component, mounted to the ComfyUI node's DOM element. + * `VisualTagArea.vue`: Visual tag display area, supporting drag-and-drop sorting and double-click weight adjustment. + * `TextEditor.vue`: Plain text editing area, supporting syntax highlighting and smart autocomplete. + * `TagSearchModal.vue`: Advanced search modal, providing filtering by category (Artist, Character, etc.). + * `SettingsModal.vue`: Settings panel for managing interface language, editing habits, etc. + * `TagManager.vue`: Data management panel for managing custom and bookmarked tags. + +* **Internationalization (i18n)**: + Uses `vue-i18n`, supporting switching between English and Chinese (`locales/en/main.json`, `locales/zh/main.json`). + +## 3. Core Feature Design + +### 3.1 Data Retrieval and Completion +* **Backend Implementation**: Leverages DuckDB's SQL engine to perform fuzzy matching via `ILIKE`. +* **Alias System**: Supports Danbooru alias searching (e.g., searching "miku" can match "hatsune_miku"). This is implemented in SQL by expanding the `alias` array using `unnest`. +* **Performance Optimization**: Parquet files are pre-sorted by `post_count`, and DuckDB queries prioritize returning high-frequency tags first. + +### 3.2 Autocomplete +* **Trigger Mechanism**: Triggered in real-time as you type in the text box, or explicitly using `Ctrl+Space`. +* **Logic**: The frontend captures the word at the cursor position and sends it to the backend's `/simple-prompt/search-tags` endpoint. The backend returns a list of matches. Configurable to include or exclude aliases. + +### 3.3 Tag Weight Adjustment +* **From UI**: Double-click a visual tag capsule or use shortcuts (`Ctrl+Up/Down`) to adjust weights. +* **Format**: Standard A1111 format `(tag:1.1)` or `(tag:0.9)`. +* **Smart Parsing**: Automatically parses weight parentheses in the text into visual numeric values. + +## 4. Data Flow Diagram + +```mermaid +graph TD + User[User Operation] -->|Input/Search| Frontend[Vue Frontend] + Frontend -->|API Request| Backend[Python Backend] + Backend -->|SQL Query| DuckDB[(DuckDB In-Memory)] + DuckDB -->|Read| Parquet1[tags.parquet] + DuckDB -->|Read| Parquet2[user_tags.parquet] + DuckDB -->|Read| Parquet3[liked_tags.parquet] + Backend -->|JSON Response| Frontend + Frontend -->|Update Prompt| ComfyUI[ComfyUI Node] +``` + +## 5. Directory Structure Explanation + +* `simple_prompt/`: Main Python backend source code. + * `core/`: Core business logic (Tags, Database). + * `api/`: API processing logic. + * `comfyui/`: ComfyUI node and Server adaptation. +* `js/`: Compiled frontend code (generated by Vite). +* `src/`: Frontend source code. +* `data/`: Directory for Parquet data files. +* `locales/`: Localization files. + +## 6. Future Roadmap +* Support for more prompt formats (e.g., NovelAI format). +* Integration of richer image previews (when hovering over a tag). diff --git a/doc/DESIGN_DOC_ZH.md b/doc/DESIGN_DOC_ZH.md new file mode 100644 index 0000000..65ff325 --- /dev/null +++ b/doc/DESIGN_DOC_ZH.md @@ -0,0 +1,97 @@ +# ComfyUI Simple Prompt 设计文档 + +## 1. 项目概述 + +**ComfyUI Simple Prompt** 是一个专为 ComfyUI 设计的高级提示词编辑节点。它旨在解决传统文本输入框在处理复杂 Prompt 时的痛点，提供类似 WebUI 的可视化编辑体验，同时保持 ComfyUI 的流畅性。 + +核心目标： +* **可视化编辑**：提供胶囊式的标签展示，易于拖拽、删除和权重调整。 +* **高性能数据检索**：利用 **DuckDB** 和 **Parquet** 文件处理海量 Danbooru 标签数据，实现毫秒级搜索和补全。 +* **无缝集成**：作为 ComfyUI 的自定义节点，完美融入现有工作流。 +* **自定义与扩展**：支持用户自定义标签、收藏标签，并允许通过网络更新数据。 + +## 2. 系统架构 + +项目采用前后端分离的架构： + +* **后端 (Python)**: 负责数据存储、检索 API 和节点逻辑。 +* **前端 (Vue 3 + TypeScript)**: 负责 UI 交互、状态管理和通过 API 与后端通信。 + +### 2.1 后端架构 (`simple_prompt/`) + +后端核心已被重构为 `simple_prompt` 包，核心业务逻辑与 ComfyUI 适配层分离。 + +* **数据存储 (DuckDB)**: + 使用 DuckDB 的 In-Memory 数据库，通过 `read_parquet` 直接查询磁盘上的 Parquet 文件。 + * **分层数据源**（优先级从高到低）： + 1. `liked_tags.parquet` (Custom Data): 用户收藏/点赞的标签。 + 2. `user_tags.parquet` (Custom Data): 用户添加的自定义标签。 + 3. `default_tags.parquet`: 预置的默认标签。 + 4. `tags.parquet`: 包含 70k+ Danbooru 标签的主数据集。 + +* **关键 API** (定义于 `simple_prompt/comfyui/server.py` & `simple_prompt/api/handlers.py`): + * `GET /simple-prompt/tags/list`: 获取标签列表（支持分页、搜索）。 + * `GET /simple-prompt/search-tags`: 高性能搜索（支持别名匹配、分类筛选）。 + * `POST /simple-prompt/add-custom-tag`: 添加/修改 Tag。 + * `POST /simple-prompt/toggle-like-tag`: 收藏/取消收藏 Tag。 + * `POST /simple-prompt/update-tags`: 从 GitHub Release 更新数据文件。 + +### 2.2 前端架构 (`src/`) + +前端使用 **Vue 3** (Composition API) 构建，利用 **Vite** 进行打包。 + +* **组件结构**: + * `App.vue`: 根组件，挂载到 ComfyUI 节点的 DOM 元素上。 + * `VisualTagArea.vue`: 可视化标签展示区，支持标签的拖拽排序、双击修改权重。 + * `TextEditor.vue`: 纯文本编辑区，支持语法高亮、自动补全 (Autocomplete)。 + * `TagSearchModal.vue`: 高级搜索模态框，提供按分类 (Artist, Character 等) 筛选的功能。 + * `SettingsModal.vue`: 设置面板，管理界面语言、编辑习惯等。 + * `TagManager.vue`: 数据管理面板，用于管理自定义和收藏的标签。 + +* **国际化 (i18n)**: + 使用 `vue-i18n`，支持中英文切换 (`locales/zh/main.json`, `locales/en/main.json`)。 + +## 3. 核心功能设计 + +### 3.1 数据检索与补全 +* **后端实现**: 利用 DuckDB 的 SQL 引擎，通过 `ILIKE` 进行模糊匹配。 +* **别名系统**: 支持 Danbooru 别名搜索（例如搜 "hatsune miku" 可以匹配到 "hatsune_miku"）。在 SQL 中通过 `unnest(alias)` 展开数组进行匹配。 +* **性能优化**: Parquet 文件按 `post_count` 预排序，DuckDB 查询时优先返回高频标签。 + +### 3.2 自动补全 (Autocomplete) +* **触发机制**: 在文本框输入时实时触发，或通过 `Ctrl+Space` 唤起。 +* **逻辑**: 前端捕获光标位置的单词，发送给后端 `/simple-prompt/search-tags`，后端返回匹配列表。支持配置是否匹配别名。 + +### 3.3 标签权重调整 +* **从 UI**: 双击可视化标签或使用快捷键 (Ctrl+Up/Down) 调整权重。 +* **格式**: 标准 A1111 格式 `(tag:1.1)` 或 `(tag:0.9)`。 +* **智能解析**: 能够自动将文本中的括号权重解析为可视化数值。 + +## 4. 数据流图 + +```mermaid +graph TD + User[用户操作] -->|输入/搜索| Frontend[Vue 前端] + Frontend -->|API 请求| Backend[Python 后端] + Backend -->|SQL Query| DuckDB[(DuckDB 内存库)] + DuckDB -->|Read| Parquet1[tags.parquet] + DuckDB -->|Read| Parquet2[user_tags.parquet] + DuckDB -->|Read| Parquet3[liked_tags.parquet] + Backend -->|JSON 响应| Frontend + Frontend -->|更新 Prompt| ComfyUI[ComfyUI 节点] +``` + +## 5. 目录结构说明 + +* `simple_prompt/`: Python 后端主要源码。 + * `core/`: 核心业务逻辑 (Tags, Database)。 + * `api/`: API 处理逻辑。 + * `comfyui/`: ComfyUI 节点与 Server 适配。 +* `js/`: 编译后的前端代码 (由 Vite 生成)。 +* `src/`: 前端源码。 +* `data/`: 存放 Parquet 数据文件。 +* `locales/`: 语言文件。 + +## 6. 未来规划 +* 支持更多的 Prompt 格式（如 NovelAI 格式）。 +* 集成更丰富的图像预览功能（当鼠标悬停在标签上时）。 diff --git a/doc/preview.png b/doc/preview.png new file mode 100644 index 0000000..98dfe45 Binary files /dev/null and b/doc/preview.png differ diff --git a/index.html b/index.html new file mode 100644 index 0000000..7988f64 --- /dev/null +++ b/index.html @@ -0,0 +1,17 @@ + + + + + + + Simple Prompt Editor - Dev Mode + + +
+ + + + + + + diff --git a/locales/en/main.json b/locales/en/main.json new file mode 100644 index 0000000..2fd88a7 --- /dev/null +++ b/locales/en/main.json @@ -0,0 +1,140 @@ +{ + "common": { + "cancel": "Cancel", + "save": "Save", + "close": "Close", + "done": "Done", + "confirm": "Confirm", + "update": "Update" + }, + "settings": { + "title": "Settings", + "sections": { + "textFormat": "Text Formatting", + "autocomplete": "Autocomplete & Search", + "editing": "Editing Behavior", + "interface": "Interface", + "data": "Data Management", + "tagManager": "Tag Manager", + "addPlaceholder": "New Tag Name (comma separated for multiple)", + "categoryManager": "Category Management", + "metaManager": "Meta Tag Management" + }, + "items": { + "convertUnderscore": "Convert Underscore to Space", + "convertUnderscoreDesc": "Automatically convert underscores to spaces", + "escapeBrackets": "Escape Brackets", + "escapeBracketsDesc": "Escape parentheses as \\(\\)", + "useAliasSearch": "Use Aliases in Search", + "useAliasSearchDesc": "Match tag aliases when searching", + "useAliasAutocomplete": "Use Aliases in Autocomplete", + "useAliasAutocompleteDesc": "Match aliases in autocomplete (may impact performance)", + "smartBackspace": "Smart Backspace", + "smartBackspaceDesc": "When deleting the last character of a tag, also remove preceding comma and space", + "language": "Language", + "languageDesc": "Interface language", + "allowEditDefaultTags": "Allow Editing Default Tags", + "allowEditDefaultTagsDesc": "Enable editing and deleting of default tags (Use with caution).", + "tagMaxLength": "Tag Max Length", + "tagMaxLengthDesc": "Maximum number of characters to display for a tag (0 to disable)", + "updateTags": "Update Tags Data", + "updateTagsDesc": "Download the latest tags.parquet from GitHub", + "updateLikedTags": "Update Liked Tags", + "updateLikedTagsDesc": "Sync liked tags data with main database.", + "updateUserTags": "Update User Tags", + "updateUserTagsDesc": "Sync custom user tags data with main database.", + "updateTagsStatus": "Update Status", + "updateNow": "Update Now", + "updating": "Updating...", + "updateSuccess": "Update successful!", + "updateError": "Update failed: ", + "currentVersion": "Current Version", + "checkingUpdate": "Checking for updates...", + "upToDate": "Latest version: ", + "updateAvailable": "Update available: ", + "unknown": "Unknown" + }, + "meta": { + "title": "Meta Presets Management", + "newPreset": "New Preset", + "editPreset": "Edit Preset", + "presetName": "Preset Name", + "presetTags": "Tags (comma separated)", + "customPresets": "Custom Presets", + "defaultPresets": "Default Presets (Read-only)", + "noPresets": "No custom presets.", + "searchPlaceholder": "Search...", + "noPresetsFound": "No presets found.", + "editInSettings": "Edit in Settings > Meta", + "presetNamePlaceholder": "Preset Name (e.g. Model A)", + "presetTagsPlaceholder": "Tags (comma separated, e.g. quality:best, rating:safe)", + "customPresetsHeading": "Custom Presets", + "defaultPresetsHeading": "Default Presets (Read-only)", + "noCustomPresets": "No custom presets." + }, + "autoSave": "Settings are saved automatically" + }, + "search": { + "title": "Tag Search", + "placeholder": "Search tags (aliases supported)...", + "filterLabel": "Category Filters:", + "loading": "Searching...", + "noResults": "No matching tags found", + "hint": "Enter at least 2 characters to search", + "categories": { + "general": "General", + "artist": "Artist", + "character": "Character", + "copyright": "Copyright", + "meta": "Meta", + "unknown": "Unknown" + }, + "addBtnTitle": "Add to prompt" + }, + "customTag": { + "title": "Add Custom Tag", + "titleEdit": "Edit Tag", + "addBtnTitle": "Add New Tag", + "nameLabel": "Tag Name", + "namePlaceholder": "e.g., my_awesome_tag", + "categoryLabel": "Category", + "countLabel": "Post Count (Optional)", + "countHint": "Estimated usage frequency", + "aliasLabel": "Aliases (Optional)", + "aliasPlaceholder": "Comma-separated, e.g., tag1, tag2", + "aliasHint": "Alternative names to search for this tag", + "errorNameRequired": "Tag name is required", + "errorGeneric": "Failed to add tag", + "successAdded": "Tag added successfully!", + "successEdited": "Tag updated successfully!" + }, + "category": { + "cannotDeleteDefault": "Cannot delete default category.", + "nameRequired": "Name and tags are required", + "namePlaceholder": "Category Name", + "addButton": "Add" + }, + "editor": { + "copy": "Copy", + "copyTitle": "Copy to Clipboard", + "placeholder": "Enter positive prompt here...", + "subtitle": "Simple Prompt Editor", + "tagCount": "Tags: ", + "autocompleteHint": "Ctrl+Space for autocomplete", + "searchBtn": "Search", + "loading": "Loading...", + "noMatches": "No matching tags", + "decreaseWeight": "Decrease weight", + "increaseWeight": "Increase weight", + "removeTag": "Remove tag", + "format": "Format", + "formatTitle": "Format Text", + "organize": "Organize", + "organizeTitle": "Organize Tags", + "autoMeta": "Auto Meta", + "qualityWords": "Quality Tags", + "selectPreset": "Select Meta Preset", + "presetSearchPlaceholder": "Search presets...", + "editInSettings": "Edit in Settings > Meta" + } +} \ No newline at end of file diff --git a/locales/zh/main.json b/locales/zh/main.json new file mode 100644 index 0000000..b7387ef --- /dev/null +++ b/locales/zh/main.json @@ -0,0 +1,140 @@ +{ + "common": { + "cancel": "取消", + "save": "保存", + "close": "关闭", + "done": "完成", + "confirm": "确认", + "update": "更新" + }, + "settings": { + "title": "设置", + "sections": { + "textFormat": "文本格式", + "autocomplete": "补全与搜索", + "editing": "编辑行为", + "interface": "界面", + "data": "数据管理", + "tagManager": "标签管理", + "addPlaceholder": "输入标签名称 (支持逗号分隔批量添加)", + "categoryManager": "分类管理", + "metaManager": "元数据管理" + }, + "items": { + "convertUnderscore": "下划线转空格", + "convertUnderscoreDesc": "自动将标签中的下划线 \"_\" 转换为空格", + "escapeBrackets": "转义括号", + "escapeBracketsDesc": "将括号 \"()\" 转义为 \"\\(\\)\"", + "useAliasSearch": "在搜索中使用别名", + "useAliasSearchDesc": "搜索时匹配标签别名", + "useAliasAutocomplete": "在自动补全中使用别名", + "useAliasAutocompleteDesc": "自动补全时匹配别名（可能影响性能）", + "smartBackspace": "智能退格", + "smartBackspaceDesc": "删除标签最后一个字符时，连同前面的逗号和空格一起删除", + "language": "语言", + "languageDesc": "界面语言", + "allowEditDefaultTags": "允许编辑默认标签", + "allowEditDefaultTagsDesc": "启用后可编辑和删除默认标签（请谨慎使用）。", + "tagMaxLength": "标签最大长度", + "tagMaxLengthDesc": "标签显示的最大字符数（0 表示不限制）", + "updateTags": "更新标签数据", + "updateTagsDesc": "从 GitHub 下载最新的 tags.parquet", + "updateLikedTags": "更新喜爱标签数据", + "updateLikedTagsDesc": "同步喜爱标签的统计数据与主数据库。", + "updateUserTags": "更新自定义标签数据", + "updateUserTagsDesc": "同步自定义标签的统计数据与主数据库。", + "updateTagsStatus": "更新状态", + "updateNow": "立即更新", + "updating": "更新中...", + "updateSuccess": "更新成功！", + "updateError": "更新失败：", + "currentVersion": "当前版本", + "checkingUpdate": "正在检查更新...", + "upToDate": "已是最新：", + "updateAvailable": "有可用更新：", + "unknown": "未知" + }, + "meta": { + "title": "Meta 预设管理", + "newPreset": "新建预设", + "editPreset": "编辑预设", + "presetName": "预设名称", + "presetTags": "标签 (逗号分隔)", + "customPresets": "自定义预设", + "defaultPresets": "默认预设 (只读)", + "noPresets": "无自定义预设", + "searchPlaceholder": "搜索...", + "noPresetsFound": "未找到预设。", + "editInSettings": "去 设置 > Meta 编辑", + "presetNamePlaceholder": "预设名称 (例如: Model A)", + "presetTagsPlaceholder": "标签 (逗号分隔, 例如: quality:best, rating:safe)", + "customPresetsHeading": "自定义预设", + "defaultPresetsHeading": "默认预设 (只读)", + "noCustomPresets": "无自定义预设" + }, + "autoSave": "设置会自动保存" + }, + "search": { + "title": "标签搜索", + "placeholder": "搜索标签（支持别名）...", + "filterLabel": "分类筛选:", + "loading": "搜索中...", + "noResults": "未找到匹配的标签", + "hint": "输入至少2个字符开始搜索", + "categories": { + "general": "通用", + "artist": "画师", + "character": "角色", + "copyright": "版权", + "meta": "元数据", + "unknown": "未知" + }, + "addBtnTitle": "添加到prompt" + }, + "customTag": { + "title": "添加自定义标签", + "titleEdit": "编辑标签", + "addBtnTitle": "添加新标签", + "nameLabel": "标签名称", + "namePlaceholder": "例如: my_awesome_tag", + "categoryLabel": "分类", + "countLabel": "帖子计数 (可选)", + "countHint": "估计的使用频率", + "aliasLabel": "别名 (可选)", + "aliasPlaceholder": "逗号分隔，例如: tag1, tag2", + "aliasHint": "搜索此标签时可使用的替代名称", + "errorNameRequired": "标签名称必填", + "errorGeneric": "添加标签失败", + "successAdded": "标签添加成功！", + "successEdited": "标签更新成功！" + }, + "category": { + "cannotDeleteDefault": "无法删除默认分类。", + "nameRequired": "名称和标签为必填项", + "namePlaceholder": "分类名称", + "addButton": "添加" + }, + "editor": { + "copy": "复制", + "copyTitle": "复制到剪贴板", + "placeholder": "在此处输入正向提示词...", + "subtitle": "Simple Prompt Editor", + "tagCount": "标签: ", + "autocompleteHint": "Ctrl+Space 触发自动补全", + "searchBtn": "搜索", + "loading": "加载中...", + "noMatches": "无匹配标签", + "decreaseWeight": "减小权重", + "increaseWeight": "增大权重", + "removeTag": "删除标签", + "format": "格式化", + "formatTitle": "格式化文本", + "organize": "整理", + "organizeTitle": "整理标签", + "autoMeta": "自动Meta", + "qualityWords": "质量标签", + "selectPreset": "选择 Meta 预设", + "presetSearchPlaceholder": "搜索预设...", + "editInSettings": "去 设置 > Meta 编辑" + } +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..cbae853 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,79 @@ +[build-system] +requires = ["setuptools>=70.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "comfyui_simple_prompt" +version = "0.2.1" +description = "Provide a simple interface to simplify prompt writing. / 提供一个简单的界面简化提示词编写。" +authors = [{ name = "0nikod", email = "dniko0@outlook.com" }] +readme = "README.md" +license = { text = "MIT license" } +requires-python = ">=3.10" +classifiers = [] +dependencies = ["duckdb>=1.4.3"] + + +[project.optional-dependencies] +dev = [ + "bump-my-version", + "coverage", # testing + "mypy", # linting + "pre-commit", # runs linting on commit + "pytest", # testing + "ruff", # linting +] + +[project.urls] +Repository = "https://github.com/0nikod/ComfyUI-Simple-Prompt" +BugTracker = "https://github.com/0nikod/ComfyUI-Simple-Prompt/issues" +Documentation = "https://github.com/0nikod/ComfyUI-Simple-Prompt/wiki" + + +[tool.comfy] +PublisherId = "dniko0" +DisplayName = "ComfyUI Simple Prompt" +Icon = "doc/preview.png" +Tags = ["prompt", "tag", "editor", "danbooru", "autocomplete"] +Repository = "https://github.com/0nikod/ComfyUI-Simple-Prompt" +includes = ["js"] + +[tool.setuptools.package-data] +"*" = ["*.*"] + +[tool.pytest.ini_options] +minversion = "8.0" +testpaths = ["tests"] + +[tool.mypy] +files = "." + +# Use strict defaults +strict = true +warn_unreachable = true +warn_no_return = true + +[[tool.mypy.overrides]] +# Don't require test functions to include types +module = "tests.*" +allow_untyped_defs = true +disable_error_code = "attr-defined" + +[tool.ruff] +# extend-exclude = ["static", "ci/templates"] +line-length = 140 +src = ["src", "tests"] +target-version = "py39" + +# Add rules to ban exec/eval +[tool.ruff.lint] +select = [ + "S102", # exec-builtin + "S307", # eval-used + "W293", + "F", # The "F" series in Ruff stands for "Pyflakes" rules, which catch various Python syntax errors and undefined names. + # See all rules here: https://docs.astral.sh/ruff/rules/#pyflakes-f +] + +[tool.ruff.lint.flake8-quotes] +inline-quotes = "double" diff --git a/readme_zh.md b/readme_zh.md new file mode 100644 index 0000000..d302311 --- /dev/null +++ b/readme_zh.md @@ -0,0 +1,82 @@ +# ComfyUI Simple Prompt + +[中文](readme_zh.md) | [English](README.md) + +这是一个功能强大的 ComfyUI 提示词编辑/管理节点，结合了文本编辑和可视化标签管理的优势。 + +![Simple Prompt Interface](./doc/preview.png) + +## ✨ 核心功能 + +* **⚡ 高性能后端**: 基于 **DuckDB** 构建，轻松处理 70k+ Danbooru 标签，实现毫秒级的搜索和自动补全。 +* **🎨 可视化与文本双模式**: + * **可视化模式**: 胶囊式标签展示，支持拖拽排序、双击调整权重，不同类型的标签（角色、画师等）显示不同颜色。 + * **文本模式**: 传统的文本输入体验，支持语法高亮和智能自动补全。 +* **🔍 高级搜索**: + * 支持标签名称和**别名**搜索（例如搜 "miku" 可以找到 "hatsune_miku"）。 + * 支持按分类（画师、风格、角色等）筛选。 + * 支持按帖子数量排序，快速找到热门标签。 +* **❤️ 个性化与数据管理**: + * **收藏标签**: 星标你常用的标签，方便快速访问。 + * **自定义标签**: 添加数据库中不存在的新标签。 + * **在线更新**: 一键从云端更新 Danbooru 标签数据库。 +* **🛠️ 深度集成**: 作为标准节点无缝融入 ComfyUI 工作流。 + +## 🚀 安装指南 + +### 方法 1: 使用 ComfyUI Manager (推荐) +1. 打开 ComfyUI Manager。 +2. 搜索 `Simple Prompt`。 +3. 点击 **Install** 安装。 + +### 方法 2: 手动安装 +1. 进入 ComfyUI 的 `custom_nodes` 目录： + ```bash + cd ComfyUI/custom_nodes/ + ``` +2. 克隆本仓库： + ```bash + git clone https://github.com/0nikod/ComfyUI-Simple-Prompt.git + cd Comfyui-Simple-Prompt + ``` +3. 安装依赖 (必须安装 DuckDB)： + ```bash + pip install -r requirements.txt + # 或者直接安装 + pip install duckdb + ``` +4. 重启 ComfyUI。 + +## 📖 使用说明 + +### 基本操作 +* **添加标签**: 在文本框输入，或点击“搜索”按钮查找标签。 +* **自动补全**: 在输入时会自动弹出建议，或按 `Ctrl + Space` 手动触发。 +* **调整权重**: + * **可视化调整**: 双击标签胶囊即可修改权重数值。 + * **快捷键**: 选中标签后，按 `Ctrl + Up` / `Ctrl + Down` 增加或减少权重。 + +### 数据管理 +* **更新数据**: 点击设置图标 -> 数据管理 -> 点击 **立即更新** 即可获取最新的标签数据。 +* **自定义标签**: 在搜索弹窗中点击 "添加自定义标签"。 + +## ⚙️ 设置选项 + +点击编辑器右上角的 **设置** (齿轮图标) 进行配置： +* **智能退格**: 删除标签时，自动清理多余的逗号和空格。 +* **下划线转换**: 添加标签时自动将 `_` 转换为 空格。 +* **语言**: 切换界面语言 (中文 / English)。 + +## 🏗️ 技术架构 + +* **前端**: Vue 3 + Vite + PrimeVue +* **后端**: Python + DuckDB + Parquet API +* **通信**: ComfyUI API Routes + +## 🤝 参与贡献 + +欢迎提交 PR 或 Issue！关于详细的架构设计，请参考 [设计文档 (中文)](./doc/DESIGN_DOC_ZH.md)。 + +## 📄 许可证 + +MIT License diff --git a/reproduce_parser.ts b/reproduce_parser.ts new file mode 100644 index 0000000..bb310c0 --- /dev/null +++ b/reproduce_parser.ts @@ -0,0 +1,135 @@ + +// Mock types +enum TagCategory { + GENERAL = 0, + ARTIST = 1, + COPYRIGHT = 2, + CHARACTER = 4, + META = 5, + SPECIAL = 6, + RATING = 7 +} + +interface TagItem { + id: string; + text: string; + weight: number; + enabled: boolean; + category: number; +} + +// Paste original functions here for testing +function textToTags(text: string): TagItem[] { + if (!text.trim()) return []; + + const tags: TagItem[] = []; + let idx = 0; + let i = 0; + let currentTag = ''; + let inWeight = false; + + while (i < text.length) { + const char = text[i]; + + // Handle escape sequences + if (char === '\\' && i + 1 < text.length) { + const nextChar = text[i + 1]; + if (nextChar === '(' || nextChar === ')') { + currentTag += nextChar; + i += 2; + continue; + } else { + currentTag += char; + i++; + continue; + } + } + + if (char === '(' && !inWeight) { + if (currentTag.trim()) { + tags.push({ id: `tag-${idx++}`, text: currentTag.trim(), weight: 1.0, enabled: true, category: TagCategory.GENERAL }); + currentTag = ''; + } + + let depth = 1; + let j = i + 1; + + while (j < text.length && depth > 0) { + if (text[j] === '\\' && j + 1 < text.length) { + j += 2; + continue; + } + if (text[j] === '(') depth++; + if (text[j] === ')') depth--; + j++; + } + + if (depth === 0) { + const content = text.substring(i + 1, j - 1); + const colonIndex = content.indexOf(':'); + + if (colonIndex > 0) { + const tagText = content.substring(0, colonIndex).trim(); + const weight = parseFloat(content.substring(colonIndex + 1).trim()); + if (!isNaN(weight) && tagText) { + tags.push({ id: `tag-${idx++}`, text: tagText, weight: weight, enabled: true, category: TagCategory.GENERAL }); + } + } else { + currentTag += text.substring(i, j); + } + i = j; + continue; + } + } + + if (char === ',' || char === '\n') { + if (currentTag.trim()) { + tags.push({ id: `tag-${idx++}`, text: currentTag.trim(), weight: 1.0, enabled: true, category: TagCategory.GENERAL }); + currentTag = ''; + } + i++; + continue; + } + + currentTag += char; + i++; + } + + if (currentTag.trim()) { + tags.push({ id: `tag-${idx++}`, text: currentTag.trim(), weight: 1.0, enabled: true, category: TagCategory.GENERAL }); + } + + return tags; +} + +function tagsToText(tags: TagItem[]): string { + const enabledTags = tags.filter(tag => tag.enabled); + if (enabledTags.length === 0) return ''; + + return enabledTags + .map(tag => { + const escapedText = tag.text.replace(/\(/g, '\\(').replace(/\)/g, '\\)'); + if (tag.weight !== 1.0) { + return `(${escapedText}:${tag.weight.toFixed(1)})`; + } + return escapedText; + }) + .join(', ') + ','; +} + +// Test Case +const input = "aaa \\(bbb\\)"; +console.log("Input:", input); +const tags = textToTags(input); +console.log("Parsed Tags:", JSON.stringify(tags, null, 2)); + +const output = tagsToText(tags); +console.log("Output:", output); + +if (output.includes("\\\\")) { + console.log("FAIL: Double escaping detected!"); +} else if (output.trim().replace(/,$/, '') === input) { + console.log("PASS: Output matches input."); +} else { + console.log("DIFF: Output different from input."); +} diff --git a/reproduce_parser_v2.ts b/reproduce_parser_v2.ts new file mode 100644 index 0000000..8284400 --- /dev/null +++ b/reproduce_parser_v2.ts @@ -0,0 +1,37 @@ + +import { textToTags, tagsToText } from './src/utils/promptParser'; + +// Test Case +const input = "aaa \\(bbb\\)"; +console.log("Input:", input); +const tags = textToTags(input); +console.log("Parsed Tags:", JSON.stringify(tags, null, 2)); + +const output = tagsToText(tags); +console.log("Output:", output); + +if (output.includes("\\\\(")) { + console.log("FAIL: Double escaping detected!"); +} else if (output.trim().replace(/,$/, '') === input) { + console.log("PASS: Output matches input."); +} else { + // It might differ slightly if we normalized spacing etc, but key is no double escape + console.log("INFO: Output is:", output); + if (!output.includes("\\\\(")) { + console.log("PASS: No double escaping."); + } +} + +// Test Weighted double escape +const inputWeighted = "(aaa \\(bbb\\):1.0)"; +console.log("\nInput Weighted:", inputWeighted); +const tagsW = textToTags(inputWeighted); +console.log("Parsed Tags Weighted:", JSON.stringify(tagsW, null, 2)); +const outputW = tagsToText(tagsW); +console.log("Output Weighted:", outputW); + +if (outputW.includes("\\\\(")) { + console.log("FAIL: Double escaping in weighted tag!"); +} else { + console.log("PASS: No double escaping in weighted tag."); +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..3da59df --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +duckdb \ No newline at end of file diff --git a/scripts/check_schema.py b/scripts/check_schema.py new file mode 100644 index 0000000..aa5c13a --- /dev/null +++ b/scripts/check_schema.py @@ -0,0 +1,21 @@ +import duckdb +import os + +BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +TAGS_PATH = os.path.join(BASE_DIR, "data", "tags.parquet") + +try: + conn = duckdb.connect(":memory:") + print(f"Inspecting {TAGS_PATH}") + if os.path.exists(TAGS_PATH): + print(conn.execute(f"DESCRIBE SELECT * FROM read_parquet('{TAGS_PATH.replace(os.sep, '/')}')").fetchall()) + else: + print("tags.parquet not found") + + USER_TAGS_PATH = os.path.join(BASE_DIR, "data", "custom", "user_tags.parquet") + print(f"Inspecting {USER_TAGS_PATH}") + if os.path.exists(USER_TAGS_PATH): + print(conn.execute(f"DESCRIBE SELECT * FROM read_parquet('{USER_TAGS_PATH.replace(os.sep, '/')}')").fetchall()) + +except Exception as e: + print(e) diff --git a/scripts/create_sample_data.py b/scripts/create_sample_data.py new file mode 100644 index 0000000..20e078c --- /dev/null +++ b/scripts/create_sample_data.py @@ -0,0 +1,74 @@ +import duckdb +import os + +# Define paths +base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +data_dir = os.path.join(base_dir, "data") +custom_dir = os.path.join(data_dir, "custom") + +os.makedirs(custom_dir, exist_ok=True) + +# Connect to in-memory db +con = duckdb.connect() + +# Create scheme for user tags +# name: VARCHAR, category: BIGINT, post_count: BIGINT, alias: VARCHAR[] + +# 1. Create user_tags.parquet +# Contains a custom tag "my_custom_tag" +con.execute(""" + CREATE TABLE user_tags ( + name VARCHAR, + category BIGINT, + post_count BIGINT, + alias VARCHAR[] + ) +""") +con.execute(""" + INSERT INTO user_tags VALUES + ('my_custom_tag', 0, 100, ['custom_alias']), + ('1girl', 4, 9999999, ['override_girl']) +""") +user_tags_path = os.path.join(custom_dir, "user_tags.parquet") +con.execute(f"COPY user_tags TO '{user_tags_path}' (FORMAT PARQUET)") +print(f"Created {user_tags_path}") + +# 2. Create liked_tags.parquet +# Contains "masterpiece" as a liked tag (will override others if priority logic works) +con.execute(""" + CREATE TABLE liked_tags ( + name VARCHAR, + category BIGINT, + post_count BIGINT, + alias VARCHAR[] + ) +""") +con.execute(""" + INSERT INTO liked_tags VALUES + ('masterpiece', 5, 5000000, ['best quality']), + ('simple background', 0, 200000, []) +""") +liked_tags_path = os.path.join(custom_dir, "liked_tags.parquet") +con.execute(f"COPY liked_tags TO '{liked_tags_path}' (FORMAT PARQUET)") +print(f"Created {liked_tags_path}") + +# 3. Create default_tags.parquet (optional but good for testing priority) +# Contains "1girl" with lower count than user_tags to test override if user_tags < default_tags (but plan says User > Default) +con.execute(""" + CREATE TABLE default_tags ( + name VARCHAR, + category BIGINT, + post_count BIGINT, + alias VARCHAR[] + ) +""") +con.execute(""" + INSERT INTO default_tags VALUES + ('1girl', 4, 500, ['girl']), + ('solo', 0, 300, []) +""") +default_tags_path = os.path.join(data_dir, "default_tags.parquet") +con.execute(f"COPY default_tags TO '{default_tags_path}' (FORMAT PARQUET)") +print(f"Created {default_tags_path}") + +print("Done creating sample data.") diff --git a/scripts/reset_default_tags.py b/scripts/reset_default_tags.py new file mode 100644 index 0000000..8f24ce3 --- /dev/null +++ b/scripts/reset_default_tags.py @@ -0,0 +1,25 @@ +import duckdb +import os + +BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA_DIR = os.path.join(BASE_DIR, "data") +DEFAULT_TAGS_PATH = os.path.join(DATA_DIR, "default_tags.parquet") + +os.makedirs(DATA_DIR, exist_ok=True) + +print(f"Creating empty default_tags.parquet at {DEFAULT_TAGS_PATH}") + +try: + conn = duckdb.connect(":memory:") + + # Schema matching SimplePrompt.py TAGS_SCHEMA + # name VARCHAR, category BIGINT, post_count BIGINT, alias VARCHAR[] + conn.execute("CREATE TABLE tags (name VARCHAR, category BIGINT, post_count BIGINT, alias VARCHAR[])") + + # Write to parquet (OVERWRITE) + save_path = DEFAULT_TAGS_PATH.replace(os.sep, "/") + conn.execute(f"COPY tags TO '{save_path}' (FORMAT PARQUET)") + print("Successfully created empty default_tags.parquet") + +except Exception as e: + print(f"Error: {e}") diff --git a/scripts/test_backend.py b/scripts/test_backend.py new file mode 100644 index 0000000..599e10c --- /dev/null +++ b/scripts/test_backend.py @@ -0,0 +1,96 @@ +import os +import duckdb +import logging + +# Setup basic logging to see output +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger("SimplePrompt") + +# Mocking the layout +BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA_DIR = os.path.join(BASE_DIR, "data") +CUSTOM_DATA_DIR = os.path.join(DATA_DIR, "custom") +TAGS_PARQUET_PATH = os.path.join(DATA_DIR, "tags.parquet") +DEFAULT_TAGS_PATH = os.path.join(DATA_DIR, "default_tags.parquet") +USER_TAGS_PATH = os.path.join(CUSTOM_DATA_DIR, "user_tags.parquet") +LIKED_TAGS_PATH = os.path.join(CUSTOM_DATA_DIR, "liked_tags.parquet") + + +def init_duckdb(): + try: + conn = duckdb.connect(":memory:") + + def get_sql_path(path): + if os.path.exists(path): + return path.replace("\\", "/") + return None + + sources = [] + + path = get_sql_path(LIKED_TAGS_PATH) + if path: + sources.append(f"SELECT name, category, post_count, alias, 1 as priority FROM read_parquet('{path}')") + + path = get_sql_path(USER_TAGS_PATH) + if path: + sources.append(f"SELECT name, category, post_count, alias, 2 as priority FROM read_parquet('{path}')") + + path = get_sql_path(DEFAULT_TAGS_PATH) + if path: + sources.append(f"SELECT name, category, post_count, alias, 3 as priority FROM read_parquet('{path}')") + + path = get_sql_path(TAGS_PARQUET_PATH) + if path: + sources.append(f"SELECT name, category, post_count, alias, 4 as priority FROM read_parquet('{path}')") + + if not sources: + print("No sources found!") + return None + + union_query = " UNION ALL ".join(sources) + print("Creating view...") + conn.execute(f"CREATE OR REPLACE VIEW all_tags_raw AS {union_query}") + + print("Deduplicating...") + conn.execute(""" + CREATE OR REPLACE VIEW tags AS + SELECT * EXCLUDE (rn) FROM ( + SELECT *, + ROW_NUMBER() OVER (PARTITION BY name ORDER BY priority ASC) as rn + FROM all_tags_raw + ) WHERE rn = 1 + """) + + print("DuckDB Initialized successfully.") + return conn + + except Exception as e: + print(f"Failed: {e}") + return None + + +def test_search(conn): + if not conn: + return + + print("\nTesting Search 'girl'...") + try: + # Simple search + res = conn.execute( + "SELECT name, priority, post_count FROM tags WHERE name ILIKE '%girl%' ORDER BY priority ASC, post_count DESC" + ).fetchall() + for r in res: + print(r) + + # Check custom tag + print("\nTesting Custom Tag 'my_custom_tag'...") + res = conn.execute("SELECT * FROM tags WHERE name = 'my_custom_tag'").fetchall() + print(res) + + except Exception as e: + print(f"Search failed: {e}") + + +conn = init_duckdb() +if conn: + test_search(conn) diff --git a/scripts/test_manager_logic.py b/scripts/test_manager_logic.py new file mode 100644 index 0000000..a6a64d8 --- /dev/null +++ b/scripts/test_manager_logic.py @@ -0,0 +1,59 @@ +import duckdb +import os +import shutil + +# Setup paths +BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA_DIR = os.path.join(BASE_DIR, "data") +CUSTOM_DATA_DIR = os.path.join(DATA_DIR, "custom") +TEST_USER_TAGS = os.path.join(CUSTOM_DATA_DIR, "test_user_tags.parquet") + +os.makedirs(CUSTOM_DATA_DIR, exist_ok=True) + +conn = duckdb.connect(":memory:") + + +def setup_data(): + # Create sample parquet + conn.execute("CREATE OR REPLACE TABLE tags (name VARCHAR, category BIGINT, post_count BIGINT, alias VARCHAR[])") + conn.execute("INSERT INTO tags VALUES ('tag1', 0, 10, []), ('tag2', 1, 20, ['alias2']), ('tag3', 0, 5, [])") + conn.execute(f"COPY tags TO '{TEST_USER_TAGS.replace(os.sep, '/')}' (FORMAT PARQUET)") + print("Test data created.") + + +def test_list(): + print("\n--- Testing LIST ---") + path_sql = TEST_USER_TAGS.replace(os.sep, "/") + res = conn.execute(f"SELECT * FROM read_parquet('{path_sql}') ORDER BY post_count DESC, name ASC").fetchall() + print("All tags:", res) + assert len(res) == 3 + assert res[0][0] == "tag2" # highest post_count + + +def test_delete(): + print("\n--- Testing DELETE 'tag2' ---") + path_sql = TEST_USER_TAGS.replace(os.sep, "/") + + conn.execute("DROP TABLE IF EXISTS temp_del_tags") + conn.execute(f"CREATE TABLE temp_del_tags AS SELECT * FROM read_parquet('{path_sql}')") + conn.execute("DELETE FROM temp_del_tags WHERE name = 'tag2'") + conn.execute(f"COPY temp_del_tags TO '{path_sql}' (FORMAT PARQUET)") + + # Verify + res = conn.execute(f"SELECT * FROM read_parquet('{path_sql}')").fetchall() + print("Remaining tags:", res) + assert len(res) == 2 + names = [r[0] for r in res] + assert "tag2" not in names + assert "tag1" in names + + +try: + setup_data() + test_list() + test_delete() + print("\nSUCCESS: All backend logic tests passed.") +finally: + # Cleanup + if os.path.exists(TEST_USER_TAGS): + os.remove(TEST_USER_TAGS) diff --git a/simple_prompt/__init__.py b/simple_prompt/__init__.py new file mode 100644 index 0000000..a332a24 --- /dev/null +++ b/simple_prompt/__init__.py @@ -0,0 +1,4 @@ +# simple_prompt package +# 核心业务逻辑包，与 ComfyUI 解耦 + +__version__ = "0.1.0" diff --git a/simple_prompt/api/__init__.py b/simple_prompt/api/__init__.py new file mode 100644 index 0000000..968c356 --- /dev/null +++ b/simple_prompt/api/__init__.py @@ -0,0 +1,2 @@ +# simple_prompt.api package +# API 处理器层（平台无关的抽象） diff --git a/simple_prompt/api/handlers.py b/simple_prompt/api/handlers.py new file mode 100644 index 0000000..562bdb3 --- /dev/null +++ b/simple_prompt/api/handlers.py @@ -0,0 +1,463 @@ +""" +API 处理器层 +提供平台无关的请求处理函数 +这些函数接收和返回纯 Python 数据结构，不依赖任何 Web 框架 +""" + +import hashlib +import logging +import os +from typing import Any, Dict, List, Optional + +from simple_prompt.core import config +from simple_prompt.core.database import get_db_connection, reinit_duckdb +from simple_prompt.core import tags as tags_module +from simple_prompt.core import categories as categories_module +from simple_prompt.core import presets as presets_module + +logger = logging.getLogger("SimplePrompt") + + +# -------------------------------------------------------------------------------- +# Tags Handlers +# -------------------------------------------------------------------------------- + + +async def handle_search_tags( + query: str, limit: int = 50, use_aliases: bool = False, categories: Optional[List[int]] = None +) -> List[Dict[str, Any]]: + """ + 处理标签搜索请求 + + Args: + query: 搜索关键词 + limit: 返回结果数量限制 + use_aliases: 是否搜索别名 + categories: 限制搜索的分类列表 + + Returns: + 匹配的标签列表 + + Raises: + RuntimeError: 数据库未初始化 + """ + conn = get_db_connection() + if not conn: + raise RuntimeError("DuckDB not initialized") + return tags_module.search_tags(conn, query, limit, use_aliases, categories) + + +async def handle_get_tags_details(names: List[str], fast: bool = False) -> Dict[str, int]: + """ + 处理获取标签详情请求 + + Args: + names: 标签名称列表 + fast: 是否使用快速模式 + + Returns: + 标签名到分类 ID 的映射 + + Raises: + RuntimeError: 数据库未初始化 + """ + conn = get_db_connection() + if not conn: + raise RuntimeError("DuckDB not initialized") + return tags_module.get_tags_details(conn, names, fast) + + +async def handle_list_tags(source: str, limit: int = 50, offset: int = 0, query: str = "") -> Dict[str, Any]: + """ + 处理列出标签请求 + + Args: + source: 数据来源 + limit: 返回数量限制 + offset: 偏移量 + query: 搜索关键词 + + Returns: + 包含 data 和 total 的字典 + + Raises: + RuntimeError: 数据库未初始化 + """ + conn = get_db_connection() + if not conn: + raise RuntimeError("DuckDB not initialized") + return tags_module.list_tags(conn, source, limit, offset, query) + + +async def handle_delete_tag(name: str, source: str) -> Dict[str, str]: + """ + 处理删除标签请求 + + Args: + name: 标签名称 + source: 数据来源 + + Returns: + 操作结果 + + Raises: + RuntimeError: 数据库未初始化 + ValueError: 参数无效 + """ + if not name or not source: + raise ValueError("Name and source are required") + + conn = get_db_connection() + if not conn: + raise RuntimeError("Database unavailable") + + path = tags_module.get_source_path_by_name(source) + if not path: + raise ValueError("Invalid source") + + success = tags_module.delete_tag(conn, name, source) + if success: + return {"status": "success", "message": f"Tag '{name}' deleted from {source}."} + else: + raise RuntimeError("Delete operation failed") + + +async def handle_add_custom_tag(data: Dict[str, Any]) -> Dict[str, str]: + """ + 处理添加自定义标签请求 + + Args: + data: 包含标签信息的字典，支持单个或批量 + + Returns: + 操作结果 + + Raises: + RuntimeError: 数据库未初始化 + ValueError: 参数无效 + """ + conn = get_db_connection() + if not conn: + raise RuntimeError("Database unavailable") + + # Detect Batch or Single + if "tags" in data and isinstance(data["tags"], list): + tags_to_process = data["tags"] + else: + tags_to_process = [data] + + if not tags_to_process: + raise ValueError("No tag data provided") + + # Get source from first item + source = tags_to_process[0].get("source", "user") + + count = tags_module.add_tags(conn, tags_to_process, source) + return {"status": "success", "message": f"{count} tag(s) added/updated."} + + +async def handle_toggle_like_tag(data: Dict[str, Any]) -> Dict[str, Any]: + """ + 处理切换点赞标签请求 + + Args: + data: 包含标签信息的字典 + + Returns: + 操作结果 + + Raises: + RuntimeError: 数据库未初始化 + ValueError: 参数无效 + """ + conn = get_db_connection() + if not conn: + raise RuntimeError("Database unavailable") + + name = data.get("name") + if not name: + raise ValueError("Name is required") + + should_like = data.get("is_liked", True) + category = data.get("category") + post_count = data.get("post_count") + alias = data.get("alias", []) + + success = tags_module.toggle_like_tag(conn, name, should_like, category, post_count, alias) + + if success: + msg = f"Tag '{name}' {'liked' if should_like else 'unliked'}." + return {"status": "success", "message": msg, "is_liked": should_like} + else: + raise RuntimeError("Toggle like operation failed") + + +# -------------------------------------------------------------------------------- +# Categories Handlers +# -------------------------------------------------------------------------------- + + +async def handle_list_categories() -> List[Dict[str, Any]]: + """ + 处理列出分类请求 + + Returns: + 分类列表 + """ + return categories_module.get_categories() + + +async def handle_save_categories(categories: List[Dict[str, Any]]) -> Dict[str, str]: + """ + 处理保存分类请求 + + Args: + categories: 分类列表 + + Returns: + 操作结果 + """ + categories_module.save_custom_categories(categories) + return {"status": "success"} + + +# -------------------------------------------------------------------------------- +# Presets Handlers +# -------------------------------------------------------------------------------- + + +async def handle_list_presets() -> Dict[str, List[Dict[str, Any]]]: + """ + 处理列出预设请求 + + Returns: + 包含 defaults 和 customs 的字典 + """ + return presets_module.get_presets() + + +async def handle_save_presets(presets: List[Dict[str, Any]]) -> Dict[str, str]: + """ + 处理保存预设请求 + + Args: + presets: 预设列表 + + Returns: + 操作结果 + """ + presets_module.save_custom_presets(presets) + return {"status": "success"} + + +# -------------------------------------------------------------------------------- +# Health & Update Handlers +# -------------------------------------------------------------------------------- + + +async def handle_health_check() -> Dict[str, Any]: + """ + 处理健康检查请求 + + Returns: + 系统状态信息 + """ + conn = get_db_connection() + return { + "status": "ok", + "duckdb": "loaded" if conn else "failed", + "sources": { + "main": os.path.exists(config.TAGS_PARQUET_PATH), + "user": os.path.exists(config.USER_TAGS_PATH), + "liked": os.path.exists(config.LIKED_TAGS_PATH), + }, + } + + +async def handle_check_update() -> Dict[str, Any]: + """ + 处理检查更新请求 + + Returns: + 更新信息 + """ + import aiohttp + + release_url = "https://api.github.com/repos/0nikod/danbooru_tag_process/releases/latest" + + local_sha256 = "" + if os.path.exists(config.TAGS_PARQUET_PATH): + try: + sha256_hash = hashlib.sha256() + with open(config.TAGS_PARQUET_PATH, "rb") as f: + for byte_block in iter(lambda: f.read(4096), b""): + sha256_hash.update(byte_block) + local_sha256 = sha256_hash.hexdigest() + except Exception as e: + logger.error(f"Error calculating local SHA256: {e}") + + async with aiohttp.ClientSession() as session: + headers = {"User-Agent": "ComfyUI-Simple-Prompt"} + async with session.get(release_url, headers=headers) as resp: + if resp.status != 200: + raise RuntimeError(f"Failed to fetch release info: {resp.status}") + release_info = await resp.json() + + target_asset = None + asset_names = ["tags_processed.parquet", "tags.parquet"] + for name in asset_names: + for asset in release_info.get("assets", []): + if asset.get("name") == name: + target_asset = asset + break + if target_asset: + break + + if not target_asset: + raise RuntimeError("Tag data file not found") + + remote_digest = target_asset.get("digest", "") + if remote_digest.startswith("sha256:"): + remote_sha256 = remote_digest[7:] + else: + remote_sha256 = remote_digest + + update_available = local_sha256 != remote_sha256 + + return { + "update_available": update_available, + "local_sha256": local_sha256, + "remote_sha256": remote_sha256, + "version": release_info.get("tag_name", "unknown"), + "name": release_info.get("name", ""), + } + + +async def handle_update_tags(url: Optional[str] = None) -> Dict[str, str]: + """ + 处理更新标签数据请求 + + Args: + url: 可选的自定义下载 URL + + Returns: + 操作结果 + """ + import aiohttp + + release_url = url or "https://api.github.com/repos/0nikod/danbooru_tag_process/releases/latest" + + async with aiohttp.ClientSession() as session: + headers = {"User-Agent": "ComfyUI-Simple-Prompt"} + async with session.get(release_url, headers=headers) as resp: + if resp.status != 200: + raise RuntimeError(f"Failed to fetch release info: {resp.status}") + release_info = await resp.json() + + download_url = None + asset_names = ["tags_processed.parquet", "tags.parquet"] + for name in asset_names: + for asset in release_info.get("assets", []): + if asset.get("name") == name: + download_url = asset.get("browser_download_url") + break + if download_url: + break + + if not download_url: + raise RuntimeError("Tag data file not found") + + async with session.get(download_url) as resp: + if resp.status != 200: + raise RuntimeError(f"Failed to download: {resp.status}") + + content = await resp.read() + os.makedirs(config.DATA_DIR, exist_ok=True) + temp_path = config.TAGS_PARQUET_PATH + ".tmp" + with open(temp_path, "wb") as f: + f.write(content) + + if os.path.exists(config.TAGS_PARQUET_PATH): + os.remove(config.TAGS_PARQUET_PATH) + os.rename(temp_path, config.TAGS_PARQUET_PATH) + + reinit_duckdb() + return {"status": "success", "message": "Tags updated successfully"} + + +async def handle_update_data(action: str) -> Dict[str, str]: + """ + 处理批量更新数据请求 + + Args: + action: 操作类型 ("update_liked", "update_user") + + Returns: + 操作结果 + """ + conn = get_db_connection() + if not conn: + raise RuntimeError("Database unavailable") + + if action == "update_liked": + target_path = config.LIKED_TAGS_PATH + elif action == "update_user": + target_path = config.USER_TAGS_PATH + else: + raise ValueError("Unknown action") + + from simple_prompt.core.database import ensure_parquet_exists + + ensure_parquet_exists(target_path, config.TAGS_SCHEMA) + target_path_sql = target_path.replace("\\", "/") + + conn.execute("DROP TABLE IF EXISTS temp_target") + conn.execute(f"CREATE TABLE temp_target AS SELECT * FROM read_parquet('{target_path_sql}')") + + # Build source views + other_sources = [] + if action == "update_liked": + if os.path.exists(config.USER_TAGS_PATH): + other_sources.append( + f"SELECT name, category, post_count, alias, 2 as p FROM read_parquet('{config.USER_TAGS_PATH.replace(os.sep, '/')}')" + ) + if os.path.exists(config.DEFAULT_TAGS_PATH): + other_sources.append( + f"SELECT name, category, post_count, alias, 3 as p FROM read_parquet('{config.DEFAULT_TAGS_PATH.replace(os.sep, '/')}')" + ) + if os.path.exists(config.TAGS_PARQUET_PATH): + other_sources.append( + f"SELECT name, category, post_count, alias, 4 as p FROM read_parquet('{config.TAGS_PARQUET_PATH.replace(os.sep, '/')}')" + ) + + if other_sources: + union_q = " UNION ALL ".join(other_sources) + conn.execute(f"CREATE OR REPLACE VIEW others_raw AS {union_q}") + conn.execute(""" + CREATE OR REPLACE VIEW others_best AS + SELECT * EXCLUDE (rn) FROM ( + SELECT *, ROW_NUMBER() OVER (PARTITION BY name ORDER BY p ASC) as rn FROM others_raw + ) WHERE rn = 1 + """) + + conn.execute(""" + UPDATE temp_target + SET category = others_best.category, + post_count = others_best.post_count, + alias = others_best.alias + FROM others_best + WHERE temp_target.name = others_best.name + """) + + conn.execute("DROP VIEW others_best") + conn.execute("DROP VIEW others_raw") + + conn.execute(f"COPY temp_target TO '{target_path_sql}' (FORMAT PARQUET)") + conn.execute("DROP TABLE temp_target") + + reinit_duckdb() + + if action == "update_liked": + return {"status": "success", "message": "Liked tags data updated from main DB."} + else: + return {"status": "success", "message": "User tags data updated."} diff --git a/simple_prompt/comfyui/__init__.py b/simple_prompt/comfyui/__init__.py new file mode 100644 index 0000000..927fb82 --- /dev/null +++ b/simple_prompt/comfyui/__init__.py @@ -0,0 +1,2 @@ +# simple_prompt.comfyui package +# ComfyUI 特定适配层 diff --git a/simple_prompt/comfyui/nodes.py b/simple_prompt/comfyui/nodes.py new file mode 100644 index 0000000..d2a7bd3 --- /dev/null +++ b/simple_prompt/comfyui/nodes.py @@ -0,0 +1,67 @@ +""" +ComfyUI 节点定义 +参考: https://docs.comfy.org/custom-nodes/backend/server_overview +""" + + +class SimplePrompt: + """ + A simple prompt node with an advanced visual tag editor. + 提供一个带有高级可视化标签编辑器的简单提示词节点。 + """ + + DESCRIPTION = "A prompt builder node with an advanced visual editor for tag management, weight adjustment, and autocomplete." + + # 节点分类（在 ComfyUI 菜单中的位置） + CATEGORY = "SimplePrompt" + + # 执行函数名称 + FUNCTION = "run" + + # 返回类型 + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("prompt",) + + # 是否为输出节点（会在执行时被调用） + OUTPUT_NODE = False + + @classmethod + def INPUT_TYPES(cls): + """ + 定义节点输入 + """ + return { + "required": { + "prompt_text": ("STRING", {"multiline": True, "dynamicPrompts": True}), + }, + } + + def run(self, prompt_text: str) -> tuple[str]: + """ + 节点执行函数 + + Args: + prompt_text: 提示词文本 + + Returns: + 包含提示词的元组 + """ + return (prompt_text,) + + @classmethod + def IS_CHANGED(cls, prompt_text: str) -> float: + """ + 判断节点是否需要重新执行 + 返回 NaN 表示每次都执行 + """ + return float("NaN") + + +# -------------------------------------------------------------------------------- +# 节点注册（符合 ComfyUI Lifecycle 规范） +# https://docs.comfy.org/custom-nodes/backend/lifecycle +# -------------------------------------------------------------------------------- + +NODE_CLASS_MAPPINGS = {"SimplePrompt": SimplePrompt} + +NODE_DISPLAY_NAME_MAPPINGS = {"SimplePrompt": "Simple Prompt"} diff --git a/simple_prompt/comfyui/server.py b/simple_prompt/comfyui/server.py new file mode 100644 index 0000000..7637d03 --- /dev/null +++ b/simple_prompt/comfyui/server.py @@ -0,0 +1,307 @@ +""" +ComfyUI PromptServer 适配层 +所有 ComfyUI 特定的 API 路由注册都在这里 + +支持热重载：设置环境变量 SIMPLE_PROMPT_DEV=1 启用开发模式 +""" + +import logging +import os + +from aiohttp import web + +logger = logging.getLogger("SimplePrompt") + +# 检测开发模式 +DEV_MODE = os.environ.get("SIMPLE_PROMPT_DEV", "").lower() in ("1", "true", "yes") + +if DEV_MODE: + logger.info("[SimplePrompt] Development mode enabled - hot reload active") + + +def _maybe_reload(): + """ + 开发模式下重新加载核心模块 + 这使得修改代码后不需要重启 ComfyUI + """ + if not DEV_MODE: + return + + import importlib + + try: + # 重新加载核心模块 + from simple_prompt.core import config, utils, database, tags, categories, presets + from simple_prompt import core + from simple_prompt.api import handlers + from simple_prompt import api + + importlib.reload(config) + importlib.reload(utils) + importlib.reload(database) + importlib.reload(tags) + importlib.reload(categories) + importlib.reload(presets) + importlib.reload(core) + importlib.reload(handlers) + importlib.reload(api) + + logger.debug("[DEV] Core modules reloaded") + except Exception as e: + logger.error(f"[DEV] Failed to reload modules: {e}") + + +def _get_handlers(): + """ + 获取处理器模块（开发模式下会先重载） + """ + _maybe_reload() + from simple_prompt.api import handlers + + return handlers + + +# -------------------------------------------------------------------------------- +# 路由注册 +# -------------------------------------------------------------------------------- + +try: + from server import PromptServer + + # ------------------------- + # Tags API + # ------------------------- + + @PromptServer.instance.routes.get("/simple-prompt/search-tags") + async def search_tags_api(request): + """搜索标签""" + handlers = _get_handlers() + + query = request.query.get("q", "") + limit = int(request.query.get("limit", 50)) + use_aliases = request.query.get("use_aliases", "false").lower() == "true" + categories_str = request.query.get("categories", "") + + categories = [] + if categories_str: + try: + categories = [int(c) for c in categories_str.split(",") if c.strip()] + except ValueError: + pass + + try: + results = await handlers.handle_search_tags(query, limit, use_aliases, categories) + return web.json_response(results) + except Exception as e: + logger.error(f"Search API error: {e}") + return web.json_response({"error": str(e)}, status=500) + + @PromptServer.instance.routes.post("/simple-prompt/get-tags-details") + async def get_tags_details_api(request): + """获取标签详情""" + handlers = _get_handlers() + + try: + data = await request.json() + names = data.get("names", []) + fast = data.get("fast", False) + + if not names: + return web.json_response({}) + + results = await handlers.handle_get_tags_details(names, fast) + return web.json_response(results) + except Exception as e: + logger.error(f"Get tags details API error: {e}") + return web.json_response({"error": str(e)}, status=500) + + @PromptServer.instance.routes.get("/simple-prompt/tags/list") + async def tags_list_api(request): + """列出标签""" + handlers = _get_handlers() + + try: + source = request.query.get("source", "user") + limit = int(request.query.get("limit", 50)) + offset = int(request.query.get("offset", 0)) + query = request.query.get("q", "") + + results = await handlers.handle_list_tags(source, limit, offset, query) + return web.json_response(results) + except Exception as e: + logger.error(f"Tags list error: {e}") + return web.json_response({"error": str(e)}, status=500) + + @PromptServer.instance.routes.delete("/simple-prompt/tags/delete") + async def tags_delete_api(request): + """删除标签""" + handlers = _get_handlers() + + try: + data = await request.json() + name = data.get("name") + source = data.get("source") + + result = await handlers.handle_delete_tag(name, source) + return web.json_response(result) + except ValueError as e: + return web.json_response({"error": str(e)}, status=400) + except Exception as e: + logger.error(f"Delete tag failed: {e}") + return web.json_response({"error": str(e)}, status=500) + + @PromptServer.instance.routes.post("/simple-prompt/add-custom-tag") + async def add_custom_tag_api(request): + """添加自定义标签""" + handlers = _get_handlers() + + try: + data = await request.json() + result = await handlers.handle_add_custom_tag(data) + return web.json_response(result) + except ValueError as e: + return web.json_response({"error": str(e)}, status=400) + except Exception as e: + logger.error(f"Add custom tag failed: {e}") + return web.json_response({"error": str(e)}, status=500) + + @PromptServer.instance.routes.post("/simple-prompt/toggle-like-tag") + async def toggle_like_tag_api(request): + """切换点赞标签""" + handlers = _get_handlers() + + try: + data = await request.json() + result = await handlers.handle_toggle_like_tag(data) + return web.json_response(result) + except ValueError as e: + return web.json_response({"error": str(e)}, status=400) + except Exception as e: + logger.error(f"Toggle like failed: {e}") + return web.json_response({"error": str(e)}, status=500) + + # ------------------------- + # Categories API + # ------------------------- + + @PromptServer.instance.routes.get("/simple-prompt/categories/list") + async def list_categories_api(request): + """列出分类""" + handlers = _get_handlers() + + try: + result = await handlers.handle_list_categories() + return web.json_response(result) + except Exception as e: + logger.error(f"List categories error: {e}") + return web.json_response({"error": str(e)}, status=500) + + @PromptServer.instance.routes.post("/simple-prompt/categories/save") + async def save_categories_api(request): + """保存分类""" + handlers = _get_handlers() + + try: + data = await request.json() + categories = data.get("categories", []) + result = await handlers.handle_save_categories(categories) + return web.json_response(result) + except Exception as e: + logger.error(f"Save categories error: {e}") + return web.json_response({"error": str(e)}, status=500) + + # ------------------------- + # Presets API + # ------------------------- + + @PromptServer.instance.routes.get("/simple-prompt/presets/list") + async def list_presets_api(request): + """列出预设""" + handlers = _get_handlers() + + try: + result = await handlers.handle_list_presets() + return web.json_response(result) + except Exception as e: + logger.error(f"List presets error: {e}") + return web.json_response({"error": str(e)}, status=500) + + @PromptServer.instance.routes.post("/simple-prompt/presets/save") + async def save_presets_api(request): + """保存预设""" + handlers = _get_handlers() + + try: + data = await request.json() + presets = data.get("presets", []) + result = await handlers.handle_save_presets(presets) + return web.json_response(result) + except Exception as e: + logger.error(f"Save presets error: {e}") + return web.json_response({"error": str(e)}, status=500) + + # ------------------------- + # Health & Update API + # ------------------------- + + @PromptServer.instance.routes.get("/simple-prompt/health") + async def health_check_api(request): + """健康检查""" + handlers = _get_handlers() + + try: + result = await handlers.handle_health_check() + return web.json_response(result) + except Exception as e: + return web.json_response({"error": str(e)}, status=500) + + @PromptServer.instance.routes.get("/simple-prompt/check-update") + async def check_update_api(request): + """检查更新""" + handlers = _get_handlers() + + try: + result = await handlers.handle_check_update() + return web.json_response(result) + except Exception as e: + logger.error(f"Check update error: {e}") + return web.json_response({"error": str(e)}, status=500) + + @PromptServer.instance.routes.post("/simple-prompt/update-tags") + async def update_tags_api(request): + """更新标签数据""" + handlers = _get_handlers() + + try: + try: + data = await request.json() + except Exception: + data = {} + + url = data.get("url") + result = await handlers.handle_update_tags(url) + return web.json_response(result) + except Exception as e: + logger.error(f"Update tags error: {e}") + return web.json_response({"error": str(e)}, status=500) + + @PromptServer.instance.routes.post("/simple-prompt/update-data") + async def update_data_api(request): + """批量更新数据""" + handlers = _get_handlers() + + try: + data = await request.json() + action = data.get("action") + result = await handlers.handle_update_data(action) + return web.json_response(result) + except ValueError as e: + return web.json_response({"error": str(e)}, status=400) + except Exception as e: + logger.error(f"Update data failed: {e}") + return web.json_response({"error": str(e)}, status=500) + + logger.info("[SimplePrompt] API routes registered successfully") + +except ImportError: + logger.warning("[SimplePrompt] PromptServer not available - running outside ComfyUI") diff --git a/simple_prompt/core/__init__.py b/simple_prompt/core/__init__.py new file mode 100644 index 0000000..faeb31e --- /dev/null +++ b/simple_prompt/core/__init__.py @@ -0,0 +1,30 @@ +# simple_prompt.core package +# 核心业务逻辑模块（与平台无关） + +from .config import ( + BASE_DIR, + DATA_DIR, + CUSTOM_DATA_DIR, + TAGS_PARQUET_PATH, + DEFAULT_TAGS_PATH, + USER_TAGS_PATH, + LIKED_TAGS_PATH, + DEFAULT_CATEGORY_PATH, + CUSTOM_CATEGORY_PATH, + DEFAULT_PRESETS_PATH, + CUSTOM_PRESETS_PATH, +) + +__all__ = [ + "BASE_DIR", + "DATA_DIR", + "CUSTOM_DATA_DIR", + "TAGS_PARQUET_PATH", + "DEFAULT_TAGS_PATH", + "USER_TAGS_PATH", + "LIKED_TAGS_PATH", + "DEFAULT_CATEGORY_PATH", + "CUSTOM_CATEGORY_PATH", + "DEFAULT_PRESETS_PATH", + "CUSTOM_PRESETS_PATH", +] diff --git a/simple_prompt/core/categories.py b/simple_prompt/core/categories.py new file mode 100644 index 0000000..9c9cdc6 --- /dev/null +++ b/simple_prompt/core/categories.py @@ -0,0 +1,106 @@ +""" +分类管理模块 +处理标签分类的加载和保存 +""" + +import logging +from typing import Any, Dict, List + +from . import config +from .utils import load_json, save_json + +logger = logging.getLogger("SimplePrompt") + +# 默认分类（硬编码备份） +DEFAULT_CATEGORIES = [ + {"id": 0, "name": "General", "color": "#0075db"}, + {"id": 1, "name": "Artist", "color": "#c00004"}, + {"id": 3, "name": "Copyright", "color": "#c000c0"}, + {"id": 4, "name": "Character", "color": "#00aa00"}, + {"id": 5, "name": "Meta", "color": "#fd9200"}, + {"id": 6, "name": "Special", "color": "#ffd700"}, # Gold + {"id": 7, "name": "Rating", "color": "#ff69b4"}, # HotPink +] + + +def get_categories() -> List[Dict[str, Any]]: + """ + 获取合并后的分类列表（默认 + 自定义） + + Returns: + 分类列表，按 ID 排序 + """ + defaults = load_json(config.DEFAULT_CATEGORY_PATH, []) + customs = load_json(config.CUSTOM_CATEGORY_PATH, []) + + # If default is empty (first run?), create it based on hardcoded defaults + if not defaults: + defaults = DEFAULT_CATEGORIES.copy() + try: + save_json(config.DEFAULT_CATEGORY_PATH, defaults) + except Exception: + pass + + # Merge: Custom overrides Default if same ID + cat_map: Dict[int, Dict[str, Any]] = {c["id"]: c for c in defaults} + for c in customs: + cat_map[c["id"]] = c + + # Return as list, sorted by ID + return sorted(cat_map.values(), key=lambda x: x["id"]) + + +def save_custom_categories(categories: List[Dict[str, Any]]) -> None: + """ + 保存自定义分类 + + Args: + categories: 分类列表 + + Raises: + Exception: 保存失败时抛出异常 + """ + save_json(config.CUSTOM_CATEGORY_PATH, categories) + + +def get_category_by_id(category_id: int) -> Dict[str, Any]: + """ + 根据 ID 获取分类信息 + + Args: + category_id: 分类 ID + + Returns: + 分类信息字典，如果不存在则返回默认值 + """ + categories = get_categories() + for cat in categories: + if cat["id"] == category_id: + return cat + return {"id": category_id, "name": "Unknown", "color": "#888888"} + + +def get_category_color(category_id: int) -> str: + """ + 获取分类颜色 + + Args: + category_id: 分类 ID + + Returns: + 颜色值（十六进制） + """ + return get_category_by_id(category_id).get("color", "#888888") + + +def get_category_name(category_id: int) -> str: + """ + 获取分类名称 + + Args: + category_id: 分类 ID + + Returns: + 分类名称 + """ + return get_category_by_id(category_id).get("name", "Unknown") diff --git a/simple_prompt/core/config.py b/simple_prompt/core/config.py new file mode 100644 index 0000000..a2fa2ba --- /dev/null +++ b/simple_prompt/core/config.py @@ -0,0 +1,55 @@ +""" +路径配置模块 +集中管理所有数据文件路径 +""" + +import os + +# 获取项目根目录（ComfyUI-Simple-Prompt） +# 路径: simple_prompt/core/config.py -> simple_prompt/core -> simple_prompt -> ComfyUI-Simple-Prompt +BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +# 数据目录 +DATA_DIR = os.path.join(BASE_DIR, "data") +CUSTOM_DATA_DIR = os.path.join(DATA_DIR, "custom") + +# 确保自定义数据目录存在 +os.makedirs(CUSTOM_DATA_DIR, exist_ok=True) + +# -------------------------------------------------------------------------------- +# Primary Data Sources (按优先级排序) +# -------------------------------------------------------------------------------- + +# Priority 1: 用户点赞的标签 +LIKED_TAGS_PATH = os.path.join(CUSTOM_DATA_DIR, "liked_tags.parquet") + +# Priority 2: 用户自定义标签 +USER_TAGS_PATH = os.path.join(CUSTOM_DATA_DIR, "user_tags.parquet") + +# Priority 3: 默认标签 +DEFAULT_TAGS_PATH = os.path.join(DATA_DIR, "default_tags.parquet") + +# Priority 4: 主仓库标签（Danbooru） +TAGS_PARQUET_PATH = os.path.join(DATA_DIR, "tags.parquet") + +# -------------------------------------------------------------------------------- +# Category & Meta Data Paths +# -------------------------------------------------------------------------------- + +# 默认分类配置 +DEFAULT_CATEGORY_PATH = os.path.join(DATA_DIR, "default_category.json") + +# 用户自定义分类配置 +CUSTOM_CATEGORY_PATH = os.path.join(CUSTOM_DATA_DIR, "custom_category.json") + +# 默认预设配置 +DEFAULT_PRESETS_PATH = os.path.join(DATA_DIR, "default_presets.json") + +# 用户自定义预设配置 +CUSTOM_PRESETS_PATH = os.path.join(CUSTOM_DATA_DIR, "custom_presets.json") + +# -------------------------------------------------------------------------------- +# Tags Schema +# -------------------------------------------------------------------------------- + +TAGS_SCHEMA = "name VARCHAR, category BIGINT, post_count BIGINT, alias VARCHAR[]" diff --git a/simple_prompt/core/database.py b/simple_prompt/core/database.py new file mode 100644 index 0000000..fede46d --- /dev/null +++ b/simple_prompt/core/database.py @@ -0,0 +1,167 @@ +""" +DuckDB 数据库管理模块 +处理数据库连接和初始化 +""" + +import logging +import os +from typing import Optional, Any + +from . import config +from .utils import get_sql_path + +logger = logging.getLogger("SimplePrompt") + +# Lazy import duckdb +try: + import duckdb +except ImportError: + logger.error("Error: 'duckdb' package not found in this Python environment.") + logger.error("Please install it using: pip install duckdb") + duckdb = None + +# 全局数据库连接 +_db_conn: Optional[Any] = None + + +def get_db_connection() -> Optional[Any]: + """ + 获取 DuckDB 连接（单例模式） + + Returns: + DuckDB 连接对象，或 None（如果 duckdb 未安装） + """ + global _db_conn + if _db_conn: + return _db_conn + + if not duckdb: + return None + + try: + _db_conn = duckdb.connect(database=":memory:") + return _db_conn + except Exception as e: + logger.error(f"Could not create DuckDB connection: {e}") + return None + + +def reset_db_connection() -> None: + """ + 重置数据库连接（用于热重载或重新初始化） + """ + global _db_conn + if _db_conn: + try: + _db_conn.close() + except Exception: + pass + _db_conn = None + + +def init_duckdb() -> None: + """ + 初始化 DuckDB 数据库 + 创建统一视图来合并多个数据源 + """ + conn = get_db_connection() + if not conn: + return + + logger.info("Initializing SimplePrompt Database...") + + sources = [] + + # Priority 1: Liked Tags + path = get_sql_path(config.LIKED_TAGS_PATH) + if path: + sources.append(f"SELECT name, category, post_count, alias, 1 as priority FROM read_parquet('{path}')") + + # Priority 2: User Tags + path = get_sql_path(config.USER_TAGS_PATH) + if path: + sources.append(f"SELECT name, category, post_count, alias, 2 as priority FROM read_parquet('{path}')") + + # Priority 3: Default Tags + path = get_sql_path(config.DEFAULT_TAGS_PATH) + if path: + sources.append(f"SELECT name, category, post_count, alias, 3 as priority FROM read_parquet('{path}')") + + # Priority 4: Repo Tags (Main) + path = get_sql_path(config.TAGS_PARQUET_PATH) + if path: + sources.append(f"SELECT name, category, post_count, alias, 4 as priority FROM read_parquet('{path}')") + else: + logger.warning(f"Main tags.parquet not found at {config.TAGS_PARQUET_PATH}") + + if not sources: + logger.warning("No tag data sources found.") + return + + # Create Unified View + try: + union_query = " UNION ALL ".join(sources) + + # 1. Create Raw View containing all duplicates + conn.execute(f"CREATE OR REPLACE VIEW all_tags_raw AS {union_query}") + + # 2. Create Deduplicated View (High priority wins) + conn.execute(""" + CREATE OR REPLACE VIEW tags AS + SELECT * EXCLUDE (rn) FROM ( + SELECT *, + ROW_NUMBER() OVER (PARTITION BY name ORDER BY priority ASC) as rn + FROM all_tags_raw + ) WHERE rn = 1 + """) + + # 3. Create Fast View (Excluding Repo Tags - Priority 4) + conn.execute(""" + CREATE OR REPLACE VIEW tags_fast AS + SELECT * EXCLUDE (rn) FROM ( + SELECT *, + ROW_NUMBER() OVER (PARTITION BY name ORDER BY priority ASC) as rn + FROM all_tags_raw + WHERE priority < 4 + ) WHERE rn = 1 + """) + + logger.info("DuckDB initialized with multi-source views (including fast view).") + + except Exception as e: + logger.error(f"DuckDB initialization failed: {e}") + + +def reinit_duckdb() -> None: + """ + 重新初始化数据库（用于数据更新后刷新视图） + """ + init_duckdb() + + +def ensure_parquet_exists(path: str, schema_sql: str) -> None: + """ + 确保 Parquet 文件存在，如果不存在则创建空文件 + + Args: + path: Parquet 文件路径 + schema_sql: 表结构的 SQL 定义 + """ + if os.path.exists(path): + return + + conn = get_db_connection() + if not conn: + return + + try: + # Create an empty table with the correct schema and write to parquet + temp_name = f"temp_init_{os.path.basename(path).replace('.', '_')}" + conn.execute(f"CREATE TABLE {temp_name} ({schema_sql})") + # Write empty table + save_path = path.replace("\\", "/") + conn.execute(f"COPY {temp_name} TO '{save_path}' (FORMAT PARQUET)") + conn.execute(f"DROP TABLE {temp_name}") + logger.info(f"Created empty parquet file: {path}") + except Exception as e: + logger.error(f"Failed to create parquet {path}: {e}") diff --git a/simple_prompt/core/presets.py b/simple_prompt/core/presets.py new file mode 100644 index 0000000..433a323 --- /dev/null +++ b/simple_prompt/core/presets.py @@ -0,0 +1,74 @@ +""" +预设管理模块 +处理 Meta 预设的加载和保存 +""" + +import logging +import os +from typing import Any, Dict, List + +from . import config +from .utils import load_json, save_json + +logger = logging.getLogger("SimplePrompt") + + +def get_presets() -> Dict[str, List[Dict[str, Any]]]: + """ + 获取预设列表（默认 + 自定义） + + Returns: + 包含 defaults 和 customs 的字典 + """ + defaults = load_json(config.DEFAULT_PRESETS_PATH, []) + customs = load_json(config.CUSTOM_PRESETS_PATH, []) + + # Ensure default file exists + if not os.path.exists(config.DEFAULT_PRESETS_PATH): + defaults = [] + try: + save_json(config.DEFAULT_PRESETS_PATH, defaults) + except Exception: + pass + + return {"defaults": defaults, "customs": customs} + + +def save_custom_presets(presets: List[Dict[str, Any]]) -> None: + """ + 保存自定义预设 + + Args: + presets: 预设列表 + + Raises: + Exception: 保存失败时抛出异常 + """ + save_json(config.CUSTOM_PRESETS_PATH, presets) + + +def get_all_presets() -> List[Dict[str, Any]]: + """ + 获取所有预设（合并默认和自定义） + + Returns: + 预设列表 + """ + data = get_presets() + return data["defaults"] + data["customs"] + + +def get_preset_by_id(preset_id: str) -> Dict[str, Any] | None: + """ + 根据 ID 获取预设 + + Args: + preset_id: 预设 ID + + Returns: + 预设信息，如果不存在则返回 None + """ + for preset in get_all_presets(): + if preset.get("id") == preset_id: + return preset + return None diff --git a/simple_prompt/core/tags.py b/simple_prompt/core/tags.py new file mode 100644 index 0000000..3ce2758 --- /dev/null +++ b/simple_prompt/core/tags.py @@ -0,0 +1,400 @@ +""" +标签管理模块 +处理标签的搜索、查询、添加、删除等操作 +""" + +import logging +import os +from typing import Any, Dict, List, Optional + +from . import config +from .database import ensure_parquet_exists, get_db_connection, reinit_duckdb + +logger = logging.getLogger("SimplePrompt") + + +def search_tags( + conn: Any, query: str, limit: int = 100, use_aliases: bool = False, categories: Optional[List[int]] = None +) -> List[Dict[str, Any]]: + """ + 搜索标签 + + Args: + conn: DuckDB 连接 + query: 搜索关键词 + limit: 返回结果数量限制 + use_aliases: 是否同时搜索别名 + categories: 限制搜索的分类 ID 列表 + + Returns: + 匹配的标签列表 + """ + if not conn: + return [] + + params: List[Any] = [] + search_term = f"%{query}%" + params.append(search_term) + + alias_clause = "" + if use_aliases: + # Searching inside the list of aliases + alias_clause = "OR EXISTS (SELECT 1 FROM unnest(alias) AS a(alias_value) WHERE alias_value ILIKE ?)" + params.append(search_term) + + category_clause = "" + if categories: + placeholders = ",".join(["?"] * len(categories)) + category_clause = f"AND category IN ({placeholders})" + params.extend(categories) + + # Add limit to params + params.append(limit) + + sql = f""" + SELECT * FROM tags + WHERE (name ILIKE ? {alias_clause}) + {category_clause} + ORDER BY priority ASC, post_count DESC + LIMIT ? + """ + + try: + res = conn.execute(sql, params) + cols = [desc[0] for desc in res.description] + results = [] + for row in res.fetchall(): + row_dict = {} + for i, val in enumerate(row): + row_dict[cols[i]] = val + results.append(row_dict) + return results + except Exception as e: + logger.error(f"Search failed: {e}") + return [] + + +def get_tags_details(conn: Any, names: List[str], fast: bool = False) -> Dict[str, int]: + """ + 获取标签详情（主要是分类信息） + + 处理空格和下划线格式： + - 输入 "blue hair" 会匹配数据库中的 "blue_hair" + - 输入 "blue_hair" 也能正常工作 + + Args: + conn: DuckDB 连接 + names: 标签名称列表 + fast: 是否使用快速视图（排除主仓库标签） + + Returns: + 标签名到分类 ID 的映射字典 + """ + if not conn or not names: + return {} + + table_name = "tags_fast" if fast else "tags" + + # Cleanup names + clean_names = [n.strip() for n in names if n.strip()] + if not clean_names: + return {} + + # Create a mapping of normalized names (underscore) to original names + normalized_to_original: Dict[str, str] = {} + all_variants: set[str] = set() + + for name in clean_names: + # Pre-clean: Remove backslashes (e.g. "aaa \(bbb\)" -> "aaa (bbb)") + cleaned_input = name.replace("\\", "") + lower_name = cleaned_input.lower() + + # Add both space and underscore variants + space_variant = lower_name.replace("_", " ") + underscore_variant = lower_name.replace(" ", "_") + + all_variants.add(space_variant) + all_variants.add(underscore_variant) + + # Map both variants to the original name AND the cleaned name + # We need to map back to 'name' (original input) so the caller knows which input matched + normalized_to_original[space_variant] = name + normalized_to_original[underscore_variant] = name + + variant_list = list(all_variants) + placeholders = ",".join(["?"] * len(variant_list)) + + sql = f""" + SELECT name, category + FROM {table_name} + WHERE name COLLATE NOCASE IN ({placeholders}) + """ + + try: + res = conn.execute(sql, variant_list).fetchall() + + result_map: Dict[str, int] = {} + for row in res: + db_name = row[0].lower() + category = row[1] + + # Map back to all original input names that could match this + original_name = normalized_to_original.get(db_name) + if original_name: + result_map[original_name] = category + + # Also add the db name itself for direct lookups + result_map[db_name] = category + + # Add the other variant too (space <-> underscore) + space_variant = db_name.replace("_", " ") + underscore_variant = db_name.replace(" ", "_") + result_map[space_variant] = category + result_map[underscore_variant] = category + + return result_map + except Exception as e: + logger.error(f"Get tags details failed: {e}") + return {} + + +def get_source_path_by_name(source: str) -> Optional[str]: + """ + 根据来源名称获取对应的文件路径 + + Args: + source: 来源名称 ("user", "liked", "default") + + Returns: + 文件路径，如果来源无效则返回 None + """ + if source == "user": + return config.USER_TAGS_PATH + elif source == "liked": + return config.LIKED_TAGS_PATH + elif source == "default": + return config.DEFAULT_TAGS_PATH + return None + + +def list_tags(conn: Any, source: str, limit: int = 50, offset: int = 0, query: str = "") -> Dict[str, Any]: + """ + 列出指定来源的标签 + + Args: + conn: DuckDB 连接 + source: 数据来源 ("user", "liked", "default") + limit: 返回数量限制 + offset: 偏移量（分页） + query: 搜索关键词 + + Returns: + 包含 data 和 total 的字典 + """ + if not conn: + return {"data": [], "total": 0} + + path = get_source_path_by_name(source) + if not path or not os.path.exists(path): + return {"data": [], "total": 0} + + path_sql = path.replace("\\", "/") + + where_clause = "" + params: List[Any] = [] + if query: + where_clause = "WHERE name ILIKE ? OR EXISTS (SELECT 1 FROM unnest(alias) AS a(alias_value) WHERE alias_value ILIKE ?)" + params = [f"%{query}%", f"%{query}%"] + + try: + # Get Total Count + count_sql = f"SELECT COUNT(*) FROM read_parquet('{path_sql}') {where_clause}" + total = conn.execute(count_sql, params).fetchone()[0] + + # Get Data + sql = f""" + SELECT * FROM read_parquet('{path_sql}') + {where_clause} + ORDER BY post_count DESC, name ASC + LIMIT {limit} OFFSET {offset} + """ + + if params: + res = conn.execute(sql, params) + else: + res = conn.execute(sql) + + cols = [desc[0] for desc in res.description] + data = [dict(zip(cols, row)) for row in res.fetchall()] + + return {"data": data, "total": total} + except Exception as e: + logger.error(f"List tags failed: {e}") + return {"data": [], "total": 0} + + +def delete_tag(conn: Any, name: str, source: str) -> bool: + """ + 删除指定来源的标签 + + Args: + conn: DuckDB 连接 + name: 标签名称 + source: 数据来源 + + Returns: + 是否成功 + """ + if not conn: + return False + + path = get_source_path_by_name(source) + if not path: + return False + + ensure_parquet_exists(path, config.TAGS_SCHEMA) + path_sql = path.replace("\\", "/") + + try: + # Delete using temp table + conn.execute("DROP TABLE IF EXISTS temp_del_tags") + conn.execute(f"CREATE TABLE temp_del_tags AS SELECT * FROM read_parquet('{path_sql}')") + conn.execute("DELETE FROM temp_del_tags WHERE name = ?", [name]) + conn.execute(f"COPY temp_del_tags TO '{path_sql}' (FORMAT PARQUET)") + conn.execute("DROP TABLE temp_del_tags") + + reinit_duckdb() # Refresh view + return True + except Exception as e: + logger.error(f"Delete tag failed: {e}") + return False + + +def add_tags(conn: Any, tags_data: List[Dict[str, Any]], source: str = "user") -> int: + """ + 添加或更新标签 + + Args: + conn: DuckDB 连接 + tags_data: 标签数据列表，每项包含 name, category, post_count, alias + source: 目标数据来源 + + Returns: + 成功添加的标签数量 + """ + if not conn or not tags_data: + return 0 + + target_path = get_source_path_by_name(source) + if not target_path: + return 0 + + ensure_parquet_exists(target_path, config.TAGS_SCHEMA) + target_path_sql = target_path.replace("\\", "/") + + try: + # Create temp table from existing parquet + conn.execute("DROP TABLE IF EXISTS temp_edit_tags") + conn.execute(f"CREATE TABLE temp_edit_tags AS SELECT * FROM read_parquet('{target_path_sql}')") + + count_added = 0 + + for item in tags_data: + name = item.get("name") + if not name: + continue + + category = int(item.get("category", 0)) + post_count = int(item.get("post_count", 0)) + alias = item.get("alias", []) + + # Delete existing if any (overwrite logic) + conn.execute("DELETE FROM temp_edit_tags WHERE name = ?", [name]) + + # Insert new tag + conn.execute("INSERT INTO temp_edit_tags VALUES (?, ?, ?, ?)", [name, category, post_count, alias]) + count_added += 1 + + # Write back to parquet + conn.execute(f"COPY temp_edit_tags TO '{target_path_sql}' (FORMAT PARQUET)") + conn.execute("DROP TABLE temp_edit_tags") + + # Refresh Main View + reinit_duckdb() + + return count_added + except Exception as e: + logger.error(f"Add tags failed: {e}") + return 0 + + +def toggle_like_tag( + conn: Any, + name: str, + should_like: bool = True, + category: Optional[int] = None, + post_count: Optional[int] = None, + alias: Optional[List[str]] = None, +) -> bool: + """ + 切换标签的点赞状态 + + Args: + conn: DuckDB 连接 + name: 标签名称 + should_like: True 为点赞，False 为取消点赞 + category: 标签分类（可选，会自动查找） + post_count: 帖子数量（可选） + alias: 别名列表（可选） + + Returns: + 是否成功 + """ + if not conn or not name: + return False + + ensure_parquet_exists(config.LIKED_TAGS_PATH, config.TAGS_SCHEMA) + liked_path_sql = config.LIKED_TAGS_PATH.replace("\\", "/") + + try: + conn.execute("DROP TABLE IF EXISTS temp_liked_tags") + conn.execute(f"CREATE TABLE temp_liked_tags AS SELECT * FROM read_parquet('{liked_path_sql}')") + + # Always remove first to avoid duplicates + conn.execute("DELETE FROM temp_liked_tags WHERE name = ?", [name]) + + if should_like: + # If data is missing, try to find it in other sources + if category is None or post_count is None: + for source_path in [config.USER_TAGS_PATH, config.DEFAULT_TAGS_PATH, config.TAGS_PARQUET_PATH]: + if os.path.exists(source_path): + src_sql = source_path.replace("\\", "/") + res = conn.execute( + f"SELECT category, post_count, alias FROM read_parquet('{src_sql}') WHERE name = ?", [name] + ).fetchone() + if res: + category = res[0] + post_count = res[1] + alias = res[2] + break + + # Default values if still missing + if category is None: + category = 0 + if post_count is None: + post_count = 0 + if alias is None: + alias = [] + + conn.execute("INSERT INTO temp_liked_tags VALUES (?, ?, ?, ?)", [name, category, post_count, alias]) + + # Write back + conn.execute(f"COPY temp_liked_tags TO '{liked_path_sql}' (FORMAT PARQUET)") + conn.execute("DROP TABLE temp_liked_tags") + + # Refresh View + reinit_duckdb() + return True + except Exception as e: + logger.error(f"Toggle like failed: {e}") + return False diff --git a/simple_prompt/core/utils.py b/simple_prompt/core/utils.py new file mode 100644 index 0000000..f8fccbc --- /dev/null +++ b/simple_prompt/core/utils.py @@ -0,0 +1,68 @@ +""" +通用工具模块 +包含 JSON 读写等通用功能 +""" + +import json +import logging +import os +from typing import Any + +logger = logging.getLogger("SimplePrompt") + + +def load_json(path: str, default: Any = None) -> Any: + """ + 加载 JSON 文件 + + Args: + path: 文件路径 + default: 文件不存在或加载失败时返回的默认值 + + Returns: + 解析后的 JSON 数据，或默认值 + """ + if not os.path.exists(path): + return default + try: + with open(path, "r", encoding="utf-8") as f: + return json.load(f) + except Exception as e: + logger.error(f"Failed to load JSON {path}: {e}") + return default + + +def save_json(path: str, data: Any) -> None: + """ + 保存数据到 JSON 文件 + + Args: + path: 文件路径 + data: 要保存的数据 + + Raises: + Exception: 保存失败时抛出异常 + """ + try: + # 确保目录存在 + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "w", encoding="utf-8") as f: + json.dump(data, f, indent=4, ensure_ascii=False) + except Exception as e: + logger.error(f"Failed to save JSON {path}: {e}") + raise e + + +def get_sql_path(path: str) -> str | None: + """ + 获取用于 SQL 查询的文件路径（将反斜杠转换为正斜杠） + + Args: + path: 文件路径 + + Returns: + 格式化后的路径，如果文件不存在则返回 None + """ + if os.path.exists(path): + return path.replace("\\", "/") + return None diff --git a/src/App.dev.vue b/src/App.dev.vue new file mode 100644 index 0000000..9a24328 --- /dev/null +++ b/src/App.dev.vue @@ -0,0 +1,89 @@ + + + + + diff --git a/src/App.vue b/src/App.vue new file mode 100644 index 0000000..c8b13d1 --- /dev/null +++ b/src/App.vue @@ -0,0 +1,108 @@ + + + + + diff --git a/src/components/AutocompleteList.vue b/src/components/AutocompleteList.vue new file mode 100644 index 0000000..4cbd55d --- /dev/null +++ b/src/components/AutocompleteList.vue @@ -0,0 +1,173 @@ + + + + + diff --git a/src/components/CustomTagModal.vue b/src/components/CustomTagModal.vue new file mode 100644 index 0000000..8cef70c --- /dev/null +++ b/src/components/CustomTagModal.vue @@ -0,0 +1,433 @@ + + + + + diff --git a/src/components/ModalWrapper.vue b/src/components/ModalWrapper.vue new file mode 100644 index 0000000..661a1ab --- /dev/null +++ b/src/components/ModalWrapper.vue @@ -0,0 +1,274 @@ + + + + + diff --git a/src/components/OtherFunctions.vue b/src/components/OtherFunctions.vue new file mode 100644 index 0000000..04c5328 --- /dev/null +++ b/src/components/OtherFunctions.vue @@ -0,0 +1,330 @@ + + + + + diff --git a/src/components/SettingsModal.vue b/src/components/SettingsModal.vue new file mode 100644 index 0000000..8915eb3 --- /dev/null +++ b/src/components/SettingsModal.vue @@ -0,0 +1,1128 @@ + + + + + diff --git a/src/components/TagItem.vue b/src/components/TagItem.vue new file mode 100644 index 0000000..01afb66 --- /dev/null +++ b/src/components/TagItem.vue @@ -0,0 +1,236 @@ + + + + + diff --git a/src/components/TagManager.vue b/src/components/TagManager.vue new file mode 100644 index 0000000..bd51a86 --- /dev/null +++ b/src/components/TagManager.vue @@ -0,0 +1,588 @@ + + + + + diff --git a/src/components/TagSearchModal.vue b/src/components/TagSearchModal.vue new file mode 100644 index 0000000..69f0855 --- /dev/null +++ b/src/components/TagSearchModal.vue @@ -0,0 +1,687 @@ + + + + + diff --git a/src/components/TextEditor.vue b/src/components/TextEditor.vue new file mode 100644 index 0000000..026c137 --- /dev/null +++ b/src/components/TextEditor.vue @@ -0,0 +1,849 @@ + + + + + diff --git a/src/components/ToolBar.vue b/src/components/ToolBar.vue new file mode 100644 index 0000000..12d4696 --- /dev/null +++ b/src/components/ToolBar.vue @@ -0,0 +1,149 @@ + + + + + \ No newline at end of file diff --git a/src/components/VisualTagArea.vue b/src/components/VisualTagArea.vue new file mode 100644 index 0000000..cd5f317 --- /dev/null +++ b/src/components/VisualTagArea.vue @@ -0,0 +1,194 @@ + + + + + diff --git a/src/dev.ts b/src/dev.ts new file mode 100644 index 0000000..a0127de --- /dev/null +++ b/src/dev.ts @@ -0,0 +1,92 @@ +/** + * Development Environment Mock + * + * This file mocks the ComfyUI environment (app, api, nodes) to allow + * standalone development of the editor modal. + */ + +import { createApp, h } from 'vue'; +import i18n from './i18n'; +import AppDev from './App.dev.vue'; +// import SimplePromptEditor from './components/SimplePromptEditor.vue'; + +console.log("[Dev] Initializing Standalone Development Environment"); + +// Mock Global Variables commonly found in ComfyUI +(window as any).app = { + ui: { + settings: { + getSettingValue: (key: string, def: any) => { + console.log(`[Dev] getSettingValue: ${key}`, def); + // Return mock settings from localStorage or default + return JSON.parse(localStorage.getItem(key) || JSON.stringify(def)); + }, + setSettingValue: (key: string, val: any) => { + console.log(`[Dev] setSettingValue: ${key}`, val); + localStorage.setItem(key, JSON.stringify(val)); + } + } + } +}; + +(window as any).api = { + apiURL: '/api' // Proxy to backend if needed +}; + +// Mock Node Object +const mockNode = { + id: 1, + type: 'SimplePrompt', + widgets: [ + { + name: 'prompt_text', + value: '1girl, smiling, (long hair:1.2), masterpiece', + callback: (v: string) => { console.log("[Dev] Widget Value Changed:", v); } + } + ], + addWidget: (type: string, name: string, value: any, callback: any) => { + console.log(`[Dev] addWidget: ${type}, ${name}`); + if (type === 'button') { + // Simulate button click to open editor immediately for debugging + const btn = document.createElement('button'); + btn.innerText = name; + btn.style.position = 'fixed'; + btn.style.top = '20px'; + btn.style.left = '20px'; + btn.style.padding = '10px 20px'; + btn.style.zIndex = '9999'; + btn.onclick = callback; + document.body.appendChild(btn); + + // Auto-click for faster dev loop + // callback(); + } + } +}; + +// Simulate Extension Initialization +async function initDev() { + // Temporarily just show a placeholder + const appContainer = document.getElementById('app'); + if (appContainer) { + // appContainer.innerHTML = '

Check console. Click "Edit Prompt" button (top-left) to simulate.

'; + const app = createApp(AppDev); + app.use(i18n); + app.mount('#app'); + } + + // Import main logic to trigger registration (but we need to mock app.registerExtension) + const mockApp = { + registerExtension: (ext: any) => { + console.log(`[Dev] Registered Extension: ${ext.name}`); + // Simulate node creation + ext.nodeCreated(mockNode, mockApp); + } + }; + + // We need to dynamically import main.ts or extract the logic + // For now, let's just expose the mockNode globally so we can test components directly + (window as any).mockNode = mockNode; +} + +initDev(); diff --git a/src/i18n.ts b/src/i18n.ts new file mode 100644 index 0000000..eabde82 --- /dev/null +++ b/src/i18n.ts @@ -0,0 +1,42 @@ +import { createI18n } from 'vue-i18n'; +import en from '../locales/en/main.json'; +import zh from '../locales/zh/main.json'; + +// Get initial locale +const getInitialLocale = () => { + // 1. Check localStorage (our app setting) + const stored = localStorage.getItem('simplePrompt.settings'); + if (stored) { + try { + const { language } = JSON.parse(stored); + if (language) return language; + } catch (e) { } + } + + // 2. Check ComfyUI Setting + if ((window as any).app?.ui?.settings) { + const comfyLang = (window as any).app.ui.settings.getSettingValue("Comfy.lang", ""); + if (comfyLang) { + if (comfyLang.startsWith('zh')) return 'zh-CN'; + return 'en'; + } + } + + // 3. Browser language + const browserLang = navigator.language; + if (browserLang.startsWith('zh')) return 'zh-CN'; + + return 'en'; +}; + +const i18n = createI18n({ + legacy: false, + locale: getInitialLocale(), + fallbackLocale: 'en', + messages: { + en: en, + 'zh-CN': zh + } +}); + +export default i18n; diff --git a/src/main.ts b/src/main.ts new file mode 100644 index 0000000..1c8a48e --- /dev/null +++ b/src/main.ts @@ -0,0 +1,97 @@ +import { app } from "../../../scripts/app.js"; +import { api } from "../../../scripts/api.js"; + +/** + * SimplePrompt Frontend Extension + * + * This extension injects the "Edit Prompt" button into the SimplePrompt node. + * It is designed to be lazily loaded and extensible. + */ +const ext = { + name: "SimplePrompt.Editor", + + async init(app) { + console.log("[SimplePrompt] Extension initialized"); + + // Load CSS + const link = document.createElement("link"); + link.rel = "stylesheet"; + link.type = "text/css"; + // Convert to absolute URL based on where the script is running + // import.meta.url points to js/main.js + // css is in js/assets/main.css + // relative path: ./assets/main.css + link.href = new URL("./assets/main.css", import.meta.url).href; + document.head.appendChild(link); + }, + + async nodeCreated(node, app) { + // Only target the SimplePrompt node + if (node.comfyClass === "SimplePrompt") { + try { + this.setupSimplePromptNode(node, app); + } catch (error) { + console.error("[SimplePrompt] Error setting up node:", error); + } + } + }, + + setupSimplePromptNode(node, app) { + // Add "Edit Prompt" button + // We use a custom widget name to identify it easily later + const editWidget = node.addWidget("button", "Edit Prompt", null, () => { + this.openEditorModal(node); + }); + + // Ensure prompt_text widget exists and is accessible + const promptWidget = node.widgets.find(w => w.name === "prompt_text"); + if (!promptWidget) { + console.warn("[SimplePrompt] prompt_text widget not found!"); + } + }, + + openEditorModal(node) { + console.log("[SimplePrompt] Opening editor for node:", node.id); + + // Get current prompt text + const promptWidget = node.widgets.find(w => w.name === "prompt_text"); + const currentPrompt = promptWidget ? promptWidget.value : ""; + + // Create container + const container = document.createElement("div"); + document.body.appendChild(container); + + // Import createApp directly (bundled) + import("vue").then(({ createApp }) => { + // Import App component + import("./App.vue").then((appModule) => { + const App = appModule.default; + + const appInstance = createApp(App, { + initialPrompt: currentPrompt, + onSave: (newPrompt) => { + if (promptWidget) { + promptWidget.value = newPrompt; + } + // Clean up + appInstance.unmount(); + document.body.removeChild(container); + }, + onClose: () => { + appInstance.unmount(); + document.body.removeChild(container); + } + }); + + // Use i18n if needed, or other plugins + import("./i18n").then((i18nModule) => { + const i18n = i18nModule.default; + appInstance.use(i18n); + appInstance.mount(container); + }); + }); + }); + } +}; + +app.registerExtension(ext); diff --git a/src/utils/caret.ts b/src/utils/caret.ts new file mode 100644 index 0000000..386957d --- /dev/null +++ b/src/utils/caret.ts @@ -0,0 +1,82 @@ +/** + * Utility to get caret coordinates in a textarea + * Based on basic mirror div technique + */ +export function getCaretCoordinates(element: HTMLTextAreaElement, position: number) { + const properties = [ + 'direction', + 'boxSizing', + 'width', + 'height', + 'overflowX', + 'overflowY', + 'borderTopWidth', + 'borderRightWidth', + 'borderBottomWidth', + 'borderLeftWidth', + 'paddingTop', + 'paddingRight', + 'paddingBottom', + 'paddingLeft', + 'fontStyle', + 'fontVariant', + 'fontWeight', + 'fontStretch', + 'fontSize', + 'fontSizeAdjust', + 'lineHeight', + 'fontFamily', + 'textAlign', + 'textTransform', + 'textIndent', + 'textDecoration', + 'letterSpacing', + 'wordSpacing', + 'tabSize', + 'MozTabSize', + ]; + + const div = document.createElement('div'); + div.id = 'input-textarea-caret-position-mirror-div'; + document.body.appendChild(div); + + const style = div.style; + const computed = window.getComputedStyle(element); + + style.whiteSpace = 'pre-wrap'; + if (element.nodeName !== 'INPUT') + style.wordWrap = 'break-word'; + + style.position = 'absolute'; + style.visibility = 'hidden'; + + properties.forEach((prop: any) => { + style[prop] = computed[prop]; + }); + + if (navigator.userAgent.toLowerCase().indexOf('firefox') > -1) { + if (element.scrollHeight > parseInt(computed.height)) + style.overflowY = 'scroll'; + } else { + style.overflow = 'hidden'; + } + + div.textContent = element.value.substring(0, position); + + if (element.nodeName === 'INPUT') + div.textContent = div.textContent.replace(/\s/g, '\u00a0'); + + const span = document.createElement('span'); + span.textContent = element.value.substring(position) || '.'; + div.appendChild(span); + + const coordinates = { + top: span.offsetTop + parseInt(computed['borderTopWidth']), + left: span.offsetLeft + parseInt(computed['borderLeftWidth']), + height: parseInt(computed['lineHeight']) + }; + + document.body.removeChild(div); + + return coordinates; +} diff --git a/src/utils/categoryService.ts b/src/utils/categoryService.ts new file mode 100644 index 0000000..e818238 --- /dev/null +++ b/src/utils/categoryService.ts @@ -0,0 +1,154 @@ +import { TagCategory, CATEGORY_COLORS } from './types'; +import { reactive } from 'vue'; + +export interface CategoryData { + value: number; + label: string; + color: string; +} + +/** + * Category Item used by UI components + */ +export interface CategoryItem { + id: number; + name: string; + color: string; +} + +export class CategoryService { + private static instance: CategoryService; + private _categories = reactive(new Map()); + private isInitialized = false; + + private constructor() { + // Initialize with defaults + this.addDefaultCategories(); + } + + public static getInstance(): CategoryService { + if (!CategoryService.instance) { + CategoryService.instance = new CategoryService(); + } + return CategoryService.instance; + } + + private addDefaultCategories() { + // Add standard Danbooru categories as baseline + this._categories.set(TagCategory.GENERAL, { value: TagCategory.GENERAL, label: 'General', color: CATEGORY_COLORS[TagCategory.GENERAL] }); + this._categories.set(TagCategory.ARTIST, { value: TagCategory.ARTIST, label: 'Artist', color: CATEGORY_COLORS[TagCategory.ARTIST] }); + this._categories.set(TagCategory.COPYRIGHT, { value: TagCategory.COPYRIGHT, label: 'Copyright', color: CATEGORY_COLORS[TagCategory.COPYRIGHT] }); + this._categories.set(TagCategory.CHARACTER, { value: TagCategory.CHARACTER, label: 'Character', color: CATEGORY_COLORS[TagCategory.CHARACTER] }); + this._categories.set(TagCategory.META, { value: TagCategory.META, label: 'Meta', color: CATEGORY_COLORS[TagCategory.META] }); + } + + public async init(): Promise { + if (this.isInitialized) return; + await this.fetchCategories(); + } + + public async fetchCategories(): Promise { + try { + const response = await fetch('/simple-prompt/categories/list'); + if (response.ok) { + const data = await response.json(); + + if (Array.isArray(data)) { + // Clear custom ones if any, but keep defaults if fetch fails next time? + // Better to just update the map. + data.forEach((cat: any) => { + this._categories.set(cat.id, { + value: cat.id, + label: cat.name, + color: cat.color + }); + }); + } + } + } catch (error) { + console.error("[CategoryService] Failed to fetch categories:", error); + } finally { + this.isInitialized = true; + } + } + + public async saveCustomCategories(categories: CategoryItem[]): Promise { + try { + const response = await fetch('/simple-prompt/categories/save', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ categories }) + }); + + if (response.ok) { + // Refresh local state after save + await this.fetchCategories(); + } else { + console.error("[CategoryService] Failed to save categories:", await response.text()); + throw new Error("Failed to save categories"); + } + } catch (error) { + console.error("[CategoryService] Error saving categories:", error); + throw error; + } + } + + public getCategory(id: number): CategoryData { + return this._categories.get(id) || { + value: id, + label: 'Unknown', + color: '#888888' + }; + } + + public getColor(id: number): string { + return this.getCategory(id).color; + } + + public getName(id: number): string { + return this.getCategory(id).label; + } + + public getAll(): CategoryData[] { + return Array.from(this._categories.values()).sort((a, b) => a.value - b.value); + } + + // Getter for template usage - returns array format expected by CustomTagModal + public get categoriesArray(): CategoryItem[] { + return this.getAll().map(cat => ({ + id: cat.value, + name: cat.label, + color: cat.color + })); + } + + public getCategoryColor(id: number): string { + return this.getColor(id); + } +} + +// Legacy compatibility export for components using `categoryService` instance +export const categoryService = { + get categories() { + return { + get value() { + return CategoryService.getInstance().categoriesArray; + } + }; + }, + getCategoryColor(id: number): string { + return CategoryService.getInstance().getColor(id); + }, + getCategoryName(id: number): string { + return CategoryService.getInstance().getName(id); + }, + async init() { + return CategoryService.getInstance().init(); + }, + async fetchCategories() { + return CategoryService.getInstance().fetchCategories(); + }, + async saveCustomCategories(categories: CategoryItem[]) { + return CategoryService.getInstance().saveCustomCategories(categories); + } +}; diff --git a/src/utils/duckdbService.ts b/src/utils/duckdbService.ts new file mode 100644 index 0000000..c1a6a21 --- /dev/null +++ b/src/utils/duckdbService.ts @@ -0,0 +1,93 @@ +export class DuckDBService { + private static instance: DuckDBService; + private isInitialized = false; + + private constructor() { } + + public static getInstance(): DuckDBService { + if (!DuckDBService.instance) { + DuckDBService.instance = new DuckDBService(); + } + return DuckDBService.instance; + } + + public async init(): Promise { + if (this.isInitialized) return; + console.log("[DuckDB] Backend initialization check (API-based)"); + // We could ping the backend here if needed + this.isInitialized = true; + } + + + + public async searchTags(query: string, limit: number = 20, useAliases: boolean = false, categories: number[] = []): Promise { + const params = new URLSearchParams({ + q: query, + limit: limit.toString(), + use_aliases: useAliases.toString(), + categories: categories.join(',') + }); + + try { + const response = await fetch(`/simple-prompt/search-tags?${params.toString()}`); + if (!response.ok) { + throw new Error(`HTTP error! status: ${response.status}`); + } + const results = await response.json(); + + // Handle the same alias matching logic as before for UI highlighting + if (useAliases && results.length > 0) { + return results.map((row: any) => { + const nameMatches = row.name && row.name.toLowerCase().includes(query.toLowerCase()); + + // If name matches, it's a direct match + if (nameMatches) { + return row; + } + + // Otherwise, check if an alias matched + const aliasArray = Array.isArray(row.alias) ? row.alias : []; + if (aliasArray.length > 0) { + const matchedAlias = aliasArray.find((a: string) => + a && a.toLowerCase().includes(query.toLowerCase()) + ); + + if (matchedAlias) { + return { ...row, matched_alias: matchedAlias }; + } + } + + return row; + }); + } + + return results; + } catch (error) { + console.error("[DuckDB] Search failed via API:", error); + return []; + } + } + + public async getTagsDetails(names: string[]): Promise> { + if (names.length === 0) return {}; + + try { + const response = await fetch('/simple-prompt/get-tags-details', { + method: 'POST', + headers: { + 'Content-Type': 'application/json' + }, + body: JSON.stringify({ names }) + }); + + if (!response.ok) { + throw new Error(`HTTP error! status: ${response.status}`); + } + return await response.json(); + } catch (error) { + console.error("[DuckDB] Get tags details failed:", error); + return {}; + } + } +} + diff --git a/src/utils/metaService.ts b/src/utils/metaService.ts new file mode 100644 index 0000000..693a960 --- /dev/null +++ b/src/utils/metaService.ts @@ -0,0 +1,83 @@ +import { ref } from 'vue'; + +class MetaService { + private static instance: MetaService; + // Preset support + public presets = ref([]); // { id, name, tags } + public defaultPresets = ref([]); + public customPresets = ref([]); + public activePresetId = ref(''); + + public loading = ref(false); + + private constructor() { } + + public static getInstance(): MetaService { + if (!MetaService.instance) { + MetaService.instance = new MetaService(); + } + return MetaService.instance; + } + + + // --- Preset API --- + + public async fetchPresets() { + try { + const response = await fetch('/simple-prompt/presets/list'); + if (response.ok) { + const data = await response.json(); + this.defaultPresets.value = data.defaults || []; + this.customPresets.value = data.customs || []; + + // Merge + this.presets.value = [...this.defaultPresets.value, ...this.customPresets.value]; + + // Restore active preset logic or default to first if none + if (!this.activePresetId.value && this.presets.value.length > 0) { + // Check storage? + // For now default to first + this.activePresetId.value = this.presets.value[0].id; + } + } + } catch (e) { + console.error("Error fetching presets:", e); + } + } + + public async saveCustomPresets(customs: any[]) { + try { + const response = await fetch('/simple-prompt/presets/save', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ presets: customs }) + }); + if (!response.ok) throw new Error("Failed"); + await this.fetchPresets(); + } catch (e) { + console.error("Error saving presets:", e); + throw e; + } + } + + public getActiveTags(): string[] { + // If we have an active preset, use it + if (this.activePresetId.value) { + const preset = this.presets.value.find(p => p.id === this.activePresetId.value); + if (preset && preset.tags) return preset.tags; + } + // Fallback to old behavior (all meta tags)? + // Or if no preset is active, use the 'metaTags' list? + // User wants "Auto Meta" to use "Selected Meta Group". + // Let's assume if Presets exist, we use them. + // If user wants "All", they can make an "All" preset? + // Or fallback to flat list if activePresetId is empty? + return []; + } + + public setActivePreset(id: string) { + this.activePresetId.value = id; + } +} + +export const metaService = MetaService.getInstance(); diff --git a/src/utils/promptParser.ts b/src/utils/promptParser.ts new file mode 100644 index 0000000..7f39446 --- /dev/null +++ b/src/utils/promptParser.ts @@ -0,0 +1,513 @@ +import type { TagItem } from './types'; +import { TagCategory } from './types'; + +/** + * Parse prompt text into TagItem array + * Matches (text:weight) format or plain comma-separated text + * Properly handles escaped parentheses \( and \) + */ +export function textToTags(text: string): TagItem[] { + if (!text.trim()) return []; + + const tags: TagItem[] = []; + let idx = 0; + let i = 0; + let currentTag = ''; + let inWeight = false; + let weightText = ''; + let tagText = ''; + + while (i < text.length) { + const char = text[i]; + + // Handle escape sequences + // matches \ followed by anything. + // We usually only care about \( and \) but let's be safe. + // If we see \, check next char. + if (char === '\\' && i + 1 < text.length) { + const nextChar = text[i + 1]; + if (nextChar === '(' || nextChar === ')') { + // It is an escaped parenthesis. + // We want to treat it as part of the tag name (literal). + // "star \(sky\)" -> tag text: "star (sky)" + currentTag += nextChar; + i += 2; + continue; + } else { + // Backslash followed by other char (e.g. \n or just \a) + // Treat backslash as literal? Or ignore? + // Mostly safe to treat as literal if not special. + // But for safety let's just add the char and skip backslash? + // Or add both? "C:\\path" -> "C:\path" + // Let's stick to strict behavior: Only escape parens. + // If not parens, treat \ as literal char. + currentTag += char; + i++; + continue; + } + } + + // Handle weighted tags: (text:weight) + if (char === '(' && !inWeight) { + // ... (rest is same) + // Start of weighted tag + // Save any accumulated plain text first + if (currentTag.trim()) { + tags.push({ + id: `tag-${idx++}`, + text: currentTag.trim(), + weight: 1.0, + enabled: true, + category: TagCategory.GENERAL + }); + currentTag = ''; + } + + // Find the closing parenthesis + let depth = 1; + let j = i + 1; + let escaped = false; + + while (j < text.length && depth > 0) { + if (text[j] === '\\' && j + 1 < text.length) { + escaped = true; + j += 2; + continue; + } + if (text[j] === '(') depth++; + if (text[j] === ')') depth--; + j++; + } + + if (depth === 0) { + // Found matching closing parenthesis + const content = text.substring(i + 1, j - 1); + const colonIndex = content.indexOf(':'); + + if (colonIndex > 0) { + // Has weight format + tagText = content.substring(0, colonIndex).trim(); + // Unescape parentheses in the tag text + tagText = tagText.replace(/\\([()])/g, '$1'); + + weightText = content.substring(colonIndex + 1).trim(); + const weight = parseFloat(weightText); + + if (!isNaN(weight) && tagText) { + tags.push({ + id: `tag-${idx++}`, + text: tagText, + weight: weight, + enabled: true, + category: TagCategory.GENERAL + }); + } + } else { + // No colon, treat whole thing as plain text with parens + currentTag += text.substring(i, j); + } + + i = j; + continue; + } + } + + // Handle comma or newline separator + if (char === ',' || char === '\n') { + if (currentTag.trim()) { + tags.push({ + id: `tag-${idx++}`, + text: currentTag.trim(), + weight: 1.0, + enabled: true, + category: TagCategory.GENERAL + }); + currentTag = ''; + } + i++; + continue; + } + + // Regular character + currentTag += char; + i++; + } + + // Don't forget the last tag + if (currentTag.trim()) { + tags.push({ + id: `tag-${idx++}`, + text: currentTag.trim(), + weight: 1.0, + enabled: true, + category: TagCategory.GENERAL + }); + } + + return tags; +} + +/** + * Convert TagItem array back to prompt text + */ +export function tagsToText(tags: TagItem[]): string { + const enabledTags = tags.filter(tag => tag.enabled); + if (enabledTags.length === 0) return ''; + + return enabledTags + .map(tag => { + // Escape parens in the text itself to avoid parser confusion, BUT avoid double escaping + const escapedText = tag.text + .replace(/(? (star \(sky\):1.2) + return `(${escapedText}:${tag.weight.toFixed(1)})`; + } + return escapedText; + }) + .join(', ') + ','; +} + +/** + * Get sort order for a category + */ +export function getSortOrder(cat: number): number { + // Row 0: Special(6) + // Row 1: Character(4), Copyright(2,3) + // Row 2: Artist(1) + // Row 3: General(0) + // Row 4: Meta(5), Rating(7) + const rowMap: Record = { + 6: 0, + 4: 1, 2: 1, 3: 1, + 1: 2, + 0: 3, + 5: 4, 7: 4 + }; + return rowMap[cat] !== undefined ? rowMap[cat] : 3; // Default to General +} + +/** + * Format tags into grouped multiline text + */ +export function formatGroupedTags(tagsList: TagItem[]): string { + if (tagsList.length === 0) return ''; + + // Group by row + const rows: Record = { 0: [], 1: [], 2: [], 3: [], 4: [] }; + tagsList.forEach(tag => { + const cat = tag.category !== undefined ? tag.category : 0; + const rowIndex = getSortOrder(cat); + rows[rowIndex].push(tag); + }); + + // Sort logic for Row 1 (Characters & Copyrights) + // Preference: Character(4) > Copyright(3) > Copyright(2) + // Using stable sort by category ID (descending) + rows[1].sort((a, b) => { + const catA = a.category !== undefined ? a.category : 0; + const catB = b.category !== undefined ? b.category : 0; + // Descending order: 4 (Character) > 3 (Copyright) > 2 (Copyright) + return catB - catA; + }); + + // Convert rows to text + const rowTexts = [0, 1, 2, 3, 4].map(idx => { + const rowTags = rows[idx]; + if (rowTags.length === 0) return ''; + return tagsToText(rowTags).replace(/,$/, ''); // Remove trailing comma from tagsToText result if we want (but usually we want comma between groups?) + // tagsToText adds a trailing comma. + // If we join by newline, we might want that trailing comma. + // Let's keep it. + }); + + // Build final output with specific blank line logic + // Users request: + // Row 0 + // Row 1 + // Row 2 + // (blank if General exists) + // Row 3 (General) + // (blank if Meta/Rating exists) + // Row 4 + + let result = ''; + // Rows 0, 1, 2 + const topPart = rowTexts.slice(0, 3).filter(t => t).join('\n'); + if (topPart) result += topPart; + + // Gap before General + if (result && rowTexts[3]) result += '\n\n'; + + if (rowTexts[3]) result += rowTexts[3]; + + // Gap after General + if (result && rowTexts[4]) result += '\n\n'; + + if (rowTexts[4]) result += rowTexts[4]; + + return result; +} + +/** + * Apply Auto Meta tags with minimal changes to the original text. + * + * Logic: + * 1. Find and remove existing occurrences of metaTags from the text + * 2. Find where ratingTags start (if any) + * 3. Insert meta block (in preset order) just before rating section (or at end) + * 4. Preserve all other text including newlines + * + * @param text Original prompt text + * @param metaTags List of meta tag names to ensure are present + * @param ratingTagNames List of known rating tag names (to identify where to insert) + * @returns Modified text + */ +export function applyAutoMeta( + text: string, + metaTags: string[], + ratingTagNames: string[] +): string { + if (!metaTags.length) return text; + + // Normalize for matching (lowercase) + const metaLower = new Set(metaTags.map(t => t.toLowerCase())); + const ratingLower = new Set(ratingTagNames.map(t => t.toLowerCase())); + + // Parse to find tag positions and existing weights + // We need to track: { tagText, startIndex, endIndex, weight } + interface TagMatch { + text: string; + fullMatch: string; + start: number; + end: number; + weight: number; + } + + const foundTags: TagMatch[] = []; + + // Regex patterns: + // 1. Weighted tag: (text:weight) + // 2. Plain tag: text between commas/newlines + // We need to carefully find positions. + + // Split by delimiters but keep track of positions + let i = 0; + let currentStart = 0; + let currentTag = ''; + + const processTag = (tagStr: string, fullMatch: string, start: number, end: number) => { + const trimmed = tagStr.trim(); + if (!trimmed) return; + + // Check for weight syntax: (text:weight) + const weightMatch = trimmed.match(/^\((.+):(\d+\.?\d*)\)$/); + let tagText = trimmed; + let weight = 1.0; + + if (weightMatch) { + tagText = weightMatch[1].trim(); + weight = parseFloat(weightMatch[2]); + } + + // Unescape parentheses for matching + const unescaped = tagText.replace(/\\([()])/g, '$1'); + + foundTags.push({ + text: unescaped, + fullMatch, + start, + end, + weight + }); + }; + + while (i < text.length) { + const char = text[i]; + + // Handle weighted tag start + if (char === '(' && currentTag.trim() === '') { + // Find matching close + let depth = 1; + let j = i + 1; + while (j < text.length && depth > 0) { + if (text[j] === '\\' && j + 1 < text.length) { + j += 2; + continue; + } + if (text[j] === '(') depth++; + if (text[j] === ')') depth--; + j++; + } + + if (depth === 0) { + const fullMatch = text.substring(i, j); + const content = text.substring(i + 1, j - 1); + processTag(fullMatch, fullMatch, i, j); + i = j; + currentStart = j; + currentTag = ''; + continue; + } + } + + // Handle escape + if (char === '\\' && i + 1 < text.length) { + currentTag += char + text[i + 1]; + i += 2; + continue; + } + + // Handle delimiter + if (char === ',' || char === '\n') { + if (currentTag.trim()) { + processTag(currentTag.trim(), currentTag, currentStart, i); + } + i++; + currentStart = i; + currentTag = ''; + continue; + } + + currentTag += char; + i++; + } + + // Last tag + if (currentTag.trim()) { + processTag(currentTag.trim(), currentTag, currentStart, i); + } + + // Identify which tags are meta and which are rating + const metaToRemove: TagMatch[] = []; + const existingMetaWeights = new Map(); + let ratingStartIndex = text.length; + + for (const tag of foundTags) { + const lower = tag.text.toLowerCase(); + + if (metaLower.has(lower)) { + metaToRemove.push(tag); + existingMetaWeights.set(lower, tag.weight); + } + + if (ratingLower.has(lower)) { + if (tag.start < ratingStartIndex) { + ratingStartIndex = tag.start; + } + } + } + + // Build result by removing meta tags + // Sort removals by position descending to avoid index shifting + const removals = [...metaToRemove].sort((a, b) => b.start - a.start); + + let result = text; + for (const removal of removals) { + // Find the actual range to remove (including surrounding comma/space) + let removeStart = removal.start; + let removeEnd = removal.end; + + // Look backward for comma/space + while (removeStart > 0 && /[\s,]/.test(result[removeStart - 1])) { + removeStart--; + } + + // Look forward for comma/space (but not newlines - preserve structure) + while (removeEnd < result.length && result[removeEnd] === ' ') { + removeEnd++; + } + if (removeEnd < result.length && result[removeEnd] === ',') { + removeEnd++; + // Also consume trailing space after comma + while (removeEnd < result.length && result[removeEnd] === ' ') { + removeEnd++; + } + } + + // If we're removing from the start of a line, don't leave orphan comma + // If there's content before and after, we need a comma + const before = result.substring(0, removeStart); + const after = result.substring(removeEnd); + + result = before + after; + + // Update ratingStartIndex if needed + const removedLength = removeEnd - removeStart; + if (ratingStartIndex > removal.start) { + ratingStartIndex -= removedLength; + } + } + + // Clean up any double commas or trailing commas before newlines + result = result.replace(/,\s*,/g, ','); + result = result.replace(/,\s*\n/g, ',\n'); + result = result.replace(/\n\s*,/g, '\n'); + + // Build meta block string + const metaBlock = metaTags.map(m => { + const weight = existingMetaWeights.get(m.toLowerCase()) || 1.0; + // Escape parens + const escaped = m.replace(/\(/g, '\\(').replace(/\)/g, '\\)'); + if (weight !== 1.0) { + return `(${escaped}:${weight.toFixed(1)})`; + } + return escaped; + }).join(', ') + ','; + + // Find insertion point (before rating section, or at end) + // We need to recalculate ratingStartIndex in the modified string + // Simpler approach: find first rating tag in modified text + + const modifiedTags = textToTags(result); + let insertIndex = result.length; + + for (let idx = 0; idx < modifiedTags.length; idx++) { + const lower = modifiedTags[idx].text.toLowerCase(); + if (ratingLower.has(lower)) { + // Find this tag's position in result + // This is complex... let's use a simpler approach + // Find the first occurrence of any rating tag + for (const ratingName of ratingTagNames) { + const escaped = ratingName.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); + const regex = new RegExp(`(^|[,\\n\\s])${escaped}([,\\n]|$)`, 'i'); + const match = result.match(regex); + if (match && match.index !== undefined) { + const pos = match.index + match[1].length; + if (pos < insertIndex) { + insertIndex = pos; + } + } + } + break; + } + } + + // Find proper insertion point (before rating, after last content) + // If there's content before insert point that doesn't end with newline, add newline + let prefix = result.substring(0, insertIndex).trimEnd(); + let suffix = result.substring(insertIndex); + + // Ensure proper separation + if (prefix && !prefix.endsWith(',') && !prefix.endsWith('\n')) { + prefix += ','; + } + + // Add newlines for separation if needed + if (prefix && !prefix.endsWith('\n\n')) { + if (prefix.endsWith('\n')) { + prefix += '\n'; + } else { + prefix += '\n\n'; + } + } + + if (suffix.trim() && !suffix.startsWith('\n')) { + suffix = '\n\n' + suffix.trimStart(); + } + + return prefix + metaBlock + suffix; +} diff --git a/src/utils/settings.ts b/src/utils/settings.ts new file mode 100644 index 0000000..96929f7 --- /dev/null +++ b/src/utils/settings.ts @@ -0,0 +1,60 @@ +import { reactive, watch } from 'vue'; + +export interface UserSettings { + convertUnderscoreToSpace: boolean; + escapeBrackets: boolean; + useAliasesInSearch: boolean; + useAliasesInAutocomplete: boolean; + smartBackspace: boolean; + language: 'en' | 'zh-CN'; + allowEditDefaultTags: boolean; + autoMetaEnabled: boolean; + tagMaxLength: number; +} + +const STORAGE_KEY = 'simplePrompt.settings'; + +const defaultSettings: UserSettings = { + convertUnderscoreToSpace: true, + escapeBrackets: false, + useAliasesInSearch: true, + useAliasesInAutocomplete: false, // Default false as per new spec + smartBackspace: false, // Default false - user must opt-in + language: 'en', // Default fallback + allowEditDefaultTags: false, + autoMetaEnabled: false, + tagMaxLength: 50, +}; + +export const settings = reactive({ ...defaultSettings }); + +// Load from localStorage +export function loadSettings() { + try { + const stored = localStorage.getItem(STORAGE_KEY); + if (stored) { + const parsed = JSON.parse(stored); + Object.assign(settings, { ...defaultSettings, ...parsed }); + } else { + // Detect language if not stored + const browserLang = navigator.language; + if (browserLang.startsWith('zh')) { + settings.language = 'zh-CN'; + } + } + } catch (e) { + console.error("Failed to load settings:", e); + } +} + +// Save to localStorage when changed +watch(settings, (newSettings) => { + try { + localStorage.setItem(STORAGE_KEY, JSON.stringify(newSettings)); + } catch (e) { + console.error("Failed to save settings:", e); + } +}, { deep: true }); + +// Initialize immediately +loadSettings(); diff --git a/src/utils/types.ts b/src/utils/types.ts new file mode 100644 index 0000000..a611f0a --- /dev/null +++ b/src/utils/types.ts @@ -0,0 +1,42 @@ +/** + * Tag Category Enum (Danbooru standard) + */ +export enum TagCategory { + GENERAL = 0, + ARTIST = 1, + COPYRIGHT = 3, + CHARACTER = 4, + META = 5 +} + +/** + * Category Interface + */ +export interface Category { + value: number; + label: string; + color: string; +} + +/** + * Category Colors (Danbooru standard) + * @deprecated Use CategoryService instead + */ +export const CATEGORY_COLORS: Record = { + [TagCategory.GENERAL]: '#0075db', // Blue + [TagCategory.ARTIST]: '#c00004', // Red + [TagCategory.COPYRIGHT]: '#c000c0', // Purple + [TagCategory.CHARACTER]: '#00aa00', // Green + [TagCategory.META]: '#fd9200' // Orange +}; + +/** + * Tag Item Data Structure + */ +export interface TagItem { + id: string; // Unique ID + text: string; // Tag text, e.g. "1girl" + weight: number; // Weight, default 1.0 + enabled: boolean; // Is enabled + category?: number; // Changed from TagCategory to facilitate dynamic IDs +} diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_search.py b/tests/test_search.py new file mode 100644 index 0000000..f8764b5 --- /dev/null +++ b/tests/test_search.py @@ -0,0 +1,174 @@ +import unittest +import duckdb +import tempfile +import os +import json +from simple_prompt.core.tags import search_tags, get_tags_details +from simple_prompt.core.utils import load_json, save_json + + +class TestSearchTags(unittest.TestCase): + def setUp(self): + self.conn = duckdb.connect(":memory:") + self.conn.execute(""" + CREATE TABLE tags ( + id INTEGER, + name VARCHAR, + post_count INTEGER, + category INTEGER, + alias VARCHAR[], + priority INTEGER + ) + """) + self.conn.execute(""" + INSERT INTO tags VALUES + (1, 'cat', 100, 0, ['kitty'], 4), + (2, 'dog', 80, 0, ['puppy'], 4), + (3, '1girl', 1000, 4, [], 4) + """) + + def test_search_basic(self): + results = search_tags(self.conn, "cat") + self.assertEqual(len(results), 1) + self.assertEqual(results[0]["name"], "cat") + + def test_search_alias(self): + results = search_tags(self.conn, "kitty", use_aliases=True) + self.assertEqual(len(results), 1) + self.assertEqual(results[0]["name"], "cat") + + def test_search_alias_disabled(self): + results = search_tags(self.conn, "kitty", use_aliases=False) + self.assertEqual(len(results), 0) + + def test_search_case_insensitive(self): + results = search_tags(self.conn, "CAT") + self.assertEqual(len(results), 1) + + def test_category_filter(self): + results = search_tags(self.conn, "cat", categories=[0]) + self.assertEqual(len(results), 1) + + results = search_tags(self.conn, "cat", categories=[1]) + self.assertEqual(len(results), 0) + + def test_search_injection_attempt(self): + # This string would cause issues if not parameterized: ' OR '1'='1 + results = search_tags(self.conn, "' OR '1'='1") + self.assertEqual(len(results), 0) + + +class TestGetTagsDetails(unittest.TestCase): + """Tests for get_tags_details function.""" + + def setUp(self): + self.conn = duckdb.connect(":memory:") + # Create tags table matching the expected schema + self.conn.execute(""" + CREATE TABLE tags ( + name VARCHAR, + category INTEGER, + post_count INTEGER, + alias VARCHAR[], + priority INTEGER + ) + """) + self.conn.execute(""" + INSERT INTO tags VALUES + ('blue_hair', 0, 500, [], 4), + ('red_eyes', 0, 300, [], 4), + ('hatsune_miku', 4, 10000, ['miku'], 4) + """) + # Create tags_fast view (same as tags for testing) + self.conn.execute("CREATE VIEW tags_fast AS SELECT * FROM tags") + + def test_basic_lookup(self): + """Test basic tag detail lookup.""" + result = get_tags_details(self.conn, ["blue_hair"]) + self.assertIn("blue_hair", result) + self.assertEqual(result["blue_hair"], 0) + + def test_space_to_underscore_conversion(self): + """Test that 'blue hair' matches 'blue_hair' in database.""" + result = get_tags_details(self.conn, ["blue hair"]) + # Should match and return category + self.assertIn("blue hair", result) + self.assertEqual(result["blue hair"], 0) + + def test_underscore_to_space_conversion(self): + """Test that both formats are returned in result.""" + result = get_tags_details(self.conn, ["blue_hair"]) + # Both variants should be in result + self.assertIn("blue_hair", result) + self.assertIn("blue hair", result) + + def test_empty_input(self): + """Test empty input returns empty dict.""" + result = get_tags_details(self.conn, []) + self.assertEqual(result, {}) + + def test_whitespace_only_input(self): + """Test whitespace-only names are filtered.""" + result = get_tags_details(self.conn, [" ", ""]) + self.assertEqual(result, {}) + + def test_multiple_tags(self): + """Test querying multiple tags at once.""" + result = get_tags_details(self.conn, ["blue_hair", "red_eyes", "hatsune_miku"]) + self.assertIn("blue_hair", result) + self.assertIn("red_eyes", result) + self.assertIn("hatsune_miku", result) + + def test_nonexistent_tag(self): + """Test that nonexistent tags are not in result.""" + result = get_tags_details(self.conn, ["nonexistent_tag_xyz"]) + self.assertNotIn("nonexistent_tag_xyz", result) + + def test_fast_mode(self): + """Test fast=True uses tags_fast view.""" + result = get_tags_details(self.conn, ["blue_hair"], fast=True) + self.assertIn("blue_hair", result) + + def test_none_connection(self): + """Test None connection returns empty dict.""" + result = get_tags_details(None, ["blue_hair"]) + self.assertEqual(result, {}) + + +class TestJsonHelpers(unittest.TestCase): + """Tests for load_json and save_json helper functions.""" + + def setUp(self): + self.temp_dir = tempfile.mkdtemp() + + def tearDown(self): + # Cleanup temp files + for f in os.listdir(self.temp_dir): + os.remove(os.path.join(self.temp_dir, f)) + os.rmdir(self.temp_dir) + + def test_save_and_load_json(self): + """Test saving and loading JSON data.""" + test_data = {"key": "value", "number": 42, "list": [1, 2, 3]} + path = os.path.join(self.temp_dir, "test.json") + + save_json(path, test_data) + loaded = load_json(path) + + self.assertEqual(loaded, test_data) + + def test_load_nonexistent_file(self): + """Test loading nonexistent file returns default.""" + path = os.path.join(self.temp_dir, "nonexistent.json") + result = load_json(path, default={"default": True}) + self.assertEqual(result, {"default": True}) + + def test_load_default_none(self): + """Test loading nonexistent file with None default.""" + path = os.path.join(self.temp_dir, "nonexistent.json") + result = load_json(path) + self.assertIsNone(result) + + +if __name__ == "__main__": + unittest.main() diff --git a/tsconfig.json b/tsconfig.json new file mode 100644 index 0000000..44e24ae --- /dev/null +++ b/tsconfig.json @@ -0,0 +1,22 @@ +{ + "compilerOptions": { + "target": "ES2022", + "useDefineForClassFields": true, + "module": "ESNext", + "resolveJsonModule": true, + "paths": { + "@/*": [ + "./src/*" + ] + }, + "typeRoots": [ + "src/types", + "node_modules/@types" + ] + }, + "exclude": [ + "node_modules", + "dist", + "js" + ] +} \ No newline at end of file diff --git a/vite.config.mts b/vite.config.mts new file mode 100644 index 0000000..12d03b8 --- /dev/null +++ b/vite.config.mts @@ -0,0 +1,58 @@ +import { fileURLToPath, URL } from 'node:url' +import { defineConfig } from 'vite' +import vue from '@vitejs/plugin-vue' +import vueDevTools from 'vite-plugin-vue-devtools' + +export default defineConfig({ + plugins: [ + vue(), + vueDevTools() + ], + define: { + 'process.env': { + NODE_ENV: JSON.stringify('production') + } + }, + resolve: { + alias: { + '@': fileURLToPath(new URL('./src', import.meta.url)) + }, + }, + build: { + lib: { + entry: './src/main.ts', + formats: ['es'], + fileName: 'main' + }, + rollupOptions: { + external: [ + '../../../scripts/app.js', + '../../../scripts/api.js', + '../../../scripts/domWidget.js', + '../../../scripts/utils.js', + /^primevue\/?.*/, + /^@primevue\/themes\/?.*/, + ], + output: { + dir: 'js', + assetFileNames: 'assets/[name].[ext]', + entryFileNames: 'main.js' + } + }, + outDir: 'js', + sourcemap: false, + assetsInlineLimit: 0, + cssCodeSplit: false + }, + server: { + fs: { + allow: ['..'] + }, + proxy: { + '/simple-prompt': { + target: 'http://127.0.0.1:8188', + changeOrigin: true, + } + } + } +}) \ No newline at end of file diff --git a/vite.debug.config.mts b/vite.debug.config.mts new file mode 100644 index 0000000..01b1f02 --- /dev/null +++ b/vite.debug.config.mts @@ -0,0 +1,22 @@ +import { defineConfig } from 'vite' + +export default defineConfig({ + build: { + lib: { + entry: './src/main.ts', + formats: ['es'], + fileName: 'main' + }, + rollupOptions: { + external: [ + '../../../scripts/app.js', + '../../../scripts/api.js', + ], + output: { + dir: 'js', + } + }, + outDir: 'js', + emptyOutDir: false + } +})