diff --git a/.dockerignore b/.dockerignore index 30bd2e24..965adc9e 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,9 +1,9 @@ # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and WebStorm # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 -# github acions +# github actions +.git .github/ .*ignore -.git/ # User-specific stuff .idea/ # Byte-compiled / optimized / DLL files @@ -15,10 +15,10 @@ env/ venv*/ ENV/ .conda/ -README*.md dashboard/ data/ changelogs/ tests/ .ruff_cache/ -.astrbot \ No newline at end of file +.astrbot +astrbot.lock \ No newline at end of file diff --git a/.github/ISSUE_TEMPLATE/PLUGIN_PUBLISH.yml b/.github/ISSUE_TEMPLATE/PLUGIN_PUBLISH.yml index 7957178c..c24bcf6d 100644 --- a/.github/ISSUE_TEMPLATE/PLUGIN_PUBLISH.yml +++ b/.github/ISSUE_TEMPLATE/PLUGIN_PUBLISH.yml @@ -16,7 +16,7 @@ body: 请将插件信息填写到下方的 JSON 代码块中。其中 `tags`（插件标签）和 `social_link`（社交链接）选填。 - 不熟悉 JSON ？现在可以从 [这里](https://plugins.astrbot.app/#/submit) 获取你的 JSON 啦！获取到了记得复制粘贴过来哦！ + 不熟悉 JSON ？可以从 [此站](https://plugins.astrbot.app) 右下角提交。 - type: textarea id: plugin-info @@ -26,12 +26,13 @@ body: value: | ```json { - "name": "插件名", - "desc": "插件介绍", + "name": "插件名，请以 astrbot_plugin_ 开头", + "display_name": "用于展示的插件名，方便人类阅读", + "desc": "插件的简短介绍", "author": "作者名", "repo": "插件仓库链接", "tags": [], - "social_link": "" + "social_link": "", } ``` validations: diff --git a/.github/ISSUE_TEMPLATE/bug-report.yml b/.github/ISSUE_TEMPLATE/bug-report.yml index 7506d0ec..77eeb3be 100644 --- a/.github/ISSUE_TEMPLATE/bug-report.yml +++ b/.github/ISSUE_TEMPLATE/bug-report.yml @@ -1,46 +1,44 @@ -name: '🐛 报告 Bug' +name: '🐛 Report Bug / 报告 Bug' title: '[Bug]' -description: 提交报告帮助我们改进。 +description: Submit bug report to help us improve. / 提交报告帮助我们改进。 labels: [ 'bug' ] body: - type: markdown attributes: value: | - 感谢您抽出时间报告问题！请准确解释您的问题。如果可能，请提供一个可复现的片段（这有助于更快地解决问题）。 + Thank you for taking the time to report this issue! Please describe your problem accurately. If possible, please provide a reproducible snippet (this will help resolve the issue more quickly). Please note that issues that are not detailed or have no logs will be closed immediately. Thank you for your understanding. / 感谢您抽出时间报告问题！请准确解释您的问题。如果可能，请提供一个可复现的片段（这有助于更快地解决问题）。请注意，不详细 / 没有日志的 issue 会被直接关闭，谢谢理解。 - type: textarea attributes: - label: 发生了什么 - description: 描述你遇到的异常 + label: What happened / 发生了什么 + description: Description placeholder: > - 一个清晰且具体的描述这个异常是什么。 + Please provide a clear and specific description of what this exception is. Please note that issues that are not detailed or have no logs will be closed immediately. Thank you for your understanding. / 一个清晰且具体的描述这个异常是什么。请注意，不详细 / 没有日志的 issue 会被直接关闭，谢谢理解。 validations: required: true - type: textarea attributes: - label: 如何复现？ + label: Reproduce / 如何复现？ description: > - 复现该问题的步骤 + The steps to reproduce the issue. / 复现该问题的步骤 placeholder: > - 如: 1. 打开 '...' + Example: 1. Open '...' validations: required: true - type: textarea attributes: - label: AstrBot 版本、部署方式（如 Windows Docker Desktop 部署）、使用的提供商、使用的消息平台适配器 - description: > - 请提供您的 AstrBot 版本和部署方式。 + label: AstrBot version, deployment method (e.g., Windows Docker Desktop deployment), provider used, and messaging platform used. / AstrBot 版本、部署方式（如 Windows Docker Desktop 部署）、使用的提供商、使用的消息平台适配器 placeholder: > - 如: 3.1.8 Docker, 3.1.7 Windows启动器 + Example: 4.5.7 Docker, 3.1.7 Windows Launcher validations: required: true - type: dropdown attributes: - label: 操作系统 + label: OS description: | - 你在哪个操作系统上遇到了这个问题？ + On which operating system did you encounter this problem? / 你在哪个操作系统上遇到了这个问题？ multiple: false options: - 'Windows' @@ -53,30 +51,30 @@ body: - type: textarea attributes: - label: 报错日志 + label: Logs / 报错日志 description: > - 如报错日志、截图等。请提供完整的 Debug 级别的日志，不要介意它很长！ + Please provide complete Debug-level logs, such as error logs and screenshots. Don't worry if they're long! Please note that issues with insufficient details or no logs will be closed immediately. Thank you for your understanding. / 如报错日志、截图等。请提供完整的 Debug 级别的日志，不要介意它很长！请注意，不详细 / 没有日志的 issue 会被直接关闭，谢谢理解。 placeholder: > - 请提供完整的报错日志或截图。 + Please provide a complete error log or screenshot. / 请提供完整的报错日志或截图。 validations: required: true - type: checkboxes attributes: - label: 你愿意提交 PR 吗？ + label: Are you willing to submit a PR? / 你愿意提交 PR 吗？ description: > - 这不是必需的，但我们很乐意在贡献过程中为您提供指导特别是如果你已经很好地理解了如何实现修复。 + This is not required, but we would be happy to provide guidance during the contribution process, especially if you already have a good understanding of how to implement the fix. / 这不是必需的，但我们很乐意在贡献过程中为您提供指导特别是如果你已经很好地理解了如何实现修复。 options: - - label: 是的，我愿意提交 PR! + - label: Yes! - type: checkboxes attributes: label: Code of Conduct options: - label: > - 我已阅读并同意遵守该项目的 [行为准则](https://docs.github.com/zh/site-policy/github-terms/github-community-code-of-conduct)。 + I have read and agree to abide by the project's [Code of Conduct](https://docs.github.com/zh/site-policy/github-terms/github-community-code-of-conduct)。 required: true - type: markdown attributes: - value: "感谢您填写我们的表单！" + value: "Thank you for filling out our form! / 感谢您填写我们的表单！" diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index e58a301e..70bb8f30 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -1,44 +1,25 @@ - - - -fixes #XYZ - ---- - -### Motivation / 动机 - - - + + ### Modifications / 改动点 -### Verification Steps / 验证步骤 - - - +- [x] This is NOT a breaking change. / 这不是一个破坏性变更。 + ### Screenshots or Test Results / 运行截图或测试结果 - - -### Compatibility & Breaking Changes / 兼容性与破坏性变更 - - - - -- [ ] 这是一个破坏性变更 (Breaking Change)。/ This is a breaking change. -- [ ] 这不是一个破坏性变更。/ This is NOT a breaking change. + --- ### Checklist / 检查清单 - + - [ ] 😊 如果 PR 中有新加入的功能，已经通过 Issue / 邮件等方式和作者讨论过。/ If there are new features added in the PR, I have discussed it with the authors through issues/emails, etc. - [ ] 👀 我的更改经过了良好的测试，**并已在上方提供了“验证步骤”和“运行截图”**。/ My changes have been well-tested, **and "Verification Steps" and "Screenshots" have been provided above**. diff --git a/.github/workflows/dashboard_ci.yml b/.github/workflows/dashboard_ci.yml index 20562784..f02085f8 100644 --- a/.github/workflows/dashboard_ci.yml +++ b/.github/workflows/dashboard_ci.yml @@ -36,7 +36,7 @@ jobs: zip -r dist.zip dist - name: Archive production artifacts - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v5 with: name: dist-without-markdown path: | diff --git a/.github/workflows/docker-image.yml b/.github/workflows/docker-image.yml index ecc098c3..e7174958 100644 --- a/.github/workflows/docker-image.yml +++ b/.github/workflows/docker-image.yml @@ -3,18 +3,125 @@ name: Docker Image CI/CD on: push: tags: - - 'v*' + - "v*" + schedule: + # Run at 00:00 UTC every day + - cron: "0 0 * * *" workflow_dispatch: jobs: - publish-docker: + build-nightly-image: + if: github.event_name == 'schedule' runs-on: ubuntu-latest + env: + DOCKER_HUB_USERNAME: ${{ secrets.DOCKER_HUB_USERNAME }} + GHCR_OWNER: soulter + HAS_GHCR_TOKEN: ${{ secrets.GHCR_GITHUB_TOKEN != '' }} steps: - - name: Pull The Codes + - name: Checkout uses: actions/checkout@v5 with: - fetch-depth: 0 # Must be 0 so we can fetch tags + fetch-depth: 1 + fetch-tag: true + + - name: Check for new commits today + if: github.event_name == 'schedule' + id: check-commits + run: | + # Get commits from the last 24 hours + commits=$(git log --since="24 hours ago" --oneline) + if [ -z "$commits" ]; then + echo "No commits in the last 24 hours, skipping build" + echo "has_commits=false" >> $GITHUB_OUTPUT + else + echo "Found commits in the last 24 hours:" + echo "$commits" + echo "has_commits=true" >> $GITHUB_OUTPUT + fi + + - name: Exit if no commits + if: github.event_name == 'schedule' && steps.check-commits.outputs.has_commits == 'false' + run: exit 0 + + - name: Build Dashboard + run: | + cd dashboard + npm install + npm run build + mkdir -p dist/assets + echo $(git rev-parse HEAD) > dist/assets/version + cd .. + mkdir -p data + cp -r dashboard/dist data/ + + - name: Determine test image tags + id: test-meta + run: | + short_sha=$(echo "${GITHUB_SHA}" | cut -c1-12) + build_date=$(date +%Y%m%d) + echo "short_sha=$short_sha" >> $GITHUB_OUTPUT + echo "build_date=$build_date" >> $GITHUB_OUTPUT + + - name: Set QEMU + uses: docker/setup-qemu-action@v3 + + - name: Set Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Log in to DockerHub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKER_HUB_USERNAME }} + password: ${{ secrets.DOCKER_HUB_PASSWORD }} + + - name: Login to GitHub Container Registry + if: env.HAS_GHCR_TOKEN == 'true' + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ env.GHCR_OWNER }} + password: ${{ secrets.GHCR_GITHUB_TOKEN }} + + - name: Build nightly image tags list + id: test-tags + run: | + TAGS="${{ env.DOCKER_HUB_USERNAME }}/astrbot:nightly-latest + ${{ env.DOCKER_HUB_USERNAME }}/astrbot:nightly-${{ steps.test-meta.outputs.build_date }}-${{ steps.test-meta.outputs.short_sha }}" + if [ "${{ env.HAS_GHCR_TOKEN }}" = "true" ]; then + TAGS="$TAGS + ghcr.io/${{ env.GHCR_OWNER }}/astrbot:nightly-latest + ghcr.io/${{ env.GHCR_OWNER }}/astrbot:nightly-${{ steps.test-meta.outputs.build_date }}-${{ steps.test-meta.outputs.short_sha }}" + fi + echo "tags<> $GITHUB_OUTPUT + echo "$TAGS" >> $GITHUB_OUTPUT + echo "EOF" >> $GITHUB_OUTPUT + + - name: Build and Push Nightly Image + uses: docker/build-push-action@v6 + with: + context: . + platforms: linux/amd64,linux/arm64 + push: true + tags: ${{ steps.test-tags.outputs.tags }} + + - name: Post build notifications + run: echo "Test Docker image has been built and pushed successfully" + + build-release-image: + if: github.event_name == 'workflow_dispatch' || (github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')) + runs-on: ubuntu-latest + env: + DOCKER_HUB_USERNAME: ${{ secrets.DOCKER_HUB_USERNAME }} + GHCR_OWNER: soulter + HAS_GHCR_TOKEN: ${{ secrets.GHCR_GITHUB_TOKEN != '' }} + + steps: + - name: Checkout + uses: actions/checkout@v5 + with: + fetch-depth: 1 + fetch-tag: true - name: Get latest tag (only on manual trigger) id: get-latest-tag @@ -27,21 +134,22 @@ jobs: if: github.event_name == 'workflow_dispatch' run: git checkout ${{ steps.get-latest-tag.outputs.latest_tag }} - - name: Check if version is pre-release - id: check-prerelease + - name: Compute release metadata + id: release-meta run: | - if [ "${{ github.event_name }}" == "workflow_dispatch" ]; then + if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then version="${{ steps.get-latest-tag.outputs.latest_tag }}" else - version="${{ github.ref_name }}" + version="${GITHUB_REF#refs/tags/}" fi if [[ "$version" == *"beta"* ]] || [[ "$version" == *"alpha"* ]]; then echo "is_prerelease=true" >> $GITHUB_OUTPUT - echo "Version $version is a pre-release, will not push latest tag" + echo "Version $version marked as pre-release" else echo "is_prerelease=false" >> $GITHUB_OUTPUT - echo "Version $version is a stable release, will push latest tag" + echo "Version $version marked as stable" fi + echo "version=$version" >> $GITHUB_OUTPUT - name: Build Dashboard run: | @@ -67,23 +175,24 @@ jobs: password: ${{ secrets.DOCKER_HUB_PASSWORD }} - name: Login to GitHub Container Registry + if: env.HAS_GHCR_TOKEN == 'true' uses: docker/login-action@v3 with: registry: ghcr.io - username: Soulter + username: ${{ env.GHCR_OWNER }} password: ${{ secrets.GHCR_GITHUB_TOKEN }} - - name: Build and Push Docker to DockerHub and Github GHCR + - name: Build and Push Release Image uses: docker/build-push-action@v6 with: context: . platforms: linux/amd64,linux/arm64 push: true tags: | - ${{ steps.check-prerelease.outputs.is_prerelease == 'false' && format('{0}/astrbot:latest', secrets.DOCKER_HUB_USERNAME) || '' }} - ${{ secrets.DOCKER_HUB_USERNAME }}/astrbot:${{ github.event_name == 'workflow_dispatch' && steps.get-latest-tag.outputs.latest_tag || github.ref_name }} - ${{ steps.check-prerelease.outputs.is_prerelease == 'false' && 'ghcr.io/soulter/astrbot:latest' || '' }} - ghcr.io/soulter/astrbot:${{ github.event_name == 'workflow_dispatch' && steps.get-latest-tag.outputs.latest_tag || github.ref_name }} + ${{ steps.release-meta.outputs.is_prerelease == 'false' && format('{0}/astrbot:latest', env.DOCKER_HUB_USERNAME) || '' }} + ${{ steps.release-meta.outputs.is_prerelease == 'false' && env.HAS_GHCR_TOKEN == 'true' && format('ghcr.io/{0}/astrbot:latest', env.GHCR_OWNER) || '' }} + ${{ format('{0}/astrbot:{1}', env.DOCKER_HUB_USERNAME, steps.release-meta.outputs.version) }} + ${{ env.HAS_GHCR_TOKEN == 'true' && format('ghcr.io/{0}/astrbot:{1}', env.GHCR_OWNER, steps.release-meta.outputs.version) || '' }} - name: Post build notifications - run: echo "Docker image has been built and pushed successfully" + run: echo "Release Docker image has been built and pushed successfully" diff --git a/.gitignore b/.gitignore index 8006fef8..366a1030 100644 --- a/.gitignore +++ b/.gitignore @@ -1,35 +1,49 @@ +# Python related __pycache__ -botpy.log -.vscode +.mypy_cache .venv* -.idea -data_v2.db -data_v3.db -configs/session -configs/config.yaml -**/.DS_Store -temp -cmd_config.json -data -cookies.json -logs/ -addons/plugins +.conda/ +uv.lock .coverage +# IDE and editors +.vscode +.idea +# Logs and temporary files +botpy.log +logs/ +temp +cookies.json + +# Data files +data_v2.db +data_v3.db +data +configs/session +configs/config.yaml +cmd_config.json + +# Plugins and packages +addons/plugins +packages/python_interpreter/workplace tests/astrbot_plugin_openai -chroma + +# Dashboard dashboard/node_modules/ dashboard/dist/ -.DS_Store package-lock.json package.json -venv/* -packages/python_interpreter/workplace -.venv/* -.conda/ -.idea -pytest.ini -.astrbot -uv.lock \ No newline at end of file +# Operating System +**/.DS_Store +.DS_Store + +# AstrBot specific +.astrbot +astrbot.lock + +# Other +chroma +venv/* +pytest.ini diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index e0a0e314..8611e269 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -6,8 +6,20 @@ ci: autoupdate_schedule: weekly autoupdate_commit_msg: ":balloon: pre-commit autoupdate" repos: - - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.11.2 - hooks: - - id: ruff - - id: ruff-format +- repo: https://github.com/astral-sh/ruff-pre-commit + # Ruff version. + rev: v0.14.1 + hooks: + # Run the linter. + - id: ruff-check + types_or: [ python, pyi ] + args: [ --fix ] + # Run the formatter. + - id: ruff-format + types_or: [ python, pyi ] + +- repo: https://github.com/asottile/pyupgrade + rev: v3.21.0 + hooks: + - id: pyupgrade + args: [--py310-plus] diff --git a/Dockerfile b/Dockerfile index df48b2be..f143cdd6 100644 --- a/Dockerfile +++ b/Dockerfile @@ -12,19 +12,21 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ ca-certificates \ bash \ ffmpeg \ + curl \ + gnupg \ + git \ && apt-get clean \ - && rm -rf /var/lib/apt/lists/* + && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* -RUN apt-get update && apt-get install -y curl gnupg && \ - curl -fsSL https://deb.nodesource.com/setup_lts.x | bash - && \ - apt-get install -y nodejs && \ - rm -rf /var/lib/apt/lists/* +RUN apt-get update && apt-get install -y curl gnupg \ + && curl -fsSL https://deb.nodesource.com/setup_lts.x | bash - \ + && apt-get install -y nodejs -RUN python -m pip install uv +RUN python -m pip install uv \ + && echo "3.11" > .python-version RUN uv pip install -r requirements.txt --no-cache-dir --system RUN uv pip install socksio uv pilk --no-cache-dir --system EXPOSE 6185 -EXPOSE 6186 -CMD [ "python", "main.py" ] +CMD ["python", "main.py"] diff --git a/Dockerfile_with_node b/Dockerfile_with_node deleted file mode 100644 index 3bd37468..00000000 --- a/Dockerfile_with_node +++ /dev/null @@ -1,35 +0,0 @@ -FROM python:3.10-slim - -WORKDIR /AstrBot - -COPY . /AstrBot/ - -RUN apt-get update && apt-get install -y --no-install-recommends \ - gcc \ - build-essential \ - python3-dev \ - libffi-dev \ - libssl-dev \ - curl \ - unzip \ - ca-certificates \ - bash \ - && apt-get clean \ - && rm -rf /var/lib/apt/lists/* - -# Installation of Node.js -ENV NVM_DIR="/root/.nvm" -RUN curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.2/install.sh | bash && \ - . "$NVM_DIR/nvm.sh" && \ - nvm install 22 && \ - nvm use 22 -RUN /bin/bash -c ". \"$NVM_DIR/nvm.sh\" && node -v && npm -v" - -RUN python -m pip install uv -RUN uv pip install -r requirements.txt --no-cache-dir --system -RUN uv pip install socksio uv pyffmpeg --no-cache-dir --system - -EXPOSE 6185 -EXPOSE 6186 - -CMD ["python", "main.py"] diff --git a/README.md b/README.md index 762ae9dd..6d60f80a 100644 --- a/README.md +++ b/README.md @@ -4,24 +4,35 @@
-Soulter%2FAstrBot | Trendshift +
-[![GitHub release (latest by date)](https://img.shields.io/github/v/release/Soulter/AstrBot?style=for-the-badge&color=76bad9)](https://github.com/Soulter/AstrBot/releases/latest) +
+Soulter%2FAstrBot | Trendshift +Featured｜HelloGitHub +
+ +
+ +
+ python Docker pull -QQ_community -Telegram_community -![Dynamic JSON Badge](https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fapi.soulter.top%2Fastrbot%2Fplugin-num&query=%24.result&suffix=%E4%B8%AA&style=for-the-badge&label=%E6%8F%92%E4%BB%B6%E5%B8%82%E5%9C%BA&cacheSeconds=3600) +QQ_community +Telegram_community + +
-English ｜ -日本語 ｜ +
+ +English ｜ +日本語 ｜ 文档 ｜ Blog ｜ 路线图 ｜ -问题提交 +问题提交
-AstrBot 是一个开源的一站式 Agent 聊天机器人平台及开发框架。 +AstrBot 是一个开源的一站式 Agent 聊天机器人平台，可无缝接入主流即时通讯软件，为个人、开发者和团队打造可靠、可扩展的对话式智能基础设施。无论是个人 AI 伙伴、智能客服、自动化助手，还是企业知识库，AstrBot 都能在你的即时通讯软件平台的工作流中快速构建生产可用的 AI 应用。 ## 主要功能 @@ -31,7 +42,7 @@ AstrBot 是一个开源的一站式 Agent 聊天机器人平台及开发框架 4. **插件扩展**。深度优化的插件机制，支持[开发插件](https://astrbot.app/dev/plugin.html)扩展功能，社区插件生态丰富。 5. **WebUI**。可视化配置和管理机器人，功能齐全。 -## 部署方式 +## 部署方式 #### Docker 部署(推荐 🥳) @@ -61,7 +72,7 @@ AstrBot 已由雨云官方上架至云应用平台，可一键部署。 社区贡献的部署方式。 -[![Run on Repl.it](https://repl.it/badge/github/Soulter/AstrBot)](https://repl.it/github/Soulter/AstrBot) +[![Run on Repl.it](https://repl.it/badge/github/AstrBotDevs/AstrBot)](https://repl.it/github/AstrBotDevs/AstrBot) #### Windows 一键安装器部署 @@ -108,83 +119,73 @@ uv run main.py Discord_community -## ⚡ 消息平台支持情况 +## 支持的消息平台 **官方维护** -| 平台 | 支持性 | -| -------- | ------- | -| QQ(官方平台) | ✔ | -| QQ(OneBot) | ✔ | -| Telegram | ✔ | -| 企微应用 | ✔ | -| 企微智能机器人 | ✔ | -| 微信客服 | ✔ | -| 微信公众号 | ✔ | -| 飞书 | ✔ | -| 钉钉 | ✔ | -| Slack | ✔ | -| Discord | ✔ | -| Satori | ✔ | -| Misskey | ✔ | -| Whatsapp | 将支持 | -| LINE | 将支持 | +- QQ (官方平台 & OneBot) +- Telegram +- 企微应用 & 企微智能机器人 +- 微信客服 & 微信公众号 +- 飞书 +- 钉钉 +- Slack +- Discord +- Satori +- Misskey +- Whatsapp (将支持) +- LINE (将支持) **社区维护** -| 平台 | 支持性 | -| -------- | ------- | -| [KOOK](https://github.com/wuyan1003/astrbot_plugin_kook_adapter) | ✔ | -| [VoceChat](https://github.com/HikariFroya/astrbot_plugin_vocechat) | ✔ | -| [Bilibili 私信](https://github.com/Hina-Chat/astrbot_plugin_bilibili_adapter) | ✔ | -| [wxauto](https://github.com/luosheng520qaq/wxauto-repost-onebotv11) | ✔ | +- [KOOK](https://github.com/wuyan1003/astrbot_plugin_kook_adapter) +- [VoceChat](https://github.com/HikariFroya/astrbot_plugin_vocechat) +- [Bilibili 私信](https://github.com/Hina-Chat/astrbot_plugin_bilibili_adapter) +- [wxauto](https://github.com/luosheng520qaq/wxauto-repost-onebotv11) -## ⚡ 提供商支持情况 +## 支持的模型服务 **大模型服务** -| 名称 | 支持性 | 备注 | -| -------- | ------- | ------- | -| OpenAI | ✔ | 支持任何兼容 OpenAI API 的服务 | -| Anthropic | ✔ | | -| Google Gemini | ✔ | | -| Moonshot AI | ✔ | | -| 智谱 AI | ✔ | | -| DeepSeek | ✔ | | -| Ollama | ✔ | 本地部署 DeepSeek 等开源语言模型 | -| LM Studio | ✔ | 本地部署 DeepSeek 等开源语言模型 | -| [优云智算](https://www.compshare.cn/?ytag=GPU_YY-gh_astrbot&referral_code=FV7DcGowN4hB5UuXKgpE74) | ✔ | | -| [302.AI](https://share.302.ai/rr1M3l) | ✔ | | -| [小马算力](https://www.tokenpony.cn/3YPyf) | ✔ | | -| 硅基流动 | ✔ | | -| PPIO 派欧云 | ✔ | | -| ModelScope | ✔ | | -| OneAPI | ✔ | | -| Dify | ✔ | | -| 阿里云百炼应用 | ✔ | | -| Coze | ✔ | | +- OpenAI 及兼容服务 +- Anthropic +- Google Gemini +- Moonshot AI +- 智谱 AI +- DeepSeek +- Ollama (本地部署) +- LM Studio (本地部署) +- [优云智算](https://www.compshare.cn/?ytag=GPU_YY-gh_astrbot&referral_code=FV7DcGowN4hB5UuXKgpE74) +- [302.AI](https://share.302.ai/rr1M3l) +- [小马算力](https://www.tokenpony.cn/3YPyf) +- [硅基流动](https://docs.siliconflow.cn/cn/usercases/use-siliconcloud-in-astrbot) +- [PPIO 派欧云](https://ppio.com/user/register?invited_by=AIOONE) +- ModelScope +- OneAPI + +**LLMOps 平台** + +- Dify +- 阿里云百炼应用 +- Coze **语音转文本服务** -| 名称 | 支持性 | 备注 | -| -------- | ------- | ------- | -| Whisper | ✔ | 支持 API、本地部署 | -| SenseVoice | ✔ | 本地部署 | +- OpenAI Whisper +- SenseVoice **文本转语音服务** -| 名称 | 支持性 | 备注 | -| -------- | ------- | ------- | -| OpenAI TTS | ✔ | | -| Gemini TTS | ✔ | | -| GSVI | ✔ | GPT-Sovits-Inference | -| GPT-SoVITs | ✔ | GPT-Sovits | -| FishAudio | ✔ | | -| Edge TTS | ✔ | Edge 浏览器的免费 TTS | -| 阿里云百炼 TTS | ✔ | | -| Azure TTS | ✔ | | -| Minimax TTS | ✔ | | -| 火山引擎 TTS | ✔ | | +- OpenAI TTS +- Gemini TTS +- GPT-Sovits-Inference +- GPT-Sovits +- FishAudio +- Edge TTS +- 阿里云百炼 TTS +- Azure TTS +- Minimax TTS +- 火山引擎 TTS ## ❤️ 贡献 @@ -199,7 +200,7 @@ uv run main.py AstrBot 使用 `ruff` 进行代码格式化和检查。 ```bash -git clone https://github.com/Soulter/AstrBot +git clone https://github.com/AstrBotDevs/AstrBot pip install pre-commit pre-commit install ``` @@ -218,12 +219,12 @@ pre-commit install ## ⭐ Star History -> [!TIP] +> [!TIP] > 如果本项目对您的生活 / 工作产生了帮助，或者您关注本项目的未来发展，请给项目 Star，这是我们维护这个开源项目的动力 <3
-[![Star History Chart](https://api.star-history.com/svg?repos=soulter/astrbot&type=Date)](https://star-history.com/#soulter/astrbot&Date) +[![Star History Chart](https://api.star-history.com/svg?repos=astrbotdevs/astrbot&type=Date)](https://star-history.com/#astrbotdevs/astrbot&Date)
diff --git a/README_en.md b/README_en.md index 249dd5e7..520cfebe 100644 --- a/README_en.md +++ b/README_en.md @@ -1,182 +1,233 @@ -

- -![6e1279651f16d7fdf4727558b72bbaf1](https://github.com/user-attachments/assets/ead4c551-fc3c-48f7-a6f7-afbfdb820512) +![AstrBot-Logo-Simplified](https://github.com/user-attachments/assets/ffd99b6b-3272-4682-beaa-6fe74250f7d9)

-_✨ Easy-to-use Multi-platform LLM Chatbot & Development Framework ✨_ +
+
Soulter%2FAstrBot | Trendshift - -[![GitHub release (latest by date)](https://img.shields.io/github/v/release/Soulter/AstrBot)](https://github.com/Soulter/AstrBot/releases/latest) -python -Docker pull -Static Badge -[![wakatime](https://wakatime.com/badge/user/915e5316-99c6-4563-a483-ef186cf000c9/project/018e705a-a1a7-409a-a849-3013485e6c8e.svg)](https://wakatime.com/badge/user/915e5316-99c6-4563-a483-ef186cf000c9/project/018e705a-a1a7-409a-a849-3013485e6c8e) -![Dynamic JSON Badge](https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fapi.soulter.top%2Fastrbot%2Fstats&query=v&label=7%E6%97%A5%E6%B6%88%E6%81%AF%E4%B8%8A%E8%A1%8C%E9%87%8F&cacheSeconds=3600) -[![codecov](https://codecov.io/gh/Soulter/AstrBot/graph/badge.svg?token=FF3P5967B8)](https://codecov.io/gh/Soulter/AstrBot) - -Documentation ｜ -Issue Tracking +Featured｜HelloGitHub
-AstrBot is a loosely coupled, asynchronous chatbot and development framework that supports multi-platform deployment, featuring an easy-to-use plugin system and comprehensive Large Language Model (LLM) integration capabilities. +
-## ✨ Key Features +
+ +python +Docker pull +QQ_community +Telegram_community + +
-1. **LLM Conversations** - Supports various LLMs including OpenAI API, Google Gemini, Llama, Deepseek, ChatGLM, etc. Enables local model deployment via Ollama/LLMTuner. Features multi-turn dialogues, personality contexts, multimodal capabilities (image understanding), and speech-to-text (Whisper). -2. **Multi-platform Integration** - Supports QQ (OneBot), QQ Channels, WeChat (Gewechat), Feishu, and Telegram. Planned support for DingTalk, Discord, WhatsApp, and Xiaomi Smart Speakers. Includes rate limiting, whitelisting, keyword filtering, and Baidu content moderation. -3. **Agent Capabilities** - Native support for code execution, natural language TODO lists, web search. Integrates with [Dify Platform](https://dify.ai/) for easy access to Dify assistants/knowledge bases/workflows. -4. **Plugin System** - Optimized plugin mechanism with minimal development effort. Supports multiple installed plugins. -5. **Web Dashboard** - Visual configuration management, plugin controls, logging, and WebChat interface for direct LLM interaction. -6. **High Stability & Modularity** - Event bus and pipeline architecture ensures high modularization and loose coupling. +
-> [!TIP] -> Dashboard Demo: [https://demo.astrbot.app/](https://demo.astrbot.app/) -> Username: `astrbot`, Password: `astrbot` (LLM not configured for chat page) +中文 ｜ +日本語 ｜ +Documentation ｜ +Blog ｜ +Roadmap ｜ +Issue Tracker +
-## ✨ Deployment +AstrBot is an open-source all-in-one Agent chatbot platform and development framework. -#### Docker Deployment +## Key Features -See docs: [Deploy with Docker](https://astrbot.app/deploy/astrbot/docker.html#docker-deployment) +1. **LLM Conversations**. Supports integration with various large language model services. Features include multimodal capabilities, tool calling, MCP, native knowledge base, character personas, and more. +2. **Multi-Platform Support**. Integrates with QQ, WeChat Work, WeChat Official Accounts, Feishu, Telegram, DingTalk, Discord, KOOK, and other platforms. Supports rate limiting, whitelisting, and Baidu content moderation. +3. **Agent Capabilities**. Fully optimized agentic features including multi-turn tool calling, built-in sandboxed code executor, web search, and more. +4. **Plugin Extensions**. Deeply optimized plugin mechanism supporting [plugin development](https://astrbot.app/dev/plugin.html) to extend functionality, with a rich community plugin ecosystem. +5. **Web UI**. Visual configuration and management of your bot with comprehensive features. -#### Windows Installer +## Deployment Methods -Requires Python (>3.10). See docs: [Windows Installer Guide](https://astrbot.app/deploy/astrbot/windows.html) +#### Docker Deployment (Recommended 🥳) -#### Replit Deployment +We recommend deploying AstrBot using Docker or Docker Compose. -[![Run on Repl.it](https://repl.it/badge/github/Soulter/AstrBot)](https://repl.it/github/Soulter/AstrBot) +Please refer to the official documentation: [Deploy AstrBot with Docker](https://astrbot.app/deploy/astrbot/docker.html#%E4%BD%BF%E7%94%A8-docker-%E9%83%A8%E7%BD%B2-astrbot). + +#### BT-Panel Deployment + +AstrBot has partnered with BT-Panel and is now available in their marketplace. + +Please refer to the official documentation: [BT-Panel Deployment](https://astrbot.app/deploy/astrbot/btpanel.html). + +#### 1Panel Deployment + +AstrBot has been officially listed on the 1Panel marketplace. + +Please refer to the official documentation: [1Panel Deployment](https://astrbot.app/deploy/astrbot/1panel.html). + +#### Deploy on RainYun + +AstrBot has been officially listed on RainYun's cloud application platform with one-click deployment. + +[![Deploy on RainYun](https://rainyun-apps.cn-nb1.rains3.com/materials/deploy-on-rainyun-en.svg)](https://app.rainyun.com/apps/rca/store/5994?ref=NjU1ODg0) + +#### Deploy on Replit + +Community-contributed deployment method. + +[![Run on Repl.it](https://repl.it/badge/github/AstrBotDevs/AstrBot)](https://repl.it/github/AstrBotDevs/AstrBot) + +#### Windows One-Click Installer + +Please refer to the official documentation: [Deploy AstrBot with Windows One-Click Installer](https://astrbot.app/deploy/astrbot/windows.html). #### CasaOS Deployment -Community-contributed method. -See docs: [CasaOS Deployment](https://astrbot.app/deploy/astrbot/casaos.html) +Community-contributed deployment method. + +Please refer to the official documentation: [CasaOS Deployment](https://astrbot.app/deploy/astrbot/casaos.html). #### Manual Deployment -See docs: [Source Code Deployment](https://astrbot.app/deploy/astrbot/cli.html) +First, install uv: -## ⚡ Platform Support +```bash +pip install uv +``` -| Platform | Status | Details | Message Types | -| -------------------------------------------------------------- | ------ | ------------------- | ------------------- | -| QQ (Official Bot) | ✔ | Private/Group chats | Text, Images | -| QQ (OneBot) | ✔ | Private/Group chats | Text, Images, Voice | -| WeChat (Personal) | ✔ | Private/Group chats | Text, Images, Voice | -| [Telegram](https://github.com/Soulter/astrbot_plugin_telegram) | ✔ | Private/Group chats | Text, Images | -| [WeChat Work](https://github.com/Soulter/astrbot_plugin_wecom) | ✔ | Private chats | Text, Images, Voice | -| Feishu | ✔ | Group chats | Text, Images | -| WeChat Open Platform | 🚧 | Planned | - | -| Discord | 🚧 | Planned | - | -| WhatsApp | 🚧 | Planned | - | -| Xiaomi Speakers | 🚧 | Planned | - | +Install AstrBot via Git Clone: -## Provider Support Status +```bash +git clone https://github.com/AstrBotDevs/AstrBot && cd AstrBot +uv run main.py +``` -| Name | Support | Type | Notes | -|---------------------------|---------|------------------------|-----------------------------------------------------------------------| -| OpenAI API | ✔ | Text Generation | Supports all OpenAI API-compatible services including DeepSeek, Google Gemini, GLM, Moonshot, Alibaba Cloud Bailian, Silicon Flow, xAI, etc. | -| Claude API | ✔ | Text Generation | | -| Google Gemini API | ✔ | Text Generation | | -| Dify | ✔ | LLMOps | | -| DashScope (Alibaba Cloud) | ✔ | LLMOps | | -| Ollama | ✔ | Model Loader | Local deployment for open-source LLMs (DeepSeek, Llama, etc.) | -| LM Studio | ✔ | Model Loader | Local deployment for open-source LLMs (DeepSeek, Llama, etc.) | -| LLMTuner | ✔ | Model Loader | Local loading of fine-tuned models (e.g. LoRA) | -| OneAPI | ✔ | LLM Distribution | | -| Whisper | ✔ | Speech-to-Text | Supports API and local deployment | -| SenseVoice | ✔ | Speech-to-Text | Local deployment | -| OpenAI TTS API | ✔ | Text-to-Speech | | -| Fishaudio | ✔ | Text-to-Speech | Project involving GPT-Sovits author | +Or refer to the official documentation: [Deploy AstrBot from Source](https://astrbot.app/deploy/astrbot/cli.html). -# 🦌 Roadmap +## 🌍 Community -> [!TIP] -> Suggestions welcome via Issues <3 +### QQ Groups -- [ ] Ensure feature parity across all platform adapters -- [ ] Optimize plugin APIs -- [ ] Add default TTS services (e.g., GPT-Sovits) -- [ ] Enhance chat features with persistent memory -- [ ] i18n Planning +- Group 1: 322154837 +- Group 3: 630166526 +- Group 5: 822130018 +- Group 6: 753075035 +- Developer Group: 975206796 -## ❤️ Contributions +### Telegram Group -All Issues/PRs welcome! Simply submit your changes to this project :) +Telegram_community -For major features, please discuss via Issues first. +### Discord Server -## 🌟 Support +Discord_community -- Star this project! -- Support via [Afdian](https://afdian.com/a/soulter) -- WeChat support: [QR Code](https://drive.soulter.top/f/pYfA/d903f4fa49a496fda3f16d2be9e023b5.png) +## Supported Messaging Platforms -## ✨ Demos +**Officially Maintained** -> [!NOTE] -> Code executor file I/O currently tested with Napcat(QQ)/Lagrange(QQ) +- QQ (Official Platform & OneBot) +- Telegram +- WeChat Work Application & WeChat Work Intelligent Bot +- WeChat Customer Service & WeChat Official Accounts +- Feishu (Lark) +- DingTalk +- Slack +- Discord +- Satori +- Misskey +- WhatsApp (Coming Soon) +- LINE (Coming Soon) -
+**Community Maintained** - +- [KOOK](https://github.com/wuyan1003/astrbot_plugin_kook_adapter) +- [VoceChat](https://github.com/HikariFroya/astrbot_plugin_vocechat) +- [Bilibili Direct Messages](https://github.com/Hina-Chat/astrbot_plugin_bilibili_adapter) +- [wxauto](https://github.com/luosheng520qaq/wxauto-repost-onebotv11) -_✨ Docker-based Sandboxed Code Executor (Beta) ✨_ +## Supported Model Services - +**LLM Services** -_✨ Multimodal Input, Web Search, Text-to-Image ✨_ +- OpenAI and Compatible Services +- Anthropic +- Google Gemini +- Moonshot AI +- Zhipu AI +- DeepSeek +- Ollama (Self-hosted) +- LM Studio (Self-hosted) +- [CompShare](https://www.compshare.cn/?ytag=GPU_YY-gh_astrbot&referral_code=FV7DcGowN4hB5UuXKgpE74) +- [302.AI](https://share.302.ai/rr1M3l) +- [TokenPony](https://www.tokenpony.cn/3YPyf) +- [SiliconFlow](https://docs.siliconflow.cn/cn/usecases/use-siliconcloud-in-astrbot) +- [PPIO Cloud](https://ppio.com/user/register?invited_by=AIOONE) +- ModelScope +- OneAPI - +**LLMOps Platforms** -_✨ Natural Language TODO Lists ✨_ +- Dify +- Alibaba Cloud Bailian Applications +- Coze - - +**Speech-to-Text Services** -_✨ Plugin System Showcase ✨_ +- OpenAI Whisper +- SenseVoice - +**Text-to-Speech Services** -_✨ Web Dashboard ✨_ +- OpenAI TTS +- Gemini TTS +- GPT-Sovits-Inference +- GPT-Sovits +- FishAudio +- Edge TTS +- Alibaba Cloud Bailian TTS +- Azure TTS +- Minimax TTS +- Volcano Engine TTS -![webchat](https://drive.soulter.top/f/vlsA/ezgif-5-fb044b2542.gif) +## ❤️ Contributing -_✨ Built-in Web Chat Interface ✨_ +Issues and Pull Requests are always welcome! Feel free to submit your changes to this project :) -
+### How to Contribute + +You can contribute by reviewing issues or helping with pull request reviews. Any issues or PRs are welcome to encourage community participation. Of course, these are just suggestions—you can contribute in any way you like. For adding new features, please discuss through an Issue first. + +### Development Environment + +AstrBot uses `ruff` for code formatting and linting. + +```bash +git clone https://github.com/AstrBotDevs/AstrBot +pip install pre-commit +pre-commit install +``` + +## ❤️ Special Thanks + +Special thanks to all Contributors and plugin developers for their contributions to AstrBot ❤️ + + + + + +Additionally, the birth of this project would not have been possible without the help of the following open-source projects: + +- [NapNeko/NapCatQQ](https://github.com/NapNeko/NapCatQQ) - The amazing cat framework ## ⭐ Star History -> [!TIP] -> If this project helps you, please give it a star <3 +> [!TIP] +> If this project has helped you in your life or work, or if you're interested in its future development, please give the project a Star. It's the driving force behind maintaining this open-source project <3
- -[![Star History Chart](https://api.star-history.com/svg?repos=soulter/astrbot&type=Date)](https://star-history.com/#soulter/astrbot&Date) + +[![Star History Chart](https://api.star-history.com/svg?repos=astrbotdevs/astrbot&type=Date)](https://star-history.com/#astrbotdevs/astrbot&Date)
-## Disclaimer - -1. Licensed under `AGPL-v3`. -2. WeChat integration uses [Gewechat](https://github.com/Devo919/Gewechat). Use at your own risk with non-critical accounts. -3. Users must comply with local laws and regulations. - - - + _私は、高性能ですから!_ - diff --git a/README_ja.md b/README_ja.md index 8e648c8e..3fce0187 100644 --- a/README_ja.md +++ b/README_ja.md @@ -1,167 +1,233 @@ -

- -![6e1279651f16d7fdf4727558b72bbaf1](https://github.com/user-attachments/assets/ead4c551-fc3c-48f7-a6f7-afbfdb820512) +![AstrBot-Logo-Simplified](https://github.com/user-attachments/assets/ffd99b6b-3272-4682-beaa-6fe74250f7d9)

-_✨ 簡単に使えるマルチプラットフォーム LLM チャットボットおよび開発フレームワーク ✨_ +
+
Soulter%2FAstrBot | Trendshift - -[![GitHub release (latest by date)](https://img.shields.io/github/v/release/Soulter/AstrBot)](https://github.com/Soulter/AstrBot/releases/latest) -python -Docker pull -Static Badge -[![wakatime](https://wakatime.com/badge/user/915e5316-99c6-4563-a483-ef186cf000c9/project/018e705a-a1a7-409a-a849-3013485e6c8e.svg)](https://wakatime.com/badge/user/915e5316-99c6-4563-a483-ef186cf000c9/project/018e705a-a1a7-409a-a849-3013485e6c8e) -![Dynamic JSON Badge](https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fapi.soulter.top%2Fastrbot%2Fstats&query=v&label=7%E6%97%A5%E6%B6%88%E6%81%AF%E4%B8%8A%E8%A1%8C%E9%87%8F&cacheSeconds=3600) -[![codecov](https://codecov.io/gh/Soulter/AstrBot/graph/badge.svg?token=FF3P5967B8)](https://codecov.io/gh/Soulter/AstrBot) - -ドキュメントを見る ｜ -問題を報告する +Featured｜HelloGitHub
-AstrBot は、疎結合、非同期、複数のメッセージプラットフォームに対応したデプロイ、使いやすいプラグインシステム、および包括的な大規模言語モデル（LLM）接続機能を備えたチャットボットおよび開発フレームワークです。 +
-## ✨ 主な機能 +
+ +python +Docker pull +QQ_community +Telegram_community + +
-1. **大規模言語モデルの対話**。OpenAI API、Google Gemini、Llama、Deepseek、ChatGLM など、さまざまな大規模言語モデルをサポートし、Ollama、LLMTuner を介してローカルにデプロイされた大規模モデルをサポートします。多輪対話、人格シナリオ、多モーダル機能を備え、画像理解、音声からテキストへの変換（Whisper）をサポートします。 -2. **複数のメッセージプラットフォームの接続**。QQ（OneBot）、QQ チャンネル、Feishu、Telegram への接続をサポートします。今後、DingTalk、Discord、WhatsApp、Xiaoai 音響をサポートする予定です。レート制限、ホワイトリスト、キーワードフィルタリング、Baidu コンテンツ監査をサポートします。 -3. **エージェント**。一部のエージェント機能をネイティブにサポートし、コードエグゼキューター、自然言語タスク、ウェブ検索などを提供します。[Dify プラットフォーム](https://dify.ai/)と連携し、Dify スマートアシスタント、ナレッジベース、Dify ワークフローを簡単に接続できます。 -4. **プラグインの拡張**。深く最適化されたプラグインメカニズムを備え、[プラグインの開発](https://astrbot.app/dev/plugin.html)をサポートし、機能を拡張できます。複数のプラグインのインストールをサポートします。 -5. **ビジュアル管理パネル**。設定の視覚的な変更、プラグイン管理、ログの表示などをサポートし、設定の難易度を低減します。WebChat を統合し、パネル上で大規模モデルと対話できます。 -6. **高い安定性と高いモジュール性**。イベントバスとパイプラインに基づくアーキテクチャ設計により、高度にモジュール化され、低結合です。 +
-> [!TIP] -> 管理パネルのオンラインデモを体験する: [https://demo.astrbot.app/](https://demo.astrbot.app/) -> -> ユーザー名: `astrbot`, パスワード: `astrbot`。LLM が設定されていないため、チャットページで大規模モデルを使用することはできません。（デモのログインパスワードを変更しないでください 😭） +中文 ｜ +English ｜ +ドキュメント ｜ +Blog ｜ +ロードマップ ｜ +Issue +
-## ✨ 使用方法 +AstrBot は、オープンソースのオールインワン Agent チャットボットプラットフォーム及び開発フレームワークです。 -#### Docker デプロイ +## 主な機能 -公式ドキュメント [Docker を使用して AstrBot をデプロイする](https://astrbot.app/deploy/astrbot/docker.html#%E4%BD%BF%E7%94%A8-docker-%E9%83%A8%E7%BD%B2-astrbot) を参照してください。 +1. **大規模言語モデル対話**。多様な大規模言語モデルサービスとの統合をサポート。マルチモーダル、ツール呼び出し、MCP、ネイティブナレッジベース、キャラクター設定などの機能を搭載。 +2. **マルチメッセージプラットフォームサポート**。QQ、WeChat Work、WeChat公式アカウント、Feishu、Telegram、DingTalk、Discord、KOOK などのプラットフォームと統合可能。レート制限、ホワイトリスト、Baidu コンテンツ審査をサポート。 +3. **Agent**。完全に最適化された Agentic 機能。マルチターンツール呼び出し、内蔵サンドボックスコード実行環境、Web 検索などの機能をサポート。 +4. **プラグイン拡張**。深く最適化されたプラグインメカニズムで、[プラグイン開発](https://astrbot.app/dev/plugin.html)による機能拡張をサポート。豊富なコミュニティプラグインエコシステム。 +5. **WebUI**。ビジュアル設定とボット管理、充実した機能。 -#### Windows ワンクリックインストーラーのデプロイ +## デプロイ方法 -コンピュータに Python（>3.10）がインストールされている必要があります。公式ドキュメント [Windows ワンクリックインストーラーを使用して AstrBot をデプロイする](https://astrbot.app/deploy/astrbot/windows.html) を参照してください。 +#### Docker デプロイ（推奨 🥳） -#### Replit デプロイ +Docker / Docker Compose を使用した AstrBot のデプロイを推奨します。 -[![Run on Repl.it](https://repl.it/badge/github/Soulter/AstrBot)](https://repl.it/github/Soulter/AstrBot) +公式ドキュメント [Docker を使用した AstrBot のデプロイ](https://astrbot.app/deploy/astrbot/docker.html#%E4%BD%BF%E7%94%A8-docker-%E9%83%A8%E7%BD%B2-astrbot) をご参照ください。 + +#### 宝塔パネルデプロイ + +AstrBot は宝塔パネルと提携し、宝塔パネルに公開されています。 + +公式ドキュメント [宝塔パネルデプロイ](https://astrbot.app/deploy/astrbot/btpanel.html) をご参照ください。 + +#### 1Panel デプロイ + +AstrBot は 1Panel 公式により 1Panel パネルに公開されています。 + +公式ドキュメント [1Panel デプロイ](https://astrbot.app/deploy/astrbot/1panel.html) をご参照ください。 + +#### 雨云でのデプロイ + +AstrBot は雨云公式によりクラウドアプリケーションプラットフォームに公開され、ワンクリックでデプロイ可能です。 + +[![Deploy on RainYun](https://rainyun-apps.cn-nb1.rains3.com/materials/deploy-on-rainyun-en.svg)](https://app.rainyun.com/apps/rca/store/5994?ref=NjU1ODg0) + +#### Replit でのデプロイ + +コミュニティ貢献によるデプロイ方法。 + +[![Run on Repl.it](https://repl.it/badge/github/AstrBotDevs/AstrBot)](https://repl.it/github/AstrBotDevs/AstrBot) + +#### Windows ワンクリックインストーラーデプロイ + +公式ドキュメント [Windows ワンクリックインストーラーを使用した AstrBot のデプロイ](https://astrbot.app/deploy/astrbot/windows.html) をご参照ください。 #### CasaOS デプロイ -コミュニティが提供するデプロイ方法です。 +コミュニティ貢献によるデプロイ方法。 -公式ドキュメント [ソースコードを使用して AstrBot をデプロイする](https://astrbot.app/deploy/astrbot/casaos.html) を参照してください。 +公式ドキュメント [CasaOS デプロイ](https://astrbot.app/deploy/astrbot/casaos.html) をご参照ください。 #### 手動デプロイ -公式ドキュメント [ソースコードを使用して AstrBot をデプロイする](https://astrbot.app/deploy/astrbot/cli.html) を参照してください。 +まず uv をインストールします: -## ⚡ メッセージプラットフォームのサポート状況 +```bash +pip install uv +``` -| プラットフォーム | サポート状況 | 詳細 | メッセージタイプ | -| -------- | ------- | ------- | ------ | -| QQ(公式ロボットインターフェース) | ✔ | プライベートチャット、グループチャット、QQ チャンネルプライベートチャット、グループチャット | テキスト、画像 | -| QQ(OneBot) | ✔ | プライベートチャット、グループチャット | テキスト、画像、音声 | -| WeChat(個人アカウント) | ✔ | WeChat 個人アカウントのプライベートチャット、グループチャット | テキスト、画像、音声 | -| [Telegram](https://github.com/Soulter/astrbot_plugin_telegram) | ✔ | プライベートチャット、グループチャット | テキスト、画像 | -| [WeChat(企業 WeChat)](https://github.com/Soulter/astrbot_plugin_wecom) | ✔ | プライベートチャット | テキスト、画像、音声 | -| Feishu | ✔ | グループチャット | テキスト、画像 | -| WeChat 対話オープンプラットフォーム | 🚧 | 計画中 | - | -| Discord | 🚧 | 計画中 | - | -| WhatsApp | 🚧 | 計画中 | - | -| Xiaoai 音響 | 🚧 | 計画中 | - | +Git Clone で AstrBot をインストール: -# 🦌 今後のロードマップ +```bash +git clone https://github.com/AstrBotDevs/AstrBot && cd AstrBot +uv run main.py +``` -> [!TIP] -> Issue でさらに多くの提案を歓迎します <3 +または、公式ドキュメント [ソースコードから AstrBot をデプロイ](https://astrbot.app/deploy/astrbot/cli.html) をご参照ください。 -- [ ] 現在のすべてのプラットフォームアダプターの機能の一貫性を確保し、改善する -- [ ] プラグインインターフェースの最適化 -- [ ] GPT-Sovits などの TTS サービスをデフォルトでサポート -- [ ] "チャット強化" 部分を完成させ、永続的な記憶をサポート -- [ ] i18n の計画 +## 🌍 コミュニティ -## ❤️ 貢献 +### QQ グループ -Issue や Pull Request を歓迎します！このプロジェクトに変更を加えるだけです :) +- 1群:322154837 +- 3群:630166526 +- 5群:822130018 +- 6群:753075035 +- 開発者群:975206796 -新機能の追加については、まず Issue で議論してください。 +### Telegram グループ -## 🌟 サポート +Telegram_community -- このプロジェクトに Star を付けてください！ -- [愛発電](https://afdian.com/a/soulter)で私をサポートしてください！ -- [WeChat](https://drive.soulter.top/f/pYfA/d903f4fa49a496fda3f16d2be9e023b5.png)で私をサポートしてください~ +### Discord サーバー -## ✨ デモ +Discord_community -> [!NOTE] -> コードエグゼキューターのファイル入力/出力は現在 Napcat(QQ)、Lagrange(QQ) でのみテストされています +## サポートされているメッセージプラットフォーム -
+**公式メンテナンス** - +- QQ (公式プラットフォーム & OneBot) +- Telegram +- WeChat Work アプリケーション & WeChat Work インテリジェントボット +- WeChat カスタマーサービス & WeChat 公式アカウント +- Feishu (Lark) +- DingTalk +- Slack +- Discord +- Satori +- Misskey +- WhatsApp (近日対応予定) +- LINE (近日対応予定) -_✨ Docker ベースのサンドボックス化されたコードエグゼキューター（ベータテスト中）✨_ +**コミュニティメンテナンス** - +- [KOOK](https://github.com/wuyan1003/astrbot_plugin_kook_adapter) +- [VoceChat](https://github.com/HikariFroya/astrbot_plugin_vocechat) +- [Bilibili ダイレクトメッセージ](https://github.com/Hina-Chat/astrbot_plugin_bilibili_adapter) +- [wxauto](https://github.com/luosheng520qaq/wxauto-repost-onebotv11) -_✨ 多モーダル、ウェブ検索、長文の画像変換（設定可能）✨_ +## サポートされているモデルサービス - +**大規模言語モデルサービス** -_✨ 自然言語タスク ✨_ +- OpenAI および互換サービス +- Anthropic +- Google Gemini +- Moonshot AI +- 智谱 AI +- DeepSeek +- Ollama (セルフホスト) +- LM Studio (セルフホスト) +- [優云智算](https://www.compshare.cn/?ytag=GPU_YY-gh_astrbot&referral_code=FV7DcGowN4hB5UuXKgpE74) +- [302.AI](https://share.302.ai/rr1M3l) +- [小馬算力](https://www.tokenpony.cn/3YPyf) +- [硅基流動](https://docs.siliconflow.cn/cn/usercases/use-siliconcloud-in-astrbot) +- [PPIO 派欧云](https://ppio.com/user/register?invited_by=AIOONE) +- ModelScope +- OneAPI - - +**LLMOps プラットフォーム** -_✨ プラグインシステム - 一部のプラグインの展示 ✨_ +- Dify +- Alibaba Cloud 百炼アプリケーション +- Coze - +**音声認識サービス** -_✨ 管理パネル ✨_ +- OpenAI Whisper +- SenseVoice -![webchat](https://drive.soulter.top/f/vlsA/ezgif-5-fb044b2542.gif) +**音声合成サービス** -_✨ 内蔵 Web Chat、オンラインでボットと対話 ✨_ +- OpenAI TTS +- Gemini TTS +- GPT-Sovits-Inference +- GPT-Sovits +- FishAudio +- Edge TTS +- Alibaba Cloud 百炼 TTS +- Azure TTS +- Minimax TTS +- Volcano Engine TTS -
+## ❤️ コントリビューション + +Issue や Pull Request は大歓迎です!このプロジェクトに変更を送信してください :) + +### コントリビュート方法 + +Issue を確認したり、PR(プルリクエスト)のレビューを手伝うことで貢献できます。どんな Issue や PR への参加も歓迎され、コミュニティ貢献を促進します。もちろん、これらは提案に過ぎず、どんな方法でも貢献できます。新機能の追加については、まず Issue で議論してください。 + +### 開発環境 + +AstrBot はコードのフォーマットとチェックに `ruff` を使用しています。 + +```bash +git clone https://github.com/AstrBotDevs/AstrBot +pip install pre-commit +pre-commit install +``` + +## ❤️ Special Thanks + +AstrBot への貢献をしていただいたすべてのコントリビューターとプラグイン開発者に特別な感謝を ❤️ + + + + + +また、このプロジェクトの誕生は以下のオープンソースプロジェクトの助けなしには実現できませんでした: + +- [NapNeko/NapCatQQ](https://github.com/NapNeko/NapCatQQ) - 素晴らしい猫猫フレームワーク ## ⭐ Star History > [!TIP] -> このプロジェクトがあなたの生活や仕事に役立った場合、またはこのプロジェクトの将来の発展に関心がある場合は、プロジェクトに Star を付けてください。これはこのオープンソースプロジェクトを維持するためのモチベーションです <3 +> このプロジェクトがあなたの生活や仕事に役立ったり、このプロジェクトの今後の発展に関心がある場合は、プロジェクトに Star をください。これがこのオープンソースプロジェクトを維持する原動力です <3
-[![Star History Chart](https://api.star-history.com/svg?repos=soulter/astrbot&type=Date)](https://star-history.com/#soulter/astrbot&Date) +[![Star History Chart](https://api.star-history.com/svg?repos=astrbotdevs/astrbot&type=Date)](https://star-history.com/#astrbotdevs/astrbot&Date)
-## スポンサー - -[](https://api.gitsponsors.com/api/badge/link?p=XEpbdGxlitw/RbcwiTX93UMzNK/jgDYC8NiSzamIPMoKvG2lBFmyXhSS/b0hFoWlBBMX2L5X5CxTDsUdyvcIEHTOfnkXz47UNOZvMwyt5CzbYpq0SEzsSV1OJF1cCo90qC/ZyYKYOWedal3MhZ3ikw==) - -## 免責事項 - -1. このプロジェクトは `AGPL-v3` オープンソースライセンスの下で保護されています。 -2. このプロジェクトを使用する際は、現地の法律および規制を遵守してください。 - - + _私は、高性能ですから!_ diff --git a/astrbot/api/__init__.py b/astrbot/api/__init__.py index 540171f1..5d15dedc 100644 --- a/astrbot/api/__init__.py +++ b/astrbot/api/__init__.py @@ -1,20 +1,19 @@ -from astrbot.core.config.astrbot_config import AstrBotConfig from astrbot import logger -from astrbot.core import html_renderer -from astrbot.core import sp -from astrbot.core.star.register import register_llm_tool as llm_tool -from astrbot.core.star.register import register_agent as agent -from astrbot.core.agent.tool import ToolSet, FunctionTool +from astrbot.core import html_renderer, sp +from astrbot.core.agent.tool import FunctionTool, ToolSet from astrbot.core.agent.tool_executor import BaseFunctionToolExecutor +from astrbot.core.config.astrbot_config import AstrBotConfig +from astrbot.core.star.register import register_agent as agent +from astrbot.core.star.register import register_llm_tool as llm_tool __all__ = [ "AstrBotConfig", - "logger", + "BaseFunctionToolExecutor", + "FunctionTool", + "ToolSet", + "agent", "html_renderer", "llm_tool", - "agent", + "logger", "sp", - "ToolSet", - "FunctionTool", - "BaseFunctionToolExecutor", ] diff --git a/astrbot/api/all.py b/astrbot/api/all.py index 2463dbc2..df3e1170 100644 --- a/astrbot/api/all.py +++ b/astrbot/api/all.py @@ -36,7 +36,8 @@ from astrbot.core.star.config import * # provider -from astrbot.core.provider import Provider, Personality, ProviderMetaData +from astrbot.core.provider import Provider, ProviderMetaData +from astrbot.core.db.po import Personality # platform from astrbot.core.platform import ( diff --git a/astrbot/api/event/__init__.py b/astrbot/api/event/__init__.py index 1f2fce64..2b8dd5a9 100644 --- a/astrbot/api/event/__init__.py +++ b/astrbot/api/event/__init__.py @@ -1,18 +1,17 @@ from astrbot.core.message.message_event_result import ( - MessageEventResult, - MessageChain, CommandResult, EventResultType, + MessageChain, + MessageEventResult, ResultContentType, ) - from astrbot.core.platform import AstrMessageEvent __all__ = [ - "MessageEventResult", - "MessageChain", + "AstrMessageEvent", "CommandResult", "EventResultType", - "AstrMessageEvent", + "MessageChain", + "MessageEventResult", "ResultContentType", ] diff --git a/astrbot/api/event/filter/__init__.py b/astrbot/api/event/filter/__init__.py index d63850e4..a8d2b426 100644 --- a/astrbot/api/event/filter/__init__.py +++ b/astrbot/api/event/filter/__init__.py @@ -1,51 +1,52 @@ -from astrbot.core.star.register import ( - register_command as command, - register_command_group as command_group, - register_event_message_type as event_message_type, - register_regex as regex, - register_platform_adapter_type as platform_adapter_type, - register_permission_type as permission_type, - register_custom_filter as custom_filter, - register_on_astrbot_loaded as on_astrbot_loaded, - register_on_platform_loaded as on_platform_loaded, - register_on_llm_request as on_llm_request, - register_on_llm_response as on_llm_response, - register_llm_tool as llm_tool, - register_on_decorating_result as on_decorating_result, - register_after_message_sent as after_message_sent, -) - -from astrbot.core.star.filter.event_message_type import ( - EventMessageTypeFilter, - EventMessageType, -) -from astrbot.core.star.filter.platform_adapter_type import ( - PlatformAdapterTypeFilter, - PlatformAdapterType, -) -from astrbot.core.star.filter.permission import PermissionTypeFilter, PermissionType from astrbot.core.star.filter.custom_filter import CustomFilter +from astrbot.core.star.filter.event_message_type import ( + EventMessageType, + EventMessageTypeFilter, +) +from astrbot.core.star.filter.permission import PermissionType, PermissionTypeFilter +from astrbot.core.star.filter.platform_adapter_type import ( + PlatformAdapterType, + PlatformAdapterTypeFilter, +) +from astrbot.core.star.register import register_after_message_sent as after_message_sent +from astrbot.core.star.register import register_command as command +from astrbot.core.star.register import register_command_group as command_group +from astrbot.core.star.register import register_custom_filter as custom_filter +from astrbot.core.star.register import register_event_message_type as event_message_type +from astrbot.core.star.register import register_llm_tool as llm_tool +from astrbot.core.star.register import register_on_astrbot_loaded as on_astrbot_loaded +from astrbot.core.star.register import ( + register_on_decorating_result as on_decorating_result, +) +from astrbot.core.star.register import register_on_llm_request as on_llm_request +from astrbot.core.star.register import register_on_llm_response as on_llm_response +from astrbot.core.star.register import register_on_platform_loaded as on_platform_loaded +from astrbot.core.star.register import register_permission_type as permission_type +from astrbot.core.star.register import ( + register_platform_adapter_type as platform_adapter_type, +) +from astrbot.core.star.register import register_regex as regex __all__ = [ + "CustomFilter", + "EventMessageType", + "EventMessageTypeFilter", + "PermissionType", + "PermissionTypeFilter", + "PlatformAdapterType", + "PlatformAdapterTypeFilter", + "after_message_sent", "command", "command_group", - "event_message_type", - "regex", - "platform_adapter_type", - "permission_type", - "EventMessageTypeFilter", - "EventMessageType", - "PlatformAdapterTypeFilter", - "PlatformAdapterType", - "PermissionTypeFilter", - "CustomFilter", "custom_filter", - "PermissionType", - "on_astrbot_loaded", - "on_platform_loaded", - "on_llm_request", + "event_message_type", "llm_tool", + "on_astrbot_loaded", "on_decorating_result", - "after_message_sent", + "on_llm_request", "on_llm_response", + "on_platform_loaded", + "permission_type", + "platform_adapter_type", + "regex", ] diff --git a/astrbot/api/platform/__init__.py b/astrbot/api/platform/__init__.py index 5a98c590..6a182c32 100644 --- a/astrbot/api/platform/__init__.py +++ b/astrbot/api/platform/__init__.py @@ -1,23 +1,22 @@ +from astrbot.core.message.components import * from astrbot.core.platform import ( - AstrMessageEvent, - Platform, AstrBotMessage, + AstrMessageEvent, + Group, MessageMember, MessageType, + Platform, PlatformMetadata, - Group, ) - from astrbot.core.platform.register import register_platform_adapter -from astrbot.core.message.components import * __all__ = [ - "AstrMessageEvent", - "Platform", "AstrBotMessage", + "AstrMessageEvent", + "Group", "MessageMember", "MessageType", + "Platform", "PlatformMetadata", "register_platform_adapter", - "Group", ] diff --git a/astrbot/api/provider/__init__.py b/astrbot/api/provider/__init__.py index 9b1ade50..f62b340f 100644 --- a/astrbot/api/provider/__init__.py +++ b/astrbot/api/provider/__init__.py @@ -1,17 +1,18 @@ -from astrbot.core.provider import Provider, STTProvider, Personality +from astrbot.core.db.po import Personality +from astrbot.core.provider import Provider, STTProvider from astrbot.core.provider.entities import ( + LLMResponse, + ProviderMetaData, ProviderRequest, ProviderType, - ProviderMetaData, - LLMResponse, ) __all__ = [ - "Provider", - "STTProvider", + "LLMResponse", "Personality", + "Provider", + "ProviderMetaData", "ProviderRequest", "ProviderType", - "ProviderMetaData", - "LLMResponse", + "STTProvider", ] diff --git a/astrbot/api/star/__init__.py b/astrbot/api/star/__init__.py index 1b33923f..63db07a7 100644 --- a/astrbot/api/star/__init__.py +++ b/astrbot/api/star/__init__.py @@ -1,8 +1,7 @@ +from astrbot.core.star import Context, Star, StarTools +from astrbot.core.star.config import * from astrbot.core.star.register import ( register_star as register, # 注册插件（Star） ) -from astrbot.core.star import Context, Star, StarTools -from astrbot.core.star.config import * - -__all__ = ["register", "Context", "Star", "StarTools"] +__all__ = ["Context", "Star", "StarTools", "register"] diff --git a/astrbot/api/util/__init__.py b/astrbot/api/util/__init__.py index a66206e0..1be3152d 100644 --- a/astrbot/api/util/__init__.py +++ b/astrbot/api/util/__init__.py @@ -1,7 +1,7 @@ from astrbot.core.utils.session_waiter import ( - SessionWaiter, SessionController, + SessionWaiter, session_waiter, ) -__all__ = ["SessionWaiter", "SessionController", "session_waiter"] +__all__ = ["SessionController", "SessionWaiter", "session_waiter"] diff --git a/astrbot/cli/__main__.py b/astrbot/cli/__main__.py index f2b6651f..40c46de7 100644 --- a/astrbot/cli/__main__.py +++ b/astrbot/cli/__main__.py @@ -1,11 +1,11 @@ -""" -AstrBot CLI入口 -""" +"""AstrBot CLI入口""" + +import sys import click -import sys + from . import __version__ -from .commands import init, run, plug, conf +from .commands import conf, init, plug, run logo_tmpl = r""" ___ _______.___________..______ .______ ______ .___________. diff --git a/astrbot/cli/commands/__init__.py b/astrbot/cli/commands/__init__.py index 9fa9149e..1d3e0bca 100644 --- a/astrbot/cli/commands/__init__.py +++ b/astrbot/cli/commands/__init__.py @@ -1,6 +1,6 @@ -from .cmd_init import init -from .cmd_run import run -from .cmd_plug import plug from .cmd_conf import conf +from .cmd_init import init +from .cmd_plug import plug +from .cmd_run import run -__all__ = ["init", "run", "plug", "conf"] +__all__ = ["conf", "init", "plug", "run"] diff --git a/astrbot/cli/commands/cmd_conf.py b/astrbot/cli/commands/cmd_conf.py index fea654f2..a9bd40f0 100644 --- a/astrbot/cli/commands/cmd_conf.py +++ b/astrbot/cli/commands/cmd_conf.py @@ -1,9 +1,12 @@ -import json -import click import hashlib +import json import zoneinfo -from typing import Any, Callable -from ..utils import get_astrbot_root, check_astrbot_root +from collections.abc import Callable +from typing import Any + +import click + +from ..utils import check_astrbot_root, get_astrbot_root def _validate_log_level(value: str) -> str: @@ -11,7 +14,7 @@ def _validate_log_level(value: str) -> str: value = value.upper() if value not in ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]: raise click.ClickException( - "日志级别必须是 DEBUG/INFO/WARNING/ERROR/CRITICAL 之一" + "日志级别必须是 DEBUG/INFO/WARNING/ERROR/CRITICAL 之一", ) return value @@ -73,7 +76,7 @@ def _load_config() -> dict[str, Any]: root = get_astrbot_root() if not check_astrbot_root(root): raise click.ClickException( - f"{root}不是有效的 AstrBot 根目录，如需初始化请使用 astrbot init" + f"{root}不是有效的 AstrBot 根目录，如需初始化请使用 astrbot init", ) config_path = root / "data" / "cmd_config.json" @@ -88,7 +91,7 @@ def _load_config() -> dict[str, Any]: try: return json.loads(config_path.read_text(encoding="utf-8-sig")) except json.JSONDecodeError as e: - raise click.ClickException(f"配置文件解析失败: {str(e)}") + raise click.ClickException(f"配置文件解析失败: {e!s}") def _save_config(config: dict[str, Any]) -> None: @@ -96,7 +99,8 @@ def _save_config(config: dict[str, Any]) -> None: config_path = get_astrbot_root() / "data" / "cmd_config.json" config_path.write_text( - json.dumps(config, ensure_ascii=False, indent=2), encoding="utf-8-sig" + json.dumps(config, ensure_ascii=False, indent=2), + encoding="utf-8-sig", ) @@ -108,7 +112,7 @@ def _set_nested_item(obj: dict[str, Any], path: str, value: Any) -> None: obj[part] = {} elif not isinstance(obj[part], dict): raise click.ClickException( - f"配置路径冲突: {'.'.join(parts[: parts.index(part) + 1])} 不是字典" + f"配置路径冲突: {'.'.join(parts[: parts.index(part) + 1])} 不是字典", ) obj = obj[part] obj[parts[-1]] = value @@ -140,7 +144,6 @@ def conf(): - callback_api_base: 回调接口基址 """ - pass @conf.command(name="set") @@ -148,7 +151,7 @@ def conf(): @click.argument("value") def set_config(key: str, value: str): """设置配置项的值""" - if key not in CONFIG_VALIDATORS.keys(): + if key not in CONFIG_VALIDATORS: raise click.ClickException(f"不支持的配置项: {key}") config = _load_config() @@ -170,17 +173,17 @@ def set_config(key: str, value: str): except KeyError: raise click.ClickException(f"未知的配置项: {key}") except Exception as e: - raise click.UsageError(f"设置配置失败: {str(e)}") + raise click.UsageError(f"设置配置失败: {e!s}") @conf.command(name="get") @click.argument("key", required=False) -def get_config(key: str = None): +def get_config(key: str | None = None): """获取配置项的值，不提供key则显示所有可配置项""" config = _load_config() if key: - if key not in CONFIG_VALIDATORS.keys(): + if key not in CONFIG_VALIDATORS: raise click.ClickException(f"不支持的配置项: {key}") try: @@ -191,10 +194,10 @@ def get_config(key: str = None): except KeyError: raise click.ClickException(f"未知的配置项: {key}") except Exception as e: - raise click.UsageError(f"获取配置失败: {str(e)}") + raise click.UsageError(f"获取配置失败: {e!s}") else: click.echo("当前配置:") - for key in CONFIG_VALIDATORS.keys(): + for key in CONFIG_VALIDATORS: try: value = ( "********" diff --git a/astrbot/cli/commands/cmd_init.py b/astrbot/cli/commands/cmd_init.py index d9a42f82..6c0c34b9 100644 --- a/astrbot/cli/commands/cmd_init.py +++ b/astrbot/cli/commands/cmd_init.py @@ -1,4 +1,5 @@ import asyncio +from pathlib import Path import click from filelock import FileLock, Timeout @@ -6,14 +7,14 @@ from filelock import FileLock, Timeout from ..utils import check_dashboard, get_astrbot_root -async def initialize_astrbot(astrbot_root) -> None: +async def initialize_astrbot(astrbot_root: Path) -> None: """执行 AstrBot 初始化逻辑""" dot_astrbot = astrbot_root / ".astrbot" if not dot_astrbot.exists(): click.echo(f"Current Directory: {astrbot_root}") click.echo( - "如果你确认这是 Astrbot root directory, 你需要在当前目录下创建一个 .astrbot 文件标记该目录为 AstrBot 的数据目录。" + "如果你确认这是 Astrbot root directory, 你需要在当前目录下创建一个 .astrbot 文件标记该目录为 AstrBot 的数据目录。", ) if click.confirm( f"请检查当前目录是否正确，确认正确请回车: {astrbot_root}", diff --git a/astrbot/cli/commands/cmd_plug.py b/astrbot/cli/commands/cmd_plug.py index b250ede4..a1099de1 100644 --- a/astrbot/cli/commands/cmd_plug.py +++ b/astrbot/cli/commands/cmd_plug.py @@ -1,31 +1,29 @@ import re +import shutil from pathlib import Path import click -import shutil - from ..utils import ( - get_git_repo, - build_plug_list, - manage_plugin, PluginStatus, + build_plug_list, check_astrbot_root, get_astrbot_root, + get_git_repo, + manage_plugin, ) @click.group() def plug(): """插件管理""" - pass def _get_data_path() -> Path: base = get_astrbot_root() if not check_astrbot_root(base): raise click.ClickException( - f"{base}不是有效的 AstrBot 根目录，如需初始化请使用 astrbot init" + f"{base}不是有效的 AstrBot 根目录，如需初始化请使用 astrbot init", ) return (base / "data").resolve() @@ -41,7 +39,7 @@ def display_plugins(plugins, title=None, color=None): desc = p["desc"][:30] + ("..." if len(p["desc"]) > 30 else "") click.echo( f"{p['name']:<20} {p['version']:<10} {p['status']:<10} " - f"{p['author']:<15} {desc:<30}" + f"{p['author']:<15} {desc:<30}", ) @@ -78,7 +76,7 @@ def new(name: str): f"desc: {desc}\n" f"version: {version}\n" f"author: {author}\n" - f"repo: {repo}\n" + f"repo: {repo}\n", ) # 重写 README.md @@ -86,7 +84,7 @@ def new(name: str): f.write(f"# {name}\n\n{desc}\n\n# 支持\n\n[帮助文档](https://astrbot.app)\n") # 重写 main.py - with open(plug_path / "main.py", "r", encoding="utf-8") as f: + with open(plug_path / "main.py", encoding="utf-8") as f: content = f.read() new_content = content.replace( diff --git a/astrbot/cli/commands/cmd_run.py b/astrbot/cli/commands/cmd_run.py index 38113744..9333f1b8 100644 --- a/astrbot/cli/commands/cmd_run.py +++ b/astrbot/cli/commands/cmd_run.py @@ -1,19 +1,18 @@ +import asyncio import os import sys +import traceback from pathlib import Path import click -import asyncio -import traceback - from filelock import FileLock, Timeout -from ..utils import check_dashboard, check_astrbot_root, get_astrbot_root +from ..utils import check_astrbot_root, check_dashboard, get_astrbot_root async def run_astrbot(astrbot_root: Path): """运行 AstrBot""" - from astrbot.core import logger, LogManager, LogBroker, db_helper + from astrbot.core import LogBroker, LogManager, db_helper, logger from astrbot.core.initial_loader import InitialLoader await check_dashboard(astrbot_root / "data") @@ -38,7 +37,7 @@ def run(reload: bool, port: str) -> None: if not check_astrbot_root(astrbot_root): raise click.ClickException( - f"{astrbot_root}不是有效的 AstrBot 根目录，如需初始化请使用 astrbot init" + f"{astrbot_root}不是有效的 AstrBot 根目录，如需初始化请使用 astrbot init", ) os.environ["ASTRBOT_ROOT"] = str(astrbot_root) diff --git a/astrbot/cli/utils/__init__.py b/astrbot/cli/utils/__init__.py index 9989dcf2..3830682f 100644 --- a/astrbot/cli/utils/__init__.py +++ b/astrbot/cli/utils/__init__.py @@ -1,18 +1,18 @@ from .basic import ( - get_astrbot_root, check_astrbot_root, check_dashboard, + get_astrbot_root, ) -from .plugin import get_git_repo, manage_plugin, build_plug_list, PluginStatus +from .plugin import PluginStatus, build_plug_list, get_git_repo, manage_plugin from .version_comparator import VersionComparator __all__ = [ - "get_astrbot_root", + "PluginStatus", + "VersionComparator", + "build_plug_list", "check_astrbot_root", "check_dashboard", + "get_astrbot_root", "get_git_repo", "manage_plugin", - "build_plug_list", - "VersionComparator", - "PluginStatus", ] diff --git a/astrbot/cli/utils/basic.py b/astrbot/cli/utils/basic.py index fabced48..5dbe2900 100644 --- a/astrbot/cli/utils/basic.py +++ b/astrbot/cli/utils/basic.py @@ -21,8 +21,9 @@ def get_astrbot_root() -> Path: async def check_dashboard(astrbot_root: Path) -> None: """检查是否安装了dashboard""" - from astrbot.core.utils.io import get_dashboard_version, download_dashboard from astrbot.core.config.default import VERSION + from astrbot.core.utils.io import download_dashboard, get_dashboard_version + from .version_comparator import VersionComparator try: @@ -48,19 +49,18 @@ async def check_dashboard(astrbot_root: Path) -> None: if VersionComparator.compare_version(VERSION, dashboard_version) <= 0: click.echo("管理面板已是最新版本") return - else: - try: - version = dashboard_version.split("v")[1] - click.echo(f"管理面板版本: {version}") - await download_dashboard( - path="data/dashboard.zip", - extract_path=str(astrbot_root), - version=f"v{VERSION}", - latest=False, - ) - except Exception as e: - click.echo(f"下载管理面板失败: {e}") - return + try: + version = dashboard_version.split("v")[1] + click.echo(f"管理面板版本: {version}") + await download_dashboard( + path="data/dashboard.zip", + extract_path=str(astrbot_root), + version=f"v{VERSION}", + latest=False, + ) + except Exception as e: + click.echo(f"下载管理面板失败: {e}") + return except FileNotFoundError: click.echo("初始化管理面板目录...") try: diff --git a/astrbot/cli/utils/plugin.py b/astrbot/cli/utils/plugin.py index cd1fcd97..cd76a07c 100644 --- a/astrbot/cli/utils/plugin.py +++ b/astrbot/cli/utils/plugin.py @@ -1,14 +1,14 @@ import shutil import tempfile - -import httpx -import yaml from enum import Enum from io import BytesIO from pathlib import Path from zipfile import ZipFile import click +import httpx +import yaml + from .version_comparator import VersionComparator @@ -32,7 +32,8 @@ def get_git_repo(url: str, target_path: Path, proxy: str | None = None): release_url = f"https://api.github.com/repos/{author}/{repo}/releases" try: with httpx.Client( - proxy=proxy if proxy else None, follow_redirects=True + proxy=proxy if proxy else None, + follow_redirects=True, ) as client: resp = client.get(release_url) resp.raise_for_status() @@ -55,7 +56,8 @@ def get_git_repo(url: str, target_path: Path, proxy: str | None = None): # 下载并解压 with httpx.Client( - proxy=proxy if proxy else None, follow_redirects=True + proxy=proxy if proxy else None, + follow_redirects=True, ) as client: resp = client.get(download_url) if ( @@ -89,6 +91,7 @@ def load_yaml_metadata(plugin_dir: Path) -> dict: Returns: dict: 包含元数据的字典，如果读取失败则返回空字典 + """ yaml_path = plugin_dir / "metadata.yaml" if yaml_path.exists(): @@ -107,6 +110,7 @@ def build_plug_list(plugins_dir: Path) -> list: Returns: list: 包含插件信息的字典列表 + """ # 获取本地插件信息 result = [] @@ -133,7 +137,7 @@ def build_plug_list(plugins_dir: Path) -> list: "repo": str(metadata.get("repo", "")), "status": PluginStatus.INSTALLED, "local_path": str(plugin_dir), - } + }, ) # 获取在线插件列表 @@ -153,7 +157,7 @@ def build_plug_list(plugins_dir: Path) -> list: "repo": str(plugin_info.get("repo", "")), "status": PluginStatus.NOT_INSTALLED, "local_path": None, - } + }, ) except Exception as e: click.echo(f"获取在线插件列表失败: {e}", err=True) @@ -168,7 +172,8 @@ def build_plug_list(plugins_dir: Path) -> list: ) if ( VersionComparator.compare_version( - local_plugin["version"], online_plugin["version"] + local_plugin["version"], + online_plugin["version"], ) < 0 ): @@ -186,7 +191,10 @@ def build_plug_list(plugins_dir: Path) -> list: def manage_plugin( - plugin: dict, plugins_dir: Path, is_update: bool = False, proxy: str | None = None + plugin: dict, + plugins_dir: Path, + is_update: bool = False, + proxy: str | None = None, ) -> None: """安装或更新插件 @@ -195,6 +203,7 @@ def manage_plugin( plugins_dir (Path): 插件目录 is_update (bool, optional): 是否为更新操作. 默认为 False proxy (str, optional): 代理服务器地址 + """ plugin_name = plugin["name"] repo_url = plugin["repo"] @@ -212,26 +221,26 @@ def manage_plugin( raise click.ClickException(f"插件 {plugin_name} 未安装，无法更新") # 备份现有插件 - if is_update and backup_path.exists(): + if is_update and backup_path is not None and backup_path.exists(): shutil.rmtree(backup_path) - if is_update: + if is_update and backup_path is not None: shutil.copytree(target_path, backup_path) try: click.echo( - f"正在从 {repo_url} {'更新' if is_update else '下载'}插件 {plugin_name}..." + f"正在从 {repo_url} {'更新' if is_update else '下载'}插件 {plugin_name}...", ) get_git_repo(repo_url, target_path, proxy) # 更新成功，删除备份 - if is_update and backup_path.exists(): + if is_update and backup_path is not None and backup_path.exists(): shutil.rmtree(backup_path) click.echo(f"插件 {plugin_name} {'更新' if is_update else '安装'}成功") except Exception as e: if target_path.exists(): shutil.rmtree(target_path, ignore_errors=True) - if is_update and backup_path.exists(): + if is_update and backup_path is not None and backup_path.exists(): shutil.move(backup_path, target_path) raise click.ClickException( - f"{'更新' if is_update else '安装'}插件 {plugin_name} 时出错: {e}" + f"{'更新' if is_update else '安装'}插件 {plugin_name} 时出错: {e}", ) diff --git a/astrbot/cli/utils/version_comparator.py b/astrbot/cli/utils/version_comparator.py index fecab885..0aaf8dca 100644 --- a/astrbot/cli/utils/version_comparator.py +++ b/astrbot/cli/utils/version_comparator.py @@ -1,6 +1,4 @@ -""" -拷贝自 astrbot.core.utils.version_comparator -""" +"""拷贝自 astrbot.core.utils.version_comparator""" import re @@ -42,15 +40,15 @@ class VersionComparator: for i in range(length): if v1_parts[i] > v2_parts[i]: return 1 - elif v1_parts[i] < v2_parts[i]: + if v1_parts[i] < v2_parts[i]: return -1 # 比较预发布标签 if v1_prerelease is None and v2_prerelease is not None: return 1 # 没有预发布标签的版本高于有预发布标签的版本 - elif v1_prerelease is not None and v2_prerelease is None: + if v1_prerelease is not None and v2_prerelease is None: return -1 # 有预发布标签的版本低于没有预发布标签的版本 - elif v1_prerelease is not None and v2_prerelease is not None: + if v1_prerelease is not None and v2_prerelease is not None: len_pre = max(len(v1_prerelease), len(v2_prerelease)) for i in range(len_pre): p1 = v1_prerelease[i] if i < len(v1_prerelease) else None @@ -58,21 +56,21 @@ class VersionComparator: if p1 is None and p2 is not None: return -1 - elif p1 is not None and p2 is None: + if p1 is not None and p2 is None: return 1 - elif isinstance(p1, int) and isinstance(p2, str): + if isinstance(p1, int) and isinstance(p2, str): return -1 - elif isinstance(p1, str) and isinstance(p2, int): + if isinstance(p1, str) and isinstance(p2, int): return 1 - elif isinstance(p1, int) and isinstance(p2, int): + if isinstance(p1, int) and isinstance(p2, int): if p1 > p2: return 1 - elif p1 < p2: + if p1 < p2: return -1 elif isinstance(p1, str) and isinstance(p2, str): if p1 > p2: return 1 - elif p1 < p2: + if p1 < p2: return -1 return 0 # 预发布标签完全相同 diff --git a/astrbot/core/__init__.py b/astrbot/core/__init__.py index 235a8284..30b81af6 100644 --- a/astrbot/core/__init__.py +++ b/astrbot/core/__init__.py @@ -1,12 +1,14 @@ import os -from .log import LogManager, LogBroker # noqa -from astrbot.core.utils.t2i.renderer import HtmlRenderer -from astrbot.core.utils.shared_preferences import SharedPreferences -from astrbot.core.utils.pip_installer import PipInstaller -from astrbot.core.db.sqlite import SQLiteDatabase -from astrbot.core.config.default import DB_PATH + from astrbot.core.config import AstrBotConfig +from astrbot.core.config.default import DB_PATH +from astrbot.core.db.sqlite import SQLiteDatabase from astrbot.core.file_token_service import FileTokenService +from astrbot.core.utils.pip_installer import PipInstaller +from astrbot.core.utils.shared_preferences import SharedPreferences +from astrbot.core.utils.t2i.renderer import HtmlRenderer + +from .log import LogBroker, LogManager # noqa from .utils.astrbot_path import get_astrbot_data_path # 初始化数据存储文件夹 diff --git a/astrbot/core/agent/agent.py b/astrbot/core/agent/agent.py index 061ffde0..e2206829 100644 --- a/astrbot/core/agent/agent.py +++ b/astrbot/core/agent/agent.py @@ -1,8 +1,9 @@ from dataclasses import dataclass -from .tool import FunctionTool from typing import Generic -from .run_context import TContext + from .hooks import BaseAgentRunHooks +from .run_context import TContext +from .tool import FunctionTool @dataclass diff --git a/astrbot/core/agent/handoff.py b/astrbot/core/agent/handoff.py index d2646314..85276540 100644 --- a/astrbot/core/agent/handoff.py +++ b/astrbot/core/agent/handoff.py @@ -1,14 +1,18 @@ from typing import Generic -from .tool import FunctionTool + from .agent import Agent from .run_context import TContext +from .tool import FunctionTool class HandoffTool(FunctionTool, Generic[TContext]): """Handoff tool for delegating tasks to another agent.""" def __init__( - self, agent: Agent[TContext], parameters: dict | None = None, **kwargs + self, + agent: Agent[TContext], + parameters: dict | None = None, + **kwargs, ): self.agent = agent super().__init__( diff --git a/astrbot/core/agent/hooks.py b/astrbot/core/agent/hooks.py index 884fe6bd..d834240b 100644 --- a/astrbot/core/agent/hooks.py +++ b/astrbot/core/agent/hooks.py @@ -1,12 +1,13 @@ -import mcp -from dataclasses import dataclass -from .run_context import ContextWrapper, TContext from typing import Generic -from astrbot.core.provider.entities import LLMResponse + +import mcp + from astrbot.core.agent.tool import FunctionTool +from astrbot.core.provider.entities import LLMResponse + +from .run_context import ContextWrapper, TContext -@dataclass class BaseAgentRunHooks(Generic[TContext]): async def on_agent_begin(self, run_context: ContextWrapper[TContext]): ... async def on_tool_start( @@ -23,5 +24,7 @@ class BaseAgentRunHooks(Generic[TContext]): tool_result: mcp.types.CallToolResult | None, ): ... async def on_agent_done( - self, run_context: ContextWrapper[TContext], llm_response: LLMResponse + self, + run_context: ContextWrapper[TContext], + llm_response: LLMResponse, ): ... diff --git a/astrbot/core/agent/mcp_client.py b/astrbot/core/agent/mcp_client.py index 8db9d6f2..88cab486 100644 --- a/astrbot/core/agent/mcp_client.py +++ b/astrbot/core/agent/mcp_client.py @@ -1,28 +1,44 @@ import asyncio import logging -from datetime import timedelta -from typing import Optional from contextlib import AsyncExitStack +from datetime import timedelta +from typing import Generic + +from tenacity import ( + before_sleep_log, + retry, + retry_if_exception_type, + stop_after_attempt, + wait_exponential, +) + from astrbot import logger +from astrbot.core.agent.run_context import ContextWrapper from astrbot.core.utils.log_pipe import LogPipe +from .run_context import TContext +from .tool import FunctionTool + try: + import anyio import mcp from mcp.client.sse import sse_client except (ModuleNotFoundError, ImportError): - logger.warning("警告: 缺少依赖库 'mcp'，将无法使用 MCP 服务。") + logger.warning( + "Warning: Missing 'mcp' dependency, MCP services will be unavailable." + ) try: from mcp.client.streamable_http import streamablehttp_client except (ModuleNotFoundError, ImportError): logger.warning( - "警告: 缺少依赖库 'mcp' 或者 mcp 库版本过低，无法使用 Streamable HTTP 连接方式。" + "Warning: Missing 'mcp' dependency or MCP library version too old, Streamable HTTP connection unavailable.", ) def _prepare_config(config: dict) -> dict: - """准备配置，处理嵌套格式""" - if "mcpServers" in config and config["mcpServers"]: + """Prepare configuration, handle nested format""" + if config.get("mcpServers"): first_key = next(iter(config["mcpServers"])) config = config["mcpServers"][first_key] config.pop("active", None) @@ -30,7 +46,7 @@ def _prepare_config(config: dict) -> dict: async def _quick_test_mcp_connection(config: dict) -> tuple[bool, str]: - """快速测试 MCP 服务器可达性""" + """Quick test MCP server connectivity""" import aiohttp cfg = _prepare_config(config.copy()) @@ -45,7 +61,7 @@ async def _quick_test_mcp_connection(config: dict) -> tuple[bool, str]: elif "type" in cfg: transport_type = cfg["type"] else: - raise Exception("MCP 连接配置缺少 transport 或 type 字段") + raise Exception("MCP connection config missing transport or type field") async with aiohttp.ClientSession() as session: if transport_type == "streamable_http": @@ -71,8 +87,7 @@ async def _quick_test_mcp_connection(config: dict) -> tuple[bool, str]: ) as response: if response.status == 200: return True, "" - else: - return False, f"HTTP {response.status}: {response.reason}" + return False, f"HTTP {response.status}: {response.reason}" else: async with session.get( url, @@ -84,11 +99,10 @@ async def _quick_test_mcp_connection(config: dict) -> tuple[bool, str]: ) as response: if response.status == 200: return True, "" - else: - return False, f"HTTP {response.status}: {response.reason}" + return False, f"HTTP {response.status}: {response.reason}" except asyncio.TimeoutError: - return False, f"连接超时: {timeout}秒" + return False, f"Connection timeout: {timeout} seconds" except Exception as e: return False, f"{e!s}" @@ -96,8 +110,9 @@ async def _quick_test_mcp_connection(config: dict) -> tuple[bool, str]: class MCPClient: def __init__(self): # Initialize session and client objects - self.session: Optional[mcp.ClientSession] = None + self.session: mcp.ClientSession | None = None self.exit_stack = AsyncExitStack() + self._old_exit_stacks: list[AsyncExitStack] = [] # Track old stacks for cleanup self.name: str | None = None self.active: bool = True @@ -105,21 +120,32 @@ class MCPClient: self.server_errlogs: list[str] = [] self.running_event = asyncio.Event() - async def connect_to_server(self, mcp_server_config: dict, name: str): - """连接到 MCP 服务器 + # Store connection config for reconnection + self._mcp_server_config: dict | None = None + self._server_name: str | None = None + self._reconnect_lock = asyncio.Lock() # Lock for thread-safe reconnection + self._reconnecting: bool = False # For logging and debugging - 如果 `url` 参数存在： - 1. 当 transport 指定为 `streamable_http` 时，使用 Streamable HTTP 连接方式。 - 1. 当 transport 指定为 `sse` 时，使用 SSE 连接方式。 - 2. 如果没有指定，默认使用 SSE 的方式连接到 MCP 服务。 + async def connect_to_server(self, mcp_server_config: dict, name: str): + """Connect to MCP server + + If `url` parameter exists: + 1. When transport is specified as `streamable_http`, use Streamable HTTP connection. + 2. When transport is specified as `sse`, use SSE connection. + 3. If not specified, default to SSE connection to MCP service. Args: mcp_server_config (dict): Configuration for the MCP server. See https://modelcontextprotocol.io/quickstart/server + """ + # Store config for reconnection + self._mcp_server_config = mcp_server_config + self._server_name = name + cfg = _prepare_config(mcp_server_config.copy()) def logging_callback(msg: str): - # 处理 MCP 服务的错误日志 + # Handle MCP service error logs print(f"MCP Server {name} Error: {msg}") self.server_errlogs.append(msg) @@ -133,7 +159,7 @@ class MCPClient: elif "type" in cfg: transport_type = cfg["type"] else: - raise Exception("MCP 连接配置缺少 transport 或 type 字段") + raise Exception("MCP connection config missing transport or type field") if transport_type != "streamable_http": # SSE transport method @@ -144,7 +170,7 @@ class MCPClient: sse_read_timeout=cfg.get("sse_read_timeout", 60 * 5), ) streams = await self.exit_stack.enter_async_context( - self._streams_context + self._streams_context, ) # Create a new client session @@ -154,12 +180,12 @@ class MCPClient: *streams, read_timeout_seconds=read_timeout, logging_callback=logging_callback, # type: ignore - ) + ), ) else: timeout = timedelta(seconds=cfg.get("timeout", 30)) sse_read_timeout = timedelta( - seconds=cfg.get("sse_read_timeout", 60 * 5) + seconds=cfg.get("sse_read_timeout", 60 * 5), ) self._streams_context = streamablehttp_client( url=cfg["url"], @@ -169,7 +195,7 @@ class MCPClient: terminate_on_close=cfg.get("terminate_on_close", True), ) read_s, write_s, _ = await self.exit_stack.enter_async_context( - self._streams_context + self._streams_context, ) # Create a new client session @@ -180,7 +206,7 @@ class MCPClient: write_stream=write_s, read_timeout_seconds=read_timeout, logging_callback=logging_callback, # type: ignore - ) + ), ) else: @@ -189,7 +215,7 @@ class MCPClient: ) def callback(msg: str): - # 处理 MCP 服务的错误日志 + # Handle MCP service error logs self.server_errlogs.append(msg) stdio_transport = await self.exit_stack.enter_async_context( @@ -206,7 +232,7 @@ class MCPClient: # Create a new client session self.session = await self.exit_stack.enter_async_context( - mcp.ClientSession(*stdio_transport) + mcp.ClientSession(*stdio_transport), ) await self.session.initialize() @@ -218,7 +244,142 @@ class MCPClient: self.tools = response.tools return response + async def _reconnect(self) -> None: + """Reconnect to the MCP server using the stored configuration. + + Uses asyncio.Lock to ensure thread-safe reconnection in concurrent environments. + + Raises: + Exception: raised when reconnection fails + """ + async with self._reconnect_lock: + # Check if already reconnecting (useful for logging) + if self._reconnecting: + logger.debug( + f"MCP Client {self._server_name} is already reconnecting, skipping" + ) + return + + if not self._mcp_server_config or not self._server_name: + raise Exception("Cannot reconnect: missing connection configuration") + + self._reconnecting = True + try: + logger.info( + f"Attempting to reconnect to MCP server {self._server_name}..." + ) + + # Save old exit_stack for later cleanup (don't close it now to avoid cancel scope issues) + if self.exit_stack: + self._old_exit_stacks.append(self.exit_stack) + + # Mark old session as invalid + self.session = None + + # Create new exit stack for new connection + self.exit_stack = AsyncExitStack() + + # Reconnect using stored config + await self.connect_to_server(self._mcp_server_config, self._server_name) + await self.list_tools_and_save() + + logger.info( + f"Successfully reconnected to MCP server {self._server_name}" + ) + except Exception as e: + logger.error( + f"Failed to reconnect to MCP server {self._server_name}: {e}" + ) + raise + finally: + self._reconnecting = False + + async def call_tool_with_reconnect( + self, + tool_name: str, + arguments: dict, + read_timeout_seconds: timedelta, + ) -> mcp.types.CallToolResult: + """Call MCP tool with automatic reconnection on failure, max 2 retries. + + Args: + tool_name: tool name + arguments: tool arguments + read_timeout_seconds: read timeout + + Returns: + MCP tool call result + + Raises: + ValueError: MCP session is not available + anyio.ClosedResourceError: raised after reconnection failure + """ + + @retry( + retry=retry_if_exception_type(anyio.ClosedResourceError), + stop=stop_after_attempt(2), + wait=wait_exponential(multiplier=1, min=1, max=3), + before_sleep=before_sleep_log(logger, logging.WARNING), + reraise=True, + ) + async def _call_with_retry(): + if not self.session: + raise ValueError("MCP session is not available for MCP function tools.") + + try: + return await self.session.call_tool( + name=tool_name, + arguments=arguments, + read_timeout_seconds=read_timeout_seconds, + ) + except anyio.ClosedResourceError: + logger.warning( + f"MCP tool {tool_name} call failed (ClosedResourceError), attempting to reconnect..." + ) + # Attempt to reconnect + await self._reconnect() + # Reraise the exception to trigger tenacity retry + raise + + return await _call_with_retry() + async def cleanup(self): - """Clean up resources""" - await self.exit_stack.aclose() - self.running_event.set() # Set the running event to indicate cleanup is done + """Clean up resources including old exit stacks from reconnections""" + # Set running_event first to unblock any waiting tasks + self.running_event.set() + + # Close current exit stack + try: + await self.exit_stack.aclose() + except Exception as e: + logger.debug(f"Error closing current exit stack: {e}") + + # Don't close old exit stacks as they may be in different task contexts + # They will be garbage collected naturally + # Just clear the list to release references + self._old_exit_stacks.clear() + + +class MCPTool(FunctionTool, Generic[TContext]): + """A function tool that calls an MCP service.""" + + def __init__( + self, mcp_tool: mcp.Tool, mcp_client: MCPClient, mcp_server_name: str, **kwargs + ): + super().__init__( + name=mcp_tool.name, + description=mcp_tool.description or "", + parameters=mcp_tool.inputSchema, + ) + self.mcp_tool = mcp_tool + self.mcp_client = mcp_client + self.mcp_server_name = mcp_server_name + + async def call( + self, context: ContextWrapper[TContext], **kwargs + ) -> mcp.types.CallToolResult: + return await self.mcp_client.call_tool_with_reconnect( + tool_name=self.mcp_tool.name, + arguments=kwargs, + read_timeout_seconds=timedelta(seconds=context.tool_call_timeout), + ) diff --git a/astrbot/core/agent/message.py b/astrbot/core/agent/message.py new file mode 100644 index 00000000..4c65c32f --- /dev/null +++ b/astrbot/core/agent/message.py @@ -0,0 +1,175 @@ +# Inspired by MoonshotAI/kosong, credits to MoonshotAI/kosong authors for the original implementation. +# License: Apache License 2.0 + +from typing import Any, ClassVar, Literal, cast + +from pydantic import BaseModel, GetCoreSchemaHandler +from pydantic_core import core_schema + + +class ContentPart(BaseModel): + """A part of the content in a message.""" + + __content_part_registry: ClassVar[dict[str, type["ContentPart"]]] = {} + + type: str + + def __init_subclass__(cls, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + + invalid_subclass_error_msg = f"ContentPart subclass {cls.__name__} must have a `type` field of type `str`" + + type_value = getattr(cls, "type", None) + if type_value is None or not isinstance(type_value, str): + raise ValueError(invalid_subclass_error_msg) + + cls.__content_part_registry[type_value] = cls + + @classmethod + def __get_pydantic_core_schema__( + cls, source_type: Any, handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + # If we're dealing with the base ContentPart class, use custom validation + if cls.__name__ == "ContentPart": + + def validate_content_part(value: Any) -> Any: + # if it's already an instance of a ContentPart subclass, return it + if hasattr(value, "__class__") and issubclass(value.__class__, cls): + return value + + # if it's a dict with a type field, dispatch to the appropriate subclass + if isinstance(value, dict) and "type" in value: + type_value: Any | None = cast(dict[str, Any], value).get("type") + if not isinstance(type_value, str): + raise ValueError(f"Cannot validate {value} as ContentPart") + target_class = cls.__content_part_registry[type_value] + return target_class.model_validate(value) + + raise ValueError(f"Cannot validate {value} as ContentPart") + + return core_schema.no_info_plain_validator_function(validate_content_part) + + # for subclasses, use the default schema + return handler(source_type) + + +class TextPart(ContentPart): + """ + >>> TextPart(text="Hello, world!").model_dump() + {'type': 'text', 'text': 'Hello, world!'} + """ + + type: str = "text" + text: str + + +class ImageURLPart(ContentPart): + """ + >>> ImageURLPart(image_url="http://example.com/image.jpg").model_dump() + {'type': 'image_url', 'image_url': 'http://example.com/image.jpg'} + """ + + class ImageURL(BaseModel): + url: str + """The URL of the image, can be data URI scheme like `data:image/png;base64,...`.""" + id: str | None = None + """The ID of the image, to allow LLMs to distinguish different images.""" + + type: str = "image_url" + image_url: ImageURL + + +class AudioURLPart(ContentPart): + """ + >>> AudioURLPart(audio_url=AudioURLPart.AudioURL(url="https://example.com/audio.mp3")).model_dump() + {'type': 'audio_url', 'audio_url': {'url': 'https://example.com/audio.mp3', 'id': None}} + """ + + class AudioURL(BaseModel): + url: str + """The URL of the audio, can be data URI scheme like `data:audio/aac;base64,...`.""" + id: str | None = None + """The ID of the audio, to allow LLMs to distinguish different audios.""" + + type: str = "audio_url" + audio_url: AudioURL + + +class ToolCall(BaseModel): + """ + A tool call requested by the assistant. + + >>> ToolCall( + ... id="123", + ... function=ToolCall.FunctionBody( + ... name="function", + ... arguments="{}" + ... ), + ... ).model_dump() + {'type': 'function', 'id': '123', 'function': {'name': 'function', 'arguments': '{}'}} + """ + + class FunctionBody(BaseModel): + name: str + arguments: str | None + + type: Literal["function"] = "function" + + id: str + """The ID of the tool call.""" + function: FunctionBody + """The function body of the tool call.""" + extra_content: dict[str, Any] | None = None + """Extra metadata for the tool call.""" + + def model_dump(self, **kwargs: Any) -> dict[str, Any]: + if self.extra_content is None: + kwargs.setdefault("exclude", set()).add("extra_content") + return super().model_dump(**kwargs) + + +class ToolCallPart(BaseModel): + """A part of the tool call.""" + + arguments_part: str | None = None + """A part of the arguments of the tool call.""" + + +class Message(BaseModel): + """A message in a conversation.""" + + role: Literal[ + "system", + "user", + "assistant", + "tool", + ] + + content: str | list[ContentPart] + """The content of the message.""" + + +class AssistantMessageSegment(Message): + """A message segment from the assistant.""" + + role: Literal["assistant"] = "assistant" + tool_calls: list[ToolCall] | list[dict] | None = None + + +class ToolCallMessageSegment(Message): + """A message segment representing a tool call.""" + + role: Literal["tool"] = "tool" + tool_call_id: str + + +class UserMessageSegment(Message): + """A message segment from the user.""" + + role: Literal["user"] = "user" + + +class SystemMessageSegment(Message): + """A message segment from the system.""" + + role: Literal["system"] = "system" diff --git a/astrbot/core/agent/response.py b/astrbot/core/agent/response.py index 8eb1854f..3f3430c8 100644 --- a/astrbot/core/agent/response.py +++ b/astrbot/core/agent/response.py @@ -1,5 +1,6 @@ -from dataclasses import dataclass import typing as T +from dataclasses import dataclass + from astrbot.core.message.message_event_result import MessageChain diff --git a/astrbot/core/agent/run_context.py b/astrbot/core/agent/run_context.py index a0febf8c..07e43589 100644 --- a/astrbot/core/agent/run_context.py +++ b/astrbot/core/agent/run_context.py @@ -1,18 +1,22 @@ -from dataclasses import dataclass from typing import Any, Generic + +from pydantic import Field +from pydantic.dataclasses import dataclass from typing_extensions import TypeVar -from astrbot.core.platform.astr_message_event import AstrMessageEvent +from .message import Message TContext = TypeVar("TContext", default=Any) -@dataclass +@dataclass(config={"arbitrary_types_allowed": True}) class ContextWrapper(Generic[TContext]): """A context for running an agent, which can be used to pass additional data or state.""" context: TContext - event: AstrMessageEvent + messages: list[Message] = Field(default_factory=list) + """This field stores the llm message context for the agent run, agent runners will maintain this field automatically.""" + tool_call_timeout: int = 60 # Default tool call timeout in seconds NoContext = ContextWrapper[None] diff --git a/astrbot/core/agent/runners/base.py b/astrbot/core/agent/runners/base.py index 83821ae2..f7e0913b 100644 --- a/astrbot/core/agent/runners/base.py +++ b/astrbot/core/agent/runners/base.py @@ -1,13 +1,15 @@ import abc import typing as T from enum import Enum, auto -from ..run_context import ContextWrapper, TContext -from ..response import AgentResponse -from ..hooks import BaseAgentRunHooks -from ..tool_executor import BaseFunctionToolExecutor + from astrbot.core.provider import Provider from astrbot.core.provider.entities import LLMResponse +from ..hooks import BaseAgentRunHooks +from ..response import AgentResponse +from ..run_context import ContextWrapper, TContext +from ..tool_executor import BaseFunctionToolExecutor + class AgentState(Enum): """Defines the state of the agent.""" @@ -28,31 +30,33 @@ class BaseAgentRunner(T.Generic[TContext]): agent_hooks: BaseAgentRunHooks[TContext], **kwargs: T.Any, ) -> None: - """ - Reset the agent to its initial state. + """Reset the agent to its initial state. This method should be called before starting a new run. """ ... @abc.abstractmethod async def step(self) -> T.AsyncGenerator[AgentResponse, None]: - """ - Process a single step of the agent. - """ + """Process a single step of the agent.""" + ... + + @abc.abstractmethod + async def step_until_done( + self, max_step: int + ) -> T.AsyncGenerator[AgentResponse, None]: + """Process steps until the agent is done.""" ... @abc.abstractmethod def done(self) -> bool: - """ - Check if the agent has completed its task. + """Check if the agent has completed its task. Returns True if the agent is done, False otherwise. """ ... @abc.abstractmethod def get_final_llm_resp(self) -> LLMResponse | None: - """ - Get the final observation from the agent. + """Get the final observation from the agent. This method should be called after the agent is done. """ ... diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index 33298e89..d74a4598 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -1,31 +1,33 @@ import sys import traceback import typing as T -from .base import BaseAgentRunner, AgentResponse, AgentState -from ..hooks import BaseAgentRunHooks -from ..tool_executor import BaseFunctionToolExecutor -from ..run_context import ContextWrapper, TContext -from ..response import AgentResponseData -from astrbot.core.provider.provider import Provider + +from mcp.types import ( + BlobResourceContents, + CallToolResult, + EmbeddedResource, + ImageContent, + TextContent, + TextResourceContents, +) + +from astrbot import logger from astrbot.core.message.message_event_result import ( MessageChain, ) from astrbot.core.provider.entities import ( - ProviderRequest, LLMResponse, - ToolCallMessageSegment, - AssistantMessageSegment, + ProviderRequest, ToolCallsResult, ) -from mcp.types import ( - TextContent, - ImageContent, - EmbeddedResource, - TextResourceContents, - BlobResourceContents, - CallToolResult, -) -from astrbot import logger +from astrbot.core.provider.provider import Provider + +from ..hooks import BaseAgentRunHooks +from ..message import AssistantMessageSegment, Message, ToolCallMessageSegment +from ..response import AgentResponseData +from ..run_context import ContextWrapper, TContext +from ..tool_executor import BaseFunctionToolExecutor +from .base import AgentResponse, AgentState, BaseAgentRunner if sys.version_info >= (3, 12): from typing import override @@ -53,6 +55,20 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): self.agent_hooks = agent_hooks self.run_context = run_context + messages = [] + # append existing messages in the run context + for msg in request.contexts: + messages.append(Message.model_validate(msg)) + if request.prompt is not None: + m = await request.assemble_context() + messages.append(Message.model_validate(m)) + if request.system_prompt: + messages.insert( + 0, + Message(role="system", content=request.system_prompt), + ) + self.run_context.messages = messages + def _transition_state(self, new_state: AgentState) -> None: """转换 Agent 状态""" if self._state != new_state: @@ -70,8 +86,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): @override async def step(self): - """ - Process a single step of the agent. + """Process a single step of the agent. This method should return the result of the step. """ if not self.req: @@ -95,11 +110,20 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): type="streaming_delta", data=AgentResponseData(chain=llm_response.result_chain), ) - else: + elif llm_response.completion_text: yield AgentResponse( type="streaming_delta", data=AgentResponseData( - chain=MessageChain().message(llm_response.completion_text) + chain=MessageChain().message(llm_response.completion_text), + ), + ) + elif llm_response.reasoning_content: + yield AgentResponse( + type="streaming_delta", + data=AgentResponseData( + chain=MessageChain(type="reasoning").message( + llm_response.reasoning_content, + ), ), ) continue @@ -120,8 +144,8 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): type="err", data=AgentResponseData( chain=MessageChain().message( - f"LLM 响应错误: {llm_resp.completion_text or '未知错误'}" - ) + f"LLM 响应错误: {llm_resp.completion_text or '未知错误'}", + ), ), ) @@ -129,6 +153,13 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): # 如果没有工具调用，转换到完成状态 self.final_llm_resp = llm_resp self._transition_state(AgentState.DONE) + # record the final assistant message + self.run_context.messages.append( + Message( + role="assistant", + content=llm_resp.completion_text or "", + ), + ) try: await self.agent_hooks.on_agent_done(self.run_context, llm_resp) except Exception as e: @@ -144,7 +175,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): yield AgentResponse( type="llm_result", data=AgentResponseData( - chain=MessageChain().message(llm_resp.completion_text) + chain=MessageChain().message(llm_resp.completion_text), ), ) @@ -155,13 +186,16 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): yield AgentResponse( type="tool_call", data=AgentResponseData( - chain=MessageChain().message(f"🔨 调用工具: {tool_call_name}") + chain=MessageChain(type="tool_call").message( + f"🔨 调用工具: {tool_call_name}" + ), ), ) async for result in self._handle_function_tools(self.req, llm_resp): if isinstance(result, list): tool_call_result_blocks = result elif isinstance(result, MessageChain): + result.type = "tool_call_result" yield AgentResponse( type="tool_call_result", data=AgentResponseData(chain=result), @@ -169,14 +203,28 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): # 将结果添加到上下文中 tool_calls_result = ToolCallsResult( tool_calls_info=AssistantMessageSegment( - role="assistant", - tool_calls=llm_resp.to_openai_tool_calls(), + tool_calls=llm_resp.to_openai_to_calls_model(), content=llm_resp.completion_text, ), tool_calls_result=tool_call_result_blocks, ) + # record the assistant message with tool calls + self.run_context.messages.extend( + tool_calls_result.to_openai_messages_model() + ) + self.req.append_tool_calls_result(tool_calls_result) + async def step_until_done( + self, max_step: int + ) -> T.AsyncGenerator[AgentResponse, None]: + """Process steps until the agent is done.""" + step_count = 0 + while not self.done() and step_count < max_step: + step_count += 1 + async for resp in self.step(): + yield resp + async def _handle_function_tools( self, req: ProviderRequest, @@ -205,7 +253,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): role="tool", tool_call_id=func_tool_id, content=f"error: 未找到工具 {func_tool_name}", - ) + ), ) continue @@ -214,7 +262,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): # 获取实际的 handler 函数 if func_tool.handler: logger.debug( - f"工具 {func_tool_name} 期望的参数: {func_tool.parameters}" + f"工具 {func_tool_name} 期望的参数: {func_tool.parameters}", ) if func_tool.parameters and func_tool.parameters.get("properties"): expected_params = set(func_tool.parameters["properties"].keys()) @@ -227,20 +275,21 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): # 记录被忽略的参数 ignored_params = set(func_tool_args.keys()) - set( - valid_params.keys() + valid_params.keys(), ) if ignored_params: logger.warning( - f"工具 {func_tool_name} 忽略非期望参数: {ignored_params}" + f"工具 {func_tool_name} 忽略非期望参数: {ignored_params}", ) else: # 如果没有 handler（如 MCP 工具），使用所有参数 valid_params = func_tool_args - logger.warning(f"工具 {func_tool_name} 没有 handler，使用所有参数") try: await self.agent_hooks.on_tool_start( - self.run_context, func_tool, valid_params + self.run_context, + func_tool, + valid_params, ) except Exception as e: logger.error(f"Error in on_tool_start hook: {e}", exc_info=True) @@ -262,7 +311,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): role="tool", tool_call_id=func_tool_id, content=res.content[0].text, - ) + ), ) yield MessageChain().message(res.content[0].text) elif isinstance(res.content[0], ImageContent): @@ -271,10 +320,10 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): role="tool", tool_call_id=func_tool_id, content="返回了图片(已直接发送给用户)", - ) + ), ) yield MessageChain(type="tool_direct_result").base64_image( - res.content[0].data + res.content[0].data, ) elif isinstance(res.content[0], EmbeddedResource): resource = res.content[0].resource @@ -284,7 +333,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): role="tool", tool_call_id=func_tool_id, content=resource.text, - ) + ), ) yield MessageChain().message(resource.text) elif ( @@ -297,10 +346,10 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): role="tool", tool_call_id=func_tool_id, content="返回了图片(已直接发送给用户)", - ) + ), ) yield MessageChain( - type="tool_direct_result" + type="tool_direct_result", ).base64_image(resource.blob) else: tool_call_result_blocks.append( @@ -308,41 +357,41 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): role="tool", tool_call_id=func_tool_id, content="返回的数据类型不受支持", - ) + ), ) yield MessageChain().message("返回的数据类型不受支持。") elif resp is None: # Tool 直接请求发送消息给用户 # 这里我们将直接结束 Agent Loop。 + # 发送消息逻辑在 ToolExecutor 中处理了。 + logger.warning( + f"{func_tool_name} 没有没有返回值或者将结果直接发送给用户，此工具调用不会被记录到历史中。" + ) self._transition_state(AgentState.DONE) - if res := self.run_context.event.get_result(): - if res.chain: - yield MessageChain( - chain=res.chain, type="tool_direct_result" - ) else: # 不应该出现其他类型 logger.warning( - f"Tool 返回了不支持的类型: {type(resp)}，将忽略。" + f"Tool 返回了不支持的类型: {type(resp)}，将忽略。", ) try: await self.agent_hooks.on_tool_end( - self.run_context, func_tool, func_tool_args, _final_resp + self.run_context, + func_tool, + func_tool_args, + _final_resp, ) except Exception as e: logger.error(f"Error in on_tool_end hook: {e}", exc_info=True) - - self.run_context.event.clear_result() except Exception as e: logger.warning(traceback.format_exc()) tool_call_result_blocks.append( ToolCallMessageSegment( role="tool", tool_call_id=func_tool_id, - content=f"error: {str(e)}", - ) + content=f"error: {e!s}", + ), ) # 处理函数调用响应 diff --git a/astrbot/core/agent/tool.py b/astrbot/core/agent/tool.py index ae0ab761..b6008860 100644 --- a/astrbot/core/agent/tool.py +++ b/astrbot/core/agent/tool.py @@ -1,58 +1,77 @@ -from dataclasses import dataclass +from collections.abc import Awaitable, Callable +from typing import Any, Generic + +import jsonschema +import mcp from deprecated import deprecated -from typing import Awaitable, Callable, Literal, Any, Optional -from .mcp_client import MCPClient +from pydantic import Field, model_validator +from pydantic.dataclasses import dataclass + +from .run_context import ContextWrapper, TContext + +ParametersType = dict[str, Any] +ToolExecResult = str | mcp.types.CallToolResult @dataclass -class FunctionTool: - """A class representing a function tool that can be used in function calling.""" +class ToolSchema: + """A class representing the schema of a tool for function calling.""" name: str - parameters: dict | None = None - description: str | None = None - handler: Callable[..., Awaitable[Any]] | None = None - """处理函数, 当 origin 为 mcp 时，这个为空""" - handler_module_path: str | None = None - """处理函数的模块路径，当 origin 为 mcp 时，这个为空 + """The name of the tool.""" - 必须要保留这个字段, handler 在初始化会被 functools.partial 包装，导致 handler 的 __module__ 为 functools + description: str + """The description of the tool.""" + + parameters: ParametersType + """The parameters of the tool, in JSON Schema format.""" + + @model_validator(mode="after") + def validate_parameters(self) -> "ToolSchema": + jsonschema.validate( + self.parameters, jsonschema.Draft202012Validator.META_SCHEMA + ) + return self + + +@dataclass +class FunctionTool(ToolSchema, Generic[TContext]): + """A callable tool, for function calling.""" + + handler: Callable[..., Awaitable[Any]] | None = None + """a callable that implements the tool's functionality. It should be an async function.""" + + handler_module_path: str | None = None + """ + The module path of the handler function. This is empty when the origin is mcp. + This field must be retained, as the handler will be wrapped in functools.partial during initialization, + causing the handler's __module__ to be functools """ active: bool = True - """是否激活""" - - origin: Literal["local", "mcp"] = "local" - """函数工具的来源, local 为本地函数工具, mcp 为 MCP 服务""" - - # MCP 相关字段 - mcp_server_name: str | None = None - """MCP 服务名称，当 origin 为 mcp 时有效""" - mcp_client: MCPClient | None = None - """MCP 客户端，当 origin 为 mcp 时有效""" + """ + Whether the tool is active. This field is a special field for AstrBot. + You can ignore it when integrating with other frameworks. + """ def __repr__(self): - return f"FuncTool(name={self.name}, parameters={self.parameters}, description={self.description}, active={self.active}, origin={self.origin})" + return f"FuncTool(name={self.name}, parameters={self.parameters}, description={self.description})" - def __dict__(self) -> dict[str, Any]: - """将 FunctionTool 转换为字典格式""" - return { - "name": self.name, - "parameters": self.parameters, - "description": self.description, - "active": self.active, - "origin": self.origin, - "mcp_server_name": self.mcp_server_name, - } + async def call(self, context: ContextWrapper[TContext], **kwargs) -> ToolExecResult: + """Run the tool with the given arguments. The handler field has priority.""" + raise NotImplementedError( + "FunctionTool.call() must be implemented by subclasses or set a handler." + ) +@dataclass class ToolSet: """A set of function tools that can be used in function calling. This class provides methods to add, remove, and retrieve tools, as well as - convert the tools to different API formats (OpenAI, Anthropic, Google GenAI).""" + convert the tools to different API formats (OpenAI, Anthropic, Google GenAI). + """ - def __init__(self, tools: list[FunctionTool] | None = None): - self.tools: list[FunctionTool] = tools or [] + tools: list[FunctionTool] = Field(default_factory=list) def empty(self) -> bool: """Check if the tool set is empty.""" @@ -71,7 +90,7 @@ class ToolSet: """Remove a tool by its name.""" self.tools = [tool for tool in self.tools if tool.name != name] - def get_tool(self, name: str) -> Optional[FunctionTool]: + def get_tool(self, name: str) -> FunctionTool | None: """Get a tool by its name.""" for tool in self.tools: if tool.name == name: @@ -132,10 +151,8 @@ class ToolSet: } if ( - tool.parameters - and tool.parameters.get("properties") - or not omit_empty_parameter_field - ): + tool.parameters and tool.parameters.get("properties") + ) or not omit_empty_parameter_field: func_def["function"]["parameters"] = tool.parameters result.append(func_def) @@ -185,7 +202,8 @@ class ToolSet: if "type" in schema and schema["type"] in supported_types: result["type"] = schema["type"] if "format" in schema and schema["format"] in supported_formats.get( - result["type"], set() + result["type"], + set(), ): result["format"] = schema["format"] else: @@ -222,7 +240,7 @@ class ToolSet: tools = [] for tool in self.tools: - d = { + d: dict[str, Any] = { "name": tool.name, "description": tool.description, } diff --git a/astrbot/core/agent/tool_executor.py b/astrbot/core/agent/tool_executor.py index 34a2f5e7..2704119d 100644 --- a/astrbot/core/agent/tool_executor.py +++ b/astrbot/core/agent/tool_executor.py @@ -1,11 +1,17 @@ +from collections.abc import AsyncGenerator +from typing import Any, Generic + import mcp -from typing import Any, Generic, AsyncGenerator -from .run_context import TContext, ContextWrapper + +from .run_context import ContextWrapper, TContext from .tool import FunctionTool class BaseFunctionToolExecutor(Generic[TContext]): @classmethod async def execute( - cls, tool: FunctionTool, run_context: ContextWrapper[TContext], **tool_args + cls, + tool: FunctionTool, + run_context: ContextWrapper[TContext], + **tool_args, ) -> AsyncGenerator[Any | mcp.types.CallToolResult, None]: ... diff --git a/astrbot/core/astr_agent_context.py b/astrbot/core/astr_agent_context.py index 008c3a43..5eed5de8 100644 --- a/astrbot/core/astr_agent_context.py +++ b/astrbot/core/astr_agent_context.py @@ -1,12 +1,19 @@ -from dataclasses import dataclass -from astrbot.core.provider import Provider -from astrbot.core.provider.entities import ProviderRequest +from pydantic import Field +from pydantic.dataclasses import dataclass + +from astrbot.core.agent.run_context import ContextWrapper +from astrbot.core.platform.astr_message_event import AstrMessageEvent +from astrbot.core.star.context import Context -@dataclass +@dataclass(config={"arbitrary_types_allowed": True}) class AstrAgentContext: - provider: Provider - first_provider_request: ProviderRequest - curr_provider_request: ProviderRequest - streaming: bool - tool_call_timeout: int = 60 # Default tool call timeout in seconds + context: Context + """The star context instance""" + event: AstrMessageEvent + """The message event associated with the agent context.""" + extra: dict[str, str] = Field(default_factory=dict) + """Customized extra data.""" + + +AgentContextWrapper = ContextWrapper[AstrAgentContext] diff --git a/astrbot/core/astr_agent_hooks.py b/astrbot/core/astr_agent_hooks.py new file mode 100644 index 00000000..f394fc94 --- /dev/null +++ b/astrbot/core/astr_agent_hooks.py @@ -0,0 +1,36 @@ +from typing import Any + +from mcp.types import CallToolResult + +from astrbot.core.agent.hooks import BaseAgentRunHooks +from astrbot.core.agent.run_context import ContextWrapper +from astrbot.core.agent.tool import FunctionTool +from astrbot.core.astr_agent_context import AstrAgentContext +from astrbot.core.pipeline.context_utils import call_event_hook +from astrbot.core.star.star_handler import EventType + + +class MainAgentHooks(BaseAgentRunHooks[AstrAgentContext]): + async def on_agent_done(self, run_context, llm_response): + # 执行事件钩子 + await call_event_hook( + run_context.context.event, + EventType.OnLLMResponseEvent, + llm_response, + ) + + async def on_tool_end( + self, + run_context: ContextWrapper[AstrAgentContext], + tool: FunctionTool[Any], + tool_args: dict | None, + tool_result: CallToolResult | None, + ): + run_context.context.event.clear_result() + + +class EmptyAgentHooks(BaseAgentRunHooks[AstrAgentContext]): + pass + + +MAIN_AGENT_HOOKS = MainAgentHooks() diff --git a/astrbot/core/astr_agent_run_util.py b/astrbot/core/astr_agent_run_util.py new file mode 100644 index 00000000..5deb5af4 --- /dev/null +++ b/astrbot/core/astr_agent_run_util.py @@ -0,0 +1,80 @@ +import traceback +from collections.abc import AsyncGenerator + +from astrbot.core import logger +from astrbot.core.agent.runners.tool_loop_agent_runner import ToolLoopAgentRunner +from astrbot.core.astr_agent_context import AstrAgentContext +from astrbot.core.message.message_event_result import ( + MessageChain, + MessageEventResult, + ResultContentType, +) + +AgentRunner = ToolLoopAgentRunner[AstrAgentContext] + + +async def run_agent( + agent_runner: AgentRunner, + max_step: int = 30, + show_tool_use: bool = True, + stream_to_general: bool = False, + show_reasoning: bool = False, +) -> AsyncGenerator[MessageChain | None, None]: + step_idx = 0 + astr_event = agent_runner.run_context.context.event + while step_idx < max_step: + step_idx += 1 + try: + async for resp in agent_runner.step(): + if astr_event.is_stopped(): + return + if resp.type == "tool_call_result": + msg_chain = resp.data["chain"] + if msg_chain.type == "tool_direct_result": + # tool_direct_result 用于标记 llm tool 需要直接发送给用户的内容 + await astr_event.send(resp.data["chain"]) + continue + # 对于其他情况，暂时先不处理 + continue + elif resp.type == "tool_call": + if agent_runner.streaming: + # 用来标记流式响应需要分节 + yield MessageChain(chain=[], type="break") + if show_tool_use: + await astr_event.send(resp.data["chain"]) + continue + + if stream_to_general and resp.type == "streaming_delta": + continue + + if stream_to_general or not agent_runner.streaming: + content_typ = ( + ResultContentType.LLM_RESULT + if resp.type == "llm_result" + else ResultContentType.GENERAL_RESULT + ) + astr_event.set_result( + MessageEventResult( + chain=resp.data["chain"].chain, + result_content_type=content_typ, + ), + ) + yield + astr_event.clear_result() + elif resp.type == "streaming_delta": + chain = resp.data["chain"] + if chain.type == "reasoning" and not show_reasoning: + # display the reasoning content only when configured + continue + yield resp.data["chain"] # MessageChain + if agent_runner.done(): + break + + except Exception as e: + logger.error(traceback.format_exc()) + err_msg = f"\n\nAstrBot 请求失败。\n错误类型: {type(e).__name__}\n错误信息: {e!s}\n\n请在控制台查看和分享错误详情。\n" + if agent_runner.streaming: + yield MessageChain().message(err_msg) + else: + astr_event.set_result(MessageEventResult().message(err_msg)) + return diff --git a/astrbot/core/astr_agent_tool_exec.py b/astrbot/core/astr_agent_tool_exec.py new file mode 100644 index 00000000..440dea2d --- /dev/null +++ b/astrbot/core/astr_agent_tool_exec.py @@ -0,0 +1,246 @@ +import asyncio +import inspect +import traceback +import typing as T + +import mcp + +from astrbot import logger +from astrbot.core.agent.handoff import HandoffTool +from astrbot.core.agent.mcp_client import MCPTool +from astrbot.core.agent.run_context import ContextWrapper +from astrbot.core.agent.tool import FunctionTool, ToolSet +from astrbot.core.agent.tool_executor import BaseFunctionToolExecutor +from astrbot.core.astr_agent_context import AstrAgentContext +from astrbot.core.message.message_event_result import ( + CommandResult, + MessageChain, + MessageEventResult, +) +from astrbot.core.provider.register import llm_tools + + +class FunctionToolExecutor(BaseFunctionToolExecutor[AstrAgentContext]): + @classmethod + async def execute(cls, tool, run_context, **tool_args): + """执行函数调用。 + + Args: + event (AstrMessageEvent): 事件对象, 当 origin 为 local 时必须提供。 + **kwargs: 函数调用的参数。 + + Returns: + AsyncGenerator[None | mcp.types.CallToolResult, None] + + """ + if isinstance(tool, HandoffTool): + async for r in cls._execute_handoff(tool, run_context, **tool_args): + yield r + return + + elif isinstance(tool, MCPTool): + async for r in cls._execute_mcp(tool, run_context, **tool_args): + yield r + return + + else: + async for r in cls._execute_local(tool, run_context, **tool_args): + yield r + return + + @classmethod + async def _execute_handoff( + cls, + tool: HandoffTool, + run_context: ContextWrapper[AstrAgentContext], + **tool_args, + ): + input_ = tool_args.get("input") + + # make toolset for the agent + tools = tool.agent.tools + if tools: + toolset = ToolSet() + for t in tools: + if isinstance(t, str): + _t = llm_tools.get_func(t) + if _t: + toolset.add_tool(_t) + elif isinstance(t, FunctionTool): + toolset.add_tool(t) + else: + toolset = None + + ctx = run_context.context.context + event = run_context.context.event + umo = event.unified_msg_origin + prov_id = await ctx.get_current_chat_provider_id(umo) + llm_resp = await ctx.tool_loop_agent( + event=event, + chat_provider_id=prov_id, + prompt=input_, + system_prompt=tool.agent.instructions, + tools=toolset, + max_steps=30, + run_hooks=tool.agent.run_hooks, + ) + yield mcp.types.CallToolResult( + content=[mcp.types.TextContent(type="text", text=llm_resp.completion_text)] + ) + + @classmethod + async def _execute_local( + cls, + tool: FunctionTool, + run_context: ContextWrapper[AstrAgentContext], + **tool_args, + ): + event = run_context.context.event + if not event: + raise ValueError("Event must be provided for local function tools.") + + is_override_call = False + for ty in type(tool).mro(): + if "call" in ty.__dict__ and ty.__dict__["call"] is not FunctionTool.call: + is_override_call = True + break + + # 检查 tool 下有没有 run 方法 + if not tool.handler and not hasattr(tool, "run") and not is_override_call: + raise ValueError("Tool must have a valid handler or override 'run' method.") + + awaitable = None + method_name = "" + if tool.handler: + awaitable = tool.handler + method_name = "decorator_handler" + elif is_override_call: + awaitable = tool.call + method_name = "call" + elif hasattr(tool, "run"): + awaitable = getattr(tool, "run") + method_name = "run" + if awaitable is None: + raise ValueError("Tool must have a valid handler or override 'run' method.") + + wrapper = call_local_llm_tool( + context=run_context, + handler=awaitable, + method_name=method_name, + **tool_args, + ) + while True: + try: + resp = await asyncio.wait_for( + anext(wrapper), + timeout=run_context.tool_call_timeout, + ) + if resp is not None: + if isinstance(resp, mcp.types.CallToolResult): + yield resp + else: + text_content = mcp.types.TextContent( + type="text", + text=str(resp), + ) + yield mcp.types.CallToolResult(content=[text_content]) + else: + # NOTE: Tool 在这里直接请求发送消息给用户 + # TODO: 是否需要判断 event.get_result() 是否为空? + # 如果为空,则说明没有发送消息给用户,并且返回值为空,将返回一个特殊的 TextContent,其内容如"工具没有返回内容" + if res := run_context.context.event.get_result(): + if res.chain: + try: + await event.send( + MessageChain( + chain=res.chain, + type="tool_direct_result", + ) + ) + except Exception as e: + logger.error( + f"Tool 直接发送消息失败: {e}", + exc_info=True, + ) + yield None + except asyncio.TimeoutError: + raise Exception( + f"tool {tool.name} execution timeout after {run_context.tool_call_timeout} seconds.", + ) + except StopAsyncIteration: + break + + @classmethod + async def _execute_mcp( + cls, + tool: FunctionTool, + run_context: ContextWrapper[AstrAgentContext], + **tool_args, + ): + res = await tool.call(run_context, **tool_args) + if not res: + return + yield res + + +async def call_local_llm_tool( + context: ContextWrapper[AstrAgentContext], + handler: T.Callable[..., T.Awaitable[T.Any]], + method_name: str, + *args, + **kwargs, +) -> T.AsyncGenerator[T.Any, None]: + """执行本地 LLM 工具的处理函数并处理其返回结果""" + ready_to_call = None # 一个协程或者异步生成器 + + trace_ = None + + event = context.context.event + + try: + if method_name == "run" or method_name == "decorator_handler": + ready_to_call = handler(event, *args, **kwargs) + elif method_name == "call": + ready_to_call = handler(context, *args, **kwargs) + else: + raise ValueError(f"未知的方法名: {method_name}") + except ValueError as e: + logger.error(f"调用本地 LLM 工具时出错: {e}", exc_info=True) + except TypeError: + logger.error("处理函数参数不匹配，请检查 handler 的定义。", exc_info=True) + except Exception as e: + trace_ = traceback.format_exc() + logger.error(f"调用本地 LLM 工具时出错: {e}\n{trace_}") + + if not ready_to_call: + return + + if inspect.isasyncgen(ready_to_call): + _has_yielded = False + try: + async for ret in ready_to_call: + # 这里逐步执行异步生成器, 对于每个 yield 返回的 ret, 执行下面的代码 + # 返回值只能是 MessageEventResult 或者 None（无返回值） + _has_yielded = True + if isinstance(ret, (MessageEventResult, CommandResult)): + # 如果返回值是 MessageEventResult, 设置结果并继续 + event.set_result(ret) + yield + else: + # 如果返回值是 None, 则不设置结果并继续 + # 继续执行后续阶段 + yield ret + if not _has_yielded: + # 如果这个异步生成器没有执行到 yield 分支 + yield + except Exception as e: + logger.error(f"Previous Error: {trace_}") + raise e + elif inspect.iscoroutine(ready_to_call): + # 如果只是一个协程, 直接执行 + ret = await ready_to_call + if isinstance(ret, (MessageEventResult, CommandResult)): + event.set_result(ret) + yield + else: + yield ret diff --git a/astrbot/core/astrbot_config_mgr.py b/astrbot/core/astrbot_config_mgr.py index 0ee3f4fe..3a1353ce 100644 --- a/astrbot/core/astrbot_config_mgr.py +++ b/astrbot/core/astrbot_config_mgr.py @@ -1,13 +1,14 @@ import os import uuid +from typing import TypedDict, TypeVar + from astrbot.core import AstrBotConfig, logger -from astrbot.core.utils.shared_preferences import SharedPreferences from astrbot.core.config.astrbot_config import ASTRBOT_CONFIG_PATH from astrbot.core.config.default import DEFAULT_CONFIG from astrbot.core.platform.message_session import MessageSession from astrbot.core.umop_config_router import UmopConfigRouter from astrbot.core.utils.astrbot_path import get_astrbot_config_path -from typing import TypeVar, TypedDict +from astrbot.core.utils.shared_preferences import SharedPreferences _VT = TypeVar("_VT") @@ -48,7 +49,10 @@ class AstrBotConfigManager: """获取所有的 abconf 数据""" if self.abconf_data is None: self.abconf_data = self.sp.get( - "abconf_mapping", {}, scope="global", scope_id="global" + "abconf_mapping", + {}, + scope="global", + scope_id="global", ) return self.abconf_data @@ -64,7 +68,7 @@ class AstrBotConfigManager: self.confs[uuid_] = conf else: logger.warning( - f"Config file {conf_path} for UUID {uuid_} does not exist, skipping." + f"Config file {conf_path} for UUID {uuid_} does not exist, skipping.", ) continue @@ -73,6 +77,7 @@ class AstrBotConfigManager: Returns: ConfInfo: 包含配置文件的 uuid, 路径和名称等信息, 是一个 dict 类型 + """ # uuid -> { "path": str, "name": str } abconf_data = self._get_abconf_data() @@ -103,7 +108,10 @@ class AstrBotConfigManager: ) -> None: """保存配置文件的映射关系""" abconf_data = self.sp.get( - "abconf_mapping", {}, scope="global", scope_id="global" + "abconf_mapping", + {}, + scope="global", + scope_id="global", ) random_word = abconf_name or uuid.uuid4().hex[:8] abconf_data[abconf_id] = { @@ -177,13 +185,17 @@ class AstrBotConfigManager: Raises: ValueError: 如果试图删除默认配置文件 + """ if conf_id == "default": raise ValueError("不能删除默认配置文件") # 从映射中移除 abconf_data = self.sp.get( - "abconf_mapping", {}, scope="global", scope_id="global" + "abconf_mapping", + {}, + scope="global", + scope_id="global", ) if conf_id not in abconf_data: logger.warning(f"配置文件 {conf_id} 不存在于映射中") @@ -191,7 +203,8 @@ class AstrBotConfigManager: # 获取配置文件路径 conf_path = os.path.join( - get_astrbot_config_path(), abconf_data[conf_id]["path"] + get_astrbot_config_path(), + abconf_data[conf_id]["path"], ) # 删除配置文件 @@ -224,12 +237,16 @@ class AstrBotConfigManager: Returns: bool: 更新是否成功 + """ if conf_id == "default": raise ValueError("不能更新默认配置文件的信息") abconf_data = self.sp.get( - "abconf_mapping", {}, scope="global", scope_id="global" + "abconf_mapping", + {}, + scope="global", + scope_id="global", ) if conf_id not in abconf_data: logger.warning(f"配置文件 {conf_id} 不存在于映射中") @@ -246,7 +263,10 @@ class AstrBotConfigManager: return True def g( - self, umo: str | None = None, key: str | None = None, default: _VT = None + self, + umo: str | None = None, + key: str | None = None, + default: _VT = None, ) -> _VT: """获取配置项。umo 为 None 时使用默认配置""" if umo is None: diff --git a/astrbot/core/config/__init__.py b/astrbot/core/config/__init__.py index e49ac88a..839aeef3 100644 --- a/astrbot/core/config/__init__.py +++ b/astrbot/core/config/__init__.py @@ -1,9 +1,9 @@ -from .default import DEFAULT_CONFIG, VERSION, DB_PATH from .astrbot_config import * +from .default import DB_PATH, DEFAULT_CONFIG, VERSION __all__ = [ + "DB_PATH", "DEFAULT_CONFIG", "VERSION", - "DB_PATH", "AstrBotConfig", ] diff --git a/astrbot/core/config/astrbot_config.py b/astrbot/core/config/astrbot_config.py index 5d1f6fbe..786d29c8 100644 --- a/astrbot/core/config/astrbot_config.py +++ b/astrbot/core/config/astrbot_config.py @@ -1,11 +1,12 @@ -import os +import enum import json import logging -import enum -from .default import DEFAULT_CONFIG, DEFAULT_VALUE_MAP -from typing import Dict +import os + from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from .default import DEFAULT_CONFIG, DEFAULT_VALUE_MAP + ASTRBOT_CONFIG_PATH = os.path.join(get_astrbot_data_path(), "cmd_config.json") logger = logging.getLogger("astrbot") @@ -27,7 +28,7 @@ class AstrBotConfig(dict): self, config_path: str = ASTRBOT_CONFIG_PATH, default_config: dict = DEFAULT_CONFIG, - schema: dict = None, + schema: dict | None = None, ): super().__init__() @@ -45,7 +46,7 @@ class AstrBotConfig(dict): json.dump(default_config, f, indent=4, ensure_ascii=False) object.__setattr__(self, "first_deploy", True) # 标记第一次部署 - with open(config_path, "r", encoding="utf-8-sig") as f: + with open(config_path, encoding="utf-8-sig") as f: conf_str = f.read() conf = json.loads(conf_str) @@ -65,7 +66,7 @@ class AstrBotConfig(dict): for k, v in schema.items(): if v["type"] not in DEFAULT_VALUE_MAP: raise TypeError( - f"不受支持的配置类型 {v['type']}。支持的类型有：{DEFAULT_VALUE_MAP.keys()}" + f"不受支持的配置类型 {v['type']}。支持的类型有：{DEFAULT_VALUE_MAP.keys()}", ) if "default" in v: default = v["default"] @@ -82,7 +83,7 @@ class AstrBotConfig(dict): return conf - def check_config_integrity(self, refer_conf: Dict, conf: Dict, path=""): + def check_config_integrity(self, refer_conf: dict, conf: dict, path=""): """检查配置完整性，如果有新的配置项或顺序不一致则返回 True""" has_new = False @@ -97,27 +98,28 @@ class AstrBotConfig(dict): logger.info(f"检查到配置项 {path_} 不存在，已插入默认值 {value}") new_conf[key] = value has_new = True - else: - if conf[key] is None: - # 配置项为 None，使用默认值 + elif conf[key] is None: + # 配置项为 None，使用默认值 + new_conf[key] = value + has_new = True + elif isinstance(value, dict): + # 递归检查子配置项 + if not isinstance(conf[key], dict): + # 类型不匹配，使用默认值 new_conf[key] = value has_new = True - elif isinstance(value, dict): - # 递归检查子配置项 - if not isinstance(conf[key], dict): - # 类型不匹配，使用默认值 - new_conf[key] = value - has_new = True - else: - # 递归检查并同步顺序 - child_has_new = self.check_config_integrity( - value, conf[key], path + "." + key if path else key - ) - new_conf[key] = conf[key] - has_new |= child_has_new else: - # 直接使用现有配置 + # 递归检查并同步顺序 + child_has_new = self.check_config_integrity( + value, + conf[key], + path + "." + key if path else key, + ) new_conf[key] = conf[key] + has_new |= child_has_new + else: + # 直接使用现有配置 + new_conf[key] = conf[key] # 检查是否存在参考配置中没有的配置项 for key in list(conf.keys()): @@ -140,7 +142,7 @@ class AstrBotConfig(dict): return has_new - def save_config(self, replace_config: Dict = None): + def save_config(self, replace_config: dict | None = None): """将配置写入文件 如果传入 replace_config，则将配置替换为 replace_config diff --git a/astrbot/core/config/default.py b/astrbot/core/config/default.py index de019040..cf365867 100644 --- a/astrbot/core/config/default.py +++ b/astrbot/core/config/default.py @@ -1,12 +1,10 @@ -""" -如需修改配置，请在 `data/cmd_config.json` 中修改或者在管理面板中可视化修改。 -""" +"""如需修改配置，请在 `data/cmd_config.json` 中修改或者在管理面板中可视化修改。""" import os from astrbot.core.utils.astrbot_path import get_astrbot_data_path -VERSION = "4.3.5" +VERSION = "4.6.1" DB_PATH = os.path.join(get_astrbot_data_path(), "data_v4.db") # 默认配置 @@ -70,10 +68,10 @@ DEFAULT_CONFIG = { "dequeue_context_length": 1, "streaming_response": False, "show_tool_use_status": False, - "streaming_segmented": False, "agent_runner_type": "local", "dify_runner_provider_id": "", "coze_runner_provider_id": "", + "unsupported_streaming_strategy": "realtime_segmenting", "max_agent_step": 30, "tool_call_timeout": 60, }, @@ -137,8 +135,12 @@ DEFAULT_CONFIG = { "persona": [], # deprecated "timezone": "Asia/Shanghai", "callback_api_base": "", - "default_kb_collection": "", # 默认知识库名称 + "default_kb_collection": "", # 默认知识库名称, 已经过时 "plugin_set": ["*"], # "*" 表示使用所有可用的插件, 空列表表示不使用任何插件 + "kb_names": [], # 默认知识库名称列表 + "kb_fusion_top_k": 20, # 知识库检索融合阶段返回结果数量 + "kb_final_top_k": 5, # 知识库检索最终返回结果数量 + "kb_agentic_mode": False, } @@ -324,6 +326,10 @@ CONFIG_METADATA_2 = { # "type": "string", # "options": ["fullscreen", "embedded"], # }, + "is_sandbox": { + "description": "沙箱模式", + "type": "bool", + }, "satori_api_base_url": { "description": "Satori API 终结点", "type": "string", @@ -738,6 +744,7 @@ CONFIG_METADATA_2 = { "api_base": "https://api.openai.com/v1", "timeout": 120, "model_config": {"model": "gpt-4o-mini", "temperature": 0.4}, + "custom_headers": {}, "custom_extra_body": {}, "modalities": ["text", "image", "tool_use"], "hint": "也兼容所有与 OpenAI API 兼容的服务。", @@ -753,6 +760,7 @@ CONFIG_METADATA_2 = { "api_base": "", "timeout": 120, "model_config": {"model": "gpt-4o-mini", "temperature": 0.4}, + "custom_headers": {}, "custom_extra_body": {}, "modalities": ["text", "image", "tool_use"], }, @@ -766,7 +774,9 @@ CONFIG_METADATA_2 = { "api_base": "https://api.x.ai/v1", "timeout": 120, "model_config": {"model": "grok-2-latest", "temperature": 0.4}, + "custom_headers": {}, "custom_extra_body": {}, + "xai_native_search": False, "modalities": ["text", "image", "tool_use"], }, "Anthropic": { @@ -796,6 +806,7 @@ CONFIG_METADATA_2 = { "key": ["ollama"], # ollama 的 key 默认是 ollama "api_base": "http://localhost:11434/v1", "model_config": {"model": "llama3.1-8b", "temperature": 0.4}, + "custom_headers": {}, "custom_extra_body": {}, "modalities": ["text", "image", "tool_use"], }, @@ -810,6 +821,7 @@ CONFIG_METADATA_2 = { "model_config": { "model": "llama-3.1-8b", }, + "custom_headers": {}, "custom_extra_body": {}, "modalities": ["text", "image", "tool_use"], }, @@ -826,6 +838,7 @@ CONFIG_METADATA_2 = { "model": "gemini-1.5-flash", "temperature": 0.4, }, + "custom_headers": {}, "custom_extra_body": {}, "modalities": ["text", "image", "tool_use"], }, @@ -867,6 +880,24 @@ CONFIG_METADATA_2 = { "api_base": "https://api.deepseek.com/v1", "timeout": 120, "model_config": {"model": "deepseek-chat", "temperature": 0.4}, + "custom_headers": {}, + "custom_extra_body": {}, + "modalities": ["text", "tool_use"], + }, + "Groq": { + "id": "groq_default", + "provider": "groq", + "type": "groq_chat_completion", + "provider_type": "chat_completion", + "enable": True, + "key": [], + "api_base": "https://api.groq.com/openai/v1", + "timeout": 120, + "model_config": { + "model": "openai/gpt-oss-20b", + "temperature": 0.4, + }, + "custom_headers": {}, "custom_extra_body": {}, "modalities": ["text", "tool_use"], }, @@ -880,6 +911,7 @@ CONFIG_METADATA_2 = { "api_base": "https://api.302.ai/v1", "timeout": 120, "model_config": {"model": "gpt-4.1-mini", "temperature": 0.4}, + "custom_headers": {}, "custom_extra_body": {}, "modalities": ["text", "image", "tool_use"], }, @@ -896,6 +928,7 @@ CONFIG_METADATA_2 = { "model": "deepseek-ai/DeepSeek-V3", "temperature": 0.4, }, + "custom_headers": {}, "custom_extra_body": {}, "modalities": ["text", "image", "tool_use"], }, @@ -912,6 +945,7 @@ CONFIG_METADATA_2 = { "model": "deepseek/deepseek-r1", "temperature": 0.4, }, + "custom_headers": {}, "custom_extra_body": {}, }, "小马算力": { @@ -927,6 +961,7 @@ CONFIG_METADATA_2 = { "model": "kimi-k2-instruct-0905", "temperature": 0.7, }, + "custom_headers": {}, "custom_extra_body": {}, }, "优云智算": { @@ -941,6 +976,7 @@ CONFIG_METADATA_2 = { "model_config": { "model": "moonshotai/Kimi-K2-Instruct", }, + "custom_headers": {}, "custom_extra_body": {}, "modalities": ["text", "image", "tool_use"], }, @@ -954,6 +990,7 @@ CONFIG_METADATA_2 = { "timeout": 120, "api_base": "https://api.moonshot.cn/v1", "model_config": {"model": "moonshot-v1-8k", "temperature": 0.4}, + "custom_headers": {}, "custom_extra_body": {}, "modalities": ["text", "image", "tool_use"], }, @@ -969,6 +1006,8 @@ CONFIG_METADATA_2 = { "model_config": { "model": "glm-4-flash", }, + "custom_headers": {}, + "custom_extra_body": {}, "modalities": ["text", "image", "tool_use"], }, "Dify": { @@ -1025,6 +1064,7 @@ CONFIG_METADATA_2 = { "timeout": 120, "api_base": "https://api-inference.modelscope.cn/v1", "model_config": {"model": "Qwen/Qwen3-32B", "temperature": 0.4}, + "custom_headers": {}, "custom_extra_body": {}, "modalities": ["text", "image", "tool_use"], }, @@ -1037,6 +1077,7 @@ CONFIG_METADATA_2 = { "key": [], "api_base": "https://api.fastgpt.in/api/v1", "timeout": 60, + "custom_headers": {}, "custom_extra_body": {}, }, "Whisper(API)": { @@ -1258,8 +1299,51 @@ CONFIG_METADATA_2 = { "rerank_model": "BAAI/bge-reranker-base", "timeout": 20, }, + "Xinference Rerank": { + "id": "xinference_rerank", + "type": "xinference_rerank", + "provider": "xinference", + "provider_type": "rerank", + "enable": True, + "rerank_api_key": "", + "rerank_api_base": "http://127.0.0.1:9997", + "rerank_model": "BAAI/bge-reranker-base", + "timeout": 20, + "launch_model_if_not_running": False, + }, + "阿里云百炼重排序": { + "id": "bailian_rerank", + "type": "bailian_rerank", + "provider": "bailian", + "provider_type": "rerank", + "enable": True, + "rerank_api_key": "", + "rerank_api_base": "https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank", + "rerank_model": "qwen3-rerank", + "timeout": 30, + "return_documents": False, + "instruct": "", + }, + "Xinference STT": { + "id": "xinference_stt", + "type": "xinference_stt", + "provider": "xinference", + "provider_type": "speech_to_text", + "enable": False, + "api_key": "", + "api_base": "http://127.0.0.1:9997", + "model": "whisper-large-v3", + "timeout": 180, + "launch_model_if_not_running": False, + }, }, "items": { + "xai_native_search": { + "description": "启用原生搜索功能", + "type": "bool", + "hint": "启用后，将通过 xAI 的 Chat Completions 原生 Live Search 进行联网检索（按需计费）。仅对 xAI 提供商生效。", + "condition": {"provider": "xai"}, + }, "rerank_api_base": { "description": "重排序模型 API Base URL", "type": "string", @@ -1274,6 +1358,21 @@ CONFIG_METADATA_2 = { "description": "重排序模型名称", "type": "string", }, + "return_documents": { + "description": "是否在排序结果中返回文档原文", + "type": "bool", + "hint": "默认值false，以减少网络传输开销。", + }, + "instruct": { + "description": "自定义排序任务类型说明", + "type": "string", + "hint": "仅在使用 qwen3-rerank 模型时生效。建议使用英文撰写。", + }, + "launch_model_if_not_running": { + "description": "模型未运行时自动启动", + "type": "bool", + "hint": "如果模型当前未在 Xinference 服务中运行，是否尝试自动启动它。在生产环境中建议关闭。", + }, "modalities": { "description": "模型能力", "type": "list", @@ -1283,6 +1382,12 @@ CONFIG_METADATA_2 = { "render_type": "checkbox", "hint": "模型支持的模态。如所填写的模型不支持图像，请取消勾选图像。", }, + "custom_headers": { + "description": "自定义添加请求头", + "type": "dict", + "items": {}, + "hint": "此处添加的键值对将被合并到 OpenAI SDK 的 default_headers 中，用于自定义 HTTP 请求头。值必须为字符串。", + }, "custom_extra_body": { "description": "自定义请求体参数", "type": "dict", @@ -1417,6 +1522,7 @@ CONFIG_METADATA_2 = { "description": "嵌入维度", "type": "int", "hint": "嵌入向量的维度。根据模型不同，可能需要调整，请参考具体模型的文档。此配置项请务必填写正确，否则将导致向量数据库无法正常工作。", + "_special": "get_embedding_dim", }, "embedding_model": { "description": "嵌入模型", @@ -1930,8 +2036,8 @@ CONFIG_METADATA_2 = { "show_tool_use_status": { "type": "bool", }, - "streaming_segmented": { - "type": "bool", + "unsupported_streaming_strategy": { + "type": "string", }, "agent_runner_type": { "type": "string", @@ -2070,6 +2176,10 @@ CONFIG_METADATA_2 = { "default_kb_collection": { "type": "string", }, + "kb_names": {"type": "list", "items": {"type": "string"}}, + "kb_fusion_top_k": {"type": "int", "default": 20}, + "kb_final_top_k": {"type": "int", "default": 5}, + "kb_agentic_mode": {"type": "bool"}, }, }, } @@ -2203,10 +2313,27 @@ CONFIG_METADATA_3 = { "description": "知识库", "type": "object", "items": { - "default_kb_collection": { - "description": "默认使用的知识库", - "type": "string", + "kb_names": { + "description": "知识库列表", + "type": "list", + "items": {"type": "string"}, "_special": "select_knowledgebase", + "hint": "支持多选", + }, + "kb_fusion_top_k": { + "description": "融合检索结果数", + "type": "int", + "hint": "多个知识库检索结果融合后的返回结果数量", + }, + "kb_final_top_k": { + "description": "最终返回结果数", + "type": "int", + "hint": "从知识库中检索到的结果数量，越大可能获得越多相关信息，但也可能引入噪音。建议根据实际需求调整", + }, + "kb_agentic_mode": { + "description": "Agentic 知识库检索", + "type": "bool", + "hint": "启用后，知识库检索将作为 LLM Tool，由模型自主决定何时调用知识库进行查询。需要模型支持函数调用能力。", }, }, "condition": { @@ -2305,9 +2432,15 @@ CONFIG_METADATA_3 = { "description": "流式输出", "type": "bool", }, - "provider_settings.streaming_segmented": { - "description": "不支持流式输出的平台采取分段输出", - "type": "bool", + "provider_settings.unsupported_streaming_strategy": { + "description": "不支持流式回复的平台", + "type": "string", + "options": ["realtime_segmenting", "turn_off"], + "hint": "选择在不支持流式回复的平台上的处理方式。实时分段回复会在系统接收流式响应检测到诸如标点符号等分段点时，立即发送当前已接收的内容", + "labels": ["实时分段回复", "关闭流式回复"], + "condition": { + "provider_settings.streaming_response": True, + }, }, "provider_settings.max_context_length": { "description": "最多携带对话轮数", @@ -2743,9 +2876,9 @@ CONFIG_METADATA_3_SYSTEM = { "items": {"type": "string"}, }, }, - } + }, }, - } + }, } diff --git a/astrbot/core/conversation_mgr.py b/astrbot/core/conversation_mgr.py index a6a2710f..287fe03c 100644 --- a/astrbot/core/conversation_mgr.py +++ b/astrbot/core/conversation_mgr.py @@ -1,13 +1,14 @@ -""" -AstrBot 会话-对话管理器, 维护两个本地存储, 其中一个是 json 格式的shared_preferences, 另外一个是数据库 +"""AstrBot 会话-对话管理器, 维护两个本地存储, 其中一个是 json 格式的shared_preferences, 另外一个是数据库. 在 AstrBot 中, 会话和对话是独立的, 会话用于标记对话窗口, 例如群聊"123456789"可以建立一个会话, 在一个会话中可以建立多个对话, 并且支持对话的切换和删除 """ import json +from collections.abc import Awaitable, Callable + from astrbot.core import sp -from typing import Dict, List +from astrbot.core.agent.message import AssistantMessageSegment, UserMessageSegment from astrbot.core.db import BaseDatabase from astrbot.core.db.po import Conversation, ConversationV2 @@ -16,10 +17,45 @@ class ConversationManager: """负责管理会话与 LLM 的对话，某个会话当前正在用哪个对话。""" def __init__(self, db_helper: BaseDatabase): - self.session_conversations: Dict[str, str] = {} + self.session_conversations: dict[str, str] = {} self.db = db_helper self.save_interval = 60 # 每 60 秒保存一次 + # 会话删除回调函数列表（用于级联清理，如知识库配置） + self._on_session_deleted_callbacks: list[Callable[[str], Awaitable[None]]] = [] + + def register_on_session_deleted( + self, + callback: Callable[[str], Awaitable[None]], + ) -> None: + """注册会话删除回调函数. + + 其他模块可以注册回调来响应会话删除事件，实现级联清理。 + 例如：知识库模块可以注册回调来清理会话的知识库配置。 + + Args: + callback: 回调函数，接收会话ID (unified_msg_origin) 作为参数 + + """ + self._on_session_deleted_callbacks.append(callback) + + async def _trigger_session_deleted(self, unified_msg_origin: str) -> None: + """触发会话删除回调. + + Args: + unified_msg_origin: 会话ID + + """ + for callback in self._on_session_deleted_callbacks: + try: + await callback(unified_msg_origin) + except Exception as e: + from astrbot.core import logger + + logger.error( + f"会话删除回调执行失败 (session: {unified_msg_origin}): {e}", + ) + def _convert_conv_from_v2_to_v1(self, conv_v2: ConversationV2) -> Conversation: """将 ConversationV2 对象转换为 Conversation 对象""" created_at = int(conv_v2.created_at.timestamp()) @@ -43,12 +79,13 @@ class ConversationManager: title: str | None = None, persona_id: str | None = None, ) -> str: - """新建对话，并将当前会话的对话转移到新对话 + """新建对话，并将当前会话的对话转移到新对话. Args: unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id Returns: conversation_id (str): 对话 ID, 是 uuid 格式的字符串 + """ if not platform_id: # 如果没有提供 platform_id，则从 unified_msg_origin 中解析 @@ -74,18 +111,22 @@ class ConversationManager: Args: unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id conversation_id (str): 对话 ID, 是 uuid 格式的字符串 + """ self.session_conversations[unified_msg_origin] = conversation_id await sp.session_put(unified_msg_origin, "sel_conv_id", conversation_id) async def delete_conversation( - self, unified_msg_origin: str, conversation_id: str | None = None + self, + unified_msg_origin: str, + conversation_id: str | None = None, ): """删除会话的对话，当 conversation_id 为 None 时删除会话当前的对话 Args: unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id conversation_id (str): 对话 ID, 是 uuid 格式的字符串 + """ if not conversation_id: conversation_id = self.session_conversations.get(unified_msg_origin) @@ -101,11 +142,15 @@ class ConversationManager: Args: unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id + """ await self.db.delete_conversations_by_user_id(user_id=unified_msg_origin) self.session_conversations.pop(unified_msg_origin, None) await sp.session_remove(unified_msg_origin, "sel_conv_id") + # 触发会话删除回调（级联清理） + await self._trigger_session_deleted(unified_msg_origin) + async def get_curr_conversation_id(self, unified_msg_origin: str) -> str | None: """获取会话当前的对话 ID @@ -113,6 +158,7 @@ class ConversationManager: unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id Returns: conversation_id (str): 对话 ID, 是 uuid 格式的字符串 + """ ret = self.session_conversations.get(unified_msg_origin, None) if not ret: @@ -127,13 +173,15 @@ class ConversationManager: conversation_id: str, create_if_not_exists: bool = False, ) -> Conversation | None: - """获取会话的对话 + """获取会话的对话. Args: unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id conversation_id (str): 对话 ID, 是 uuid 格式的字符串 + create_if_not_exists (bool): 如果对话不存在,是否创建一个新的对话 Returns: conversation (Conversation): 对话对象 + """ conv = await self.db.get_conversation_by_id(cid=conversation_id) if not conv and create_if_not_exists: @@ -146,18 +194,22 @@ class ConversationManager: return conv_res async def get_conversations( - self, unified_msg_origin: str | None = None, platform_id: str | None = None - ) -> List[Conversation]: - """获取对话列表 + self, + unified_msg_origin: str | None = None, + platform_id: str | None = None, + ) -> list[Conversation]: + """获取对话列表. Args: unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id，可选 platform_id (str): 平台 ID, 可选参数, 用于过滤对话 Returns: conversations (List[Conversation]): 对话对象列表 + """ convs = await self.db.get_conversations( - user_id=unified_msg_origin, platform_id=platform_id + user_id=unified_msg_origin, + platform_id=platform_id, ) convs_res = [] for conv in convs: @@ -173,7 +225,7 @@ class ConversationManager: search_query: str = "", **kwargs, ) -> tuple[list[Conversation], int]: - """获取过滤后的对话列表 + """获取过滤后的对话列表. Args: page (int): 页码, 默认为 1 @@ -182,6 +234,7 @@ class ConversationManager: search_query (str): 搜索查询字符串, 可选 Returns: conversations (list[Conversation]): 对话对象列表 + """ convs, cnt = await self.db.get_filtered_conversations( page=page, @@ -203,13 +256,14 @@ class ConversationManager: history: list[dict] | None = None, title: str | None = None, persona_id: str | None = None, - ): - """更新会话的对话 + ) -> None: + """更新会话的对话. Args: unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id conversation_id (str): 对话 ID, 是 uuid 格式的字符串 history (List[Dict]): 对话历史记录, 是一个字典列表, 每个字典包含 role 和 content 字段 + """ if not conversation_id: # 如果没有提供 conversation_id，则获取当前的 @@ -223,16 +277,20 @@ class ConversationManager: ) async def update_conversation_title( - self, unified_msg_origin: str, title: str, conversation_id: str | None = None - ): - """更新会话的对话标题 + self, + unified_msg_origin: str, + title: str, + conversation_id: str | None = None, + ) -> None: + """更新会话的对话标题. Args: unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id title (str): 对话标题 - + conversation_id (str): 对话 ID, 是 uuid 格式的字符串 Deprecated: Use `update_conversation` with `title` parameter instead. + """ await self.update_conversation( unified_msg_origin=unified_msg_origin, @@ -245,15 +303,16 @@ class ConversationManager: unified_msg_origin: str, persona_id: str, conversation_id: str | None = None, - ): - """更新会话的对话 Persona ID + ) -> None: + """更新会话的对话 Persona ID. Args: unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id persona_id (str): 对话 Persona ID - + conversation_id (str): 对话 ID, 是 uuid 格式的字符串 Deprecated: Use `update_conversation` with `persona_id` parameter instead. + """ await self.update_conversation( unified_msg_origin=unified_msg_origin, @@ -261,40 +320,85 @@ class ConversationManager: persona_id=persona_id, ) + async def add_message_pair( + self, + cid: str, + user_message: UserMessageSegment | dict, + assistant_message: AssistantMessageSegment | dict, + ) -> None: + """Add a user-assistant message pair to the conversation history. + + Args: + cid (str): Conversation ID + user_message (UserMessageSegment | dict): OpenAI-format user message object or dict + assistant_message (AssistantMessageSegment | dict): OpenAI-format assistant message object or dict + + Raises: + Exception: If the conversation with the given ID is not found + """ + conv = await self.db.get_conversation_by_id(cid=cid) + if not conv: + raise Exception(f"Conversation with id {cid} not found") + history = conv.content or [] + if isinstance(user_message, UserMessageSegment): + user_msg_dict = user_message.model_dump() + else: + user_msg_dict = user_message + if isinstance(assistant_message, AssistantMessageSegment): + assistant_msg_dict = assistant_message.model_dump() + else: + assistant_msg_dict = assistant_message + history.append(user_msg_dict) + history.append(assistant_msg_dict) + await self.db.update_conversation( + cid=cid, + content=history, + ) + async def get_human_readable_context( - self, unified_msg_origin, conversation_id, page=1, page_size=10 - ): - """获取人类可读的上下文 + self, + unified_msg_origin: str, + conversation_id: str, + page: int = 1, + page_size: int = 10, + ) -> tuple[list[str], int]: + """获取人类可读的上下文. Args: unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id conversation_id (str): 对话 ID, 是 uuid 格式的字符串 page (int): 页码 page_size (int): 每页大小 + """ conversation = await self.get_conversation(unified_msg_origin, conversation_id) + if not conversation: + return [], 0 history = json.loads(conversation.history) - contexts = [] - temp_contexts = [] + # contexts_groups 存放按顺序的段落（每个段落是一个 str 列表）， + # 之后会被展平成一个扁平的 str 列表返回。 + contexts_groups: list[list[str]] = [] + temp_contexts: list[str] = [] for record in history: if record["role"] == "user": temp_contexts.append(f"User: {record['content']}") elif record["role"] == "assistant": - if "content" in record and record["content"]: + if record.get("content"): temp_contexts.append(f"Assistant: {record['content']}") elif "tool_calls" in record: tool_calls_str = json.dumps( - record["tool_calls"], ensure_ascii=False + record["tool_calls"], + ensure_ascii=False, ) temp_contexts.append(f"Assistant: [函数调用] {tool_calls_str}") else: temp_contexts.append("Assistant: [未知的内容]") - contexts.insert(0, temp_contexts) + contexts_groups.insert(0, temp_contexts) temp_contexts = [] - # 展平 contexts 列表 - contexts = [item for sublist in contexts for item in sublist] + # 展平分组后的 contexts 列表为单层字符串列表 + contexts = [item for sublist in contexts_groups for item in sublist] # 计算分页 paged_contexts = contexts[(page - 1) * page_size : page * page_size] diff --git a/astrbot/core/core_lifecycle.py b/astrbot/core/core_lifecycle.py index 49963734..17fd5213 100644 --- a/astrbot/core/core_lifecycle.py +++ b/astrbot/core/core_lifecycle.py @@ -1,5 +1,5 @@ -""" -Astrbot 核心生命周期管理类, 负责管理 AstrBot 的启动、停止、重启等操作。 +"""Astrbot 核心生命周期管理类, 负责管理 AstrBot 的启动、停止、重启等操作. + 该类负责初始化各个组件, 包括 ProviderManager、PlatformManager、ConversationManager、PluginManager、PipelineScheduler、EventBus等。 该类还负责加载和执行插件, 以及处理事件总线的分发。 @@ -9,44 +9,45 @@ Astrbot 核心生命周期管理类, 负责管理 AstrBot 的启动、停止、 3. 执行启动完成事件钩子 """ -import traceback import asyncio -import time -import threading import os -from .event_bus import EventBus -from . import astrbot_config, html_renderer +import threading +import time +import traceback from asyncio import Queue -from typing import List -from astrbot.core.pipeline.scheduler import PipelineScheduler, PipelineContext -from astrbot.core.star import PluginManager -from astrbot.core.platform.manager import PlatformManager -from astrbot.core.star.context import Context -from astrbot.core.persona_mgr import PersonaManager -from astrbot.core.provider.manager import ProviderManager -from astrbot.core import LogBroker -from astrbot.core.db import BaseDatabase -from astrbot.core.db.migration.migra_45_to_46 import migrate_45_to_46 -from astrbot.core.updator import AstrBotUpdator -from astrbot.core import logger, sp + +from astrbot.core import LogBroker, logger, sp +from astrbot.core.astrbot_config_mgr import AstrBotConfigManager from astrbot.core.config.default import VERSION from astrbot.core.conversation_mgr import ConversationManager +from astrbot.core.db import BaseDatabase +from astrbot.core.db.migration.migra_45_to_46 import migrate_45_to_46 +from astrbot.core.db.migration.migra_webchat_session import migrate_webchat_session +from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager +from astrbot.core.persona_mgr import PersonaManager +from astrbot.core.pipeline.scheduler import PipelineContext, PipelineScheduler +from astrbot.core.platform.manager import PlatformManager from astrbot.core.platform_message_history_mgr import PlatformMessageHistoryManager +from astrbot.core.provider.manager import ProviderManager +from astrbot.core.star import PluginManager +from astrbot.core.star.context import Context +from astrbot.core.star.star_handler import EventType, star_handlers_registry, star_map from astrbot.core.umop_config_router import UmopConfigRouter -from astrbot.core.astrbot_config_mgr import AstrBotConfigManager -from astrbot.core.star.star_handler import star_handlers_registry, EventType -from astrbot.core.star.star_handler import star_map +from astrbot.core.updator import AstrBotUpdator + +from . import astrbot_config, html_renderer +from .event_bus import EventBus class AstrBotCoreLifecycle: - """ - AstrBot 核心生命周期管理类, 负责管理 AstrBot 的启动、停止、重启等操作。 + """AstrBot 核心生命周期管理类, 负责管理 AstrBot 的启动、停止、重启等操作. + 该类负责初始化各个组件, 包括 ProviderManager、PlatformManager、ConversationManager、PluginManager、PipelineScheduler、 EventBus 等。 该类还负责加载和执行插件, 以及处理事件总线的分发。 """ - def __init__(self, log_broker: LogBroker, db: BaseDatabase): + def __init__(self, log_broker: LogBroker, db: BaseDatabase) -> None: self.log_broker = log_broker # 初始化日志代理 self.astrbot_config = astrbot_config # 初始化配置 self.db = db # 初始化数据库 @@ -70,11 +71,11 @@ class AstrBotCoreLifecycle: del os.environ["no_proxy"] logger.debug("HTTP proxy cleared") - async def initialize(self): - """ - 初始化 AstrBot 核心生命周期管理类, 负责初始化各个组件, 包括 ProviderManager、PlatformManager、ConversationManager、PluginManager、PipelineScheduler、EventBus、AstrBotUpdator等。 - """ + async def initialize(self) -> None: + """初始化 AstrBot 核心生命周期管理类. + 负责初始化各个组件, 包括 ProviderManager、PlatformManager、ConversationManager、PluginManager、PipelineScheduler、EventBus、AstrBotUpdator等。 + """ # 初始化日志代理 logger.info("AstrBot v" + VERSION) if os.environ.get("TESTING", ""): @@ -91,7 +92,9 @@ class AstrBotCoreLifecycle: # 初始化 AstrBot 配置管理器 self.astrbot_config_mgr = AstrBotConfigManager( - default_config=self.astrbot_config, ucr=self.umop_config_router, sp=sp + default_config=self.astrbot_config, + ucr=self.umop_config_router, + sp=sp, ) # 4.5 to 4.6 migration for umop_config_router @@ -101,6 +104,13 @@ class AstrBotCoreLifecycle: logger.error(f"Migration from version 4.5 to 4.6 failed: {e!s}") logger.error(traceback.format_exc()) + # migration for webchat session + try: + await migrate_webchat_session(self.db) + except Exception as e: + logger.error(f"Migration for webchat session failed: {e!s}") + logger.error(traceback.format_exc()) + # 初始化事件队列 self.event_queue = Queue() @@ -110,7 +120,9 @@ class AstrBotCoreLifecycle: # 初始化供应商管理器 self.provider_manager = ProviderManager( - self.astrbot_config_mgr, self.db, self.persona_mgr + self.astrbot_config_mgr, + self.db, + self.persona_mgr, ) # 初始化平台管理器 @@ -122,6 +134,9 @@ class AstrBotCoreLifecycle: # 初始化平台消息历史管理器 self.platform_message_history_manager = PlatformMessageHistoryManager(self.db) + # 初始化知识库管理器 + self.kb_manager = KnowledgeBaseManager(self.provider_manager) + # 初始化提供给插件的上下文 self.star_context = Context( self.event_queue, @@ -133,6 +148,7 @@ class AstrBotCoreLifecycle: self.platform_message_history_manager, self.persona_mgr, self.astrbot_config_mgr, + self.kb_manager, ) # 初始化插件管理器 @@ -144,8 +160,9 @@ class AstrBotCoreLifecycle: # 根据配置实例化各个 Provider await self.provider_manager.initialize() - # 初始化消息事件流水线调度器 + await self.kb_manager.initialize() + # 初始化消息事件流水线调度器 self.pipeline_scheduler_mapping = await self.load_pipeline_scheduler() # 初始化更新器 @@ -153,14 +170,16 @@ class AstrBotCoreLifecycle: # 初始化事件总线 self.event_bus = EventBus( - self.event_queue, self.pipeline_scheduler_mapping, self.astrbot_config_mgr + self.event_queue, + self.pipeline_scheduler_mapping, + self.astrbot_config_mgr, ) # 记录启动时间 self.start_time = int(time.time()) # 初始化当前任务列表 - self.curr_tasks: List[asyncio.Task] = [] + self.curr_tasks: list[asyncio.Task] = [] # 根据配置实例化各个平台适配器 await self.platform_manager.initialize() @@ -168,13 +187,13 @@ class AstrBotCoreLifecycle: # 初始化关闭控制面板的事件 self.dashboard_shutdown_event = asyncio.Event() - def _load(self): - """加载事件总线和任务并初始化""" - + def _load(self) -> None: + """加载事件总线和任务并初始化.""" # 创建一个异步任务来执行事件总线的 dispatch() 方法 # dispatch是一个无限循环的协程, 从事件队列中获取事件并处理 event_bus_task = asyncio.create_task( - self.event_bus.dispatch(), name="event_bus" + self.event_bus.dispatch(), + name="event_bus", ) # 把插件中注册的所有协程函数注册到事件总线中并执行 @@ -185,16 +204,17 @@ class AstrBotCoreLifecycle: tasks_ = [event_bus_task, *extra_tasks] for task in tasks_: self.curr_tasks.append( - asyncio.create_task(self._task_wrapper(task), name=task.get_name()) + asyncio.create_task(self._task_wrapper(task), name=task.get_name()), ) self.start_time = int(time.time()) - async def _task_wrapper(self, task: asyncio.Task): - """异步任务包装器, 用于处理异步任务执行中出现的各种异常 + async def _task_wrapper(self, task: asyncio.Task) -> None: + """异步任务包装器, 用于处理异步任务执行中出现的各种异常. Args: task (asyncio.Task): 要执行的异步任务 + """ try: await task @@ -207,19 +227,22 @@ class AstrBotCoreLifecycle: logger.error(f"| {line}") logger.error("-------") - async def start(self): - """启动 AstrBot 核心生命周期管理类, 用load加载事件总线和任务并初始化, 执行启动完成事件钩子""" + async def start(self) -> None: + """启动 AstrBot 核心生命周期管理类. + + 用load加载事件总线和任务并初始化, 执行启动完成事件钩子 + """ self._load() logger.info("AstrBot 启动完成。") # 执行启动完成事件钩子 handlers = star_handlers_registry.get_handlers_by_event_type( - EventType.OnAstrBotLoadedEvent + EventType.OnAstrBotLoadedEvent, ) for handler in handlers: try: logger.info( - f"hook(on_astrbot_loaded) -> {star_map[handler.handler_module_path].name} - {handler.handler_name}" + f"hook(on_astrbot_loaded) -> {star_map[handler.handler_module_path].name} - {handler.handler_name}", ) await handler.handler() except BaseException: @@ -228,8 +251,8 @@ class AstrBotCoreLifecycle: # 同时运行curr_tasks中的所有任务 await asyncio.gather(*self.curr_tasks, return_exceptions=True) - async def stop(self): - """停止 AstrBot 核心生命周期管理类, 取消所有当前任务并终止各个管理器""" + async def stop(self) -> None: + """停止 AstrBot 核心生命周期管理类, 取消所有当前任务并终止各个管理器.""" # 请求停止所有正在运行的异步任务 for task in self.curr_tasks: task.cancel() @@ -240,11 +263,12 @@ class AstrBotCoreLifecycle: except Exception as e: logger.warning(traceback.format_exc()) logger.warning( - f"插件 {plugin.name} 未被正常终止 {e!s}, 可能会导致资源泄露等问题。" + f"插件 {plugin.name} 未被正常终止 {e!s}, 可能会导致资源泄露等问题。", ) await self.provider_manager.terminate() await self.platform_manager.terminate() + await self.kb_manager.terminate() self.dashboard_shutdown_event.set() # 再次遍历curr_tasks等待每个任务真正结束 @@ -256,16 +280,19 @@ class AstrBotCoreLifecycle: except Exception as e: logger.error(f"任务 {task.get_name()} 发生错误: {e}") - async def restart(self): + async def restart(self) -> None: """重启 AstrBot 核心生命周期管理类, 终止各个管理器并重新加载平台实例""" await self.provider_manager.terminate() await self.platform_manager.terminate() + await self.kb_manager.terminate() self.dashboard_shutdown_event.set() threading.Thread( - target=self.astrbot_updator._reboot, name="restart", daemon=True + target=self.astrbot_updator._reboot, + name="restart", + daemon=True, ).start() - def load_platform(self) -> List[asyncio.Task]: + def load_platform(self) -> list[asyncio.Task]: """加载平台实例并返回所有平台实例的异步任务列表""" tasks = [] platform_insts = self.platform_manager.get_insts() @@ -274,36 +301,38 @@ class AstrBotCoreLifecycle: asyncio.create_task( platform_inst.run(), name=f"{platform_inst.meta().id}({platform_inst.meta().name})", - ) + ), ) return tasks async def load_pipeline_scheduler(self) -> dict[str, PipelineScheduler]: - """加载消息事件流水线调度器 + """加载消息事件流水线调度器. Returns: dict[str, PipelineScheduler]: 平台 ID 到流水线调度器的映射 + """ mapping = {} for conf_id, ab_config in self.astrbot_config_mgr.confs.items(): scheduler = PipelineScheduler( - PipelineContext(ab_config, self.plugin_manager, conf_id) + PipelineContext(ab_config, self.plugin_manager, conf_id), ) await scheduler.initialize() mapping[conf_id] = scheduler return mapping - async def reload_pipeline_scheduler(self, conf_id: str): - """重新加载消息事件流水线调度器 + async def reload_pipeline_scheduler(self, conf_id: str) -> None: + """重新加载消息事件流水线调度器. Returns: dict[str, PipelineScheduler]: 平台 ID 到流水线调度器的映射 + """ ab_config = self.astrbot_config_mgr.confs.get(conf_id) if not ab_config: raise ValueError(f"配置文件 {conf_id} 不存在") scheduler = PipelineScheduler( - PipelineContext(ab_config, self.plugin_manager, conf_id) + PipelineContext(ab_config, self.plugin_manager, conf_id), ) await scheduler.initialize() self.pipeline_scheduler_mapping[conf_id] = scheduler diff --git a/astrbot/core/db/__init__.py b/astrbot/core/db/__init__.py index 0abd3ad4..2af0428d 100644 --- a/astrbot/core/db/__init__.py +++ b/astrbot/core/db/__init__.py @@ -1,27 +1,28 @@ import abc import datetime import typing as T -from deprecated import deprecated -from dataclasses import dataclass -from astrbot.core.db.po import ( - Stats, - PlatformStat, - ConversationV2, - PlatformMessageHistory, - Attachment, - Persona, - Preference, -) from contextlib import asynccontextmanager +from dataclasses import dataclass + +from deprecated import deprecated from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine from sqlalchemy.orm import sessionmaker +from astrbot.core.db.po import ( + Attachment, + ConversationV2, + Persona, + PlatformMessageHistory, + PlatformSession, + PlatformStat, + Preference, + Stats, +) + @dataclass class BaseDatabase(abc.ABC): - """ - 数据库基类 - """ + """数据库基类""" DATABASE_URL = "" @@ -32,12 +33,13 @@ class BaseDatabase(abc.ABC): future=True, ) self.AsyncSessionLocal = sessionmaker( - self.engine, class_=AsyncSession, expire_on_commit=False + self.engine, + class_=AsyncSession, + expire_on_commit=False, ) async def initialize(self): """初始化数据库连接""" - pass @asynccontextmanager async def get_db(self) -> T.AsyncGenerator[AsyncSession, None]: @@ -91,7 +93,9 @@ class BaseDatabase(abc.ABC): @abc.abstractmethod async def get_conversations( - self, user_id: str | None = None, platform_id: str | None = None + self, + user_id: str | None = None, + platform_id: str | None = None, ) -> list[ConversationV2]: """Get all conversations for a specific user and platform_id(optional). @@ -106,7 +110,9 @@ class BaseDatabase(abc.ABC): @abc.abstractmethod async def get_all_conversations( - self, page: int = 1, page_size: int = 20 + self, + page: int = 1, + page_size: int = 20, ) -> list[ConversationV2]: """Get all conversations with pagination.""" ... @@ -173,9 +179,12 @@ class BaseDatabase(abc.ABC): @abc.abstractmethod async def delete_platform_message_offset( - self, platform_id: str, user_id: str, offset_sec: int = 86400 + self, + platform_id: str, + user_id: str, + offset_sec: int = 86400, ) -> None: - """Delete platform message history records older than the specified offset.""" + """Delete platform message history records newer than the specified offset.""" ... @abc.abstractmethod @@ -243,7 +252,11 @@ class BaseDatabase(abc.ABC): @abc.abstractmethod async def insert_preference_or_update( - self, scope: str, scope_id: str, key: str, value: dict + self, + scope: str, + scope_id: str, + key: str, + value: dict, ) -> Preference: """Insert a new preference record.""" ... @@ -255,7 +268,10 @@ class BaseDatabase(abc.ABC): @abc.abstractmethod async def get_preferences( - self, scope: str, scope_id: str | None = None, key: str | None = None + self, + scope: str, + scope_id: str | None = None, + key: str | None = None, ) -> list[Preference]: """Get all preferences for a specific scope ID or key.""" ... @@ -298,3 +314,51 @@ class BaseDatabase(abc.ABC): ) -> tuple[list[dict], int]: """Get paginated session conversations with joined conversation and persona details, support search and platform filter.""" ... + + # ==== + # Platform Session Management + # ==== + + @abc.abstractmethod + async def create_platform_session( + self, + creator: str, + platform_id: str = "webchat", + session_id: str | None = None, + display_name: str | None = None, + is_group: int = 0, + ) -> PlatformSession: + """Create a new Platform session.""" + ... + + @abc.abstractmethod + async def get_platform_session_by_id( + self, session_id: str + ) -> PlatformSession | None: + """Get a Platform session by its ID.""" + ... + + @abc.abstractmethod + async def get_platform_sessions_by_creator( + self, + creator: str, + platform_id: str | None = None, + page: int = 1, + page_size: int = 20, + ) -> list[PlatformSession]: + """Get all Platform sessions for a specific creator (username) and optionally platform.""" + ... + + @abc.abstractmethod + async def update_platform_session( + self, + session_id: str, + display_name: str | None = None, + ) -> None: + """Update a Platform session's updated_at timestamp and optionally display_name.""" + ... + + @abc.abstractmethod + async def delete_platform_session(self, session_id: str) -> None: + """Delete a Platform session by its ID.""" + ... diff --git a/astrbot/core/db/migration/helper.py b/astrbot/core/db/migration/helper.py index 901cdc4e..d7bca306 100644 --- a/astrbot/core/db/migration/helper.py +++ b/astrbot/core/db/migration/helper.py @@ -1,27 +1,33 @@ import os -from astrbot.core.utils.astrbot_path import get_astrbot_data_path -from astrbot.core.db import BaseDatabase -from astrbot.core.config import AstrBotConfig + from astrbot.api import logger, sp +from astrbot.core.config import AstrBotConfig +from astrbot.core.db import BaseDatabase +from astrbot.core.utils.astrbot_path import get_astrbot_data_path + from .migra_3_to_4 import ( migration_conversation_table, - migration_platform_table, - migration_webchat_data, migration_persona_data, + migration_platform_table, migration_preferences, + migration_webchat_data, ) async def check_migration_needed_v4(db_helper: BaseDatabase) -> bool: - """ - 检查是否需要进行数据库迁移 + """检查是否需要进行数据库迁移 如果存在 data_v3.db 并且 preference 中没有 migration_done_v4，则需要进行迁移。 """ - data_v3_exists = os.path.exists(get_astrbot_data_path()) - if not data_v3_exists: + # 仅当 data 目录下存在旧版本数据（data_v3.db 文件）时才考虑迁移 + data_dir = get_astrbot_data_path() + data_v3_db = os.path.join(data_dir, "data_v3.db") + + if not os.path.exists(data_v3_db): return False migration_done = await db_helper.get_preference( - "global", "global", "migration_done_v4" + "global", + "global", + "migration_done_v4", ) if migration_done: return False @@ -32,9 +38,8 @@ async def do_migration_v4( db_helper: BaseDatabase, platform_id_map: dict[str, dict[str, str]], astrbot_config: AstrBotConfig, -): - """ - 执行数据库迁移 +) -> None: + """执行数据库迁移 迁移旧的 webchat_conversation 表到新的 conversation 表。 迁移旧的 platform 到新的 platform_stats 表。 """ diff --git a/astrbot/core/db/migration/migra_3_to_4.py b/astrbot/core/db/migration/migra_3_to_4.py index 4aa5082d..a75c60a1 100644 --- a/astrbot/core/db/migration/migra_3_to_4.py +++ b/astrbot/core/db/migration/migra_3_to_4.py @@ -1,15 +1,18 @@ -import json import datetime -from .. import BaseDatabase -from .sqlite_v3 import SQLiteDatabase as SQLiteV3DatabaseV3 -from .shared_preferences_v3 import sp as sp_v3 -from astrbot.core.config.default import DB_PATH +import json + +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession + from astrbot.api import logger, sp from astrbot.core.config import AstrBotConfig -from astrbot.core.platform.astr_message_event import MessageSesion -from sqlalchemy.ext.asyncio import AsyncSession +from astrbot.core.config.default import DB_PATH from astrbot.core.db.po import ConversationV2, PlatformMessageHistory -from sqlalchemy import text +from astrbot.core.platform.astr_message_event import MessageSesion + +from .. import BaseDatabase +from .shared_preferences_v3 import sp as sp_v3 +from .sqlite_v3 import SQLiteDatabase as SQLiteV3DatabaseV3 """ 1. 迁移旧的 webchat_conversation 表到新的 conversation 表。 @@ -18,7 +21,8 @@ from sqlalchemy import text def get_platform_id( - platform_id_map: dict[str, dict[str, str]], old_platform_name: str + platform_id_map: dict[str, dict[str, str]], + old_platform_name: str, ) -> str: return platform_id_map.get( old_platform_name, @@ -27,7 +31,8 @@ def get_platform_id( def get_platform_type( - platform_id_map: dict[str, dict[str, str]], old_platform_name: str + platform_id_map: dict[str, dict[str, str]], + old_platform_name: str, ) -> str: return platform_id_map.get( old_platform_name, @@ -36,13 +41,15 @@ def get_platform_type( async def migration_conversation_table( - db_helper: BaseDatabase, platform_id_map: dict[str, dict[str, str]] + db_helper: BaseDatabase, + platform_id_map: dict[str, dict[str, str]], ): db_helper_v3 = SQLiteV3DatabaseV3( - db_path=DB_PATH.replace("data_v4.db", "data_v3.db") + db_path=DB_PATH.replace("data_v4.db", "data_v3.db"), ) conversations, total_cnt = db_helper_v3.get_all_conversations( - page=1, page_size=10000000 + page=1, + page_size=10000000, ) logger.info(f"迁移 {total_cnt} 条旧的会话数据到新的表中...") @@ -61,13 +68,14 @@ async def migration_conversation_table( ) if not conv: logger.info( - f"未找到该条旧会话对应的具体数据: {conversation}, 跳过。" + f"未找到该条旧会话对应的具体数据: {conversation}, 跳过。", ) if ":" not in conv.user_id: continue session = MessageSesion.from_str(session_str=conv.user_id) platform_id = get_platform_id( - platform_id_map, session.platform_name + platform_id_map, + session.platform_name, ) session.platform_id = platform_id # 更新平台名称为新的 ID conv_v2 = ConversationV2( @@ -90,10 +98,11 @@ async def migration_conversation_table( async def migration_platform_table( - db_helper: BaseDatabase, platform_id_map: dict[str, dict[str, str]] + db_helper: BaseDatabase, + platform_id_map: dict[str, dict[str, str]], ): db_helper_v3 = SQLiteV3DatabaseV3( - db_path=DB_PATH.replace("data_v4.db", "data_v3.db") + db_path=DB_PATH.replace("data_v4.db", "data_v3.db"), ) secs_from_2023_4_10_to_now = ( datetime.datetime.now(datetime.timezone.utc) @@ -134,10 +143,12 @@ async def migration_platform_table( if cnt == 0: continue platform_id = get_platform_id( - platform_id_map, platform_stats_v3[idx].name + platform_id_map, + platform_stats_v3[idx].name, ) platform_type = get_platform_type( - platform_id_map, platform_stats_v3[idx].name + platform_id_map, + platform_stats_v3[idx].name, ) try: await dbsession.execute( @@ -149,7 +160,8 @@ async def migration_platform_table( """), { "timestamp": datetime.datetime.fromtimestamp( - bucket_end, tz=datetime.timezone.utc + bucket_end, + tz=datetime.timezone.utc, ), "platform_id": platform_id, "platform_type": platform_type, @@ -165,14 +177,16 @@ async def migration_platform_table( async def migration_webchat_data( - db_helper: BaseDatabase, platform_id_map: dict[str, dict[str, str]] + db_helper: BaseDatabase, + platform_id_map: dict[str, dict[str, str]], ): """迁移 WebChat 的历史记录到新的 PlatformMessageHistory 表中""" db_helper_v3 = SQLiteV3DatabaseV3( - db_path=DB_PATH.replace("data_v4.db", "data_v3.db") + db_path=DB_PATH.replace("data_v4.db", "data_v3.db"), ) conversations, total_cnt = db_helper_v3.get_all_conversations( - page=1, page_size=10000000 + page=1, + page_size=10000000, ) logger.info(f"迁移 {total_cnt} 条旧的 WebChat 会话数据到新的表中...") @@ -191,7 +205,7 @@ async def migration_webchat_data( ) if not conv: logger.info( - f"未找到该条旧会话对应的具体数据: {conversation}, 跳过。" + f"未找到该条旧会话对应的具体数据: {conversation}, 跳过。", ) if ":" in conv.user_id: continue @@ -218,10 +232,10 @@ async def migration_webchat_data( async def migration_persona_data( - db_helper: BaseDatabase, astrbot_config: AstrBotConfig + db_helper: BaseDatabase, + astrbot_config: AstrBotConfig, ): - """ - 迁移 Persona 数据到新的表中。 + """迁移 Persona 数据到新的表中。 旧的 Persona 数据存储在 preference 中，新的 Persona 数据存储在 persona 表中。 """ v3_persona_config: list[dict] = astrbot_config.get("persona", []) @@ -236,14 +250,15 @@ async def migration_persona_data( try: begin_dialogs = persona.get("begin_dialogs", []) mood_imitation_dialogs = persona.get("mood_imitation_dialogs", []) - mood_prompt = "" + parts = [] user_turn = True for mood_dialog in mood_imitation_dialogs: if user_turn: - mood_prompt += f"A: {mood_dialog}\n" + parts.append(f"A: {mood_dialog}\n") else: - mood_prompt += f"B: {mood_dialog}\n" + parts.append(f"B: {mood_dialog}\n") user_turn = not user_turn + mood_prompt = "".join(parts) system_prompt = persona.get("prompt", "") if mood_prompt: system_prompt += f"Here are few shots of dialogs, you need to imitate the tone of 'B' in the following dialogs to respond:\n {mood_prompt}" @@ -253,14 +268,15 @@ async def migration_persona_data( begin_dialogs=begin_dialogs, ) logger.info( - f"迁移 Persona {persona['name']}({persona_new.system_prompt[:30]}...) 到新表成功。" + f"迁移 Persona {persona['name']}({persona_new.system_prompt[:30]}...) 到新表成功。", ) except Exception as e: logger.error(f"解析 Persona 配置失败：{e}") async def migration_preferences( - db_helper: BaseDatabase, platform_id_map: dict[str, dict[str, str]] + db_helper: BaseDatabase, + platform_id_map: dict[str, dict[str, str]], ): # 1. global scope migration keys = [ @@ -329,10 +345,13 @@ async def migration_preferences( for provider_type, provider_id in perf.items(): await sp.put_async( - "umo", str(session), f"provider_perf_{provider_type}", provider_id + "umo", + str(session), + f"provider_perf_{provider_type}", + provider_id, ) logger.info( - f"迁移会话 {umo} 的提供商偏好到新表成功，平台 ID: {platform_id}" + f"迁移会话 {umo} 的提供商偏好到新表成功，平台 ID: {platform_id}", ) except Exception as e: logger.error(f"迁移会话 {umo} 的提供商偏好失败: {e}", exc_info=True) diff --git a/astrbot/core/db/migration/migra_45_to_46.py b/astrbot/core/db/migration/migra_45_to_46.py index 8a1dc5de..dc70026f 100644 --- a/astrbot/core/db/migration/migra_45_to_46.py +++ b/astrbot/core/db/migration/migra_45_to_46.py @@ -9,7 +9,7 @@ async def migrate_45_to_46(acm: AstrBotConfigManager, ucr: UmopConfigRouter): if not isinstance(abconf_data, dict): # should be unreachable logger.warning( - f"migrate_45_to_46: abconf_data is not a dict (type={type(abconf_data)}). Value: {abconf_data!r}" + f"migrate_45_to_46: abconf_data is not a dict (type={type(abconf_data)}). Value: {abconf_data!r}", ) return diff --git a/astrbot/core/db/migration/migra_webchat_session.py b/astrbot/core/db/migration/migra_webchat_session.py new file mode 100644 index 00000000..ff0b5ca6 --- /dev/null +++ b/astrbot/core/db/migration/migra_webchat_session.py @@ -0,0 +1,131 @@ +"""Migration script for WebChat sessions. + +This migration creates PlatformSession from existing platform_message_history records. + +Changes: +- Creates platform_sessions table +- Adds platform_id field (default: 'webchat') +- Adds display_name field +- Session_id format: {platform_id}_{uuid} +""" + +from sqlalchemy import func, select +from sqlmodel import col + +from astrbot.api import logger, sp +from astrbot.core.db import BaseDatabase +from astrbot.core.db.po import ConversationV2, PlatformMessageHistory, PlatformSession + + +async def migrate_webchat_session(db_helper: BaseDatabase): + """Create PlatformSession records from platform_message_history. + + This migration extracts all unique user_ids from platform_message_history + where platform_id='webchat' and creates corresponding PlatformSession records. + """ + # 检查是否已经完成迁移 + migration_done = await db_helper.get_preference( + "global", "global", "migration_done_webchat_session_1" + ) + if migration_done: + return + + logger.info("开始执行数据库迁移（WebChat 会话迁移）...") + + try: + async with db_helper.get_db() as session: + # 从 platform_message_history 创建 PlatformSession + query = ( + select( + col(PlatformMessageHistory.user_id), + col(PlatformMessageHistory.sender_name), + func.min(PlatformMessageHistory.created_at).label("earliest"), + func.max(PlatformMessageHistory.updated_at).label("latest"), + ) + .where(col(PlatformMessageHistory.platform_id) == "webchat") + .where(col(PlatformMessageHistory.sender_id) != "bot") + .group_by(col(PlatformMessageHistory.user_id)) + ) + + result = await session.execute(query) + webchat_users = result.all() + + if not webchat_users: + logger.info("没有找到需要迁移的 WebChat 数据") + await sp.put_async( + "global", "global", "migration_done_webchat_session_1", True + ) + return + + logger.info(f"找到 {len(webchat_users)} 个 WebChat 会话需要迁移") + + # 检查已存在的会话 + existing_query = select(col(PlatformSession.session_id)) + existing_result = await session.execute(existing_query) + existing_session_ids = {row[0] for row in existing_result.fetchall()} + + # 查询 Conversations 表中的 title，用于设置 display_name + # 对于每个 user_id，对应的 conversation user_id 格式为: webchat:FriendMessage:webchat!astrbot!{user_id} + user_ids_to_query = [ + f"webchat:FriendMessage:webchat!astrbot!{user_id}" + for user_id, _, _, _ in webchat_users + ] + conv_query = select( + col(ConversationV2.user_id), col(ConversationV2.title) + ).where(col(ConversationV2.user_id).in_(user_ids_to_query)) + conv_result = await session.execute(conv_query) + # 创建 user_id -> title 的映射字典 + title_map = { + user_id.replace("webchat:FriendMessage:webchat!astrbot!", ""): title + for user_id, title in conv_result.fetchall() + } + + # 批量创建 PlatformSession 记录 + sessions_to_add = [] + skipped_count = 0 + + for user_id, sender_name, created_at, updated_at in webchat_users: + # user_id 就是 webchat_conv_id (session_id) + session_id = user_id + + # sender_name 通常是 username，但可能为 None + creator = sender_name if sender_name else "guest" + + # 检查是否已经存在该会话 + if session_id in existing_session_ids: + logger.debug(f"会话 {session_id} 已存在，跳过") + skipped_count += 1 + continue + + # 从 Conversations 表中获取 display_name + display_name = title_map.get(user_id) + + # 创建新的 PlatformSession（保留原有的时间戳） + new_session = PlatformSession( + session_id=session_id, + platform_id="webchat", + creator=creator, + is_group=0, + created_at=created_at, + updated_at=updated_at, + display_name=display_name, + ) + sessions_to_add.append(new_session) + + # 批量插入 + if sessions_to_add: + session.add_all(sessions_to_add) + await session.commit() + + logger.info( + f"WebChat 会话迁移完成！成功迁移: {len(sessions_to_add)}, 跳过: {skipped_count}", + ) + else: + logger.info("没有新会话需要迁移") + + # 标记迁移完成 + await sp.put_async("global", "global", "migration_done_webchat_session_1", True) + + except Exception as e: + logger.error(f"迁移过程中发生错误: {e}", exc_info=True) + raise diff --git a/astrbot/core/db/migration/shared_preferences_v3.py b/astrbot/core/db/migration/shared_preferences_v3.py index 6a661bd3..3abcb1a6 100644 --- a/astrbot/core/db/migration/shared_preferences_v3.py +++ b/astrbot/core/db/migration/shared_preferences_v3.py @@ -1,6 +1,7 @@ import json import os from typing import TypeVar + from astrbot.core.utils.astrbot_path import get_astrbot_data_path _VT = TypeVar("_VT") @@ -16,7 +17,7 @@ class SharedPreferences: def _load_preferences(self): if os.path.exists(self.path): try: - with open(self.path, "r") as f: + with open(self.path) as f: return json.load(f) except json.JSONDecodeError: os.remove(self.path) diff --git a/astrbot/core/db/migration/sqlite_v3.py b/astrbot/core/db/migration/sqlite_v3.py index ad86c51f..a301028d 100644 --- a/astrbot/core/db/migration/sqlite_v3.py +++ b/astrbot/core/db/migration/sqlite_v3.py @@ -1,8 +1,9 @@ import sqlite3 import time -from astrbot.core.db.po import Platform, Stats -from typing import Tuple, List, Dict, Any from dataclasses import dataclass +from typing import Any + +from astrbot.core.db.po import Platform, Stats @dataclass @@ -94,7 +95,7 @@ class SQLiteDatabase: c.execute( """ PRAGMA table_info(webchat_conversation) - """ + """, ) res = c.fetchall() has_title = False @@ -108,14 +109,14 @@ class SQLiteDatabase: c.execute( """ ALTER TABLE webchat_conversation ADD COLUMN title TEXT; - """ + """, ) self.conn.commit() if not has_persona_id: c.execute( """ ALTER TABLE webchat_conversation ADD COLUMN persona_id TEXT; - """ + """, ) self.conn.commit() @@ -126,7 +127,7 @@ class SQLiteDatabase: conn.text_factory = str return conn - def _exec_sql(self, sql: str, params: Tuple = None): + def _exec_sql(self, sql: str, params: tuple = None): conn = self.conn try: c = self.conn.cursor() @@ -174,7 +175,7 @@ class SQLiteDatabase: """ SELECT * FROM platform """ - + where_clause + + where_clause, ) platform = [] @@ -194,7 +195,7 @@ class SQLiteDatabase: c.execute( """ SELECT SUM(count) FROM platform - """ + """, ) res = c.fetchone() c.close() @@ -214,7 +215,7 @@ class SQLiteDatabase: SELECT name, SUM(count), timestamp FROM platform """ + where_clause - + " GROUP BY name" + + " GROUP BY name", ) platform = [] @@ -242,7 +243,7 @@ class SQLiteDatabase: c.close() if not res: - return + return None return Conversation(*res) @@ -257,7 +258,7 @@ class SQLiteDatabase: (user_id, cid, history, updated_at, created_at), ) - def get_conversations(self, user_id: str) -> Tuple: + def get_conversations(self, user_id: str) -> tuple: try: c = self.conn.cursor() except sqlite3.ProgrammingError: @@ -280,7 +281,7 @@ class SQLiteDatabase: title = row[3] persona_id = row[4] conversations.append( - Conversation("", cid, "[]", created_at, updated_at, title, persona_id) + Conversation("", cid, "[]", created_at, updated_at, title, persona_id), ) return conversations @@ -319,8 +320,10 @@ class SQLiteDatabase: ) def get_all_conversations( - self, page: int = 1, page_size: int = 20 - ) -> Tuple[List[Dict[str, Any]], int]: + self, + page: int = 1, + page_size: int = 20, + ) -> tuple[list[dict[str, Any]], int]: """获取所有对话，支持分页，按更新时间降序排序""" try: c = self.conn.cursor() @@ -366,7 +369,7 @@ class SQLiteDatabase: "persona_id": persona_id or "", "created_at": created_at or 0, "updated_at": updated_at or 0, - } + }, ) return conversations, total_count @@ -381,12 +384,12 @@ class SQLiteDatabase: self, page: int = 1, page_size: int = 20, - platforms: List[str] = None, - message_types: List[str] = None, - search_query: str = None, - exclude_ids: List[str] = None, - exclude_platforms: List[str] = None, - ) -> Tuple[List[Dict[str, Any]], int]: + platforms: list[str] | None = None, + message_types: list[str] | None = None, + search_query: str | None = None, + exclude_ids: list[str] | None = None, + exclude_platforms: list[str] | None = None, + ) -> tuple[list[dict[str, Any]], int]: """获取筛选后的对话列表""" try: c = self.conn.cursor() @@ -422,7 +425,7 @@ class SQLiteDatabase: if search_query: search_query = search_query.encode("unicode_escape").decode("utf-8") where_clauses.append( - "(title LIKE ? OR user_id LIKE ? OR cid LIKE ? OR history LIKE ?)" + "(title LIKE ? OR user_id LIKE ? OR cid LIKE ? OR history LIKE ?)", ) search_param = f"%{search_query}%" params.extend([search_param, search_param, search_param, search_param]) @@ -482,7 +485,7 @@ class SQLiteDatabase: "persona_id": persona_id or "", "created_at": created_at or 0, "updated_at": updated_at or 0, - } + }, ) return conversations, total_count diff --git a/astrbot/core/db/po.py b/astrbot/core/db/po.py index 24a05f94..d6621d07 100644 --- a/astrbot/core/db/po.py +++ b/astrbot/core/db/po.py @@ -1,15 +1,9 @@ import uuid - -from datetime import datetime, timezone from dataclasses import dataclass, field -from sqlmodel import ( - SQLModel, - Text, - JSON, - UniqueConstraint, - Field, -) -from typing import Optional, TypedDict +from datetime import datetime, timezone +from typing import TypedDict + +from sqlmodel import JSON, Field, SQLModel, Text, UniqueConstraint class PlatformStat(SQLModel, table=True): @@ -18,7 +12,7 @@ class PlatformStat(SQLModel, table=True): Note: In astrbot v4, we moved `platform` table to here. """ - __tablename__ = "platform_stats" + __tablename__ = "platform_stats" # type: ignore id: int = Field(primary_key=True, sa_column_kwargs={"autoincrement": True}) timestamp: datetime = Field(nullable=False) @@ -37,10 +31,11 @@ class PlatformStat(SQLModel, table=True): class ConversationV2(SQLModel, table=True): - __tablename__ = "conversations" + __tablename__ = "conversations" # type: ignore inner_conversation_id: int = Field( - primary_key=True, sa_column_kwargs={"autoincrement": True} + primary_key=True, + sa_column_kwargs={"autoincrement": True}, ) conversation_id: str = Field( max_length=36, @@ -50,14 +45,14 @@ class ConversationV2(SQLModel, table=True): ) platform_id: str = Field(nullable=False) user_id: str = Field(nullable=False) - content: Optional[list] = Field(default=None, sa_type=JSON) + content: list | None = Field(default=None, sa_type=JSON) created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) updated_at: datetime = Field( default_factory=lambda: datetime.now(timezone.utc), sa_column_kwargs={"onupdate": datetime.now(timezone.utc)}, ) - title: Optional[str] = Field(default=None, max_length=255) - persona_id: Optional[str] = Field(default=None) + title: str | None = Field(default=None, max_length=255) + persona_id: str | None = Field(default=None) __table_args__ = ( UniqueConstraint( @@ -73,16 +68,18 @@ class Persona(SQLModel, table=True): It can be used to customize the behavior of LLMs. """ - __tablename__ = "personas" + __tablename__ = "personas" # type: ignore id: int | None = Field( - primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None + primary_key=True, + sa_column_kwargs={"autoincrement": True}, + default=None, ) persona_id: str = Field(max_length=255, nullable=False) system_prompt: str = Field(sa_type=Text, nullable=False) - begin_dialogs: Optional[list] = Field(default=None, sa_type=JSON) + begin_dialogs: list | None = Field(default=None, sa_type=JSON) """a list of strings, each representing a dialog to start with""" - tools: Optional[list] = Field(default=None, sa_type=JSON) + tools: list | None = Field(default=None, sa_type=JSON) """None means use ALL tools for default, empty list means no tools, otherwise a list of tool names.""" created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) updated_at: datetime = Field( @@ -101,10 +98,12 @@ class Persona(SQLModel, table=True): class Preference(SQLModel, table=True): """This class represents preferences for bots.""" - __tablename__ = "preferences" + __tablename__ = "preferences" # type: ignore id: int | None = Field( - default=None, primary_key=True, sa_column_kwargs={"autoincrement": True} + default=None, + primary_key=True, + sa_column_kwargs={"autoincrement": True}, ) scope: str = Field(nullable=False) """Scope of the preference, such as 'global', 'umo', 'plugin'.""" @@ -135,16 +134,18 @@ class PlatformMessageHistory(SQLModel, table=True): or platform-specific messages. """ - __tablename__ = "platform_message_history" + __tablename__ = "platform_message_history" # type: ignore id: int | None = Field( - primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None + primary_key=True, + sa_column_kwargs={"autoincrement": True}, + default=None, ) platform_id: str = Field(nullable=False) user_id: str = Field(nullable=False) # An id of group, user in platform - sender_id: Optional[str] = Field(default=None) # ID of the sender in the platform - sender_name: Optional[str] = Field( - default=None + sender_id: str | None = Field(default=None) # ID of the sender in the platform + sender_name: str | None = Field( + default=None, ) # Name of the sender in the platform content: dict = Field(sa_type=JSON, nullable=False) # a message chain list created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) @@ -154,16 +155,60 @@ class PlatformMessageHistory(SQLModel, table=True): ) +class PlatformSession(SQLModel, table=True): + """Platform session table for managing user sessions across different platforms. + + A session represents a chat window for a specific user on a specific platform. + Each session can have multiple conversations (对话) associated with it. + """ + + __tablename__ = "platform_sessions" # type: ignore + + inner_id: int | None = Field( + primary_key=True, + sa_column_kwargs={"autoincrement": True}, + default=None, + ) + session_id: str = Field( + max_length=100, + nullable=False, + unique=True, + default_factory=lambda: f"webchat_{uuid.uuid4()}", + ) + platform_id: str = Field(default="webchat", nullable=False) + """Platform identifier (e.g., 'webchat', 'qq', 'discord')""" + creator: str = Field(nullable=False) + """Username of the session creator""" + display_name: str | None = Field(default=None, max_length=255) + """Display name for the session""" + is_group: int = Field(default=0, nullable=False) + """0 for private chat, 1 for group chat (not implemented yet)""" + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + updated_at: datetime = Field( + default_factory=lambda: datetime.now(timezone.utc), + sa_column_kwargs={"onupdate": datetime.now(timezone.utc)}, + ) + + __table_args__ = ( + UniqueConstraint( + "session_id", + name="uix_platform_session_id", + ), + ) + + class Attachment(SQLModel, table=True): """This class represents attachments for messages in AstrBot. Attachments can be images, files, or other media types. """ - __tablename__ = "attachments" + __tablename__ = "attachments" # type: ignore inner_attachment_id: int | None = Field( - primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None + primary_key=True, + sa_column_kwargs={"autoincrement": True}, + default=None, ) attachment_id: str = Field( max_length=36, diff --git a/astrbot/core/db/sqlite.py b/astrbot/core/db/sqlite.py index f9faede1..19461861 100644 --- a/astrbot/core/db/sqlite.py +++ b/astrbot/core/db/sqlite.py @@ -1,22 +1,28 @@ import asyncio -import typing as T import threading -from datetime import datetime, timedelta +import typing as T +from datetime import datetime, timedelta, timezone + +from sqlalchemy.ext.asyncio import AsyncSession +from sqlmodel import col, delete, desc, func, or_, select, text, update + from astrbot.core.db import BaseDatabase from astrbot.core.db.po import ( - ConversationV2, - PlatformStat, - PlatformMessageHistory, Attachment, + ConversationV2, Persona, + PlatformMessageHistory, + PlatformSession, + PlatformStat, Preference, - Stats as DeprecatedStats, - Platform as DeprecatedPlatformStat, SQLModel, ) - -from sqlmodel import select, update, delete, text, func, or_, desc, col -from sqlalchemy.ext.asyncio import AsyncSession +from astrbot.core.db.po import ( + Platform as DeprecatedPlatformStat, +) +from astrbot.core.db.po import ( + Stats as DeprecatedStats, +) NOT_GIVEN = T.TypeVar("NOT_GIVEN") @@ -57,7 +63,9 @@ class SQLiteDatabase(BaseDatabase): async with session.begin(): if timestamp is None: timestamp = datetime.now().replace( - minute=0, second=0, microsecond=0 + minute=0, + second=0, + microsecond=0, ) current_hour = timestamp await session.execute( @@ -81,13 +89,13 @@ class SQLiteDatabase(BaseDatabase): session: AsyncSession result = await session.execute( select(func.count(col(PlatformStat.platform_id))).select_from( - PlatformStat - ) + PlatformStat, + ), ) count = result.scalar_one_or_none() return count if count is not None else 0 - async def get_platform_stats(self, offset_sec: int = 86400) -> T.List[PlatformStat]: + async def get_platform_stats(self, offset_sec: int = 86400) -> list[PlatformStat]: """Get platform statistics within the specified offset in seconds and group by platform_id.""" async with self.get_db() as session: session: AsyncSession @@ -138,7 +146,7 @@ class SQLiteDatabase(BaseDatabase): select(ConversationV2) .order_by(desc(ConversationV2.created_at)) .offset(offset) - .limit(page_size) + .limit(page_size), ) return result.scalars().all() @@ -157,7 +165,7 @@ class SQLiteDatabase(BaseDatabase): if platform_ids: base_query = base_query.where( - col(ConversationV2.platform_id).in_(platform_ids) + col(ConversationV2.platform_id).in_(platform_ids), ) if search_query: search_query = search_query.encode("unicode_escape").decode("utf-8") @@ -167,16 +175,16 @@ class SQLiteDatabase(BaseDatabase): col(ConversationV2.content).ilike(f"%{search_query}%"), col(ConversationV2.user_id).ilike(f"%{search_query}%"), col(ConversationV2.conversation_id).ilike(f"%{search_query}%"), - ) + ), ) if "message_types" in kwargs and len(kwargs["message_types"]) > 0: for msg_type in kwargs["message_types"]: base_query = base_query.where( - col(ConversationV2.user_id).ilike(f"%:{msg_type}:%") + col(ConversationV2.user_id).ilike(f"%:{msg_type}:%"), ) if "platforms" in kwargs and len(kwargs["platforms"]) > 0: base_query = base_query.where( - col(ConversationV2.platform_id).in_(kwargs["platforms"]) + col(ConversationV2.platform_id).in_(kwargs["platforms"]), ) # Get total count matching the filters @@ -233,7 +241,7 @@ class SQLiteDatabase(BaseDatabase): session: AsyncSession async with session.begin(): query = update(ConversationV2).where( - col(ConversationV2.conversation_id) == cid + col(ConversationV2.conversation_id) == cid, ) values = {} if title is not None: @@ -243,7 +251,7 @@ class SQLiteDatabase(BaseDatabase): if content is not None: values["content"] = content if not values: - return + return None query = query.values(**values) await session.execute(query) return await self.get_conversation_by_id(cid) @@ -254,8 +262,8 @@ class SQLiteDatabase(BaseDatabase): async with session.begin(): await session.execute( delete(ConversationV2).where( - col(ConversationV2.conversation_id) == cid - ) + col(ConversationV2.conversation_id) == cid, + ), ) async def delete_conversations_by_user_id(self, user_id: str) -> None: @@ -263,7 +271,9 @@ class SQLiteDatabase(BaseDatabase): session: AsyncSession async with session.begin(): await session.execute( - delete(ConversationV2).where(col(ConversationV2.user_id) == user_id) + delete(ConversationV2).where( + col(ConversationV2.user_id) == user_id + ), ) async def get_session_conversations( @@ -282,7 +292,7 @@ class SQLiteDatabase(BaseDatabase): select( col(Preference.scope_id).label("session_id"), func.json_extract(Preference.value, "$.val").label( - "conversation_id" + "conversation_id", ), # type: ignore col(ConversationV2.persona_id).label("persona_id"), col(ConversationV2.title).label("title"), @@ -295,7 +305,8 @@ class SQLiteDatabase(BaseDatabase): == ConversationV2.conversation_id, ) .outerjoin( - Persona, col(ConversationV2.persona_id) == Persona.persona_id + Persona, + col(ConversationV2.persona_id) == Persona.persona_id, ) .where(Preference.scope == "umo", Preference.key == "sel_conv_id") ) @@ -308,14 +319,14 @@ class SQLiteDatabase(BaseDatabase): col(Preference.scope_id).ilike(search_pattern), col(ConversationV2.title).ilike(search_pattern), col(Persona.persona_id).ilike(search_pattern), - ) + ), ) # 平台筛选 if platform: platform_pattern = f"{platform}:%" base_query = base_query.where( - col(Preference.scope_id).like(platform_pattern) + col(Preference.scope_id).like(platform_pattern), ) # 排序 @@ -336,7 +347,8 @@ class SQLiteDatabase(BaseDatabase): == ConversationV2.conversation_id, ) .outerjoin( - Persona, col(ConversationV2.persona_id) == Persona.persona_id + Persona, + col(ConversationV2.persona_id) == Persona.persona_id, ) .where(Preference.scope == "umo", Preference.key == "sel_conv_id") ) @@ -349,13 +361,13 @@ class SQLiteDatabase(BaseDatabase): col(Preference.scope_id).ilike(search_pattern), col(ConversationV2.title).ilike(search_pattern), col(Persona.persona_id).ilike(search_pattern), - ) + ), ) if platform: platform_pattern = f"{platform}:%" count_base_query = count_base_query.where( - col(Preference.scope_id).like(platform_pattern) + col(Preference.scope_id).like(platform_pattern), ) total_result = await session.execute(count_base_query) @@ -396,9 +408,12 @@ class SQLiteDatabase(BaseDatabase): return new_history async def delete_platform_message_offset( - self, platform_id, user_id, offset_sec=86400 + self, + platform_id, + user_id, + offset_sec=86400, ): - """Delete platform message history records older than the specified offset.""" + """Delete platform message history records newer than the specified offset.""" async with self.get_db() as session: session: AsyncSession async with session.begin(): @@ -408,12 +423,16 @@ class SQLiteDatabase(BaseDatabase): delete(PlatformMessageHistory).where( col(PlatformMessageHistory.platform_id) == platform_id, col(PlatformMessageHistory.user_id) == user_id, - col(PlatformMessageHistory.created_at) < cutoff_time, - ) + col(PlatformMessageHistory.created_at) >= cutoff_time, + ), ) async def get_platform_message_history( - self, platform_id, user_id, page=1, page_size=20 + self, + platform_id, + user_id, + page=1, + page_size=20, ): """Get platform message history records.""" async with self.get_db() as session: @@ -452,7 +471,11 @@ class SQLiteDatabase(BaseDatabase): return result.scalar_one_or_none() async def insert_persona( - self, persona_id, system_prompt, begin_dialogs=None, tools=None + self, + persona_id, + system_prompt, + begin_dialogs=None, + tools=None, ): """Insert a new persona record.""" async with self.get_db() as session: @@ -484,7 +507,11 @@ class SQLiteDatabase(BaseDatabase): return result.scalars().all() async def update_persona( - self, persona_id, system_prompt=None, begin_dialogs=None, tools=NOT_GIVEN + self, + persona_id, + system_prompt=None, + begin_dialogs=None, + tools=NOT_GIVEN, ): """Update a persona's system prompt or begin dialogs.""" async with self.get_db() as session: @@ -499,7 +526,7 @@ class SQLiteDatabase(BaseDatabase): if tools is not NOT_GIVEN: values["tools"] = tools if not values: - return + return None query = query.values(**values) await session.execute(query) return await self.get_persona_by_id(persona_id) @@ -510,7 +537,7 @@ class SQLiteDatabase(BaseDatabase): session: AsyncSession async with session.begin(): await session.execute( - delete(Persona).where(col(Persona.persona_id) == persona_id) + delete(Persona).where(col(Persona.persona_id) == persona_id), ) async def insert_preference_or_update(self, scope, scope_id, key, value): @@ -529,7 +556,10 @@ class SQLiteDatabase(BaseDatabase): existing_preference.value = value else: new_preference = Preference( - scope=scope, scope_id=scope_id, key=key, value=value + scope=scope, + scope_id=scope_id, + key=key, + value=value, ) session.add(new_preference) return existing_preference or new_preference @@ -568,7 +598,7 @@ class SQLiteDatabase(BaseDatabase): col(Preference.scope) == scope, col(Preference.scope_id) == scope_id, col(Preference.key) == key, - ) + ), ) await session.commit() @@ -581,7 +611,7 @@ class SQLiteDatabase(BaseDatabase): delete(Preference).where( col(Preference.scope) == scope, col(Preference.scope_id) == scope_id, - ) + ), ) await session.commit() @@ -598,7 +628,7 @@ class SQLiteDatabase(BaseDatabase): now = datetime.now() start_time = now - timedelta(seconds=offset_sec) result = await session.execute( - select(PlatformStat).where(PlatformStat.timestamp >= start_time) + select(PlatformStat).where(PlatformStat.timestamp >= start_time), ) all_datas = result.scalars().all() deprecated_stats = DeprecatedStats() @@ -608,7 +638,7 @@ class SQLiteDatabase(BaseDatabase): name=data.platform_id, count=data.count, timestamp=int(data.timestamp.timestamp()), - ) + ), ) return deprecated_stats @@ -630,7 +660,7 @@ class SQLiteDatabase(BaseDatabase): async with self.get_db() as session: session: AsyncSession result = await session.execute( - select(func.sum(PlatformStat.count)).select_from(PlatformStat) + select(func.sum(PlatformStat.count)).select_from(PlatformStat), ) total_count = result.scalar_one_or_none() return total_count if total_count is not None else 0 @@ -656,7 +686,7 @@ class SQLiteDatabase(BaseDatabase): result = await session.execute( select(PlatformStat.platform_id, func.sum(PlatformStat.count)) .where(PlatformStat.timestamp >= start_time) - .group_by(PlatformStat.platform_id) + .group_by(PlatformStat.platform_id), ) grouped_stats = result.all() deprecated_stats = DeprecatedStats() @@ -666,7 +696,7 @@ class SQLiteDatabase(BaseDatabase): name=platform_id, count=count, timestamp=int(start_time.timestamp()), - ) + ), ) return deprecated_stats @@ -680,3 +710,101 @@ class SQLiteDatabase(BaseDatabase): t.start() t.join() return result + + # ==== + # Platform Session Management + # ==== + + async def create_platform_session( + self, + creator: str, + platform_id: str = "webchat", + session_id: str | None = None, + display_name: str | None = None, + is_group: int = 0, + ) -> PlatformSession: + """Create a new Platform session.""" + kwargs = {} + if session_id: + kwargs["session_id"] = session_id + + async with self.get_db() as session: + session: AsyncSession + async with session.begin(): + new_session = PlatformSession( + creator=creator, + platform_id=platform_id, + display_name=display_name, + is_group=is_group, + **kwargs, + ) + session.add(new_session) + await session.flush() + await session.refresh(new_session) + return new_session + + async def get_platform_session_by_id( + self, session_id: str + ) -> PlatformSession | None: + """Get a Platform session by its ID.""" + async with self.get_db() as session: + session: AsyncSession + query = select(PlatformSession).where( + PlatformSession.session_id == session_id, + ) + result = await session.execute(query) + return result.scalar_one_or_none() + + async def get_platform_sessions_by_creator( + self, + creator: str, + platform_id: str | None = None, + page: int = 1, + page_size: int = 20, + ) -> list[PlatformSession]: + """Get all Platform sessions for a specific creator (username) and optionally platform.""" + async with self.get_db() as session: + session: AsyncSession + offset = (page - 1) * page_size + query = select(PlatformSession).where(PlatformSession.creator == creator) + + if platform_id: + query = query.where(PlatformSession.platform_id == platform_id) + + query = ( + query.order_by(desc(PlatformSession.updated_at)) + .offset(offset) + .limit(page_size) + ) + result = await session.execute(query) + return list(result.scalars().all()) + + async def update_platform_session( + self, + session_id: str, + display_name: str | None = None, + ) -> None: + """Update a Platform session's updated_at timestamp and optionally display_name.""" + async with self.get_db() as session: + session: AsyncSession + async with session.begin(): + values: dict[str, T.Any] = {"updated_at": datetime.now(timezone.utc)} + if display_name is not None: + values["display_name"] = display_name + + await session.execute( + update(PlatformSession) + .where(col(PlatformSession.session_id == session_id)) + .values(**values), + ) + + async def delete_platform_session(self, session_id: str) -> None: + """Delete a Platform session by its ID.""" + async with self.get_db() as session: + session: AsyncSession + async with session.begin(): + await session.execute( + delete(PlatformSession).where( + col(PlatformSession.session_id == session_id), + ), + ) diff --git a/astrbot/core/db/vec_db/base.py b/astrbot/core/db/vec_db/base.py index d6e5e56b..7440b6f2 100644 --- a/astrbot/core/db/vec_db/base.py +++ b/astrbot/core/db/vec_db/base.py @@ -10,22 +10,47 @@ class Result: class BaseVecDB: async def initialize(self): - """ - 初始化向量数据库 - """ - pass + """初始化向量数据库""" @abc.abstractmethod - async def insert(self, content: str, metadata: dict = None, id: str = None) -> int: - """ - 插入一条文本和其对应向量，自动生成 ID 并保持一致性。 + async def insert( + self, + content: str, + metadata: dict | None = None, + id: str | None = None, + ) -> int: + """插入一条文本和其对应向量，自动生成 ID 并保持一致性。""" + ... + + @abc.abstractmethod + async def insert_batch( + self, + contents: list[str], + metadatas: list[dict] | None = None, + ids: list[str] | None = None, + batch_size: int = 32, + tasks_limit: int = 3, + max_retries: int = 3, + progress_callback=None, + ) -> int: + """批量插入文本和其对应向量，自动生成 ID 并保持一致性。 + + Args: + progress_callback: 进度回调函数，接收参数 (current, total) + """ ... @abc.abstractmethod - async def retrieve(self, query: str, top_k: int = 5) -> list[Result]: - """ - 搜索最相似的文档。 + async def retrieve( + self, + query: str, + top_k: int = 5, + fetch_k: int = 20, + rerank: bool = False, + metadata_filters: dict | None = None, + ) -> list[Result]: + """搜索最相似的文档。 Args: query (str): 查询文本 top_k (int): 返回的最相似文档的数量 @@ -36,11 +61,13 @@ class BaseVecDB: @abc.abstractmethod async def delete(self, doc_id: str) -> bool: - """ - 删除指定文档。 + """删除指定文档。 Args: doc_id (str): 要删除的文档 ID Returns: bool: 删除是否成功 """ ... + + @abc.abstractmethod + async def close(self): ... diff --git a/astrbot/core/db/vec_db/faiss_impl/document_storage.py b/astrbot/core/db/vec_db/faiss_impl/document_storage.py index ee44da66..e27eb6fe 100644 --- a/astrbot/core/db/vec_db/faiss_impl/document_storage.py +++ b/astrbot/core/db/vec_db/faiss_impl/document_storage.py @@ -1,59 +1,232 @@ -import aiosqlite +import json import os +from contextlib import asynccontextmanager +from datetime import datetime + +from sqlalchemy import Column, Text +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine +from sqlalchemy.orm import sessionmaker +from sqlmodel import Field, MetaData, SQLModel, col, func, select, text + +from astrbot.core import logger + + +class BaseDocModel(SQLModel, table=False): + metadata = MetaData() + + +class Document(BaseDocModel, table=True): + """SQLModel for documents table.""" + + __tablename__ = "documents" # type: ignore + + id: int | None = Field( + default=None, + primary_key=True, + sa_column_kwargs={"autoincrement": True}, + ) + doc_id: str = Field(nullable=False) + text: str = Field(nullable=False) + metadata_: str | None = Field(default=None, sa_column=Column("metadata", Text)) + created_at: datetime | None = Field(default=None) + updated_at: datetime | None = Field(default=None) class DocumentStorage: def __init__(self, db_path: str): self.db_path = db_path - self.connection = None + self.DATABASE_URL = f"sqlite+aiosqlite:///{db_path}" + self.engine: AsyncEngine | None = None + self.async_session_maker: sessionmaker | None = None self.sqlite_init_path = os.path.join( - os.path.dirname(__file__), "sqlite_init.sql" + os.path.dirname(__file__), + "sqlite_init.sql", ) async def initialize(self): """Initialize the SQLite database and create the documents table if it doesn't exist.""" - if not os.path.exists(self.db_path): - await self.connect() - async with self.connection.cursor() as cursor: - with open(self.sqlite_init_path, "r", encoding="utf-8") as f: - sql_script = f.read() - await cursor.executescript(sql_script) - await self.connection.commit() - else: - await self.connect() + await self.connect() + async with self.engine.begin() as conn: # type: ignore + # Create tables using SQLModel + await conn.run_sync(BaseDocModel.metadata.create_all) + + try: + await conn.execute( + text( + "ALTER TABLE documents ADD COLUMN kb_doc_id TEXT " + "GENERATED ALWAYS AS (json_extract(metadata, '$.kb_doc_id')) STORED", + ), + ) + await conn.execute( + text( + "ALTER TABLE documents ADD COLUMN user_id TEXT " + "GENERATED ALWAYS AS (json_extract(metadata, '$.user_id')) STORED", + ), + ) + + # Create indexes + await conn.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_documents_kb_doc_id ON documents(kb_doc_id)", + ), + ) + await conn.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_documents_user_id ON documents(user_id)", + ), + ) + except BaseException: + pass + + await conn.commit() async def connect(self): """Connect to the SQLite database.""" - self.connection = await aiosqlite.connect(self.db_path) + if self.engine is None: + self.engine = create_async_engine( + self.DATABASE_URL, + echo=False, + future=True, + ) + self.async_session_maker = sessionmaker( + self.engine, # type: ignore + class_=AsyncSession, + expire_on_commit=False, + ) # type: ignore - async def get_documents(self, metadata_filters: dict, ids: list = None): + @asynccontextmanager + async def get_session(self): + """Context manager for database sessions.""" + async with self.async_session_maker() as session: # type: ignore + yield session + + async def get_documents( + self, + metadata_filters: dict, + ids: list | None = None, + offset: int | None = 0, + limit: int | None = 100, + ) -> list[dict]: """Retrieve documents by metadata filters and ids. Args: metadata_filters (dict): The metadata filters to apply. + ids (list | None): Optional list of document IDs to filter. + offset (int | None): Offset for pagination. + limit (int | None): Limit for pagination. Returns: - list: The list of document IDs(primary key, not doc_id) that match the filters. - """ - # metadata filter -> SQL WHERE clause - where_clauses = [] - values = [] - for key, val in metadata_filters.items(): - where_clauses.append(f"json_extract(metadata, '$.{key}') = ?") - values.append(val) - if ids is not None and len(ids) > 0: - ids = [str(i) for i in ids if i != -1] - where_clauses.append("id IN ({})".format(",".join("?" * len(ids)))) - values.extend(ids) - where_sql = " AND ".join(where_clauses) or "1=1" + list: The list of documents that match the filters. - result = [] - async with self.connection.cursor() as cursor: - sql = "SELECT * FROM documents WHERE " + where_sql - await cursor.execute(sql, values) - for row in await cursor.fetchall(): - result.append(await self.tuple_to_dict(row)) - return result + """ + if self.engine is None: + logger.warning( + "Database connection is not initialized, returning empty result", + ) + return [] + + async with self.get_session() as session: + query = select(Document) + + for key, val in metadata_filters.items(): + query = query.where( + text(f"json_extract(metadata, '$.{key}') = :filter_{key}"), + ).params(**{f"filter_{key}": val}) + + if ids is not None and len(ids) > 0: + valid_ids = [int(i) for i in ids if i != -1] + if valid_ids: + query = query.where(col(Document.id).in_(valid_ids)) + + if limit is not None: + query = query.limit(limit) + if offset is not None: + query = query.offset(offset) + + result = await session.execute(query) + documents = result.scalars().all() + + return [self._document_to_dict(doc) for doc in documents] + + async def insert_document(self, doc_id: str, text: str, metadata: dict) -> int: + """Insert a single document and return its integer ID. + + Args: + doc_id (str): The document ID (UUID string). + text (str): The document text. + metadata (dict): The document metadata. + + Returns: + int: The integer ID of the inserted document. + + """ + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session, session.begin(): + document = Document( + doc_id=doc_id, + text=text, + metadata_=json.dumps(metadata), + created_at=datetime.now(), + updated_at=datetime.now(), + ) + session.add(document) + await session.flush() # Flush to get the ID + return document.id # type: ignore + + async def insert_documents_batch( + self, + doc_ids: list[str], + texts: list[str], + metadatas: list[dict], + ) -> list[int]: + """Batch insert documents and return their integer IDs. + + Args: + doc_ids (list[str]): List of document IDs (UUID strings). + texts (list[str]): List of document texts. + metadatas (list[dict]): List of document metadata. + + Returns: + list[int]: List of integer IDs of the inserted documents. + + """ + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session, session.begin(): + import json + + documents = [] + for doc_id, text, metadata in zip(doc_ids, texts, metadatas): + document = Document( + doc_id=doc_id, + text=text, + metadata_=json.dumps(metadata), + created_at=datetime.now(), + updated_at=datetime.now(), + ) + documents.append(document) + session.add(document) + + await session.flush() # Flush to get all IDs + return [doc.id for doc in documents] # type: ignore + + async def delete_document_by_doc_id(self, doc_id: str): + """Delete a document by its doc_id. + + Args: + doc_id (str): The doc_id of the document to delete. + + """ + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session, session.begin(): + query = select(Document).where(col(Document.doc_id) == doc_id) + result = await session.execute(query) + document = result.scalar_one_or_none() + + if document: + await session.delete(document) async def get_document_by_doc_id(self, doc_id: str): """Retrieve a document by its doc_id. @@ -62,40 +235,134 @@ class DocumentStorage: doc_id (str): The doc_id of the document to retrieve. Returns: - dict: The document data. + dict: The document data or None if not found. + """ - async with self.connection.cursor() as cursor: - await cursor.execute("SELECT * FROM documents WHERE doc_id = ?", (doc_id,)) - row = await cursor.fetchone() - if row: - return await self.tuple_to_dict(row) - else: - return None + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session: + query = select(Document).where(col(Document.doc_id) == doc_id) + result = await session.execute(query) + document = result.scalar_one_or_none() + + if document: + return self._document_to_dict(document) + return None async def update_document_by_doc_id(self, doc_id: str, new_text: str): - """Retrieve a document by its doc_id. + """Update a document by its doc_id. Args: doc_id (str): The doc_id. new_text (str): The new text to update the document with. + """ - async with self.connection.cursor() as cursor: - await cursor.execute( - "UPDATE documents SET text = ? WHERE doc_id = ?", (new_text, doc_id) + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session, session.begin(): + query = select(Document).where(col(Document.doc_id) == doc_id) + result = await session.execute(query) + document = result.scalar_one_or_none() + + if document: + document.text = new_text + document.updated_at = datetime.now() + session.add(document) + + async def delete_documents(self, metadata_filters: dict): + """Delete documents by their metadata filters. + + Args: + metadata_filters (dict): The metadata filters to apply. + + """ + if self.engine is None: + logger.warning( + "Database connection is not initialized, skipping delete operation", ) - await self.connection.commit() + return + + async with self.get_session() as session, session.begin(): + query = select(Document) + + for key, val in metadata_filters.items(): + query = query.where( + text(f"json_extract(metadata, '$.{key}') = :filter_{key}"), + ).params(**{f"filter_{key}": val}) + + result = await session.execute(query) + documents = result.scalars().all() + + for doc in documents: + await session.delete(doc) + + async def count_documents(self, metadata_filters: dict | None = None) -> int: + """Count documents in the database. + + Args: + metadata_filters (dict | None): Metadata filters to apply. + + Returns: + int: The count of documents. + + """ + if self.engine is None: + logger.warning("Database connection is not initialized, returning 0") + return 0 + + async with self.get_session() as session: + query = select(func.count(col(Document.id))) + + if metadata_filters: + for key, val in metadata_filters.items(): + query = query.where( + text(f"json_extract(metadata, '$.{key}') = :filter_{key}"), + ).params(**{f"filter_{key}": val}) + + result = await session.execute(query) + count = result.scalar_one_or_none() + return count if count is not None else 0 async def get_user_ids(self) -> list[str]: """Retrieve all user IDs from the documents table. Returns: list: A list of user IDs. + """ - async with self.connection.cursor() as cursor: - await cursor.execute("SELECT DISTINCT user_id FROM documents") - rows = await cursor.fetchall() + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session: + query = text( + "SELECT DISTINCT user_id FROM documents WHERE user_id IS NOT NULL", + ) + result = await session.execute(query) + rows = result.fetchall() return [row[0] for row in rows] + def _document_to_dict(self, document: Document) -> dict: + """Convert a Document model to a dictionary. + + Args: + document (Document): The document to convert. + + Returns: + dict: The converted dictionary. + + """ + return { + "id": document.id, + "doc_id": document.doc_id, + "text": document.text, + "metadata": document.metadata_, + "created_at": document.created_at.isoformat() + if isinstance(document.created_at, datetime) + else document.created_at, + "updated_at": document.updated_at.isoformat() + if isinstance(document.updated_at, datetime) + else document.updated_at, + } + async def tuple_to_dict(self, row): """Convert a tuple to a dictionary. @@ -104,6 +371,9 @@ class DocumentStorage: Returns: dict: The converted dictionary. + + Note: This method is kept for backward compatibility but is no longer used internally. + """ return { "id": row[0], @@ -116,6 +386,7 @@ class DocumentStorage: async def close(self): """Close the connection to the SQLite database.""" - if self.connection: - await self.connection.close() - self.connection = None + if self.engine: + await self.engine.dispose() + self.engine = None + self.async_session_maker = None diff --git a/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py b/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py index 18ee3189..24f1c323 100644 --- a/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py +++ b/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py @@ -2,14 +2,15 @@ try: import faiss except ModuleNotFoundError: raise ImportError( - "faiss 未安装。请使用 'pip install faiss-cpu' 或 'pip install faiss-gpu' 安装。" + "faiss 未安装。请使用 'pip install faiss-cpu' 或 'pip install faiss-gpu' 安装。", ) import os + import numpy as np class EmbeddingStorage: - def __init__(self, dimension: int, path: str = None): + def __init__(self, dimension: int, path: str | None = None): self.dimension = dimension self.path = path self.index = None @@ -18,7 +19,6 @@ class EmbeddingStorage: else: base_index = faiss.IndexFlatL2(dimension) self.index = faiss.IndexIDMap(base_index) - self.storage = {} async def insert(self, vector: np.ndarray, id: int): """插入向量 @@ -28,13 +28,32 @@ class EmbeddingStorage: id (int): 向量的ID Raises: ValueError: 如果向量的维度与存储的维度不匹配 + """ + assert self.index is not None, "FAISS index is not initialized." if vector.shape[0] != self.dimension: raise ValueError( - f"向量维度不匹配, 期望: {self.dimension}, 实际: {vector.shape[0]}" + f"向量维度不匹配, 期望: {self.dimension}, 实际: {vector.shape[0]}", ) self.index.add_with_ids(vector.reshape(1, -1), np.array([id])) - self.storage[id] = vector + await self.save_index() + + async def insert_batch(self, vectors: np.ndarray, ids: list[int]): + """批量插入向量 + + Args: + vectors (np.ndarray): 要插入的向量数组 + ids (list[int]): 向量的ID列表 + Raises: + ValueError: 如果向量的维度与存储的维度不匹配 + + """ + assert self.index is not None, "FAISS index is not initialized." + if vectors.shape[1] != self.dimension: + raise ValueError( + f"向量维度不匹配, 期望: {self.dimension}, 实际: {vectors.shape[1]}", + ) + self.index.add_with_ids(vectors, np.array(ids)) await self.save_index() async def search(self, vector: np.ndarray, k: int) -> tuple: @@ -45,15 +64,30 @@ class EmbeddingStorage: k (int): 返回的最相似向量的数量 Returns: tuple: (距离, 索引) + """ + assert self.index is not None, "FAISS index is not initialized." faiss.normalize_L2(vector) distances, indices = self.index.search(vector, k) return distances, indices + async def delete(self, ids: list[int]): + """删除向量 + + Args: + ids (list[int]): 要删除的向量ID列表 + + """ + assert self.index is not None, "FAISS index is not initialized." + id_array = np.array(ids, dtype=np.int64) + self.index.remove_ids(id_array) + await self.save_index() + async def save_index(self): """保存索引 Args: path (str): 保存索引的路径 + """ faiss.write_index(self.index, self.path) diff --git a/astrbot/core/db/vec_db/faiss_impl/vec_db.py b/astrbot/core/db/vec_db/faiss_impl/vec_db.py index 7c2ae1c0..14221f1e 100644 --- a/astrbot/core/db/vec_db/faiss_impl/vec_db.py +++ b/astrbot/core/db/vec_db/faiss_impl/vec_db.py @@ -1,17 +1,18 @@ +import time import uuid -import json + import numpy as np + +from astrbot import logger +from astrbot.core.provider.provider import EmbeddingProvider, RerankProvider + +from ..base import BaseVecDB, Result from .document_storage import DocumentStorage from .embedding_storage import EmbeddingStorage -from ..base import Result, BaseVecDB -from astrbot.core.provider.provider import EmbeddingProvider -from astrbot.core.provider.provider import RerankProvider class FaissVecDB(BaseVecDB): - """ - A class to represent a vector database. - """ + """A class to represent a vector database.""" def __init__( self, @@ -25,7 +26,8 @@ class FaissVecDB(BaseVecDB): self.embedding_provider = embedding_provider self.document_storage = DocumentStorage(doc_store_path) self.embedding_storage = EmbeddingStorage( - embedding_provider.get_dim(), index_store_path + embedding_provider.get_dim(), + index_store_path, ) self.embedding_provider = embedding_provider self.rerank_provider = rerank_provider @@ -34,28 +36,69 @@ class FaissVecDB(BaseVecDB): await self.document_storage.initialize() async def insert( - self, content: str, metadata: dict | None = None, id: str | None = None + self, + content: str, + metadata: dict | None = None, + id: str | None = None, ) -> int: - """ - 插入一条文本和其对应向量，自动生成 ID 并保持一致性。 - """ + """插入一条文本和其对应向量，自动生成 ID 并保持一致性。""" metadata = metadata or {} str_id = id or str(uuid.uuid4()) # 使用 UUID 作为原始 ID vector = await self.embedding_provider.get_embedding(content) vector = np.array(vector, dtype=np.float32) - async with self.document_storage.connection.cursor() as cursor: - await cursor.execute( - "INSERT INTO documents (doc_id, text, metadata) VALUES (?, ?, ?)", - (str_id, content, json.dumps(metadata)), - ) - await self.document_storage.connection.commit() - result = await self.document_storage.get_document_by_doc_id(str_id) - int_id = result["id"] - # 插入向量到 FAISS - await self.embedding_storage.insert(vector, int_id) - return int_id + # 使用 DocumentStorage 的方法插入文档 + int_id = await self.document_storage.insert_document(str_id, content, metadata) + + # 插入向量到 FAISS + await self.embedding_storage.insert(vector, int_id) + return int_id + + async def insert_batch( + self, + contents: list[str], + metadatas: list[dict] | None = None, + ids: list[str] | None = None, + batch_size: int = 32, + tasks_limit: int = 3, + max_retries: int = 3, + progress_callback=None, + ) -> list[int]: + """批量插入文本和其对应向量，自动生成 ID 并保持一致性。 + + Args: + progress_callback: 进度回调函数，接收参数 (current, total) + + """ + metadatas = metadatas or [{} for _ in contents] + ids = ids or [str(uuid.uuid4()) for _ in contents] + + start = time.time() + logger.debug(f"Generating embeddings for {len(contents)} contents...") + vectors = await self.embedding_provider.get_embeddings_batch( + contents, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + progress_callback=progress_callback, + ) + end = time.time() + logger.debug( + f"Generated embeddings for {len(contents)} contents in {end - start:.2f} seconds.", + ) + + # 使用 DocumentStorage 的批量插入方法 + int_ids = await self.document_storage.insert_documents_batch( + ids, + contents, + metadatas, + ) + + # 批量插入向量到 FAISS + vectors_array = np.array(vectors).astype("float32") + await self.embedding_storage.insert_batch(vectors_array, int_ids) + return int_ids async def retrieve( self, @@ -65,8 +108,7 @@ class FaissVecDB(BaseVecDB): rerank: bool = False, metadata_filters: dict | None = None, ) -> list[Result]: - """ - 搜索最相似的文档。 + """搜索最相似的文档。 Args: query (str): 查询文本 @@ -77,6 +119,7 @@ class FaissVecDB(BaseVecDB): Returns: List[Result]: 查询结果 + """ embedding = await self.embedding_provider.get_embedding(query) scores, indices = await self.embedding_storage.search( @@ -89,7 +132,8 @@ class FaissVecDB(BaseVecDB): scores[0] = 1.0 - (scores[0] / 2.0) # NOTE: maybe the size is less than k. fetched_docs = await self.document_storage.get_documents( - metadata_filters=metadata_filters or {}, ids=indices[0] + metadata_filters=metadata_filters or {}, + ids=indices[0], ) if not fetched_docs: return [] @@ -110,7 +154,9 @@ class FaissVecDB(BaseVecDB): documents = [doc.data["text"] for doc in top_k_results] reranked_results = await self.rerank_provider.rerank(query, documents) reranked_results = sorted( - reranked_results, key=lambda x: x.relevance_score, reverse=True + reranked_results, + key=lambda x: x.relevance_score, + reverse=True, ) top_k_results = [ top_k_results[reranked_result.index] @@ -119,23 +165,40 @@ class FaissVecDB(BaseVecDB): return top_k_results - async def delete(self, doc_id: int): - """ - 删除一条文档 - """ - await self.document_storage.connection.execute( - "DELETE FROM documents WHERE doc_id = ?", (doc_id,) - ) - await self.document_storage.connection.commit() + async def delete(self, doc_id: str): + """删除一条文档块（chunk）""" + # 获得对应的 int id + result = await self.document_storage.get_document_by_doc_id(doc_id) + int_id = result["id"] if result else None + if int_id is None: + return + + # 使用 DocumentStorage 的删除方法 + await self.document_storage.delete_document_by_doc_id(doc_id) + await self.embedding_storage.delete([int_id]) async def close(self): await self.document_storage.close() - async def count_documents(self) -> int: + async def count_documents(self, metadata_filter: dict | None = None) -> int: + """计算文档数量 + + Args: + metadata_filter (dict | None): 元数据过滤器 + """ - 计算文档数量 - """ - async with self.document_storage.connection.cursor() as cursor: - await cursor.execute("SELECT COUNT(*) FROM documents") - count = await cursor.fetchone() - return count[0] if count else 0 + count = await self.document_storage.count_documents( + metadata_filters=metadata_filter or {}, + ) + return count + + async def delete_documents(self, metadata_filters: dict): + """根据元数据过滤器删除文档""" + docs = await self.document_storage.get_documents( + metadata_filters=metadata_filters, + offset=None, + limit=None, + ) + doc_ids: list[int] = [doc["id"] for doc in docs] + await self.embedding_storage.delete(doc_ids) + await self.document_storage.delete_documents(metadata_filters=metadata_filters) diff --git a/astrbot/core/event_bus.py b/astrbot/core/event_bus.py index 2ae70939..749df753 100644 --- a/astrbot/core/event_bus.py +++ b/astrbot/core/event_bus.py @@ -1,5 +1,4 @@ -""" -事件总线, 用于处理事件的分发和处理 +"""事件总线, 用于处理事件的分发和处理 事件总线是一个异步队列, 用于接收各种消息事件, 并将其发送到Scheduler调度器进行处理 其中包含了一个无限循环的调度函数, 用于从事件队列中获取新的事件, 并创建一个新的异步任务来执行管道调度器的处理逻辑 @@ -13,10 +12,12 @@ class: import asyncio from asyncio import Queue -from astrbot.core.pipeline.scheduler import PipelineScheduler + from astrbot.core import logger -from .platform import AstrMessageEvent from astrbot.core.astrbot_config_mgr import AstrBotConfigManager +from astrbot.core.pipeline.scheduler import PipelineScheduler + +from .platform import AstrMessageEvent class EventBus: @@ -46,14 +47,15 @@ class EventBus: Args: event (AstrMessageEvent): 事件对象 + """ # 如果有发送者名称: [平台名] 发送者名称/发送者ID: 消息概要 if event.get_sender_name(): logger.info( - f"[{conf_name}] [{event.get_platform_id()}({event.get_platform_name()})] {event.get_sender_name()}/{event.get_sender_id()}: {event.get_message_outline()}" + f"[{conf_name}] [{event.get_platform_id()}({event.get_platform_name()})] {event.get_sender_name()}/{event.get_sender_id()}: {event.get_message_outline()}", ) # 没有发送者名称: [平台名] 发送者ID: 消息概要 else: logger.info( - f"[{conf_name}] [{event.get_platform_id()}({event.get_platform_name()})] {event.get_sender_id()}: {event.get_message_outline()}" + f"[{conf_name}] [{event.get_platform_id()}({event.get_platform_name()})] {event.get_sender_id()}: {event.get_message_outline()}", ) diff --git a/astrbot/core/exceptions.py b/astrbot/core/exceptions.py new file mode 100644 index 00000000..e637d493 --- /dev/null +++ b/astrbot/core/exceptions.py @@ -0,0 +1,9 @@ +from __future__ import annotations + + +class AstrBotError(Exception): + """Base exception for all AstrBot errors.""" + + +class ProviderNotFoundError(AstrBotError): + """Raised when a specified provider is not found.""" diff --git a/astrbot/core/file_token_service.py b/astrbot/core/file_token_service.py index ce5e2349..ea97759c 100644 --- a/astrbot/core/file_token_service.py +++ b/astrbot/core/file_token_service.py @@ -1,9 +1,9 @@ import asyncio import os -import uuid -import time -from urllib.parse import urlparse, unquote import platform +import time +import uuid +from urllib.parse import unquote, urlparse class FileTokenService: @@ -23,7 +23,12 @@ class FileTokenService: for token in expired_tokens: self.staged_files.pop(token, None) - async def register_file(self, file_path: str, timeout: float = None) -> str: + async def check_token_expired(self, file_token: str) -> bool: + async with self.lock: + await self._cleanup_expired_tokens() + return file_token not in self.staged_files + + async def register_file(self, file_path: str, timeout: float | None = None) -> str: """向令牌服务注册一个文件。 Args: @@ -35,8 +40,8 @@ class FileTokenService: Raises: FileNotFoundError: 当路径不存在时抛出 - """ + """ # 处理 file:/// try: parsed_uri = urlparse(file_path) @@ -56,7 +61,7 @@ class FileTokenService: if not os.path.exists(local_path): raise FileNotFoundError( - f"文件不存在: {local_path} (原始输入: {file_path})" + f"文件不存在: {local_path} (原始输入: {file_path})", ) file_token = str(uuid.uuid4()) @@ -79,6 +84,7 @@ class FileTokenService: Raises: KeyError: 当令牌不存在或已过期时抛出 FileNotFoundError: 当文件本身已被删除时抛出 + """ async with self.lock: await self._cleanup_expired_tokens() diff --git a/astrbot/core/initial_loader.py b/astrbot/core/initial_loader.py index b3eeb6e8..f54d1864 100644 --- a/astrbot/core/initial_loader.py +++ b/astrbot/core/initial_loader.py @@ -1,5 +1,4 @@ -""" -AstrBot 启动器，负责初始化和启动核心组件和仪表板服务器。 +"""AstrBot 启动器，负责初始化和启动核心组件和仪表板服务器。 工作流程: 1. 初始化核心生命周期, 传递数据库和日志代理实例到核心生命周期 @@ -8,10 +7,10 @@ AstrBot 启动器，负责初始化和启动核心组件和仪表板服务器。 import asyncio import traceback -from astrbot.core import logger + +from astrbot.core import LogBroker, logger from astrbot.core.core_lifecycle import AstrBotCoreLifecycle from astrbot.core.db import BaseDatabase -from astrbot.core import LogBroker from astrbot.dashboard.server import AstrBotDashboard @@ -39,12 +38,18 @@ class InitialLoader: webui_dir = self.webui_dir self.dashboard_server = AstrBotDashboard( - core_lifecycle, self.db, core_lifecycle.dashboard_shutdown_event, webui_dir + core_lifecycle, + self.db, + core_lifecycle.dashboard_shutdown_event, + webui_dir, ) - task = asyncio.gather( - core_task, self.dashboard_server.run() - ) # 启动核心任务和仪表板服务器 + coro = self.dashboard_server.run() + if coro: + # 启动核心任务和仪表板服务器 + task = asyncio.gather(core_task, coro) + else: + task = core_task try: await task # 整个AstrBot在这里运行 except asyncio.CancelledError: diff --git a/astrbot/core/knowledge_base/chunking/__init__.py b/astrbot/core/knowledge_base/chunking/__init__.py new file mode 100644 index 00000000..805ddc24 --- /dev/null +++ b/astrbot/core/knowledge_base/chunking/__init__.py @@ -0,0 +1,9 @@ +"""文档分块模块""" + +from .base import BaseChunker +from .fixed_size import FixedSizeChunker + +__all__ = [ + "BaseChunker", + "FixedSizeChunker", +] diff --git a/astrbot/core/knowledge_base/chunking/base.py b/astrbot/core/knowledge_base/chunking/base.py new file mode 100644 index 00000000..a45d86ad --- /dev/null +++ b/astrbot/core/knowledge_base/chunking/base.py @@ -0,0 +1,25 @@ +"""文档分块器基类 + +定义了文档分块处理的抽象接口。 +""" + +from abc import ABC, abstractmethod + + +class BaseChunker(ABC): + """分块器基类 + + 所有分块器都应该继承此类并实现 chunk 方法。 + """ + + @abstractmethod + async def chunk(self, text: str, **kwargs) -> list[str]: + """将文本分块 + + Args: + text: 输入文本 + + Returns: + list[str]: 分块后的文本列表 + + """ diff --git a/astrbot/core/knowledge_base/chunking/fixed_size.py b/astrbot/core/knowledge_base/chunking/fixed_size.py new file mode 100644 index 00000000..5439f070 --- /dev/null +++ b/astrbot/core/knowledge_base/chunking/fixed_size.py @@ -0,0 +1,59 @@ +"""固定大小分块器 + +按照固定的字符数将文本分块,支持重叠区域。 +""" + +from .base import BaseChunker + + +class FixedSizeChunker(BaseChunker): + """固定大小分块器 + + 按照固定的字符数分块,并支持块之间的重叠。 + """ + + def __init__(self, chunk_size: int = 512, chunk_overlap: int = 50): + """初始化分块器 + + Args: + chunk_size: 块的大小(字符数) + chunk_overlap: 块之间的重叠字符数 + + """ + self.chunk_size = chunk_size + self.chunk_overlap = chunk_overlap + + async def chunk(self, text: str, **kwargs) -> list[str]: + """固定大小分块 + + Args: + text: 输入文本 + chunk_size: 每个文本块的最大大小 + chunk_overlap: 每个文本块之间的重叠部分大小 + + Returns: + list[str]: 分块后的文本列表 + + """ + chunk_size = kwargs.get("chunk_size", self.chunk_size) + chunk_overlap = kwargs.get("chunk_overlap", self.chunk_overlap) + + chunks = [] + start = 0 + text_len = len(text) + + while start < text_len: + end = start + chunk_size + chunk = text[start:end] + + if chunk: + chunks.append(chunk) + + # 移动窗口,保留重叠部分 + start = end - chunk_overlap + + # 防止无限循环: 如果重叠过大,直接移到end + if start >= end or chunk_overlap >= chunk_size: + start = end + + return chunks diff --git a/astrbot/core/knowledge_base/chunking/recursive.py b/astrbot/core/knowledge_base/chunking/recursive.py new file mode 100644 index 00000000..3f4aabb5 --- /dev/null +++ b/astrbot/core/knowledge_base/chunking/recursive.py @@ -0,0 +1,161 @@ +from collections.abc import Callable + +from .base import BaseChunker + + +class RecursiveCharacterChunker(BaseChunker): + def __init__( + self, + chunk_size: int = 500, + chunk_overlap: int = 100, + length_function: Callable[[str], int] = len, + is_separator_regex: bool = False, + separators: list[str] | None = None, + ): + """初始化递归字符文本分割器 + + Args: + chunk_size: 每个文本块的最大大小 + chunk_overlap: 每个文本块之间的重叠部分大小 + length_function: 计算文本长度的函数 + is_separator_regex: 分隔符是否为正则表达式 + separators: 用于分割文本的分隔符列表，按优先级排序 + + """ + self.chunk_size = chunk_size + self.chunk_overlap = chunk_overlap + self.length_function = length_function + self.is_separator_regex = is_separator_regex + + # 默认分隔符列表，按优先级从高到低 + self.separators = separators or [ + "\n\n", # 段落 + "\n", # 换行 + "。", # 中文句子 + "，", # 中文逗号 + ". ", # 句子 + ", ", # 逗号分隔 + " ", # 单词 + "", # 字符 + ] + + async def chunk(self, text: str, **kwargs) -> list[str]: + """递归地将文本分割成块 + + Args: + text: 要分割的文本 + chunk_size: 每个文本块的最大大小 + chunk_overlap: 每个文本块之间的重叠部分大小 + + Returns: + 分割后的文本块列表 + + """ + if not text: + return [] + + overlap = kwargs.get("chunk_overlap", self.chunk_overlap) + chunk_size = kwargs.get("chunk_size", self.chunk_size) + + text_length = self.length_function(text) + if text_length <= chunk_size: + return [text] + + for separator in self.separators: + if separator == "": + return self._split_by_character(text, chunk_size, overlap) + + if separator in text: + splits = text.split(separator) + # 重新添加分隔符（除了最后一个片段） + splits = [s + separator for s in splits[:-1]] + [splits[-1]] + splits = [s for s in splits if s] + if len(splits) == 1: + continue + + # 递归合并分割后的文本块 + final_chunks = [] + current_chunk = [] + current_chunk_length = 0 + + for split in splits: + split_length = self.length_function(split) + + # 如果单个分割部分已经超过了chunk_size，需要递归分割 + if split_length > chunk_size: + # 先处理当前积累的块 + if current_chunk: + combined_text = "".join(current_chunk) + final_chunks.extend( + await self.chunk( + combined_text, + chunk_size=chunk_size, + chunk_overlap=overlap, + ), + ) + current_chunk = [] + current_chunk_length = 0 + + # 递归分割过大的部分 + final_chunks.extend( + await self.chunk( + split, + chunk_size=chunk_size, + chunk_overlap=overlap, + ), + ) + # 如果添加这部分会使当前块超过chunk_size + elif current_chunk_length + split_length > chunk_size: + # 合并当前块并添加到结果中 + combined_text = "".join(current_chunk) + final_chunks.append(combined_text) + + # 处理重叠部分 + overlap_start = max(0, len(combined_text) - overlap) + if overlap_start > 0: + overlap_text = combined_text[overlap_start:] + current_chunk = [overlap_text, split] + current_chunk_length = ( + self.length_function(overlap_text) + split_length + ) + else: + current_chunk = [split] + current_chunk_length = split_length + else: + # 添加到当前块 + current_chunk.append(split) + current_chunk_length += split_length + + # 处理剩余的块 + if current_chunk: + final_chunks.append("".join(current_chunk)) + + return final_chunks + + return [text] + + def _split_by_character( + self, + text: str, + chunk_size: int | None = None, + overlap: int | None = None, + ) -> list[str]: + """按字符级别分割文本 + + Args: + text: 要分割的文本 + + Returns: + 分割后的文本块列表 + + """ + chunk_size = chunk_size or self.chunk_size + overlap = overlap or self.chunk_overlap + result = [] + for i in range(0, len(text), chunk_size - overlap): + end = min(i + chunk_size, len(text)) + result.append(text[i:end]) + if end == len(text): + break + + return result diff --git a/astrbot/core/knowledge_base/kb_db_sqlite.py b/astrbot/core/knowledge_base/kb_db_sqlite.py new file mode 100644 index 00000000..5e1db842 --- /dev/null +++ b/astrbot/core/knowledge_base/kb_db_sqlite.py @@ -0,0 +1,301 @@ +from contextlib import asynccontextmanager +from pathlib import Path + +from sqlalchemy import delete, func, select, text, update +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine +from sqlmodel import col, desc + +from astrbot.core import logger +from astrbot.core.db.vec_db.faiss_impl import FaissVecDB +from astrbot.core.knowledge_base.models import ( + BaseKBModel, + KBDocument, + KBMedia, + KnowledgeBase, +) + + +class KBSQLiteDatabase: + def __init__(self, db_path: str = "data/knowledge_base/kb.db") -> None: + """初始化知识库数据库 + + Args: + db_path: 数据库文件路径, 默认为 data/knowledge_base/kb.db + + """ + self.db_path = db_path + self.DATABASE_URL = f"sqlite+aiosqlite:///{db_path}" + self.inited = False + + # 确保目录存在 + Path(db_path).parent.mkdir(parents=True, exist_ok=True) + + # 创建异步引擎 + self.engine = create_async_engine( + self.DATABASE_URL, + echo=False, + pool_pre_ping=True, + pool_recycle=3600, + ) + + # 创建会话工厂 + self.async_session = async_sessionmaker( + self.engine, + class_=AsyncSession, + expire_on_commit=False, + ) + + @asynccontextmanager + async def get_db(self): + """获取数据库会话 + + 用法: + async with kb_db.get_db() as session: + # 执行数据库操作 + result = await session.execute(stmt) + """ + async with self.async_session() as session: + yield session + + async def initialize(self) -> None: + """初始化数据库,创建表并配置 SQLite 参数""" + async with self.engine.begin() as conn: + # 创建所有知识库相关表 + await conn.run_sync(BaseKBModel.metadata.create_all) + + # 配置 SQLite 性能优化参数 + await conn.execute(text("PRAGMA journal_mode=WAL")) + await conn.execute(text("PRAGMA synchronous=NORMAL")) + await conn.execute(text("PRAGMA cache_size=20000")) + await conn.execute(text("PRAGMA temp_store=MEMORY")) + await conn.execute(text("PRAGMA mmap_size=134217728")) + await conn.execute(text("PRAGMA optimize")) + await conn.commit() + + self.inited = True + + async def migrate_to_v1(self) -> None: + """执行知识库数据库 v1 迁移 + + 创建所有必要的索引以优化查询性能 + """ + async with self.get_db() as session: + session: AsyncSession + async with session.begin(): + # 创建知识库表索引 + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_kb_kb_id " + "ON knowledge_bases(kb_id)", + ), + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_kb_name " + "ON knowledge_bases(kb_name)", + ), + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_kb_created_at " + "ON knowledge_bases(created_at)", + ), + ) + + # 创建文档表索引 + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_doc_doc_id " + "ON kb_documents(doc_id)", + ), + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_doc_kb_id " + "ON kb_documents(kb_id)", + ), + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_doc_name " + "ON kb_documents(doc_name)", + ), + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_doc_type " + "ON kb_documents(file_type)", + ), + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_doc_created_at " + "ON kb_documents(created_at)", + ), + ) + + # 创建多媒体表索引 + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_media_media_id " + "ON kb_media(media_id)", + ), + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_media_doc_id " + "ON kb_media(doc_id)", + ), + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_media_kb_id ON kb_media(kb_id)", + ), + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_media_type " + "ON kb_media(media_type)", + ), + ) + + await session.commit() + + async def close(self) -> None: + """关闭数据库连接""" + await self.engine.dispose() + logger.info(f"知识库数据库已关闭: {self.db_path}") + + async def get_kb_by_id(self, kb_id: str) -> KnowledgeBase | None: + """根据 ID 获取知识库""" + async with self.get_db() as session: + stmt = select(KnowledgeBase).where(col(KnowledgeBase.kb_id) == kb_id) + result = await session.execute(stmt) + return result.scalar_one_or_none() + + async def get_kb_by_name(self, kb_name: str) -> KnowledgeBase | None: + """根据名称获取知识库""" + async with self.get_db() as session: + stmt = select(KnowledgeBase).where(col(KnowledgeBase.kb_name) == kb_name) + result = await session.execute(stmt) + return result.scalar_one_or_none() + + async def list_kbs(self, offset: int = 0, limit: int = 100) -> list[KnowledgeBase]: + """列出所有知识库""" + async with self.get_db() as session: + stmt = ( + select(KnowledgeBase) + .offset(offset) + .limit(limit) + .order_by(desc(KnowledgeBase.created_at)) + ) + result = await session.execute(stmt) + return list(result.scalars().all()) + + async def count_kbs(self) -> int: + """统计知识库数量""" + async with self.get_db() as session: + stmt = select(func.count(col(KnowledgeBase.id))) + result = await session.execute(stmt) + return result.scalar() or 0 + + # ===== 文档查询 ===== + + async def get_document_by_id(self, doc_id: str) -> KBDocument | None: + """根据 ID 获取文档""" + async with self.get_db() as session: + stmt = select(KBDocument).where(col(KBDocument.doc_id) == doc_id) + result = await session.execute(stmt) + return result.scalar_one_or_none() + + async def list_documents_by_kb( + self, + kb_id: str, + offset: int = 0, + limit: int = 100, + ) -> list[KBDocument]: + """列出知识库的所有文档""" + async with self.get_db() as session: + stmt = ( + select(KBDocument) + .where(col(KBDocument.kb_id) == kb_id) + .offset(offset) + .limit(limit) + .order_by(desc(KBDocument.created_at)) + ) + result = await session.execute(stmt) + return list(result.scalars().all()) + + async def count_documents_by_kb(self, kb_id: str) -> int: + """统计知识库的文档数量""" + async with self.get_db() as session: + stmt = select(func.count(col(KBDocument.id))).where( + col(KBDocument.kb_id) == kb_id, + ) + result = await session.execute(stmt) + return result.scalar() or 0 + + async def get_document_with_metadata(self, doc_id: str) -> dict | None: + async with self.get_db() as session: + stmt = ( + select(KBDocument, KnowledgeBase) + .join(KnowledgeBase, col(KBDocument.kb_id) == col(KnowledgeBase.kb_id)) + .where(col(KBDocument.doc_id) == doc_id) + ) + result = await session.execute(stmt) + row = result.first() + + if not row: + return None + + return { + "document": row[0], + "knowledge_base": row[1], + } + + async def delete_document_by_id(self, doc_id: str, vec_db: FaissVecDB): + """删除单个文档及其相关数据""" + # 在知识库表中删除 + async with self.get_db() as session, session.begin(): + # 删除文档记录 + delete_stmt = delete(KBDocument).where(col(KBDocument.doc_id) == doc_id) + await session.execute(delete_stmt) + await session.commit() + + # 在 vec db 中删除相关向量 + await vec_db.delete_documents(metadata_filters={"kb_doc_id": doc_id}) + + # ===== 多媒体查询 ===== + + async def list_media_by_doc(self, doc_id: str) -> list[KBMedia]: + """列出文档的所有多媒体资源""" + async with self.get_db() as session: + stmt = select(KBMedia).where(col(KBMedia.doc_id) == doc_id) + result = await session.execute(stmt) + return list(result.scalars().all()) + + async def get_media_by_id(self, media_id: str) -> KBMedia | None: + """根据 ID 获取多媒体资源""" + async with self.get_db() as session: + stmt = select(KBMedia).where(col(KBMedia.media_id) == media_id) + result = await session.execute(stmt) + return result.scalar_one_or_none() + + async def update_kb_stats(self, kb_id: str, vec_db: FaissVecDB) -> None: + """更新知识库统计信息""" + chunk_cnt = await vec_db.count_documents() + + async with self.get_db() as session, session.begin(): + update_stmt = ( + update(KnowledgeBase) + .where(col(KnowledgeBase.kb_id) == kb_id) + .values( + doc_count=select(func.count(col(KBDocument.id))) + .where(col(KBDocument.kb_id) == kb_id) + .scalar_subquery(), + chunk_count=chunk_cnt, + ) + ) + + await session.execute(update_stmt) + await session.commit() diff --git a/astrbot/core/knowledge_base/kb_helper.py b/astrbot/core/knowledge_base/kb_helper.py new file mode 100644 index 00000000..4adfb60b --- /dev/null +++ b/astrbot/core/knowledge_base/kb_helper.py @@ -0,0 +1,642 @@ +import asyncio +import json +import re +import time +import uuid +from pathlib import Path + +import aiofiles + +from astrbot.core import logger +from astrbot.core.db.vec_db.base import BaseVecDB +from astrbot.core.db.vec_db.faiss_impl.vec_db import FaissVecDB +from astrbot.core.provider.manager import ProviderManager +from astrbot.core.provider.provider import ( + EmbeddingProvider, + RerankProvider, +) +from astrbot.core.provider.provider import ( + Provider as LLMProvider, +) + +from .chunking.base import BaseChunker +from .chunking.recursive import RecursiveCharacterChunker +from .kb_db_sqlite import KBSQLiteDatabase +from .models import KBDocument, KBMedia, KnowledgeBase +from .parsers.url_parser import extract_text_from_url +from .parsers.util import select_parser +from .prompts import TEXT_REPAIR_SYSTEM_PROMPT + + +class RateLimiter: + """一个简单的速率限制器""" + + def __init__(self, max_rpm: int): + self.max_per_minute = max_rpm + self.interval = 60.0 / max_rpm if max_rpm > 0 else 0 + self.last_call_time = 0 + + async def __aenter__(self): + if self.interval == 0: + return + + now = time.monotonic() + elapsed = now - self.last_call_time + + if elapsed < self.interval: + await asyncio.sleep(self.interval - elapsed) + + self.last_call_time = time.monotonic() + + async def __aexit__(self, exc_type, exc_val, exc_tb): + pass + + +async def _repair_and_translate_chunk_with_retry( + chunk: str, + repair_llm_service: LLMProvider, + rate_limiter: RateLimiter, + max_retries: int = 2, +) -> list[str]: + """ + Repairs, translates, and optionally re-chunks a single text chunk using the small LLM, with rate limiting. + """ + # 为了防止 LLM 上下文污染，在 user_prompt 中也加入明确的指令 + user_prompt = f"""IGNORE ALL PREVIOUS INSTRUCTIONS. Your ONLY task is to process the following text chunk according to the system prompt provided. + +Text chunk to process: +--- +{chunk} +--- +""" + for attempt in range(max_retries + 1): + try: + async with rate_limiter: + response = await repair_llm_service.text_chat( + prompt=user_prompt, system_prompt=TEXT_REPAIR_SYSTEM_PROMPT + ) + + llm_output = response.completion_text + + if "" in llm_output: + return [] # Signal to discard this chunk + + # More robust regex to handle potential LLM formatting errors (spaces, newlines in tags) + matches = re.findall( + r"<\s*repaired_text\s*>\s*(.*?)\s*<\s*/\s*repaired_text\s*>", + llm_output, + re.DOTALL, + ) + + if matches: + # Further cleaning to ensure no empty strings are returned + return [m.strip() for m in matches if m.strip()] + else: + # If no valid tags and not explicitly discarded, discard it to be safe. + return [] + except Exception as e: + logger.warning( + f" - LLM call failed on attempt {attempt + 1}/{max_retries + 1}. Error: {str(e)}" + ) + + logger.error( + f" - Failed to process chunk after {max_retries + 1} attempts. Using original text." + ) + return [chunk] + + +class KBHelper: + vec_db: BaseVecDB + kb: KnowledgeBase + + def __init__( + self, + kb_db: KBSQLiteDatabase, + kb: KnowledgeBase, + provider_manager: ProviderManager, + kb_root_dir: str, + chunker: BaseChunker, + ): + self.kb_db = kb_db + self.kb = kb + self.prov_mgr = provider_manager + self.kb_root_dir = kb_root_dir + self.chunker = chunker + + self.kb_dir = Path(self.kb_root_dir) / self.kb.kb_id + self.kb_medias_dir = Path(self.kb_dir) / "medias" / self.kb.kb_id + self.kb_files_dir = Path(self.kb_dir) / "files" / self.kb.kb_id + + self.kb_medias_dir.mkdir(parents=True, exist_ok=True) + self.kb_files_dir.mkdir(parents=True, exist_ok=True) + + async def initialize(self): + await self._ensure_vec_db() + + async def get_ep(self) -> EmbeddingProvider: + if not self.kb.embedding_provider_id: + raise ValueError(f"知识库 {self.kb.kb_name} 未配置 Embedding Provider") + ep: EmbeddingProvider = await self.prov_mgr.get_provider_by_id( + self.kb.embedding_provider_id, + ) # type: ignore + if not ep: + raise ValueError( + f"无法找到 ID 为 {self.kb.embedding_provider_id} 的 Embedding Provider", + ) + return ep + + async def get_rp(self) -> RerankProvider | None: + if not self.kb.rerank_provider_id: + return None + rp: RerankProvider = await self.prov_mgr.get_provider_by_id( + self.kb.rerank_provider_id, + ) # type: ignore + if not rp: + raise ValueError( + f"无法找到 ID 为 {self.kb.rerank_provider_id} 的 Rerank Provider", + ) + return rp + + async def _ensure_vec_db(self) -> FaissVecDB: + if not self.kb.embedding_provider_id: + raise ValueError(f"知识库 {self.kb.kb_name} 未配置 Embedding Provider") + + ep = await self.get_ep() + rp = await self.get_rp() + + vec_db = FaissVecDB( + doc_store_path=str(self.kb_dir / "doc.db"), + index_store_path=str(self.kb_dir / "index.faiss"), + embedding_provider=ep, + rerank_provider=rp, + ) + await vec_db.initialize() + self.vec_db = vec_db + return vec_db + + async def delete_vec_db(self): + """删除知识库的向量数据库和所有相关文件""" + import shutil + + await self.terminate() + if self.kb_dir.exists(): + shutil.rmtree(self.kb_dir) + + async def terminate(self): + if self.vec_db: + await self.vec_db.close() + + async def upload_document( + self, + file_name: str, + file_content: bytes | None, + file_type: str, + chunk_size: int = 512, + chunk_overlap: int = 50, + batch_size: int = 32, + tasks_limit: int = 3, + max_retries: int = 3, + progress_callback=None, + pre_chunked_text: list[str] | None = None, + ) -> KBDocument: + """上传并处理文档（带原子性保证和失败清理） + + 流程: + 1. 保存原始文件 + 2. 解析文档内容 + 3. 提取多媒体资源 + 4. 分块处理 + 5. 生成向量并存储 + 6. 保存元数据（事务） + 7. 更新统计 + + Args: + progress_callback: 进度回调函数，接收参数 (stage, current, total) + - stage: 当前阶段 ('parsing', 'chunking', 'embedding') + - current: 当前进度 + - total: 总数 + + """ + await self._ensure_vec_db() + doc_id = str(uuid.uuid4()) + media_paths: list[Path] = [] + file_size = 0 + + # file_path = self.kb_files_dir / f"{doc_id}.{file_type}" + # async with aiofiles.open(file_path, "wb") as f: + # await f.write(file_content) + + try: + chunks_text = [] + saved_media = [] + + if pre_chunked_text is not None: + # 如果提供了预分块文本，直接使用 + chunks_text = pre_chunked_text + file_size = sum(len(chunk) for chunk in chunks_text) + logger.info(f"使用预分块文本进行上传，共 {len(chunks_text)} 个块。") + else: + # 否则，执行标准的文件解析和分块流程 + if file_content is None: + raise ValueError( + "当未提供 pre_chunked_text 时，file_content 不能为空。" + ) + + file_size = len(file_content) + + # 阶段1: 解析文档 + if progress_callback: + await progress_callback("parsing", 0, 100) + + parser = await select_parser(f".{file_type}") + parse_result = await parser.parse(file_content, file_name) + text_content = parse_result.text + media_items = parse_result.media + + if progress_callback: + await progress_callback("parsing", 100, 100) + + # 保存媒体文件 + for media_item in media_items: + media = await self._save_media( + doc_id=doc_id, + media_type=media_item.media_type, + file_name=media_item.file_name, + content=media_item.content, + mime_type=media_item.mime_type, + ) + saved_media.append(media) + media_paths.append(Path(media.file_path)) + + # 阶段2: 分块 + if progress_callback: + await progress_callback("chunking", 0, 100) + + chunks_text = await self.chunker.chunk( + text_content, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + ) + contents = [] + metadatas = [] + for idx, chunk_text in enumerate(chunks_text): + contents.append(chunk_text) + metadatas.append( + { + "kb_id": self.kb.kb_id, + "kb_doc_id": doc_id, + "chunk_index": idx, + }, + ) + + if progress_callback: + await progress_callback("chunking", 100, 100) + + # 阶段3: 生成向量（带进度回调） + async def embedding_progress_callback(current, total): + if progress_callback: + await progress_callback("embedding", current, total) + + await self.vec_db.insert_batch( + contents=contents, + metadatas=metadatas, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + progress_callback=embedding_progress_callback, + ) + + # 保存文档的元数据 + doc = KBDocument( + doc_id=doc_id, + kb_id=self.kb.kb_id, + doc_name=file_name, + file_type=file_type, + file_size=file_size, + # file_path=str(file_path), + file_path="", + chunk_count=len(chunks_text), + media_count=0, + ) + async with self.kb_db.get_db() as session: + async with session.begin(): + session.add(doc) + for media in saved_media: + session.add(media) + await session.commit() + + await session.refresh(doc) + + vec_db: FaissVecDB = self.vec_db # type: ignore + await self.kb_db.update_kb_stats(kb_id=self.kb.kb_id, vec_db=vec_db) + await self.refresh_kb() + await self.refresh_document(doc_id) + return doc + except Exception as e: + logger.error(f"上传文档失败: {e}") + # if file_path.exists(): + # file_path.unlink() + + for media_path in media_paths: + try: + if media_path.exists(): + media_path.unlink() + except Exception as me: + logger.warning(f"清理多媒体文件失败 {media_path}: {me}") + + raise e + + async def list_documents( + self, + offset: int = 0, + limit: int = 100, + ) -> list[KBDocument]: + """列出知识库的所有文档""" + docs = await self.kb_db.list_documents_by_kb(self.kb.kb_id, offset, limit) + return docs + + async def get_document(self, doc_id: str) -> KBDocument | None: + """获取单个文档""" + doc = await self.kb_db.get_document_by_id(doc_id) + return doc + + async def delete_document(self, doc_id: str): + """删除单个文档及其相关数据""" + await self.kb_db.delete_document_by_id( + doc_id=doc_id, + vec_db=self.vec_db, # type: ignore + ) + await self.kb_db.update_kb_stats( + kb_id=self.kb.kb_id, + vec_db=self.vec_db, # type: ignore + ) + await self.refresh_kb() + + async def delete_chunk(self, chunk_id: str, doc_id: str): + """删除单个文本块及其相关数据""" + vec_db: FaissVecDB = self.vec_db # type: ignore + await vec_db.delete(chunk_id) + await self.kb_db.update_kb_stats( + kb_id=self.kb.kb_id, + vec_db=self.vec_db, # type: ignore + ) + await self.refresh_kb() + await self.refresh_document(doc_id) + + async def refresh_kb(self): + if self.kb: + kb = await self.kb_db.get_kb_by_id(self.kb.kb_id) + if kb: + self.kb = kb + + async def refresh_document(self, doc_id: str) -> None: + """更新文档的元数据""" + doc = await self.get_document(doc_id) + if not doc: + raise ValueError(f"无法找到 ID 为 {doc_id} 的文档") + chunk_count = await self.get_chunk_count_by_doc_id(doc_id) + doc.chunk_count = chunk_count + async with self.kb_db.get_db() as session: + async with session.begin(): + session.add(doc) + await session.commit() + await session.refresh(doc) + + async def get_chunks_by_doc_id( + self, + doc_id: str, + offset: int = 0, + limit: int = 100, + ) -> list[dict]: + """获取文档的所有块及其元数据""" + vec_db: FaissVecDB = self.vec_db # type: ignore + chunks = await vec_db.document_storage.get_documents( + metadata_filters={"kb_doc_id": doc_id}, + offset=offset, + limit=limit, + ) + result = [] + for chunk in chunks: + chunk_md = json.loads(chunk["metadata"]) + result.append( + { + "chunk_id": chunk["doc_id"], + "doc_id": chunk_md["kb_doc_id"], + "kb_id": chunk_md["kb_id"], + "chunk_index": chunk_md["chunk_index"], + "content": chunk["text"], + "char_count": len(chunk["text"]), + }, + ) + return result + + async def get_chunk_count_by_doc_id(self, doc_id: str) -> int: + """获取文档的块数量""" + vec_db: FaissVecDB = self.vec_db # type: ignore + count = await vec_db.count_documents(metadata_filter={"kb_doc_id": doc_id}) + return count + + async def _save_media( + self, + doc_id: str, + media_type: str, + file_name: str, + content: bytes, + mime_type: str, + ) -> KBMedia: + """保存多媒体资源""" + media_id = str(uuid.uuid4()) + ext = Path(file_name).suffix + + # 保存文件 + file_path = self.kb_medias_dir / doc_id / f"{media_id}{ext}" + file_path.parent.mkdir(parents=True, exist_ok=True) + async with aiofiles.open(file_path, "wb") as f: + await f.write(content) + + media = KBMedia( + media_id=media_id, + doc_id=doc_id, + kb_id=self.kb.kb_id, + media_type=media_type, + file_name=file_name, + file_path=str(file_path), + file_size=len(content), + mime_type=mime_type, + ) + + return media + + async def upload_from_url( + self, + url: str, + chunk_size: int = 512, + chunk_overlap: int = 50, + batch_size: int = 32, + tasks_limit: int = 3, + max_retries: int = 3, + progress_callback=None, + enable_cleaning: bool = False, + cleaning_provider_id: str | None = None, + ) -> KBDocument: + """从 URL 上传并处理文档（带原子性保证和失败清理） + Args: + url: 要提取内容的网页 URL + chunk_size: 文本块大小 + chunk_overlap: 文本块重叠大小 + batch_size: 批处理大小 + tasks_limit: 并发任务限制 + max_retries: 最大重试次数 + progress_callback: 进度回调函数，接收参数 (stage, current, total) + - stage: 当前阶段 ('extracting', 'cleaning', 'parsing', 'chunking', 'embedding') + - current: 当前进度 + - total: 总数 + Returns: + KBDocument: 上传的文档对象 + Raises: + ValueError: 如果 URL 为空或无法提取内容 + IOError: 如果网络请求失败 + """ + # 获取 Tavily API 密钥 + config = self.prov_mgr.acm.default_conf + tavily_keys = config.get("provider_settings", {}).get( + "websearch_tavily_key", [] + ) + if not tavily_keys: + raise ValueError( + "Error: Tavily API key is not configured in provider_settings." + ) + + # 阶段1: 从 URL 提取内容 + if progress_callback: + await progress_callback("extracting", 0, 100) + + try: + text_content = await extract_text_from_url(url, tavily_keys) + except Exception as e: + logger.error(f"Failed to extract content from URL {url}: {e}") + raise OSError(f"Failed to extract content from URL {url}: {e}") from e + + if not text_content: + raise ValueError(f"No content extracted from URL: {url}") + + if progress_callback: + await progress_callback("extracting", 100, 100) + + # 阶段2: (可选)清洗内容并分块 + final_chunks = await self._clean_and_rechunk_content( + content=text_content, + url=url, + progress_callback=progress_callback, + enable_cleaning=enable_cleaning, + cleaning_provider_id=cleaning_provider_id, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + ) + + if enable_cleaning and not final_chunks: + raise ValueError( + "内容清洗后未提取到有效文本。请尝试关闭内容清洗功能，或更换更高性能的LLM模型后重试。" + ) + + # 创建一个虚拟文件名 + file_name = url.split("/")[-1] or f"document_from_{url}" + if not Path(file_name).suffix: + file_name += ".url" + + # 复用现有的 upload_document 方法，但传入预分块文本 + return await self.upload_document( + file_name=file_name, + file_content=None, + file_type="url", # 使用 'url' 作为特殊文件类型 + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + progress_callback=progress_callback, + pre_chunked_text=final_chunks, + ) + + async def _clean_and_rechunk_content( + self, + content: str, + url: str, + progress_callback=None, + enable_cleaning: bool = False, + cleaning_provider_id: str | None = None, + repair_max_rpm: int = 60, + chunk_size: int = 512, + chunk_overlap: int = 50, + ) -> list[str]: + """ + 对从 URL 获取的内容进行清洗、修复、翻译和重新分块。 + """ + if not enable_cleaning: + # 如果不启用清洗，则使用从前端传递的参数进行分块 + logger.info( + f"内容清洗未启用，使用指定参数进行分块: chunk_size={chunk_size}, chunk_overlap={chunk_overlap}" + ) + return await self.chunker.chunk( + content, chunk_size=chunk_size, chunk_overlap=chunk_overlap + ) + + if not cleaning_provider_id: + logger.warning( + "启用了内容清洗，但未提供 cleaning_provider_id，跳过清洗并使用默认分块。" + ) + return await self.chunker.chunk(content) + + if progress_callback: + await progress_callback("cleaning", 0, 100) + + try: + # 获取指定的 LLM Provider + llm_provider = await self.prov_mgr.get_provider_by_id(cleaning_provider_id) + if not llm_provider or not isinstance(llm_provider, LLMProvider): + raise ValueError( + f"无法找到 ID 为 {cleaning_provider_id} 的 LLM Provider 或类型不正确" + ) + + # 初步分块 + # 优化分隔符，优先按段落分割，以获得更高质量的文本块 + text_splitter = RecursiveCharacterChunker( + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + separators=["\n\n", "\n", " "], # 优先使用段落分隔符 + ) + initial_chunks = await text_splitter.chunk(content) + logger.info(f"初步分块完成，生成 {len(initial_chunks)} 个块用于修复。") + + # 并发处理所有块 + rate_limiter = RateLimiter(repair_max_rpm) + tasks = [ + _repair_and_translate_chunk_with_retry( + chunk, llm_provider, rate_limiter + ) + for chunk in initial_chunks + ] + + repaired_results = await asyncio.gather(*tasks, return_exceptions=True) + + final_chunks = [] + for i, result in enumerate(repaired_results): + if isinstance(result, Exception): + logger.warning(f"块 {i} 处理异常: {str(result)}. 回退到原始块。") + final_chunks.append(initial_chunks[i]) + elif isinstance(result, list): + final_chunks.extend(result) + + logger.info( + f"文本修复完成: {len(initial_chunks)} 个原始块 -> {len(final_chunks)} 个最终块。" + ) + + if progress_callback: + await progress_callback("cleaning", 100, 100) + + return final_chunks + + except Exception as e: + logger.error(f"使用 Provider '{cleaning_provider_id}' 清洗内容失败: {e}") + # 清洗失败，返回默认分块结果，保证流程不中断 + return await self.chunker.chunk(content) diff --git a/astrbot/core/knowledge_base/kb_mgr.py b/astrbot/core/knowledge_base/kb_mgr.py new file mode 100644 index 00000000..2219cc00 --- /dev/null +++ b/astrbot/core/knowledge_base/kb_mgr.py @@ -0,0 +1,330 @@ +import traceback +from pathlib import Path + +from astrbot.core import logger +from astrbot.core.provider.manager import ProviderManager + +# from .chunking.fixed_size import FixedSizeChunker +from .chunking.recursive import RecursiveCharacterChunker +from .kb_db_sqlite import KBSQLiteDatabase +from .kb_helper import KBHelper +from .models import KBDocument, KnowledgeBase +from .retrieval.manager import RetrievalManager, RetrievalResult +from .retrieval.rank_fusion import RankFusion +from .retrieval.sparse_retriever import SparseRetriever + +FILES_PATH = "data/knowledge_base" +DB_PATH = Path(FILES_PATH) / "kb.db" +"""Knowledge Base storage root directory""" +CHUNKER = RecursiveCharacterChunker() + + +class KnowledgeBaseManager: + kb_db: KBSQLiteDatabase + retrieval_manager: RetrievalManager + + def __init__( + self, + provider_manager: ProviderManager, + ): + Path(DB_PATH).parent.mkdir(parents=True, exist_ok=True) + self.provider_manager = provider_manager + self._session_deleted_callback_registered = False + + self.kb_insts: dict[str, KBHelper] = {} + + async def initialize(self): + """初始化知识库模块""" + try: + logger.info("正在初始化知识库模块...") + + # 初始化数据库 + await self._init_kb_database() + + # 初始化检索管理器 + sparse_retriever = SparseRetriever(self.kb_db) + rank_fusion = RankFusion(self.kb_db) + self.retrieval_manager = RetrievalManager( + sparse_retriever=sparse_retriever, + rank_fusion=rank_fusion, + kb_db=self.kb_db, + ) + await self.load_kbs() + + except ImportError as e: + logger.error(f"知识库模块导入失败: {e}") + logger.warning("请确保已安装所需依赖: pypdf, aiofiles, Pillow, rank-bm25") + except Exception as e: + logger.error(f"知识库模块初始化失败: {e}") + logger.error(traceback.format_exc()) + + async def _init_kb_database(self): + self.kb_db = KBSQLiteDatabase(DB_PATH.as_posix()) + await self.kb_db.initialize() + await self.kb_db.migrate_to_v1() + logger.info(f"KnowledgeBase database initialized: {DB_PATH}") + + async def load_kbs(self): + """加载所有知识库实例""" + kb_records = await self.kb_db.list_kbs() + for record in kb_records: + kb_helper = KBHelper( + kb_db=self.kb_db, + kb=record, + provider_manager=self.provider_manager, + kb_root_dir=FILES_PATH, + chunker=CHUNKER, + ) + await kb_helper.initialize() + self.kb_insts[record.kb_id] = kb_helper + + async def create_kb( + self, + kb_name: str, + description: str | None = None, + emoji: str | None = None, + embedding_provider_id: str | None = None, + rerank_provider_id: str | None = None, + chunk_size: int | None = None, + chunk_overlap: int | None = None, + top_k_dense: int | None = None, + top_k_sparse: int | None = None, + top_m_final: int | None = None, + ) -> KBHelper: + """创建新的知识库实例""" + kb = KnowledgeBase( + kb_name=kb_name, + description=description, + emoji=emoji or "📚", + embedding_provider_id=embedding_provider_id, + rerank_provider_id=rerank_provider_id, + chunk_size=chunk_size if chunk_size is not None else 512, + chunk_overlap=chunk_overlap if chunk_overlap is not None else 50, + top_k_dense=top_k_dense if top_k_dense is not None else 50, + top_k_sparse=top_k_sparse if top_k_sparse is not None else 50, + top_m_final=top_m_final if top_m_final is not None else 5, + ) + async with self.kb_db.get_db() as session: + session.add(kb) + await session.commit() + await session.refresh(kb) + + kb_helper = KBHelper( + kb_db=self.kb_db, + kb=kb, + provider_manager=self.provider_manager, + kb_root_dir=FILES_PATH, + chunker=CHUNKER, + ) + await kb_helper.initialize() + self.kb_insts[kb.kb_id] = kb_helper + return kb_helper + + async def get_kb(self, kb_id: str) -> KBHelper | None: + """获取知识库实例""" + if kb_id in self.kb_insts: + return self.kb_insts[kb_id] + + async def get_kb_by_name(self, kb_name: str) -> KBHelper | None: + """通过名称获取知识库实例""" + for kb_helper in self.kb_insts.values(): + if kb_helper.kb.kb_name == kb_name: + return kb_helper + return None + + async def delete_kb(self, kb_id: str) -> bool: + """删除知识库实例""" + kb_helper = await self.get_kb(kb_id) + if not kb_helper: + return False + + await kb_helper.delete_vec_db() + async with self.kb_db.get_db() as session: + await session.delete(kb_helper.kb) + await session.commit() + + self.kb_insts.pop(kb_id, None) + return True + + async def list_kbs(self) -> list[KnowledgeBase]: + """列出所有知识库实例""" + kbs = [kb_helper.kb for kb_helper in self.kb_insts.values()] + return kbs + + async def update_kb( + self, + kb_id: str, + kb_name: str, + description: str | None = None, + emoji: str | None = None, + embedding_provider_id: str | None = None, + rerank_provider_id: str | None = None, + chunk_size: int | None = None, + chunk_overlap: int | None = None, + top_k_dense: int | None = None, + top_k_sparse: int | None = None, + top_m_final: int | None = None, + ) -> KBHelper | None: + """更新知识库实例""" + kb_helper = await self.get_kb(kb_id) + if not kb_helper: + return None + + kb = kb_helper.kb + if kb_name is not None: + kb.kb_name = kb_name + if description is not None: + kb.description = description + if emoji is not None: + kb.emoji = emoji + if embedding_provider_id is not None: + kb.embedding_provider_id = embedding_provider_id + kb.rerank_provider_id = rerank_provider_id # 允许设置为 None + if chunk_size is not None: + kb.chunk_size = chunk_size + if chunk_overlap is not None: + kb.chunk_overlap = chunk_overlap + if top_k_dense is not None: + kb.top_k_dense = top_k_dense + if top_k_sparse is not None: + kb.top_k_sparse = top_k_sparse + if top_m_final is not None: + kb.top_m_final = top_m_final + async with self.kb_db.get_db() as session: + session.add(kb) + await session.commit() + await session.refresh(kb) + + return kb_helper + + async def retrieve( + self, + query: str, + kb_names: list[str], + top_k_fusion: int = 20, + top_m_final: int = 5, + ) -> dict | None: + """从指定知识库中检索相关内容""" + kb_ids = [] + kb_id_helper_map = {} + for kb_name in kb_names: + if kb_helper := await self.get_kb_by_name(kb_name): + kb_ids.append(kb_helper.kb.kb_id) + kb_id_helper_map[kb_helper.kb.kb_id] = kb_helper + + if not kb_ids: + return {} + + results = await self.retrieval_manager.retrieve( + query=query, + kb_ids=kb_ids, + kb_id_helper_map=kb_id_helper_map, + top_k_fusion=top_k_fusion, + top_m_final=top_m_final, + ) + if not results: + return None + + context_text = self._format_context(results) + + results_dict = [ + { + "chunk_id": r.chunk_id, + "doc_id": r.doc_id, + "kb_id": r.kb_id, + "kb_name": r.kb_name, + "doc_name": r.doc_name, + "chunk_index": r.metadata.get("chunk_index", 0), + "content": r.content, + "score": r.score, + "char_count": r.metadata.get("char_count", 0), + } + for r in results + ] + + return { + "context_text": context_text, + "results": results_dict, + } + + def _format_context(self, results: list[RetrievalResult]) -> str: + """格式化知识上下文 + + Args: + results: 检索结果列表 + + Returns: + str: 格式化的上下文文本 + + """ + lines = ["以下是相关的知识库内容,请参考这些信息回答用户的问题:\n"] + + for i, result in enumerate(results, 1): + lines.append(f"【知识 {i}】") + lines.append(f"来源: {result.kb_name} / {result.doc_name}") + lines.append(f"内容: {result.content}") + lines.append(f"相关度: {result.score:.2f}") + lines.append("") + + return "\n".join(lines) + + async def terminate(self): + """终止所有知识库实例,关闭数据库连接""" + for kb_id, kb_helper in self.kb_insts.items(): + try: + await kb_helper.terminate() + except Exception as e: + logger.error(f"关闭知识库 {kb_id} 失败: {e}") + + self.kb_insts.clear() + + # 关闭元数据数据库 + if hasattr(self, "kb_db") and self.kb_db: + try: + await self.kb_db.close() + except Exception as e: + logger.error(f"关闭知识库元数据数据库失败: {e}") + + async def upload_from_url( + self, + kb_id: str, + url: str, + chunk_size: int = 512, + chunk_overlap: int = 50, + batch_size: int = 32, + tasks_limit: int = 3, + max_retries: int = 3, + progress_callback=None, + ) -> KBDocument: + """从 URL 上传文档到指定的知识库 + + Args: + kb_id: 知识库 ID + url: 要提取内容的网页 URL + chunk_size: 文本块大小 + chunk_overlap: 文本块重叠大小 + batch_size: 批处理大小 + tasks_limit: 并发任务限制 + max_retries: 最大重试次数 + progress_callback: 进度回调函数 + + Returns: + KBDocument: 上传的文档对象 + + Raises: + ValueError: 如果知识库不存在或 URL 为空 + IOError: 如果网络请求失败 + """ + kb_helper = await self.get_kb(kb_id) + if not kb_helper: + raise ValueError(f"Knowledge base with id {kb_id} not found.") + + return await kb_helper.upload_from_url( + url=url, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + progress_callback=progress_callback, + ) diff --git a/astrbot/core/knowledge_base/models.py b/astrbot/core/knowledge_base/models.py new file mode 100644 index 00000000..da919a38 --- /dev/null +++ b/astrbot/core/knowledge_base/models.py @@ -0,0 +1,120 @@ +import uuid +from datetime import datetime, timezone + +from sqlmodel import Field, MetaData, SQLModel, Text, UniqueConstraint + + +class BaseKBModel(SQLModel, table=False): + metadata = MetaData() + + +class KnowledgeBase(BaseKBModel, table=True): + """知识库表 + + 存储知识库的基本信息和统计数据。 + """ + + __tablename__ = "knowledge_bases" # type: ignore + + id: int | None = Field( + primary_key=True, + sa_column_kwargs={"autoincrement": True}, + default=None, + ) + kb_id: str = Field( + max_length=36, + nullable=False, + unique=True, + default_factory=lambda: str(uuid.uuid4()), + index=True, + ) + kb_name: str = Field(max_length=100, nullable=False) + description: str | None = Field(default=None, sa_type=Text) + emoji: str | None = Field(default="📚", max_length=10) + embedding_provider_id: str | None = Field(default=None, max_length=100) + rerank_provider_id: str | None = Field(default=None, max_length=100) + # 分块配置参数 + chunk_size: int | None = Field(default=512, nullable=True) + chunk_overlap: int | None = Field(default=50, nullable=True) + # 检索配置参数 + top_k_dense: int | None = Field(default=50, nullable=True) + top_k_sparse: int | None = Field(default=50, nullable=True) + top_m_final: int | None = Field(default=5, nullable=True) + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + updated_at: datetime = Field( + default_factory=lambda: datetime.now(timezone.utc), + sa_column_kwargs={"onupdate": datetime.now(timezone.utc)}, + ) + doc_count: int = Field(default=0, nullable=False) + chunk_count: int = Field(default=0, nullable=False) + + __table_args__ = ( + UniqueConstraint( + "kb_name", + name="uix_kb_name", + ), + ) + + +class KBDocument(BaseKBModel, table=True): + """文档表 + + 存储上传到知识库的文档元数据。 + """ + + __tablename__ = "kb_documents" # type: ignore + + id: int | None = Field( + primary_key=True, + sa_column_kwargs={"autoincrement": True}, + default=None, + ) + doc_id: str = Field( + max_length=36, + nullable=False, + unique=True, + default_factory=lambda: str(uuid.uuid4()), + index=True, + ) + kb_id: str = Field(max_length=36, nullable=False, index=True) + doc_name: str = Field(max_length=255, nullable=False) + file_type: str = Field(max_length=20, nullable=False) + file_size: int = Field(nullable=False) + file_path: str = Field(max_length=512, nullable=False) + chunk_count: int = Field(default=0, nullable=False) + media_count: int = Field(default=0, nullable=False) + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + updated_at: datetime = Field( + default_factory=lambda: datetime.now(timezone.utc), + sa_column_kwargs={"onupdate": datetime.now(timezone.utc)}, + ) + + +class KBMedia(BaseKBModel, table=True): + """多媒体资源表 + + 存储从文档中提取的图片、视频等多媒体资源。 + """ + + __tablename__ = "kb_media" # type: ignore + + id: int | None = Field( + primary_key=True, + sa_column_kwargs={"autoincrement": True}, + default=None, + ) + media_id: str = Field( + max_length=36, + nullable=False, + unique=True, + default_factory=lambda: str(uuid.uuid4()), + index=True, + ) + doc_id: str = Field(max_length=36, nullable=False, index=True) + kb_id: str = Field(max_length=36, nullable=False, index=True) + media_type: str = Field(max_length=20, nullable=False) + file_name: str = Field(max_length=255, nullable=False) + file_path: str = Field(max_length=512, nullable=False) + file_size: int = Field(nullable=False) + mime_type: str = Field(max_length=100, nullable=False) + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) diff --git a/astrbot/core/knowledge_base/parsers/__init__.py b/astrbot/core/knowledge_base/parsers/__init__.py new file mode 100644 index 00000000..184f2fd4 --- /dev/null +++ b/astrbot/core/knowledge_base/parsers/__init__.py @@ -0,0 +1,13 @@ +"""文档解析器模块""" + +from .base import BaseParser, MediaItem, ParseResult +from .pdf_parser import PDFParser +from .text_parser import TextParser + +__all__ = [ + "BaseParser", + "MediaItem", + "PDFParser", + "ParseResult", + "TextParser", +] diff --git a/astrbot/core/knowledge_base/parsers/base.py b/astrbot/core/knowledge_base/parsers/base.py new file mode 100644 index 00000000..4ffca9c6 --- /dev/null +++ b/astrbot/core/knowledge_base/parsers/base.py @@ -0,0 +1,51 @@ +"""文档解析器基类和数据结构 + +定义了文档解析器的抽象接口和相关数据类。 +""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass + + +@dataclass +class MediaItem: + """多媒体项 + + 表示从文档中提取的多媒体资源。 + """ + + media_type: str # image, video + file_name: str + content: bytes + mime_type: str + + +@dataclass +class ParseResult: + """解析结果 + + 包含解析后的文本内容和提取的多媒体资源。 + """ + + text: str + media: list[MediaItem] + + +class BaseParser(ABC): + """文档解析器基类 + + 所有文档解析器都应该继承此类并实现 parse 方法。 + """ + + @abstractmethod + async def parse(self, file_content: bytes, file_name: str) -> ParseResult: + """解析文档 + + Args: + file_content: 文件内容 + file_name: 文件名 + + Returns: + ParseResult: 解析结果 + + """ diff --git a/astrbot/core/knowledge_base/parsers/markitdown_parser.py b/astrbot/core/knowledge_base/parsers/markitdown_parser.py new file mode 100644 index 00000000..9ef34793 --- /dev/null +++ b/astrbot/core/knowledge_base/parsers/markitdown_parser.py @@ -0,0 +1,26 @@ +import io +import os + +from markitdown_no_magika import MarkItDown, StreamInfo + +from astrbot.core.knowledge_base.parsers.base import ( + BaseParser, + ParseResult, +) + + +class MarkitdownParser(BaseParser): + """解析 docx, xls, xlsx 格式""" + + async def parse(self, file_content: bytes, file_name: str) -> ParseResult: + md = MarkItDown(enable_plugins=False) + bio = io.BytesIO(file_content) + stream_info = StreamInfo( + extension=os.path.splitext(file_name)[1].lower(), + filename=file_name, + ) + result = md.convert(bio, stream_info=stream_info) + return ParseResult( + text=result.markdown, + media=[], + ) diff --git a/astrbot/core/knowledge_base/parsers/pdf_parser.py b/astrbot/core/knowledge_base/parsers/pdf_parser.py new file mode 100644 index 00000000..aeeea930 --- /dev/null +++ b/astrbot/core/knowledge_base/parsers/pdf_parser.py @@ -0,0 +1,101 @@ +"""PDF 文件解析器 + +支持解析 PDF 文件中的文本和图片资源。 +""" + +import io + +from pypdf import PdfReader + +from astrbot.core.knowledge_base.parsers.base import ( + BaseParser, + MediaItem, + ParseResult, +) + + +class PDFParser(BaseParser): + """PDF 文档解析器 + + 提取 PDF 中的文本内容和嵌入的图片资源。 + """ + + async def parse(self, file_content: bytes, file_name: str) -> ParseResult: + """解析 PDF 文件 + + Args: + file_content: 文件内容 + file_name: 文件名 + + Returns: + ParseResult: 包含文本和图片的解析结果 + + """ + pdf_file = io.BytesIO(file_content) + reader = PdfReader(pdf_file) + + text_parts = [] + media_items = [] + + # 提取文本 + for page in reader.pages: + text = page.extract_text() + if text: + text_parts.append(text) + + # 提取图片 + image_counter = 0 + for page_num, page in enumerate(reader.pages): + try: + # 安全检查 Resources + if "/Resources" not in page: + continue + + resources = page["/Resources"] + if not resources or "/XObject" not in resources: # type: ignore + continue + + xobjects = resources["/XObject"].get_object() # type: ignore + if not xobjects: + continue + + for obj_name in xobjects: + try: + obj = xobjects[obj_name] + + if obj.get("/Subtype") != "/Image": + continue + + # 提取图片数据 + image_data = obj.get_data() + + # 确定格式 + filter_type = obj.get("/Filter", "") + if filter_type == "/DCTDecode": + ext = "jpg" + mime_type = "image/jpeg" + elif filter_type == "/FlateDecode": + ext = "png" + mime_type = "image/png" + else: + ext = "png" + mime_type = "image/png" + + image_counter += 1 + media_items.append( + MediaItem( + media_type="image", + file_name=f"page_{page_num}_img_{image_counter}.{ext}", + content=image_data, + mime_type=mime_type, + ), + ) + except Exception: + # 单个图片提取失败不影响整体 + continue + except Exception: + # 页面处理失败不影响其他页面 + continue + + full_text = "\n\n".join(text_parts) + return ParseResult(text=full_text, media=media_items) diff --git a/astrbot/core/knowledge_base/parsers/text_parser.py b/astrbot/core/knowledge_base/parsers/text_parser.py new file mode 100644 index 00000000..bed2d09b --- /dev/null +++ b/astrbot/core/knowledge_base/parsers/text_parser.py @@ -0,0 +1,42 @@ +"""文本文件解析器 + +支持解析 TXT 和 Markdown 文件。 +""" + +from astrbot.core.knowledge_base.parsers.base import BaseParser, ParseResult + + +class TextParser(BaseParser): + """TXT/MD 文本解析器 + + 支持多种字符编码的自动检测。 + """ + + async def parse(self, file_content: bytes, file_name: str) -> ParseResult: + """解析文本文件 + + 尝试使用多种编码解析文件内容。 + + Args: + file_content: 文件内容 + file_name: 文件名 + + Returns: + ParseResult: 解析结果,不包含多媒体资源 + + Raises: + ValueError: 如果无法解码文件 + + """ + # 尝试多种编码 + for encoding in ["utf-8", "gbk", "gb2312", "gb18030"]: + try: + text = file_content.decode(encoding) + break + except UnicodeDecodeError: + continue + else: + raise ValueError(f"无法解码文件: {file_name}") + + # 文本文件无多媒体资源 + return ParseResult(text=text, media=[]) diff --git a/astrbot/core/knowledge_base/parsers/url_parser.py b/astrbot/core/knowledge_base/parsers/url_parser.py new file mode 100644 index 00000000..f68e2e0c --- /dev/null +++ b/astrbot/core/knowledge_base/parsers/url_parser.py @@ -0,0 +1,103 @@ +import asyncio + +import aiohttp + + +class URLExtractor: + """URL 内容提取器，封装了 Tavily API 调用和密钥管理""" + + def __init__(self, tavily_keys: list[str]): + """ + 初始化 URL 提取器 + + Args: + tavily_keys: Tavily API 密钥列表 + """ + if not tavily_keys: + raise ValueError("Error: Tavily API keys are not configured.") + + self.tavily_keys = tavily_keys + self.tavily_key_index = 0 + self.tavily_key_lock = asyncio.Lock() + + async def _get_tavily_key(self) -> str: + """并发安全的从列表中获取并轮换Tavily API密钥。""" + async with self.tavily_key_lock: + key = self.tavily_keys[self.tavily_key_index] + self.tavily_key_index = (self.tavily_key_index + 1) % len(self.tavily_keys) + return key + + async def extract_text_from_url(self, url: str) -> str: + """ + 使用 Tavily API 从 URL 提取主要文本内容。 + 这是 web_searcher 插件中 tavily_extract_web_page 方法的简化版本， + 专门为知识库模块设计，不依赖 AstrMessageEvent。 + + Args: + url: 要提取内容的网页 URL + + Returns: + 提取的文本内容 + + Raises: + ValueError: 如果 URL 为空或 API 密钥未配置 + IOError: 如果请求失败或返回错误 + """ + if not url: + raise ValueError("Error: url must be a non-empty string.") + + tavily_key = await self._get_tavily_key() + api_url = "https://api.tavily.com/extract" + headers = { + "Authorization": f"Bearer {tavily_key}", + "Content-Type": "application/json", + } + + payload = { + "urls": [url], + "extract_depth": "basic", # 使用基础提取深度 + } + + try: + async with aiohttp.ClientSession(trust_env=True) as session: + async with session.post( + api_url, + json=payload, + headers=headers, + timeout=30.0, # 增加超时时间，因为内容提取可能需要更长时间 + ) as response: + if response.status != 200: + reason = await response.text() + raise OSError( + f"Tavily web extraction failed: {reason}, status: {response.status}" + ) + + data = await response.json() + results = data.get("results", []) + + if not results: + raise ValueError(f"No content extracted from URL: {url}") + + # 返回第一个结果的内容 + return results[0].get("raw_content", "") + + except aiohttp.ClientError as e: + raise OSError(f"Failed to fetch URL {url}: {e}") from e + except Exception as e: + raise OSError(f"Failed to extract content from URL {url}: {e}") from e + + +# 为了向后兼容，提供一个简单的函数接口 +async def extract_text_from_url(url: str, tavily_keys: list[str]) -> str: + """ + 简单的函数接口，用于从 URL 提取文本内容 + + Args: + url: 要提取内容的网页 URL + tavily_keys: Tavily API 密钥列表 + + Returns: + 提取的文本内容 + """ + extractor = URLExtractor(tavily_keys) + return await extractor.extract_text_from_url(url) diff --git a/astrbot/core/knowledge_base/parsers/util.py b/astrbot/core/knowledge_base/parsers/util.py new file mode 100644 index 00000000..7a446320 --- /dev/null +++ b/astrbot/core/knowledge_base/parsers/util.py @@ -0,0 +1,13 @@ +from .base import BaseParser + + +async def select_parser(ext: str) -> BaseParser: + if ext in {".md", ".txt", ".markdown", ".xlsx", ".docx", ".xls"}: + from .markitdown_parser import MarkitdownParser + + return MarkitdownParser() + if ext == ".pdf": + from .pdf_parser import PDFParser + + return PDFParser() + raise ValueError(f"暂时不支持的文件格式: {ext}") diff --git a/astrbot/core/knowledge_base/prompts.py b/astrbot/core/knowledge_base/prompts.py new file mode 100644 index 00000000..7874fa5f --- /dev/null +++ b/astrbot/core/knowledge_base/prompts.py @@ -0,0 +1,65 @@ +TEXT_REPAIR_SYSTEM_PROMPT = """You are a meticulous digital archivist. Your mission is to reconstruct a clean, readable article from raw, noisy text chunks. + +**Core Task:** +1. **Analyze:** Examine the text chunk to separate "signal" (substantive information) from "noise" (UI elements, ads, navigation, footers). +2. **Process:** Clean and repair the signal. **Do not translate it.** Keep the original language. + +**Crucial Rules:** +- **NEVER discard a chunk if it contains ANY valuable information.** Your primary duty is to salvage content. +- **If a chunk contains multiple distinct topics, split them.** Enclose each topic in its own `` tag. +- Your output MUST be ONLY `...` tags or a single `` tag. + +--- +**Example 1: Chunk with Noise and Signal** + +*Input Chunk:* +"Home | About | Products | **The Llama is a domesticated South American camelid.** | © 2025 ACME Corp." + +*Your Thought Process:* +1. "Home | About | Products..." and "© 2025 ACME Corp." are noise. +2. "The Llama is a domesticated..." is the signal. +3. I must extract the signal and wrap it. + +*Your Output:* + +The Llama is a domesticated South American camelid. + + +--- +**Example 2: Chunk with ONLY Noise** + +*Input Chunk:* +"Next Page > | Subscribe to our newsletter | Follow us on X" + +*Your Thought Process:* +1. This entire chunk is noise. There is no signal. +2. I must discard this. + +*Your Output:* + + +--- +**Example 3: Chunk with Multiple Topics (Requires Splitting)** + +*Input Chunk:* +"## Chapter 1: The Sun +The Sun is the star at the center of the Solar System. + +## Chapter 2: The Moon +The Moon is Earth's only natural satellite." + +*Your Thought Process:* +1. This chunk contains two distinct topics. +2. I must process them separately to maintain semantic integrity. +3. I will create two `` blocks. + +*Your Output:* + +## Chapter 1: The Sun +The Sun is the star at the center of the Solar System. + + +## Chapter 2: The Moon +The Moon is Earth's only natural satellite. + +""" diff --git a/astrbot/core/knowledge_base/retrieval/__init__.py b/astrbot/core/knowledge_base/retrieval/__init__.py new file mode 100644 index 00000000..f5d196cb --- /dev/null +++ b/astrbot/core/knowledge_base/retrieval/__init__.py @@ -0,0 +1,14 @@ +"""检索模块""" + +from .manager import RetrievalManager, RetrievalResult +from .rank_fusion import FusedResult, RankFusion +from .sparse_retriever import SparseResult, SparseRetriever + +__all__ = [ + "FusedResult", + "RankFusion", + "RetrievalManager", + "RetrievalResult", + "SparseResult", + "SparseRetriever", +] diff --git a/astrbot/core/knowledge_base/retrieval/hit_stopwords.txt b/astrbot/core/knowledge_base/retrieval/hit_stopwords.txt new file mode 100644 index 00000000..84b26283 --- /dev/null +++ b/astrbot/core/knowledge_base/retrieval/hit_stopwords.txt @@ -0,0 +1,767 @@ +——— +》）， +）÷（１－ +”， +）、 +＝（ +: +→ +℃ +& +* +一一 +~~~~ +’ +. +『 +.一 +./ +-- +』 +＝″ +【 +［＊］ +｝＞ +［⑤］］ +［①Ｄ］ +ｃ］ +ｎｇ昉 +＊ +// +［ +］ +［②ｅ］ +［②ｇ］ +＝｛ +} +，也 +‘ +Ａ +［①⑥］ +［②Ｂ］ +［①ａ］ +［④ａ］ +［①③］ +［③ｈ］ +③］ +１． +－－ +［②ｂ］ +’‘ +××× +［①⑧］ +０：２ +＝［ +［⑤ｂ］ +［②ｃ］ +［④ｂ］ +［②③］ +［③ａ］ +［④ｃ］ +［①⑤］ +［①⑦］ +［①ｇ］ +∈［ +［①⑨］ +［①④］ +［①ｃ］ +［②ｆ］ +［②⑧］ +［②①］ +［①Ｃ］ +［③ｃ］ +［③ｇ］ +［②⑤］ +［②②］ +一. +［①ｈ］ +.数 +［］ +［①Ｂ］ +数/ +［①ｉ］ +［③ｅ］ +［①①］ +［④ｄ］ +［④ｅ］ +［③ｂ］ +［⑤ａ］ +［①Ａ］ +［②⑧］ +［②⑦］ +［①ｄ］ +［②ｊ］ +〕〔 +］［ +:// +′∈ +［②④ +［⑤ｅ］ +１２％ +ｂ］ +... +................... +…………………………………………………③ +ＺＸＦＩＴＬ +［③Ｆ］ +」 +［①ｏ］ +］∧′＝［ +∪φ∈ +′｜ +｛－ +②ｃ +｝ +［③①］ +Ｒ．Ｌ． +［①Ｅ］ +Ψ +－［＊］－ +↑ +.日 +［②ｄ］ +［② +［②⑦］ +［②②］ +［③ｅ］ +［①ｉ］ +［①Ｂ］ +［①ｈ］ +［①ｄ］ +［①ｇ］ +［①②］ +［②ａ］ +ｆ］ +［⑩］ +ａ］ +［①ｅ］ +［②ｈ］ +［②⑥］ +［③ｄ］ +［②⑩］ +ｅ］ +〉 +】 +元／吨 +［②⑩］ +２．３％ +５：０ +［①］ +:: +［②］ +［③］ +［④］ +［⑤］ +［⑥］ +［⑦］ +［⑧］ +［⑨］ +…… +—— +? +、 +。 +“ +” +《 +》 +！ +， +： +； +？ +． +, +． +' +? +· +——— +── +? +— +< +> +（ +） +〔 +〕 +[ +] +( +) +- ++ +～ +× +／ +/ +① +② +③ +④ +⑤ +⑥ +⑦ +⑧ +⑨ +⑩ +Ⅲ +В +" +; +# +@ +γ +μ +φ +φ． +× +Δ +■ +▲ +sub +exp +sup +sub +Lex +＃ +％ +＆ +＇ +＋ +＋ξ +＋＋ +－ +－β +＜ +＜± +＜Δ +＜λ +＜φ +＜＜ += +＝ +＝☆ +＝－ +＞ +＞λ +＿ +～± +～＋ +［⑤ｆ］ +［⑤ｄ］ +［②ｉ］ +≈ +［②Ｇ］ +［①ｆ］ +ＬＩ +㈧ +［－ +...... +〉 +［③⑩］ +第二 +一番 +一直 +一个 +一些 +许多 +种 +有的是 +也就是说 +末##末 +啊 +阿 +哎 +哎呀 +哎哟 +唉 +俺 +俺们 +按 +按照 +吧 +吧哒 +把 +罢了 +被 +本 +本着 +比 +比方 +比如 +鄙人 +彼 +彼此 +边 +别 +别的 +别说 +并 +并且 +不比 +不成 +不单 +不但 +不独 +不管 +不光 +不过 +不仅 +不拘 +不论 +不怕 +不然 +不如 +不特 +不惟 +不问 +不只 +朝 +朝着 +趁 +趁着 +乘 +冲 +除 +除此之外 +除非 +除了 +此 +此间 +此外 +从 +从而 +打 +待 +但 +但是 +当 +当着 +到 +得 +的 +的话 +等 +等等 +地 +第 +叮咚 +对 +对于 +多 +多少 +而 +而况 +而且 +而是 +而外 +而言 +而已 +尔后 +反过来 +反过来说 +反之 +非但 +非徒 +否则 +嘎 +嘎登 +该 +赶 +个 +各 +各个 +各位 +各种 +各自 +给 +根据 +跟 +故 +故此 +固然 +关于 +管 +归 +果然 +果真 +过 +哈 +哈哈 +呵 +和 +何 +何处 +何况 +何时 +嘿 +哼 +哼唷 +呼哧 +乎 +哗 +还是 +还有 +换句话说 +换言之 +或 +或是 +或者 +极了 +及 +及其 +及至 +即 +即便 +即或 +即令 +即若 +即使 +几 +几时 +己 +既 +既然 +既是 +继而 +加之 +假如 +假若 +假使 +鉴于 +将 +较 +较之 +叫 +接着 +结果 +借 +紧接着 +进而 +尽 +尽管 +经 +经过 +就 +就是 +就是说 +据 +具体地说 +具体说来 +开始 +开外 +靠 +咳 +可 +可见 +可是 +可以 +况且 +啦 +来 +来着 +离 +例如 +哩 +连 +连同 +两者 +了 +临 +另 +另外 +另一方面 +论 +嘛 +吗 +慢说 +漫说 +冒 +么 +每 +每当 +们 +莫若 +某 +某个 +某些 +拿 +哪 +哪边 +哪儿 +哪个 +哪里 +哪年 +哪怕 +哪天 +哪些 +哪样 +那 +那边 +那儿 +那个 +那会儿 +那里 +那么 +那么些 +那么样 +那时 +那些 +那样 +乃 +乃至 +呢 +能 +你 +你们 +您 +宁 +宁可 +宁肯 +宁愿 +哦 +呕 +啪达 +旁人 +呸 +凭 +凭借 +其 +其次 +其二 +其他 +其它 +其一 +其余 +其中 +起 +起见 +起见 +岂但 +恰恰相反 +前后 +前者 +且 +然而 +然后 +然则 +让 +人家 +任 +任何 +任凭 +如 +如此 +如果 +如何 +如其 +如若 +如上所述 +若 +若非 +若是 +啥 +上下 +尚且 +设若 +设使 +甚而 +甚么 +甚至 +省得 +时候 +什么 +什么样 +使得 +是 +是的 +首先 +谁 +谁知 +顺 +顺着 +似的 +虽 +虽然 +虽说 +虽则 +随 +随着 +所 +所以 +他 +他们 +他人 +它 +它们 +她 +她们 +倘 +倘或 +倘然 +倘若 +倘使 +腾 +替 +通过 +同 +同时 +哇 +万一 +往 +望 +为 +为何 +为了 +为什么 +为着 +喂 +嗡嗡 +我 +我们 +呜 +呜呼 +乌乎 +无论 +无宁 +毋宁 +嘻 +吓 +相对而言 +像 +向 +向着 +嘘 +呀 +焉 +沿 +沿着 +要 +要不 +要不然 +要不是 +要么 +要是 +也 +也罢 +也好 +一 +一般 +一旦 +一方面 +一来 +一切 +一样 +一则 +依 +依照 +矣 +以 +以便 +以及 +以免 +以至 +以至于 +以致 +抑或 +因 +因此 +因而 +因为 +哟 +用 +由 +由此可见 +由于 +有 +有的 +有关 +有些 +又 +于 +于是 +于是乎 +与 +与此同时 +与否 +与其 +越是 +云云 +哉 +再说 +再者 +在 +在下 +咱 +咱们 +则 +怎 +怎么 +怎么办 +怎么样 +怎样 +咋 +照 +照着 +者 +这 +这边 +这儿 +这个 +这会儿 +这就是说 +这里 +这么 +这么点儿 +这么些 +这么样 +这时 +这些 +这样 +正如 +吱 +之 +之类 +之所以 +之一 +只是 +只限 +只要 +只有 +至 +至于 +诸位 +着 +着呢 +自 +自从 +自个儿 +自各儿 +自己 +自家 +自身 +综上所述 +总的来看 +总的来说 +总的说来 +总而言之 +总之 +纵 +纵令 +纵然 +纵使 +遵照 +作为 +兮 +呃 +呗 +咚 +咦 +喏 +啐 +喔唷 +嗬 +嗯 +嗳 \ No newline at end of file diff --git a/astrbot/core/knowledge_base/retrieval/manager.py b/astrbot/core/knowledge_base/retrieval/manager.py new file mode 100644 index 00000000..9a42cd6c --- /dev/null +++ b/astrbot/core/knowledge_base/retrieval/manager.py @@ -0,0 +1,276 @@ +"""检索管理器 + +协调稠密检索、稀疏检索和 Rerank,提供统一的检索接口 +""" + +import time +from dataclasses import dataclass + +from astrbot import logger +from astrbot.core.db.vec_db.base import Result +from astrbot.core.db.vec_db.faiss_impl import FaissVecDB +from astrbot.core.knowledge_base.kb_db_sqlite import KBSQLiteDatabase +from astrbot.core.knowledge_base.retrieval.rank_fusion import RankFusion +from astrbot.core.knowledge_base.retrieval.sparse_retriever import SparseRetriever +from astrbot.core.provider.provider import RerankProvider + +from ..kb_helper import KBHelper + + +@dataclass +class RetrievalResult: + """检索结果""" + + chunk_id: str + doc_id: str + doc_name: str + kb_id: str + kb_name: str + content: str + score: float + metadata: dict + + +class RetrievalManager: + """检索管理器 + + 职责: + - 协调稠密检索、稀疏检索和 Rerank + - 结果融合和排序 + """ + + def __init__( + self, + sparse_retriever: SparseRetriever, + rank_fusion: RankFusion, + kb_db: KBSQLiteDatabase, + ): + """初始化检索管理器 + + Args: + vec_db_factory: 向量数据库工厂 + sparse_retriever: 稀疏检索器 + rank_fusion: 结果融合器 + kb_db: 知识库数据库实例 + + """ + self.sparse_retriever = sparse_retriever + self.rank_fusion = rank_fusion + self.kb_db = kb_db + + async def retrieve( + self, + query: str, + kb_ids: list[str], + kb_id_helper_map: dict[str, KBHelper], + top_k_fusion: int = 20, + top_m_final: int = 5, + ) -> list[RetrievalResult]: + """混合检索 + + 流程: + 1. 稠密检索 (向量相似度) + 2. 稀疏检索 (BM25) + 3. 结果融合 (RRF) + 4. Rerank 重排序 + + Args: + query: 查询文本 + kb_ids: 知识库 ID 列表 + top_m_final: 最终返回数量 + enable_rerank: 是否启用 Rerank + + Returns: + List[RetrievalResult]: 检索结果列表 + + """ + if not kb_ids: + return [] + + kb_options: dict = {} + new_kb_ids = [] + for kb_id in kb_ids: + kb_helper = kb_id_helper_map.get(kb_id) + if kb_helper: + kb = kb_helper.kb + kb_options[kb_id] = { + "top_k_dense": kb.top_k_dense or 50, + "top_k_sparse": kb.top_k_sparse or 50, + "top_m_final": kb.top_m_final or 5, + "vec_db": kb_helper.vec_db, + "rerank_provider_id": kb.rerank_provider_id, + } + new_kb_ids.append(kb_id) + else: + logger.warning(f"知识库 ID {kb_id} 实例未找到, 已跳过该知识库的检索") + + kb_ids = new_kb_ids + + # 1. 稠密检索 + time_start = time.time() + dense_results = await self._dense_retrieve( + query=query, + kb_ids=kb_ids, + kb_options=kb_options, + ) + time_end = time.time() + logger.debug( + f"Dense retrieval across {len(kb_ids)} bases took {time_end - time_start:.2f}s and returned {len(dense_results)} results.", + ) + + # 2. 稀疏检索 + time_start = time.time() + sparse_results = await self.sparse_retriever.retrieve( + query=query, + kb_ids=kb_ids, + kb_options=kb_options, + ) + time_end = time.time() + logger.debug( + f"Sparse retrieval across {len(kb_ids)} bases took {time_end - time_start:.2f}s and returned {len(sparse_results)} results.", + ) + + # 3. 结果融合 + time_start = time.time() + fused_results = await self.rank_fusion.fuse( + dense_results=dense_results, + sparse_results=sparse_results, + top_k=top_k_fusion, + ) + time_end = time.time() + logger.debug( + f"Rank fusion took {time_end - time_start:.2f}s and returned {len(fused_results)} results.", + ) + + # 4. 转换为 RetrievalResult (获取元数据) + retrieval_results = [] + for fr in fused_results: + metadata_dict = await self.kb_db.get_document_with_metadata(fr.doc_id) + if metadata_dict: + retrieval_results.append( + RetrievalResult( + chunk_id=fr.chunk_id, + doc_id=fr.doc_id, + doc_name=metadata_dict["document"].doc_name, + kb_id=fr.kb_id, + kb_name=metadata_dict["knowledge_base"].kb_name, + content=fr.content, + score=fr.score, + metadata={ + "chunk_index": fr.chunk_index, + "char_count": len(fr.content), + }, + ), + ) + + # 5. Rerank + first_rerank = None + for kb_id in kb_ids: + vec_db: FaissVecDB = kb_options[kb_id]["vec_db"] + rerank_pi = kb_options[kb_id]["rerank_provider_id"] + if ( + vec_db + and vec_db.rerank_provider + and rerank_pi + and rerank_pi == vec_db.rerank_provider.meta().id + ): + first_rerank = vec_db.rerank_provider + break + if first_rerank and retrieval_results: + retrieval_results = await self._rerank( + query=query, + results=retrieval_results, + top_k=top_m_final, + rerank_provider=first_rerank, + ) + + return retrieval_results[:top_m_final] + + async def _dense_retrieve( + self, + query: str, + kb_ids: list[str], + kb_options: dict, + ): + """稠密检索 (向量相似度) + + 为每个知识库使用独立的向量数据库进行检索,然后合并结果。 + + Args: + query: 查询文本 + kb_ids: 知识库 ID 列表 + top_k: 返回结果数量 + + Returns: + List[Result]: 检索结果列表 + + """ + all_results: list[Result] = [] + for kb_id in kb_ids: + if kb_id not in kb_options: + continue + try: + vec_db: FaissVecDB = kb_options[kb_id]["vec_db"] + dense_k = int(kb_options[kb_id]["top_k_dense"]) + vec_results = await vec_db.retrieve( + query=query, + k=dense_k, + fetch_k=dense_k * 2, + rerank=False, # 稠密检索阶段不进行 rerank + metadata_filters={"kb_id": kb_id}, + ) + + all_results.extend(vec_results) + except Exception as e: + from astrbot.core import logger + + logger.warning(f"知识库 {kb_id} 稠密检索失败: {e}") + continue + + # 按相似度排序并返回 top_k + all_results.sort(key=lambda x: x.similarity, reverse=True) + # return all_results[: len(all_results) // len(kb_ids)] + return all_results + + async def _rerank( + self, + query: str, + results: list[RetrievalResult], + top_k: int, + rerank_provider: RerankProvider, + ) -> list[RetrievalResult]: + """Rerank 重排序 + + Args: + query: 查询文本 + results: 检索结果列表 + top_k: 返回结果数量 + + Returns: + List[RetrievalResult]: 重排序后的结果列表 + + """ + if not results: + return [] + + # 准备文档列表 + docs = [r.content for r in results] + + # 调用 Rerank Provider + rerank_results = await rerank_provider.rerank( + query=query, + documents=docs, + ) + + # 更新分数并重新排序 + reranked_list = [] + for rerank_result in rerank_results: + idx = rerank_result.index + if idx < len(results): + result = results[idx] + result.score = rerank_result.relevance_score + reranked_list.append(result) + + reranked_list.sort(key=lambda x: x.score, reverse=True) + + return reranked_list[:top_k] diff --git a/astrbot/core/knowledge_base/retrieval/rank_fusion.py b/astrbot/core/knowledge_base/retrieval/rank_fusion.py new file mode 100644 index 00000000..26203f94 --- /dev/null +++ b/astrbot/core/knowledge_base/retrieval/rank_fusion.py @@ -0,0 +1,142 @@ +"""检索结果融合器 + +使用 Reciprocal Rank Fusion (RRF) 算法融合稠密检索和稀疏检索的结果 +""" + +import json +from dataclasses import dataclass + +from astrbot.core.db.vec_db.base import Result +from astrbot.core.knowledge_base.kb_db_sqlite import KBSQLiteDatabase +from astrbot.core.knowledge_base.retrieval.sparse_retriever import SparseResult + + +@dataclass +class FusedResult: + """融合后的检索结果""" + + chunk_id: str + chunk_index: int + doc_id: str + kb_id: str + content: str + score: float + + +class RankFusion: + """检索结果融合器 + + 职责: + - 融合稠密检索和稀疏检索的结果 + - 使用 Reciprocal Rank Fusion (RRF) 算法 + """ + + def __init__(self, kb_db: KBSQLiteDatabase, k: int = 60): + """初始化结果融合器 + + Args: + kb_db: 知识库数据库实例 + k: RRF 参数,用于平滑排名 + + """ + self.kb_db = kb_db + self.k = k + + async def fuse( + self, + dense_results: list[Result], + sparse_results: list[SparseResult], + top_k: int = 20, + ) -> list[FusedResult]: + """融合稠密和稀疏检索结果 + + RRF 公式: + score(doc) = sum(1 / (k + rank_i)) + + Args: + dense_results: 稠密检索结果 + sparse_results: 稀疏检索结果 + top_k: 返回结果数量 + + Returns: + List[FusedResult]: 融合后的结果列表 + + """ + # 1. 构建排名映射 + dense_ranks = { + r.data["doc_id"]: (idx + 1) for idx, r in enumerate(dense_results) + } # 这里的 doc_id 实际上是 chunk_id + sparse_ranks = {r.chunk_id: (idx + 1) for idx, r in enumerate(sparse_results)} + + # 2. 收集所有唯一的 ID + # 需要统一为 chunk_id + all_chunk_ids = set() + vec_doc_id_to_dense: dict[str, Result] = {} # vec_doc_id -> Result + chunk_id_to_sparse: dict[str, SparseResult] = {} # chunk_id -> SparseResult + + # 处理稀疏检索结果 + for r in sparse_results: + all_chunk_ids.add(r.chunk_id) + chunk_id_to_sparse[r.chunk_id] = r + + # 处理稠密检索结果 (需要转换 vec_doc_id 到 chunk_id) + for r in dense_results: + vec_doc_id = r.data["doc_id"] + all_chunk_ids.add(vec_doc_id) + vec_doc_id_to_dense[vec_doc_id] = r + + # 3. 计算 RRF 分数 + rrf_scores: dict[str, float] = {} + + for identifier in all_chunk_ids: + score = 0.0 + + # 来自稠密检索的贡献 + if identifier in dense_ranks: + score += 1.0 / (self.k + dense_ranks[identifier]) + + # 来自稀疏检索的贡献 + if identifier in sparse_ranks: + score += 1.0 / (self.k + sparse_ranks[identifier]) + + rrf_scores[identifier] = score + + # 4. 排序 + sorted_ids = sorted( + rrf_scores.keys(), + key=lambda cid: rrf_scores[cid], + reverse=True, + )[:top_k] + + # 5. 构建融合结果 + fused_results = [] + for identifier in sorted_ids: + # 优先从稀疏检索获取完整信息 + if identifier in chunk_id_to_sparse: + sr = chunk_id_to_sparse[identifier] + fused_results.append( + FusedResult( + chunk_id=sr.chunk_id, + chunk_index=sr.chunk_index, + doc_id=sr.doc_id, + kb_id=sr.kb_id, + content=sr.content, + score=rrf_scores[identifier], + ), + ) + elif identifier in vec_doc_id_to_dense: + # 从向量检索获取信息,需要从数据库获取块的详细信息 + vec_result = vec_doc_id_to_dense[identifier] + chunk_md = json.loads(vec_result.data["metadata"]) + fused_results.append( + FusedResult( + chunk_id=identifier, + chunk_index=chunk_md["chunk_index"], + doc_id=chunk_md["kb_doc_id"], + kb_id=chunk_md["kb_id"], + content=vec_result.data["text"], + score=rrf_scores[identifier], + ), + ) + + return fused_results diff --git a/astrbot/core/knowledge_base/retrieval/sparse_retriever.py b/astrbot/core/knowledge_base/retrieval/sparse_retriever.py new file mode 100644 index 00000000..ea5da1c9 --- /dev/null +++ b/astrbot/core/knowledge_base/retrieval/sparse_retriever.py @@ -0,0 +1,136 @@ +"""稀疏检索器 + +使用 BM25 算法进行基于关键词的文档检索 +""" + +import json +import os +from dataclasses import dataclass + +import jieba +from rank_bm25 import BM25Okapi + +from astrbot.core.db.vec_db.faiss_impl import FaissVecDB +from astrbot.core.knowledge_base.kb_db_sqlite import KBSQLiteDatabase + + +@dataclass +class SparseResult: + """稀疏检索结果""" + + chunk_index: int + chunk_id: str + doc_id: str + kb_id: str + content: str + score: float + + +class SparseRetriever: + """BM25 稀疏检索器 + + 职责: + - 基于关键词的文档检索 + - 使用 BM25 算法计算相关度 + """ + + def __init__(self, kb_db: KBSQLiteDatabase): + """初始化稀疏检索器 + + Args: + kb_db: 知识库数据库实例 + + """ + self.kb_db = kb_db + self._index_cache = {} # 缓存 BM25 索引 + + with open( + os.path.join(os.path.dirname(__file__), "hit_stopwords.txt"), + encoding="utf-8", + ) as f: + self.hit_stopwords = { + word.strip() for word in set(f.read().splitlines()) if word.strip() + } + + async def retrieve( + self, + query: str, + kb_ids: list[str], + kb_options: dict, + ) -> list[SparseResult]: + """执行稀疏检索 + + Args: + query: 查询文本 + kb_ids: 知识库 ID 列表 + kb_options: 每个知识库的检索选项 + + Returns: + List[SparseResult]: 检索结果列表 + + """ + # 1. 获取所有相关块 + top_k_sparse = 0 + chunks = [] + for kb_id in kb_ids: + vec_db: FaissVecDB = kb_options.get(kb_id, {}).get("vec_db") + if not vec_db: + continue + result = await vec_db.document_storage.get_documents( + metadata_filters={}, + limit=None, + offset=None, + ) + chunk_mds = [json.loads(doc["metadata"]) for doc in result] + result = [ + { + "chunk_id": doc["doc_id"], + "chunk_index": chunk_md["chunk_index"], + "doc_id": chunk_md["kb_doc_id"], + "kb_id": kb_id, + "text": doc["text"], + } + for doc, chunk_md in zip(result, chunk_mds) + ] + chunks.extend(result) + top_k_sparse += kb_options.get(kb_id, {}).get("top_k_sparse", 50) + + if not chunks: + return [] + + # 2. 准备文档和索引 + corpus = [chunk["text"] for chunk in chunks] + tokenized_corpus = [list(jieba.cut(doc)) for doc in corpus] + tokenized_corpus = [ + [word for word in doc if word not in self.hit_stopwords] + for doc in tokenized_corpus + ] + + # 3. 构建 BM25 索引 + bm25 = BM25Okapi(tokenized_corpus) + + # 4. 执行检索 + tokenized_query = list(jieba.cut(query)) + tokenized_query = [ + word for word in tokenized_query if word not in self.hit_stopwords + ] + scores = bm25.get_scores(tokenized_query) + + # 5. 排序并返回 Top-K + results = [] + for idx, score in enumerate(scores): + chunk = chunks[idx] + results.append( + SparseResult( + chunk_id=chunk["chunk_id"], + chunk_index=chunk["chunk_index"], + doc_id=chunk["doc_id"], + kb_id=chunk["kb_id"], + content=chunk["text"], + score=float(score), + ), + ) + + results.sort(key=lambda x: x.score, reverse=True) + # return results[: len(results) // len(kb_ids)] + return results[:top_k_sparse] diff --git a/astrbot/core/log.py b/astrbot/core/log.py index 3a1c5037..376f5ffd 100644 --- a/astrbot/core/log.py +++ b/astrbot/core/log.py @@ -1,5 +1,4 @@ -""" -日志系统, 用于支持核心组件和插件的日志记录, 提供了日志订阅功能 +"""日志系统, 用于支持核心组件和插件的日志记录, 提供了日志订阅功能 const: CACHED_SIZE: 日志缓存大小, 用于限制缓存的日志数量 @@ -21,14 +20,14 @@ function: 4. 订阅者可以使用 register() 方法注册到 LogBroker, 订阅日志流 """ -import logging -import colorlog import asyncio +import logging import os import sys -from collections import deque from asyncio import Queue -from typing import List +from collections import deque + +import colorlog # 日志缓存大小 CACHED_SIZE = 200 @@ -52,6 +51,7 @@ def is_plugin_path(pathname): Returns: bool: 如果路径来自插件目录，则返回 True，否则返回 False + """ if not pathname: return False @@ -68,6 +68,7 @@ def get_short_level_name(level_name): Returns: str: 四个字母的日志级别缩写 + """ level_map = { "DEBUG": "DBUG", @@ -87,13 +88,14 @@ class LogBroker: def __init__(self): self.log_cache = deque(maxlen=CACHED_SIZE) # 环形缓冲区, 保存最近的日志 - self.subscribers: List[Queue] = [] # 订阅者列表 + self.subscribers: list[Queue] = [] # 订阅者列表 def register(self) -> Queue: """注册新的订阅者, 并给每个订阅者返回一个带有日志缓存的队列 Returns: Queue: 订阅者的队列, 可用于接收日志消息 + """ q = Queue(maxsize=CACHED_SIZE + 10) self.subscribers.append(q) @@ -104,6 +106,7 @@ class LogBroker: Args: q (Queue): 需要取消订阅的队列 + """ self.subscribers.remove(q) @@ -113,6 +116,7 @@ class LogBroker: Args: log_entry (dict): 日志消息, 包含日志级别和日志内容. example: {"level": "INFO", "data": "This is a log message.", "time": "2023-10-01 12:00:00"} + """ self.log_cache.append(log_entry) for q in self.subscribers: @@ -138,6 +142,7 @@ class LogQueueHandler(logging.Handler): Args: record (logging.LogRecord): 日志记录对象, 包含日志信息 + """ log_entry = self.format(record) self.log_broker.publish( @@ -145,7 +150,7 @@ class LogQueueHandler(logging.Handler): "level": record.levelname, "time": record.asctime, "data": log_entry, - } + }, ) @@ -164,6 +169,7 @@ class LogManager: Returns: logging.Logger: 返回配置好的日志记录器 + """ logger = logging.getLogger(log_name) # 检查该logger或父级logger是否已经有处理器, 如果已经有处理器, 直接返回该logger, 避免重复配置 @@ -171,10 +177,10 @@ class LogManager: return logger # 如果logger没有处理器 console_handler = logging.StreamHandler( - sys.stdout + sys.stdout, ) # 创建一个StreamHandler用于控制台输出 console_handler.setLevel( - logging.DEBUG + logging.DEBUG, ) # 将日志级别设置为DEBUG(最低级别, 显示所有日志), *如果插件没有设置级别, 默认为DEBUG # 创建彩色日志格式化器, 输出日志格式为: [时间] [插件标签] [日志级别] [文件名:行号]: 日志消息 @@ -195,7 +201,8 @@ class LogManager: class FileNameFilter(logging.Filter): """文件名过滤器类, 用于修改日志记录的文件名格式 - 例如: 将文件路径 /path/to/file.py 转换为 file. 格式""" + 例如: 将文件路径 /path/to/file.py 转换为 file. 格式 + """ # 获取这个文件和父文件夹的名字：. 并且去除 .py def filter(self, record): @@ -231,6 +238,7 @@ class LogManager: Args: logger (logging.Logger): 日志记录器 log_broker (LogBroker): 日志代理类, 用于缓存和分发日志消息 + """ handler = LogQueueHandler(log_broker) handler.setLevel(logging.DEBUG) @@ -240,7 +248,7 @@ class LogManager: # 为队列处理器设置相同格式的formatter handler.setFormatter( logging.Formatter( - "[%(asctime)s] [%(short_levelname)s] %(plugin_tag)s[%(filename)s:%(lineno)d]: %(message)s" - ) + "[%(asctime)s] [%(short_levelname)s] %(plugin_tag)s[%(filename)s:%(lineno)d]: %(message)s", + ), ) logger.addHandler(handler) diff --git a/astrbot/core/message/components.py b/astrbot/core/message/components.py index d9ec4b41..43e3bf0e 100644 --- a/astrbot/core/message/components.py +++ b/astrbot/core/message/components.py @@ -1,5 +1,4 @@ -""" -MIT License +"""MIT License Copyright (c) 2021 Lxns-Network @@ -26,7 +25,6 @@ import asyncio import base64 import json import os -import typing as T import uuid from enum import Enum @@ -38,60 +36,36 @@ from astrbot.core.utils.io import download_file, download_image_by_url, file_to_ class ComponentType(str, Enum): - Plain = "Plain" # 纯文本消息 - Face = "Face" # QQ表情 - Record = "Record" # 语音 - Video = "Video" # 视频 - At = "At" # At - Node = "Node" # 转发消息的一个节点 - Nodes = "Nodes" # 转发消息的多个节点 - Poke = "Poke" # QQ 戳一戳 - Image = "Image" # 图片 - Reply = "Reply" # 回复 - Forward = "Forward" # 转发消息 - File = "File" # 文件 + # Basic Segment Types + Plain = "Plain" # plain text message + Image = "Image" # image + Record = "Record" # audio + Video = "Video" # video + File = "File" # file attachment + # IM-specific Segment Types + Face = "Face" # Emoji segment for Tencent QQ platform + At = "At" # mention a user in IM apps + Node = "Node" # a node in a forwarded message + Nodes = "Nodes" # a forwarded message consisting of multiple nodes + Poke = "Poke" # a poke message for Tencent QQ platform + Reply = "Reply" # a reply message segment + Forward = "Forward" # a forwarded message segment RPS = "RPS" # TODO Dice = "Dice" # TODO Shake = "Shake" # TODO - Anonymous = "Anonymous" # TODO Share = "Share" Contact = "Contact" # TODO Location = "Location" # TODO Music = "Music" - RedBag = "RedBag" - Xml = "Xml" Json = "Json" - CardImage = "CardImage" - TTS = "TTS" Unknown = "Unknown" - WechatEmoji = "WechatEmoji" # Wechat 下的 emoji 表情包 class BaseMessageComponent(BaseModel): type: ComponentType - def toString(self): - output = f"[CQ:{self.type.lower()}" - for k, v in self.__dict__.items(): - if k == "type" or v is None: - continue - if k == "_type": - k = "type" - if isinstance(v, bool): - v = 1 if v else 0 - output += ",%s=%s" % ( - k, - str(v) - .replace("&", "&") - .replace(",", ",") - .replace("[", "[") - .replace("]", "]"), - ) - output += "]" - return output - def toDict(self): data = {} for k, v in self.__dict__.items(): @@ -110,18 +84,11 @@ class BaseMessageComponent(BaseModel): class Plain(BaseMessageComponent): type = ComponentType.Plain text: str - convert: T.Optional[bool] = True # 若为 False 则直接发送未转换 CQ 码的消息 + convert: bool | None = True def __init__(self, text: str, convert: bool = True, **_): super().__init__(text=text, convert=convert, **_) - def toString(self): # 没有 [CQ:plain] 这种东西，所以直接导出纯文本 - if not self.convert: - return self.text - return ( - self.text.replace("&", "&").replace("[", "[").replace("]", "]") - ) - def toDict(self): return {"type": "text", "data": {"text": self.text.strip()}} @@ -139,17 +106,17 @@ class Face(BaseMessageComponent): class Record(BaseMessageComponent): type = ComponentType.Record - file: T.Optional[str] = "" - magic: T.Optional[bool] = False - url: T.Optional[str] = "" - cache: T.Optional[bool] = True - proxy: T.Optional[bool] = True - timeout: T.Optional[int] = 0 + file: str | None = "" + magic: bool | None = False + url: str | None = "" + cache: bool | None = True + proxy: bool | None = True + timeout: int | None = 0 # 额外 - path: T.Optional[str] + path: str | None - def __init__(self, file: T.Optional[str], **_): - for k in _.keys(): + def __init__(self, file: str | None, **_): + for k in _: if k == "url": pass # Protocol.warn(f"go-cqhttp doesn't support send {self.type} by {k}") @@ -174,15 +141,16 @@ class Record(BaseMessageComponent): Returns: str: 语音的本地路径，以绝对路径表示。 + """ if not self.file: raise Exception(f"not a valid file: {self.file}") if self.file.startswith("file:///"): return self.file[8:] - elif self.file.startswith("http"): + if self.file.startswith("http"): file_path = await download_image_by_url(self.file) return os.path.abspath(file_path) - elif self.file.startswith("base64://"): + if self.file.startswith("base64://"): bs64_data = self.file.removeprefix("base64://") image_bytes = base64.b64decode(bs64_data) temp_dir = os.path.join(get_astrbot_data_path(), "temp") @@ -190,16 +158,16 @@ class Record(BaseMessageComponent): with open(file_path, "wb") as f: f.write(image_bytes) return os.path.abspath(file_path) - elif os.path.exists(self.file): + if os.path.exists(self.file): return os.path.abspath(self.file) - else: - raise Exception(f"not a valid file: {self.file}") + raise Exception(f"not a valid file: {self.file}") async def convert_to_base64(self) -> str: """将语音统一转换为 base64 编码。这个方法避免了手动判断语音数据类型，直接返回语音数据的 base64 编码。 Returns: str: 语音的 base64 编码，不以 base64:// 或者 data:image/jpeg;base64, 开头。 + """ # convert to base64 if not self.file: @@ -219,14 +187,14 @@ class Record(BaseMessageComponent): return bs64_data async def register_to_file_service(self) -> str: - """ - 将语音注册到文件服务。 + """将语音注册到文件服务。 Returns: str: 注册后的URL Raises: Exception: 如果未配置 callback_api_base + """ callback_host = astrbot_config.get("callback_api_base") @@ -245,10 +213,10 @@ class Record(BaseMessageComponent): class Video(BaseMessageComponent): type = ComponentType.Video file: str - cover: T.Optional[str] = "" - c: T.Optional[int] = 2 + cover: str | None = "" + c: int | None = 2 # 额外 - path: T.Optional[str] = "" + path: str | None = "" def __init__(self, file: str, **_): super().__init__(file=file, **_) @@ -268,32 +236,31 @@ class Video(BaseMessageComponent): Returns: str: 视频的本地路径，以绝对路径表示。 + """ url = self.file if url and url.startswith("file:///"): return url[8:] - elif url and url.startswith("http"): + if url and url.startswith("http"): download_dir = os.path.join(get_astrbot_data_path(), "temp") video_file_path = os.path.join(download_dir, f"{uuid.uuid4().hex}") await download_file(url, video_file_path) if os.path.exists(video_file_path): return os.path.abspath(video_file_path) - else: - raise Exception(f"download failed: {url}") - elif os.path.exists(url): + raise Exception(f"download failed: {url}") + if os.path.exists(url): return os.path.abspath(url) - else: - raise Exception(f"not a valid file: {url}") + raise Exception(f"not a valid file: {url}") async def register_to_file_service(self): - """ - 将视频注册到文件服务。 + """将视频注册到文件服务。 Returns: str: 注册后的URL Raises: Exception: 如果未配置 callback_api_base + """ callback_host = astrbot_config.get("callback_api_base") @@ -330,8 +297,8 @@ class Video(BaseMessageComponent): class At(BaseMessageComponent): type = ComponentType.At - qq: T.Union[int, str] # 此处str为all时代表所有人 - name: T.Optional[str] = "" + qq: int | str # 此处str为all时代表所有人 + name: str | None = "" def __init__(self, **_): super().__init__(**_) @@ -371,20 +338,12 @@ class Shake(BaseMessageComponent): # TODO super().__init__(**_) -class Anonymous(BaseMessageComponent): # TODO - type = ComponentType.Anonymous - ignore: T.Optional[bool] = False - - def __init__(self, **_): - super().__init__(**_) - - class Share(BaseMessageComponent): type = ComponentType.Share url: str title: str - content: T.Optional[str] = "" - image: T.Optional[str] = "" + content: str | None = "" + image: str | None = "" def __init__(self, **_): super().__init__(**_) @@ -393,7 +352,7 @@ class Share(BaseMessageComponent): class Contact(BaseMessageComponent): # TODO type = ComponentType.Contact _type: str # type 字段冲突 - id: T.Optional[int] = 0 + id: int | None = 0 def __init__(self, **_): super().__init__(**_) @@ -403,8 +362,8 @@ class Location(BaseMessageComponent): # TODO type = ComponentType.Location lat: float lon: float - title: T.Optional[str] = "" - content: T.Optional[str] = "" + title: str | None = "" + content: str | None = "" def __init__(self, **_): super().__init__(**_) @@ -413,12 +372,12 @@ class Location(BaseMessageComponent): # TODO class Music(BaseMessageComponent): type = ComponentType.Music _type: str - id: T.Optional[int] = 0 - url: T.Optional[str] = "" - audio: T.Optional[str] = "" - title: T.Optional[str] = "" - content: T.Optional[str] = "" - image: T.Optional[str] = "" + id: int | None = 0 + url: str | None = "" + audio: str | None = "" + title: str | None = "" + content: str | None = "" + image: str | None = "" def __init__(self, **_): # for k in _.keys(): @@ -429,18 +388,18 @@ class Music(BaseMessageComponent): class Image(BaseMessageComponent): type = ComponentType.Image - file: T.Optional[str] = "" - _type: T.Optional[str] = "" - subType: T.Optional[int] = 0 - url: T.Optional[str] = "" - cache: T.Optional[bool] = True - id: T.Optional[int] = 40000 - c: T.Optional[int] = 2 + file: str | None = "" + _type: str | None = "" + subType: int | None = 0 + url: str | None = "" + cache: bool | None = True + id: int | None = 40000 + c: int | None = 2 # 额外 - path: T.Optional[str] = "" - file_unique: T.Optional[str] = "" # 某些平台可能有图片缓存的唯一标识 + path: str | None = "" + file_unique: str | None = "" # 某些平台可能有图片缓存的唯一标识 - def __init__(self, file: T.Optional[str], **_): + def __init__(self, file: str | None, **_): super().__init__(file=file, **_) @staticmethod @@ -470,16 +429,17 @@ class Image(BaseMessageComponent): Returns: str: 图片的本地路径，以绝对路径表示。 + """ url = self.url or self.file if not url: raise ValueError("No valid file or URL provided") if url.startswith("file:///"): return url[8:] - elif url.startswith("http"): + if url.startswith("http"): image_file_path = await download_image_by_url(url) return os.path.abspath(image_file_path) - elif url.startswith("base64://"): + if url.startswith("base64://"): bs64_data = url.removeprefix("base64://") image_bytes = base64.b64decode(bs64_data) temp_dir = os.path.join(get_astrbot_data_path(), "temp") @@ -487,16 +447,16 @@ class Image(BaseMessageComponent): with open(image_file_path, "wb") as f: f.write(image_bytes) return os.path.abspath(image_file_path) - elif os.path.exists(url): + if os.path.exists(url): return os.path.abspath(url) - else: - raise Exception(f"not a valid file: {url}") + raise Exception(f"not a valid file: {url}") async def convert_to_base64(self) -> str: """将这个图片统一转换为 base64 编码。这个方法避免了手动判断图片数据类型，直接返回图片数据的 base64 编码。 Returns: str: 图片的 base64 编码，不以 base64:// 或者 data:image/jpeg;base64, 开头。 + """ # convert to base64 url = self.url or self.file @@ -517,14 +477,14 @@ class Image(BaseMessageComponent): return bs64_data async def register_to_file_service(self) -> str: - """ - 将图片注册到文件服务。 + """将图片注册到文件服务。 Returns: str: 注册后的URL Raises: Exception: 如果未配置 callback_api_base + """ callback_host = astrbot_config.get("callback_api_base") @@ -542,42 +502,34 @@ class Image(BaseMessageComponent): class Reply(BaseMessageComponent): type = ComponentType.Reply - id: T.Union[str, int] + id: str | int """所引用的消息 ID""" - chain: T.Optional[T.List["BaseMessageComponent"]] = [] + chain: list["BaseMessageComponent"] | None = [] """被引用的消息段列表""" - sender_id: T.Optional[int] | T.Optional[str] = 0 + sender_id: int | None | str = 0 """被引用的消息对应的发送者的 ID""" - sender_nickname: T.Optional[str] = "" + sender_nickname: str | None = "" """被引用的消息对应的发送者的昵称""" - time: T.Optional[int] = 0 + time: int | None = 0 """被引用的消息发送时间""" - message_str: T.Optional[str] = "" + message_str: str | None = "" """被引用的消息解析后的纯文本消息字符串""" - text: T.Optional[str] = "" + text: str | None = "" """deprecated""" - qq: T.Optional[int] = 0 + qq: int | None = 0 """deprecated""" - seq: T.Optional[int] = 0 + seq: int | None = 0 """deprecated""" def __init__(self, **_): super().__init__(**_) -class RedBag(BaseMessageComponent): - type = ComponentType.RedBag - title: str - - def __init__(self, **_): - super().__init__(**_) - - class Poke(BaseMessageComponent): type: str = ComponentType.Poke - id: T.Optional[int] = 0 - qq: T.Optional[int] = 0 + id: int | None = 0 + qq: int | None = 0 def __init__(self, type: str, **_): type = f"Poke:{type}" @@ -596,12 +548,12 @@ class Node(BaseMessageComponent): """群合并转发消息""" type = ComponentType.Node - id: T.Optional[int] = 0 # 忽略 - name: T.Optional[str] = "" # qq昵称 - uin: T.Optional[str] = "0" # qq号 - content: T.Optional[list[BaseMessageComponent]] = [] - seq: T.Optional[T.Union[str, list]] = "" # 忽略 - time: T.Optional[int] = 0 # 忽略 + id: int | None = 0 # 忽略 + name: str | None = "" # qq昵称 + uin: str | None = "0" # qq号 + content: list[BaseMessageComponent] | None = [] + seq: str | list | None = "" # 忽略 + time: int | None = 0 # 忽略 def __init__(self, content: list[BaseMessageComponent], **_): if isinstance(content, Node): @@ -619,7 +571,7 @@ class Node(BaseMessageComponent): { "type": comp.type.lower(), "data": {"file": f"base64://{bs64}"}, - } + }, ) elif isinstance(comp, Plain): # For Plain segments, we need to handle the plain differently @@ -648,9 +600,9 @@ class Node(BaseMessageComponent): class Nodes(BaseMessageComponent): type = ComponentType.Nodes - nodes: T.List[Node] + nodes: list[Node] - def __init__(self, nodes: T.List[Node], **_): + def __init__(self, nodes: list[Node], **_): super().__init__(nodes=nodes, **_) def toDict(self): @@ -672,19 +624,10 @@ class Nodes(BaseMessageComponent): return ret -class Xml(BaseMessageComponent): - type = ComponentType.Xml - data: str - resid: T.Optional[int] = 0 - - def __init__(self, **_): - super().__init__(**_) - - class Json(BaseMessageComponent): type = ComponentType.Json - data: T.Union[str, dict] - resid: T.Optional[int] = 0 + data: str | dict + resid: int | None = 0 def __init__(self, data, **_): if isinstance(data, dict): @@ -692,50 +635,18 @@ class Json(BaseMessageComponent): super().__init__(data=data, **_) -class CardImage(BaseMessageComponent): - type = ComponentType.CardImage - file: str - cache: T.Optional[bool] = True - minwidth: T.Optional[int] = 400 - minheight: T.Optional[int] = 400 - maxwidth: T.Optional[int] = 500 - maxheight: T.Optional[int] = 500 - source: T.Optional[str] = "" - icon: T.Optional[str] = "" - - def __init__(self, **_): - super().__init__(**_) - - @staticmethod - def fromFileSystem(path, **_): - return CardImage(file=f"file:///{os.path.abspath(path)}", **_) - - -class TTS(BaseMessageComponent): - type = ComponentType.TTS - text: str - - def __init__(self, **_): - super().__init__(**_) - - class Unknown(BaseMessageComponent): type = ComponentType.Unknown text: str - def toString(self): - return "" - class File(BaseMessageComponent): - """ - 文件消息段 - """ + """文件消息段""" type = ComponentType.File - name: T.Optional[str] = "" # 名字 - file_: T.Optional[str] = "" # 本地路径 - url: T.Optional[str] = "" # url + name: str | None = "" # 名字 + file_: str | None = "" # 本地路径 + url: str | None = "" # url def __init__(self, name: str, file: str = "", url: str = ""): """文件消息段。""" @@ -743,11 +654,11 @@ class File(BaseMessageComponent): @property def file(self) -> str: - """ - 获取文件路径，如果文件不存在但有URL，则同步下载文件 + """获取文件路径，如果文件不存在但有URL，则同步下载文件 Returns: str: 文件路径 + """ if self.file_ and os.path.exists(self.file_): return os.path.abspath(self.file_) @@ -757,19 +668,16 @@ class File(BaseMessageComponent): loop = asyncio.get_event_loop() if loop.is_running(): logger.warning( - ( - "不可以在异步上下文中同步等待下载! " - "这个警告通常发生于某些逻辑试图通过 .file 获取文件消息段的文件内容。" - "请使用 await get_file() 代替直接获取 .file 字段" - ) + "不可以在异步上下文中同步等待下载! " + "这个警告通常发生于某些逻辑试图通过 .file 获取文件消息段的文件内容。" + "请使用 await get_file() 代替直接获取 .file 字段", ) return "" - else: - # 等待下载完成 - loop.run_until_complete(self._download_file()) + # 等待下载完成 + loop.run_until_complete(self._download_file()) - if self.file_ and os.path.exists(self.file_): - return os.path.abspath(self.file_) + if self.file_ and os.path.exists(self.file_): + return os.path.abspath(self.file_) except Exception as e: logger.error(f"文件下载失败: {e}") @@ -777,11 +685,11 @@ class File(BaseMessageComponent): @file.setter def file(self, value: str): - """ - 向前兼容, 设置file属性, 传入的参数可能是文件路径或URL + """向前兼容, 设置file属性, 传入的参数可能是文件路径或URL Args: value (str): 文件路径或URL + """ if value.startswith("http://") or value.startswith("https://"): self.url = value @@ -796,6 +704,7 @@ class File(BaseMessageComponent): 注意，如果为 True，也可能返回文件路径。 Returns: str: 文件路径或者 http 下载链接 + """ if allow_return_url and self.url: return self.url @@ -818,14 +727,14 @@ class File(BaseMessageComponent): self.file_ = os.path.abspath(file_path) async def register_to_file_service(self): - """ - 将文件注册到文件服务。 + """将文件注册到文件服务。 Returns: str: 注册后的URL Raises: Exception: 如果未配置 callback_api_base + """ callback_host = astrbot_config.get("callback_api_base") @@ -863,41 +772,38 @@ class File(BaseMessageComponent): class WechatEmoji(BaseMessageComponent): type = ComponentType.WechatEmoji - md5: T.Optional[str] = "" - md5_len: T.Optional[int] = 0 - cdnurl: T.Optional[str] = "" + md5: str | None = "" + md5_len: int | None = 0 + cdnurl: str | None = "" def __init__(self, **_): super().__init__(**_) ComponentTypes = { + # Basic Message Segments "plain": Plain, "text": Plain, - "face": Face, + "image": Image, "record": Record, "video": Video, + "file": File, + # IM-specific Message Segments + "face": Face, "at": At, "rps": RPS, "dice": Dice, "shake": Shake, - "anonymous": Anonymous, "share": Share, "contact": Contact, "location": Location, "music": Music, - "image": Image, "reply": Reply, - "redbag": RedBag, "poke": Poke, "forward": Forward, "node": Node, "nodes": Nodes, - "xml": Xml, "json": Json, - "cardimage": CardImage, - "tts": TTS, "unknown": Unknown, - "file": File, "WechatEmoji": WechatEmoji, } diff --git a/astrbot/core/message/message_event_result.py b/astrbot/core/message/message_event_result.py index 7bfdd34c..ed4e25f4 100644 --- a/astrbot/core/message/message_event_result.py +++ b/astrbot/core/message/message_event_result.py @@ -1,15 +1,16 @@ import enum - -from typing import List, Optional, Union, AsyncGenerator +from collections.abc import AsyncGenerator from dataclasses import dataclass, field + +from typing_extensions import deprecated + from astrbot.core.message.components import ( - BaseMessageComponent, - Plain, - Image, At, AtAll, + BaseMessageComponent, + Image, + Plain, ) -from typing_extensions import deprecated @dataclass @@ -20,18 +21,18 @@ class MessageChain: Attributes: `chain` (list): 用于顺序存储各个组件。 `use_t2i_` (bool): 用于标记是否使用文本转图片服务。默认为 None，即跟随用户的设置。当设置为 True 时，将会使用文本转图片服务。 + """ - chain: List[BaseMessageComponent] = field(default_factory=list) - use_t2i_: Optional[bool] = None # None 为跟随用户设置 - type: Optional[str] = None + chain: list[BaseMessageComponent] = field(default_factory=list) + use_t2i_: bool | None = None # None 为跟随用户设置 + type: str | None = None """消息链承载的消息的类型。可选，用于让消息平台区分不同业务场景的消息链。""" def message(self, message: str): """添加一条文本消息到消息链 `chain` 中。 Example: - CommandResult().message("Hello ").message("world!") # 输出 Hello world! @@ -39,11 +40,10 @@ class MessageChain: self.chain.append(Plain(message)) return self - def at(self, name: str, qq: Union[str, int]): + def at(self, name: str, qq: str | int): """添加一条 At 消息到消息链 `chain` 中。 Example: - CommandResult().at("张三", "12345678910") # 输出 @张三 @@ -55,7 +55,6 @@ class MessageChain: """添加一条 AtAll 消息到消息链 `chain` 中。 Example: - CommandResult().at_all() # 输出 @所有人 @@ -68,7 +67,6 @@ class MessageChain: """添加一条错误消息到消息链 `chain` 中 Example: - CommandResult().error("解析失败") """ @@ -82,7 +80,6 @@ class MessageChain: 如果需要发送本地图片，请使用 `file_image` 方法。 Example: - CommandResult().image("https://example.com/image.jpg") """ @@ -96,6 +93,7 @@ class MessageChain: 如果需要发送网络图片，请使用 `url_image` 方法。 CommandResult().image("image.jpg") + """ self.chain.append(Image.fromFileSystem(path)) return self @@ -114,6 +112,7 @@ class MessageChain: Args: use_t2i (bool): 是否使用文本转图片服务。默认为 None，即跟随用户的设置。当设置为 True 时，将会使用文本转图片服务。 + """ self.use_t2i_ = use_t2i return self @@ -125,7 +124,7 @@ class MessageChain: def squash_plain(self): """将消息链中的所有 Plain 消息段聚合到第一个 Plain 消息段中。""" if not self.chain: - return + return None new_chain = [] first_plain = None @@ -153,6 +152,7 @@ class EventResultType(enum.Enum): Attributes: CONTINUE: 事件将会继续传播 STOP: 事件将会终止传播 + """ CONTINUE = enum.auto() @@ -181,17 +181,18 @@ class MessageEventResult(MessageChain): `chain` (list): 用于顺序存储各个组件。 `use_t2i_` (bool): 用于标记是否使用文本转图片服务。默认为 None，即跟随用户的设置。当设置为 True 时，将会使用文本转图片服务。 `result_type` (EventResultType): 事件处理的结果类型。 + """ - result_type: Optional[EventResultType] = field( - default_factory=lambda: EventResultType.CONTINUE + result_type: EventResultType | None = field( + default_factory=lambda: EventResultType.CONTINUE, ) - result_content_type: Optional[ResultContentType] = field( - default_factory=lambda: ResultContentType.GENERAL_RESULT + result_content_type: ResultContentType | None = field( + default_factory=lambda: ResultContentType.GENERAL_RESULT, ) - async_stream: Optional[AsyncGenerator] = None + async_stream: AsyncGenerator | None = None """异步流""" def stop_event(self) -> "MessageEventResult": @@ -205,9 +206,7 @@ class MessageEventResult(MessageChain): return self def is_stopped(self) -> bool: - """ - 是否终止事件传播。 - """ + """是否终止事件传播。""" return self.result_type == EventResultType.STOP def set_async_stream(self, stream: AsyncGenerator) -> "MessageEventResult": @@ -220,6 +219,7 @@ class MessageEventResult(MessageChain): Args: result_type (EventResultType): 事件处理的结果类型。 + """ self.result_content_type = typ return self diff --git a/astrbot/core/persona_mgr.py b/astrbot/core/persona_mgr.py index add3c74b..5d1743ab 100644 --- a/astrbot/core/persona_mgr.py +++ b/astrbot/core/persona_mgr.py @@ -1,8 +1,8 @@ +from astrbot import logger +from astrbot.core.astrbot_config_mgr import AstrBotConfigManager from astrbot.core.db import BaseDatabase from astrbot.core.db.po import Persona, Personality -from astrbot.core.astrbot_config_mgr import AstrBotConfigManager from astrbot.core.platform.message_session import MessageSession -from astrbot import logger DEFAULT_PERSONALITY = Personality( prompt="You are a helpful and friendly assistant.", @@ -41,12 +41,14 @@ class PersonaManager: return persona async def get_default_persona_v3( - self, umo: str | MessageSession | None = None + self, + umo: str | MessageSession | None = None, ) -> Personality: """获取默认 persona""" cfg = self.acm.get_conf(umo) default_persona_id = cfg.get("provider_settings", {}).get( - "default_personality", "default" + "default_personality", + "default", ) if not default_persona_id or default_persona_id == "default": return DEFAULT_PERSONALITY @@ -66,16 +68,19 @@ class PersonaManager: async def update_persona( self, persona_id: str, - system_prompt: str = None, - begin_dialogs: list[str] = None, - tools: list[str] = None, + system_prompt: str | None = None, + begin_dialogs: list[str] | None = None, + tools: list[str] | None = None, ): """更新指定 persona 的信息。tools 参数为 None 时表示使用所有工具，空列表表示不使用任何工具""" existing_persona = await self.db.get_persona_by_id(persona_id) if not existing_persona: raise ValueError(f"Persona with ID {persona_id} does not exist.") persona = await self.db.update_persona( - persona_id, system_prompt, begin_dialogs, tools=tools + persona_id, + system_prompt, + begin_dialogs, + tools=tools, ) if persona: for i, p in enumerate(self.personas): @@ -100,7 +105,10 @@ class PersonaManager: if await self.db.get_persona_by_id(persona_id): raise ValueError(f"Persona with ID {persona_id} already exists.") new_persona = await self.db.insert_persona( - persona_id, system_prompt, begin_dialogs, tools=tools + persona_id, + system_prompt, + begin_dialogs, + tools=tools, ) self.personas.append(new_persona) self.get_v3_persona_data() @@ -115,6 +123,7 @@ class PersonaManager: - list[dict]: 包含 persona 配置的字典列表。 - list[Personality]: 包含 Personality 对象的列表。 - Personality: 默认选择的 Personality 对象。 + """ v3_persona_config = [ { @@ -136,7 +145,7 @@ class PersonaManager: if begin_dialogs: if len(begin_dialogs) % 2 != 0: logger.error( - f"{persona_cfg['name']} 人格情景预设对话格式不对，条数应该为偶数。" + f"{persona_cfg['name']} 人格情景预设对话格式不对，条数应该为偶数。", ) begin_dialogs = [] user_turn = True @@ -146,7 +155,7 @@ class PersonaManager: "role": "user" if user_turn else "assistant", "content": dialog, "_no_save": None, # 不持久化到 db - } + }, ) user_turn = not user_turn diff --git a/astrbot/core/pipeline/__init__.py b/astrbot/core/pipeline/__init__.py index 29a324a1..75fef84d 100644 --- a/astrbot/core/pipeline/__init__.py +++ b/astrbot/core/pipeline/__init__.py @@ -27,15 +27,15 @@ STAGES_ORDER = [ ] __all__ = [ - "WakingCheckStage", - "WhitelistCheckStage", - "SessionStatusCheckStage", - "RateLimitStage", "ContentSafetyCheckStage", + "EventResultType", + "MessageEventResult", "PreProcessStage", "ProcessStage", - "ResultDecorateStage", + "RateLimitStage", "RespondStage", - "MessageEventResult", - "EventResultType", + "ResultDecorateStage", + "SessionStatusCheckStage", + "WakingCheckStage", + "WhitelistCheckStage", ] diff --git a/astrbot/core/pipeline/content_safety_check/stage.py b/astrbot/core/pipeline/content_safety_check/stage.py index e6ecd995..c477cc23 100644 --- a/astrbot/core/pipeline/content_safety_check/stage.py +++ b/astrbot/core/pipeline/content_safety_check/stage.py @@ -1,9 +1,11 @@ -from typing import Union, AsyncGenerator -from ..stage import Stage, register_stage -from ..context import PipelineContext -from astrbot.core.platform.astr_message_event import AstrMessageEvent -from astrbot.core.message.message_event_result import MessageEventResult +from collections.abc import AsyncGenerator + from astrbot.core import logger +from astrbot.core.message.message_event_result import MessageEventResult +from astrbot.core.platform.astr_message_event import AstrMessageEvent + +from ..context import PipelineContext +from ..stage import Stage, register_stage from .strategies.strategy import StrategySelector @@ -19,8 +21,10 @@ class ContentSafetyCheckStage(Stage): self.strategy_selector = StrategySelector(config) async def process( - self, event: AstrMessageEvent, check_text: str | None = None - ) -> Union[None, AsyncGenerator[None, None]]: + self, + event: AstrMessageEvent, + check_text: str | None = None, + ) -> None | AsyncGenerator[None, None]: """检查内容安全""" text = check_text if check_text else event.get_message_str() ok, info = self.strategy_selector.check(text) @@ -28,8 +32,8 @@ class ContentSafetyCheckStage(Stage): if event.is_at_or_wake_command: event.set_result( MessageEventResult().message( - "你的消息或者大模型的响应中包含不适当的内容，已被屏蔽。" - ) + "你的消息或者大模型的响应中包含不适当的内容，已被屏蔽。", + ), ) yield event.stop_event() diff --git a/astrbot/core/pipeline/content_safety_check/strategies/__init__.py b/astrbot/core/pipeline/content_safety_check/strategies/__init__.py index 5701f063..f0a34e73 100644 --- a/astrbot/core/pipeline/content_safety_check/strategies/__init__.py +++ b/astrbot/core/pipeline/content_safety_check/strategies/__init__.py @@ -1,8 +1,7 @@ import abc -from typing import Tuple class ContentSafetyStrategy(abc.ABC): @abc.abstractmethod - def check(self, content: str) -> Tuple[bool, str]: + def check(self, content: str) -> tuple[bool, str]: raise NotImplementedError diff --git a/astrbot/core/pipeline/content_safety_check/strategies/baidu_aip.py b/astrbot/core/pipeline/content_safety_check/strategies/baidu_aip.py index 26284e1a..bfa82de0 100644 --- a/astrbot/core/pipeline/content_safety_check/strategies/baidu_aip.py +++ b/astrbot/core/pipeline/content_safety_check/strategies/baidu_aip.py @@ -1,9 +1,8 @@ -""" -使用此功能应该先 pip install baidu-aip -""" +"""使用此功能应该先 pip install baidu-aip""" + +from aip import AipContentCensor from . import ContentSafetyStrategy -from aip import AipContentCensor class BaiduAipStrategy(ContentSafetyStrategy): @@ -19,12 +18,12 @@ class BaiduAipStrategy(ContentSafetyStrategy): return False, "" if res["conclusionType"] == 1: return True, "" - else: - if "data" not in res: - return False, "" - count = len(res["data"]) - info = f"百度审核服务发现 {count} 处违规：\n" - for i in res["data"]: - info += f"{i['msg']}；\n" - info += "\n判断结果：" + res["conclusion"] - return False, info + if "data" not in res: + return False, "" + count = len(res["data"]) + parts = [f"百度审核服务发现 {count} 处违规：\n"] + for i in res["data"]: + parts.append(f"{i['msg']}；\n") + parts.append("\n判断结果：" + res["conclusion"]) + info = "".join(parts) + return False, info diff --git a/astrbot/core/pipeline/content_safety_check/strategies/keywords.py b/astrbot/core/pipeline/content_safety_check/strategies/keywords.py index c65faa00..53ad900f 100644 --- a/astrbot/core/pipeline/content_safety_check/strategies/keywords.py +++ b/astrbot/core/pipeline/content_safety_check/strategies/keywords.py @@ -1,4 +1,5 @@ import re + from . import ContentSafetyStrategy diff --git a/astrbot/core/pipeline/content_safety_check/strategies/strategy.py b/astrbot/core/pipeline/content_safety_check/strategies/strategy.py index af960328..c971ef26 100644 --- a/astrbot/core/pipeline/content_safety_check/strategies/strategy.py +++ b/astrbot/core/pipeline/content_safety_check/strategies/strategy.py @@ -1,16 +1,16 @@ -from . import ContentSafetyStrategy -from typing import List, Tuple from astrbot import logger +from . import ContentSafetyStrategy + class StrategySelector: def __init__(self, config: dict) -> None: - self.enabled_strategies: List[ContentSafetyStrategy] = [] + self.enabled_strategies: list[ContentSafetyStrategy] = [] if config["internal_keywords"]["enable"]: from .keywords import KeywordsStrategy self.enabled_strategies.append( - KeywordsStrategy(config["internal_keywords"]["extra_keywords"]) + KeywordsStrategy(config["internal_keywords"]["extra_keywords"]), ) if config["baidu_aip"]["enable"]: try: @@ -23,10 +23,10 @@ class StrategySelector: config["baidu_aip"]["app_id"], config["baidu_aip"]["api_key"], config["baidu_aip"]["secret_key"], - ) + ), ) - def check(self, content: str) -> Tuple[bool, str]: + def check(self, content: str) -> tuple[bool, str]: for strategy in self.enabled_strategies: ok, info = strategy.check(content) if not ok: diff --git a/astrbot/core/pipeline/context.py b/astrbot/core/pipeline/context.py index 803626aa..a6cd567e 100644 --- a/astrbot/core/pipeline/context.py +++ b/astrbot/core/pipeline/context.py @@ -1,7 +1,9 @@ from dataclasses import dataclass + from astrbot.core.config import AstrBotConfig from astrbot.core.star import PluginManager -from .context_utils import call_handler, call_event_hook + +from .context_utils import call_event_hook, call_handler @dataclass diff --git a/astrbot/core/pipeline/context_utils.py b/astrbot/core/pipeline/context_utils.py index e7ac120b..73d28c5d 100644 --- a/astrbot/core/pipeline/context_utils.py +++ b/astrbot/core/pipeline/context_utils.py @@ -1,11 +1,12 @@ import inspect import traceback import typing as T + from astrbot import logger -from astrbot.core.star.star_handler import star_handlers_registry, EventType -from astrbot.core.star.star import star_map -from astrbot.core.message.message_event_result import MessageEventResult, CommandResult +from astrbot.core.message.message_event_result import CommandResult, MessageEventResult from astrbot.core.platform.astr_message_event import AstrMessageEvent +from astrbot.core.star.star import star_map +from astrbot.core.star.star_handler import EventType, star_handlers_registry async def call_handler( @@ -26,6 +27,7 @@ async def call_handler( Returns: AsyncGenerator[None, None]: 异步生成器，用于在管道中传递控制流 + """ ready_to_call = None # 一个协程或者异步生成器 @@ -80,14 +82,17 @@ async def call_event_hook( Returns: bool: 如果事件被终止，返回 True - #""" + # + + """ handlers = star_handlers_registry.get_handlers_by_event_type( - hook_type, plugins_name=event.plugins_name + hook_type, + plugins_name=event.plugins_name, ) for handler in handlers: try: logger.debug( - f"hook({hook_type.name}) -> {star_map[handler.handler_module_path].name} - {handler.handler_name}" + f"hook({hook_type.name}) -> {star_map[handler.handler_module_path].name} - {handler.handler_name}", ) await handler.handler(event, *args, **kwargs) except BaseException: @@ -95,7 +100,7 @@ async def call_event_hook( if event.is_stopped(): logger.info( - f"{star_map[handler.handler_module_path].name} - {handler.handler_name} 终止了事件传播。" + f"{star_map[handler.handler_module_path].name} - {handler.handler_name} 终止了事件传播。", ) return True diff --git a/astrbot/core/pipeline/preprocess_stage/stage.py b/astrbot/core/pipeline/preprocess_stage/stage.py index 5c075687..a69d07ff 100644 --- a/astrbot/core/pipeline/preprocess_stage/stage.py +++ b/astrbot/core/pipeline/preprocess_stage/stage.py @@ -1,12 +1,14 @@ -import traceback import asyncio import random -from typing import Union, AsyncGenerator -from ..stage import Stage, register_stage -from ..context import PipelineContext -from astrbot.core.platform.astr_message_event import AstrMessageEvent +import traceback +from collections.abc import AsyncGenerator + from astrbot.core import logger -from astrbot.core.message.components import Plain, Record, Image +from astrbot.core.message.components import Image, Plain, Record +from astrbot.core.platform.astr_message_event import AstrMessageEvent + +from ..context import PipelineContext +from ..stage import Stage, register_stage @register_stage @@ -20,8 +22,9 @@ class PreProcessStage(Stage): self.platform_settings: dict = self.config.get("platform_settings", {}) async def process( - self, event: AstrMessageEvent - ) -> Union[None, AsyncGenerator[None, None]]: + self, + event: AstrMessageEvent, + ) -> None | AsyncGenerator[None, None]: """在处理事件之前的预处理""" # 平台特异配置：platform_specific..pre_ack_emoji supported = {"telegram", "lark"} @@ -68,7 +71,7 @@ class PreProcessStage(Stage): stt_provider = ctx.get_using_stt_provider(event.unified_msg_origin) if not stt_provider: logger.warning( - f"会话 {event.unified_msg_origin} 未配置语音转文本模型。" + f"会话 {event.unified_msg_origin} 未配置语音转文本模型。", ) return message_chain = event.get_messages() diff --git a/astrbot/core/pipeline/process_stage/method/llm_request.py b/astrbot/core/pipeline/process_stage/method/llm_request.py index 85eaf9ff..bd9e4ce3 100644 --- a/astrbot/core/pipeline/process_stage/method/llm_request.py +++ b/astrbot/core/pipeline/process_stage/method/llm_request.py @@ -1,15 +1,14 @@ -""" -本地 Agent 模式的 LLM 调用 Stage -""" +"""本地 Agent 模式的 LLM 调用 Stage""" import asyncio import copy import json -import traceback -from datetime import timedelta -from typing import AsyncGenerator, Union -from astrbot.core.conversation_mgr import Conversation +from collections.abc import AsyncGenerator + from astrbot.core import logger +from astrbot.core.agent.tool import ToolSet +from astrbot.core.astr_agent_context import AstrAgentContext +from astrbot.core.conversation_mgr import Conversation from astrbot.core.message.components import Image from astrbot.core.message.message_event_result import ( MessageChain, @@ -22,291 +21,18 @@ from astrbot.core.provider.entities import ( LLMResponse, ProviderRequest, ) -from astrbot.core.agent.hooks import BaseAgentRunHooks -from astrbot.core.agent.runners.tool_loop_agent_runner import ToolLoopAgentRunner -from astrbot.core.agent.run_context import ContextWrapper -from astrbot.core.agent.tool import ToolSet, FunctionTool -from astrbot.core.agent.tool_executor import BaseFunctionToolExecutor -from astrbot.core.agent.handoff import HandoffTool from astrbot.core.star.session_llm_manager import SessionServiceManager -from astrbot.core.star.star_handler import EventType +from astrbot.core.star.star_handler import EventType, star_map from astrbot.core.utils.metrics import Metric -from ...context import PipelineContext, call_event_hook, call_handler +from astrbot.core.utils.session_lock import session_lock_manager + +from ....astr_agent_context import AgentContextWrapper +from ....astr_agent_hooks import MAIN_AGENT_HOOKS +from ....astr_agent_run_util import AgentRunner, run_agent +from ....astr_agent_tool_exec import FunctionToolExecutor +from ...context import PipelineContext, call_event_hook from ..stage import Stage -from astrbot.core.provider.register import llm_tools -from astrbot.core.star.star_handler import star_map -from astrbot.core.astr_agent_context import AstrAgentContext - -try: - import mcp -except (ModuleNotFoundError, ImportError): - logger.warning("警告: 缺少依赖库 'mcp'，将无法使用 MCP 服务。") - - -AgentContextWrapper = ContextWrapper[AstrAgentContext] -AgentRunner = ToolLoopAgentRunner[AstrAgentContext] - - -class FunctionToolExecutor(BaseFunctionToolExecutor[AstrAgentContext]): - @classmethod - async def execute(cls, tool, run_context, **tool_args): - """执行函数调用。 - - Args: - event (AstrMessageEvent): 事件对象, 当 origin 为 local 时必须提供。 - **kwargs: 函数调用的参数。 - - Returns: - AsyncGenerator[None | mcp.types.CallToolResult, None] - """ - if isinstance(tool, HandoffTool): - async for r in cls._execute_handoff(tool, run_context, **tool_args): - yield r - return - - if tool.origin == "local": - async for r in cls._execute_local(tool, run_context, **tool_args): - yield r - return - - elif tool.origin == "mcp": - async for r in cls._execute_mcp(tool, run_context, **tool_args): - yield r - return - - raise Exception(f"Unknown function origin: {tool.origin}") - - @classmethod - async def _execute_handoff( - cls, - tool: HandoffTool, - run_context: ContextWrapper[AstrAgentContext], - **tool_args, - ): - input_ = tool_args.get("input", "agent") - agent_runner = AgentRunner() - - # make toolset for the agent - tools = tool.agent.tools - if tools: - toolset = ToolSet() - for t in tools: - if isinstance(t, str): - _t = llm_tools.get_func(t) - if _t: - toolset.add_tool(_t) - elif isinstance(t, FunctionTool): - toolset.add_tool(t) - else: - toolset = None - - request = ProviderRequest( - prompt=input_, - system_prompt=tool.description or "", - image_urls=[], # 暂时不传递原始 agent 的上下文 - contexts=[], # 暂时不传递原始 agent 的上下文 - func_tool=toolset, - ) - astr_agent_ctx = AstrAgentContext( - provider=run_context.context.provider, - first_provider_request=run_context.context.first_provider_request, - curr_provider_request=request, - streaming=run_context.context.streaming, - ) - - logger.debug(f"正在将任务委托给 Agent: {tool.agent.name}, input: {input_}") - await run_context.event.send( - MessageChain().message("✨ 正在将任务委托给 Agent: " + tool.agent.name) - ) - - await agent_runner.reset( - provider=run_context.context.provider, - request=request, - run_context=AgentContextWrapper( - context=astr_agent_ctx, event=run_context.event - ), - tool_executor=FunctionToolExecutor(), - agent_hooks=tool.agent.run_hooks or BaseAgentRunHooks[AstrAgentContext](), - streaming=run_context.context.streaming, - ) - - async for _ in run_agent(agent_runner, 15, True): - pass - - if agent_runner.done(): - llm_response = agent_runner.get_final_llm_resp() - - if not llm_response: - text_content = mcp.types.TextContent( - type="text", - text=f"error when deligate task to {tool.agent.name}", - ) - yield mcp.types.CallToolResult(content=[text_content]) - return - - logger.debug( - f"Agent {tool.agent.name} 任务完成, response: {llm_response.completion_text}" - ) - - result = ( - f"Agent {tool.agent.name} respond with: {llm_response.completion_text}\n\n" - "Note: If the result is error or need user provide more information, please provide more information to the agent(you can ask user for more information first)." - ) - - text_content = mcp.types.TextContent( - type="text", - text=result, - ) - yield mcp.types.CallToolResult(content=[text_content]) - else: - text_content = mcp.types.TextContent( - type="text", - text=f"error when deligate task to {tool.agent.name}", - ) - yield mcp.types.CallToolResult(content=[text_content]) - return - - @classmethod - async def _execute_local( - cls, - tool: FunctionTool, - run_context: ContextWrapper[AstrAgentContext], - **tool_args, - ): - if not run_context.event: - raise ValueError("Event must be provided for local function tools.") - - # 检查 tool 下有没有 run 方法 - if not tool.handler and not hasattr(tool, "run"): - raise ValueError("Tool must have a valid handler or 'run' method.") - awaitable = tool.handler or getattr(tool, "run") - - wrapper = call_handler( - event=run_context.event, - handler=awaitable, - **tool_args, - ) - # async for resp in wrapper: - while True: - try: - resp = await asyncio.wait_for( - anext(wrapper), - timeout=run_context.context.tool_call_timeout, - ) - if resp is not None: - if isinstance(resp, mcp.types.CallToolResult): - yield resp - else: - text_content = mcp.types.TextContent( - type="text", - text=str(resp), - ) - yield mcp.types.CallToolResult(content=[text_content]) - else: - # NOTE: Tool 在这里直接请求发送消息给用户 - # TODO: 是否需要判断 event.get_result() 是否为空? - # 如果为空,则说明没有发送消息给用户,并且返回值为空,将返回一个特殊的 TextContent,其内容如"工具没有返回内容" - yield None - except asyncio.TimeoutError: - raise Exception( - f"tool {tool.name} execution timeout after {run_context.context.tool_call_timeout} seconds." - ) - except StopAsyncIteration: - break - - @classmethod - async def _execute_mcp( - cls, - tool: FunctionTool, - run_context: ContextWrapper[AstrAgentContext], - **tool_args, - ): - if not tool.mcp_client: - raise ValueError("MCP client is not available for MCP function tools.") - - session = tool.mcp_client.session - if not session: - raise ValueError("MCP session is not available for MCP function tools.") - res = await session.call_tool( - name=tool.name, - arguments=tool_args, - read_timeout_seconds=timedelta( - seconds=run_context.context.tool_call_timeout - ), - ) - if not res: - return - yield res - - -class MainAgentHooks(BaseAgentRunHooks[AstrAgentContext]): - async def on_agent_done(self, run_context, llm_response): - # 执行事件钩子 - await call_event_hook( - run_context.event, EventType.OnLLMResponseEvent, llm_response - ) - - -MAIN_AGENT_HOOKS = MainAgentHooks() - - -async def run_agent( - agent_runner: AgentRunner, max_step: int = 30, show_tool_use: bool = True -) -> AsyncGenerator[MessageChain, None]: - step_idx = 0 - astr_event = agent_runner.run_context.event - while step_idx < max_step: - step_idx += 1 - try: - async for resp in agent_runner.step(): - if astr_event.is_stopped(): - return - if resp.type == "tool_call_result": - msg_chain = resp.data["chain"] - if msg_chain.type == "tool_direct_result": - # tool_direct_result 用于标记 llm tool 需要直接发送给用户的内容 - resp.data["chain"].type = "tool_call_result" - await astr_event.send(resp.data["chain"]) - continue - # 对于其他情况，暂时先不处理 - continue - elif resp.type == "tool_call": - if agent_runner.streaming: - # 用来标记流式响应需要分节 - yield MessageChain(chain=[], type="break") - if show_tool_use or astr_event.get_platform_name() == "webchat": - resp.data["chain"].type = "tool_call" - await astr_event.send(resp.data["chain"]) - continue - - if not agent_runner.streaming: - content_typ = ( - ResultContentType.LLM_RESULT - if resp.type == "llm_result" - else ResultContentType.GENERAL_RESULT - ) - astr_event.set_result( - MessageEventResult( - chain=resp.data["chain"].chain, - result_content_type=content_typ, - ) - ) - yield - astr_event.clear_result() - else: - if resp.type == "streaming_delta": - yield resp.data["chain"] # MessageChain - if agent_runner.done(): - break - - except Exception as e: - logger.error(traceback.format_exc()) - err_msg = f"\n\nAstrBot 请求失败。\n错误类型: {type(e).__name__}\n错误信息: {str(e)}\n\n请在控制台查看和分享错误详情。\n" - if agent_runner.streaming: - yield MessageChain().message(err_msg) - else: - astr_event.set_result(MessageEventResult().message(err_msg)) - return +from ..utils import KNOWLEDGE_BASE_QUERY_TOOL, retrieve_knowledge_base class LLMRequestSubStage(Stage): @@ -322,16 +48,21 @@ class LLMRequestSubStage(Stage): self.max_context_length - 1, ) self.streaming_response: bool = settings["streaming_response"] + self.unsupported_streaming_strategy: str = settings[ + "unsupported_streaming_strategy" + ] self.max_step: int = settings.get("max_agent_step", 30) self.tool_call_timeout: int = settings.get("tool_call_timeout", 60) if isinstance(self.max_step, bool): # workaround: #2622 self.max_step = 30 self.show_tool_use: bool = settings.get("show_tool_use_status", True) + self.show_reasoning = settings.get("display_reasoning_text", False) + self.kb_agentic_mode: bool = conf.get("kb_agentic_mode", False) for bwp in self.bot_wake_prefixs: if self.provider_wake_prefix.startswith(bwp): logger.info( - f"识别 LLM 聊天额外唤醒前缀 {self.provider_wake_prefix} 以机器人唤醒前缀 {bwp} 开头，已自动去除。" + f"识别 LLM 聊天额外唤醒前缀 {self.provider_wake_prefix} 以机器人唤醒前缀 {bwp} 开头，已自动去除。", ) self.provider_wake_prefix = self.provider_wake_prefix[len(bwp) :] @@ -365,97 +96,68 @@ class LLMRequestSubStage(Stage): raise RuntimeError("无法创建新的对话。") return conversation - async def process( - self, event: AstrMessageEvent, _nested: bool = False - ) -> Union[None, AsyncGenerator[None, None]]: - req: ProviderRequest | None = None - - if not self.ctx.astrbot_config["provider_settings"]["enable"]: - logger.debug("未启用 LLM 能力，跳过处理。") - return - - # 检查会话级别的LLM启停状态 - if not SessionServiceManager.should_process_llm_request(event): - logger.debug(f"会话 {event.unified_msg_origin} 禁用了 LLM，跳过处理。") - return - - provider = self._select_provider(event) - if provider is None: - return - if not isinstance(provider, Provider): - logger.error(f"选择的提供商类型无效（{type(provider)}），跳过 LLM 请求处理。") - return - - if event.get_extra("provider_request"): - req = event.get_extra("provider_request") - assert isinstance(req, ProviderRequest), ( - "provider_request 必须是 ProviderRequest 类型。" - ) - - if req.conversation: - req.contexts = json.loads(req.conversation.history) - - else: - req = ProviderRequest(prompt="", image_urls=[]) - if sel_model := event.get_extra("selected_model"): - req.model = sel_model - if self.provider_wake_prefix: - if not event.message_str.startswith(self.provider_wake_prefix): + async def _apply_kb( + self, + event: AstrMessageEvent, + req: ProviderRequest, + ): + """Apply knowledge base context to the provider request""" + if not self.kb_agentic_mode: + if req.prompt is None: + return + try: + kb_result = await retrieve_knowledge_base( + query=req.prompt, + umo=event.unified_msg_origin, + context=self.ctx.plugin_manager.context, + ) + if not kb_result: return - req.prompt = event.message_str[len(self.provider_wake_prefix) :] - # func_tool selection 现在已经转移到 packages/astrbot 插件中进行选择。 - # req.func_tool = self.ctx.plugin_manager.context.get_llm_tool_manager() - for comp in event.message_obj.message: - if isinstance(comp, Image): - image_path = await comp.convert_to_file_path() - req.image_urls.append(image_path) + if req.system_prompt is not None: + req.system_prompt += ( + f"\n\n[Related Knowledge Base Results]:\n{kb_result}" + ) + except Exception as e: + logger.error(f"Error occurred while retrieving knowledge base: {e}") + else: + if req.func_tool is None: + req.func_tool = ToolSet() + req.func_tool.add_tool(KNOWLEDGE_BASE_QUERY_TOOL) - conversation = await self._get_session_conv(event) - req.conversation = conversation - req.contexts = json.loads(conversation.history) + def _truncate_contexts( + self, + contexts: list[dict], + ) -> list[dict]: + """截断上下文列表，确保不超过最大长度""" + if self.max_context_length == -1: + return contexts - event.set_extra("provider_request", req) + if len(contexts) // 2 <= self.max_context_length: + return contexts - if not req.prompt and not req.image_urls: - return + truncated_contexts = contexts[ + -(self.max_context_length - self.dequeue_context_length + 1) * 2 : + ] + # 找到第一个role 为 user 的索引，确保上下文格式正确 + index = next( + ( + i + for i, item in enumerate(truncated_contexts) + if item.get("role") == "user" + ), + None, + ) + if index is not None and index > 0: + truncated_contexts = truncated_contexts[index:] - # 执行请求 LLM 前事件钩子。 - if await call_event_hook(event, EventType.OnLLMRequestEvent, req): - return + return truncated_contexts - if isinstance(req.contexts, str): - req.contexts = json.loads(req.contexts) - - # max context length - if ( - self.max_context_length != -1 # -1 为不限制 - and len(req.contexts) // 2 > self.max_context_length - ): - logger.debug("上下文长度超过限制，将截断。") - req.contexts = req.contexts[ - -(self.max_context_length - self.dequeue_context_length + 1) * 2 : - ] - # 找到第一个role 为 user 的索引，确保上下文格式正确 - index = next( - ( - i - for i, item in enumerate(req.contexts) - if item.get("role") == "user" - ), - None, - ) - if index is not None and index > 0: - req.contexts = req.contexts[index:] - - # session_id - if not req.session_id: - req.session_id = event.unified_msg_origin - - # fix messages - req.contexts = self.fix_messages(req.contexts) - - # check provider modalities - # 如果提供商不支持图像/工具使用，但请求中包含图像/工具列表，则清空。图片转述等的检测和调用发生在这之前，因此这里可以这样处理。 + def _modalities_fix( + self, + provider: Provider, + req: ProviderRequest, + ): + """检查提供商的模态能力，清理请求中的不支持内容""" if req.image_urls: provider_cfg = provider.provider_config.get("modalities", ["image"]) if "image" not in provider_cfg: @@ -466,10 +168,16 @@ class LLMRequestSubStage(Stage): # 如果模型不支持工具使用，但请求中包含工具列表，则清空。 if "tool_use" not in provider_cfg: logger.debug( - f"用户设置提供商 {provider} 不支持工具使用，清空工具列表。" + f"用户设置提供商 {provider} 不支持工具使用，清空工具列表。", ) req.func_tool = None - # 插件可用性设置 + + def _plugin_tool_fix( + self, + event: AstrMessageEvent, + req: ProviderRequest, + ): + """根据事件中的插件设置，过滤请求中的工具列表""" if event.plugins_name is not None and req.func_tool: new_tool_set = ToolSet() for tool in req.func_tool.tools: @@ -483,85 +191,18 @@ class LLMRequestSubStage(Stage): new_tool_set.add_tool(tool) req.func_tool = new_tool_set - # 备份 req.contexts - backup_contexts = copy.deepcopy(req.contexts) - - # run agent - agent_runner = AgentRunner() - logger.debug( - f"handle provider[id: {provider.provider_config['id']}] request: {req}" - ) - astr_agent_ctx = AstrAgentContext( - provider=provider, - first_provider_request=req, - curr_provider_request=req, - streaming=self.streaming_response, - tool_call_timeout=self.tool_call_timeout, - ) - await agent_runner.reset( - provider=provider, - request=req, - run_context=AgentContextWrapper(context=astr_agent_ctx, event=event), - tool_executor=FunctionToolExecutor(), - agent_hooks=MAIN_AGENT_HOOKS, - streaming=self.streaming_response, - ) - - if self.streaming_response: - # 流式响应 - event.set_result( - MessageEventResult() - .set_result_content_type(ResultContentType.STREAMING_RESULT) - .set_async_stream( - run_agent(agent_runner, self.max_step, self.show_tool_use) - ) - ) - yield - if agent_runner.done(): - if final_llm_resp := agent_runner.get_final_llm_resp(): - if final_llm_resp.completion_text: - chain = ( - MessageChain().message(final_llm_resp.completion_text).chain - ) - elif final_llm_resp.result_chain: - chain = final_llm_resp.result_chain.chain - else: - chain = MessageChain().chain - event.set_result( - MessageEventResult( - chain=chain, - result_content_type=ResultContentType.STREAMING_FINISH, - ) - ) - else: - async for _ in run_agent(agent_runner, self.max_step, self.show_tool_use): - yield - - # 恢复备份的 contexts - req.contexts = backup_contexts - - await self._save_to_history(event, req, agent_runner.get_final_llm_resp()) - - # 异步处理 WebChat 特殊情况 - if event.get_platform_name() == "webchat": - asyncio.create_task(self._handle_webchat(event, req, provider)) - - asyncio.create_task( - Metric.upload( - llm_tick=1, - model_name=agent_runner.provider.get_model(), - provider_type=agent_runner.provider.meta().type, - ) - ) - async def _handle_webchat( - self, event: AstrMessageEvent, req: ProviderRequest, prov: Provider + self, + event: AstrMessageEvent, + req: ProviderRequest, + prov: Provider, ): """处理 WebChat 平台的特殊情况，包括第一次 LLM 对话时总结对话内容生成 title""" if not req.conversation: return conversation = await self.conv_manager.get_conversation( - event.unified_msg_origin, req.conversation.cid + event.unified_msg_origin, + req.conversation.cid, ) if conversation and not req.conversation.title: messages = json.loads(conversation.history) @@ -597,9 +238,6 @@ class LLMRequestSubStage(Stage): ), ) if llm_resp and llm_resp.completion_text: - logger.debug( - f"WebChat 对话标题生成响应: {llm_resp.completion_text.strip()}" - ) title = llm_resp.completion_text.strip() if not title or "" in title: return @@ -627,6 +265,9 @@ class LLMRequestSubStage(Stage): logger.debug("LLM 响应为空，不保存记录。") return + if req.contexts is None: + req.contexts = [] + # 历史上下文 messages = copy.deepcopy(req.contexts) # 这一轮对话请求的用户输入 @@ -641,10 +282,12 @@ class LLMRequestSubStage(Stage): messages.append({"role": "assistant", "content": llm_response.completion_text}) messages = list(filter(lambda item: "_no_save" not in item, messages)) await self.conv_manager.update_conversation( - event.unified_msg_origin, req.conversation.cid, history=messages + event.unified_msg_origin, + req.conversation.cid, + history=messages, ) - def fix_messages(self, messages: list[dict]) -> list[dict]: + def _fix_messages(self, messages: list[dict]) -> list[dict]: """验证并且修复上下文""" fixed_messages = [] for message in messages: @@ -659,3 +302,184 @@ class LLMRequestSubStage(Stage): else: fixed_messages.append(message) return fixed_messages + + async def process( + self, + event: AstrMessageEvent, + _nested: bool = False, + ) -> None | AsyncGenerator[None, None]: + req: ProviderRequest | None = None + + if not self.ctx.astrbot_config["provider_settings"]["enable"]: + logger.debug("未启用 LLM 能力，跳过处理。") + return + + # 检查会话级别的LLM启停状态 + if not SessionServiceManager.should_process_llm_request(event): + logger.debug(f"会话 {event.unified_msg_origin} 禁用了 LLM，跳过处理。") + return + + provider = self._select_provider(event) + if provider is None: + return + if not isinstance(provider, Provider): + logger.error(f"选择的提供商类型无效({type(provider)})，跳过 LLM 请求处理。") + return + + streaming_response = self.streaming_response + if (enable_streaming := event.get_extra("enable_streaming")) is not None: + streaming_response = bool(enable_streaming) + + logger.debug("ready to request llm provider") + async with session_lock_manager.acquire_lock(event.unified_msg_origin): + logger.debug("acquired session lock for llm request") + if event.get_extra("provider_request"): + req = event.get_extra("provider_request") + assert isinstance(req, ProviderRequest), ( + "provider_request 必须是 ProviderRequest 类型。" + ) + + if req.conversation: + req.contexts = json.loads(req.conversation.history) + + else: + req = ProviderRequest() + req.prompt = "" + req.image_urls = [] + if sel_model := event.get_extra("selected_model"): + req.model = sel_model + if self.provider_wake_prefix and not event.message_str.startswith( + self.provider_wake_prefix + ): + return + + req.prompt = event.message_str[len(self.provider_wake_prefix) :] + # func_tool selection 现在已经转移到 packages/astrbot 插件中进行选择。 + # req.func_tool = self.ctx.plugin_manager.context.get_llm_tool_manager() + for comp in event.message_obj.message: + if isinstance(comp, Image): + image_path = await comp.convert_to_file_path() + req.image_urls.append(image_path) + + conversation = await self._get_session_conv(event) + req.conversation = conversation + req.contexts = json.loads(conversation.history) + + event.set_extra("provider_request", req) + + if not req.prompt and not req.image_urls: + return + + # call event hook + if await call_event_hook(event, EventType.OnLLMRequestEvent, req): + return + + # apply knowledge base feature + await self._apply_kb(event, req) + + # fix contexts json str + if isinstance(req.contexts, str): + req.contexts = json.loads(req.contexts) + + # truncate contexts to fit max length + if req.contexts: + req.contexts = self._truncate_contexts(req.contexts) + self._fix_messages(req.contexts) + + # session_id + if not req.session_id: + req.session_id = event.unified_msg_origin + + # check provider modalities, if provider does not support image/tool_use, clear them in request. + self._modalities_fix(provider, req) + + # filter tools, only keep tools from this pipeline's selected plugins + self._plugin_tool_fix(event, req) + + stream_to_general = ( + self.unsupported_streaming_strategy == "turn_off" + and not event.platform_meta.support_streaming_message + ) + # 备份 req.contexts + backup_contexts = copy.deepcopy(req.contexts) + + # run agent + agent_runner = AgentRunner() + logger.debug( + f"handle provider[id: {provider.provider_config['id']}] request: {req}", + ) + astr_agent_ctx = AstrAgentContext( + context=self.ctx.plugin_manager.context, + event=event, + ) + await agent_runner.reset( + provider=provider, + request=req, + run_context=AgentContextWrapper( + context=astr_agent_ctx, + tool_call_timeout=self.tool_call_timeout, + ), + tool_executor=FunctionToolExecutor(), + agent_hooks=MAIN_AGENT_HOOKS, + streaming=streaming_response, + ) + + if streaming_response and not stream_to_general: + # 流式响应 + event.set_result( + MessageEventResult() + .set_result_content_type(ResultContentType.STREAMING_RESULT) + .set_async_stream( + run_agent( + agent_runner, + self.max_step, + self.show_tool_use, + show_reasoning=self.show_reasoning, + ), + ), + ) + yield + if agent_runner.done(): + if final_llm_resp := agent_runner.get_final_llm_resp(): + if final_llm_resp.completion_text: + chain = ( + MessageChain() + .message(final_llm_resp.completion_text) + .chain + ) + elif final_llm_resp.result_chain: + chain = final_llm_resp.result_chain.chain + else: + chain = MessageChain().chain + event.set_result( + MessageEventResult( + chain=chain, + result_content_type=ResultContentType.STREAMING_FINISH, + ), + ) + else: + async for _ in run_agent( + agent_runner, + self.max_step, + self.show_tool_use, + stream_to_general, + show_reasoning=self.show_reasoning, + ): + yield + + # 恢复备份的 contexts + req.contexts = backup_contexts + + await self._save_to_history(event, req, agent_runner.get_final_llm_resp()) + + # 异步处理 WebChat 特殊情况 + if event.get_platform_name() == "webchat": + asyncio.create_task(self._handle_webchat(event, req, provider)) + + asyncio.create_task( + Metric.upload( + llm_tick=1, + model_name=agent_runner.provider.get_model(), + provider_type=agent_runner.provider.meta().type, + ), + ) diff --git a/astrbot/core/pipeline/process_stage/method/star_request.py b/astrbot/core/pipeline/process_stage/method/star_request.py index 42990aae..ff8120b1 100644 --- a/astrbot/core/pipeline/process_stage/method/star_request.py +++ b/astrbot/core/pipeline/process_stage/method/star_request.py @@ -1,16 +1,17 @@ -""" -本地 Agent 模式的 AstrBot 插件调用 Stage -""" +"""本地 Agent 模式的 AstrBot 插件调用 Stage""" + +import traceback +from collections.abc import AsyncGenerator +from typing import Any + +from astrbot.core import logger +from astrbot.core.message.message_event_result import MessageEventResult +from astrbot.core.platform.astr_message_event import AstrMessageEvent +from astrbot.core.star.star import star_map +from astrbot.core.star.star_handler import StarHandlerMetadata from ...context import PipelineContext, call_handler from ..stage import Stage -from typing import Dict, Any, List, AsyncGenerator, Union -from astrbot.core.platform.astr_message_event import AstrMessageEvent -from astrbot.core.message.message_event_result import MessageEventResult -from astrbot.core import logger -from astrbot.core.star.star_handler import StarHandlerMetadata -from astrbot.core.star.star import star_map -import traceback class StarRequestSubStage(Stage): @@ -21,13 +22,14 @@ class StarRequestSubStage(Stage): self.ctx = ctx async def process( - self, event: AstrMessageEvent - ) -> Union[None, AsyncGenerator[None, None]]: - activated_handlers: List[StarHandlerMetadata] = event.get_extra( - "activated_handlers" + self, + event: AstrMessageEvent, + ) -> None | AsyncGenerator[None, None]: + activated_handlers: list[StarHandlerMetadata] = event.get_extra( + "activated_handlers", ) - handlers_parsed_params: Dict[str, Dict[str, Any]] = event.get_extra( - "handlers_parsed_params" + handlers_parsed_params: dict[str, dict[str, Any]] = event.get_extra( + "handlers_parsed_params", ) if not handlers_parsed_params: handlers_parsed_params = {} @@ -37,7 +39,7 @@ class StarRequestSubStage(Stage): md = star_map.get(handler.handler_module_path) if not md: logger.warning( - f"Cannot find plugin for given handler module path: {handler.handler_module_path}" + f"Cannot find plugin for given handler module path: {handler.handler_module_path}", ) continue logger.debug(f"plugin -> {md.name} - {handler.handler_name}") diff --git a/astrbot/core/pipeline/process_stage/stage.py b/astrbot/core/pipeline/process_stage/stage.py index f653a9fb..9f0b5f92 100644 --- a/astrbot/core/pipeline/process_stage/stage.py +++ b/astrbot/core/pipeline/process_stage/stage.py @@ -1,12 +1,14 @@ -from typing import List, Union, AsyncGenerator -from ..stage import Stage, register_stage +from collections.abc import AsyncGenerator + +from astrbot.core import logger +from astrbot.core.platform.astr_message_event import AstrMessageEvent +from astrbot.core.provider.entities import ProviderRequest +from astrbot.core.star.star_handler import StarHandlerMetadata + from ..context import PipelineContext +from ..stage import Stage, register_stage from .method.llm_request import LLMRequestSubStage from .method.star_request import StarRequestSubStage -from astrbot.core.platform.astr_message_event import AstrMessageEvent -from astrbot.core.star.star_handler import StarHandlerMetadata -from astrbot.core.provider.entities import ProviderRequest -from astrbot.core import logger @register_stage @@ -22,11 +24,12 @@ class ProcessStage(Stage): await self.star_request_sub_stage.initialize(ctx) async def process( - self, event: AstrMessageEvent - ) -> Union[None, AsyncGenerator[None, None]]: + self, + event: AstrMessageEvent, + ) -> None | AsyncGenerator[None, None]: """处理事件""" - activated_handlers: List[StarHandlerMetadata] = event.get_extra( - "activated_handlers" + activated_handlers: list[StarHandlerMetadata] = event.get_extra( + "activated_handlers", ) # 有插件 Handler 被激活 if activated_handlers: diff --git a/astrbot/core/pipeline/process_stage/utils.py b/astrbot/core/pipeline/process_stage/utils.py new file mode 100644 index 00000000..24e052e1 --- /dev/null +++ b/astrbot/core/pipeline/process_stage/utils.py @@ -0,0 +1,125 @@ +from pydantic import Field +from pydantic.dataclasses import dataclass + +from astrbot.api import logger, sp +from astrbot.core.agent.run_context import ContextWrapper +from astrbot.core.agent.tool import FunctionTool, ToolExecResult +from astrbot.core.astr_agent_context import AstrAgentContext +from astrbot.core.star.context import Context + + +@dataclass +class KnowledgeBaseQueryTool(FunctionTool[AstrAgentContext]): + name: str = "astr_kb_search" + description: str = ( + "Query the knowledge base for facts or relevant context. " + "Use this tool when the user's question requires factual information, " + "definitions, background knowledge, or previously indexed content. " + "Only send short keywords or a concise question as the query." + ) + parameters: dict = Field( + default_factory=lambda: { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "A concise keyword query for the knowledge base.", + }, + }, + "required": ["query"], + } + ) + + async def call( + self, context: ContextWrapper[AstrAgentContext], **kwargs + ) -> ToolExecResult: + query = kwargs.get("query", "") + if not query: + return "error: Query parameter is empty." + result = await retrieve_knowledge_base( + query=kwargs.get("query", ""), + umo=context.context.event.unified_msg_origin, + context=context.context.context, + ) + if not result: + return "No relevant knowledge found." + return result + + +async def retrieve_knowledge_base( + query: str, + umo: str, + context: Context, +) -> str | None: + """Inject knowledge base context into the provider request + + Args: + umo: Unique message object (session ID) + p_ctx: Pipeline context + """ + kb_mgr = context.kb_manager + config = context.get_config(umo=umo) + + # 1. 优先读取会话级配置 + session_config = await sp.session_get(umo, "kb_config", default={}) + + if session_config and "kb_ids" in session_config: + # 会话级配置 + kb_ids = session_config.get("kb_ids", []) + + # 如果配置为空列表，明确表示不使用知识库 + if not kb_ids: + logger.info(f"[知识库] 会话 {umo} 已被配置为不使用知识库") + return + + top_k = session_config.get("top_k", 5) + + # 将 kb_ids 转换为 kb_names + kb_names = [] + invalid_kb_ids = [] + for kb_id in kb_ids: + kb_helper = await kb_mgr.get_kb(kb_id) + if kb_helper: + kb_names.append(kb_helper.kb.kb_name) + else: + logger.warning(f"[知识库] 知识库不存在或未加载: {kb_id}") + invalid_kb_ids.append(kb_id) + + if invalid_kb_ids: + logger.warning( + f"[知识库] 会话 {umo} 配置的以下知识库无效: {invalid_kb_ids}", + ) + + if not kb_names: + return + + logger.debug(f"[知识库] 使用会话级配置，知识库数量: {len(kb_names)}") + else: + kb_names = config.get("kb_names", []) + top_k = config.get("kb_final_top_k", 5) + logger.debug(f"[知识库] 使用全局配置，知识库数量: {len(kb_names)}") + + top_k_fusion = config.get("kb_fusion_top_k", 20) + + if not kb_names: + return + + logger.debug(f"[知识库] 开始检索知识库，数量: {len(kb_names)}, top_k={top_k}") + kb_context = await kb_mgr.retrieve( + query=query, + kb_names=kb_names, + top_k_fusion=top_k_fusion, + top_m_final=top_k, + ) + + if not kb_context: + return + + formatted = kb_context.get("context_text", "") + if formatted: + results = kb_context.get("results", []) + logger.debug(f"[知识库] 为会话 {umo} 注入了 {len(results)} 条相关知识块") + return formatted + + +KNOWLEDGE_BASE_QUERY_TOOL = KnowledgeBaseQueryTool() diff --git a/astrbot/core/pipeline/rate_limit_check/stage.py b/astrbot/core/pipeline/rate_limit_check/stage.py index b36a2fbd..64e21dd7 100644 --- a/astrbot/core/pipeline/rate_limit_check/stage.py +++ b/astrbot/core/pipeline/rate_limit_check/stage.py @@ -1,18 +1,19 @@ import asyncio -from datetime import datetime, timedelta from collections import defaultdict, deque -from typing import DefaultDict, Deque, Union, AsyncGenerator -from ..stage import Stage, register_stage -from ..context import PipelineContext -from astrbot.core.platform.astr_message_event import AstrMessageEvent +from collections.abc import AsyncGenerator +from datetime import datetime, timedelta + from astrbot.core import logger from astrbot.core.config.astrbot_config import RateLimitStrategy +from astrbot.core.platform.astr_message_event import AstrMessageEvent + +from ..context import PipelineContext +from ..stage import Stage, register_stage @register_stage class RateLimitStage(Stage): - """ - 检查是否需要限制消息发送的限流器。 + """检查是否需要限制消息发送的限流器。 使用 Fixed Window 算法。 如果触发限流，将 stall 流水线，直到下一个时间窗口来临时自动唤醒。 @@ -20,32 +21,30 @@ class RateLimitStage(Stage): def __init__(self): # 存储每个会话的请求时间队列 - self.event_timestamps: DefaultDict[str, Deque[datetime]] = defaultdict(deque) + self.event_timestamps: defaultdict[str, deque[datetime]] = defaultdict(deque) # 为每个会话设置一个锁，避免并发冲突 - self.locks: DefaultDict[str, asyncio.Lock] = defaultdict(asyncio.Lock) + self.locks: defaultdict[str, asyncio.Lock] = defaultdict(asyncio.Lock) # 限流参数 self.rate_limit_count: int = 0 self.rate_limit_time: timedelta = timedelta(0) async def initialize(self, ctx: PipelineContext) -> None: - """ - 初始化限流器，根据配置设置限流参数。 - """ + """初始化限流器，根据配置设置限流参数。""" self.rate_limit_count = ctx.astrbot_config["platform_settings"]["rate_limit"][ "count" ] self.rate_limit_time = timedelta( - seconds=ctx.astrbot_config["platform_settings"]["rate_limit"]["time"] + seconds=ctx.astrbot_config["platform_settings"]["rate_limit"]["time"], ) self.rl_strategy = ctx.astrbot_config["platform_settings"]["rate_limit"][ "strategy" ] # stall or discard async def process( - self, event: AstrMessageEvent - ) -> Union[None, AsyncGenerator[None, None]]: - """ - 检查并处理限流逻辑。如果触发限流，流水线会 stall 并在窗口期后自动恢复。 + self, + event: AstrMessageEvent, + ) -> None | AsyncGenerator[None, None]: + """检查并处理限流逻辑。如果触发限流，流水线会 stall 并在窗口期后自动恢复。 Args: event (AstrMessageEvent): 当前消息事件。 @@ -53,6 +52,7 @@ class RateLimitStage(Stage): Returns: MessageEventResult: 继续或停止事件处理的结果。 + """ session_id = event.session_id now = datetime.now() @@ -66,32 +66,33 @@ class RateLimitStage(Stage): if len(timestamps) < self.rate_limit_count: timestamps.append(now) break - else: - next_window_time = timestamps[0] + self.rate_limit_time - stall_duration = (next_window_time - now).total_seconds() + 0.3 + next_window_time = timestamps[0] + self.rate_limit_time + stall_duration = (next_window_time - now).total_seconds() + 0.3 - match self.rl_strategy: - case RateLimitStrategy.STALL.value: - logger.info( - f"会话 {session_id} 被限流。根据限流策略，此会话处理将被暂停 {stall_duration:.2f} 秒。" - ) - await asyncio.sleep(stall_duration) - now = datetime.now() - case RateLimitStrategy.DISCARD.value: - logger.info( - f"会话 {session_id} 被限流。根据限流策略，此请求已被丢弃，直到限额于 {stall_duration:.2f} 秒后重置。" - ) - return event.stop_event() + match self.rl_strategy: + case RateLimitStrategy.STALL.value: + logger.info( + f"会话 {session_id} 被限流。根据限流策略，此会话处理将被暂停 {stall_duration:.2f} 秒。", + ) + await asyncio.sleep(stall_duration) + now = datetime.now() + case RateLimitStrategy.DISCARD.value: + logger.info( + f"会话 {session_id} 被限流。根据限流策略，此请求已被丢弃，直到限额于 {stall_duration:.2f} 秒后重置。", + ) + return event.stop_event() def _remove_expired_timestamps( - self, timestamps: Deque[datetime], now: datetime + self, + timestamps: deque[datetime], + now: datetime, ) -> None: - """ - 移除时间窗口外的时间戳。 + """移除时间窗口外的时间戳。 Args: timestamps (Deque[datetime]): 当前会话的时间戳队列。 now (datetime): 当前时间，用于计算过期时间。 + """ expiry_threshold: datetime = now - self.rate_limit_time while timestamps and timestamps[0] < expiry_threshold: diff --git a/astrbot/core/pipeline/respond/stage.py b/astrbot/core/pipeline/respond/stage.py index dc6a67e2..76064956 100644 --- a/astrbot/core/pipeline/respond/stage.py +++ b/astrbot/core/pipeline/respond/stage.py @@ -1,17 +1,18 @@ -import random import asyncio import math +import random +from collections.abc import AsyncGenerator + import astrbot.core.message.components as Comp -from typing import Union, AsyncGenerator -from ..stage import register_stage, Stage -from ..context import PipelineContext, call_event_hook -from astrbot.core.platform.astr_message_event import AstrMessageEvent -from astrbot.core.message.message_event_result import MessageChain, ResultContentType from astrbot.core import logger from astrbot.core.message.components import BaseMessageComponent, ComponentType +from astrbot.core.message.message_event_result import MessageChain, ResultContentType +from astrbot.core.platform.astr_message_event import AstrMessageEvent from astrbot.core.star.star_handler import EventType from astrbot.core.utils.path_util import path_Mapping -from astrbot.core.utils.session_lock import session_lock_manager + +from ..context import PipelineContext, call_event_hook +from ..stage import Stage, register_stage @register_stage @@ -19,7 +20,7 @@ class RespondStage(Stage): # 组件类型到其非空判断函数的映射 _component_validators = { Comp.Plain: lambda comp: bool( - comp.text and comp.text.strip() + comp.text and comp.text.strip(), ), # 纯文本消息需要strip Comp.Face: lambda comp: comp.id is not None, # QQ表情 Comp.Record: lambda comp: bool(comp.file), # 语音 @@ -58,7 +59,7 @@ class RespondStage(Stage): "segmented_reply" ]["interval_method"] self.log_base = float( - ctx.astrbot_config["platform_settings"]["segmented_reply"]["log_base"] + ctx.astrbot_config["platform_settings"]["segmented_reply"]["log_base"], ) interval_str: str = ctx.astrbot_config["platform_settings"]["segmented_reply"][ "interval" @@ -86,17 +87,16 @@ class RespondStage(Stage): wc = await self._word_cnt(comp.text) i = math.log(wc + 1, self.log_base) return random.uniform(i, i + 0.5) - else: - return random.uniform(1, 1.75) - else: - # random - return random.uniform(self.interval[0], self.interval[1]) + return random.uniform(1, 1.75) + # random + return random.uniform(self.interval[0], self.interval[1]) async def _is_empty_message_chain(self, chain: list[BaseMessageComponent]): """检查消息链是否为空 Args: chain (list[BaseMessageComponent]): 包含消息对象的列表 + """ if not chain: return True @@ -150,8 +150,9 @@ class RespondStage(Stage): return extracted async def process( - self, event: AstrMessageEvent - ) -> Union[None, AsyncGenerator[None, None]]: + self, + event: AstrMessageEvent, + ) -> None | AsyncGenerator[None, None]: result = event.get_result() if result is None: return @@ -159,7 +160,7 @@ class RespondStage(Stage): return logger.info( - f"Prepare to send - {event.get_sender_name()}/{event.get_sender_id()}: {event._outline_chain(result.chain)}" + f"Prepare to send - {event.get_sender_name()}/{event.get_sender_id()}: {event._outline_chain(result.chain)}", ) if result.result_content_type == ResultContentType.STREAMING_RESULT: @@ -167,13 +168,17 @@ class RespondStage(Stage): logger.warning("async_stream 为空，跳过发送。") return # 流式结果直接交付平台适配器处理 - use_fallback = self.config.get("provider_settings", {}).get( - "streaming_segmented", False + realtime_segmenting = ( + self.config.get("provider_settings", {}).get( + "unsupported_streaming_strategy", + "realtime_segmenting", + ) + == "realtime_segmenting" ) logger.info(f"应用流式输出({event.get_platform_id()})") - await event.send_streaming(result.async_stream, use_fallback) + await event.send_streaming(result.async_stream, realtime_segmenting) return - elif len(result.chain) > 0: + if len(result.chain) > 0: # 检查路径映射 if mappings := self.platform_settings.get("path_mapping", []): for idx, component in enumerate(result.chain): @@ -212,24 +217,23 @@ class RespondStage(Stage): if not result.chain or len(result.chain) == 0: # may fix #2670 logger.warning( - f"实际消息链为空, 跳过发送阶段。header_chain: {header_comps}, actual_chain: {result.chain}" + f"实际消息链为空, 跳过发送阶段。header_chain: {header_comps}, actual_chain: {result.chain}", ) return - async with session_lock_manager.acquire_lock(event.unified_msg_origin): - for comp in result.chain: - i = await self._calc_comp_interval(comp) - await asyncio.sleep(i) - try: - if comp.type in need_separately: - await event.send(MessageChain([comp])) - else: - await event.send(MessageChain([*header_comps, comp])) - header_comps.clear() - except Exception as e: - logger.error( - f"发送消息链失败: chain = {MessageChain([comp])}, error = {e}", - exc_info=True, - ) + for comp in result.chain: + i = await self._calc_comp_interval(comp) + await asyncio.sleep(i) + try: + if comp.type in need_separately: + await event.send(MessageChain([comp])) + else: + await event.send(MessageChain([*header_comps, comp])) + header_comps.clear() + except Exception as e: + logger.error( + f"发送消息链失败: chain = {MessageChain([comp])}, error = {e}", + exc_info=True, + ) else: if all( comp.type in {ComponentType.Reply, ComponentType.At} @@ -237,7 +241,7 @@ class RespondStage(Stage): ): # may fix #2670 logger.warning( - f"消息链全为 Reply 和 At 消息段, 跳过发送阶段。chain: {result.chain}" + f"消息链全为 Reply 和 At 消息段, 跳过发送阶段。chain: {result.chain}", ) return sep_comps = self._extract_comp( diff --git a/astrbot/core/pipeline/result_decorate/stage.py b/astrbot/core/pipeline/result_decorate/stage.py index c1f893ba..5dfb52f6 100644 --- a/astrbot/core/pipeline/result_decorate/stage.py +++ b/astrbot/core/pipeline/result_decorate/stage.py @@ -1,7 +1,7 @@ import re import time import traceback -from typing import AsyncGenerator, Union +from collections.abc import AsyncGenerator from astrbot.core import file_token_service, html_renderer, logger from astrbot.core.message.components import At, File, Image, Node, Plain, Record, Reply @@ -30,8 +30,7 @@ class ResultDecorateStage(Stage): self.t2i_word_threshold = ctx.astrbot_config["t2i_word_threshold"] try: self.t2i_word_threshold = int(self.t2i_word_threshold) - if self.t2i_word_threshold < 50: - self.t2i_word_threshold = 50 + self.t2i_word_threshold = max(self.t2i_word_threshold, 50) except BaseException: self.t2i_word_threshold = 150 self.t2i_strategy = ctx.astrbot_config["t2i_strategy"] @@ -46,7 +45,7 @@ class ResultDecorateStage(Stage): self.words_count_threshold = int( ctx.astrbot_config["platform_settings"]["segmented_reply"][ "words_count_threshold" - ] + ], ) self.enable_segmented_reply = ctx.astrbot_config["platform_settings"][ "segmented_reply" @@ -71,8 +70,9 @@ class ResultDecorateStage(Stage): await self.content_safe_check_stage.initialize(ctx) async def process( - self, event: AstrMessageEvent - ) -> Union[None, AsyncGenerator[None, None]]: + self, + event: AstrMessageEvent, + ) -> None | AsyncGenerator[None, None]: result = event.get_result() if result is None or not result.chain: return @@ -94,34 +94,36 @@ class ResultDecorateStage(Stage): if isinstance(comp, Plain): text += comp.text async for _ in self.content_safe_check_stage.process( - event, check_text=text + event, + check_text=text, ): yield # 发送消息前事件钩子 handlers = star_handlers_registry.get_handlers_by_event_type( - EventType.OnDecoratingResultEvent, plugins_name=event.plugins_name + EventType.OnDecoratingResultEvent, + plugins_name=event.plugins_name, ) for handler in handlers: try: logger.debug( - f"hook(on_decorating_result) -> {star_map[handler.handler_module_path].name} - {handler.handler_name}" + f"hook(on_decorating_result) -> {star_map[handler.handler_module_path].name} - {handler.handler_name}", ) if is_stream: logger.warning( - "启用流式输出时，依赖发送消息前事件钩子的插件可能无法正常工作" + "启用流式输出时，依赖发送消息前事件钩子的插件可能无法正常工作", ) await handler.handler(event) if event.get_result() is None or not event.get_result().chain: logger.debug( - f"hook(on_decorating_result) -> {star_map[handler.handler_module_path].name} - {handler.handler_name} 将消息结果清空。" + f"hook(on_decorating_result) -> {star_map[handler.handler_module_path].name} - {handler.handler_name} 将消息结果清空。", ) except BaseException: logger.error(traceback.format_exc()) if event.is_stopped(): logger.info( - f"{star_map[handler.handler_module_path].name} - {handler.handler_name} 终止了事件传播。" + f"{star_map[handler.handler_module_path].name} - {handler.handler_name} 终止了事件传播。", ) return @@ -160,7 +162,9 @@ class ResultDecorateStage(Stage): new_chain.append(comp) continue split_response = re.findall( - self.regex, comp.text, re.DOTALL | re.MULTILINE + self.regex, + comp.text, + re.DOTALL | re.MULTILINE, ) if not split_response: new_chain.append(comp) @@ -177,7 +181,7 @@ class ResultDecorateStage(Stage): # TTS tts_provider = self.ctx.plugin_manager.context.get_using_tts_provider( - event.unified_msg_origin + event.unified_msg_origin, ) if ( @@ -187,7 +191,7 @@ class ResultDecorateStage(Stage): ): if not tts_provider: logger.warning( - f"会话 {event.unified_msg_origin} 未配置文本转语音模型。" + f"会话 {event.unified_msg_origin} 未配置文本转语音模型。", ) else: new_chain = [] @@ -199,7 +203,7 @@ class ResultDecorateStage(Stage): logger.info(f"TTS 结果: {audio_path}") if not audio_path: logger.error( - f"由于 TTS 音频文件未找到，消息段转语音失败: {comp.text}" + f"由于 TTS 音频文件未找到，消息段转语音失败: {comp.text}", ) new_chain.append(comp) continue @@ -217,7 +221,7 @@ class ResultDecorateStage(Stage): url = None if use_file_service and callback_api_base: token = await file_token_service.register_file( - audio_path + audio_path, ) url = f"{callback_api_base}/api/file/{token}" logger.debug(f"已注册：{url}") @@ -226,7 +230,7 @@ class ResultDecorateStage(Stage): Record( file=url or audio_path, url=url or audio_path, - ) + ), ) if dual_output: new_chain.append(comp) @@ -242,12 +246,13 @@ class ResultDecorateStage(Stage): elif ( result.use_t2i_ is None and self.ctx.astrbot_config["t2i"] ) or result.use_t2i_: - plain_str = "" + parts = [] for comp in result.chain: if isinstance(comp, Plain): - plain_str += "\n\n" + comp.text + parts.append("\n\n" + comp.text) else: break + plain_str = "".join(parts) if plain_str and len(plain_str) > self.t2i_word_threshold: render_start = time.time() try: @@ -262,7 +267,7 @@ class ResultDecorateStage(Stage): return if time.time() - render_start > 3: logger.warning( - "文本转图片耗时超过了 3 秒，如果觉得很慢可以使用 /t2i 关闭文本转图片模式。" + "文本转图片耗时超过了 3 秒，如果觉得很慢可以使用 /t2i 关闭文本转图片模式。", ) if url: if url.startswith("http"): @@ -286,7 +291,9 @@ class ResultDecorateStage(Stage): word_cnt += len(comp.text) if word_cnt > self.forward_threshold: node = Node( - uin=event.get_self_id(), name="AstrBot", content=[*result.chain] + uin=event.get_self_id(), + name="AstrBot", + content=[*result.chain], ) result.chain = [node] @@ -298,7 +305,8 @@ class ResultDecorateStage(Stage): and event.get_message_type() != MessageType.FRIEND_MESSAGE ): result.chain.insert( - 0, At(qq=event.get_sender_id(), name=event.get_sender_name()) + 0, + At(qq=event.get_sender_id(), name=event.get_sender_name()), ) if len(result.chain) > 1 and isinstance(result.chain[1], Plain): result.chain[1].text = "\n" + result.chain[1].text diff --git a/astrbot/core/pipeline/scheduler.py b/astrbot/core/pipeline/scheduler.py index 7a38ec03..5c461a1e 100644 --- a/astrbot/core/pipeline/scheduler.py +++ b/astrbot/core/pipeline/scheduler.py @@ -1,9 +1,11 @@ -from . import STAGES_ORDER -from .stage import registered_stages -from .context import PipelineContext -from typing import AsyncGenerator -from astrbot.core.platform import AstrMessageEvent +from collections.abc import AsyncGenerator + from astrbot.core import logger +from astrbot.core.platform import AstrMessageEvent + +from . import STAGES_ORDER +from .context import PipelineContext +from .stage import registered_stages class PipelineScheduler: @@ -11,7 +13,7 @@ class PipelineScheduler: def __init__(self, context: PipelineContext): registered_stages.sort( - key=lambda x: STAGES_ORDER.index(x.__name__) + key=lambda x: STAGES_ORDER.index(x.__name__), ) # 按照顺序排序 self.ctx = context # 上下文对象 self.stages = [] # 存储阶段实例 @@ -29,12 +31,13 @@ class PipelineScheduler: Args: event (AstrMessageEvent): 事件对象 from_stage (int): 从第几个阶段开始执行, 默认从0开始 + """ for i in range(from_stage, len(self.stages)): stage = self.stages[i] # 获取当前要执行的阶段 # logger.debug(f"执行阶段 {stage.__class__.__name__}") coroutine = stage.process( - event + event, ) # 调用阶段的process方法, 返回协程或者异步生成器 if isinstance(coroutine, AsyncGenerator): @@ -43,7 +46,7 @@ class PipelineScheduler: # 此处是前置处理完成后的暂停点(yield), 下面开始执行后续阶段 if event.is_stopped(): logger.debug( - f"阶段 {stage.__class__.__name__} 已终止事件传播。" + f"阶段 {stage.__class__.__name__} 已终止事件传播。", ) break @@ -53,7 +56,7 @@ class PipelineScheduler: # 此处是后续所有阶段处理完毕后返回的点, 执行后置处理 if event.is_stopped(): logger.debug( - f"阶段 {stage.__class__.__name__} 已终止事件传播。" + f"阶段 {stage.__class__.__name__} 已终止事件传播。", ) break else: @@ -70,6 +73,7 @@ class PipelineScheduler: Args: event (AstrMessageEvent): 事件对象 + """ await self._process_stages(event) diff --git a/astrbot/core/pipeline/session_status_check/stage.py b/astrbot/core/pipeline/session_status_check/stage.py index 3c451e26..7feeeb86 100644 --- a/astrbot/core/pipeline/session_status_check/stage.py +++ b/astrbot/core/pipeline/session_status_check/stage.py @@ -1,9 +1,11 @@ -from ..stage import Stage, register_stage -from ..context import PipelineContext -from typing import AsyncGenerator, Union +from collections.abc import AsyncGenerator + +from astrbot.core import logger from astrbot.core.platform.astr_message_event import AstrMessageEvent from astrbot.core.star.session_llm_manager import SessionServiceManager -from astrbot.core import logger + +from ..context import PipelineContext +from ..stage import Stage, register_stage @register_stage @@ -15,19 +17,21 @@ class SessionStatusCheckStage(Stage): self.conv_mgr = ctx.plugin_manager.context.conversation_manager async def process( - self, event: AstrMessageEvent - ) -> Union[None, AsyncGenerator[None, None]]: + self, + event: AstrMessageEvent, + ) -> None | AsyncGenerator[None, None]: # 检查会话是否整体启用 if not SessionServiceManager.is_session_enabled(event.unified_msg_origin): logger.debug(f"会话 {event.unified_msg_origin} 已被关闭，已终止事件传播。") # workaround for #2309 conv_id = await self.conv_mgr.get_curr_conversation_id( - event.unified_msg_origin + event.unified_msg_origin, ) if not conv_id: await self.conv_mgr.new_conversation( - event.unified_msg_origin, platform_id=event.get_platform_id() + event.unified_msg_origin, + platform_id=event.get_platform_id(), ) event.stop_event() diff --git a/astrbot/core/pipeline/stage.py b/astrbot/core/pipeline/stage.py index c4550495..74aca4ef 100644 --- a/astrbot/core/pipeline/stage.py +++ b/astrbot/core/pipeline/stage.py @@ -1,10 +1,13 @@ from __future__ import annotations + import abc -from typing import List, AsyncGenerator, Union, Type +from collections.abc import AsyncGenerator + from astrbot.core.platform.astr_message_event import AstrMessageEvent + from .context import PipelineContext -registered_stages: List[Type[Stage]] = [] # 维护了所有已注册的 Stage 实现类类型 +registered_stages: list[type[Stage]] = [] # 维护了所有已注册的 Stage 实现类类型 def register_stage(cls): @@ -22,18 +25,21 @@ class Stage(abc.ABC): Args: ctx (PipelineContext): 消息管道上下文对象, 包括配置和插件管理器 + """ raise NotImplementedError @abc.abstractmethod async def process( - self, event: AstrMessageEvent - ) -> Union[None, AsyncGenerator[None, None]]: + self, + event: AstrMessageEvent, + ) -> None | AsyncGenerator[None, None]: """处理事件 Args: event (AstrMessageEvent): 事件对象，包含事件的相关信息 Returns: Union[None, AsyncGenerator[None, None]]: 处理结果，可能是 None 或者异步生成器, 如果为 None 则表示不需要继续处理, 如果为异步生成器则表示需要继续处理(进入下一个阶段) + """ raise NotImplementedError diff --git a/astrbot/core/pipeline/waking_check/stage.py b/astrbot/core/pipeline/waking_check/stage.py index de6ad5e3..81491911 100644 --- a/astrbot/core/pipeline/waking_check/stage.py +++ b/astrbot/core/pipeline/waking_check/stage.py @@ -1,11 +1,11 @@ -from typing import AsyncGenerator, Union +from collections.abc import AsyncGenerator from astrbot import logger from astrbot.core.message.components import At, AtAll, Reply from astrbot.core.message.message_event_result import MessageChain, MessageEventResult from astrbot.core.platform.astr_message_event import AstrMessageEvent -from astrbot.core.star.filter.permission import PermissionTypeFilter from astrbot.core.star.filter.command_group import CommandGroupFilter +from astrbot.core.star.filter.permission import PermissionTypeFilter from astrbot.core.star.session_plugin_manager import SessionPluginManager from astrbot.core.star.star import star_map from astrbot.core.star.star_handler import EventType, star_handlers_registry @@ -30,10 +30,12 @@ class WakingCheckStage(Stage): Args: ctx (PipelineContext): 消息管道上下文对象, 包括配置和插件管理器 + """ self.ctx = ctx self.no_permission_reply = self.ctx.astrbot_config["platform_settings"].get( - "no_permission_reply", True + "no_permission_reply", + True, ) # 私聊是否需要 wake_prefix 才能唤醒机器人 self.friend_message_needs_wake_prefix = self.ctx.astrbot_config[ @@ -41,15 +43,18 @@ class WakingCheckStage(Stage): ].get("friend_message_needs_wake_prefix", False) # 是否忽略机器人自己发送的消息 self.ignore_bot_self_message = self.ctx.astrbot_config["platform_settings"].get( - "ignore_bot_self_message", False + "ignore_bot_self_message", + False, ) self.ignore_at_all = self.ctx.astrbot_config["platform_settings"].get( - "ignore_at_all", False + "ignore_at_all", + False, ) async def process( - self, event: AstrMessageEvent - ) -> Union[None, AsyncGenerator[None, None]]: + self, + event: AstrMessageEvent, + ) -> None | AsyncGenerator[None, None]: if ( self.ignore_bot_self_message and event.get_self_id() == event.get_sender_id() @@ -123,7 +128,8 @@ class WakingCheckStage(Stage): logger.debug(f"enabled_plugins_name: {enabled_plugins_name}") for handler in star_handlers_registry.get_handlers_by_event_type( - EventType.AdapterMessageEvent, plugins_name=event.plugins_name + EventType.AdapterMessageEvent, + plugins_name=event.plugins_name, ): # filter 需满足 AND 逻辑关系 passed = True @@ -138,15 +144,14 @@ class WakingCheckStage(Stage): if not filter.filter(event, self.ctx.astrbot_config): permission_not_pass = True permission_filter_raise_error = filter.raise_error - else: - if not filter.filter(event, self.ctx.astrbot_config): - passed = False - break + elif not filter.filter(event, self.ctx.astrbot_config): + passed = False + break except Exception as e: await event.send( MessageEventResult().message( - f"插件 {star_map[handler.handler_module_path].name}: {e}" - ) + f"插件 {star_map[handler.handler_module_path].name}: {e}", + ), ) event.stop_event() passed = False @@ -159,11 +164,11 @@ class WakingCheckStage(Stage): if self.no_permission_reply: await event.send( MessageChain().message( - f"您(ID: {event.get_sender_id()})的权限不足以使用此指令。通过 /sid 获取 ID 并请管理员添加。" - ) + f"您(ID: {event.get_sender_id()})的权限不足以使用此指令。通过 /sid 获取 ID 并请管理员添加。", + ), ) logger.info( - f"触发 {star_map[handler.handler_module_path].name} 时, 用户(ID={event.get_sender_id()}) 权限不足。" + f"触发 {star_map[handler.handler_module_path].name} 时, 用户(ID={event.get_sender_id()}) 权限不足。", ) event.stop_event() return @@ -185,7 +190,8 @@ class WakingCheckStage(Stage): # 根据会话配置过滤插件处理器 activated_handlers = SessionPluginManager.filter_handlers_by_session( - event, activated_handlers + event, + activated_handlers, ) event.set_extra("activated_handlers", activated_handlers) diff --git a/astrbot/core/pipeline/whitelist_check/stage.py b/astrbot/core/pipeline/whitelist_check/stage.py index b140d23b..ea9c5522 100644 --- a/astrbot/core/pipeline/whitelist_check/stage.py +++ b/astrbot/core/pipeline/whitelist_check/stage.py @@ -1,9 +1,11 @@ -from ..stage import Stage, register_stage -from ..context import PipelineContext -from typing import AsyncGenerator, Union +from collections.abc import AsyncGenerator + +from astrbot.core import logger from astrbot.core.platform.astr_message_event import AstrMessageEvent from astrbot.core.platform.message_type import MessageType -from astrbot.core import logger + +from ..context import PipelineContext +from ..stage import Stage, register_stage @register_stage @@ -27,8 +29,9 @@ class WhitelistCheckStage(Stage): self.wl_log = ctx.astrbot_config["platform_settings"]["id_whitelist_log"] async def process( - self, event: AstrMessageEvent - ) -> Union[None, AsyncGenerator[None, None]]: + self, + event: AstrMessageEvent, + ) -> None | AsyncGenerator[None, None]: if not self.enable_whitelist_check: # 白名单检查未启用 return @@ -60,6 +63,6 @@ class WhitelistCheckStage(Stage): ): if self.wl_log: logger.info( - f"会话 ID {event.unified_msg_origin} 不在会话白名单中，已终止事件传播。请在配置文件中添加该会话 ID 到白名单。" + f"会话 ID {event.unified_msg_origin} 不在会话白名单中，已终止事件传播。请在配置文件中添加该会话 ID 到白名单。", ) event.stop_event() diff --git a/astrbot/core/platform/__init__.py b/astrbot/core/platform/__init__.py index 4007b2d9..30b94723 100644 --- a/astrbot/core/platform/__init__.py +++ b/astrbot/core/platform/__init__.py @@ -1,14 +1,14 @@ -from .platform import Platform from .astr_message_event import AstrMessageEvent +from .astrbot_message import AstrBotMessage, Group, MessageMember, MessageType +from .platform import Platform from .platform_metadata import PlatformMetadata -from .astrbot_message import AstrBotMessage, MessageMember, MessageType, Group __all__ = [ - "Platform", - "AstrMessageEvent", - "PlatformMetadata", "AstrBotMessage", + "AstrMessageEvent", + "Group", "MessageMember", "MessageType", - "Group", + "Platform", + "PlatformMetadata", ] diff --git a/astrbot/core/platform/astr_message_event.py b/astrbot/core/platform/astr_message_event.py index e948ed5b..6402aeae 100644 --- a/astrbot/core/platform/astr_message_event.py +++ b/astrbot/core/platform/astr_message_event.py @@ -1,32 +1,31 @@ import abc import asyncio -import re import hashlib +import re import uuid - -from typing import List, Union, Optional, AsyncGenerator, TypeVar, Any +from collections.abc import AsyncGenerator +from typing import Any from astrbot import logger from astrbot.core.db.po import Conversation from astrbot.core.message.components import ( - Plain, - Image, - BaseMessageComponent, - Face, At, AtAll, + BaseMessageComponent, + Face, Forward, + Image, + Plain, Reply, ) -from astrbot.core.message.message_event_result import MessageEventResult, MessageChain +from astrbot.core.message.message_event_result import MessageChain, MessageEventResult from astrbot.core.platform.message_type import MessageType from astrbot.core.provider.entities import ProviderRequest from astrbot.core.utils.metrics import Metric -from .astrbot_message import AstrBotMessage, Group -from .platform_metadata import PlatformMetadata -from .message_session import MessageSession, MessageSesion # noqa -_VT = TypeVar("_VT") +from .astrbot_message import AstrBotMessage, Group +from .message_session import MessageSesion, MessageSession # noqa +from .platform_metadata import PlatformMetadata class AstrMessageEvent(abc.ABC): @@ -76,7 +75,8 @@ class AstrMessageEvent(abc.ABC): def get_platform_name(self): """获取这个事件所属的平台的类型（如 aiocqhttp, slack, discord 等）。 - NOTE: 用户可能会同时运行多个相同类型的平台适配器。""" + NOTE: 用户可能会同时运行多个相同类型的平台适配器。 + """ return self.platform_meta.name def get_platform_id(self): @@ -87,135 +87,103 @@ class AstrMessageEvent(abc.ABC): return self.platform_meta.id def get_message_str(self) -> str: - """ - 获取消息字符串。 - """ + """获取消息字符串。""" return self.message_str - def _outline_chain(self, chain: Optional[List[BaseMessageComponent]]) -> str: - outline = "" + def _outline_chain(self, chain: list[BaseMessageComponent] | None) -> str: if not chain: - return outline + return "" + + parts = [] for i in chain: if isinstance(i, Plain): - outline += i.text + parts.append(i.text) elif isinstance(i, Image): - outline += "[图片]" + parts.append("[图片]") elif isinstance(i, Face): - outline += f"[表情:{i.id}]" + parts.append(f"[表情:{i.id}]") elif isinstance(i, At): - outline += f"[At:{i.qq}]" + parts.append(f"[At:{i.qq}]") elif isinstance(i, AtAll): - outline += "[At:全体成员]" + parts.append("[At:全体成员]") elif isinstance(i, Forward): # 转发消息 - outline += "[转发消息]" + parts.append("[转发消息]") elif isinstance(i, Reply): # 引用回复 if i.message_str: - outline += f"[引用消息({i.sender_nickname}: {i.message_str})]" + parts.append(f"[引用消息({i.sender_nickname}: {i.message_str})]") else: - outline += "[引用消息]" + parts.append("[引用消息]") else: - outline += f"[{i.type}]" - outline += " " - return outline + parts.append(f"[{i.type}]") + parts.append(" ") + return "".join(parts) def get_message_outline(self) -> str: - """ - 获取消息概要。 + """获取消息概要。 除了文本消息外，其他消息类型会被转换为对应的占位符。如图片消息会被转换为 [图片]。 """ return self._outline_chain(self.message_obj.message) - def get_messages(self) -> List[BaseMessageComponent]: - """ - 获取消息链。 - """ + def get_messages(self) -> list[BaseMessageComponent]: + """获取消息链。""" return self.message_obj.message def get_message_type(self) -> MessageType: - """ - 获取消息类型。 - """ + """获取消息类型。""" return self.message_obj.type def get_session_id(self) -> str: - """ - 获取会话id。 - """ + """获取会话id。""" return self.session_id def get_group_id(self) -> str: - """ - 获取群组id。如果不是群组消息，返回空字符串。 - """ + """获取群组id。如果不是群组消息，返回空字符串。""" return self.message_obj.group_id def get_self_id(self) -> str: - """ - 获取机器人自身的id。 - """ + """获取机器人自身的id。""" return self.message_obj.self_id def get_sender_id(self) -> str: - """ - 获取消息发送者的id。 - """ + """获取消息发送者的id。""" return self.message_obj.sender.user_id def get_sender_name(self) -> str: - """ - 获取消息发送者的名称。(可能会返回空字符串) - """ + """获取消息发送者的名称。(可能会返回空字符串)""" return self.message_obj.sender.nickname def set_extra(self, key, value): - """ - 设置额外的信息。 - """ + """设置额外的信息。""" self._extras[key] = value - def get_extra( - self, key: str | None = None, default: _VT = None - ) -> Any: - """ - 获取额外的信息。 - """ + def get_extra(self, key: str | None = None, default=None) -> Any: + """获取额外的信息。""" if key is None: return self._extras return self._extras.get(key, default) def clear_extra(self): - """ - 清除额外的信息。 - """ + """清除额外的信息。""" logger.info(f"清除 {self.get_platform_name()} 的额外信息: {self._extras}") self._extras.clear() def is_private_chat(self) -> bool: - """ - 是否是私聊。 - """ + """是否是私聊。""" return self.message_obj.type.value == (MessageType.FRIEND_MESSAGE).value def is_wake_up(self) -> bool: - """ - 是否是唤醒机器人的事件。 - """ + """是否是唤醒机器人的事件。""" return self.is_wake def is_admin(self) -> bool: - """ - 是否是管理员。 - """ + """是否是管理员。""" return self.role == "admin" async def process_buffer(self, buffer: str, pattern: re.Pattern) -> str: - """ - 将消息缓冲区中的文本按指定正则表达式分割后发送至消息平台，作为不支持流式输出平台的Fallback。 - """ + """将消息缓冲区中的文本按指定正则表达式分割后发送至消息平台，作为不支持流式输出平台的Fallback。""" while True: match = re.search(pattern, buffer) if not match: @@ -227,14 +195,16 @@ class AstrMessageEvent(abc.ABC): return buffer async def send_streaming( - self, generator: AsyncGenerator[MessageChain, None], use_fallback: bool = False + self, + generator: AsyncGenerator[MessageChain, None], + use_fallback: bool = False, ): """发送流式消息到消息平台，使用异步生成器。 目前仅支持: telegram，qq official 私聊。 Fallback仅支持 aiocqhttp。 """ asyncio.create_task( - Metric.upload(msg_event_tick=1, adapter_name=self.platform_meta.name) + Metric.upload(msg_event_tick=1, adapter_name=self.platform_meta.name), ) self._has_send_oper = True @@ -244,7 +214,7 @@ class AstrMessageEvent(abc.ABC): async def _post_send(self): """调度器会在执行 send() 后调用该方法 deprecated in v3.5.18""" - def set_result(self, result: Union[MessageEventResult, str]): + def set_result(self, result: MessageEventResult | str): """设置消息事件的结果。 Note: @@ -264,6 +234,7 @@ class AstrMessageEvent(abc.ABC): event.set_result(MessageEventResult().set_console_log("数量已增加", logging.DEBUG).set_result_type(EventResultType.CONTINUE)) return ``` + """ if isinstance(result, str): result = MessageEventResult().message(result) @@ -287,41 +258,32 @@ class AstrMessageEvent(abc.ABC): self._result.continue_event() def is_stopped(self) -> bool: - """ - 是否终止事件传播。 - """ + """是否终止事件传播。""" if self._result is None: return False # 默认是继续传播 return self._result.is_stopped() def should_call_llm(self, call_llm: bool): - """ - 是否在此消息事件中禁止默认的 LLM 请求。 + """是否在此消息事件中禁止默认的 LLM 请求。 只会阻止 AstrBot 默认的 LLM 请求链路，不会阻止插件中的 LLM 请求。 """ self.call_llm = call_llm def get_result(self) -> MessageEventResult: - """ - 获取消息事件的结果。 - """ + """获取消息事件的结果。""" return self._result def clear_result(self): - """ - 清除消息事件的结果。 - """ + """清除消息事件的结果。""" self._result = None """消息链相关""" def make_result(self) -> MessageEventResult: - """ - 创建一个空的消息事件结果。 + """创建一个空的消息事件结果。 Example: - ```python # 纯文本回复 yield event.make_result().message("Hi") @@ -329,18 +291,16 @@ class AstrMessageEvent(abc.ABC): yield event.make_result().url_image("https://example.com/image.jpg") yield event.make_result().file_image("image.jpg") ``` + """ return MessageEventResult() def plain_result(self, text: str) -> MessageEventResult: - """ - 创建一个空的消息事件结果，只包含一条文本消息。 - """ + """创建一个空的消息事件结果，只包含一条文本消息。""" return MessageEventResult().message(text) def image_result(self, url_or_path: str) -> MessageEventResult: - """ - 创建一个空的消息事件结果，只包含一条图片消息。 + """创建一个空的消息事件结果，只包含一条图片消息。 根据开头是否包含 http 来判断是网络图片还是本地图片。 """ @@ -348,10 +308,8 @@ class AstrMessageEvent(abc.ABC): return MessageEventResult().url_image(url_or_path) return MessageEventResult().file_image(url_or_path) - def chain_result(self, chain: List[BaseMessageComponent]) -> MessageEventResult: - """ - 创建一个空的消息事件结果，包含指定的消息链。 - """ + def chain_result(self, chain: list[BaseMessageComponent]) -> MessageEventResult: + """创建一个空的消息事件结果，包含指定的消息链。""" mer = MessageEventResult() mer.chain = chain return mer @@ -363,13 +321,12 @@ class AstrMessageEvent(abc.ABC): prompt: str, func_tool_manager=None, session_id: str = None, - image_urls: List[str] = [], - contexts: List = [], + image_urls: list[str] | None = None, + contexts: list | None = None, system_prompt: str = "", - conversation: Conversation = None, + conversation: Conversation | None = None, ) -> ProviderRequest: - """ - 创建一个 LLM 请求。 + """创建一个 LLM 请求。 Examples: ```py @@ -388,8 +345,12 @@ class AstrMessageEvent(abc.ABC): func_tool_manager: 函数工具管理器，用于调用函数工具。用 self.context.get_llm_tool_manager() 获取。 conversation: 可选。如果指定，将在指定的对话中进行 LLM 请求。对话的人格会被用于 LLM 请求，并且结果将会被记录到对话中。 - """ + """ + if image_urls is None: + image_urls = [] + if contexts is None: + contexts = [] if len(contexts) > 0 and conversation: conversation = None @@ -410,20 +371,22 @@ class AstrMessageEvent(abc.ABC): Args: message (MessageChain): 消息链，具体使用方式请参考文档。 + """ # Leverage BLAKE2 hash function to generate a non-reversible hash of the sender ID for privacy. hash_obj = hashlib.blake2b(self.get_sender_id().encode("utf-8"), digest_size=16) sid = str(uuid.UUID(bytes=hash_obj.digest())) asyncio.create_task( Metric.upload( - msg_event_tick=1, adapter_name=self.platform_meta.name, sid=sid - ) + msg_event_tick=1, + adapter_name=self.platform_meta.name, + sid=sid, + ), ) self._has_send_oper = True async def react(self, emoji: str): - """ - 对消息添加表情回应。 + """对消息添加表情回应。 默认实现为发送一条包含该表情的消息。 注意：此实现并不一定符合所有平台的原生“表情回应”行为。 @@ -431,11 +394,10 @@ class AstrMessageEvent(abc.ABC): """ await self.send(MessageChain([Plain(emoji)])) - async def get_group(self, group_id: str = None, **kwargs) -> Optional[Group]: + async def get_group(self, group_id: str | None = None, **kwargs) -> Group | None: """获取一个群聊的数据, 如果不填写 group_id: 如果是私聊消息，返回 None。如果是群聊消息，返回当前群聊的数据。 适配情况: - aiocqhttp(OneBotv11) """ - ... diff --git a/astrbot/core/platform/astrbot_message.py b/astrbot/core/platform/astrbot_message.py index 1808c291..0ada1850 100644 --- a/astrbot/core/platform/astrbot_message.py +++ b/astrbot/core/platform/astrbot_message.py @@ -1,14 +1,15 @@ import time -from typing import List from dataclasses import dataclass + from astrbot.core.message.components import BaseMessageComponent + from .message_type import MessageType @dataclass class MessageMember: user_id: str # 发送者id - nickname: str = None + nickname: str | None = None def __str__(self): # 使用 f-string 来构建返回的字符串表示形式 @@ -22,15 +23,15 @@ class MessageMember: class Group: group_id: str """群号""" - group_name: str = None + group_name: str | None = None """群名称""" - group_avatar: str = None + group_avatar: str | None = None """群头像""" - group_owner: str = None + group_owner: str | None = None """群主 id""" - group_admins: List[str] = None + group_admins: list[str] | None = None """群管理员 id""" - members: List[MessageMember] = None + members: list[MessageMember] | None = None """所有群成员""" def __str__(self): @@ -47,9 +48,7 @@ class Group: class AstrBotMessage: - """ - AstrBot 的消息对象 - """ + """AstrBot 的消息对象""" type: MessageType # 消息类型 self_id: str # 机器人的识别id @@ -57,7 +56,7 @@ class AstrBotMessage: message_id: str # 消息id group: Group # 群组 sender: MessageMember # 发送者 - message: List[BaseMessageComponent] # 消息链使用 Nakuru 的消息链格式 + message: list[BaseMessageComponent] # 消息链使用 Nakuru 的消息链格式 message_str: str # 最直观的纯文本消息字符串 raw_message: object timestamp: int # 消息时间戳 @@ -71,8 +70,7 @@ class AstrBotMessage: @property def group_id(self) -> str: - """ - 向后兼容的 group_id 属性 + """向后兼容的 group_id 属性 群组id，如果为私聊，则为空 """ if self.group: diff --git a/astrbot/core/platform/manager.py b/astrbot/core/platform/manager.py index 7090c669..9ff89202 100644 --- a/astrbot/core/platform/manager.py +++ b/astrbot/core/platform/manager.py @@ -1,18 +1,19 @@ -import traceback import asyncio -from astrbot.core.config.astrbot_config import AstrBotConfig -from .platform import Platform -from typing import List +import traceback from asyncio import Queue -from .register import platform_cls_map + from astrbot.core import logger -from astrbot.core.star.star_handler import star_handlers_registry, star_map, EventType +from astrbot.core.config.astrbot_config import AstrBotConfig +from astrbot.core.star.star_handler import EventType, star_handlers_registry, star_map + +from .platform import Platform +from .register import platform_cls_map from .sources.webchat.webchat_adapter import WebChatAdapter class PlatformManager: def __init__(self, config: AstrBotConfig, event_queue: Queue): - self.platform_insts: List[Platform] = [] + self.platform_insts: list[Platform] = [] """加载的 Platform 的实例""" self._inst_map = {} @@ -36,7 +37,7 @@ class PlatformManager: webchat_inst = WebChatAdapter({}, self.settings, self.event_queue) self.platform_insts.append(webchat_inst) asyncio.create_task( - self._task_wrapper(asyncio.create_task(webchat_inst.run(), name="webchat")) + self._task_wrapper(asyncio.create_task(webchat_inst.run(), name="webchat")), ) async def load_platform(self, platform_config: dict): @@ -47,7 +48,7 @@ class PlatformManager: return logger.info( - f"载入 {platform_config['type']}({platform_config['id']}) 平台适配器 ..." + f"载入 {platform_config['type']}({platform_config['id']}) 平台适配器 ...", ) match platform_config["type"]: case "aiocqhttp": @@ -106,14 +107,14 @@ class PlatformManager: ) except (ImportError, ModuleNotFoundError) as e: logger.error( - f"加载平台适配器 {platform_config['type']} 失败，原因：{e}。请检查依赖库是否安装。提示：可以在 管理面板->控制台->安装Pip库 中安装依赖库。" + f"加载平台适配器 {platform_config['type']} 失败，原因：{e}。请检查依赖库是否安装。提示：可以在 管理面板->控制台->安装Pip库 中安装依赖库。", ) except Exception as e: logger.error(f"加载平台适配器 {platform_config['type']} 失败，原因：{e}。") if platform_config["type"] not in platform_cls_map: logger.error( - f"未找到适用于 {platform_config['type']}({platform_config['id']}) 平台适配器，请检查是否已经安装或者名称填写错误" + f"未找到适用于 {platform_config['type']}({platform_config['id']}) 平台适配器，请检查是否已经安装或者名称填写错误", ) return cls_type = platform_cls_map[platform_config["type"]] @@ -129,16 +130,16 @@ class PlatformManager: asyncio.create_task( inst.run(), name=f"platform_{platform_config['type']}_{platform_config['id']}", - ) - ) + ), + ), ) handlers = star_handlers_registry.get_handlers_by_event_type( - EventType.OnPlatformLoadedEvent + EventType.OnPlatformLoadedEvent, ) for handler in handlers: try: logger.info( - f"hook(on_platform_loaded) -> {star_map[handler.handler_module_path].name} - {handler.handler_name}" + f"hook(on_platform_loaded) -> {star_map[handler.handler_module_path].name} - {handler.handler_name}", ) await handler.handler() except Exception: @@ -180,7 +181,7 @@ class PlatformManager: inst for inst in self.platform_insts if inst.client_self_id == client_id - ) + ), ) except Exception: logger.warning(f"可能未完全移除 {platform_id} 平台适配器") diff --git a/astrbot/core/platform/message_session.py b/astrbot/core/platform/message_session.py index bf5a72a9..bca5300b 100644 --- a/astrbot/core/platform/message_session.py +++ b/astrbot/core/platform/message_session.py @@ -1,17 +1,19 @@ -from astrbot.core.platform.message_type import MessageType from dataclasses import dataclass +from astrbot.core.platform.message_type import MessageType + @dataclass class MessageSession: """描述一条消息在 AstrBot 中对应的会话的唯一标识。 - 如果您需要实例化 MessageSession，请不要给 platform_id 赋值（或者同时给 platform_name 和 platform_id 赋值相同值）。它会在 __post_init__ 中自动设置为 platform_name 的值。""" + 如果您需要实例化 MessageSession，请不要给 platform_id 赋值（或者同时给 platform_name 和 platform_id 赋值相同值）。它会在 __post_init__ 中自动设置为 platform_name 的值。 + """ platform_name: str """平台适配器实例的唯一标识符。自 AstrBot v4.0.0 起，该字段实际为 platform_id。""" message_type: MessageType session_id: str - platform_id: str = None + platform_id: str | None = None def __str__(self): return f"{self.platform_id}:{self.message_type.value}:{self.session_id}" diff --git a/astrbot/core/platform/platform.py b/astrbot/core/platform/platform.py index c109f29b..3f36e17f 100644 --- a/astrbot/core/platform/platform.py +++ b/astrbot/core/platform/platform.py @@ -1,13 +1,16 @@ import abc import uuid -from typing import Awaitable, Any from asyncio import Queue -from .platform_metadata import PlatformMetadata -from .astr_message_event import AstrMessageEvent +from collections.abc import Awaitable +from typing import Any + from astrbot.core.message.message_event_result import MessageChain -from .message_session import MessageSesion from astrbot.core.utils.metrics import Metric +from .astr_message_event import AstrMessageEvent +from .message_session import MessageSesion +from .platform_metadata import PlatformMetadata + class Platform(abc.ABC): def __init__(self, event_queue: Queue): @@ -18,42 +21,31 @@ class Platform(abc.ABC): @abc.abstractmethod def run(self) -> Awaitable[Any]: - """ - 得到一个平台的运行实例，需要返回一个协程对象。 - """ + """得到一个平台的运行实例，需要返回一个协程对象。""" raise NotImplementedError async def terminate(self): - """ - 终止一个平台的运行实例。 - """ - ... + """终止一个平台的运行实例。""" @abc.abstractmethod def meta(self) -> PlatformMetadata: - """ - 得到一个平台的元数据。 - """ + """得到一个平台的元数据。""" raise NotImplementedError async def send_by_session( - self, session: MessageSesion, message_chain: MessageChain + self, + session: MessageSesion, + message_chain: MessageChain, ) -> Awaitable[Any]: - """ - 通过会话发送消息。该方法旨在让插件能够直接通过**可持久化的会话数据**发送消息，而不需要保存 event 对象。 + """通过会话发送消息。该方法旨在让插件能够直接通过**可持久化的会话数据**发送消息，而不需要保存 event 对象。 异步方法。 """ await Metric.upload(msg_event_tick=1, adapter_name=self.meta().name) def commit_event(self, event: AstrMessageEvent): - """ - 提交一个事件到事件队列。 - """ + """提交一个事件到事件队列。""" self._event_queue.put_nowait(event) def get_client(self): - """ - 获取平台的客户端对象。 - """ - pass + """获取平台的客户端对象。""" diff --git a/astrbot/core/platform/platform_metadata.py b/astrbot/core/platform/platform_metadata.py index 37f8527a..c63bd82b 100644 --- a/astrbot/core/platform/platform_metadata.py +++ b/astrbot/core/platform/platform_metadata.py @@ -7,12 +7,15 @@ class PlatformMetadata: """平台的名称，即平台的类型，如 aiocqhttp, discord, slack""" description: str """平台的描述""" - id: str = None + id: str | None = None """平台的唯一标识符，用于配置中识别特定平台""" - default_config_tmpl: dict = None + default_config_tmpl: dict | None = None """平台的默认配置模板""" - adapter_display_name: str = None + adapter_display_name: str | None = None """显示在 WebUI 配置页中的平台名称，如空则是 name""" - logo_path: str = None + logo_path: str | None = None """平台适配器的 logo 文件路径（相对于插件目录）""" + + support_streaming_message: bool = True + """平台是否支持真实流式传输""" diff --git a/astrbot/core/platform/register.py b/astrbot/core/platform/register.py index 97c33a43..c1721c5c 100644 --- a/astrbot/core/platform/register.py +++ b/astrbot/core/platform/register.py @@ -1,19 +1,20 @@ -from typing import List, Dict, Type -from .platform_metadata import PlatformMetadata from astrbot.core import logger -platform_registry: List[PlatformMetadata] = [] +from .platform_metadata import PlatformMetadata + +platform_registry: list[PlatformMetadata] = [] """维护了通过装饰器注册的平台适配器""" -platform_cls_map: Dict[str, Type] = {} +platform_cls_map: dict[str, type] = {} """维护了平台适配器名称和适配器类的映射""" def register_platform_adapter( adapter_name: str, desc: str, - default_config_tmpl: dict = None, - adapter_display_name: str = None, - logo_path: str = None, + default_config_tmpl: dict | None = None, + adapter_display_name: str | None = None, + logo_path: str | None = None, + support_streaming_message: bool = True, ): """用于注册平台适配器的带参装饰器。 @@ -24,7 +25,7 @@ def register_platform_adapter( def decorator(cls): if adapter_name in platform_cls_map: raise ValueError( - f"平台适配器 {adapter_name} 已经注册过了，可能发生了适配器命名冲突。" + f"平台适配器 {adapter_name} 已经注册过了，可能发生了适配器命名冲突。", ) # 添加必备选项 @@ -42,6 +43,7 @@ def register_platform_adapter( default_config_tmpl=default_config_tmpl, adapter_display_name=adapter_display_name, logo_path=logo_path, + support_streaming_message=support_streaming_message, ) platform_registry.append(pm) platform_cls_map[adapter_name] = cls diff --git a/astrbot/core/platform/sources/aiocqhttp/aiocqhttp_message_event.py b/astrbot/core/platform/sources/aiocqhttp/aiocqhttp_message_event.py index b8bb723d..ce8fd56d 100644 --- a/astrbot/core/platform/sources/aiocqhttp/aiocqhttp_message_event.py +++ b/astrbot/core/platform/sources/aiocqhttp/aiocqhttp_message_event.py @@ -1,24 +1,31 @@ import asyncio import re -from typing import AsyncGenerator, Dict, List +from collections.abc import AsyncGenerator + from aiocqhttp import CQHttp, Event + from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.message_components import ( + BaseMessageComponent, + File, Image, Node, Nodes, Plain, Record, Video, - File, - BaseMessageComponent, ) from astrbot.api.platform import Group, MessageMember class AiocqhttpMessageEvent(AstrMessageEvent): def __init__( - self, message_str, message_obj, platform_meta, session_id, bot: CQHttp + self, + message_str, + message_obj, + platform_meta, + session_id, + bot: CQHttp, ): super().__init__(message_str, message_obj, platform_meta, session_id) self.bot = bot @@ -35,16 +42,15 @@ class AiocqhttpMessageEvent(AstrMessageEvent): "file": f"base64://{bs64}", }, } - elif isinstance(segment, File): + if isinstance(segment, File): # For File segments, we need to handle the file differently d = await segment.to_dict() return d - elif isinstance(segment, Video): + if isinstance(segment, Video): d = await segment.to_dict() return d - else: - # For other segments, we simply convert them to a dict by calling toDict - return segment.toDict() + # For other segments, we simply convert them to a dict by calling toDict + return segment.toDict() @staticmethod async def _parse_onebot_json(message_chain: MessageChain): @@ -78,7 +84,7 @@ class AiocqhttpMessageEvent(AstrMessageEvent): await bot.send(event=event, message=messages) else: raise ValueError( - f"无法发送消息：缺少有效的数字 session_id({session_id}) 或 event({event})" + f"无法发送消息：缺少有效的数字 session_id({session_id}) 或 event({event})", ) @classmethod @@ -88,7 +94,7 @@ class AiocqhttpMessageEvent(AstrMessageEvent): message_chain: MessageChain, event: Event | None = None, is_group: bool = False, - session_id: str = None, + session_id: str | None = None, ): """发送消息至 QQ 协议端（aiocqhttp）。 @@ -98,8 +104,8 @@ class AiocqhttpMessageEvent(AstrMessageEvent): event (Event | None, optional): aiocqhttp 事件对象. is_group (bool, optional): 是否为群消息. session_id (str | None, optional): 会话 ID（群号或 QQ 号 - """ + """ # 转发消息、文件消息不能和普通消息混在一起发送 send_one_by_one = any( isinstance(seg, (Node, Nodes, File)) for seg in message_chain.chain @@ -152,7 +158,9 @@ class AiocqhttpMessageEvent(AstrMessageEvent): await super().send(message) async def send_streaming( - self, generator: AsyncGenerator, use_fallback: bool = False + self, + generator: AsyncGenerator, + use_fallback: bool = False, ): if not use_fallback: buffer = None @@ -162,7 +170,7 @@ class AiocqhttpMessageEvent(AstrMessageEvent): else: buffer.chain.extend(chain.chain) if not buffer: - return + return None buffer.squash_plain() await self.send(buffer) return await super().send_streaming(generator, use_fallback) @@ -198,7 +206,7 @@ class AiocqhttpMessageEvent(AstrMessageEvent): group_id=group_id, ) - members: List[Dict] = await self.bot.call_action( + members: list[dict] = await self.bot.call_action( "get_group_member_list", group_id=group_id, ) diff --git a/astrbot/core/platform/sources/aiocqhttp/aiocqhttp_platform_adapter.py b/astrbot/core/platform/sources/aiocqhttp/aiocqhttp_platform_adapter.py index d1992b6c..91a16c69 100644 --- a/astrbot/core/platform/sources/aiocqhttp/aiocqhttp_platform_adapter.py +++ b/astrbot/core/platform/sources/aiocqhttp/aiocqhttp_platform_adapter.py @@ -1,33 +1,42 @@ -import time import asyncio -import logging -import uuid import itertools -from typing import Awaitable, Any +import logging +import time +import uuid +from collections.abc import Awaitable +from typing import Any + from aiocqhttp import CQHttp, Event +from aiocqhttp.exceptions import ActionFailed + +from astrbot.api import logger +from astrbot.api.event import MessageChain +from astrbot.api.message_components import * from astrbot.api.platform import ( - Platform, AstrBotMessage, MessageMember, MessageType, + Platform, PlatformMetadata, ) -from astrbot.api.event import MessageChain -from .aiocqhttp_message_event import * # noqa: F403 -from astrbot.api.message_components import * # noqa: F403 -from astrbot.api import logger -from .aiocqhttp_message_event import AiocqhttpMessageEvent from astrbot.core.platform.astr_message_event import MessageSesion + from ...register import register_platform_adapter -from aiocqhttp.exceptions import ActionFailed +from .aiocqhttp_message_event import * +from .aiocqhttp_message_event import AiocqhttpMessageEvent @register_platform_adapter( - "aiocqhttp", "适用于 OneBot V11 标准的消息平台适配器，支持反向 WebSockets。" + "aiocqhttp", + "适用于 OneBot V11 标准的消息平台适配器，支持反向 WebSockets。", + support_streaming_message=False, ) class AiocqhttpAdapter(Platform): def __init__( - self, platform_config: dict, platform_settings: dict, event_queue: asyncio.Queue + self, + platform_config: dict, + platform_settings: dict, + event_queue: asyncio.Queue, ) -> None: super().__init__(event_queue) @@ -41,6 +50,7 @@ class AiocqhttpAdapter(Platform): name="aiocqhttp", description="适用于 OneBot 标准的消息平台适配器，支持反向 WebSockets。", id=self.config.get("id"), + support_streaming_message=False, ) self.bot = CQHttp( @@ -48,7 +58,7 @@ class AiocqhttpAdapter(Platform): import_name="aiocqhttp", api_timeout_sec=180, access_token=platform_config.get( - "ws_reverse_token" + "ws_reverse_token", ), # 以防旧版本配置不存在 ) @@ -81,7 +91,9 @@ class AiocqhttpAdapter(Platform): logger.info("aiocqhttp(OneBot v11) 适配器已连接。") async def send_by_session( - self, session: MessageSesion, message_chain: MessageChain + self, + session: MessageSesion, + message_chain: MessageChain, ): is_group = session.message_type == MessageType.GROUP_MESSAGE if is_group: @@ -97,14 +109,14 @@ class AiocqhttpAdapter(Platform): ) await super().send_by_session(session, message_chain) - async def convert_message(self, event: Event) -> AstrBotMessage: + async def convert_message(self, event: Event) -> AstrBotMessage | None: logger.debug(f"[aiocqhttp] RawMessage {event}") if event["post_type"] == "message": abm = await self._convert_handle_message_event(event) if abm.sender.user_id == "2854196310": # 屏蔽 QQ 管家的消息 - return + return None elif event["post_type"] == "notice": abm = await self._convert_handle_notice_event(event) elif event["post_type"] == "request": @@ -118,7 +130,7 @@ class AiocqhttpAdapter(Platform): abm.self_id = str(event.self_id) abm.sender = MessageMember(user_id=str(event.user_id), nickname=event.user_id) abm.type = MessageType.OTHER_MESSAGE - if "group_id" in event and event["group_id"]: + if event.get("group_id"): abm.type = MessageType.GROUP_MESSAGE abm.group_id = str(event.group_id) else: @@ -144,7 +156,7 @@ class AiocqhttpAdapter(Platform): abm.self_id = str(event.self_id) abm.sender = MessageMember(user_id=str(event.user_id), nickname=event.user_id) abm.type = MessageType.OTHER_MESSAGE - if "group_id" in event and event["group_id"]: + if event.get("group_id"): abm.group_id = str(event.group_id) abm.type = MessageType.GROUP_MESSAGE else: @@ -167,12 +179,14 @@ class AiocqhttpAdapter(Platform): if "sub_type" in event: if event["sub_type"] == "poke" and "target_id" in event: - abm.message.append(Poke(qq=str(event["target_id"]), type="poke")) # noqa: F405 + abm.message.append(Poke(qq=str(event["target_id"]), type="poke")) return abm async def _convert_handle_message_event( - self, event: Event, get_reply=True + self, + event: Event, + get_reply=True, ) -> AstrBotMessage: """OneBot V11 消息类事件 @@ -207,13 +221,13 @@ class AiocqhttpAdapter(Platform): message_str = "" if not isinstance(event.message, list): - err = f"aiocqhttp: 无法识别的消息类型: {str(event.message)}，此条消息将被忽略。如果您在使用 go-cqhttp，请将其配置文件中的 message.post-format 更改为 array。" + err = f"aiocqhttp: 无法识别的消息类型: {event.message!s}，此条消息将被忽略。如果您在使用 go-cqhttp，请将其配置文件中的 message.post-format 更改为 array。" logger.critical(err) try: - self.bot.send(event, err) + await self.bot.send(event, err) except BaseException as e: logger.error(f"回复消息失败: {e}") - return + return None # 按消息段类型类型适配 for t, m_group in itertools.groupby(event.message, key=lambda x: x["type"]): @@ -224,7 +238,7 @@ class AiocqhttpAdapter(Platform): # 如果文本段为空，则跳过 continue message_str += current_text - a = ComponentTypes[t](text=current_text) # noqa: F405 + a = ComponentTypes[t](text=current_text) abm.message.append(a) elif t == "file": @@ -264,7 +278,7 @@ class AiocqhttpAdapter(Platform): elif t == "reply": for m in m_group: if not get_reply: - a = ComponentTypes[t](**m["data"]) # noqa: F405 + a = ComponentTypes[t](**m["data"]) abm.message.append(a) else: try: @@ -277,11 +291,12 @@ class AiocqhttpAdapter(Platform): new_event = Event.from_payload(reply_event_data) if not new_event: logger.error( - f"无法从回复消息数据构造 Event 对象: {reply_event_data}" + f"无法从回复消息数据构造 Event 对象: {reply_event_data}", ) continue abm_reply = await self._convert_handle_message_event( - new_event, get_reply=False + new_event, + get_reply=False, ) reply_seg = Reply( @@ -298,10 +313,12 @@ class AiocqhttpAdapter(Platform): abm.message.append(reply_seg) except BaseException as e: logger.error(f"获取引用消息失败: {e}。") - a = ComponentTypes[t](**m["data"]) # noqa: F405 + a = ComponentTypes[t](**m["data"]) abm.message.append(a) elif t == "at": first_at_self_processed = False + # Accumulate @ mention text for efficient concatenation + at_parts = [] for m in m_group: try: @@ -324,7 +341,8 @@ class AiocqhttpAdapter(Platform): no_cache=False, ) nickname = at_info.get("nick", "") or at_info.get( - "nickname", "" + "nickname", + "", ) is_at_self = str(m["data"]["qq"]) in {abm.self_id, "all"} @@ -332,7 +350,7 @@ class AiocqhttpAdapter(Platform): At( qq=m["data"]["qq"], name=nickname, - ) + ), ) if is_at_self and not first_at_self_processed: @@ -340,16 +358,18 @@ class AiocqhttpAdapter(Platform): first_at_self_processed = True else: # 非第一个@机器人或@其他用户，添加到message_str - message_str += f" @{nickname}({m['data']['qq']}) " + at_parts.append(f" @{nickname}({m['data']['qq']}) ") else: abm.message.append(At(qq=str(m["data"]["qq"]), name="")) except ActionFailed as e: logger.error(f"获取 @ 用户信息失败: {e}，此消息段将被忽略。") except BaseException as e: logger.error(f"获取 @ 用户信息失败: {e}，此消息段将被忽略。") + + message_str += "".join(at_parts) else: for m in m_group: - a = ComponentTypes[t](**m["data"]) # noqa: F405 + a = ComponentTypes[t](**m["data"]) abm.message.append(a) abm.timestamp = int(time.time()) @@ -361,7 +381,7 @@ class AiocqhttpAdapter(Platform): def run(self) -> Awaitable[Any]: if not self.host or not self.port: logger.warning( - "aiocqhttp: 未配置 ws_reverse_host 或 ws_reverse_port，将使用默认值：http://0.0.0.0:6199" + "aiocqhttp: 未配置 ws_reverse_host 或 ws_reverse_port，将使用默认值：http://0.0.0.0:6199", ) self.host = "0.0.0.0" self.port = 6199 diff --git a/astrbot/core/platform/sources/dingtalk/dingtalk_adapter.py b/astrbot/core/platform/sources/dingtalk/dingtalk_adapter.py index e61e2385..18be49a1 100644 --- a/astrbot/core/platform/sources/dingtalk/dingtalk_adapter.py +++ b/astrbot/core/platform/sources/dingtalk/dingtalk_adapter.py @@ -1,26 +1,28 @@ import asyncio import os +import threading import uuid + import aiohttp import dingtalk_stream -import threading +from dingtalk_stream import AckMessage +from astrbot import logger +from astrbot.api.event import MessageChain +from astrbot.api.message_components import At, Image, Plain from astrbot.api.platform import ( - Platform, AstrBotMessage, MessageMember, MessageType, + Platform, PlatformMetadata, ) -from astrbot.api.event import MessageChain -from astrbot.api.message_components import Image, Plain, At from astrbot.core.platform.astr_message_event import MessageSesion -from .dingtalk_event import DingtalkMessageEvent -from ...register import register_platform_adapter -from astrbot import logger -from dingtalk_stream import AckMessage -from astrbot.core.utils.io import download_file from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.io import download_file + +from ...register import register_platform_adapter +from .dingtalk_event import DingtalkMessageEvent class MyEventHandler(dingtalk_stream.EventHandler): @@ -35,10 +37,15 @@ class MyEventHandler(dingtalk_stream.EventHandler): return AckMessage.STATUS_OK, "OK" -@register_platform_adapter("dingtalk", "钉钉机器人官方 API 适配器") +@register_platform_adapter( + "dingtalk", "钉钉机器人官方 API 适配器", support_streaming_message=False +) class DingtalkPlatformAdapter(Platform): def __init__( - self, platform_config: dict, platform_settings: dict, event_queue: asyncio.Queue + self, + platform_config: dict, + platform_settings: dict, + event_queue: asyncio.Queue, ) -> None: super().__init__(event_queue) @@ -64,12 +71,23 @@ class DingtalkPlatformAdapter(Platform): client = dingtalk_stream.DingTalkStreamClient(credential, logger=logger) client.register_all_event_handler(MyEventHandler()) client.register_callback_handler( - dingtalk_stream.ChatbotMessage.TOPIC, self.client + dingtalk_stream.ChatbotMessage.TOPIC, + self.client, ) self.client_ = client # 用于 websockets 的 client + def _id_to_sid(self, dingtalk_id: str | None) -> str | None: + if not dingtalk_id: + return dingtalk_id + prefix = "$:LWCP_v1:$" + if dingtalk_id.startswith(prefix): + return dingtalk_id[len(prefix) :] + return dingtalk_id + async def send_by_session( - self, session: MessageSesion, message_chain: MessageChain + self, + session: MessageSesion, + message_chain: MessageChain, ): raise NotImplementedError("钉钉机器人适配器不支持 send_by_session") @@ -78,10 +96,12 @@ class DingtalkPlatformAdapter(Platform): name="dingtalk", description="钉钉机器人官方 API 适配器", id=self.config.get("id"), + support_streaming_message=False, ) async def convert_msg( - self, message: dingtalk_stream.ChatbotMessage + self, + message: dingtalk_stream.ChatbotMessage, ) -> AstrBotMessage: abm = AstrBotMessage() abm.message = [] @@ -93,15 +113,19 @@ class DingtalkPlatformAdapter(Platform): else MessageType.FRIEND_MESSAGE ) abm.sender = MessageMember( - user_id=message.sender_id, nickname=message.sender_nick + user_id=self._id_to_sid(message.sender_id), + nickname=message.sender_nick, ) - abm.self_id = message.chatbot_user_id + abm.self_id = self._id_to_sid(message.chatbot_user_id) abm.message_id = message.message_id abm.raw_message = message if abm.type == MessageType.GROUP_MESSAGE: - if message.is_in_at_list: - abm.message.append(At(qq=abm.self_id)) + # 处理所有被 @ 的用户（包括机器人自己，因 at_users 已包含） + if message.at_users: + for user in message.at_users: + if id := self._id_to_sid(user.dingtalk_id): + abm.message.append(At(qq=id)) abm.group_id = message.conversation_id if self.unique_session: abm.session_id = abm.sender.user_id @@ -136,7 +160,10 @@ class DingtalkPlatformAdapter(Platform): return abm # 别忘了返回转换后的消息对象 async def download_ding_file( - self, download_code: str, robot_code: str, ext: str + self, + download_code: str, + robot_code: str, + ext: str, ) -> str: """下载钉钉文件 @@ -156,20 +183,22 @@ class DingtalkPlatformAdapter(Platform): } temp_dir = os.path.join(get_astrbot_data_path(), "temp") f_path = os.path.join(temp_dir, f"dingtalk_file_{uuid.uuid4()}.{ext}") - async with aiohttp.ClientSession() as session: - async with session.post( + async with ( + aiohttp.ClientSession() as session, + session.post( "https://api.dingtalk.com/v1.0/robot/messageFiles/download", headers=headers, json=payload, - ) as resp: - if resp.status != 200: - logger.error( - f"下载钉钉文件失败: {resp.status}, {await resp.text()}" - ) - return None - resp_data = await resp.json() - download_url = resp_data["data"]["downloadUrl"] - await download_file(download_url, f_path) + ) as resp, + ): + if resp.status != 200: + logger.error( + f"下载钉钉文件失败: {resp.status}, {await resp.text()}", + ) + return None + resp_data = await resp.json() + download_url = resp_data["data"]["downloadUrl"] + await download_file(download_url, f_path) return f_path async def get_access_token(self) -> str: @@ -184,7 +213,7 @@ class DingtalkPlatformAdapter(Platform): ) as resp: if resp.status != 200: logger.error( - f"获取钉钉机器人 access_token 失败: {resp.status}, {await resp.text()}" + f"获取钉钉机器人 access_token 失败: {resp.status}, {await resp.text()}", ) return None return (await resp.json())["data"]["accessToken"] diff --git a/astrbot/core/platform/sources/dingtalk/dingtalk_event.py b/astrbot/core/platform/sources/dingtalk/dingtalk_event.py index 1e6ddd49..a1cd9c1a 100644 --- a/astrbot/core/platform/sources/dingtalk/dingtalk_event.py +++ b/astrbot/core/platform/sources/dingtalk/dingtalk_event.py @@ -1,8 +1,10 @@ import asyncio + import dingtalk_stream + import astrbot.api.message_components as Comp -from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot import logger +from astrbot.api.event import AstrMessageEvent, MessageChain class DingtalkMessageEvent(AstrMessageEvent): @@ -18,7 +20,9 @@ class DingtalkMessageEvent(AstrMessageEvent): self.client = client async def send_with_client( - self, client: dingtalk_stream.ChatbotHandler, message: MessageChain + self, + client: dingtalk_stream.ChatbotHandler, + message: MessageChain, ): for segment in message.chain: if isinstance(segment, Comp.Plain): @@ -69,7 +73,7 @@ class DingtalkMessageEvent(AstrMessageEvent): else: buffer.chain.extend(chain.chain) if not buffer: - return + return None buffer.squash_plain() await self.send(buffer) return await super().send_streaming(generator, use_fallback) diff --git a/astrbot/core/platform/sources/discord/client.py b/astrbot/core/platform/sources/discord/client.py index 78894491..5d29e342 100644 --- a/astrbot/core/platform/sources/discord/client.py +++ b/astrbot/core/platform/sources/discord/client.py @@ -1,7 +1,9 @@ -import discord -from astrbot import logger import sys +import discord + +from astrbot import logger + if sys.version_info >= (3, 12): from typing import override else: @@ -12,7 +14,7 @@ else: class DiscordBotClient(discord.Bot): """Discord客户端封装""" - def __init__(self, token: str, proxy: str = None): + def __init__(self, token: str, proxy: str | None = None): self.token = token self.proxy = proxy @@ -41,7 +43,8 @@ class DiscordBotClient(discord.Bot): await self.on_ready_once_callback() except Exception as e: logger.error( - f"[Discord] on_ready_once_callback 执行失败: {e}", exc_info=True + f"[Discord] on_ready_once_callback 执行失败: {e}", + exc_info=True, ) def _create_message_data(self, message: discord.Message) -> dict: @@ -84,7 +87,7 @@ class DiscordBotClient(discord.Bot): return logger.debug( - f"[Discord] 收到原始消息 from {message.author.name}: {message.content}" + f"[Discord] 收到原始消息 from {message.author.name}: {message.content}", ) if self.on_message_received: @@ -103,12 +106,12 @@ class DiscordBotClient(discord.Bot): command_name = interaction_data.get("name", "") if options := interaction_data.get("options", []): params = " ".join( - [f"{opt['name']}:{opt.get('value', '')}" for opt in options] + [f"{opt['name']}:{opt.get('value', '')}" for opt in options], ) return f"/{command_name} {params}" return f"/{command_name}" - elif interaction_type == discord.InteractionType.component: + if interaction_type == discord.InteractionType.component: custom_id = interaction_data.get("custom_id", "") component_type = interaction_data.get("component_type", "") return f"component:{custom_id}:{component_type}" diff --git a/astrbot/core/platform/sources/discord/components.py b/astrbot/core/platform/sources/discord/components.py index 07e71216..d3e69e76 100644 --- a/astrbot/core/platform/sources/discord/components.py +++ b/astrbot/core/platform/sources/discord/components.py @@ -1,5 +1,5 @@ import discord -from typing import List + from astrbot.api.message_components import BaseMessageComponent @@ -11,14 +11,14 @@ class DiscordEmbed(BaseMessageComponent): def __init__( self, - title: str = None, - description: str = None, - color: int = None, - url: str = None, - thumbnail: str = None, - image: str = None, - footer: str = None, - fields: List[dict] = None, + title: str | None = None, + description: str | None = None, + color: int | None = None, + url: str | None = None, + thumbnail: str | None = None, + image: str | None = None, + footer: str | None = None, + fields: list[dict] | None = None, ): self.title = title self.description = description @@ -66,10 +66,10 @@ class DiscordButton(BaseMessageComponent): def __init__( self, label: str, - custom_id: str = None, + custom_id: str | None = None, style: str = "primary", - emoji: str = None, - url: str = None, + emoji: str | None = None, + url: str | None = None, disabled: bool = False, ): self.label = label @@ -96,7 +96,9 @@ class DiscordView(BaseMessageComponent): type: str = "discord_view" def __init__( - self, components: List[BaseMessageComponent] = None, timeout: float = None + self, + components: list[BaseMessageComponent] = None, + timeout: float = None, ): self.components = components or [] self.timeout = timeout @@ -108,7 +110,9 @@ class DiscordView(BaseMessageComponent): for component in self.components: if isinstance(component, DiscordButton): button_style = getattr( - discord.ButtonStyle, component.style, discord.ButtonStyle.primary + discord.ButtonStyle, + component.style, + discord.ButtonStyle.primary, ) if component.url: diff --git a/astrbot/core/platform/sources/discord/discord_platform_adapter.py b/astrbot/core/platform/sources/discord/discord_platform_adapter.py index 6764eda6..49b886de 100644 --- a/astrbot/core/platform/sources/discord/discord_platform_adapter.py +++ b/astrbot/core/platform/sources/discord/discord_platform_adapter.py @@ -1,30 +1,32 @@ import asyncio -import discord -import sys import re +import sys +from typing import Any + +import discord from discord.abc import Messageable from discord.channel import DMChannel + +from astrbot import logger +from astrbot.api.event import MessageChain +from astrbot.api.message_components import File, Image, Plain from astrbot.api.platform import ( - Platform, AstrBotMessage, MessageMember, - PlatformMetadata, MessageType, + Platform, + PlatformMetadata, + register_platform_adapter, ) -from astrbot.api.event import MessageChain -from astrbot.api.message_components import Plain, Image, File from astrbot.core.platform.astr_message_event import MessageSesion -from astrbot.api.platform import register_platform_adapter -from astrbot import logger -from .client import DiscordBotClient -from .discord_platform_event import DiscordPlatformEvent - -from typing import Any, Tuple from astrbot.core.star.filter.command import CommandFilter from astrbot.core.star.filter.command_group import CommandGroupFilter from astrbot.core.star.star import star_map from astrbot.core.star.star_handler import StarHandlerMetadata, star_handlers_registry +from .client import DiscordBotClient +from .discord_platform_event import DiscordPlatformEvent + if sys.version_info >= (3, 12): from typing import override else: @@ -32,10 +34,15 @@ else: # 注册平台适配器 -@register_platform_adapter("discord", "Discord 适配器 (基于 Pycord)") +@register_platform_adapter( + "discord", "Discord 适配器 (基于 Pycord)", support_streaming_message=False +) class DiscordPlatformAdapter(Platform): def __init__( - self, platform_config: dict, platform_settings: dict, event_queue: asyncio.Queue + self, + platform_config: dict, + platform_settings: dict, + event_queue: asyncio.Queue, ) -> None: super().__init__(event_queue) self.config = platform_config @@ -51,7 +58,9 @@ class DiscordPlatformAdapter(Platform): @override async def send_by_session( - self, session: MessageSesion, message_chain: MessageChain + self, + session: MessageSesion, + message_chain: MessageChain, ): """通过会话发送消息""" # 创建一个 message_obj 以便在 event 中使用 @@ -71,18 +80,19 @@ class DiscordPlatformAdapter(Platform): message_obj.group_id = self._get_channel_id(channel) else: logger.warning( - f"[Discord] Can't get channel info for {channel_id_str}, will guess message type." + f"[Discord] Can't get channel info for {channel_id_str}, will guess message type.", ) message_obj.type = MessageType.GROUP_MESSAGE message_obj.group_id = session.session_id message_obj.message_str = message_chain.get_plain_text() message_obj.sender = MessageMember( - user_id=str(self.client_self_id), nickname=self.client.user.display_name + user_id=str(self.client_self_id), + nickname=self.client.user.display_name, ) message_obj.self_id = self.client_self_id message_obj.session_id = session.session_id - message_obj.message = message_chain + message_obj.message = message_chain.chain # 创建临时事件对象来发送消息 temp_event = DiscordPlatformEvent( @@ -103,6 +113,7 @@ class DiscordPlatformAdapter(Platform): "Discord 适配器", id=self.config.get("id"), default_config_tmpl=self.config, + support_streaming_message=False, ) @override @@ -149,7 +160,9 @@ class DiscordPlatformAdapter(Platform): logger.error(f"[Discord] 适配器运行时发生意外错误: {e}", exc_info=True) def _get_message_type( - self, channel: Messageable, guild_id: int | None = None + self, + channel: Messageable, + guild_id: int | None = None, ) -> MessageType: """根据 channel 对象和 guild_id 判断消息类型""" if guild_id is not None: @@ -201,7 +214,8 @@ class DiscordPlatformAdapter(Platform): abm.group_id = self._get_channel_id(message.channel) abm.message_str = content abm.sender = MessageMember( - user_id=str(message.author.id), nickname=message.author.display_name + user_id=str(message.author.id), + nickname=message.author.display_name, ) message_chain = [] if abm.message_str: @@ -209,14 +223,14 @@ class DiscordPlatformAdapter(Platform): if message.attachments: for attachment in message.attachments: if attachment.content_type and attachment.content_type.startswith( - "image/" + "image/", ): message_chain.append( - Image(file=attachment.url, filename=attachment.filename) + Image(file=attachment.url, filename=attachment.filename), ) else: message_chain.append( - File(name=attachment.filename, url=attachment.url) + File(name=attachment.filename, url=attachment.url), ) abm.message = message_chain abm.raw_message = message @@ -260,7 +274,7 @@ class DiscordPlatformAdapter(Platform): if hasattr(message.raw_message, "guild") and message.raw_message.guild: try: bot_member = message.raw_message.guild.get_member( - self.client.user.id + self.client.user.id, ) except Exception: bot_member = None @@ -346,7 +360,7 @@ class DiscordPlatformAdapter(Platform): description="指令的所有参数", type=discord.SlashCommandOptionType.string, required=False, - ) + ), ] # 创建SlashCommand @@ -362,7 +376,7 @@ class DiscordPlatformAdapter(Platform): if registered_commands: logger.info( - f"[Discord] 准备同步 {len(registered_commands)} 个指令: {', '.join(registered_commands)}" + f"[Discord] 准备同步 {len(registered_commands)} 个指令: {', '.join(registered_commands)}", ) else: logger.info("[Discord] 没有发现可注册的指令。") @@ -375,7 +389,9 @@ class DiscordPlatformAdapter(Platform): def _create_dynamic_callback(self, cmd_name: str): """为每个指令动态创建一个异步回调函数""" - async def dynamic_callback(ctx: discord.ApplicationContext, params: str = None): + async def dynamic_callback( + ctx: discord.ApplicationContext, params: str | None = None + ): # 将平台特定的前缀'/'剥离，以适配通用的CommandFilter logger.debug(f"[Discord] 回调函数触发: {cmd_name}") logger.debug(f"[Discord] 回调函数参数: {ctx}") @@ -387,7 +403,7 @@ class DiscordPlatformAdapter(Platform): logger.debug( f"[Discord] 斜杠指令 '{cmd_name}' 被触发。 " f"原始参数: '{params}'. " - f"构建的指令字符串: '{message_str_for_filter}'" + f"构建的指令字符串: '{message_str_for_filter}'", ) # 尝试立即响应，防止超时 @@ -404,7 +420,8 @@ class DiscordPlatformAdapter(Platform): abm.group_id = self._get_channel_id(ctx.channel) abm.message_str = message_str_for_filter abm.sender = MessageMember( - user_id=str(ctx.author.id), nickname=ctx.author.display_name + user_id=str(ctx.author.id), + nickname=ctx.author.display_name, ) abm.message = [Plain(text=message_str_for_filter)] abm.raw_message = ctx.interaction @@ -419,8 +436,9 @@ class DiscordPlatformAdapter(Platform): @staticmethod def _extract_command_info( - event_filter: Any, handler_metadata: StarHandlerMetadata - ) -> Tuple[str, str, CommandFilter] | None: + event_filter: Any, + handler_metadata: StarHandlerMetadata, + ) -> tuple[str, str, CommandFilter] | None: """从事件过滤器中提取指令信息""" cmd_name = None # is_group = False diff --git a/astrbot/core/platform/sources/discord/discord_platform_event.py b/astrbot/core/platform/sources/discord/discord_platform_event.py index 2c8d055f..82eb9f14 100644 --- a/astrbot/core/platform/sources/discord/discord_platform_event.py +++ b/astrbot/core/platform/sources/discord/discord_platform_event.py @@ -1,29 +1,26 @@ import asyncio -import discord import base64 +import binascii +from collections.abc import AsyncGenerator from io import BytesIO from pathlib import Path -from typing import Optional -import sys +import discord + +from astrbot import logger from astrbot.api.event import AstrMessageEvent, MessageChain -from astrbot.api.platform import AstrBotMessage, PlatformMetadata, At from astrbot.api.message_components import ( - Plain, - Image, - File, BaseMessageComponent, + File, + Image, + Plain, Reply, ) -from astrbot import logger +from astrbot.api.platform import AstrBotMessage, At, PlatformMetadata + from .client import DiscordBotClient from .components import DiscordEmbed, DiscordView -if sys.version_info >= (3, 12): - from typing import override -else: - from typing_extensions import override - # 自定义Discord视图组件（兼容旧版本） class DiscordViewComponent(BaseMessageComponent): @@ -41,16 +38,14 @@ class DiscordPlatformEvent(AstrMessageEvent): platform_meta: PlatformMetadata, session_id: str, client: DiscordBotClient, - interaction_followup_webhook: Optional[discord.Webhook] = None, + interaction_followup_webhook: discord.Webhook | None = None, ): super().__init__(message_str, message_obj, platform_meta, session_id) self.client = client self.interaction_followup_webhook = interaction_followup_webhook - @override async def send(self, message: MessageChain): """发送消息到Discord平台""" - # 解析消息链为 Discord 所需的对象 try: ( @@ -90,20 +85,34 @@ class DiscordPlatformEvent(AstrMessageEvent): channel = await self._get_channel() if not channel: return - else: - await channel.send(**kwargs) + await channel.send(**kwargs) except Exception as e: logger.error(f"[Discord] 发送消息时发生未知错误: {e}", exc_info=True) await super().send(message) - async def _get_channel(self) -> Optional[discord.abc.Messageable]: + async def send_streaming( + self, generator: AsyncGenerator[MessageChain, None], use_fallback: bool = False + ): + buffer = None + async for chain in generator: + if not buffer: + buffer = chain + else: + buffer.chain.extend(chain.chain) + if not buffer: + return None + buffer.squash_plain() + await self.send(buffer) + return await super().send_streaming(generator, use_fallback) + + async def _get_channel(self) -> discord.abc.Messageable | None: """获取当前事件对应的频道对象""" try: channel_id = int(self.session_id) return self.client.get_channel( - channel_id + channel_id, ) or await self.client.fetch_channel(channel_id) except (ValueError, discord.errors.NotFound, discord.errors.Forbidden): logger.error(f"[Discord] 无法获取频道 {self.session_id}") @@ -112,20 +121,20 @@ class DiscordPlatformEvent(AstrMessageEvent): async def _parse_to_discord( self, message: MessageChain, - ) -> tuple[str, list[discord.File], Optional[discord.ui.View], list[discord.Embed]]: + ) -> tuple[str, list[discord.File], discord.ui.View | None, list[discord.Embed]]: """将 MessageChain 解析为 Discord 发送所需的内容""" - content = "" + content_parts = [] files = [] view = None embeds = [] reference_message_id = None for i in message.chain: # 遍历消息链 if isinstance(i, Plain): # 如果是文字类型的 - content += i.text + content_parts.append(i.text) elif isinstance(i, Reply): reference_message_id = i.id elif isinstance(i, At): - content += f"<@{i.qq}>" + content_parts.append(f"<@{i.qq}>") elif isinstance(i, Image): logger.debug(f"[Discord] 开始处理 Image 组件: {i}") try: @@ -146,13 +155,14 @@ class DiscordPlatformEvent(AstrMessageEvent): continue # 2. File URI - elif file_content.startswith("file:///"): + if file_content.startswith("file:///"): logger.debug(f"[Discord] 处理 File URI: {file_content}") path = Path(file_content[8:]) if await asyncio.to_thread(path.exists): file_bytes = await asyncio.to_thread(path.read_bytes) discord_file = discord.File( - BytesIO(file_bytes), filename=filename or path.name + BytesIO(file_bytes), + filename=filename or path.name, ) else: logger.warning(f"[Discord] 图片文件不存在: {path}") @@ -166,7 +176,8 @@ class DiscordPlatformEvent(AstrMessageEvent): b64_data += "=" * (4 - missing_padding) img_bytes = base64.b64decode(b64_data) discord_file = discord.File( - BytesIO(img_bytes), filename=filename or "image.png" + BytesIO(img_bytes), + filename=filename or "image.png", ) # 4. 裸 Base64 或本地路径 @@ -179,17 +190,19 @@ class DiscordPlatformEvent(AstrMessageEvent): b64_data += "=" * (4 - missing_padding) img_bytes = base64.b64decode(b64_data) discord_file = discord.File( - BytesIO(img_bytes), filename=filename or "image.png" + BytesIO(img_bytes), + filename=filename or "image.png", ) - except (ValueError, TypeError, base64.binascii.Error): + except (ValueError, TypeError, binascii.Error): logger.debug( - f"[Discord] 裸 Base64 解码失败，作为本地路径处理: {file_content}" + f"[Discord] 裸 Base64 解码失败，作为本地路径处理: {file_content}", ) path = Path(file_content) if await asyncio.to_thread(path.exists): file_bytes = await asyncio.to_thread(path.read_bytes) discord_file = discord.File( - BytesIO(file_bytes), filename=filename or path.name + BytesIO(file_bytes), + filename=filename or path.name, ) else: logger.warning(f"[Discord] 图片文件不存在: {path}") @@ -212,11 +225,11 @@ class DiscordPlatformEvent(AstrMessageEvent): if await asyncio.to_thread(path.exists): file_bytes = await asyncio.to_thread(path.read_bytes) files.append( - discord.File(BytesIO(file_bytes), filename=i.name) + discord.File(BytesIO(file_bytes), filename=i.name), ) else: logger.warning( - f"[Discord] 获取文件失败，路径不存在: {file_path_str}" + f"[Discord] 获取文件失败，路径不存在: {file_path_str}", ) else: logger.warning(f"[Discord] 获取文件失败: {i.name}") @@ -235,6 +248,7 @@ class DiscordPlatformEvent(AstrMessageEvent): else: logger.debug(f"[Discord] 忽略了不支持的消息组件: {i.type}") + content = "".join(content_parts) if len(content) > 2000: logger.warning("[Discord] 消息内容超过2000字符，将被截断。") content = content[:2000] @@ -244,7 +258,8 @@ class DiscordPlatformEvent(AstrMessageEvent): """对原消息添加反应""" try: if hasattr(self.message_obj, "raw_message") and hasattr( - self.message_obj.raw_message, "add_reaction" + self.message_obj.raw_message, + "add_reaction", ): await self.message_obj.raw_message.add_reaction(emoji) except Exception as e: @@ -279,7 +294,8 @@ class DiscordPlatformEvent(AstrMessageEvent): def is_mentioned(self) -> bool: """判断机器人是否被@""" if hasattr(self.message_obj, "raw_message") and hasattr( - self.message_obj.raw_message, "mentions" + self.message_obj.raw_message, + "mentions", ): return any( mention.id == int(self.message_obj.self_id) @@ -290,7 +306,8 @@ class DiscordPlatformEvent(AstrMessageEvent): def get_mention_clean_content(self) -> str: """获取去除@后的清洁内容""" if hasattr(self.message_obj, "raw_message") and hasattr( - self.message_obj.raw_message, "clean_content" + self.message_obj.raw_message, + "clean_content", ): return self.message_obj.raw_message.clean_content return self.message_str diff --git a/astrbot/core/platform/sources/lark/lark_adapter.py b/astrbot/core/platform/sources/lark/lark_adapter.py index 4a7ca096..e6e6d4d2 100644 --- a/astrbot/core/platform/sources/lark/lark_adapter.py +++ b/astrbot/core/platform/sources/lark/lark_adapter.py @@ -1,30 +1,37 @@ -import base64 import asyncio +import base64 import json import re import uuid -import astrbot.api.message_components as Comp -from astrbot.api.platform import ( - Platform, - AstrBotMessage, - MessageMember, - MessageType, - PlatformMetadata, -) -from astrbot.api.event import MessageChain -from astrbot.core.platform.astr_message_event import MessageSesion -from .lark_event import LarkMessageEvent -from ...register import register_platform_adapter -from astrbot import logger import lark_oapi as lark from lark_oapi.api.im.v1 import * +import astrbot.api.message_components as Comp +from astrbot import logger +from astrbot.api.event import MessageChain +from astrbot.api.platform import ( + AstrBotMessage, + MessageMember, + MessageType, + Platform, + PlatformMetadata, +) +from astrbot.core.platform.astr_message_event import MessageSesion -@register_platform_adapter("lark", "飞书机器人官方 API 适配器") +from ...register import register_platform_adapter +from .lark_event import LarkMessageEvent + + +@register_platform_adapter( + "lark", "飞书机器人官方 API 适配器", support_streaming_message=False +) class LarkPlatformAdapter(Platform): def __init__( - self, platform_config: dict, platform_settings: dict, event_queue: asyncio.Queue + self, + platform_config: dict, + platform_settings: dict, + event_queue: asyncio.Queue, ) -> None: super().__init__(event_queue) @@ -65,14 +72,16 @@ class LarkPlatformAdapter(Platform): ) async def send_by_session( - self, session: MessageSesion, message_chain: MessageChain + self, + session: MessageSesion, + message_chain: MessageChain, ): res = await LarkMessageEvent._convert_to_lark(message_chain, self.lark_api) wrapped = { "zh_cn": { "title": "", "content": res, - } + }, } if session.message_type == MessageType.GROUP_MESSAGE: @@ -91,7 +100,7 @@ class LarkPlatformAdapter(Platform): .content(json.dumps(wrapped)) .msg_type("post") .uuid(str(uuid.uuid4())) - .build() + .build(), ) .build() ) @@ -108,6 +117,7 @@ class LarkPlatformAdapter(Platform): name="lark", description="飞书机器人官方 API 适配器", id=self.config.get("id"), + support_streaming_message=False, ) async def convert_msg(self, event: lark.im.v1.P2ImMessageReceiveV1): @@ -160,7 +170,7 @@ class LarkPlatformAdapter(Platform): content_json_b = _ls elif message.message_type == "image": content_json_b = [ - {"tag": "img", "image_key": content_json_b["image_key"], "style": []} + {"tag": "img", "image_key": content_json_b["image_key"], "style": []}, ] if message.message_type in ("post", "image"): @@ -200,11 +210,10 @@ class LarkPlatformAdapter(Platform): abm.session_id = abm.group_id else: abm.session_id = abm.sender.user_id + elif abm.type == MessageType.GROUP_MESSAGE: + abm.session_id = f"{abm.sender.user_id}%{abm.group_id}" # 也保留群组id else: - if abm.type == MessageType.GROUP_MESSAGE: - abm.session_id = f"{abm.sender.user_id}%{abm.group_id}" # 也保留群组id - else: - abm.session_id = abm.sender.user_id + abm.session_id = abm.sender.user_id logger.debug(abm) await self.handle_msg(abm) diff --git a/astrbot/core/platform/sources/lark/lark_event.py b/astrbot/core/platform/sources/lark/lark_event.py index 2174c497..04204d35 100644 --- a/astrbot/core/platform/sources/lark/lark_event.py +++ b/astrbot/core/platform/sources/lark/lark_event.py @@ -1,27 +1,34 @@ +import base64 import json import os import uuid -import base64 -import lark_oapi as lark from io import BytesIO -from typing import List -from astrbot.api.event import AstrMessageEvent, MessageChain -from astrbot.api.message_components import Plain, Image as AstrBotImage, At -from astrbot.core.utils.io import download_image_by_url + +import lark_oapi as lark from lark_oapi.api.im.v1 import * + from astrbot import logger +from astrbot.api.event import AstrMessageEvent, MessageChain +from astrbot.api.message_components import At, Plain +from astrbot.api.message_components import Image as AstrBotImage from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.io import download_image_by_url class LarkMessageEvent(AstrMessageEvent): def __init__( - self, message_str, message_obj, platform_meta, session_id, bot: lark.Client + self, + message_str, + message_obj, + platform_meta, + session_id, + bot: lark.Client, ): super().__init__(message_str, message_obj, platform_meta, session_id) self.bot = bot @staticmethod - async def _convert_to_lark(message: MessageChain, lark_client: lark.Client) -> List: + async def _convert_to_lark(message: MessageChain, lark_client: lark.Client) -> list: ret = [] _stage = [] for comp in message.chain: @@ -58,7 +65,7 @@ class LarkMessageEvent(AstrMessageEvent): CreateImageRequestBody.builder() .image_type("message") .image(image_file) - .build() + .build(), ) .build() ) @@ -83,7 +90,7 @@ class LarkMessageEvent(AstrMessageEvent): "zh_cn": { "title": "", "content": res, - } + }, } request = ( @@ -95,7 +102,7 @@ class LarkMessageEvent(AstrMessageEvent): .msg_type("post") .uuid(str(uuid.uuid4())) .reply_in_thread(False) - .build() + .build(), ) .build() ) @@ -114,14 +121,14 @@ class LarkMessageEvent(AstrMessageEvent): .request_body( CreateMessageReactionRequestBody.builder() .reaction_type(Emoji.builder().emoji_type(emoji).build()) - .build() + .build(), ) .build() ) response = await self.bot.im.v1.message_reaction.acreate(request) if not response.success(): logger.error(f"发送飞书表情回应失败({response.code}): {response.msg}") - return None + return async def send_streaming(self, generator, use_fallback: bool = False): buffer = None @@ -131,7 +138,7 @@ class LarkMessageEvent(AstrMessageEvent): else: buffer.chain.extend(chain.chain) if not buffer: - return + return None buffer.squash_plain() await self.send(buffer) return await super().send_streaming(generator, use_fallback) diff --git a/astrbot/core/platform/sources/misskey/misskey_adapter.py b/astrbot/core/platform/sources/misskey/misskey_adapter.py index 981d05c8..ddeec93b 100644 --- a/astrbot/core/platform/sources/misskey/misskey_adapter.py +++ b/astrbot/core/platform/sources/misskey/misskey_adapter.py @@ -1,7 +1,10 @@ import asyncio +import os import random -from typing import Dict, Any, Optional, Awaitable, List +from collections.abc import Awaitable +from typing import Any +import astrbot.api.message_components as Comp from astrbot.api import logger from astrbot.api.event import MessageChain from astrbot.api.platform import ( @@ -11,42 +14,46 @@ from astrbot.api.platform import ( register_platform_adapter, ) from astrbot.core.platform.astr_message_event import MessageSession -import astrbot.api.message_components as Comp from .misskey_api import MisskeyAPI -import os try: import magic # type: ignore except Exception: magic = None +from astrbot.core.utils.astrbot_path import get_astrbot_data_path + from .misskey_event import MisskeyPlatformEvent from .misskey_utils import ( - serialize_message_chain, - resolve_message_visibility, - is_valid_user_session_id, - is_valid_room_session_id, add_at_mention_if_needed, - process_files, - extract_sender_info, - create_base_message, - process_at_mention, - format_poll, - cache_user_info, cache_room_info, + cache_user_info, + create_base_message, + extract_sender_info, + format_poll, + is_valid_room_session_id, + is_valid_user_session_id, + process_at_mention, + process_files, + resolve_message_visibility, + serialize_message_chain, ) -from astrbot.core.utils.astrbot_path import get_astrbot_data_path # Constants MAX_FILE_UPLOAD_COUNT = 16 DEFAULT_UPLOAD_CONCURRENCY = 3 -@register_platform_adapter("misskey", "Misskey 平台适配器") +@register_platform_adapter( + "misskey", "Misskey 平台适配器", support_streaming_message=False +) class MisskeyPlatformAdapter(Platform): def __init__( - self, platform_config: dict, platform_settings: dict, event_queue: asyncio.Queue + self, + platform_config: dict, + platform_settings: dict, + event_queue: asyncio.Queue, ) -> None: super().__init__(event_queue) self.config = platform_config or {} @@ -55,7 +62,8 @@ class MisskeyPlatformAdapter(Platform): self.access_token = self.config.get("misskey_token", "") self.max_message_length = self.config.get("max_message_length", 3000) self.default_visibility = self.config.get( - "misskey_default_visibility", "public" + "misskey_default_visibility", + "public", ) self.local_only = self.config.get("misskey_local_only", False) self.enable_chat = self.config.get("misskey_enable_chat", True) @@ -64,7 +72,7 @@ class MisskeyPlatformAdapter(Platform): # download / security related options (exposed to platform_config) self.allow_insecure_downloads = bool( - self.config.get("misskey_allow_insecure_downloads", False) + self.config.get("misskey_allow_insecure_downloads", False), ) # parse download timeout and chunk size safely _dt = self.config.get("misskey_download_timeout") @@ -87,7 +95,7 @@ class MisskeyPlatformAdapter(Platform): self.unique_session = platform_settings["unique_session"] - self.api: Optional[MisskeyAPI] = None + self.api: MisskeyAPI | None = None self._running = False self.client_self_id = "" self._bot_username = "" @@ -114,6 +122,7 @@ class MisskeyPlatformAdapter(Platform): description="Misskey 平台适配器", id=self.config.get("id", "misskey"), default_config_tmpl=default_config, + support_streaming_message=False, ) async def run(self): @@ -136,7 +145,7 @@ class MisskeyPlatformAdapter(Platform): self.client_self_id = str(user_info.get("id", "")) self._bot_username = user_info.get("username", "") logger.info( - f"[Misskey] 已连接用户: {self._bot_username} (ID: {self.client_self_id})" + f"[Misskey] 已连接用户: {self._bot_username} (ID: {self.client_self_id})", ) except Exception as e: logger.error(f"[Misskey] 获取用户信息失败: {e}") @@ -153,12 +162,17 @@ class MisskeyPlatformAdapter(Platform): if self.enable_chat: streaming.add_message_handler("newChatMessage", self._handle_chat_message) streaming.add_message_handler( - "messaging:newChatMessage", self._handle_chat_message + "messaging:newChatMessage", + self._handle_chat_message, ) streaming.add_message_handler("_debug", self._debug_handler) async def _send_text_only_message( - self, session_id: str, text: str, session, message_chain + self, + session_id: str, + text: str, + session, + message_chain, ): """发送纯文本消息（无文件上传）""" if not self.api: @@ -168,7 +182,7 @@ class MisskeyPlatformAdapter(Platform): from .misskey_utils import extract_user_id_from_session_id user_id = extract_user_id_from_session_id(session_id) - payload: Dict[str, Any] = {"toUserId": user_id, "text": text} + payload: dict[str, Any] = {"toUserId": user_id, "text": text} await self.api.send_message(payload) elif session_id and is_valid_room_session_id(session_id): from .misskey_utils import extract_room_id_from_session_id @@ -180,14 +194,17 @@ class MisskeyPlatformAdapter(Platform): return await super().send_by_session(session, message_chain) def _process_poll_data( - self, message: AstrBotMessage, poll: Dict[str, Any], message_parts: List[str] + self, + message: AstrBotMessage, + poll: dict[str, Any], + message_parts: list[str], ): """处理投票数据，将其添加到消息中""" try: if not isinstance(message.raw_message, dict): message.raw_message = {} message.raw_message["poll"] = poll - setattr(message, "poll", poll) + message.poll = poll except Exception: pass @@ -196,25 +213,26 @@ class MisskeyPlatformAdapter(Platform): message.message.append(Comp.Plain(poll_text)) message_parts.append(poll_text) - def _extract_additional_fields(self, session, message_chain) -> Dict[str, Any]: + def _extract_additional_fields(self, session, message_chain) -> dict[str, Any]: """从会话和消息链中提取额外字段""" fields = {"cw": None, "poll": None, "renote_id": None, "channel_id": None} for comp in message_chain.chain: if hasattr(comp, "cw") and getattr(comp, "cw", None): - fields["cw"] = getattr(comp, "cw") + fields["cw"] = comp.cw break if hasattr(session, "extra_data") and isinstance( - getattr(session, "extra_data", None), dict + getattr(session, "extra_data", None), + dict, ): - extra_data = getattr(session, "extra_data") + extra_data = session.extra_data fields.update( { "poll": extra_data.get("poll"), "renote_id": extra_data.get("renote_id"), "channel_id": extra_data.get("channel_id"), - } + }, ) return fields @@ -237,7 +255,7 @@ class MisskeyPlatformAdapter(Platform): if await streaming.connect(): logger.info( - f"[Misskey] WebSocket 已连接 (尝试 #{connection_attempts})" + f"[Misskey] WebSocket 已连接 (尝试 #{connection_attempts})", ) connection_attempts = 0 await streaming.subscribe_channel("main") @@ -250,34 +268,34 @@ class MisskeyPlatformAdapter(Platform): await streaming.listen() else: logger.error( - f"[Misskey] WebSocket 连接失败 (尝试 #{connection_attempts})" + f"[Misskey] WebSocket 连接失败 (尝试 #{connection_attempts})", ) except Exception as e: logger.error( - f"[Misskey] WebSocket 异常 (尝试 #{connection_attempts}): {e}" + f"[Misskey] WebSocket 异常 (尝试 #{connection_attempts}): {e}", ) if self._running: jitter = random.uniform(0, 1.0) sleep_time = backoff_delay + jitter logger.info( - f"[Misskey] {sleep_time:.1f}秒后重连 (下次尝试 #{connection_attempts + 1})" + f"[Misskey] {sleep_time:.1f}秒后重连 (下次尝试 #{connection_attempts + 1})", ) await asyncio.sleep(sleep_time) backoff_delay = min(backoff_delay * backoff_multiplier, max_backoff) - async def _handle_notification(self, data: Dict[str, Any]): + async def _handle_notification(self, data: dict[str, Any]): try: notification_type = data.get("type") logger.debug( - f"[Misskey] 收到通知事件: type={notification_type}, user_id={data.get('userId', 'unknown')}" + f"[Misskey] 收到通知事件: type={notification_type}, user_id={data.get('userId', 'unknown')}", ) if notification_type in ["mention", "reply", "quote"]: note = data.get("note") if note and self._is_bot_mentioned(note): logger.info( - f"[Misskey] 处理贴文提及: {note.get('text', '')[:50]}..." + f"[Misskey] 处理贴文提及: {note.get('text', '')[:50]}...", ) message = await self.convert_message(note) event = MisskeyPlatformEvent( @@ -291,14 +309,14 @@ class MisskeyPlatformAdapter(Platform): except Exception as e: logger.error(f"[Misskey] 处理通知失败: {e}") - async def _handle_chat_message(self, data: Dict[str, Any]): + async def _handle_chat_message(self, data: dict[str, Any]): try: sender_id = str( - data.get("fromUserId", "") or data.get("fromUser", {}).get("id", "") + data.get("fromUserId", "") or data.get("fromUser", {}).get("id", ""), ) room_id = data.get("toRoomId") logger.debug( - f"[Misskey] 收到聊天事件: sender_id={sender_id}, room_id={room_id}, is_self={sender_id == self.client_self_id}" + f"[Misskey] 收到聊天事件: sender_id={sender_id}, room_id={room_id}, is_self={sender_id == self.client_self_id}", ) if sender_id == self.client_self_id: return @@ -306,7 +324,7 @@ class MisskeyPlatformAdapter(Platform): if room_id: raw_text = data.get("text", "") logger.debug( - f"[Misskey] 检查群聊消息: '{raw_text}', 机器人用户名: '{self._bot_username}'" + f"[Misskey] 检查群聊消息: '{raw_text}', 机器人用户名: '{self._bot_username}'", ) message = await self.convert_room_message(data) @@ -326,13 +344,13 @@ class MisskeyPlatformAdapter(Platform): except Exception as e: logger.error(f"[Misskey] 处理聊天消息失败: {e}") - async def _debug_handler(self, data: Dict[str, Any]): + async def _debug_handler(self, data: dict[str, Any]): event_type = data.get("type", "unknown") logger.debug( - f"[Misskey] 收到未处理事件: type={event_type}, channel={data.get('channel', 'unknown')}" + f"[Misskey] 收到未处理事件: type={event_type}, channel={data.get('channel', 'unknown')}", ) - def _is_bot_mentioned(self, note: Dict[str, Any]) -> bool: + def _is_bot_mentioned(self, note: dict[str, Any]) -> bool: text = note.get("text", "") if not text: return False @@ -352,7 +370,9 @@ class MisskeyPlatformAdapter(Platform): return False async def send_by_session( - self, session: MessageSession, message_chain: MessageChain + self, + session: MessageSession, + message_chain: MessageChain, ) -> Awaitable[Any]: if not self.api: logger.error("[Misskey] API 客户端未初始化") @@ -394,30 +414,33 @@ class MisskeyPlatformAdapter(Platform): if not has_file_components: logger.warning("[Misskey] 消息内容为空且无文件组件，跳过发送") return await super().send_by_session(session, message_chain) - else: - text = "" + text = "" if len(text) > self.max_message_length: text = text[: self.max_message_length] + "..." - file_ids: List[str] = [] - fallback_urls: List[str] = [] + file_ids: list[str] = [] + fallback_urls: list[str] = [] if not self.enable_file_upload: return await self._send_text_only_message( - session_id, text, session, message_chain + session_id, + text, + session, + message_chain, ) MAX_UPLOAD_CONCURRENCY = 10 upload_concurrency = int( self.config.get( - "misskey_upload_concurrency", DEFAULT_UPLOAD_CONCURRENCY - ) + "misskey_upload_concurrency", + DEFAULT_UPLOAD_CONCURRENCY, + ), ) upload_concurrency = min(upload_concurrency, MAX_UPLOAD_CONCURRENCY) sem = asyncio.Semaphore(upload_concurrency) - async def _upload_comp(comp) -> Optional[object]: + async def _upload_comp(comp) -> object | None: """组件上传函数：处理 URL（下载后上传）或本地文件（直接上传）""" from .misskey_utils import ( resolve_component_url_or_path, @@ -432,14 +455,16 @@ class MisskeyPlatformAdapter(Platform): # 解析组件的 URL 或本地路径 url_candidate, local_path = await resolve_component_url_or_path( - comp + comp, ) if not url_candidate and not local_path: return None preferred_name = getattr(comp, "name", None) or getattr( - comp, "file", None + comp, + "file", + None, ) # URL 上传：下载后本地上传 @@ -479,7 +504,7 @@ class MisskeyPlatformAdapter(Platform): if local_path and isinstance(local_path, str): data_temp = os.path.join(get_astrbot_data_path(), "temp") if local_path.startswith(data_temp) and os.path.exists( - local_path + local_path, ): try: os.remove(local_path) @@ -508,7 +533,7 @@ class MisskeyPlatformAdapter(Platform): if len(file_components) > MAX_FILE_UPLOAD_COUNT: logger.warning( - f"[Misskey] 文件数量超过限制 ({len(file_components)} > {MAX_FILE_UPLOAD_COUNT})，只上传前{MAX_FILE_UPLOAD_COUNT}个文件" + f"[Misskey] 文件数量超过限制 ({len(file_components)} > {MAX_FILE_UPLOAD_COUNT})，只上传前{MAX_FILE_UPLOAD_COUNT}个文件", ) file_components = file_components[:MAX_FILE_UPLOAD_COUNT] @@ -540,7 +565,7 @@ class MisskeyPlatformAdapter(Platform): if fallback_urls: appended = "\n" + "\n".join(fallback_urls) text = (text or "") + appended - payload: Dict[str, Any] = {"toRoomId": room_id, "text": text} + payload: dict[str, Any] = {"toRoomId": room_id, "text": text} if file_ids: payload["fileIds"] = file_ids await self.api.send_room_message(payload) @@ -555,13 +580,13 @@ class MisskeyPlatformAdapter(Platform): if fallback_urls: appended = "\n" + "\n".join(fallback_urls) text = (text or "") + appended - payload: Dict[str, Any] = {"toUserId": user_id, "text": text} + payload: dict[str, Any] = {"toUserId": user_id, "text": text} if file_ids: # 聊天消息只支持单个文件，使用 fileId 而不是 fileIds payload["fileId"] = file_ids[0] if len(file_ids) > 1: logger.warning( - f"[Misskey] 聊天消息只支持单个文件，忽略其余 {len(file_ids) - 1} 个文件" + f"[Misskey] 聊天消息只支持单个文件，忽略其余 {len(file_ids) - 1} 个文件", ) await self.api.send_message(payload) else: @@ -581,7 +606,7 @@ class MisskeyPlatformAdapter(Platform): default_visibility=self.default_visibility, ) logger.debug( - f"[Misskey] 解析可见性: visibility={visibility}, visible_user_ids={visible_user_ids}, session_id={session_id}, user_id_for_cache={user_id_for_cache}" + f"[Misskey] 解析可见性: visibility={visibility}, visible_user_ids={visible_user_ids}, session_id={session_id}, user_id_for_cache={user_id_for_cache}", ) fields = self._extract_additional_fields(session, message_chain) @@ -610,7 +635,7 @@ class MisskeyPlatformAdapter(Platform): return await super().send_by_session(session, message_chain) - async def convert_message(self, raw_data: Dict[str, Any]) -> AstrBotMessage: + async def convert_message(self, raw_data: dict[str, Any]) -> AstrBotMessage: """将 Misskey 贴文数据转换为 AstrBotMessage 对象""" sender_info = extract_sender_info(raw_data, is_chat=False) message = create_base_message( @@ -621,7 +646,11 @@ class MisskeyPlatformAdapter(Platform): unique_session=self.unique_session, ) cache_user_info( - self._user_cache, sender_info, raw_data, self.client_self_id, is_chat=False + self._user_cache, + sender_info, + raw_data, + self.client_self_id, + is_chat=False, ) message_parts = [] @@ -629,7 +658,10 @@ class MisskeyPlatformAdapter(Platform): if raw_text: text_parts, processed_text = process_at_mention( - message, raw_text, self._bot_username, self.client_self_id + message, + raw_text, + self._bot_username, + self.client_self_id, ) message_parts.extend(text_parts) @@ -652,7 +684,7 @@ class MisskeyPlatformAdapter(Platform): ) return message - async def convert_chat_message(self, raw_data: Dict[str, Any]) -> AstrBotMessage: + async def convert_chat_message(self, raw_data: dict[str, Any]) -> AstrBotMessage: """将 Misskey 聊天消息数据转换为 AstrBotMessage 对象""" sender_info = extract_sender_info(raw_data, is_chat=True) message = create_base_message( @@ -663,7 +695,11 @@ class MisskeyPlatformAdapter(Platform): unique_session=self.unique_session, ) cache_user_info( - self._user_cache, sender_info, raw_data, self.client_self_id, is_chat=True + self._user_cache, + sender_info, + raw_data, + self.client_self_id, + is_chat=True, ) raw_text = raw_data.get("text", "") @@ -676,7 +712,7 @@ class MisskeyPlatformAdapter(Platform): message.message_str = raw_text if raw_text else "" return message - async def convert_room_message(self, raw_data: Dict[str, Any]) -> AstrBotMessage: + async def convert_room_message(self, raw_data: dict[str, Any]) -> AstrBotMessage: """将 Misskey 群聊消息数据转换为 AstrBotMessage 对象""" sender_info = extract_sender_info(raw_data, is_chat=True) room_id = raw_data.get("toRoomId", "") @@ -690,7 +726,11 @@ class MisskeyPlatformAdapter(Platform): ) cache_user_info( - self._user_cache, sender_info, raw_data, self.client_self_id, is_chat=False + self._user_cache, + sender_info, + raw_data, + self.client_self_id, + is_chat=False, ) cache_room_info(self._user_cache, raw_data, self.client_self_id) @@ -700,7 +740,10 @@ class MisskeyPlatformAdapter(Platform): if raw_text: if self._bot_username and f"@{self._bot_username}" in raw_text: text_parts, processed_text = process_at_mention( - message, raw_text, self._bot_username, self.client_self_id + message, + raw_text, + self._bot_username, + self.client_self_id, ) message_parts.extend(text_parts) else: diff --git a/astrbot/core/platform/sources/misskey/misskey_api.py b/astrbot/core/platform/sources/misskey/misskey_api.py index 4b920508..06dc6304 100644 --- a/astrbot/core/platform/sources/misskey/misskey_api.py +++ b/astrbot/core/platform/sources/misskey/misskey_api.py @@ -1,18 +1,20 @@ +import asyncio import json import random -import asyncio -from typing import Any, Optional, Dict, List, Callable, Awaitable import uuid +from collections.abc import Awaitable, Callable +from typing import Any try: import aiohttp import websockets except ImportError as e: raise ImportError( - "aiohttp and websockets are required for Misskey API. Please install them with: pip install aiohttp websockets" + "aiohttp and websockets are required for Misskey API. Please install them with: pip install aiohttp websockets", ) from e from astrbot.api import logger + from .misskey_utils import FileIDExtractor # Constants @@ -23,54 +25,47 @@ HTTP_OK = 200 class APIError(Exception): """Misskey API 基础异常""" - pass - class APIConnectionError(APIError): """网络连接异常""" - pass - class APIRateLimitError(APIError): """API 频率限制异常""" - pass - class AuthenticationError(APIError): """认证失败异常""" - pass - class WebSocketError(APIError): """WebSocket 连接异常""" - pass - class StreamingClient: def __init__(self, instance_url: str, access_token: str): self.instance_url = instance_url.rstrip("/") self.access_token = access_token - self.websocket: Optional[Any] = None + self.websocket: Any | None = None self.is_connected = False - self.message_handlers: Dict[str, Callable] = {} - self.channels: Dict[str, str] = {} - self.desired_channels: Dict[str, Optional[Dict]] = {} + self.message_handlers: dict[str, Callable] = {} + self.channels: dict[str, str] = {} + self.desired_channels: dict[str, dict | None] = {} self._running = False self._last_pong = None async def connect(self) -> bool: try: ws_url = self.instance_url.replace("https://", "wss://").replace( - "http://", "ws://" + "http://", + "ws://", ) ws_url += f"/streaming?i={self.access_token}" self.websocket = await websockets.connect( - ws_url, ping_interval=30, ping_timeout=10 + ws_url, + ping_interval=30, + ping_timeout=10, ) self.is_connected = True self._running = True @@ -84,7 +79,7 @@ class StreamingClient: await self.subscribe_channel(channel_type, params) except Exception as e: logger.warning( - f"[Misskey WebSocket] 重新订阅 {channel_type} 失败: {e}" + f"[Misskey WebSocket] 重新订阅 {channel_type} 失败: {e}", ) except Exception: pass @@ -104,7 +99,9 @@ class StreamingClient: logger.info("[Misskey WebSocket] 连接已断开") async def subscribe_channel( - self, channel_type: str, params: Optional[Dict] = None + self, + channel_type: str, + params: dict | None = None, ) -> str: if not self.is_connected or not self.websocket: raise WebSocketError("WebSocket 未连接") @@ -136,7 +133,9 @@ class StreamingClient: self.desired_channels.pop(channel_type, None) def add_message_handler( - self, event_type: str, handler: Callable[[Dict], Awaitable[None]] + self, + event_type: str, + handler: Callable[[dict], Awaitable[None]], ): self.message_handlers[event_type] = handler @@ -166,7 +165,7 @@ class StreamingClient: pass except websockets.exceptions.ConnectionClosed as e: logger.warning( - f"[Misskey WebSocket] 连接已关闭 (代码: {e.code}, 原因: {e.reason})" + f"[Misskey WebSocket] 连接已关闭 (代码: {e.code}, 原因: {e.reason})", ) self.is_connected = False try: @@ -188,11 +187,11 @@ class StreamingClient: except Exception: pass - async def _handle_message(self, data: Dict[str, Any]): + async def _handle_message(self, data: dict[str, Any]): message_type = data.get("type") body = data.get("body", {}) - def _build_channel_summary(message_type: Optional[str], body: Any) -> str: + def _build_channel_summary(message_type: str | None, body: Any) -> str: try: if not isinstance(body, dict): return f"[Misskey WebSocket] 收到消息类型: {message_type}" @@ -228,7 +227,7 @@ class StreamingClient: event_body = body.get("body", {}) logger.debug( - f"[Misskey WebSocket] 频道消息: {channel_id}, 事件类型: {event_type}" + f"[Misskey WebSocket] 频道消息: {channel_id}, 事件类型: {event_type}", ) if channel_id in self.channels: @@ -243,7 +242,7 @@ class StreamingClient: await self.message_handlers[event_type](event_body) else: logger.debug( - f"[Misskey WebSocket] 未找到处理器: {handler_key} 或 {event_type}" + f"[Misskey WebSocket] 未找到处理器: {handler_key} 或 {event_type}", ) if "_debug" in self.message_handlers: await self.message_handlers["_debug"]( @@ -251,7 +250,7 @@ class StreamingClient: "type": event_type, "body": event_body, "channel": channel_type, - } + }, ) elif message_type in self.message_handlers: @@ -269,14 +268,14 @@ def retry_async( backoff_base: float = 1.0, max_backoff: float = 30.0, ): - """ - 智能异步重试装饰器 + """智能异步重试装饰器 Args: max_retries: 最大重试次数 retryable_exceptions: 可重试的异常类型 backoff_base: 退避基数 max_backoff: 最大退避时间 + """ def decorator(func): @@ -291,7 +290,7 @@ def retry_async( last_exc = e if attempt == max_retries: logger.error( - f"[Misskey API] {func_name} 重试 {max_retries} 次后仍失败: {e}" + f"[Misskey API] {func_name} 重试 {max_retries} 次后仍失败: {e}", ) break @@ -308,7 +307,7 @@ def retry_async( logger.warning( f"[Misskey API] {func_name} 第 {attempt} 次重试失败: {e}，" - f"{sleep_time:.1f}s后重试" + f"{sleep_time:.1f}s后重试", ) await asyncio.sleep(sleep_time) continue @@ -334,12 +333,12 @@ class MisskeyAPI: allow_insecure_downloads: bool = False, download_timeout: int = 15, chunk_size: int = 64 * 1024, - max_download_bytes: Optional[int] = None, + max_download_bytes: int | None = None, ): self.instance_url = instance_url.rstrip("/") self.access_token = access_token - self._session: Optional[aiohttp.ClientSession] = None - self.streaming: Optional[StreamingClient] = None + self._session: aiohttp.ClientSession | None = None + self.streaming: StreamingClient | None = None # download options self.allow_insecure_downloads = allow_insecure_downloads self.download_timeout = download_timeout @@ -381,39 +380,40 @@ class MisskeyAPI: if status == 400: logger.error(f"[Misskey API] 请求参数错误: {endpoint} (HTTP {status})") raise APIError(f"Bad request for {endpoint}") - elif status == 401: + if status == 401: logger.error(f"[Misskey API] 未授权访问: {endpoint} (HTTP {status})") raise AuthenticationError(f"Unauthorized access for {endpoint}") - elif status == 403: + if status == 403: logger.error(f"[Misskey API] 访问被禁止: {endpoint} (HTTP {status})") raise AuthenticationError(f"Forbidden access for {endpoint}") - elif status == 404: + if status == 404: logger.error(f"[Misskey API] 资源不存在: {endpoint} (HTTP {status})") raise APIError(f"Resource not found for {endpoint}") - elif status == 413: + if status == 413: logger.error(f"[Misskey API] 请求体过大: {endpoint} (HTTP {status})") raise APIError(f"Request entity too large for {endpoint}") - elif status == 429: + if status == 429: logger.warning(f"[Misskey API] 请求频率限制: {endpoint} (HTTP {status})") raise APIRateLimitError(f"Rate limit exceeded for {endpoint}") - elif status == 500: + if status == 500: logger.error(f"[Misskey API] 服务器内部错误: {endpoint} (HTTP {status})") raise APIConnectionError(f"Internal server error for {endpoint}") - elif status == 502: + if status == 502: logger.error(f"[Misskey API] 网关错误: {endpoint} (HTTP {status})") raise APIConnectionError(f"Bad gateway for {endpoint}") - elif status == 503: + if status == 503: logger.error(f"[Misskey API] 服务不可用: {endpoint} (HTTP {status})") raise APIConnectionError(f"Service unavailable for {endpoint}") - elif status == 504: + if status == 504: logger.error(f"[Misskey API] 网关超时: {endpoint} (HTTP {status})") raise APIConnectionError(f"Gateway timeout for {endpoint}") - else: - logger.error(f"[Misskey API] 未知错误: {endpoint} (HTTP {status})") - raise APIConnectionError(f"HTTP {status} for {endpoint}") + logger.error(f"[Misskey API] 未知错误: {endpoint} (HTTP {status})") + raise APIConnectionError(f"HTTP {status} for {endpoint}") async def _process_response( - self, response: aiohttp.ClientResponse, endpoint: str + self, + response: aiohttp.ClientResponse, + endpoint: str, ) -> Any: """处理 API 响应""" if response.status == HTTP_OK: @@ -429,7 +429,7 @@ class MisskeyAPI: ) if notifications_data: logger.debug( - f"[Misskey API] 获取到 {len(notifications_data)} 条新通知" + f"[Misskey API] 获取到 {len(notifications_data)} 条新通知", ) else: logger.debug(f"[Misskey API] 请求成功: {endpoint}") @@ -441,11 +441,11 @@ class MisskeyAPI: try: error_text = await response.text() logger.error( - f"[Misskey API] 请求失败: {endpoint} - HTTP {response.status}, 响应: {error_text}" + f"[Misskey API] 请求失败: {endpoint} - HTTP {response.status}, 响应: {error_text}", ) except Exception: logger.error( - f"[Misskey API] 请求失败: {endpoint} - HTTP {response.status}" + f"[Misskey API] 请求失败: {endpoint} - HTTP {response.status}", ) self._handle_response_status(response.status, endpoint) @@ -456,7 +456,9 @@ class MisskeyAPI: retryable_exceptions=(APIConnectionError, APIRateLimitError), ) async def _make_request( - self, endpoint: str, data: Optional[Dict[str, Any]] = None + self, + endpoint: str, + data: dict[str, Any] | None = None, ) -> Any: url = f"{self.instance_url}/api/{endpoint}" payload = {"i": self.access_token} @@ -472,24 +474,24 @@ class MisskeyAPI: async def create_note( self, - text: Optional[str] = None, + text: str | None = None, visibility: str = "public", - reply_id: Optional[str] = None, - visible_user_ids: Optional[List[str]] = None, - file_ids: Optional[List[str]] = None, + reply_id: str | None = None, + visible_user_ids: list[str] | None = None, + file_ids: list[str] | None = None, local_only: bool = False, - cw: Optional[str] = None, - poll: Optional[Dict[str, Any]] = None, - renote_id: Optional[str] = None, - channel_id: Optional[str] = None, - reaction_acceptance: Optional[str] = None, - no_extract_mentions: Optional[bool] = None, - no_extract_hashtags: Optional[bool] = None, - no_extract_emojis: Optional[bool] = None, - media_ids: Optional[List[str]] = None, - ) -> Dict[str, Any]: + cw: str | None = None, + poll: dict[str, Any] | None = None, + renote_id: str | None = None, + channel_id: str | None = None, + reaction_acceptance: str | None = None, + no_extract_mentions: bool | None = None, + no_extract_hashtags: bool | None = None, + no_extract_emojis: bool | None = None, + media_ids: list[str] | None = None, + ) -> dict[str, Any]: """Create a note (wrapper for notes/create). All additional fields are optional and passed through to the API.""" - data: Dict[str, Any] = {} + data: dict[str, Any] = {} if text is not None: data["text"] = text @@ -537,9 +539,9 @@ class MisskeyAPI: async def upload_file( self, file_path: str, - name: Optional[str] = None, - folder_id: Optional[str] = None, - ) -> Dict[str, Any]: + name: str | None = None, + folder_id: str | None = None, + ) -> dict[str, Any]: """Upload a file to Misskey drive/files/create and return a dict containing id and raw result.""" if not file_path: raise APIError("No file path provided for upload") @@ -565,7 +567,7 @@ class MisskeyAPI: result = await self._process_response(resp, "drive/files/create") file_id = FileIDExtractor.extract_file_id(result) logger.debug( - f"[Misskey API] 本地文件上传成功: {filename} -> {file_id}" + f"[Misskey API] 本地文件上传成功: {filename} -> {file_id}", ) return {"id": file_id, "raw": result} finally: @@ -574,7 +576,7 @@ class MisskeyAPI: logger.error(f"[Misskey API] 文件上传网络错误: {e}") raise APIConnectionError(f"Upload failed: {e}") from e - async def find_files_by_hash(self, md5_hash: str) -> List[Dict[str, Any]]: + async def find_files_by_hash(self, md5_hash: str) -> list[dict[str, Any]]: """Find files by MD5 hash""" if not md5_hash: raise APIError("No MD5 hash provided for find-by-hash") @@ -585,7 +587,7 @@ class MisskeyAPI: logger.debug(f"[Misskey API] find-by-hash 请求: md5={md5_hash}") result = await self._make_request("drive/files/find-by-hash", data) logger.debug( - f"[Misskey API] find-by-hash 响应: 找到 {len(result) if isinstance(result, list) else 0} 个文件" + f"[Misskey API] find-by-hash 响应: 找到 {len(result) if isinstance(result, list) else 0} 个文件", ) return result if isinstance(result, list) else [] except Exception as e: @@ -593,13 +595,15 @@ class MisskeyAPI: raise async def find_files_by_name( - self, name: str, folder_id: Optional[str] = None - ) -> List[Dict[str, Any]]: + self, + name: str, + folder_id: str | None = None, + ) -> list[dict[str, Any]]: """Find files by name""" if not name: raise APIError("No name provided for find") - data: Dict[str, Any] = {"name": name} + data: dict[str, Any] = {"name": name} if folder_id: data["folderId"] = folder_id @@ -607,7 +611,7 @@ class MisskeyAPI: logger.debug(f"[Misskey API] find 请求: name={name}, folder_id={folder_id}") result = await self._make_request("drive/files/find", data) logger.debug( - f"[Misskey API] find 响应: 找到 {len(result) if isinstance(result, list) else 0} 个文件" + f"[Misskey API] find 响应: 找到 {len(result) if isinstance(result, list) else 0} 个文件", ) return result if isinstance(result, list) else [] except Exception as e: @@ -617,11 +621,11 @@ class MisskeyAPI: async def find_files( self, limit: int = 10, - folder_id: Optional[str] = None, - type: Optional[str] = None, - ) -> List[Dict[str, Any]]: + folder_id: str | None = None, + type: str | None = None, + ) -> list[dict[str, Any]]: """List files with optional filters""" - data: Dict[str, Any] = {"limit": limit} + data: dict[str, Any] = {"limit": limit} if folder_id is not None: data["folderId"] = folder_id if type is not None: @@ -629,11 +633,11 @@ class MisskeyAPI: try: logger.debug( - f"[Misskey API] 列表文件请求: limit={limit}, folder_id={folder_id}, type={type}" + f"[Misskey API] 列表文件请求: limit={limit}, folder_id={folder_id}, type={type}", ) result = await self._make_request("drive/files", data) logger.debug( - f"[Misskey API] 列表文件响应: 找到 {len(result) if isinstance(result, list) else 0} 个文件" + f"[Misskey API] 列表文件响应: 找到 {len(result) if isinstance(result, list) else 0} 个文件", ) return result if isinstance(result, list) else [] except Exception as e: @@ -641,27 +645,34 @@ class MisskeyAPI: raise async def _download_with_existing_session( - self, url: str, ssl_verify: bool = True - ) -> Optional[bytes]: + self, + url: str, + ssl_verify: bool = True, + ) -> bytes | None: """使用现有会话下载文件""" if not (hasattr(self, "session") and self.session): raise APIConnectionError("No existing session available") async with self.session.get( - url, timeout=aiohttp.ClientTimeout(total=15), ssl=ssl_verify + url, + timeout=aiohttp.ClientTimeout(total=15), + ssl=ssl_verify, ) as response: if response.status == 200: return await response.read() return None async def _download_with_temp_session( - self, url: str, ssl_verify: bool = True - ) -> Optional[bytes]: + self, + url: str, + ssl_verify: bool = True, + ) -> bytes | None: """使用临时会话下载文件""" connector = aiohttp.TCPConnector(ssl=ssl_verify) async with aiohttp.ClientSession(connector=connector) as temp_session: async with temp_session.get( - url, timeout=aiohttp.ClientTimeout(total=15) + url, + timeout=aiohttp.ClientTimeout(total=15), ) as response: if response.status == 200: return await response.read() @@ -670,13 +681,12 @@ class MisskeyAPI: async def upload_and_find_file( self, url: str, - name: Optional[str] = None, - folder_id: Optional[str] = None, + name: str | None = None, + folder_id: str | None = None, max_wait_time: float = 30.0, check_interval: float = 2.0, - ) -> Optional[Dict[str, Any]]: - """ - 简化的文件上传：尝试 URL 上传，失败则下载后本地上传 + ) -> dict[str, Any] | None: + """简化的文件上传：尝试 URL 上传，失败则下载后本地上传 Args: url: 文件URL @@ -687,28 +697,31 @@ class MisskeyAPI: Returns: 包含文件ID和元信息的字典，失败时返回None + """ if not url: raise APIError("URL不能为空") # 通过本地上传获取即时文件 ID（下载文件 → 上传 → 返回 ID） try: - import tempfile import os + import tempfile # SSL 验证下载，失败则重试不验证 SSL tmp_bytes = None try: tmp_bytes = await self._download_with_existing_session( - url, ssl_verify=True + url, + ssl_verify=True, ) or await self._download_with_temp_session(url, ssl_verify=True) except Exception as ssl_error: logger.debug( - f"[Misskey API] SSL 验证下载失败: {ssl_error}，重试不验证 SSL" + f"[Misskey API] SSL 验证下载失败: {ssl_error}，重试不验证 SSL", ) try: tmp_bytes = await self._download_with_existing_session( - url, ssl_verify=False + url, + ssl_verify=False, ) or await self._download_with_temp_session(url, ssl_verify=False) except Exception: pass @@ -732,13 +745,15 @@ class MisskeyAPI: return None - async def get_current_user(self) -> Dict[str, Any]: + async def get_current_user(self) -> dict[str, Any]: """获取当前用户信息""" return await self._make_request("i", {}) async def send_message( - self, user_id_or_payload: Any, text: Optional[str] = None - ) -> Dict[str, Any]: + self, + user_id_or_payload: Any, + text: str | None = None, + ) -> dict[str, Any]: """发送聊天消息。 Accepts either (user_id: str, text: str) or a single dict payload prepared by caller. @@ -754,8 +769,10 @@ class MisskeyAPI: return result async def send_room_message( - self, room_id_or_payload: Any, text: Optional[str] = None - ) -> Dict[str, Any]: + self, + room_id_or_payload: Any, + text: str | None = None, + ) -> dict[str, Any]: """发送房间消息。 Accepts either (room_id: str, text: str) or a single dict payload. @@ -771,10 +788,13 @@ class MisskeyAPI: return result async def get_messages( - self, user_id: str, limit: int = 10, since_id: Optional[str] = None - ) -> List[Dict[str, Any]]: + self, + user_id: str, + limit: int = 10, + since_id: str | None = None, + ) -> list[dict[str, Any]]: """获取聊天消息历史""" - data: Dict[str, Any] = {"userId": user_id, "limit": limit} + data: dict[str, Any] = {"userId": user_id, "limit": limit} if since_id: data["sinceId"] = since_id @@ -785,10 +805,12 @@ class MisskeyAPI: return [] async def get_mentions( - self, limit: int = 10, since_id: Optional[str] = None - ) -> List[Dict[str, Any]]: + self, + limit: int = 10, + since_id: str | None = None, + ) -> list[dict[str, Any]]: """获取提及通知""" - data: Dict[str, Any] = {"limit": limit} + data: dict[str, Any] = {"limit": limit} if since_id: data["sinceId"] = since_id data["includeTypes"] = ["mention", "reply", "quote"] @@ -796,23 +818,21 @@ class MisskeyAPI: result = await self._make_request("i/notifications", data) if isinstance(result, list): return result - elif isinstance(result, dict) and "notifications" in result: + if isinstance(result, dict) and "notifications" in result: return result["notifications"] - else: - logger.warning(f"[Misskey API] 提及通知响应格式异常: {type(result)}") - return [] + logger.warning(f"[Misskey API] 提及通知响应格式异常: {type(result)}") + return [] async def send_message_with_media( self, message_type: str, target_id: str, - text: Optional[str] = None, - media_urls: Optional[List[str]] = None, - local_files: Optional[List[str]] = None, + text: str | None = None, + media_urls: list[str] | None = None, + local_files: list[str] | None = None, **kwargs, - ) -> Dict[str, Any]: - """ - 通用消息发送函数：统一处理文本+媒体发送 + ) -> dict[str, Any]: + """通用消息发送函数：统一处理文本+媒体发送 Args: message_type: 消息类型 ('chat', 'room', 'note') @@ -827,6 +847,7 @@ class MisskeyAPI: Raises: APIError: 参数错误或发送失败 + """ if not text and not media_urls and not local_files: raise APIError("消息内容不能为空：需要文本或媒体文件") @@ -843,10 +864,14 @@ class MisskeyAPI: # 根据消息类型发送 return await self._dispatch_message( - message_type, target_id, text, file_ids, **kwargs + message_type, + target_id, + text, + file_ids, + **kwargs, ) - async def _process_media_urls(self, urls: List[str]) -> List[str]: + async def _process_media_urls(self, urls: list[str]) -> list[str]: """处理远程媒体文件URL列表，返回文件ID列表""" file_ids = [] for url in urls: @@ -863,7 +888,7 @@ class MisskeyAPI: continue return file_ids - async def _process_local_files(self, file_paths: List[str]) -> List[str]: + async def _process_local_files(self, file_paths: list[str]) -> list[str]: """处理本地文件路径列表，返回文件ID列表""" file_ids = [] for file_path in file_paths: @@ -883,10 +908,10 @@ class MisskeyAPI: self, message_type: str, target_id: str, - text: Optional[str], - file_ids: List[str], + text: str | None, + file_ids: list[str], **kwargs, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """根据消息类型分发到对应的发送方法""" if message_type == "chat": # 聊天消息使用 fileId (单数) @@ -907,7 +932,7 @@ class MisskeyAPI: return {"multiple": True, "results": results} return await self.send_message(payload) - elif message_type == "room": + if message_type == "room": # 房间消息使用 fileId (单数) payload = {"toRoomId": target_id} if text: @@ -926,7 +951,7 @@ class MisskeyAPI: return {"multiple": True, "results": results} return await self.send_room_message(payload) - elif message_type == "note": + if message_type == "note": # 发帖使用 fileIds (复数) note_kwargs = { "text": text, @@ -936,5 +961,4 @@ class MisskeyAPI: note_kwargs.update(kwargs) return await self.create_note(**note_kwargs) - else: - raise APIError(f"不支持的消息类型: {message_type}") + raise APIError(f"不支持的消息类型: {message_type}") diff --git a/astrbot/core/platform/sources/misskey/misskey_event.py b/astrbot/core/platform/sources/misskey/misskey_event.py index cd737f78..7975f0ec 100644 --- a/astrbot/core/platform/sources/misskey/misskey_event.py +++ b/astrbot/core/platform/sources/misskey/misskey_event.py @@ -1,19 +1,20 @@ import asyncio import re -from typing import AsyncGenerator +from collections.abc import AsyncGenerator + from astrbot.api import logger from astrbot.api.event import AstrMessageEvent, MessageChain -from astrbot.api.platform import PlatformMetadata, AstrBotMessage from astrbot.api.message_components import Plain +from astrbot.api.platform import AstrBotMessage, PlatformMetadata from .misskey_utils import ( - serialize_message_chain, - resolve_visibility_from_raw_message, - is_valid_user_session_id, - is_valid_room_session_id, add_at_mention_if_needed, - extract_user_id_from_session_id, extract_room_id_from_session_id, + extract_user_id_from_session_id, + is_valid_room_session_id, + is_valid_user_session_id, + resolve_visibility_from_raw_message, + serialize_message_chain, ) @@ -43,7 +44,7 @@ class MisskeyPlatformEvent(AstrMessageEvent): """发送消息，使用适配器的完整上传和发送逻辑""" try: logger.debug( - f"[MisskeyEvent] send 方法被调用，消息链包含 {len(message.chain)} 个组件" + f"[MisskeyEvent] send 方法被调用，消息链包含 {len(message.chain)} 个组件", ) # 使用适配器的 send_by_session 方法，它包含文件上传逻辑 @@ -65,7 +66,7 @@ class MisskeyPlatformEvent(AstrMessageEvent): ) logger.debug( - f"[MisskeyEvent] 检查适配器方法: hasattr(self.client, 'send_by_session') = {hasattr(self.client, 'send_by_session')}" + f"[MisskeyEvent] 检查适配器方法: hasattr(self.client, 'send_by_session') = {hasattr(self.client, 'send_by_session')}", ) # 调用适配器的 send_by_session 方法 @@ -88,25 +89,27 @@ class MisskeyPlatformEvent(AstrMessageEvent): user_info = { "username": user_data.get("username", ""), "nickname": user_data.get( - "name", user_data.get("username", "") + "name", + user_data.get("username", ""), ), } content = add_at_mention_if_needed(content, user_info, has_at) # 根据会话类型选择发送方式 if hasattr(self.client, "send_message") and is_valid_user_session_id( - self.session_id + self.session_id, ): user_id = extract_user_id_from_session_id(self.session_id) await self.client.send_message(user_id, content) elif hasattr( - self.client, "send_room_message" + self.client, + "send_room_message", ) and is_valid_room_session_id(self.session_id): room_id = extract_room_id_from_session_id(self.session_id) await self.client.send_room_message(room_id, content) elif original_message_id and hasattr(self.client, "create_note"): visibility, visible_user_ids = resolve_visibility_from_raw_message( - raw_message + raw_message, ) await self.client.create_note( content, @@ -124,7 +127,9 @@ class MisskeyPlatformEvent(AstrMessageEvent): logger.error(f"[MisskeyEvent] 发送失败: {e}") async def send_streaming( - self, generator: AsyncGenerator[MessageChain, None], use_fallback: bool = False + self, + generator: AsyncGenerator[MessageChain, None], + use_fallback: bool = False, ): if not use_fallback: buffer = None @@ -134,7 +139,7 @@ class MisskeyPlatformEvent(AstrMessageEvent): else: buffer.chain.extend(chain.chain) if not buffer: - return + return None buffer.squash_plain() await self.send(buffer) return await super().send_streaming(generator, use_fallback) diff --git a/astrbot/core/platform/sources/misskey/misskey_utils.py b/astrbot/core/platform/sources/misskey/misskey_utils.py index ebc95d8d..290acd64 100644 --- a/astrbot/core/platform/sources/misskey/misskey_utils.py +++ b/astrbot/core/platform/sources/misskey/misskey_utils.py @@ -1,6 +1,7 @@ """Misskey 平台适配器通用工具函数""" -from typing import Dict, Any, List, Tuple, Optional, Union +from typing import Any + import astrbot.api.message_components as Comp from astrbot.api.platform import AstrBotMessage, MessageMember, MessageType @@ -9,7 +10,7 @@ class FileIDExtractor: """从 API 响应中提取文件 ID 的帮助类（无状态）。""" @staticmethod - def extract_file_id(result: Any) -> Optional[str]: + def extract_file_id(result: Any) -> str | None: if not isinstance(result, dict): return None @@ -34,8 +35,10 @@ class MessagePayloadBuilder: @staticmethod def build_chat_payload( - user_id: str, text: Optional[str], file_id: Optional[str] = None - ) -> Dict[str, Any]: + user_id: str, + text: str | None, + file_id: str | None = None, + ) -> dict[str, Any]: payload = {"toUserId": user_id} if text: payload["text"] = text @@ -45,8 +48,10 @@ class MessagePayloadBuilder: @staticmethod def build_room_payload( - room_id: str, text: Optional[str], file_id: Optional[str] = None - ) -> Dict[str, Any]: + room_id: str, + text: str | None, + file_id: str | None = None, + ) -> dict[str, Any]: payload = {"toRoomId": room_id} if text: payload["text"] = text @@ -56,9 +61,11 @@ class MessagePayloadBuilder: @staticmethod def build_note_payload( - text: Optional[str], file_ids: Optional[List[str]] = None, **kwargs - ) -> Dict[str, Any]: - payload: Dict[str, Any] = {} + text: str | None, + file_ids: list[str] | None = None, + **kwargs, + ) -> dict[str, Any]: + payload: dict[str, Any] = {} if text: payload["text"] = text if file_ids: @@ -67,7 +74,7 @@ class MessagePayloadBuilder: return payload -def serialize_message_chain(chain: List[Any]) -> Tuple[str, bool]: +def serialize_message_chain(chain: list[Any]) -> tuple[str, bool]: """将消息链序列化为文本字符串""" text_parts = [] has_at = False @@ -76,27 +83,25 @@ def serialize_message_chain(chain: List[Any]) -> Tuple[str, bool]: nonlocal has_at if isinstance(component, Comp.Plain): return component.text - elif isinstance(component, Comp.File): + if isinstance(component, Comp.File): # 为文件组件返回占位符，但适配器仍会处理原组件 return "[文件]" - elif isinstance(component, Comp.Image): + if isinstance(component, Comp.Image): # 为图片组件返回占位符，但适配器仍会处理原组件 return "[图片]" - elif isinstance(component, Comp.At): + if isinstance(component, Comp.At): has_at = True # 优先使用name字段（用户名），如果没有则使用qq字段 # 这样可以避免在Misskey中生成 @ 这样的无效提及 if hasattr(component, "name") and component.name: return f"@{component.name}" - else: - return f"@{component.qq}" - elif hasattr(component, "text"): + return f"@{component.qq}" + if hasattr(component, "text"): text = getattr(component, "text", "") if "@" in text: has_at = True return text - else: - return str(component) + return str(component) for component in chain: if isinstance(component, Comp.Node) and component.content: @@ -113,12 +118,12 @@ def serialize_message_chain(chain: List[Any]) -> Tuple[str, bool]: def resolve_message_visibility( - user_id: Optional[str] = None, - user_cache: Optional[Dict[str, Any]] = None, - self_id: Optional[str] = None, - raw_message: Optional[Dict[str, Any]] = None, + user_id: str | None = None, + user_cache: dict[str, Any] | None = None, + self_id: str | None = None, + raw_message: dict[str, Any] | None = None, default_visibility: str = "public", -) -> Tuple[str, Optional[List[str]]]: +) -> tuple[str, list[str] | None]: """解析 Misskey 消息的可见性设置 可以从 user_cache 或 raw_message 中解析，支持两种调用方式： @@ -169,13 +174,14 @@ def resolve_message_visibility( # 保留旧函数名作为向后兼容的别名 def resolve_visibility_from_raw_message( - raw_message: Dict[str, Any], self_id: Optional[str] = None -) -> Tuple[str, Optional[List[str]]]: + raw_message: dict[str, Any], + self_id: str | None = None, +) -> tuple[str, list[str] | None]: """从原始消息数据中解析可见性设置（已弃用，使用 resolve_message_visibility 替代）""" return resolve_message_visibility(raw_message=raw_message, self_id=self_id) -def is_valid_user_session_id(session_id: Union[str, Any]) -> bool: +def is_valid_user_session_id(session_id: str | Any) -> bool: """检查 session_id 是否是有效的聊天用户 session_id (仅限chat%前缀)""" if not isinstance(session_id, str) or "%" not in session_id: return False @@ -189,7 +195,7 @@ def is_valid_user_session_id(session_id: Union[str, Any]) -> bool: ) -def is_valid_room_session_id(session_id: Union[str, Any]) -> bool: +def is_valid_room_session_id(session_id: str | Any) -> bool: """检查 session_id 是否是有效的房间 session_id (仅限room%前缀)""" if not isinstance(session_id, str) or "%" not in session_id: return False @@ -203,7 +209,7 @@ def is_valid_room_session_id(session_id: Union[str, Any]) -> bool: ) -def is_valid_chat_session_id(session_id: Union[str, Any]) -> bool: +def is_valid_chat_session_id(session_id: str | Any) -> bool: """检查 session_id 是否是有效的聊天 session_id (仅限chat%前缀)""" if not isinstance(session_id, str) or "%" not in session_id: return False @@ -236,7 +242,9 @@ def extract_room_id_from_session_id(session_id: str) -> str: def add_at_mention_if_needed( - text: str, user_info: Optional[Dict[str, Any]], has_at: bool = False + text: str, + user_info: dict[str, Any] | None, + has_at: bool = False, ) -> str: """如果需要且没有@用户，则添加@用户 @@ -258,7 +266,7 @@ def add_at_mention_if_needed( return text -def create_file_component(file_info: Dict[str, Any]) -> Tuple[Any, str]: +def create_file_component(file_info: dict[str, Any]) -> tuple[Any, str]: """创建文件组件和描述文本""" file_url = file_info.get("url", "") file_name = file_info.get("name", "未知文件") @@ -266,16 +274,17 @@ def create_file_component(file_info: Dict[str, Any]) -> Tuple[Any, str]: if file_type.startswith("image/"): return Comp.Image(url=file_url, file=file_name), f"图片[{file_name}]" - elif file_type.startswith("audio/"): + if file_type.startswith("audio/"): return Comp.Record(url=file_url, file=file_name), f"音频[{file_name}]" - elif file_type.startswith("video/"): + if file_type.startswith("video/"): return Comp.Video(url=file_url, file=file_name), f"视频[{file_name}]" - else: - return Comp.File(name=file_name, url=file_url), f"文件[{file_name}]" + return Comp.File(name=file_name, url=file_url), f"文件[{file_name}]" def process_files( - message: AstrBotMessage, files: list, include_text_parts: bool = True + message: AstrBotMessage, + files: list, + include_text_parts: bool = True, ) -> list: """处理文件列表，添加到消息组件中并返回文本描述""" file_parts = [] @@ -287,7 +296,7 @@ def process_files( return file_parts -def format_poll(poll: Dict[str, Any]) -> str: +def format_poll(poll: dict[str, Any]) -> str: """将 Misskey 的 poll 对象格式化为可读字符串。""" if not poll or not isinstance(poll, dict): return "" @@ -304,8 +313,9 @@ def format_poll(poll: Dict[str, Any]) -> str: def extract_sender_info( - raw_data: Dict[str, Any], is_chat: bool = False -) -> Dict[str, Any]: + raw_data: dict[str, Any], + is_chat: bool = False, +) -> dict[str, Any]: """提取发送者信息""" if is_chat: sender = raw_data.get("fromUser", {}) @@ -323,11 +333,11 @@ def extract_sender_info( def create_base_message( - raw_data: Dict[str, Any], - sender_info: Dict[str, Any], + raw_data: dict[str, Any], + sender_info: dict[str, Any], client_self_id: str, is_chat: bool = False, - room_id: Optional[str] = None, + room_id: str | None = None, unique_session: bool = False, ) -> AstrBotMessage: """创建基础消息对象""" @@ -366,8 +376,11 @@ def create_base_message( def process_at_mention( - message: AstrBotMessage, raw_text: str, bot_username: str, client_self_id: str -) -> Tuple[List[str], str]: + message: AstrBotMessage, + raw_text: str, + bot_username: str, + client_self_id: str, +) -> tuple[list[str], str]: """处理@提及逻辑，返回消息部分列表和处理后的文本""" message_parts = [] @@ -382,16 +395,15 @@ def process_at_mention( message.message.append(Comp.Plain(remaining_text)) message_parts.append(remaining_text) return message_parts, remaining_text - else: - message.message.append(Comp.Plain(raw_text)) - message_parts.append(raw_text) - return message_parts, raw_text + message.message.append(Comp.Plain(raw_text)) + message_parts.append(raw_text) + return message_parts, raw_text def cache_user_info( - user_cache: Dict[str, Any], - sender_info: Dict[str, Any], - raw_data: Dict[str, Any], + user_cache: dict[str, Any], + sender_info: dict[str, Any], + raw_data: dict[str, Any], client_self_id: str, is_chat: bool = False, ): @@ -417,7 +429,9 @@ def cache_user_info( def cache_room_info( - user_cache: Dict[str, Any], raw_data: Dict[str, Any], client_self_id: str + user_cache: dict[str, Any], + raw_data: dict[str, Any], + client_self_id: str, ): """缓存房间信息""" room_data = raw_data.get("toRoom") @@ -437,7 +451,7 @@ def cache_room_info( async def resolve_component_url_or_path( comp: Any, -) -> Tuple[Optional[str], Optional[str]]: +) -> tuple[str | None, str | None]: """尝试从组件解析可上传的远程 URL 或本地路径。 返回 (url_candidate, local_path)。两者可能都为 None。 @@ -468,8 +482,7 @@ async def resolve_component_url_or_path( if value.startswith("http"): url_candidate = value break - else: - local_path = value + local_path = value except Exception: continue @@ -491,9 +504,8 @@ async def resolve_component_url_or_path( if value.startswith("http"): url_candidate = value break - else: - local_path = value - break + local_path = value + break except Exception: continue @@ -503,7 +515,7 @@ async def resolve_component_url_or_path( return url_candidate, local_path -def summarize_component_for_log(comp: Any) -> Dict[str, Any]: +def summarize_component_for_log(comp: Any) -> dict[str, Any]: """生成适合日志的组件属性字典（尽量不抛异常）。""" attrs = {} for a in ("file", "url", "path", "src", "source", "name"): @@ -519,15 +531,15 @@ def summarize_component_for_log(comp: Any) -> Dict[str, Any]: async def upload_local_with_retries( api: Any, local_path: str, - preferred_name: Optional[str], - folder_id: Optional[str], -) -> Optional[str]: + preferred_name: str | None, + folder_id: str | None, +) -> str | None: """尝试本地上传，返回 file id 或 None。如果文件类型不允许则直接失败。""" try: res = await api.upload_file(local_path, preferred_name, folder_id) if isinstance(res, dict): fid = res.get("id") or (res.get("raw") or {}).get("createdFile", {}).get( - "id" + "id", ) if fid: return str(fid) diff --git a/astrbot/core/platform/sources/qqofficial/qqofficial_message_event.py b/astrbot/core/platform/sources/qqofficial/qqofficial_message_event.py index 2096237c..fe149664 100644 --- a/astrbot/core/platform/sources/qqofficial/qqofficial_message_event.py +++ b/astrbot/core/platform/sources/qqofficial/qqofficial_message_event.py @@ -1,25 +1,26 @@ +import asyncio +import base64 +import os +import random +import uuid + +import aiofiles import botpy import botpy.message import botpy.types import botpy.types.message -import asyncio -import base64 -import aiofiles -from astrbot.core.utils.io import file_to_base64, download_image_by_url -from astrbot.core.utils.tencent_record_helper import wav_to_tencent_silk -from astrbot.core.utils.astrbot_path import get_astrbot_data_path -from astrbot.api.event import AstrMessageEvent, MessageChain -from astrbot.api.platform import AstrBotMessage, PlatformMetadata -from astrbot.api.message_components import Plain, Image, Record from botpy import Client from botpy.http import Route -from astrbot.api import logger -from botpy.types.message import Media from botpy.types import message -from typing import Optional -import random -import uuid -import os +from botpy.types.message import Media + +from astrbot.api import logger +from astrbot.api.event import AstrMessageEvent, MessageChain +from astrbot.api.message_components import Image, Plain, Record +from astrbot.api.platform import AstrBotMessage, PlatformMetadata +from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.io import download_image_by_url, file_to_base64 +from astrbot.core.utils.tencent_record_helper import wav_to_tencent_silk class QQOfficialMessageEvent(AstrMessageEvent): @@ -75,9 +76,9 @@ class QQOfficialMessageEvent(AstrMessageEvent): return await super().send_streaming(generator, use_fallback) - async def _post_send(self, stream: dict = None): + async def _post_send(self, stream: dict | None = None): if not self.send_buffer: - return + return None source = self.message_obj.raw_message assert isinstance( @@ -103,7 +104,7 @@ class QQOfficialMessageEvent(AstrMessageEvent): and not image_path and not record_file_path ): - return + return None payload = { "content": plain_text, @@ -119,29 +120,38 @@ class QQOfficialMessageEvent(AstrMessageEvent): case botpy.message.GroupMessage: if image_base64: media = await self.upload_group_and_c2c_image( - image_base64, 1, group_openid=source.group_openid + image_base64, + 1, + group_openid=source.group_openid, ) payload["media"] = media payload["msg_type"] = 7 if record_file_path: # group record msg media = await self.upload_group_and_c2c_record( - record_file_path, 3, group_openid=source.group_openid + record_file_path, + 3, + group_openid=source.group_openid, ) payload["media"] = media payload["msg_type"] = 7 ret = await self.bot.api.post_group_message( - group_openid=source.group_openid, **payload + group_openid=source.group_openid, + **payload, ) case botpy.message.C2CMessage: if image_base64: media = await self.upload_group_and_c2c_image( - image_base64, 1, openid=source.author.user_openid + image_base64, + 1, + openid=source.author.user_openid, ) payload["media"] = media payload["msg_type"] = 7 if record_file_path: # c2c record media = await self.upload_group_and_c2c_record( - record_file_path, 3, openid=source.author.user_openid + record_file_path, + 3, + openid=source.author.user_openid, ) payload["media"] = media payload["msg_type"] = 7 @@ -153,14 +163,16 @@ class QQOfficialMessageEvent(AstrMessageEvent): ) else: ret = await self.post_c2c_message( - openid=source.author.user_openid, **payload + openid=source.author.user_openid, + **payload, ) logger.debug(f"Message sent to C2C: {ret}") case botpy.message.Message: if image_path: payload["file_image"] = image_path ret = await self.bot.api.post_message( - channel_id=source.channel_id, **payload + channel_id=source.channel_id, + **payload, ) case botpy.message.DirectMessage: if image_path: @@ -174,7 +186,10 @@ class QQOfficialMessageEvent(AstrMessageEvent): return ret async def upload_group_and_c2c_image( - self, image_base64: str, file_type: int, **kwargs + self, + image_base64: str, + file_type: int, + **kwargs, ) -> botpy.types.message.Media: payload = { "file_data": image_base64, @@ -185,7 +200,7 @@ class QQOfficialMessageEvent(AstrMessageEvent): payload["openid"] = kwargs["openid"] route = Route("POST", "/v2/users/{openid}/files", openid=kwargs["openid"]) return await self.bot.api._http.request(route, json=payload) - elif "group_openid" in kwargs: + if "group_openid" in kwargs: payload["group_openid"] = kwargs["group_openid"] route = Route( "POST", @@ -195,11 +210,13 @@ class QQOfficialMessageEvent(AstrMessageEvent): return await self.bot.api._http.request(route, json=payload) async def upload_group_and_c2c_record( - self, file_source: str, file_type: int, srv_send_msg: bool = False, **kwargs - ) -> Optional[Media]: - """ - 上传媒体文件 - """ + self, + file_source: str, + file_type: int, + srv_send_msg: bool = False, + **kwargs, + ) -> Media | None: + """上传媒体文件""" # 构建基础payload payload = {"file_type": file_type, "srv_send_msg": srv_send_msg} @@ -248,17 +265,17 @@ class QQOfficialMessageEvent(AstrMessageEvent): self, openid: str, msg_type: int = 0, - content: str = None, - embed: message.Embed = None, - ark: message.Ark = None, - message_reference: message.Reference = None, - media: message.Media = None, - msg_id: str = None, + content: str | None = None, + embed: message.Embed | None = None, + ark: message.Ark | None = None, + message_reference: message.Reference | None = None, + media: message.Media | None = None, + msg_id: str | None = None, msg_seq: str = 1, - event_id: str = None, - markdown: message.MarkdownPayload = None, - keyboard: message.Keyboard = None, - stream: dict = None, + event_id: str | None = None, + markdown: message.MarkdownPayload | None = None, + keyboard: message.Keyboard | None = None, + stream: dict | None = None, ) -> message.Message: payload = locals() payload.pop("self", None) @@ -291,11 +308,13 @@ class QQOfficialMessageEvent(AstrMessageEvent): record_wav_path = await i.convert_to_file_path() # wav 路径 temp_dir = os.path.join(get_astrbot_data_path(), "temp") record_tecent_silk_path = os.path.join( - temp_dir, f"{uuid.uuid4()}.silk" + temp_dir, + f"{uuid.uuid4()}.silk", ) try: duration = await wav_to_tencent_silk( - record_wav_path, record_tecent_silk_path + record_wav_path, + record_tecent_silk_path, ) if duration > 0: record_file_path = record_tecent_silk_path diff --git a/astrbot/core/platform/sources/qqofficial/qqofficial_platform_adapter.py b/astrbot/core/platform/sources/qqofficial/qqofficial_platform_adapter.py index d5285f75..96be734f 100644 --- a/astrbot/core/platform/sources/qqofficial/qqofficial_platform_adapter.py +++ b/astrbot/core/platform/sources/qqofficial/qqofficial_platform_adapter.py @@ -1,30 +1,31 @@ from __future__ import annotations -import botpy -import logging -import time import asyncio +import logging +import os +import time + +import botpy import botpy.message import botpy.types import botpy.types.message -import os - from botpy import Client + +from astrbot import logger +from astrbot.api.event import MessageChain +from astrbot.api.message_components import At, Image, Plain from astrbot.api.platform import ( - Platform, AstrBotMessage, MessageMember, MessageType, + Platform, PlatformMetadata, ) -from astrbot import logger -from astrbot.api.event import MessageChain -from typing import Union, List -from astrbot.api.message_components import Image, Plain, At -from astrbot.core.platform.astr_message_event import MessageSesion -from .qqofficial_message_event import QQOfficialMessageEvent -from ...register import register_platform_adapter from astrbot.core.message.components import BaseMessageComponent +from astrbot.core.platform.astr_message_event import MessageSesion + +from ...register import register_platform_adapter +from .qqofficial_message_event import QQOfficialMessageEvent # remove logger handler for handler in logging.root.handlers[:]: @@ -33,13 +34,14 @@ for handler in logging.root.handlers[:]: # QQ 机器人官方框架 class botClient(Client): - def set_platform(self, platform: "QQOfficialPlatformAdapter"): + def set_platform(self, platform: QQOfficialPlatformAdapter): self.platform = platform # 收到群消息 async def on_group_at_message_create(self, message: botpy.message.GroupMessage): abm = QQOfficialPlatformAdapter._parse_from_qqofficial( - message, MessageType.GROUP_MESSAGE + message, + MessageType.GROUP_MESSAGE, ) abm.session_id = ( abm.sender.user_id if self.platform.unique_session else message.group_openid @@ -49,7 +51,8 @@ class botClient(Client): # 收到频道消息 async def on_at_message_create(self, message: botpy.message.Message): abm = QQOfficialPlatformAdapter._parse_from_qqofficial( - message, MessageType.GROUP_MESSAGE + message, + MessageType.GROUP_MESSAGE, ) abm.session_id = ( abm.sender.user_id if self.platform.unique_session else message.channel_id @@ -59,7 +62,8 @@ class botClient(Client): # 收到私聊消息 async def on_direct_message_create(self, message: botpy.message.DirectMessage): abm = QQOfficialPlatformAdapter._parse_from_qqofficial( - message, MessageType.FRIEND_MESSAGE + message, + MessageType.FRIEND_MESSAGE, ) abm.session_id = abm.sender.user_id self._commit(abm) @@ -67,7 +71,8 @@ class botClient(Client): # 收到 C2C 消息 async def on_c2c_message_create(self, message: botpy.message.C2CMessage): abm = QQOfficialPlatformAdapter._parse_from_qqofficial( - message, MessageType.FRIEND_MESSAGE + message, + MessageType.FRIEND_MESSAGE, ) abm.session_id = abm.sender.user_id self._commit(abm) @@ -80,14 +85,17 @@ class botClient(Client): self.platform.meta(), abm.session_id, self.platform.client, - ) + ), ) @register_platform_adapter("qq_official", "QQ 机器人官方 API 适配器") class QQOfficialPlatformAdapter(Platform): def __init__( - self, platform_config: dict, platform_settings: dict, event_queue: asyncio.Queue + self, + platform_config: dict, + platform_settings: dict, + event_queue: asyncio.Queue, ) -> None: super().__init__(event_queue) @@ -107,7 +115,8 @@ class QQOfficialPlatformAdapter(Platform): ) else: self.intents = botpy.Intents( - public_guild_messages=True, direct_message=guild_dm + public_guild_messages=True, + direct_message=guild_dm, ) self.client = botClient( intents=self.intents, @@ -120,7 +129,9 @@ class QQOfficialPlatformAdapter(Platform): self.test_mode = os.environ.get("TEST_MODE", "off") == "on" async def send_by_session( - self, session: MessageSesion, message_chain: MessageChain + self, + session: MessageSesion, + message_chain: MessageChain, ): raise NotImplementedError("QQ 机器人官方 API 适配器不支持 send_by_session") @@ -133,7 +144,7 @@ class QQOfficialPlatformAdapter(Platform): @staticmethod def _parse_from_qqofficial( - message: Union[botpy.message.Message, botpy.message.GroupMessage], + message: botpy.message.Message | botpy.message.GroupMessage, message_type: MessageType, ): abm = AstrBotMessage() @@ -142,10 +153,11 @@ class QQOfficialPlatformAdapter(Platform): abm.raw_message = message abm.message_id = message.id abm.tag = "qq_official" - msg: List[BaseMessageComponent] = [] + msg: list[BaseMessageComponent] = [] if isinstance(message, botpy.message.GroupMessage) or isinstance( - message, botpy.message.C2CMessage + message, + botpy.message.C2CMessage, ): if isinstance(message, botpy.message.GroupMessage): abm.sender = MessageMember(message.author.member_openid, "") @@ -167,7 +179,8 @@ class QQOfficialPlatformAdapter(Platform): abm.message = msg elif isinstance(message, botpy.message.Message) or isinstance( - message, botpy.message.DirectMessage + message, + botpy.message.DirectMessage, ): try: abm.self_id = str(message.mentions[0].id) @@ -175,7 +188,8 @@ class QQOfficialPlatformAdapter(Platform): abm.self_id = "" plain_content = message.content.replace( - "<@!" + str(abm.self_id) + ">", "" + "<@!" + str(abm.self_id) + ">", + "", ).strip() if message.attachments: @@ -189,7 +203,8 @@ class QQOfficialPlatformAdapter(Platform): abm.message = msg abm.message_str = plain_content abm.sender = MessageMember( - str(message.author.id), str(message.author.username) + str(message.author.id), + str(message.author.username), ) msg.append(At(qq="qq_official")) msg.append(Plain(plain_content)) diff --git a/astrbot/core/platform/sources/qqofficial_webhook/qo_webhook_adapter.py b/astrbot/core/platform/sources/qqofficial_webhook/qo_webhook_adapter.py index cc12e976..2b8c0b42 100644 --- a/astrbot/core/platform/sources/qqofficial_webhook/qo_webhook_adapter.py +++ b/astrbot/core/platform/sources/qqofficial_webhook/qo_webhook_adapter.py @@ -1,19 +1,21 @@ -import botpy -import logging import asyncio +import logging + +import botpy import botpy.message import botpy.types import botpy.types.message - from botpy import Client -from astrbot.api.platform import Platform, AstrBotMessage, MessageType, PlatformMetadata -from astrbot.api.event import MessageChain -from astrbot.core.platform.astr_message_event import MessageSesion -from .qo_webhook_event import QQOfficialWebhookMessageEvent -from ...register import register_platform_adapter -from .qo_webhook_server import QQOfficialWebhook -from ..qqofficial.qqofficial_platform_adapter import QQOfficialPlatformAdapter + from astrbot import logger +from astrbot.api.event import MessageChain +from astrbot.api.platform import AstrBotMessage, MessageType, Platform, PlatformMetadata +from astrbot.core.platform.astr_message_event import MessageSesion + +from ...register import register_platform_adapter +from ..qqofficial.qqofficial_platform_adapter import QQOfficialPlatformAdapter +from .qo_webhook_event import QQOfficialWebhookMessageEvent +from .qo_webhook_server import QQOfficialWebhook # remove logger handler for handler in logging.root.handlers[:]: @@ -28,7 +30,8 @@ class botClient(Client): # 收到群消息 async def on_group_at_message_create(self, message: botpy.message.GroupMessage): abm = QQOfficialPlatformAdapter._parse_from_qqofficial( - message, MessageType.GROUP_MESSAGE + message, + MessageType.GROUP_MESSAGE, ) abm.session_id = ( abm.sender.user_id if self.platform.unique_session else message.group_openid @@ -38,7 +41,8 @@ class botClient(Client): # 收到频道消息 async def on_at_message_create(self, message: botpy.message.Message): abm = QQOfficialPlatformAdapter._parse_from_qqofficial( - message, MessageType.GROUP_MESSAGE + message, + MessageType.GROUP_MESSAGE, ) abm.session_id = ( abm.sender.user_id if self.platform.unique_session else message.channel_id @@ -48,7 +52,8 @@ class botClient(Client): # 收到私聊消息 async def on_direct_message_create(self, message: botpy.message.DirectMessage): abm = QQOfficialPlatformAdapter._parse_from_qqofficial( - message, MessageType.FRIEND_MESSAGE + message, + MessageType.FRIEND_MESSAGE, ) abm.session_id = abm.sender.user_id self._commit(abm) @@ -56,7 +61,8 @@ class botClient(Client): # 收到 C2C 消息 async def on_c2c_message_create(self, message: botpy.message.C2CMessage): abm = QQOfficialPlatformAdapter._parse_from_qqofficial( - message, MessageType.FRIEND_MESSAGE + message, + MessageType.FRIEND_MESSAGE, ) abm.session_id = abm.sender.user_id self._commit(abm) @@ -64,15 +70,22 @@ class botClient(Client): def _commit(self, abm: AstrBotMessage): self.platform.commit_event( QQOfficialWebhookMessageEvent( - abm.message_str, abm, self.platform.meta(), abm.session_id, self - ) + abm.message_str, + abm, + self.platform.meta(), + abm.session_id, + self, + ), ) @register_platform_adapter("qq_official_webhook", "QQ 机器人官方 API 适配器(Webhook)") class QQOfficialWebhookPlatformAdapter(Platform): def __init__( - self, platform_config: dict, platform_settings: dict, event_queue: asyncio.Queue + self, + platform_config: dict, + platform_settings: dict, + event_queue: asyncio.Queue, ) -> None: super().__init__(event_queue) @@ -83,7 +96,9 @@ class QQOfficialWebhookPlatformAdapter(Platform): self.unique_session = platform_settings["unique_session"] intents = botpy.Intents( - public_messages=True, public_guild_messages=True, direct_message=True + public_messages=True, + public_guild_messages=True, + direct_message=True, ) self.client = botClient( intents=intents, # 已经无用 @@ -93,7 +108,9 @@ class QQOfficialWebhookPlatformAdapter(Platform): self.client.set_platform(self) async def send_by_session( - self, session: MessageSesion, message_chain: MessageChain + self, + session: MessageSesion, + message_chain: MessageChain, ): raise NotImplementedError("QQ 机器人官方 API 适配器不支持 send_by_session") @@ -106,7 +123,9 @@ class QQOfficialWebhookPlatformAdapter(Platform): async def run(self): self.webhook_helper = QQOfficialWebhook( - self.config, self._event_queue, self.client + self.config, + self._event_queue, + self.client, ) await self.webhook_helper.initialize() await self.webhook_helper.start_polling() diff --git a/astrbot/core/platform/sources/qqofficial_webhook/qo_webhook_event.py b/astrbot/core/platform/sources/qqofficial_webhook/qo_webhook_event.py index 4c0bf832..306db5e5 100644 --- a/astrbot/core/platform/sources/qqofficial_webhook/qo_webhook_event.py +++ b/astrbot/core/platform/sources/qqofficial_webhook/qo_webhook_event.py @@ -1,5 +1,7 @@ -from astrbot.api.platform import AstrBotMessage, PlatformMetadata from botpy import Client + +from astrbot.api.platform import AstrBotMessage, PlatformMetadata + from ..qqofficial.qqofficial_message_event import QQOfficialMessageEvent diff --git a/astrbot/core/platform/sources/qqofficial_webhook/qo_webhook_server.py b/astrbot/core/platform/sources/qqofficial_webhook/qo_webhook_server.py index 4a2eae74..65b7c701 100644 --- a/astrbot/core/platform/sources/qqofficial_webhook/qo_webhook_server.py +++ b/astrbot/core/platform/sources/qqofficial_webhook/qo_webhook_server.py @@ -1,10 +1,12 @@ -import quart -import logging import asyncio -from botpy import BotAPI, BotHttp, Client, Token, BotWebSocket, ConnectionSession -from astrbot.api import logger +import logging + +import quart +from botpy import BotAPI, BotHttp, BotWebSocket, Client, ConnectionSession, Token from cryptography.hazmat.primitives.asymmetric import ed25519 +from astrbot.api import logger + # remove logger handler for handler in logging.root.handlers[:]: logging.root.removeHandler(handler) @@ -27,7 +29,9 @@ class QQOfficialWebhook: self.server = quart.Quart(__name__) self.server.add_url_rule( - "/astrbot-qo-webhook/callback", view_func=self.callback, methods=["POST"] + "/astrbot-qo-webhook/callback", + view_func=self.callback, + methods=["POST"], ) self.client = botpy_client self.event_queue = event_queue @@ -62,7 +66,8 @@ class QQOfficialWebhook: seed = await self.repeat_seed(self.secret) private_key = ed25519.Ed25519PrivateKey.from_private_bytes(seed) msg = validation_payload.get("event_ts", "") + validation_payload.get( - "plain_token", "" + "plain_token", + "", ) # sign signature = private_key.sign(msg.encode()).hex() @@ -99,7 +104,7 @@ class QQOfficialWebhook: async def start_polling(self): logger.info( - f"将在 {self.callback_server_host}:{self.port} 端口启动 QQ 官方机器人 webhook 适配器。" + f"将在 {self.callback_server_host}:{self.port} 端口启动 QQ 官方机器人 webhook 适配器。", ) await self.server.run_task( host=self.callback_server_host, diff --git a/astrbot/core/platform/sources/satori/satori_adapter.py b/astrbot/core/platform/sources/satori/satori_adapter.py index a3f4f53e..fd90804f 100644 --- a/astrbot/core/platform/sources/satori/satori_adapter.py +++ b/astrbot/core/platform/sources/satori/satori_adapter.py @@ -1,13 +1,22 @@ import asyncio import json import time +from xml.etree import ElementTree as ET + import websockets -from websockets.asyncio.client import connect -from typing import Optional from aiohttp import ClientSession, ClientTimeout -from websockets.asyncio.client import ClientConnection +from websockets.asyncio.client import ClientConnection, connect + from astrbot.api import logger from astrbot.api.event import MessageChain +from astrbot.api.message_components import ( + At, + File, + Image, + Plain, + Record, + Reply, +) from astrbot.api.platform import ( AstrBotMessage, MessageMember, @@ -17,35 +26,30 @@ from astrbot.api.platform import ( register_platform_adapter, ) from astrbot.core.platform.astr_message_event import MessageSession -from astrbot.api.message_components import ( - Plain, - Image, - At, - File, - Record, - Reply, -) -from xml.etree import ElementTree as ET @register_platform_adapter( - "satori", - "Satori 协议适配器", + "satori", "Satori 协议适配器", support_streaming_message=False ) class SatoriPlatformAdapter(Platform): def __init__( - self, platform_config: dict, platform_settings: dict, event_queue: asyncio.Queue + self, + platform_config: dict, + platform_settings: dict, + event_queue: asyncio.Queue, ) -> None: super().__init__(event_queue) self.config = platform_config self.settings = platform_settings self.api_base_url = self.config.get( - "satori_api_base_url", "http://localhost:5140/satori/v1" + "satori_api_base_url", + "http://localhost:5140/satori/v1", ) self.token = self.config.get("satori_token", "") self.endpoint = self.config.get( - "satori_endpoint", "ws://localhost:5140/satori/v1/events" + "satori_endpoint", + "ws://localhost:5140/satori/v1/events", ) self.auto_reconnect = self.config.get("satori_auto_reconnect", True) self.heartbeat_interval = self.config.get("satori_heartbeat_interval", 10) @@ -55,23 +59,28 @@ class SatoriPlatformAdapter(Platform): name="satori", description="Satori 通用协议适配器", id=self.config["id"], + support_streaming_message=False, ) - self.ws: Optional[ClientConnection] = None - self.session: Optional[ClientSession] = None + self.ws: ClientConnection | None = None + self.session: ClientSession | None = None self.sequence = 0 self.logins = [] self.running = False - self.heartbeat_task: Optional[asyncio.Task] = None + self.heartbeat_task: asyncio.Task | None = None self.ready_received = False async def send_by_session( - self, session: MessageSession, message_chain: MessageChain + self, + session: MessageSession, + message_chain: MessageChain, ): from .satori_event import SatoriPlatformEvent await SatoriPlatformEvent.send_with_adapter( - self, message_chain, session.session_id + self, + message_chain, + session.session_id, ) await super().send_by_session(session, message_chain) @@ -85,10 +94,9 @@ class SatoriPlatformAdapter(Platform): try: if hasattr(ws, "closed"): return ws.closed - elif hasattr(ws, "close_code"): + if hasattr(ws, "close_code"): return ws.close_code is not None - else: - return False + return False except AttributeError: return False @@ -240,7 +248,7 @@ class SatoriPlatformAdapter(Platform): user_id = user.get("id", "") user_name = user.get("name", "") logger.info( - f"Satori 连接成功 - Bot {i + 1}: platform={platform}, user_id={user_id}, user_name={user_name}" + f"Satori 连接成功 - Bot {i + 1}: platform={platform}, user_id={user_id}, user_name={user_name}", ) if "sn" in body: @@ -282,7 +290,12 @@ class SatoriPlatformAdapter(Platform): return abm = await self.convert_satori_message( - message, user, channel, guild, login, timestamp + message, + user, + channel, + guild, + login, + timestamp, ) if abm: await self.handle_msg(abm) @@ -295,10 +308,10 @@ class SatoriPlatformAdapter(Platform): message: dict, user: dict, channel: dict, - guild: Optional[dict], + guild: dict | None, login: dict, - timestamp: Optional[int] = None, - ) -> Optional[AstrBotMessage]: + timestamp: int | None = None, + ) -> AstrBotMessage | None: try: abm = AstrBotMessage() abm.message_id = message.get("id", "") @@ -438,7 +451,7 @@ class SatoriPlatformAdapter(Platform): return prefixes - async def _extract_quote_element(self, content: str) -> Optional[dict]: + async def _extract_quote_element(self, content: str) -> dict | None: """提取标签信息""" try: # 处理命名空间前缀问题 @@ -451,7 +464,7 @@ class SatoriPlatformAdapter(Platform): [ f'xmlns:{prefix}="http://temp.uri/{prefix}"' for prefix in prefixes - ] + ], ) # 包装内容 @@ -483,14 +496,17 @@ class SatoriPlatformAdapter(Platform): inner_content += quote_element.text for child in quote_element: inner_content += ET.tostring( - child, encoding="unicode", method="xml" + child, + encoding="unicode", + method="xml", ) if child.tail: inner_content += child.tail # 构造移除了标签的内容 content_without_quote = content.replace( - ET.tostring(quote_element, encoding="unicode", method="xml"), "" + ET.tostring(quote_element, encoding="unicode", method="xml"), + "", ) return { @@ -506,7 +522,7 @@ class SatoriPlatformAdapter(Platform): logger.error(f"提取标签时发生错误: {e}") return None - async def _extract_quote_with_regex(self, content: str) -> Optional[dict]: + async def _extract_quote_with_regex(self, content: str) -> dict | None: """使用正则表达式提取quote标签信息""" import re @@ -529,7 +545,7 @@ class SatoriPlatformAdapter(Platform): "content_without_quote": content_without_quote, } - async def _convert_quote_message(self, quote: dict) -> Optional[AstrBotMessage]: + async def _convert_quote_message(self, quote: dict) -> AstrBotMessage | None: """转换引用消息""" try: quote_abm = AstrBotMessage() @@ -587,7 +603,7 @@ class SatoriPlatformAdapter(Platform): [ f'xmlns:{prefix}="http://temp.uri/{prefix}"' for prefix in prefixes - ] + ], ) # 包装内容 @@ -747,13 +763,15 @@ class SatoriPlatformAdapter(Platform): try: async with self.session.request( - method, url, json=data, headers=headers + method, + url, + json=data, + headers=headers, ) as response: if response.status == 200: result = await response.json() return result - else: - return {} + return {} except Exception as e: logger.error(f"Satori HTTP 请求异常: {e}") return {} diff --git a/astrbot/core/platform/sources/satori/satori_event.py b/astrbot/core/platform/sources/satori/satori_event.py index 78325c9a..81a0d222 100644 --- a/astrbot/core/platform/sources/satori/satori_event.py +++ b/astrbot/core/platform/sources/satori/satori_event.py @@ -1,19 +1,20 @@ from typing import TYPE_CHECKING + from astrbot.api import logger from astrbot.api.event import AstrMessageEvent, MessageChain -from astrbot.api.platform import AstrBotMessage, PlatformMetadata from astrbot.api.message_components import ( - Plain, - Image, At, File, - Record, - Video, - Reply, Forward, + Image, Node, Nodes, + Plain, + Record, + Reply, + Video, ) +from astrbot.api.platform import AstrBotMessage, PlatformMetadata if TYPE_CHECKING: from .satori_adapter import SatoriPlatformAdapter @@ -53,14 +54,17 @@ class SatoriPlatformEvent(AstrMessageEvent): @classmethod async def send_with_adapter( - cls, adapter: "SatoriPlatformAdapter", message: MessageChain, session_id: str + cls, + adapter: "SatoriPlatformAdapter", + message: MessageChain, + session_id: str, ): try: content_parts = [] for component in message.chain: component_content = await cls._convert_component_to_satori_static( - component + component, ) if component_content: content_parts.append(component_content) @@ -92,12 +96,15 @@ class SatoriPlatformEvent(AstrMessageEvent): user_id = user.get("id", "") if user else "" result = await adapter.send_http_request( - "POST", "/message.create", data, platform, user_id + "POST", + "/message.create", + data, + platform, + user_id, ) if result: return result - else: - return None + return None except Exception as e: logger.error(f"Satori 消息发送异常: {e}") @@ -140,7 +147,11 @@ class SatoriPlatformEvent(AstrMessageEvent): data = {"channel_id": channel_id, "content": content} result = await self.adapter.send_http_request( - "POST", "/message.create", data, platform, user_id + "POST", + "/message.create", + data, + platform, + user_id, ) if not result: logger.error("Satori 消息发送失败") @@ -178,9 +189,9 @@ class SatoriPlatformEvent(AstrMessageEvent): img_chain = MessageChain( [ Plain( - text=f'' - ) - ] + text=f'', + ), + ], ) await self.send(img_chain) except Exception as e: @@ -209,10 +220,10 @@ class SatoriPlatformEvent(AstrMessageEvent): ) return text - elif isinstance(component, At): + if isinstance(component, At): if component.qq: return f'' - elif component.name: + if component.name: return f'' elif isinstance(component, Image): @@ -264,7 +275,7 @@ class SatoriPlatformEvent(AstrMessageEvent): if node.content: for content_component in node.content: component_content = await self._convert_component_to_satori( - content_component + content_component, ) if component_content: content_parts.append(component_content) @@ -302,10 +313,10 @@ class SatoriPlatformEvent(AstrMessageEvent): ) return text - elif isinstance(component, At): + if isinstance(component, At): if component.qq: return f'' - elif component.name: + if component.name: return f'' elif isinstance(component, Image): @@ -358,7 +369,7 @@ class SatoriPlatformEvent(AstrMessageEvent): if node.content: for content_component in node.content: component_content = await cls._convert_component_to_satori_static( - content_component + content_component, ) if component_content: content_parts.append(component_content) @@ -395,8 +406,7 @@ class SatoriPlatformEvent(AstrMessageEvent): if node_parts: return f"{''.join(node_parts)}" - else: - return "" + return "" except Exception as e: logger.error(f"转换合并转发消息失败: {e}") @@ -415,8 +425,7 @@ class SatoriPlatformEvent(AstrMessageEvent): if node_parts: return f"{''.join(node_parts)}" - else: - return "" + return "" except Exception as e: logger.error(f"转换合并转发消息失败: {e}") diff --git a/astrbot/core/platform/sources/slack/client.py b/astrbot/core/platform/sources/slack/client.py index 7877e4f5..0411f73a 100644 --- a/astrbot/core/platform/sources/slack/client.py +++ b/astrbot/core/platform/sources/slack/client.py @@ -1,14 +1,16 @@ -import json -import hmac -import hashlib import asyncio +import hashlib +import hmac +import json import logging -from typing import Callable, Optional -from quart import Quart, request, Response -from slack_sdk.web.async_client import AsyncWebClient +from collections.abc import Callable + +from quart import Quart, Response, request from slack_sdk.socket_mode.aiohttp import SocketModeClient from slack_sdk.socket_mode.request import SocketModeRequest from slack_sdk.socket_mode.response import SocketModeResponse +from slack_sdk.web.async_client import AsyncWebClient + from astrbot.api import logger @@ -22,7 +24,7 @@ class SlackWebhookClient: host: str = "0.0.0.0", port: int = 3000, path: str = "/slack/events", - event_handler: Optional[Callable] = None, + event_handler: Callable | None = None, ): self.web_client = web_client self.signing_secret = signing_secret @@ -93,7 +95,7 @@ class SlackWebhookClient: async def start(self): """启动 Webhook 服务器""" logger.info( - f"Slack Webhook 服务器启动中，监听 {self.host}:{self.port}{self.path}..." + f"Slack Webhook 服务器启动中，监听 {self.host}:{self.port}{self.path}...", ) await self.app.run_task( @@ -119,7 +121,7 @@ class SlackSocketClient: self, web_client: AsyncWebClient, app_token: str, - event_handler: Optional[Callable] = None, + event_handler: Callable | None = None, ): self.web_client = web_client self.app_token = app_token diff --git a/astrbot/core/platform/sources/slack/slack_adapter.py b/astrbot/core/platform/sources/slack/slack_adapter.py index 7e75f3c2..d5427deb 100644 --- a/astrbot/core/platform/sources/slack/slack_adapter.py +++ b/astrbot/core/platform/sources/slack/slack_adapter.py @@ -1,34 +1,43 @@ -import time import asyncio -import uuid -import aiohttp -import re import base64 -from typing import Awaitable, Any -from slack_sdk.web.async_client import AsyncWebClient +import re +import time +import uuid +from collections.abc import Awaitable +from typing import Any + +import aiohttp from slack_sdk.socket_mode.request import SocketModeRequest +from slack_sdk.web.async_client import AsyncWebClient + +from astrbot.api import logger +from astrbot.api.event import MessageChain +from astrbot.api.message_components import * from astrbot.api.platform import ( - Platform, AstrBotMessage, MessageMember, MessageType, + Platform, PlatformMetadata, ) -from astrbot.api.event import MessageChain -from .slack_event import SlackMessageEvent -from .client import SlackWebhookClient, SlackSocketClient -from astrbot.api.message_components import * # noqa: F403 -from astrbot.api import logger from astrbot.core.platform.astr_message_event import MessageSesion + from ...register import register_platform_adapter +from .client import SlackSocketClient, SlackWebhookClient +from .slack_event import SlackMessageEvent @register_platform_adapter( - "slack", "适用于 Slack 的消息平台适配器，支持 Socket Mode 和 Webhook Mode。" + "slack", + "适用于 Slack 的消息平台适配器，支持 Socket Mode 和 Webhook Mode。", + support_streaming_message=False, ) class SlackAdapter(Platform): def __init__( - self, platform_config: dict, platform_settings: dict, event_queue: asyncio.Queue + self, + platform_config: dict, + platform_settings: dict, + event_queue: asyncio.Queue, ) -> None: super().__init__(event_queue) @@ -43,7 +52,8 @@ class SlackAdapter(Platform): self.webhook_host = platform_config.get("slack_webhook_host", "0.0.0.0") self.webhook_port = platform_config.get("slack_webhook_port", 3000) self.webhook_path = platform_config.get( - "slack_webhook_path", "/astrbot-slack-webhook/callback" + "slack_webhook_path", + "/astrbot-slack-webhook/callback", ) if not self.bot_token: @@ -59,6 +69,7 @@ class SlackAdapter(Platform): name="slack", description="适用于 Slack 的消息平台适配器，支持 Socket Mode 和 Webhook Mode。", id=self.config.get("id"), + support_streaming_message=False, ) # 初始化 Slack Web Client @@ -69,10 +80,13 @@ class SlackAdapter(Platform): self.bot_self_id = None async def send_by_session( - self, session: MessageSesion, message_chain: MessageChain + self, + session: MessageSesion, + message_chain: MessageChain, ): - blocks, text = SlackMessageEvent._parse_slack_blocks( - message_chain=message_chain, web_client=self.web_client + blocks, text = await SlackMessageEvent._parse_slack_blocks( + message_chain=message_chain, + web_client=self.web_client, ) try: @@ -150,7 +164,7 @@ class SlackAdapter(Platform): abm.message = [] # 优先使用 blocks 字段解析消息 - if "blocks" in event and event["blocks"]: + if event.get("blocks"): abm.message = self._parse_blocks(event["blocks"]) # 更新 message_str abm.message_str = "" @@ -166,7 +180,8 @@ class SlackAdapter(Platform): mentioned_user = await self.web_client.users_info(user=mention) user_data = mentioned_user["user"] user_name = user_data.get("real_name") or user_data.get( - "name", mention + "name", + mention, ) abm.message.append(At(qq=mention, name=user_name)) except Exception: @@ -189,7 +204,7 @@ class SlackAdapter(Platform): else: # TODO: 下载鉴权 abm.message.append( - File(name=file_name, file=file_url, url=file_url) + File(name=file_name, file=file_url, url=file_url), ) abm.raw_message = event @@ -209,39 +224,41 @@ class SlackAdapter(Platform): if element.get("type") == "rich_text_section": # 处理富文本段落 section_elements = element.get("elements", []) - text_content = "" - + text_parts = [] for section_element in section_elements: element_type = section_element.get("type", "") if element_type == "text": # 普通文本 - text_content += section_element.get("text", "") + text_parts.append(section_element.get("text", "")) elif element_type == "user": # @用户提及 user_id = section_element.get("user_id", "") if user_id: # 将之前的文本内容先添加到组件中 + text_content = "".join(text_parts) if text_content.strip(): message_components.append( - Plain(text=text_content) + Plain(text=text_content), ) - text_content = "" + text_parts = [] # 添加@提及组件 message_components.append(At(qq=user_id, name="")) elif element_type == "channel": # #频道提及 channel_id = section_element.get("channel_id", "") - text_content += f"#{channel_id}" + text_parts.append(f"#{channel_id}") elif element_type == "link": # 链接 url = section_element.get("url", "") link_text = section_element.get("text", url) - text_content += f"[{link_text}]({url})" + text_parts.append(f"[{link_text}]({url})") elif element_type == "emoji": # 表情符号 emoji_name = section_element.get("name", "") - text_content += f":{emoji_name}:" + text_parts.append(f":{emoji_name}:") + + text_content = "".join(text_parts) if text_content.strip(): message_components.append(Plain(text=text_content)) @@ -307,11 +324,10 @@ class SlackAdapter(Platform): content = await resp.read() base64_content = base64.b64encode(content).decode("utf-8") return base64_content - else: - logger.error( - f"Failed to download slack file: {resp.status} {await resp.text()}" - ) - raise Exception(f"下载文件失败: {resp.status}") + logger.error( + f"Failed to download slack file: {resp.status} {await resp.text()}", + ) + raise Exception(f"下载文件失败: {resp.status}") async def run(self) -> Awaitable[Any]: self.bot_self_id = await self.get_bot_user_id() @@ -323,7 +339,9 @@ class SlackAdapter(Platform): # 创建 Socket 客户端 self.socket_client = SlackSocketClient( - self.web_client, self.app_token, self._handle_socket_event + self.web_client, + self.app_token, + self._handle_socket_event, ) logger.info("Slack 适配器 (Socket Mode) 启动中...") @@ -344,13 +362,13 @@ class SlackAdapter(Platform): ) logger.info( - f"Slack 适配器 (Webhook Mode) 启动中，监听 {self.webhook_host}:{self.webhook_port}{self.webhook_path}..." + f"Slack 适配器 (Webhook Mode) 启动中，监听 {self.webhook_host}:{self.webhook_port}{self.webhook_path}...", ) await self.webhook_client.start() else: raise ValueError( - f"不支持的连接模式: {self.connection_mode}，请使用 'socket' 或 'webhook'" + f"不支持的连接模式: {self.connection_mode}，请使用 'socket' 或 'webhook'", ) async def _handle_webhook_event(self, event_data: dict): diff --git a/astrbot/core/platform/sources/slack/slack_event.py b/astrbot/core/platform/sources/slack/slack_event.py index 86f9f976..c918abba 100644 --- a/astrbot/core/platform/sources/slack/slack_event.py +++ b/astrbot/core/platform/sources/slack/slack_event.py @@ -1,16 +1,18 @@ import asyncio import re -from typing import AsyncGenerator +from collections.abc import AsyncGenerator + from slack_sdk.web.async_client import AsyncWebClient + +from astrbot.api import logger from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.message_components import ( + BaseMessageComponent, + File, Image, Plain, - File, - BaseMessageComponent, ) from astrbot.api.platform import Group, MessageMember -from astrbot.api import logger class SlackMessageEvent(AstrMessageEvent): @@ -27,12 +29,13 @@ class SlackMessageEvent(AstrMessageEvent): @staticmethod async def _from_segment_to_slack_block( - segment: BaseMessageComponent, web_client: AsyncWebClient + segment: BaseMessageComponent, + web_client: AsyncWebClient, ) -> dict: """将消息段转换为 Slack 块格式""" if isinstance(segment, Plain): return {"type": "section", "text": {"type": "mrkdwn", "text": segment.text}} - elif isinstance(segment, Image): + if isinstance(segment, Image): # upload file url = segment.url or segment.file if url.startswith("http"): @@ -61,7 +64,7 @@ class SlackMessageEvent(AstrMessageEvent): }, "alt_text": "图片", } - elif isinstance(segment, File): + if isinstance(segment, File): # upload file url = segment.url or segment.file response = await web_client.files_upload_v2( @@ -82,12 +85,12 @@ class SlackMessageEvent(AstrMessageEvent): "text": f"文件: <{file_url}|{segment.name or '文件'}>", }, } - else: - return {"type": "section", "text": {"type": "mrkdwn", "text": str(segment)}} + return {"type": "section", "text": {"type": "mrkdwn", "text": str(segment)}} @staticmethod async def _parse_slack_blocks( - message_chain: MessageChain, web_client: AsyncWebClient + message_chain: MessageChain, + web_client: AsyncWebClient, ): """解析成 Slack 块格式""" blocks = [] @@ -103,27 +106,29 @@ class SlackMessageEvent(AstrMessageEvent): { "type": "section", "text": {"type": "mrkdwn", "text": text_content}, - } + }, ) text_content = "" # 添加其他类型的块 block = await SlackMessageEvent._from_segment_to_slack_block( - segment, web_client + segment, + web_client, ) blocks.append(block) # 如果最后还有文本内容 if text_content.strip(): blocks.append( - {"type": "section", "text": {"type": "mrkdwn", "text": text_content}} + {"type": "section", "text": {"type": "mrkdwn", "text": text_content}}, ) return blocks, "" if blocks else text_content async def send(self, message: MessageChain): blocks, text = await SlackMessageEvent._parse_slack_blocks( - message, self.web_client + message, + self.web_client, ) try: @@ -143,28 +148,33 @@ class SlackMessageEvent(AstrMessageEvent): ) except Exception: # 如果块发送失败，尝试只发送文本 - fallback_text = "" + parts = [] for segment in message.chain: if isinstance(segment, Plain): - fallback_text += segment.text + parts.append(segment.text) elif isinstance(segment, File): - fallback_text += f" [文件: {segment.name}] " + parts.append(f" [文件: {segment.name}] ") elif isinstance(segment, Image): - fallback_text += " [图片] " + parts.append(" [图片] ") + fallback_text = "".join(parts) if self.get_group_id(): await self.web_client.chat_postMessage( - channel=self.get_group_id(), text=fallback_text + channel=self.get_group_id(), + text=fallback_text, ) else: await self.web_client.chat_postMessage( - channel=self.get_sender_id(), text=fallback_text + channel=self.get_sender_id(), + text=fallback_text, ) await super().send(message) async def send_streaming( - self, generator: AsyncGenerator, use_fallback: bool = False + self, + generator: AsyncGenerator, + use_fallback: bool = False, ): if not use_fallback: buffer = None @@ -174,7 +184,7 @@ class SlackMessageEvent(AstrMessageEvent): else: buffer.chain.extend(chain.chain) if not buffer: - return + return None buffer.squash_plain() await self.send(buffer) return await super().send_streaming(generator, use_fallback) @@ -211,7 +221,7 @@ class SlackMessageEvent(AstrMessageEvent): # 获取频道成员 members_response = await self.web_client.conversations_members( - channel=channel_id + channel=channel_id, ) members = [] @@ -224,7 +234,7 @@ class SlackMessageEvent(AstrMessageEvent): user_id=member_id, nickname=user_data.get("real_name") or user_data.get("name", member_id), - ) + ), ) except Exception: # 如果获取用户信息失败，使用默认信息 diff --git a/astrbot/core/platform/sources/telegram/tg_adapter.py b/astrbot/core/platform/sources/telegram/tg_adapter.py index 68ee6a98..88a9f7dc 100644 --- a/astrbot/core/platform/sources/telegram/tg_adapter.py +++ b/astrbot/core/platform/sources/telegram/tg_adapter.py @@ -37,7 +37,10 @@ else: @register_platform_adapter("telegram", "telegram 适配器") class TelegramPlatformAdapter(Platform): def __init__( - self, platform_config: dict, platform_settings: dict, event_queue: asyncio.Queue + self, + platform_config: dict, + platform_settings: dict, + event_queue: asyncio.Queue, ) -> None: super().__init__(event_queue) self.config = platform_config @@ -45,13 +48,15 @@ class TelegramPlatformAdapter(Platform): self.client_self_id = uuid.uuid4().hex[:8] base_url = self.config.get( - "telegram_api_base_url", "https://api.telegram.org/bot" + "telegram_api_base_url", + "https://api.telegram.org/bot", ) if not base_url: base_url = "https://api.telegram.org/bot" file_base_url = self.config.get( - "telegram_file_base_url", "https://api.telegram.org/file/bot" + "telegram_file_base_url", + "https://api.telegram.org/file/bot", ) if not file_base_url: file_base_url = "https://api.telegram.org/file/bot" @@ -59,10 +64,12 @@ class TelegramPlatformAdapter(Platform): self.base_url = base_url self.enable_command_register = self.config.get( - "telegram_command_register", True + "telegram_command_register", + True, ) self.enable_command_refresh = self.config.get( - "telegram_command_auto_refresh", True + "telegram_command_auto_refresh", + True, ) self.last_command_hash = None @@ -85,11 +92,15 @@ class TelegramPlatformAdapter(Platform): @override async def send_by_session( - self, session: MessageSesion, message_chain: MessageChain + self, + session: MessageSesion, + message_chain: MessageChain, ): from_username = session.session_id await TelegramPlatformEvent.send_with_client( - self.client, message_chain, from_username + self.client, + message_chain, + from_username, ) await super().send_by_session(session, message_chain) @@ -131,7 +142,7 @@ class TelegramPlatformAdapter(Platform): if commands: current_hash = hash( - tuple((cmd.command, cmd.description) for cmd in commands) + tuple((cmd.command, cmd.description) for cmd in commands), ) if current_hash == self.last_command_hash: return @@ -153,7 +164,9 @@ class TelegramPlatformAdapter(Platform): continue for event_filter in handler_metadata.event_filters: cmd_info = self._extract_command_info( - event_filter, handler_metadata, skip_commands + event_filter, + handler_metadata, + skip_commands, ) if cmd_info: cmd_name, description = cmd_info @@ -164,7 +177,9 @@ class TelegramPlatformAdapter(Platform): @staticmethod def _extract_command_info( - event_filter, handler_metadata, skip_commands: set + event_filter, + handler_metadata, + skip_commands: set, ) -> tuple[str, str] | None: """从事件过滤器中提取指令信息""" cmd_name = None @@ -199,11 +214,12 @@ class TelegramPlatformAdapter(Platform): async def start(self, update: Update, context: ContextTypes.DEFAULT_TYPE): if not update.effective_chat: logger.warning( - "Received a start command without an effective chat, skipping /start reply." + "Received a start command without an effective chat, skipping /start reply.", ) return await context.bot.send_message( - chat_id=update.effective_chat.id, text=self.config["start_message"] + chat_id=update.effective_chat.id, + text=self.config["start_message"], ) async def message_handler(self, update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -213,7 +229,10 @@ class TelegramPlatformAdapter(Platform): await self.handle_msg(abm) async def convert_message( - self, update: Update, context: ContextTypes.DEFAULT_TYPE, get_reply=True + self, + update: Update, + context: ContextTypes.DEFAULT_TYPE, + get_reply=True, ) -> AstrBotMessage | None: """转换 Telegram 的消息对象为 AstrBotMessage 对象。 @@ -244,7 +263,8 @@ class TelegramPlatformAdapter(Platform): logger.warning("[Telegram] Received a message without a from_user.") return None message.sender = MessageMember( - str(_from_user.id), _from_user.username or "Unknown" + str(_from_user.id), + _from_user.username or "Unknown", ) message.self_id = str(context.bot.username) message.raw_message = update @@ -274,7 +294,7 @@ class TelegramPlatformAdapter(Platform): message_str=reply_abm.message_str, text=reply_abm.message_str, qq=reply_abm.sender.user_id, - ) + ), ) if update.message.text: @@ -320,7 +340,7 @@ class TelegramPlatformAdapter(Platform): if message.message_str.strip() == "/start": await self.start(update, context) - return + return None elif update.message.voice: file = await update.message.voice.get_file() @@ -358,7 +378,7 @@ class TelegramPlatformAdapter(Platform): file_path = file.file_path if file_path is None: logger.warning( - f"Telegram document file_path is None, cannot save the file {file_name}." + f"Telegram document file_path is None, cannot save the file {file_name}.", ) else: message.message.append(Comp.File(file=file_path, name=file_name)) @@ -369,7 +389,7 @@ class TelegramPlatformAdapter(Platform): file_path = file.file_path if file_path is None: logger.warning( - f"Telegram video file_path is None, cannot save the file {file_name}." + f"Telegram video file_path is None, cannot save the file {file_name}.", ) else: message.message.append(Comp.Video(file=file_path, path=file.file_path)) diff --git a/astrbot/core/platform/sources/telegram/tg_event.py b/astrbot/core/platform/sources/telegram/tg_event.py index 2da7aafe..34fd86ad 100644 --- a/astrbot/core/platform/sources/telegram/tg_event.py +++ b/astrbot/core/platform/sources/telegram/tg_event.py @@ -1,22 +1,24 @@ +import asyncio import os import re -import asyncio + import telegramify_markdown +from telegram import ReactionTypeCustomEmoji, ReactionTypeEmoji +from telegram.ext import ExtBot + +from astrbot import logger from astrbot.api.event import AstrMessageEvent, MessageChain -from astrbot.api.platform import AstrBotMessage, PlatformMetadata, MessageType from astrbot.api.message_components import ( - Plain, - Image, - Reply, At, File, + Image, + Plain, Record, + Reply, ) -from telegram.ext import ExtBot -from astrbot.core.utils.io import download_file -from astrbot import logger +from astrbot.api.platform import AstrBotMessage, MessageType, PlatformMetadata from astrbot.core.utils.astrbot_path import get_astrbot_data_path -from telegram import ReactionTypeEmoji, ReactionTypeCustomEmoji +from astrbot.core.utils.io import download_file class TelegramPlatformEvent(AstrMessageEvent): @@ -68,7 +70,10 @@ class TelegramPlatformEvent(AstrMessageEvent): @classmethod async def send_with_client( - cls, client: ExtBot, message: MessageChain, user_name: str + cls, + client: ExtBot, + message: MessageChain, + user_name: str, ): image_path = None @@ -104,14 +109,18 @@ class TelegramPlatformEvent(AstrMessageEvent): for chunk in chunks: try: md_text = telegramify_markdown.markdownify( - chunk, max_line_length=None, normalize_whitespace=False + chunk, + max_line_length=None, + normalize_whitespace=False, ) await client.send_message( - text=md_text, parse_mode="MarkdownV2", **payload + text=md_text, + parse_mode="MarkdownV2", + **payload, ) except Exception as e: logger.warning( - f"MarkdownV2 send failed: {e}. Using plain text instead." + f"MarkdownV2 send failed: {e}. Using plain text instead.", ) await client.send_message(text=chunk, **payload) elif isinstance(i, Image): @@ -137,8 +146,7 @@ class TelegramPlatformEvent(AstrMessageEvent): await super().send(message) async def react(self, emoji: str | None, big: bool = False): - """ - 给原消息添加 Telegram 反应： + """给原消息添加 Telegram 反应： - 普通 emoji：传入 '👍'、'😂' 等 - 自定义表情：传入其 custom_emoji_id（纯数字字符串） - 取消本机器人的反应：传入 None 或空字符串 @@ -216,7 +224,9 @@ class TelegramPlatformEvent(AstrMessageEvent): i.file = path await self.client.send_document( - document=i.file, filename=i.name, **payload + document=i.file, + filename=i.name, + **payload, ) continue elif isinstance(i, Record): @@ -263,7 +273,9 @@ class TelegramPlatformEvent(AstrMessageEvent): if delta and current_content != delta: try: markdown_text = telegramify_markdown.markdownify( - delta, max_line_length=None, normalize_whitespace=False + delta, + max_line_length=None, + normalize_whitespace=False, ) await self.client.edit_message_text( text=markdown_text, @@ -274,7 +286,9 @@ class TelegramPlatformEvent(AstrMessageEvent): except Exception as e: logger.warning(f"Markdown转换失败，使用普通文本: {e!s}") await self.client.edit_message_text( - text=delta, chat_id=payload["chat_id"], message_id=message_id + text=delta, + chat_id=payload["chat_id"], + message_id=message_id, ) except Exception as e: logger.warning(f"编辑消息失败(streaming): {e!s}") diff --git a/astrbot/core/platform/sources/webchat/webchat_adapter.py b/astrbot/core/platform/sources/webchat/webchat_adapter.py index faec122a..ff5482f5 100644 --- a/astrbot/core/platform/sources/webchat/webchat_adapter.py +++ b/astrbot/core/platform/sources/webchat/webchat_adapter.py @@ -1,24 +1,27 @@ -import time import asyncio -import uuid import os -from typing import Awaitable, Any, Callable +import time +import uuid +from collections.abc import Awaitable, Callable +from typing import Any + +from astrbot import logger +from astrbot.core.message.components import Image, Plain, Record +from astrbot.core.message.message_event_result import MessageChain from astrbot.core.platform import ( - Platform, AstrBotMessage, MessageMember, MessageType, + Platform, PlatformMetadata, ) -from astrbot.core.message.message_event_result import MessageChain -from astrbot.core.message.components import Plain, Image, Record # noqa: F403 -from astrbot import logger -from .webchat_queue_mgr import webchat_queue_mgr, WebChatQueueMgr -from .webchat_event import WebChatMessageEvent from astrbot.core.platform.astr_message_event import MessageSesion -from ...register import register_platform_adapter from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from ...register import register_platform_adapter +from .webchat_event import WebChatMessageEvent +from .webchat_queue_mgr import WebChatQueueMgr, webchat_queue_mgr + class QueueListener: def __init__(self, webchat_queue_mgr: WebChatQueueMgr, callback: Callable) -> None: @@ -35,7 +38,7 @@ class QueueListener: await self.callback(data) except Exception as e: logger.error( - f"Error processing message from conversation {conversation_id}: {e}" + f"Error processing message from conversation {conversation_id}: {e}", ) break @@ -66,7 +69,10 @@ class QueueListener: @register_platform_adapter("webchat", "webchat") class WebChatAdapter(Platform): def __init__( - self, platform_config: dict, platform_settings: dict, event_queue: asyncio.Queue + self, + platform_config: dict, + platform_settings: dict, + event_queue: asyncio.Queue, ) -> None: super().__init__(event_queue) @@ -77,11 +83,15 @@ class WebChatAdapter(Platform): os.makedirs(self.imgs_dir, exist_ok=True) self.metadata = PlatformMetadata( - name="webchat", description="webchat", id="webchat" + name="webchat", + description="webchat", + id="webchat", ) async def send_by_session( - self, session: MessageSesion, message_chain: MessageChain + self, + session: MessageSesion, + message_chain: MessageChain, ): await WebChatMessageEvent._send(message_chain, session.session_id) await super().send_by_session(session, message_chain) @@ -106,13 +116,13 @@ class WebChatAdapter(Platform): if isinstance(payload["image_url"], list): for img in payload["image_url"]: abm.message.append( - Image.fromFileSystem(os.path.join(self.imgs_dir, img)) + Image.fromFileSystem(os.path.join(self.imgs_dir, img)), ) else: abm.message.append( Image.fromFileSystem( - os.path.join(self.imgs_dir, payload["image_url"]) - ) + os.path.join(self.imgs_dir, payload["image_url"]), + ), ) if payload["audio_url"]: if isinstance(payload["audio_url"], list): @@ -153,6 +163,9 @@ class WebChatAdapter(Platform): _, _, payload = message.raw_message # type: ignore message_event.set_extra("selected_provider", payload.get("selected_provider")) message_event.set_extra("selected_model", payload.get("selected_model")) + message_event.set_extra( + "enable_streaming", payload.get("enable_streaming", True) + ) self.commit_event(message_event) diff --git a/astrbot/core/platform/sources/webchat/webchat_event.py b/astrbot/core/platform/sources/webchat/webchat_event.py index 3bf1c0a2..4ced79b1 100644 --- a/astrbot/core/platform/sources/webchat/webchat_event.py +++ b/astrbot/core/platform/sources/webchat/webchat_event.py @@ -1,11 +1,13 @@ +import base64 import os import uuid -import base64 + from astrbot.api import logger from astrbot.api.event import AstrMessageEvent, MessageChain -from astrbot.api.message_components import Plain, Image, Record -from astrbot.core.utils.io import download_image_by_url +from astrbot.api.message_components import Image, Plain, Record from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.io import download_image_by_url + from .webchat_queue_mgr import webchat_queue_mgr imgs_dir = os.path.join(get_astrbot_data_path(), "webchat", "imgs") @@ -26,7 +28,7 @@ class WebChatMessageEvent(AstrMessageEvent): "type": "end", "data": "", "streaming": False, - } # end means this request is finished + }, # end means this request is finished ) return "" @@ -41,7 +43,7 @@ class WebChatMessageEvent(AstrMessageEvent): "data": data, "streaming": streaming, "chain_type": message.type, - } + }, ) elif isinstance(comp, Image): # save image to local @@ -70,7 +72,7 @@ class WebChatMessageEvent(AstrMessageEvent): "cid": cid, "data": data, "streaming": streaming, - } + }, ) elif isinstance(comp, Record): # save record to local @@ -94,7 +96,7 @@ class WebChatMessageEvent(AstrMessageEvent): "cid": cid, "data": data, "streaming": streaming, - } + }, ) else: logger.debug(f"webchat 忽略: {comp.type}") @@ -107,6 +109,7 @@ class WebChatMessageEvent(AstrMessageEvent): async def send_streaming(self, generator, use_fallback: bool = False): final_data = "" + reasoning_content = "" cid = self.session_id.split("!")[-1] web_chat_back_queue = webchat_queue_mgr.get_or_create_back_queue(cid) async for chain in generator: @@ -118,20 +121,28 @@ class WebChatMessageEvent(AstrMessageEvent): "data": final_data, "streaming": True, "cid": cid, - } + }, ) final_data = "" continue - final_data += await WebChatMessageEvent._send( - chain, session_id=self.session_id, streaming=True + + r = await WebChatMessageEvent._send( + chain, + session_id=self.session_id, + streaming=True, ) + if chain.type == "reasoning": + reasoning_content += chain.get_plain_text() + else: + final_data += r await web_chat_back_queue.put( { "type": "complete", # complete means we return the final result "data": final_data, + "reasoning": reasoning_content, "streaming": True, "cid": cid, - } + }, ) await super().send_streaming(generator, use_fallback) diff --git a/astrbot/core/platform/sources/wechatpadpro/wechatpadpro_adapter.py b/astrbot/core/platform/sources/wechatpadpro/wechatpadpro_adapter.py index 6b835ecb..e8629ec1 100644 --- a/astrbot/core/platform/sources/wechatpadpro/wechatpadpro_adapter.py +++ b/astrbot/core/platform/sources/wechatpadpro/wechatpadpro_adapter.py @@ -2,24 +2,24 @@ import asyncio import base64 import json import os -import traceback import time -from typing import Optional +import traceback import aiohttp import anyio import websockets + from astrbot import logger -from astrbot.api.message_components import Plain, Image, At, Record +from astrbot.api.message_components import At, Image, Plain, Record from astrbot.api.platform import Platform, PlatformMetadata from astrbot.core.message.message_event_result import MessageChain +from astrbot.core.platform.astr_message_event import MessageSesion from astrbot.core.platform.astrbot_message import ( AstrBotMessage, MessageMember, MessageType, ) from astrbot.core.utils.astrbot_path import get_astrbot_data_path -from astrbot.core.platform.astr_message_event import MessageSesion from ...register import register_platform_adapter from .wechatpadpro_message_event import WeChatPadProMessageEvent @@ -28,14 +28,19 @@ try: from .xml_data_parser import GeweDataParser except ImportError as e: logger.warning( - f"警告: 可能未安装 defusedxml 依赖库，将导致无法解析微信的 表情包、引用 类型的消息: {str(e)}" + f"警告: 可能未安装 defusedxml 依赖库，将导致无法解析微信的 表情包、引用 类型的消息: {e!s}", ) -@register_platform_adapter("wechatpadpro", "WeChatPadPro 消息平台适配器") +@register_platform_adapter( + "wechatpadpro", "WeChatPadPro 消息平台适配器", support_streaming_message=False +) class WeChatPadProAdapter(Platform): def __init__( - self, platform_config: dict, platform_settings: dict, event_queue: asyncio.Queue + self, + platform_config: dict, + platform_settings: dict, + event_queue: asyncio.Queue, ) -> None: super().__init__(event_queue) self._shutdown_event = None @@ -48,6 +53,7 @@ class WeChatPadProAdapter(Platform): name="wechatpadpro", description="WeChatPadPro 消息平台适配器", id=self.config.get("id", "wechatpadpro"), + support_streaming_message=False, ) # 保存配置信息 @@ -55,16 +61,19 @@ class WeChatPadProAdapter(Platform): self.host = self.config.get("host") self.port = self.config.get("port") self.active_mesasge_poll: bool = self.config.get( - "wpp_active_message_poll", False + "wpp_active_message_poll", + False, ) self.active_message_poll_interval: int = self.config.get( - "wpp_active_message_poll_interval", 5 + "wpp_active_message_poll_interval", + 5, ) self.base_url = f"http://{self.host}:{self.port}" self.auth_key = None # 用于保存生成的授权码 self.wxid = None # 用于保存登录成功后的 wxid self.credentials_file = os.path.join( - get_astrbot_data_path(), "wechatpadpro_credentials.json" + get_astrbot_data_path(), + "wechatpadpro_credentials.json", ) # 持久化文件路径 self.ws_handle_task = None @@ -81,9 +90,7 @@ class WeChatPadProAdapter(Platform): self.max_text_cache = 100 async def run(self) -> None: - """ - 启动平台适配器的运行实例。 - """ + """启动平台适配器的运行实例。""" logger.info("WeChatPadPro 适配器正在启动...") if loaded_credentials := self.load_credentials(): @@ -132,12 +139,10 @@ class WeChatPadProAdapter(Platform): logger.info("WeChatPadPro 适配器已停止。") def load_credentials(self): - """ - 从文件中加载 auth_key 和 wxid。 - """ + """从文件中加载 auth_key 和 wxid。""" if os.path.exists(self.credentials_file): try: - with open(self.credentials_file, "r") as f: + with open(self.credentials_file) as f: credentials = json.load(f) logger.info("成功加载 WeChatPadPro 凭据。") return credentials @@ -146,9 +151,7 @@ class WeChatPadProAdapter(Platform): return None def save_credentials(self): - """ - 将 auth_key 和 wxid 保存到文件。 - """ + """将 auth_key 和 wxid 保存到文件。""" credentials = { "auth_key": self.auth_key, "wxid": self.wxid, @@ -163,9 +166,7 @@ class WeChatPadProAdapter(Platform): logger.error(f"保存 WeChatPadPro 凭据失败: {e}") async def check_online_status(self): - """ - 检查 WeChatPadPro 设备是否在线。 - """ + """检查 WeChatPadPro 设备是否在线。""" if not self.auth_key: return False url = f"{self.base_url}/login/GetLoginStatus" @@ -182,25 +183,23 @@ class WeChatPadProAdapter(Platform): logger.info("WeChatPadPro 设备当前在线。") return True # login_state == 3 为离线状态 - elif login_state == 3: + if login_state == 3: logger.info("WeChatPadPro 设备不在线。") return False - else: - logger.error(f"未知的在线状态: {response_data}") - return False + logger.error(f"未知的在线状态: {response_data}") + return False # Code == 300 为微信退出状态。 - elif response.status == 200 and response_data.get("Code") == 300: + if response.status == 200 and response_data.get("Code") == 300: logger.info("WeChatPadPro 设备已退出。") return False - elif response.status == 200 and response_data.get("Code") == -2: + if response.status == 200 and response_data.get("Code") == -2: # 该链接不存在 self.auth_key = None return False - else: - logger.error( - f"检查在线状态失败: {response.status}, {response_data}" - ) - return False + logger.error( + f"检查在线状态失败: {response.status}, {response_data}", + ) + return False except aiohttp.ClientConnectorError as e: logger.error(f"连接到 WeChatPadPro 服务失败: {e}") @@ -221,9 +220,7 @@ class WeChatPadProAdapter(Platform): return None async def generate_auth_key(self): - """ - 生成授权码。 - """ + """生成授权码。""" url = f"{self.base_url}/admin/GenAuthKey1" params = {"key": self.admin_key} payload = {"Count": 1, "Days": 365} # 生成一个有效期365天的授权码 @@ -235,7 +232,7 @@ class WeChatPadProAdapter(Platform): async with session.post(url, params=params, json=payload) as response: if response.status != 200: logger.error( - f"生成授权码失败: {response.status}, {await response.text()}" + f"生成授权码失败: {response.status}, {await response.text()}", ) return @@ -248,7 +245,7 @@ class WeChatPadProAdapter(Platform): logger.info("成功获取授权码") else: logger.error( - f"生成授权码成功但未找到授权码: {response_data}" + f"生成授权码成功但未找到授权码: {response_data}", ) else: logger.error(f"生成授权码失败: {response_data}") @@ -258,9 +255,7 @@ class WeChatPadProAdapter(Platform): logger.error(f"生成授权码时发生错误: {e}") async def get_login_qr_code(self): - """ - 获取登录二维码地址。 - """ + """获取登录二维码地址。""" url = f"{self.base_url}/login/GetLoginQrCodeNew" params = {"key": self.auth_key} payload = {} # 根据文档，这个接口的 body 可以为空 @@ -272,26 +267,24 @@ class WeChatPadProAdapter(Platform): if response.status == 200 and response_data.get("Code") == 200: # 二维码地址在 Data.QrCodeUrl 字段中 if response_data.get("Data") and response_data["Data"].get( - "QrCodeUrl" + "QrCodeUrl", ): return response_data["Data"]["QrCodeUrl"] - else: - logger.error( - f"获取登录二维码成功但未找到二维码地址: {response_data}" - ) - return None - elif "该 key 无效" in response_data.get("Text"): logger.error( - "授权码无效，已经清除。请重新启动 AstrBot 或者本消息适配器。原因也可能是 WeChatPadPro 的 MySQL 服务没有启动成功，请检查 WeChatPadPro 服务的日志。" + f"获取登录二维码成功但未找到二维码地址: {response_data}", + ) + return None + if "该 key 无效" in response_data.get("Text"): + logger.error( + "授权码无效，已经清除。请重新启动 AstrBot 或者本消息适配器。原因也可能是 WeChatPadPro 的 MySQL 服务没有启动成功，请检查 WeChatPadPro 服务的日志。", ) self.auth_key = None self.save_credentials() return None - else: - logger.error( - f"获取登录二维码失败: {response.status}, {response_data}" - ) - return None + logger.error( + f"获取登录二维码失败: {response.status}, {response_data}", + ) + return None except aiohttp.ClientConnectorError as e: logger.error(f"连接到 WeChatPadPro 服务失败: {e}") return None @@ -300,8 +293,7 @@ class WeChatPadProAdapter(Platform): return None async def check_login_status(self): - """ - 循环检测扫码状态。 + """循环检测扫码状态。 尝试 6 次后跳出循环，添加倒计时。 返回 True 如果登录成功，否则返回 False。 """ @@ -325,31 +317,31 @@ class WeChatPadProAdapter(Platform): ): status = response_data["Data"]["state"] logger.info( - f"第 {attempts + 1} 次尝试，当前登录状态: {status}，还剩{countdown - attempts * 5}秒" + f"第 {attempts + 1} 次尝试，当前登录状态: {status}，还剩{countdown - attempts * 5}秒", ) if status == 2: # 状态 2 表示登录成功 self.wxid = response_data["Data"].get("wxid") self.wxnewpass = response_data["Data"].get( - "wxnewpass" + "wxnewpass", ) logger.info( - f"登录成功，wxid: {self.wxid}, wxnewpass: {self.wxnewpass}" + f"登录成功，wxid: {self.wxid}, wxnewpass: {self.wxnewpass}", ) self.save_credentials() # 登录成功后保存凭据 return True - elif status == -2: # 二维码过期 + if status == -2: # 二维码过期 logger.error("二维码已过期，请重新获取。") return False else: logger.error( - f"检测登录状态成功但未找到登录状态: {response_data}" + f"检测登录状态成功但未找到登录状态: {response_data}", ) elif response_data.get("Code") == 300: # "不存在状态" pass else: logger.info( - f"检测登录状态失败: {response.status}, {response_data}" + f"检测登录状态失败: {response.status}, {response_data}", ) except aiohttp.ClientConnectorError as e: @@ -368,13 +360,11 @@ class WeChatPadProAdapter(Platform): return False async def connect_websocket(self): - """ - 建立 WebSocket 连接并处理接收到的消息。 - """ + """建立 WebSocket 连接并处理接收到的消息。""" os.environ["no_proxy"] = f"localhost,127.0.0.1,{self.host}" ws_url = f"ws://{self.host}:{self.port}/ws/GetSyncMsg?key={self.auth_key}" logger.info( - f"正在连接 WebSocket: ws://{self.host}:{self.port}/ws/GetSyncMsg?key=***" + f"正在连接 WebSocket: ws://{self.host}:{self.port}/ws/GetSyncMsg?key=***", ) while True: try: @@ -389,7 +379,8 @@ class WeChatPadProAdapter(Platform): while True: try: message = await asyncio.wait_for( - websocket.recv(), timeout=wait_time + websocket.recv(), + timeout=wait_time, ) # logger.debug(message) # 不显示原始消息内容 asyncio.create_task(self.handle_websocket_message(message)) @@ -404,14 +395,12 @@ class WeChatPadProAdapter(Platform): break except Exception as e: logger.error( - f"WebSocket 连接失败: {e}, 请检查WeChatPadPro服务状态，或尝试重启WeChatPadPro适配器。" + f"WebSocket 连接失败: {e}, 请检查WeChatPadPro服务状态，或尝试重启WeChatPadPro适配器。", ) await asyncio.sleep(5) async def handle_websocket_message(self, message: str): - """ - 处理从 WebSocket 接收到的消息。 - """ + """处理从 WebSocket 接收到的消息。""" logger.debug(f"收到 WebSocket 消息: {message}") try: message_data = json.loads(message) @@ -441,9 +430,7 @@ class WeChatPadProAdapter(Platform): logger.error(f"处理 WebSocket 消息时发生错误: {e}") async def convert_message(self, raw_message: dict) -> AstrBotMessage | None: - """ - 将 WeChatPadPro 原始消息转换为 AstrBotMessage。 - """ + """将 WeChatPadPro 原始消息转换为 AstrBotMessage。""" abm = AstrBotMessage() abm.raw_message = raw_message abm.message_id = str(raw_message.get("msg_id")) @@ -452,7 +439,7 @@ class WeChatPadProAdapter(Platform): if int(time.time()) - abm.timestamp > 180: logger.warning( - f"忽略 3 分钟前的旧消息：消息时间戳 {abm.timestamp} 超过当前时间 {int(time.time())}。" + f"忽略 3 分钟前的旧消息：消息时间戳 {abm.timestamp} 超过当前时间 {int(time.time())}。", ) return None @@ -476,7 +463,12 @@ class WeChatPadProAdapter(Platform): # 先判断群聊/私聊并设置基本属性 if await self._process_chat_type( - abm, raw_message, from_user_name, to_user_name, content, push_content + abm, + raw_message, + from_user_name, + to_user_name, + content, + push_content, ): # 再根据消息类型处理消息内容 await self._process_message_content(abm, raw_message, msg_type, content) @@ -493,9 +485,7 @@ class WeChatPadProAdapter(Platform): content: str, push_content: str, ): - """ - 判断消息是群聊还是私聊，并设置 AstrBotMessage 的基本属性。 - """ + """判断消息是群聊还是私聊，并设置 AstrBotMessage 的基本属性。""" if from_user_name == "weixin": return False at_me = False @@ -510,7 +500,8 @@ class WeChatPadProAdapter(Platform): # 获取群聊发送者的nickname if sender_wxid: accurate_nickname = await self._get_group_member_nickname( - abm.group_id, sender_wxid + abm.group_id, + sender_wxid, ) if accurate_nickname: abm.sender.nickname = accurate_nickname @@ -539,11 +530,11 @@ class WeChatPadProAdapter(Platform): return True async def _get_group_member_nickname( - self, group_id: str, member_wxid: str - ) -> Optional[str]: - """ - 通过接口获取群成员的昵称。 - """ + self, + group_id: str, + member_wxid: str, + ) -> str | None: + """通过接口获取群成员的昵称。""" url = f"{self.base_url}/group/GetChatroomMemberDetail" params = {"key": self.auth_key} payload = { @@ -565,11 +556,11 @@ class WeChatPadProAdapter(Platform): if member.get("user_name") == member_wxid: return member.get("nick_name") logger.warning( - f"在群 {group_id} 中未找到成员 {member_wxid} 的昵称" + f"在群 {group_id} 中未找到成员 {member_wxid} 的昵称", ) else: logger.error( - f"获取群成员详情失败: {response.status}, {response_data}" + f"获取群成员详情失败: {response.status}, {response_data}", ) return None except aiohttp.ClientConnectorError as e: @@ -580,7 +571,10 @@ class WeChatPadProAdapter(Platform): return None async def _download_raw_image( - self, from_user_name: str, to_user_name: str, msg_id: int + self, + from_user_name: str, + to_user_name: str, + msg_id: int, ): """下载原始图片。""" url = f"{self.base_url}/message/GetMsgBigImg" @@ -598,9 +592,8 @@ class WeChatPadProAdapter(Platform): async with session.post(url, params=params, json=payload) as response: if response.status == 200: return await response.json() - else: - logger.error(f"下载图片失败: {response.status}") - return None + logger.error(f"下载图片失败: {response.status}") + return None except aiohttp.ClientConnectorError as e: logger.error(f"连接到 WeChatPadPro 服务失败: {e}") return None @@ -609,7 +602,11 @@ class WeChatPadProAdapter(Platform): return None async def download_voice( - self, to_user_name: str, new_msg_id: str, bufid: str, length: int + self, + to_user_name: str, + new_msg_id: str, + bufid: str, + length: int, ): """下载原始音频。""" url = f"{self.base_url}/message/GetMsgVoice" @@ -635,11 +632,13 @@ class WeChatPadProAdapter(Platform): return None async def _process_message_content( - self, abm: AstrBotMessage, raw_message: dict, msg_type: int, content: str + self, + abm: AstrBotMessage, + raw_message: dict, + msg_type: int, + content: str, ): - """ - 根据消息类型处理消息内容，填充 AstrBotMessage 的 message 列表。 - """ + """根据消息类型处理消息内容，填充 AstrBotMessage 的 message 列表。""" if msg_type == 1: # 文本消息 abm.message_str = content if abm.type == MessageType.GROUP_MESSAGE: @@ -671,10 +670,12 @@ class WeChatPadProAdapter(Platform): if at_me: # 被@了，在消息开头插入At组件（参考gewechat的做法） bot_nickname = await self._get_group_member_nickname( - abm.group_id, abm.self_id + abm.group_id, + abm.self_id, ) abm.message.insert( - 0, At(qq=abm.self_id, name=bot_nickname or abm.self_id) + 0, + At(qq=abm.self_id, name=bot_nickname or abm.self_id), ) # 只有当消息内容不仅仅是@时才添加Plain组件 @@ -727,7 +728,9 @@ class WeChatPadProAdapter(Platform): to_user_name = raw_message.get("to_user_name", {}).get("str", "") msg_id = raw_message.get("msg_id") image_resp = await self._download_raw_image( - from_user_name, to_user_name, msg_id + from_user_name, + to_user_name, + msg_id, ) image_bs64_data = ( image_resp.get("Data", {}).get("Data", {}).get("Buffer", None) @@ -789,7 +792,8 @@ class WeChatPadProAdapter(Platform): voice_bs64_data = base64.b64decode(voice_bs64_data) temp_dir = os.path.join(get_astrbot_data_path(), "temp") file_path = os.path.join( - temp_dir, f"wechatpadpro_voice_{abm.message_id}.silk" + temp_dir, + f"wechatpadpro_voice_{abm.message_id}.silk", ) async with await anyio.open_file(file_path, "wb") as f: @@ -819,9 +823,7 @@ class WeChatPadProAdapter(Platform): logger.warning(f"收到未处理的消息类型: {msg_type}。") async def terminate(self): - """ - 终止一个平台的运行实例。 - """ + """终止一个平台的运行实例。""" logger.info("终止 WeChatPadPro 适配器。") try: if self.ws_handle_task: @@ -831,13 +833,13 @@ class WeChatPadProAdapter(Platform): pass def meta(self) -> PlatformMetadata: - """ - 得到一个平台的元数据。 - """ + """得到一个平台的元数据。""" return self.metadata async def send_by_session( - self, session: MessageSesion, message_chain: MessageChain + self, + session: MessageSesion, + message_chain: MessageChain, ): dummy_message_obj = AstrBotMessage() dummy_message_obj.session_id = session.session_id @@ -864,9 +866,7 @@ class WeChatPadProAdapter(Platform): await sending_event.send(message_chain) async def get_contact_list(self): - """ - 获取联系人列表。 - """ + """获取联系人列表。""" url = f"{self.base_url}/friend/GetContactList" params = {"key": self.auth_key} payload = {"CurrentChatRoomContactSeq": 0, "CurrentWxcontactSeq": 0} @@ -884,9 +884,8 @@ class WeChatPadProAdapter(Platform): .get("contactUsernameList", []) ) return contact_list - else: - logger.error(f"获取联系人列表失败: {result}") - return None + logger.error(f"获取联系人列表失败: {result}") + return None except aiohttp.ClientConnectorError as e: logger.error(f"连接到 WeChatPadPro 服务失败: {e}") return None @@ -895,11 +894,11 @@ class WeChatPadProAdapter(Platform): return None async def get_contact_details_list( - self, room_wx_id_list: list[str] = None, user_names: list[str] = None - ) -> Optional[dict]: - """ - 获取联系人详情列表。 - """ + self, + room_wx_id_list: list[str] = None, + user_names: list[str] = None, + ) -> dict | None: + """获取联系人详情列表。""" if room_wx_id_list is None: room_wx_id_list = [] if user_names is None: @@ -917,9 +916,8 @@ class WeChatPadProAdapter(Platform): if result.get("Code") == 200 and result.get("Data"): contact_list = result.get("Data", {}).get("contactList", {}) return contact_list - else: - logger.error(f"获取联系人详情列表失败: {result}") - return None + logger.error(f"获取联系人详情列表失败: {result}") + return None except aiohttp.ClientConnectorError as e: logger.error(f"连接到 WeChatPadPro 服务失败: {e}") return None diff --git a/astrbot/core/platform/sources/wechatpadpro/wechatpadpro_message_event.py b/astrbot/core/platform/sources/wechatpadpro/wechatpadpro_message_event.py index 2bd3a1b8..08ab2701 100644 --- a/astrbot/core/platform/sources/wechatpadpro/wechatpadpro_message_event.py +++ b/astrbot/core/platform/sources/wechatpadpro/wechatpadpro_message_event.py @@ -1,6 +1,7 @@ import asyncio import base64 import io +from collections.abc import AsyncGenerator from typing import TYPE_CHECKING import aiohttp @@ -10,8 +11,8 @@ from astrbot import logger from astrbot.core.message.components import ( Image, Plain, - WechatEmoji, Record, + WechatEmoji, ) # Import Image from astrbot.core.message.message_event_result import MessageChain from astrbot.core.platform.astr_message_event import AstrMessageEvent @@ -50,14 +51,29 @@ class WeChatPadProMessageEvent(AstrMessageEvent): await self._send_voice(session, comp) await super().send(message) + async def send_streaming( + self, generator: AsyncGenerator[MessageChain, None], use_fallback: bool = False + ): + buffer = None + async for chain in generator: + if not buffer: + buffer = chain + else: + buffer.chain.extend(chain.chain) + if not buffer: + return None + buffer.squash_plain() + await self.send(buffer) + return await super().send_streaming(generator, use_fallback) + async def _send_image(self, session: aiohttp.ClientSession, comp: Image): b64 = await comp.convert_to_base64() raw = self._validate_base64(b64) b64c = self._compress_image(raw) payload = { "MsgItem": [ - {"ImageContent": b64c, "MsgType": 3, "ToUserName": self.session_id} - ] + {"ImageContent": b64c, "MsgType": 3, "ToUserName": self.session_id}, + ], } url = f"{self.adapter.base_url}/message/SendImageNewMessage" await self._post(session, url, payload) @@ -66,7 +82,8 @@ class WeChatPadProMessageEvent(AstrMessageEvent): if ( self.message_obj.type == MessageType.GROUP_MESSAGE # 确保是群聊消息 and self.adapter.settings.get( - "reply_with_mention", False + "reply_with_mention", + False, ) # 检查适配器设置是否启用 reply_with_mention and self.message_obj.sender # 确保有发送者信息 and ( @@ -91,8 +108,8 @@ class WeChatPadProMessageEvent(AstrMessageEvent): "MsgType": 1, "TextContent": message_text, "ToUserName": session_id, - } - ] + }, + ], } url = f"{self.adapter.base_url}/message/SendTextMessage" await self._post(session, url, payload) @@ -104,8 +121,8 @@ class WeChatPadProMessageEvent(AstrMessageEvent): "EmojiMd5": comp.md5, "EmojiSize": comp.md5_len, "ToUserName": self.session_id, - } - ] + }, + ], } url = f"{self.adapter.base_url}/message/SendEmojiMessage" await self._post(session, url, payload) diff --git a/astrbot/core/platform/sources/wechatpadpro/xml_data_parser.py b/astrbot/core/platform/sources/wechatpadpro/xml_data_parser.py index 054ca1b4..09924edb 100644 --- a/astrbot/core/platform/sources/wechatpadpro/xml_data_parser.py +++ b/astrbot/core/platform/sources/wechatpadpro/xml_data_parser.py @@ -1,10 +1,13 @@ from defusedxml import ElementTree as eT + from astrbot.api import logger from astrbot.api.message_components import ( - WechatEmoji as Emoji, - Plain, - Image, BaseMessageComponent, + Image, + Plain, +) +from astrbot.api.message_components import ( + WechatEmoji as Emoji, ) @@ -15,7 +18,7 @@ class GeweDataParser: is_private_chat: bool = False, cached_texts=None, cached_images=None, - raw_message: dict = None, + raw_message: dict | None = None, downloader=None, ): self._xml = None @@ -47,9 +50,7 @@ class GeweDataParser: raise async def parse_mutil_49(self) -> list[BaseMessageComponent] | None: - """ - 处理 msg_type == 49 的多种 appmsg 类型（目前支持 type==57） - """ + """处理 msg_type == 49 的多种 appmsg 类型（目前支持 type==57）""" try: appmsg_type = self._format_to_xml().findtext(".//appmsg/type") if appmsg_type == "57": @@ -59,9 +60,7 @@ class GeweDataParser: return None async def parse_reply(self) -> list[BaseMessageComponent]: - """ - 处理 type == 57 的引用消息：支持文本（1）、图片（3）、嵌套49（49） - """ + """处理 type == 57 的引用消息：支持文本（1）、图片（3）、嵌套49（49）""" components = [] try: @@ -96,7 +95,9 @@ class GeweDataParser: ) if cdn_url and self.downloader: image_resp = await self.downloader( - self.from_user_name, self.to_user_name, self.msg_id + self.from_user_name, + self.to_user_name, + self.msg_id, ) quoted_image_b64 = ( image_resp.get("Data", {}) @@ -111,11 +112,11 @@ class GeweDataParser: [ Image.fromBase64(quoted_image_b64), Plain(f"[引用] {nickname}: [引用的图片]"), - ] + ], ) else: components.append( - Plain(f"[引用] {nickname}: [引用的图片 - 未能获取]") + Plain(f"[引用] {nickname}: [引用的图片 - 未能获取]"), ) case 49: # 嵌套引用 @@ -143,9 +144,7 @@ class GeweDataParser: return components def parse_emoji(self) -> Emoji | None: - """ - 处理 msg_type == 47 的表情消息（emoji） - """ + """处理 msg_type == 47 的表情消息（emoji）""" try: emoji_element = self._format_to_xml().find(".//emoji") if emoji_element is not None: diff --git a/astrbot/core/platform/sources/wecom/wecom_adapter.py b/astrbot/core/platform/sources/wecom/wecom_adapter.py index 50341a8a..1ea4c8e2 100644 --- a/astrbot/core/platform/sources/wecom/wecom_adapter.py +++ b/astrbot/core/platform/sources/wecom/wecom_adapter.py @@ -41,10 +41,14 @@ class WecomServer: self.port = int(config.get("port")) self.callback_server_host = config.get("callback_server_host", "0.0.0.0") self.server.add_url_rule( - "/callback/command", view_func=self.verify, methods=["GET"] + "/callback/command", + view_func=self.verify, + methods=["GET"], ) self.server.add_url_rule( - "/callback/command", view_func=self.callback_command, methods=["POST"] + "/callback/command", + view_func=self.callback_command, + methods=["POST"], ) self.event_queue = event_queue @@ -94,7 +98,7 @@ class WecomServer: async def start_polling(self): logger.info( - f"将在 {self.callback_server_host}:{self.port} 端口启动 企业微信 适配器。" + f"将在 {self.callback_server_host}:{self.port} 端口启动 企业微信 适配器。", ) await self.server.run_task( host=self.callback_server_host, @@ -106,24 +110,27 @@ class WecomServer: await self.shutdown_event.wait() -@register_platform_adapter("wecom", "wecom 适配器") +@register_platform_adapter("wecom", "wecom 适配器", support_streaming_message=False) class WecomPlatformAdapter(Platform): def __init__( - self, platform_config: dict, platform_settings: dict, event_queue: asyncio.Queue + self, + platform_config: dict, + platform_settings: dict, + event_queue: asyncio.Queue, ) -> None: super().__init__(event_queue) self.config = platform_config self.settingss = platform_settings self.client_self_id = uuid.uuid4().hex[:8] self.api_base_url = platform_config.get( - "api_base_url", "https://qyapi.weixin.qq.com/cgi-bin/" + "api_base_url", + "https://qyapi.weixin.qq.com/cgi-bin/", ) if not self.api_base_url: self.api_base_url = "https://qyapi.weixin.qq.com/cgi-bin/" - if self.api_base_url.endswith("/"): - self.api_base_url = self.api_base_url[:-1] + self.api_base_url = self.api_base_url.removesuffix("/") if not self.api_base_url.endswith("/cgi-bin"): self.api_base_url += "/cgi-bin" @@ -165,7 +172,8 @@ class WecomPlatformAdapter(Platform): return None msg_new = await asyncio.get_event_loop().run_in_executor( - None, get_latest_msg_item + None, + get_latest_msg_item, ) if msg_new: await self.convert_wechat_kf_message(msg_new) @@ -176,7 +184,9 @@ class WecomPlatformAdapter(Platform): @override async def send_by_session( - self, session: MessageSesion, message_chain: MessageChain + self, + session: MessageSesion, + message_chain: MessageChain, ): await super().send_by_session(session, message_chain) @@ -186,6 +196,7 @@ class WecomPlatformAdapter(Platform): "wecom", "wecom 适配器", id=self.config.get("id", "wecom"), + support_streaming_message=False, ) @override @@ -195,10 +206,11 @@ class WecomPlatformAdapter(Platform): try: acc_list = ( await loop.run_in_executor( - None, self.wechat_kf_api.get_account_list + None, + self.wechat_kf_api.get_account_list, ) ).get("account_list", []) - logger.debug(f"获取到微信客服列表: {str(acc_list)}") + logger.debug(f"获取到微信客服列表: {acc_list!s}") for acc in acc_list: name = acc.get("name", None) if name != self.kf_name: @@ -206,7 +218,7 @@ class WecomPlatformAdapter(Platform): open_kfid = acc.get("open_kfid", None) if not open_kfid: logger.error("获取微信客服失败，open_kfid 为空。") - logger.debug(f"Found open_kfid: {str(open_kfid)}") + logger.debug(f"Found open_kfid: {open_kfid!s}") kf_url = ( await loop.run_in_executor( None, @@ -216,7 +228,7 @@ class WecomPlatformAdapter(Platform): ) ).get("url", "") logger.info( - f"请打开以下链接，在微信扫码以获取客服微信: https://api.cl2wm.cn/api/qrcode/code?text={kf_url}" + f"请打开以下链接，在微信扫码以获取客服微信: https://api.cl2wm.cn/api/qrcode/code?text={kf_url}", ) except Exception as e: logger.error(e) @@ -256,7 +268,9 @@ class WecomPlatformAdapter(Platform): assert isinstance(msg, VoiceMessage) resp: Response = await asyncio.get_event_loop().run_in_executor( - None, self.client.media.download, msg.media_id + None, + self.client.media.download, + msg.media_id, ) temp_dir = os.path.join(get_astrbot_data_path(), "temp") path = os.path.join(temp_dir, f"wecom_{msg.media_id}.amr") @@ -294,8 +308,8 @@ class WecomPlatformAdapter(Platform): await self.handle_msg(abm) async def convert_wechat_kf_message(self, msg: dict) -> AstrBotMessage | None: - msgtype = msg.get("msgtype", None) - external_userid = msg.get("external_userid", None) + msgtype = msg.get("msgtype") + external_userid = msg.get("external_userid") abm = AstrBotMessage() abm.raw_message = msg abm.raw_message["_wechat_kf_flag"] = None # 方便处理 @@ -312,7 +326,9 @@ class WecomPlatformAdapter(Platform): elif msgtype == "image": media_id = msg.get("image", {}).get("media_id", "") resp: Response = await asyncio.get_event_loop().run_in_executor( - None, self.client.media.download, media_id + None, + self.client.media.download, + media_id, ) path = f"data/temp/wechat_kf_{media_id}.jpg" with open(path, "wb") as f: @@ -321,7 +337,9 @@ class WecomPlatformAdapter(Platform): elif msgtype == "voice": media_id = msg.get("voice", {}).get("media_id", "") resp: Response = await asyncio.get_event_loop().run_in_executor( - None, self.client.media.download, media_id + None, + self.client.media.download, + media_id, ) temp_dir = os.path.join(get_astrbot_data_path(), "temp") diff --git a/astrbot/core/platform/sources/wecom/wecom_event.py b/astrbot/core/platform/sources/wecom/wecom_event.py index e8078a9a..ba9ad9a4 100644 --- a/astrbot/core/platform/sources/wecom/wecom_event.py +++ b/astrbot/core/platform/sources/wecom/wecom_event.py @@ -1,22 +1,23 @@ +import asyncio import os import uuid -import asyncio -from astrbot.api.event import AstrMessageEvent, MessageChain -from astrbot.api.platform import AstrBotMessage, PlatformMetadata -from astrbot.api.message_components import Plain, Image, Record + from wechatpy.enterprise import WeChatClient -from .wecom_kf_message import WeChatKFMessage from astrbot.api import logger +from astrbot.api.event import AstrMessageEvent, MessageChain +from astrbot.api.message_components import Image, Plain, Record +from astrbot.api.platform import AstrBotMessage, PlatformMetadata from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from .wecom_kf_message import WeChatKFMessage + try: import pydub except Exception: logger.warning( - "检测到 pydub 库未安装，企业微信将无法语音收发。如需使用语音，请前往管理面板 -> 控制台 -> 安装 Pip 库安装 pydub。" + "检测到 pydub 库未安装，企业微信将无法语音收发。如需使用语音，请前往管理面板 -> 控制台 -> 安装 Pip 库安装 pydub。", ) - pass class WecomPlatformEvent(AstrMessageEvent): @@ -33,7 +34,9 @@ class WecomPlatformEvent(AstrMessageEvent): @staticmethod async def send_with_client( - client: WeChatClient, message: MessageChain, user_name: str + client: WeChatClient, + message: MessageChain, + user_name: str, ): pass @@ -44,44 +47,44 @@ class WecomPlatformEvent(AstrMessageEvent): plain (str): 要分割的长文本 Returns: list[str]: 分割后的文本列表 + """ if len(plain) <= 2048: return [plain] - else: - result = [] - start = 0 - while start < len(plain): - # 剩下的字符串长度<2048时结束 - if start + 2048 >= len(plain): - result.append(plain[start:]) + result = [] + start = 0 + while start < len(plain): + # 剩下的字符串长度<2048时结束 + if start + 2048 >= len(plain): + result.append(plain[start:]) + break + + # 向前搜索分割标点符号 + end = min(start + 2048, len(plain)) + cut_position = end + for i in range(end, start, -1): + if i < len(plain) and plain[i - 1] in [ + "。", + "！", + "？", + ".", + "!", + "?", + "\n", + ";", + "；", + ]: + cut_position = i break - # 向前搜索分割标点符号 - end = min(start + 2048, len(plain)) + # 没找到合适的位置分割, 直接切分 + if cut_position == end and end < len(plain): cut_position = end - for i in range(end, start, -1): - if i < len(plain) and plain[i - 1] in [ - "。", - "！", - "？", - ".", - "!", - "?", - "\n", - ";", - "；", - ]: - cut_position = i - break - # 没找到合适的位置分割, 直接切分 - if cut_position == end and end < len(plain): - cut_position = end + result.append(plain[start:cut_position]) + start = cut_position - result.append(plain[start:cut_position]) - start = cut_position - - return result + return result async def send(self, message: MessageChain): message_obj = self.message_obj @@ -111,7 +114,7 @@ class WecomPlatformEvent(AstrMessageEvent): except Exception as e: logger.error(f"微信客服上传图片失败: {e}") await self.send( - MessageChain().message(f"微信客服上传图片失败: {e}") + MessageChain().message(f"微信客服上传图片失败: {e}"), ) return logger.debug(f"微信客服上传图片返回: {response}") @@ -126,7 +129,8 @@ class WecomPlatformEvent(AstrMessageEvent): temp_dir = os.path.join(get_astrbot_data_path(), "temp") record_path_amr = os.path.join(temp_dir, f"{uuid.uuid4()}.amr") pydub.AudioSegment.from_wav(record_path).export( - record_path_amr, format="amr" + record_path_amr, + format="amr", ) with open(record_path_amr, "rb") as f: @@ -135,7 +139,7 @@ class WecomPlatformEvent(AstrMessageEvent): except Exception as e: logger.error(f"微信客服上传语音失败: {e}") await self.send( - MessageChain().message(f"微信客服上传语音失败: {e}") + MessageChain().message(f"微信客服上传语音失败: {e}"), ) return logger.info(f"微信客服上传语音返回: {response}") @@ -154,7 +158,9 @@ class WecomPlatformEvent(AstrMessageEvent): plain_chunks = await self.split_plain(comp.text) for chunk in plain_chunks: self.client.message.send_text( - message_obj.self_id, message_obj.session_id, chunk + message_obj.self_id, + message_obj.session_id, + chunk, ) await asyncio.sleep(0.5) # Avoid sending too fast elif isinstance(comp, Image): @@ -166,7 +172,7 @@ class WecomPlatformEvent(AstrMessageEvent): except Exception as e: logger.error(f"企业微信上传图片失败: {e}") await self.send( - MessageChain().message(f"企业微信上传图片失败: {e}") + MessageChain().message(f"企业微信上传图片失败: {e}"), ) return logger.debug(f"企业微信上传图片返回: {response}") @@ -181,7 +187,8 @@ class WecomPlatformEvent(AstrMessageEvent): temp_dir = os.path.join(get_astrbot_data_path(), "temp") record_path_amr = os.path.join(temp_dir, f"{uuid.uuid4()}.amr") pydub.AudioSegment.from_wav(record_path).export( - record_path_amr, format="amr" + record_path_amr, + format="amr", ) with open(record_path_amr, "rb") as f: @@ -190,7 +197,7 @@ class WecomPlatformEvent(AstrMessageEvent): except Exception as e: logger.error(f"企业微信上传语音失败: {e}") await self.send( - MessageChain().message(f"企业微信上传语音失败: {e}") + MessageChain().message(f"企业微信上传语音失败: {e}"), ) return logger.info(f"企业微信上传语音返回: {response}") @@ -212,7 +219,7 @@ class WecomPlatformEvent(AstrMessageEvent): else: buffer.chain.extend(chain.chain) if not buffer: - return + return None buffer.squash_plain() await self.send(buffer) return await super().send_streaming(generator, use_fallback) diff --git a/astrbot/core/platform/sources/wecom/wecom_kf.py b/astrbot/core/platform/sources/wecom/wecom_kf.py index 11866797..51f4ee14 100644 --- a/astrbot/core/platform/sources/wecom/wecom_kf.py +++ b/astrbot/core/platform/sources/wecom/wecom_kf.py @@ -1,7 +1,4 @@ -# -*- coding: utf-8 -*- - -""" -The MIT License (MIT) +"""The MIT License (MIT) Copyright (c) 2014-2020 messense @@ -28,15 +25,13 @@ from wechatpy.client.api.base import BaseWeChatAPI class WeChatKF(BaseWeChatAPI): - """ - 微信客服接口 + """微信客服接口 https://work.weixin.qq.com/api/doc/90000/90135/94670 """ def sync_msg(self, token, open_kfid, cursor="", limit=1000): - """ - 微信客户发送的消息、接待人员在企业微信回复的消息、发送消息接口发送失败事件（如被用户拒收） + """微信客户发送的消息、接待人员在企业微信回复的消息、发送消息接口发送失败事件（如被用户拒收） 、客户点击菜单消息的回复消息，可以通过该接口获取具体的消息内容和事件。不支持读取通过发送消息接口发送的消息。 支持的消息类型：文本、图片、语音、视频、文件、位置、链接、名片、小程序、事件。 @@ -57,8 +52,7 @@ class WeChatKF(BaseWeChatAPI): return self._post("kf/sync_msg", data=data) def get_service_state(self, open_kfid, external_userid): - """ - 获取会话状态 + """获取会话状态 ID 状态 说明 0 未处理 新会话接入。可选择：1.直接用API自动回复消息。2.放进待接入池等待接待人员接待。3.指定接待人员进行接待 @@ -78,10 +72,13 @@ class WeChatKF(BaseWeChatAPI): return self._post("kf/service_state/get", data=data) def trans_service_state( - self, open_kfid, external_userid, service_state, servicer_userid="" + self, + open_kfid, + external_userid, + service_state, + servicer_userid="", ): - """ - 变更会话状态 + """变更会话状态 :param open_kfid: 客服帐号ID :param external_userid: 微信客户的external_userid @@ -98,8 +95,7 @@ class WeChatKF(BaseWeChatAPI): return self._post("kf/service_state/trans", data=data) def get_servicer_list(self, open_kfid): - """ - 获取接待人员列表 + """获取接待人员列表 :param open_kfid: 客服帐号ID :return: 接口调用结果 @@ -110,8 +106,7 @@ class WeChatKF(BaseWeChatAPI): return self._get("kf/servicer/list", params=data) def add_servicer(self, open_kfid, userid_list): - """ - 添加接待人员 + """添加接待人员 添加指定客服帐号的接待人员。 :param open_kfid: 客服帐号ID @@ -128,8 +123,7 @@ class WeChatKF(BaseWeChatAPI): return self._post("kf/servicer/add", data=data) def del_servicer(self, open_kfid, userid_list): - """ - 删除接待人员 + """删除接待人员 从客服帐号删除接待人员 :param open_kfid: 客服帐号ID @@ -146,8 +140,7 @@ class WeChatKF(BaseWeChatAPI): return self._post("kf/servicer/del", data=data) def batchget_customer(self, external_userid_list): - """ - 客户基本信息获取 + """客户基本信息获取 :param external_userid_list: external_userid列表 :return: 接口调用结果 @@ -161,16 +154,14 @@ class WeChatKF(BaseWeChatAPI): return self._post("kf/customer/batchget", data=data) def get_account_list(self): - """ - 获取客服帐号列表 + """获取客服帐号列表 :return: 接口调用结果 """ return self._get("kf/account/list") def add_contact_way(self, open_kfid, scene): - """ - 获取客服帐号链接 + """获取客服帐号链接 :param open_kfid: 客服帐号ID :param scene: 场景值，字符串类型，由开发者自定义。不多于32字节;字符串取值范围(正则表达式)：[0-9a-zA-Z_-]* @@ -180,18 +171,21 @@ class WeChatKF(BaseWeChatAPI): return self._post("kf/add_contact_way", data=data) def get_upgrade_service_config(self): - """ - 获取配置的专员与客户群 + """获取配置的专员与客户群 :return: 接口调用结果 """ return self._get("kf/customer/get_upgrade_service_config") def upgrade_service( - self, open_kfid, external_userid, service_type, member=None, groupchat=None + self, + open_kfid, + external_userid, + service_type, + member=None, + groupchat=None, ): - """ - 为客户升级为专员或客户群服务 + """为客户升级为专员或客户群服务 :param open_kfid: 客服帐号ID :param external_userid: 微信客户的external_userid @@ -200,7 +194,6 @@ class WeChatKF(BaseWeChatAPI): :param groupchat: 推荐的客户群，type等于2时有效 :return: 接口调用结果 """ - data = { "open_kfid": open_kfid, "external_userid": external_userid, @@ -213,20 +206,17 @@ class WeChatKF(BaseWeChatAPI): return self._post("kf/customer/upgrade_service", data=data) def cancel_upgrade_service(self, open_kfid, external_userid): - """ - 为客户取消推荐 + """为客户取消推荐 :param open_kfid: 客服帐号ID :param external_userid: 微信客户的external_userid :return: 接口调用结果 """ - data = {"open_kfid": open_kfid, "external_userid": external_userid} return self._post("kf/customer/cancel_upgrade_service", data=data) def send_msg_on_event(self, code, msgtype, msg_content, msgid=None): - """ - 当特定的事件回调消息包含code字段，可以此code为凭证，调用该接口给用户发送相应事件场景下的消息，如客服欢迎语。 + """当特定的事件回调消息包含code字段，可以此code为凭证，调用该接口给用户发送相应事件场景下的消息，如客服欢迎语。 支持发送消息类型：文本、菜单消息。 :param code: 事件响应消息对应的code。通过事件回调下发，仅可使用一次。 @@ -236,7 +226,6 @@ class WeChatKF(BaseWeChatAPI): 字符串取值范围(正则表达式)：[0-9a-zA-Z_-]* :return: 接口调用结果 """ - data = {"code": code, "msgtype": msgtype} if msgid: data["msgid"] = msgid @@ -244,8 +233,7 @@ class WeChatKF(BaseWeChatAPI): return self._post("kf/send_msg_on_event", data=data) def get_corp_statistic(self, start_time, end_time, open_kfid=None): - """ - 获取「客户数据统计」企业汇总数据 + """获取「客户数据统计」企业汇总数据 :param start_time: 开始时间 :param end_time: 结束时间 @@ -256,10 +244,13 @@ class WeChatKF(BaseWeChatAPI): return self._post("kf/get_corp_statistic", data=data) def get_servicer_statistic( - self, start_time, end_time, open_kfid=None, servicer_userid=None + self, + start_time, + end_time, + open_kfid=None, + servicer_userid=None, ): - """ - 获取「客户数据统计」接待人员明细数据 + """获取「客户数据统计」接待人员明细数据 :param start_time: 开始时间 :param end_time: 结束时间 @@ -276,8 +267,7 @@ class WeChatKF(BaseWeChatAPI): return self._post("kf/get_servicer_statistic", data=data) def account_update(self, open_kfid, name, media_id): - """ - 修改客服账号 + """修改客服账号 :param open_kfid: 客服帐号ID :param name: 客服名称 diff --git a/astrbot/core/platform/sources/wecom/wecom_kf_message.py b/astrbot/core/platform/sources/wecom/wecom_kf_message.py index 42fc20d6..d839134a 100644 --- a/astrbot/core/platform/sources/wecom/wecom_kf_message.py +++ b/astrbot/core/platform/sources/wecom/wecom_kf_message.py @@ -1,5 +1,4 @@ -""" -The MIT License (MIT) +"""The MIT License (MIT) Copyright (c) 2014-2020 messense @@ -23,13 +22,11 @@ SOFTWARE. """ from optionaldict import optionaldict - from wechatpy.client.api.base import BaseWeChatAPI class WeChatKFMessage(BaseWeChatAPI): - """ - 发送微信客服消息 + """发送微信客服消息 https://work.weixin.qq.com/api/doc/90000/90135/94677 @@ -46,8 +43,7 @@ class WeChatKFMessage(BaseWeChatAPI): """ def send(self, user_id, open_kfid, msgid="", msg=None): - """ - 当微信客户处于“新接入待处理”或“由智能助手接待”状态下，可调用该接口给用户发送消息。 + """当微信客户处于“新接入待处理”或“由智能助手接待”状态下，可调用该接口给用户发送消息。 注意仅当微信客户在主动发送消息给客服后的48小时内，企业可发送消息给客户，最多可发送5条消息；若用户继续发送消息，企业可再次下发消息。 支持发送消息类型：文本、图片、语音、视频、文件、图文、小程序、菜单消息、地理位置。 @@ -127,7 +123,13 @@ class WeChatKFMessage(BaseWeChatAPI): ) def send_msgmenu( - self, user_id, open_kfid, head_content, menu_list, tail_content, msgid="" + self, + user_id, + open_kfid, + head_content, + menu_list, + tail_content, + msgid="", ): return self.send( user_id, @@ -144,7 +146,14 @@ class WeChatKFMessage(BaseWeChatAPI): ) def send_location( - self, user_id, open_kfid, name, address, latitude, longitude, msgid="" + self, + user_id, + open_kfid, + name, + address, + latitude, + longitude, + msgid="", ): return self.send( user_id, @@ -162,7 +171,14 @@ class WeChatKFMessage(BaseWeChatAPI): ) def send_miniprogram( - self, user_id, open_kfid, appid, title, thumb_media_id, pagepath, msgid="" + self, + user_id, + open_kfid, + appid, + title, + thumb_media_id, + pagepath, + msgid="", ): return self.send( user_id, diff --git a/astrbot/core/platform/sources/wecom_ai_bot/WXBizJsonMsgCrypt.py b/astrbot/core/platform/sources/wecom_ai_bot/WXBizJsonMsgCrypt.py index 5332942b..2df09a76 100644 --- a/astrbot/core/platform/sources/wecom_ai_bot/WXBizJsonMsgCrypt.py +++ b/astrbot/core/platform/sources/wecom_ai_bot/WXBizJsonMsgCrypt.py @@ -1,5 +1,4 @@ #!/usr/bin/env python -# -*- encoding:utf-8 -*- """对企业微信发送给企业后台的消息加解密示例代码. @copyright: Copyright (c) 1998-2020 Tencent Inc. @@ -7,15 +6,16 @@ """ # ------------------------------------------------------------------------ -import logging import base64 -import random import hashlib -import time -import struct -from Crypto.Cipher import AES -import socket import json +import logging +import secrets +import socket +import struct +import time + +from Crypto.Cipher import AES from . import ierror @@ -31,7 +31,7 @@ class FormatException(Exception): def throw_exception(message, exception_class=FormatException): - """my define raise exception function""" + """My define raise exception function""" raise exception_class(message) @@ -136,9 +136,15 @@ class PKCS7Encoder: return decrypted[:-pad] -class Prpcrypt(object): +class Prpcrypt: """提供接收和推送给企业微信消息的加解密接口""" + # 16位随机字符串的范围常量 + # randbelow(RANDOM_RANGE) 返回 [0, 8999999999999999]（两端都包含，即包含0和8999999999999999） + # 加上 MIN_RANDOM_VALUE 后得到 [1000000000000000, 9999999999999999]（两端都包含）即16位数字 + MIN_RANDOM_VALUE = 1000000000000000 # 最小值: 1000000000000000 (16位) + RANDOM_RANGE = 9000000000000000 # 范围大小: 确保最大值为 9999999999999999 (16位) + def __init__(self, key): # self.key = base64.b64decode(key+"=") self.key = key @@ -207,10 +213,12 @@ class Prpcrypt(object): """随机生成16位字符串 @return: 16位字符串 """ - return str(random.randint(1000000000000000, 9999999999999999)).encode() + return str( + secrets.randbelow(self.RANDOM_RANGE) + self.MIN_RANDOM_VALUE + ).encode() -class WXBizJsonMsgCrypt(object): +class WXBizJsonMsgCrypt: # 构造函数 def __init__(self, sToken, sEncodingAESKey, sReceiveId): try: diff --git a/astrbot/core/platform/sources/wecom_ai_bot/__init__.py b/astrbot/core/platform/sources/wecom_ai_bot/__init__.py index 7da90003..2f87b88b 100644 --- a/astrbot/core/platform/sources/wecom_ai_bot/__init__.py +++ b/astrbot/core/platform/sources/wecom_ai_bot/__init__.py @@ -1,6 +1,4 @@ -""" -企业微信智能机器人平台适配器包 -""" +"""企业微信智能机器人平台适配器包""" from .wecomai_adapter import WecomAIBotAdapter from .wecomai_api import WecomAIBotAPIClient @@ -9,9 +7,9 @@ from .wecomai_server import WecomAIBotServer from .wecomai_utils import WecomAIBotConstants __all__ = [ - "WecomAIBotAdapter", "WecomAIBotAPIClient", + "WecomAIBotAdapter", + "WecomAIBotConstants", "WecomAIBotMessageEvent", "WecomAIBotServer", - "WecomAIBotConstants", ] diff --git a/astrbot/core/platform/sources/wecom_ai_bot/ierror.py b/astrbot/core/platform/sources/wecom_ai_bot/ierror.py index cc1bf221..0df14a50 100644 --- a/astrbot/core/platform/sources/wecom_ai_bot/ierror.py +++ b/astrbot/core/platform/sources/wecom_ai_bot/ierror.py @@ -1,5 +1,4 @@ #!/usr/bin/env python -# -*- coding: utf-8 -*- ######################################################################### # Author: jonyqin # Created Time: Thu 11 Sep 2014 01:53:58 PM CST diff --git a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_adapter.py b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_adapter.py index 830d8de5..9c13cfef 100644 --- a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_adapter.py +++ b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_adapter.py @@ -1,38 +1,37 @@ -""" -企业微信智能机器人平台适配器 +"""企业微信智能机器人平台适配器 基于企业微信智能机器人 API 的消息平台适配器，支持 HTTP 回调 参考webchat_adapter.py的队列机制，实现异步消息处理和流式响应 """ -import time import asyncio -import uuid -import hashlib import base64 -from typing import Awaitable, Any, Dict, Optional, Callable - +import hashlib +import time +import uuid +from collections.abc import Awaitable, Callable +from typing import Any +from astrbot.api import logger +from astrbot.api.event import MessageChain +from astrbot.api.message_components import At, Image, Plain from astrbot.api.platform import ( - Platform, AstrBotMessage, MessageMember, MessageType, + Platform, PlatformMetadata, ) -from astrbot.api.event import MessageChain -from astrbot.api.message_components import Plain, At, Image -from astrbot.api import logger from astrbot.core.platform.astr_message_event import MessageSesion -from ...register import register_platform_adapter +from ...register import register_platform_adapter from .wecomai_api import ( WecomAIBotAPIClient, WecomAIBotMessageParser, WecomAIBotStreamMessageBuilder, ) from .wecomai_event import WecomAIBotMessageEvent +from .wecomai_queue_mgr import WecomAIQueueMgr from .wecomai_server import WecomAIBotServer -from .wecomai_queue_mgr import wecomai_queue_mgr, WecomAIQueueMgr from .wecomai_utils import ( WecomAIBotConstants, format_session_id, @@ -45,7 +44,9 @@ class WecomAIQueueListener: """企业微信智能机器人队列监听器，参考webchat的QueueListener设计""" def __init__( - self, queue_mgr: WecomAIQueueMgr, callback: Callable[[dict], Awaitable[None]] + self, + queue_mgr: WecomAIQueueMgr, + callback: Callable[[dict], Awaitable[None]], ) -> None: self.queue_mgr = queue_mgr self.callback = callback @@ -90,13 +91,17 @@ class WecomAIQueueListener: @register_platform_adapter( - "wecom_ai_bot", "企业微信智能机器人适配器，支持 HTTP 回调接收消息" + "wecom_ai_bot", + "企业微信智能机器人适配器，支持 HTTP 回调接收消息", ) class WecomAIBotAdapter(Platform): """企业微信智能机器人适配器""" def __init__( - self, platform_config: dict, platform_settings: dict, event_queue: asyncio.Queue + self, + platform_config: dict, + platform_settings: dict, + event_queue: asyncio.Queue, ) -> None: super().__init__(event_queue) @@ -110,10 +115,12 @@ class WecomAIBotAdapter(Platform): self.host = self.config.get("callback_server_host", "0.0.0.0") self.bot_name = self.config.get("wecom_ai_bot_name", "") self.initial_respond_text = self.config.get( - "wecomaibot_init_respond_text", "💭 思考中..." + "wecomaibot_init_respond_text", + "💭 思考中...", ) self.friend_message_welcome_text = self.config.get( - "wecomaibot_friend_message_welcome_text", "" + "wecomaibot_friend_message_welcome_text", + "", ) # 平台元数据 @@ -137,9 +144,13 @@ class WecomAIBotAdapter(Platform): # 事件循环和关闭信号 self.shutdown_event = asyncio.Event() + # 队列管理器 + self.queue_mgr = WecomAIQueueMgr() + # 队列监听器 self.queue_listener = WecomAIQueueListener( - wecomai_queue_mgr, self._handle_queued_message + self.queue_mgr, + self._handle_queued_message, ) async def _handle_queued_message(self, data: dict): @@ -151,8 +162,10 @@ class WecomAIBotAdapter(Platform): logger.error(f"处理队列消息时发生异常: {e}") async def _process_message( - self, message_data: Dict[str, Any], callback_params: Dict[str, str] - ) -> Optional[str]: + self, + message_data: dict[str, Any], + callback_params: dict[str, str], + ) -> str | None: """处理接收到的消息 Args: @@ -161,6 +174,7 @@ class WecomAIBotAdapter(Platform): Returns: 加密后的响应消息，无需响应时返回 None + """ msgtype = message_data.get("msgtype") if not msgtype: @@ -173,15 +187,22 @@ class WecomAIBotAdapter(Platform): # create a brand-new unique stream_id for this message session stream_id = f"{session_id}_{generate_random_string(10)}" await self._enqueue_message( - message_data, callback_params, stream_id, session_id + message_data, + callback_params, + stream_id, + session_id, ) - wecomai_queue_mgr.set_pending_response(stream_id, callback_params) + self.queue_mgr.set_pending_response(stream_id, callback_params) resp = WecomAIBotStreamMessageBuilder.make_text_stream( - stream_id, self.initial_respond_text, False + stream_id, + self.initial_respond_text, + False, ) return await self.api_client.encrypt_message( - resp, callback_params["nonce"], callback_params["timestamp"] + resp, + callback_params["nonce"], + callback_params["timestamp"], ) except Exception as e: logger.error("处理消息时发生异常: %s", e) @@ -189,12 +210,14 @@ class WecomAIBotAdapter(Platform): elif msgtype == "stream": # wechat server is requesting for updates of a stream stream_id = message_data["stream"]["id"] - if not wecomai_queue_mgr.has_back_queue(stream_id): + if not self.queue_mgr.has_back_queue(stream_id): logger.error(f"Cannot find back queue for stream_id: {stream_id}") # 返回结束标志，告诉微信服务器流已结束 end_message = WecomAIBotStreamMessageBuilder.make_text_stream( - stream_id, "", True + stream_id, + "", + True, ) resp = await self.api_client.encrypt_message( end_message, @@ -202,10 +225,10 @@ class WecomAIBotAdapter(Platform): callback_params["timestamp"], ) return resp - queue = wecomai_queue_mgr.get_or_create_back_queue(stream_id) + queue = self.queue_mgr.get_or_create_back_queue(stream_id) if queue.empty(): logger.debug( - f"No new messages in back queue for stream_id: {stream_id}" + f"No new messages in back queue for stream_id: {stream_id}", ) return None @@ -222,12 +245,11 @@ class WecomAIBotAdapter(Platform): elif msg["type"] == "end": # stream end finish = True - wecomai_queue_mgr.remove_queues(stream_id) + self.queue_mgr.remove_queues(stream_id) break - else: - pass + logger.debug( - f"Aggregated content: {latest_plain_content}, image: {len(image_base64)}, finish: {finish}" + f"Aggregated content: {latest_plain_content}, image: {len(image_base64)}, finish: {finish}", ) if latest_plain_content or image_base64: msg_items = [] @@ -240,12 +262,15 @@ class WecomAIBotAdapter(Platform): { "msgtype": WecomAIBotConstants.MSG_TYPE_IMAGE, "image": {"base64": img_b64, "md5": img_md5}, - } + }, ) image_base64 = [] plain_message = WecomAIBotStreamMessageBuilder.make_mixed_stream( - stream_id, latest_plain_content, msg_items, finish + stream_id, + latest_plain_content, + msg_items, + finish, ) encrypted_message = await self.api_client.encrypt_message( plain_message, @@ -254,7 +279,7 @@ class WecomAIBotAdapter(Platform): ) if encrypted_message: logger.debug( - f"Stream message sent successfully, stream_id: {stream_id}" + f"Stream message sent successfully, stream_id: {stream_id}", ) else: logger.error("消息加密失败") @@ -266,7 +291,7 @@ class WecomAIBotAdapter(Platform): # 用户进入会话，发送欢迎消息 try: resp = WecomAIBotStreamMessageBuilder.make_text( - self.friend_message_welcome_text + self.friend_message_welcome_text, ) return await self.api_client.encrypt_message( resp, @@ -276,23 +301,22 @@ class WecomAIBotAdapter(Platform): except Exception as e: logger.error("处理欢迎消息时发生异常: %s", e) return None - pass - def _extract_session_id(self, message_data: Dict[str, Any]) -> str: + def _extract_session_id(self, message_data: dict[str, Any]) -> str: """从消息数据中提取会话ID""" user_id = message_data.get("from", {}).get("userid", "default_user") return format_session_id("wecomai", user_id) async def _enqueue_message( self, - message_data: Dict[str, Any], - callback_params: Dict[str, str], + message_data: dict[str, Any], + callback_params: dict[str, str], stream_id: str, session_id: str, ): """将消息放入队列进行异步处理""" - input_queue = wecomai_queue_mgr.get_or_create_queue(stream_id) - _ = wecomai_queue_mgr.get_or_create_back_queue(stream_id) + input_queue = self.queue_mgr.get_or_create_queue(stream_id) + _ = self.queue_mgr.get_or_create_back_queue(stream_id) message_payload = { "message_data": message_data, "callback_params": callback_params, @@ -320,7 +344,7 @@ class WecomAIBotAdapter(Platform): content = WecomAIBotMessageParser.parse_text_message(message_data) elif msgtype == WecomAIBotConstants.MSG_TYPE_IMAGE: _img_url_to_process.append( - WecomAIBotMessageParser.parse_image_message(message_data) + WecomAIBotMessageParser.parse_image_message(message_data), ) elif msgtype == WecomAIBotConstants.MSG_TYPE_MIXED: # 提取混合消息中的文本内容 @@ -390,7 +414,9 @@ class WecomAIBotAdapter(Platform): return abm async def send_by_session( - self, session: MessageSesion, message_chain: MessageChain + self, + session: MessageSesion, + message_chain: MessageChain, ): """通过会话发送消息""" # 企业微信智能机器人主要通过回调响应，这里记录日志 @@ -429,6 +455,7 @@ class WecomAIBotAdapter(Platform): platform_meta=self.meta(), session_id=message.session_id, api_client=self.api_client, + queue_mgr=self.queue_mgr, ) self.commit_event(message_event) diff --git a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_api.py b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_api.py index 540bf06b..6c448a97 100644 --- a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_api.py +++ b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_api.py @@ -1,19 +1,20 @@ -""" -企业微信智能机器人 API 客户端 +"""企业微信智能机器人 API 客户端 处理消息加密解密、API 调用等 """ -import json import base64 import hashlib -from typing import Dict, Any, Optional, Tuple, Union -from Crypto.Cipher import AES -import aiohttp +import json +from typing import Any + +import aiohttp +from Crypto.Cipher import AES -from .WXBizJsonMsgCrypt import WXBizJsonMsgCrypt -from .wecomai_utils import WecomAIBotConstants from astrbot import logger +from .wecomai_utils import WecomAIBotConstants +from .WXBizJsonMsgCrypt import WXBizJsonMsgCrypt + class WecomAIBotAPIClient: """企业微信智能机器人 API 客户端""" @@ -24,14 +25,19 @@ class WecomAIBotAPIClient: Args: token: 企业微信机器人 Token encoding_aes_key: 消息加密密钥 + """ self.token = token self.encoding_aes_key = encoding_aes_key self.wxcpt = WXBizJsonMsgCrypt(token, encoding_aes_key, "") # receiveid 为空串 async def decrypt_message( - self, encrypted_data: bytes, msg_signature: str, timestamp: str, nonce: str - ) -> Tuple[int, Optional[Dict[str, Any]]]: + self, + encrypted_data: bytes, + msg_signature: str, + timestamp: str, + nonce: str, + ) -> tuple[int, dict[str, Any] | None]: """解密企业微信消息 Args: @@ -42,10 +48,14 @@ class WecomAIBotAPIClient: Returns: (错误码, 解密后的消息数据字典) + """ try: ret, decrypted_msg = self.wxcpt.DecryptMsg( - encrypted_data, msg_signature, timestamp, nonce + encrypted_data, + msg_signature, + timestamp, + nonce, ) if ret != WecomAIBotConstants.SUCCESS: @@ -70,8 +80,11 @@ class WecomAIBotAPIClient: return WecomAIBotConstants.DECRYPT_ERROR, None async def encrypt_message( - self, plain_message: str, nonce: str, timestamp: str - ) -> Optional[str]: + self, + plain_message: str, + nonce: str, + timestamp: str, + ) -> str | None: """加密消息 Args: @@ -81,6 +94,7 @@ class WecomAIBotAPIClient: Returns: 加密后的消息，失败时返回 None + """ try: ret, encrypted_msg = self.wxcpt.EncryptMsg(plain_message, nonce, timestamp) @@ -97,7 +111,11 @@ class WecomAIBotAPIClient: return None def verify_url( - self, msg_signature: str, timestamp: str, nonce: str, echostr: str + self, + msg_signature: str, + timestamp: str, + nonce: str, + echostr: str, ) -> str: """验证回调 URL @@ -109,10 +127,14 @@ class WecomAIBotAPIClient: Returns: 验证结果字符串 + """ try: ret, echo_result = self.wxcpt.VerifyURL( - msg_signature, timestamp, nonce, echostr + msg_signature, + timestamp, + nonce, + echostr, ) if ret != WecomAIBotConstants.SUCCESS: @@ -127,8 +149,10 @@ class WecomAIBotAPIClient: return "verify fail" async def process_encrypted_image( - self, image_url: str, aes_key_base64: Optional[str] = None - ) -> Tuple[bool, Union[bytes, str]]: + self, + image_url: str, + aes_key_base64: str | None = None, + ) -> tuple[bool, bytes | str]: """下载并解密加密图片 Args: @@ -137,6 +161,7 @@ class WecomAIBotAPIClient: Returns: (是否成功, 图片数据或错误信息) + """ try: # 下载图片 @@ -161,7 +186,7 @@ class WecomAIBotAPIClient: # Base64 解码密钥 aes_key = base64.b64decode( - aes_key_base64 + "=" * (-len(aes_key_base64) % 4) + aes_key_base64 + "=" * (-len(aes_key_base64) % 4), ) if len(aes_key) != 32: raise ValueError("无效的 AES 密钥长度: 应为 32 字节") @@ -183,17 +208,17 @@ class WecomAIBotAPIClient: return True, decrypted_data except aiohttp.ClientError as e: - error_msg = f"图片下载失败: {str(e)}" + error_msg = f"图片下载失败: {e!s}" logger.error(error_msg) return False, error_msg except ValueError as e: - error_msg = f"参数错误: {str(e)}" + error_msg = f"参数错误: {e!s}" logger.error(error_msg) return False, error_msg except Exception as e: - error_msg = f"图片处理异常: {str(e)}" + error_msg = f"图片处理异常: {e!s}" logger.error(error_msg) return False, error_msg @@ -212,6 +237,7 @@ class WecomAIBotStreamMessageBuilder: Returns: JSON 格式的流消息字符串 + """ plain = { "msgtype": WecomAIBotConstants.MSG_TYPE_STREAM, @@ -221,7 +247,9 @@ class WecomAIBotStreamMessageBuilder: @staticmethod def make_image_stream( - stream_id: str, image_data: bytes, finish: bool = False + stream_id: str, + image_data: bytes, + finish: bool = False, ) -> str: """构建图片流消息 @@ -232,6 +260,7 @@ class WecomAIBotStreamMessageBuilder: Returns: JSON 格式的流消息字符串 + """ image_md5 = hashlib.md5(image_data).hexdigest() image_base64 = base64.b64encode(image_data).decode("utf-8") @@ -245,7 +274,7 @@ class WecomAIBotStreamMessageBuilder: { "msgtype": WecomAIBotConstants.MSG_TYPE_IMAGE, "image": {"base64": image_base64, "md5": image_md5}, - } + }, ], }, } @@ -253,7 +282,10 @@ class WecomAIBotStreamMessageBuilder: @staticmethod def make_mixed_stream( - stream_id: str, content: str, msg_items: list, finish: bool = False + stream_id: str, + content: str, + msg_items: list, + finish: bool = False, ) -> str: """构建混合类型流消息 @@ -265,6 +297,7 @@ class WecomAIBotStreamMessageBuilder: Returns: JSON 格式的流消息字符串 + """ plain = { "msgtype": WecomAIBotConstants.MSG_TYPE_STREAM, @@ -283,6 +316,7 @@ class WecomAIBotStreamMessageBuilder: Returns: JSON 格式的文本消息字符串 + """ plain = {"msgtype": "text", "text": {"content": content}} return json.dumps(plain, ensure_ascii=False) @@ -292,7 +326,7 @@ class WecomAIBotMessageParser: """企业微信智能机器人消息解析器""" @staticmethod - def parse_text_message(data: Dict[str, Any]) -> Optional[str]: + def parse_text_message(data: dict[str, Any]) -> str | None: """解析文本消息 Args: @@ -300,6 +334,7 @@ class WecomAIBotMessageParser: Returns: 文本内容，解析失败返回 None + """ try: return data.get("text", {}).get("content") @@ -308,7 +343,7 @@ class WecomAIBotMessageParser: return None @staticmethod - def parse_image_message(data: Dict[str, Any]) -> Optional[str]: + def parse_image_message(data: dict[str, Any]) -> str | None: """解析图片消息 Args: @@ -316,6 +351,7 @@ class WecomAIBotMessageParser: Returns: 图片 URL，解析失败返回 None + """ try: return data.get("image", {}).get("url") @@ -324,7 +360,7 @@ class WecomAIBotMessageParser: return None @staticmethod - def parse_stream_message(data: Dict[str, Any]) -> Optional[Dict[str, Any]]: + def parse_stream_message(data: dict[str, Any]) -> dict[str, Any] | None: """解析流消息 Args: @@ -332,6 +368,7 @@ class WecomAIBotMessageParser: Returns: 流消息数据，解析失败返回 None + """ try: stream_data = data.get("stream", {}) @@ -346,7 +383,7 @@ class WecomAIBotMessageParser: return None @staticmethod - def parse_mixed_message(data: Dict[str, Any]) -> Optional[list]: + def parse_mixed_message(data: dict[str, Any]) -> list | None: """解析混合消息 Args: @@ -354,6 +391,7 @@ class WecomAIBotMessageParser: Returns: 消息项列表，解析失败返回 None + """ try: return data.get("mixed", {}).get("msg_item", []) @@ -362,7 +400,7 @@ class WecomAIBotMessageParser: return None @staticmethod - def parse_event_message(data: Dict[str, Any]) -> Optional[Dict[str, Any]]: + def parse_event_message(data: dict[str, Any]) -> dict[str, Any] | None: """解析事件消息 Args: @@ -370,6 +408,7 @@ class WecomAIBotMessageParser: Returns: 事件数据，解析失败返回 None + """ try: return data.get("event", {}) diff --git a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_event.py b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_event.py index 2d7ec91c..0091783a 100644 --- a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_event.py +++ b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_event.py @@ -1,16 +1,14 @@ -""" -企业微信智能机器人事件处理模块，处理消息事件的发送和接收 -""" +"""企业微信智能机器人事件处理模块，处理消息事件的发送和接收""" +from astrbot.api import logger from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.message_components import ( Image, Plain, ) -from astrbot.api import logger from .wecomai_api import WecomAIBotAPIClient -from .wecomai_queue_mgr import wecomai_queue_mgr +from .wecomai_queue_mgr import WecomAIQueueMgr class WecomAIBotMessageEvent(AstrMessageEvent): @@ -23,6 +21,7 @@ class WecomAIBotMessageEvent(AstrMessageEvent): platform_meta, session_id: str, api_client: WecomAIBotAPIClient, + queue_mgr: WecomAIQueueMgr, ): """初始化消息事件 @@ -32,17 +31,20 @@ class WecomAIBotMessageEvent(AstrMessageEvent): platform_meta: 平台元数据 session_id: 会话 ID api_client: API 客户端 + """ super().__init__(message_str, message_obj, platform_meta, session_id) self.api_client = api_client + self.queue_mgr = queue_mgr @staticmethod async def _send( message_chain: MessageChain, stream_id: str, + queue_mgr: WecomAIQueueMgr, streaming: bool = False, ): - back_queue = wecomai_queue_mgr.get_or_create_back_queue(stream_id) + back_queue = queue_mgr.get_or_create_back_queue(stream_id) if not message_chain: await back_queue.put( @@ -50,7 +52,7 @@ class WecomAIBotMessageEvent(AstrMessageEvent): "type": "end", "data": "", "streaming": False, - } + }, ) return "" @@ -64,7 +66,7 @@ class WecomAIBotMessageEvent(AstrMessageEvent): "data": data, "streaming": streaming, "session_id": stream_id, - } + }, ) elif isinstance(comp, Image): # 处理图片消息 @@ -77,7 +79,7 @@ class WecomAIBotMessageEvent(AstrMessageEvent): "image_data": image_base64, "streaming": streaming, "session_id": stream_id, - } + }, ) else: logger.warning("图片数据为空，跳过") @@ -95,7 +97,7 @@ class WecomAIBotMessageEvent(AstrMessageEvent): "wecom_ai_bot platform event raw_message should be a dict" ) stream_id = raw.get("stream_id", self.session_id) - await WecomAIBotMessageEvent._send(message, stream_id) + await WecomAIBotMessageEvent._send(message, stream_id, self.queue_mgr) await super().send(message) async def send_streaming(self, generator, use_fallback=False): @@ -106,7 +108,7 @@ class WecomAIBotMessageEvent(AstrMessageEvent): "wecom_ai_bot platform event raw_message should be a dict" ) stream_id = raw.get("stream_id", self.session_id) - back_queue = wecomai_queue_mgr.get_or_create_back_queue(stream_id) + back_queue = self.queue_mgr.get_or_create_back_queue(stream_id) # 企业微信智能机器人不支持增量发送，因此我们需要在这里将增量内容累积起来，积累发送 increment_plain = "" @@ -127,7 +129,7 @@ class WecomAIBotMessageEvent(AstrMessageEvent): "data": final_data, "streaming": True, "session_id": self.session_id, - } + }, ) final_data = "" continue @@ -135,6 +137,7 @@ class WecomAIBotMessageEvent(AstrMessageEvent): final_data += await WecomAIBotMessageEvent._send( chain, stream_id=stream_id, + queue_mgr=self.queue_mgr, streaming=True, ) @@ -144,6 +147,6 @@ class WecomAIBotMessageEvent(AstrMessageEvent): "data": final_data, "streaming": True, "session_id": self.session_id, - } + }, ) await super().send_streaming(generator, use_fallback) diff --git a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_queue_mgr.py b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_queue_mgr.py index 1367301c..3a982bdf 100644 --- a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_queue_mgr.py +++ b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_queue_mgr.py @@ -1,11 +1,11 @@ -""" -企业微信智能机器人队列管理器 +"""企业微信智能机器人队列管理器 参考 webchat_queue_mgr.py，为企业微信智能机器人实现队列机制 支持异步消息处理和流式响应 """ import asyncio -from typing import Dict, Any, Optional +from typing import Any + from astrbot.api import logger @@ -13,13 +13,13 @@ class WecomAIQueueMgr: """企业微信智能机器人队列管理器""" def __init__(self) -> None: - self.queues: Dict[str, asyncio.Queue] = {} + self.queues: dict[str, asyncio.Queue] = {} """StreamID 到输入队列的映射 - 用于接收用户消息""" - self.back_queues: Dict[str, asyncio.Queue] = {} + self.back_queues: dict[str, asyncio.Queue] = {} """StreamID 到输出队列的映射 - 用于发送机器人响应""" - self.pending_responses: Dict[str, Dict[str, Any]] = {} + self.pending_responses: dict[str, dict[str, Any]] = {} """待处理的响应缓存，用于流式响应""" def get_or_create_queue(self, session_id: str) -> asyncio.Queue: @@ -30,6 +30,7 @@ class WecomAIQueueMgr: Returns: 输入队列实例 + """ if session_id not in self.queues: self.queues[session_id] = asyncio.Queue() @@ -44,6 +45,7 @@ class WecomAIQueueMgr: Returns: 输出队列实例 + """ if session_id not in self.back_queues: self.back_queues[session_id] = asyncio.Queue() @@ -55,6 +57,7 @@ class WecomAIQueueMgr: Args: session_id: 会话ID + """ if session_id in self.queues: del self.queues[session_id] @@ -76,6 +79,7 @@ class WecomAIQueueMgr: Returns: 是否存在队列 + """ return session_id in self.queues @@ -87,15 +91,17 @@ class WecomAIQueueMgr: Returns: 是否存在输出队列 + """ return session_id in self.back_queues - def set_pending_response(self, session_id: str, callback_params: Dict[str, str]): + def set_pending_response(self, session_id: str, callback_params: dict[str, str]): """设置待处理的响应参数 Args: session_id: 会话ID callback_params: 回调参数（nonce, timestamp等） + """ self.pending_responses[session_id] = { "callback_params": callback_params, @@ -103,7 +109,7 @@ class WecomAIQueueMgr: } logger.debug(f"[WecomAI] 设置待处理响应: {session_id}") - def get_pending_response(self, session_id: str) -> Optional[Dict[str, Any]]: + def get_pending_response(self, session_id: str) -> dict[str, Any] | None: """获取待处理的响应参数 Args: @@ -111,6 +117,7 @@ class WecomAIQueueMgr: Returns: 响应参数，如果不存在则返回None + """ return self.pending_responses.get(session_id) @@ -119,6 +126,7 @@ class WecomAIQueueMgr: Args: max_age_seconds: 最大存活时间（秒） + """ current_time = asyncio.get_event_loop().time() expired_sessions = [] @@ -131,18 +139,15 @@ class WecomAIQueueMgr: del self.pending_responses[session_id] logger.debug(f"[WecomAI] 清理过期响应: {session_id}") - def get_stats(self) -> Dict[str, int]: + def get_stats(self) -> dict[str, int]: """获取队列统计信息 Returns: 统计信息字典 + """ return { "input_queues": len(self.queues), "output_queues": len(self.back_queues), "pending_responses": len(self.pending_responses), } - - -# 全局队列管理器实例 -wecomai_queue_mgr = WecomAIQueueMgr() diff --git a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_server.py b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_server.py index bbb69d04..35acd906 100644 --- a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_server.py +++ b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_server.py @@ -1,12 +1,13 @@ -""" -企业微信智能机器人 HTTP 服务器 +"""企业微信智能机器人 HTTP 服务器 处理企业微信智能机器人的 HTTP 回调请求 """ import asyncio -from typing import Dict, Any, Optional, Callable +from collections.abc import Callable +from typing import Any import quart + from astrbot.api import logger from .wecomai_api import WecomAIBotAPIClient @@ -21,9 +22,7 @@ class WecomAIBotServer: host: str, port: int, api_client: WecomAIBotAPIClient, - message_handler: Optional[ - Callable[[Dict[str, Any], Dict[str, str]], Any] - ] = None, + message_handler: Callable[[dict[str, Any], dict[str, str]], Any] | None = None, ): """初始化服务器 @@ -32,6 +31,7 @@ class WecomAIBotServer: port: 监听端口 api_client: API客户端实例 message_handler: 消息处理回调函数 + """ self.host = host self.port = port @@ -45,7 +45,6 @@ class WecomAIBotServer: def _setup_routes(self): """设置 Quart 路由""" - # 使用 Quart 的 add_url_rule 方法添加路由 self.app.add_url_rule( "/webhook/wecom-ai-bot", @@ -98,7 +97,7 @@ class WecomAIBotServer: assert nonce is not None logger.debug( - f"收到消息回调，msg_signature={msg_signature}, timestamp={timestamp}, nonce={nonce}" + f"收到消息回调，msg_signature={msg_signature}, timestamp={timestamp}, nonce={nonce}", ) try: @@ -111,7 +110,10 @@ class WecomAIBotServer: # 解密消息 ret_code, message_data = await self.api_client.decrypt_message( - post_data, msg_signature, timestamp, nonce + post_data, + msg_signature, + timestamp, + nonce, ) if ret_code != WecomAIBotConstants.SUCCESS or not message_data: @@ -123,7 +125,8 @@ class WecomAIBotServer: if self.message_handler: try: response = await self.message_handler( - message_data, {"nonce": nonce, "timestamp": timestamp} + message_data, + {"nonce": nonce, "timestamp": timestamp}, ) except Exception as e: logger.error("消息处理器执行异常: %s", e) @@ -131,8 +134,7 @@ class WecomAIBotServer: if response: return response, 200, {"Content-Type": "text/plain"} - else: - return "success", 200, {"Content-Type": "text/plain"} + return "success", 200, {"Content-Type": "text/plain"} except Exception as e: logger.error("处理消息时发生异常: %s", e) diff --git a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_utils.py b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_utils.py index dccb2e26..f7cbe380 100644 --- a/astrbot/core/platform/sources/wecom_ai_bot/wecomai_utils.py +++ b/astrbot/core/platform/sources/wecom_ai_bot/wecomai_utils.py @@ -1,16 +1,17 @@ -""" -企业微信智能机器人工具模块 +"""企业微信智能机器人工具模块 提供常量定义、工具函数和辅助方法 """ -import string -import random -import hashlib -import base64 -import aiohttp import asyncio +import base64 +import hashlib +import secrets +import string +from typing import Any + +import aiohttp from Crypto.Cipher import AES -from typing import Any, Tuple + from astrbot.api import logger @@ -49,9 +50,10 @@ def generate_random_string(length: int = 10) -> str: Returns: 随机字符串 + """ letters = string.ascii_letters + string.digits - return "".join(random.choice(letters) for _ in range(length)) + return "".join(secrets.choice(letters) for _ in range(length)) def calculate_image_md5(image_data: bytes) -> str: @@ -62,6 +64,7 @@ def calculate_image_md5(image_data: bytes) -> str: Returns: MD5 哈希值（十六进制字符串） + """ return hashlib.md5(image_data).hexdigest() @@ -74,6 +77,7 @@ def encode_image_base64(image_data: bytes) -> str: Returns: Base64 编码的字符串 + """ return base64.b64encode(image_data).decode("utf-8") @@ -87,11 +91,12 @@ def format_session_id(session_type: str, session_id: str) -> str: Returns: 格式化后的会话 ID + """ return f"wecom_ai_bot_{session_type}_{session_id}" -def parse_session_id(formatted_session_id: str) -> Tuple[str, str]: +def parse_session_id(formatted_session_id: str) -> tuple[str, str]: """解析格式化的会话 ID Args: @@ -99,6 +104,7 @@ def parse_session_id(formatted_session_id: str) -> Tuple[str, str]: Returns: (会话类型, 原始会话ID) + """ parts = formatted_session_id.split("_", 3) if ( @@ -120,6 +126,7 @@ def safe_json_loads(json_str: str, default: Any = None) -> Any: Returns: 解析结果或默认值 + """ import json @@ -139,13 +146,15 @@ def format_error_response(error_code: int, error_msg: str) -> str: Returns: 格式化的错误响应字符串 + """ return f"Error {error_code}: {error_msg}" async def process_encrypted_image( - image_url: str, aes_key_base64: str -) -> Tuple[bool, str]: + image_url: str, + aes_key_base64: str, +) -> tuple[bool, str]: """下载并解密加密图片 Args: @@ -155,6 +164,7 @@ async def process_encrypted_image( Returns: Tuple[bool, str]: status 为 True 时 data 是解密后的图片数据的 base64 编码， status 为 False 时 data 是错误信息 + """ # 1. 下载加密图片 logger.info("开始下载加密图片: %s", image_url) @@ -165,7 +175,7 @@ async def process_encrypted_image( encrypted_data = await response.read() logger.info("图片下载成功，大小: %d 字节", len(encrypted_data)) except (aiohttp.ClientError, asyncio.TimeoutError) as e: - error_msg = f"下载图片失败: {str(e)}" + error_msg = f"下载图片失败: {e!s}" logger.error(error_msg) return False, error_msg diff --git a/astrbot/core/platform/sources/weixin_official_account/weixin_offacc_adapter.py b/astrbot/core/platform/sources/weixin_official_account/weixin_offacc_adapter.py index c67c2037..d1309374 100644 --- a/astrbot/core/platform/sources/weixin_official_account/weixin_offacc_adapter.py +++ b/astrbot/core/platform/sources/weixin_official_account/weixin_offacc_adapter.py @@ -1,28 +1,28 @@ +import asyncio import sys import uuid -import asyncio -import quart +import quart +from requests import Response +from wechatpy import WeChatClient, parse_message +from wechatpy.crypto import WeChatCrypto +from wechatpy.exceptions import InvalidSignatureException +from wechatpy.messages import BaseMessage, ImageMessage, TextMessage, VoiceMessage +from wechatpy.utils import check_signature + +from astrbot.api.event import MessageChain +from astrbot.api.message_components import Image, Plain, Record from astrbot.api.platform import ( - Platform, AstrBotMessage, MessageMember, - PlatformMetadata, MessageType, + Platform, + PlatformMetadata, + register_platform_adapter, ) -from astrbot.api.event import MessageChain -from astrbot.api.message_components import Plain, Image, Record -from astrbot.core.platform.astr_message_event import MessageSesion -from astrbot.api.platform import register_platform_adapter from astrbot.core import logger -from requests import Response +from astrbot.core.platform.astr_message_event import MessageSesion -from wechatpy.utils import check_signature -from wechatpy.crypto import WeChatCrypto -from wechatpy import WeChatClient -from wechatpy.messages import TextMessage, ImageMessage, VoiceMessage, BaseMessage -from wechatpy.exceptions import InvalidSignatureException -from wechatpy import parse_message from .weixin_offacc_event import WeixinOfficialAccountPlatformEvent if sys.version_info >= (3, 12): @@ -40,10 +40,14 @@ class WecomServer: self.encoding_aes_key = config.get("encoding_aes_key") self.appid = config.get("appid") self.server.add_url_rule( - "/callback/command", view_func=self.verify, methods=["GET"] + "/callback/command", + view_func=self.verify, + methods=["GET"], ) self.server.add_url_rule( - "/callback/command", view_func=self.callback_command, methods=["POST"] + "/callback/command", + view_func=self.callback_command, + methods=["POST"], ) self.crypto = WeChatCrypto(self.token, self.encoding_aes_key, self.appid) @@ -97,7 +101,7 @@ class WecomServer: async def start_polling(self): logger.info( - f"将在 {self.callback_server_host}:{self.port} 端口启动 微信公众平台 适配器。" + f"将在 {self.callback_server_host}:{self.port} 端口启动 微信公众平台 适配器。", ) await self.server.run_task( host=self.callback_server_host, @@ -109,25 +113,30 @@ class WecomServer: await self.shutdown_event.wait() -@register_platform_adapter("weixin_official_account", "微信公众平台 适配器") +@register_platform_adapter( + "weixin_official_account", "微信公众平台 适配器", support_streaming_message=False +) class WeixinOfficialAccountPlatformAdapter(Platform): def __init__( - self, platform_config: dict, platform_settings: dict, event_queue: asyncio.Queue + self, + platform_config: dict, + platform_settings: dict, + event_queue: asyncio.Queue, ) -> None: super().__init__(event_queue) self.config = platform_config self.settingss = platform_settings self.client_self_id = uuid.uuid4().hex[:8] self.api_base_url = platform_config.get( - "api_base_url", "https://api.weixin.qq.com/cgi-bin/" + "api_base_url", + "https://api.weixin.qq.com/cgi-bin/", ) self.active_send_mode = self.config.get("active_send_mode", False) if not self.api_base_url: self.api_base_url = "https://api.weixin.qq.com/cgi-bin/" - if self.api_base_url.endswith("/"): - self.api_base_url = self.api_base_url[:-1] + self.api_base_url = self.api_base_url.removesuffix("/") if not self.api_base_url.endswith("/cgi-bin"): self.api_base_url += "/cgi-bin" @@ -161,7 +170,8 @@ class WeixinOfficialAccountPlatformAdapter(Platform): await self.convert_message(msg, future) # I love shield so much! result = await asyncio.wait_for( - asyncio.shield(future), 60 + asyncio.shield(future), + 60, ) # wait for 60s logger.debug(f"Got future result: {result}") self.wexin_event_workers.pop(msg.id, None) @@ -175,7 +185,9 @@ class WeixinOfficialAccountPlatformAdapter(Platform): @override async def send_by_session( - self, session: MessageSesion, message_chain: MessageChain + self, + session: MessageSesion, + message_chain: MessageChain, ): await super().send_by_session(session, message_chain) @@ -185,6 +197,7 @@ class WeixinOfficialAccountPlatformAdapter(Platform): "weixin_official_account", "微信公众平台 适配器", id=self.config.get("id", "weixin_official_account"), + support_streaming_message=False, ) @override @@ -192,7 +205,9 @@ class WeixinOfficialAccountPlatformAdapter(Platform): await self.server.start_polling() async def convert_message( - self, msg, future: asyncio.Future = None + self, + msg, + future: asyncio.Future = None, ) -> AstrBotMessage | None: abm = AstrBotMessage() if isinstance(msg, TextMessage): @@ -224,7 +239,9 @@ class WeixinOfficialAccountPlatformAdapter(Platform): assert isinstance(msg, VoiceMessage) resp: Response = await asyncio.get_event_loop().run_in_executor( - None, self.client.media.download, msg.media_id + None, + self.client.media.download, + msg.media_id, ) path = f"data/temp/wecom_{msg.media_id}.amr" with open(path, "wb") as f: @@ -238,7 +255,7 @@ class WeixinOfficialAccountPlatformAdapter(Platform): audio.export(path_wav, format="wav") except Exception as e: logger.error( - f"转换音频失败: {e}。如果没有安装 pydub 和 ffmpeg 请先安装。" + f"转换音频失败: {e}。如果没有安装 pydub 和 ffmpeg 请先安装。", ) path_wav = path return diff --git a/astrbot/core/platform/sources/weixin_official_account/weixin_offacc_event.py b/astrbot/core/platform/sources/weixin_official_account/weixin_offacc_event.py index 4077cc1a..d138fc80 100644 --- a/astrbot/core/platform/sources/weixin_official_account/weixin_offacc_event.py +++ b/astrbot/core/platform/sources/weixin_official_account/weixin_offacc_event.py @@ -1,21 +1,20 @@ -import uuid import asyncio -from astrbot.api.event import AstrMessageEvent, MessageChain -from astrbot.api.platform import AstrBotMessage, PlatformMetadata -from astrbot.api.message_components import Plain, Image, Record -from wechatpy import WeChatClient -from wechatpy.replies import TextReply, ImageReply, VoiceReply +import uuid +from wechatpy import WeChatClient +from wechatpy.replies import ImageReply, TextReply, VoiceReply from astrbot.api import logger +from astrbot.api.event import AstrMessageEvent, MessageChain +from astrbot.api.message_components import Image, Plain, Record +from astrbot.api.platform import AstrBotMessage, PlatformMetadata try: import pydub except Exception: logger.warning( - "检测到 pydub 库未安装，微信公众平台将无法语音收发。如需使用语音，请前往管理面板 -> 控制台 -> 安装 Pip 库安装 pydub。" + "检测到 pydub 库未安装，微信公众平台将无法语音收发。如需使用语音，请前往管理面板 -> 控制台 -> 安装 Pip 库安装 pydub。", ) - pass class WeixinOfficialAccountPlatformEvent(AstrMessageEvent): @@ -32,7 +31,9 @@ class WeixinOfficialAccountPlatformEvent(AstrMessageEvent): @staticmethod async def send_with_client( - client: WeChatClient, message: MessageChain, user_name: str + client: WeChatClient, + message: MessageChain, + user_name: str, ): pass @@ -43,44 +44,44 @@ class WeixinOfficialAccountPlatformEvent(AstrMessageEvent): plain (str): 要分割的长文本 Returns: list[str]: 分割后的文本列表 + """ if len(plain) <= 2048: return [plain] - else: - result = [] - start = 0 - while start < len(plain): - # 剩下的字符串长度<2048时结束 - if start + 2048 >= len(plain): - result.append(plain[start:]) + result = [] + start = 0 + while start < len(plain): + # 剩下的字符串长度<2048时结束 + if start + 2048 >= len(plain): + result.append(plain[start:]) + break + + # 向前搜索分割标点符号 + end = min(start + 2048, len(plain)) + cut_position = end + for i in range(end, start, -1): + if i < len(plain) and plain[i - 1] in [ + "。", + "！", + "？", + ".", + "!", + "?", + "\n", + ";", + "；", + ]: + cut_position = i break - # 向前搜索分割标点符号 - end = min(start + 2048, len(plain)) + # 没找到合适的位置分割, 直接切分 + if cut_position == end and end < len(plain): cut_position = end - for i in range(end, start, -1): - if i < len(plain) and plain[i - 1] in [ - "。", - "！", - "？", - ".", - "!", - "?", - "\n", - ";", - "；", - ]: - cut_position = i - break - # 没找到合适的位置分割, 直接切分 - if cut_position == end and end < len(plain): - cut_position = end + result.append(plain[start:cut_position]) + start = cut_position - result.append(plain[start:cut_position]) - start = cut_position - - return result + return result async def send(self, message: MessageChain): message_obj = self.message_obj @@ -111,7 +112,7 @@ class WeixinOfficialAccountPlatformEvent(AstrMessageEvent): except Exception as e: logger.error(f"微信公众平台上传图片失败: {e}") await self.send( - MessageChain().message(f"微信公众平台上传图片失败: {e}") + MessageChain().message(f"微信公众平台上传图片失败: {e}"), ) return logger.debug(f"微信公众平台上传图片返回: {response}") @@ -136,7 +137,8 @@ class WeixinOfficialAccountPlatformEvent(AstrMessageEvent): # 转成amr record_path_amr = f"data/temp/{uuid.uuid4()}.amr" pydub.AudioSegment.from_wav(record_path).export( - record_path_amr, format="amr" + record_path_amr, + format="amr", ) with open(record_path_amr, "rb") as f: @@ -145,7 +147,7 @@ class WeixinOfficialAccountPlatformEvent(AstrMessageEvent): except Exception as e: logger.error(f"微信公众平台上传语音失败: {e}") await self.send( - MessageChain().message(f"微信公众平台上传语音失败: {e}") + MessageChain().message(f"微信公众平台上传语音失败: {e}"), ) return logger.info(f"微信公众平台上传语音返回: {response}") @@ -178,7 +180,7 @@ class WeixinOfficialAccountPlatformEvent(AstrMessageEvent): else: buffer.chain.extend(chain.chain) if not buffer: - return + return None buffer.squash_plain() await self.send(buffer) return await super().send_streaming(generator, use_fallback) diff --git a/astrbot/core/platform_message_history_mgr.py b/astrbot/core/platform_message_history_mgr.py index 16e59a5c..0e079e89 100644 --- a/astrbot/core/platform_message_history_mgr.py +++ b/astrbot/core/platform_message_history_mgr.py @@ -11,8 +11,8 @@ class PlatformMessageHistoryManager: platform_id: str, user_id: str, content: list[dict], # TODO: parse from message chain - sender_id: str = None, - sender_name: str = None, + sender_id: str | None = None, + sender_name: str | None = None, ): """Insert a new platform message history record.""" await self.db.insert_platform_message_history( @@ -43,5 +43,7 @@ class PlatformMessageHistoryManager: async def delete(self, platform_id: str, user_id: str, offset_sec: int = 86400): """Delete platform message history records older than the specified offset.""" await self.db.delete_platform_message_offset( - platform_id=platform_id, user_id=user_id, offset_sec=offset_sec + platform_id=platform_id, + user_id=user_id, + offset_sec=offset_sec, ) diff --git a/astrbot/core/provider/__init__.py b/astrbot/core/provider/__init__.py index ed7135fe..812e0217 100644 --- a/astrbot/core/provider/__init__.py +++ b/astrbot/core/provider/__init__.py @@ -1,5 +1,4 @@ -from .provider import Provider, Personality, STTProvider - from .entities import ProviderMetaData +from .provider import Provider, STTProvider -__all__ = ["Provider", "Personality", "ProviderMetaData", "STTProvider"] +__all__ = ["Provider", "ProviderMetaData", "STTProvider"] diff --git a/astrbot/core/provider/entites.py b/astrbot/core/provider/entites.py index dbbbca92..af97c4ab 100644 --- a/astrbot/core/provider/entites.py +++ b/astrbot/core/provider/entites.py @@ -1,19 +1,19 @@ from astrbot.core.provider.entities import ( + AssistantMessageSegment, + LLMResponse, + ProviderMetaData, ProviderRequest, ProviderType, - ProviderMetaData, - ToolCallsResult, - AssistantMessageSegment, ToolCallMessageSegment, - LLMResponse, + ToolCallsResult, ) __all__ = [ + "AssistantMessageSegment", + "LLMResponse", + "ProviderMetaData", "ProviderRequest", "ProviderType", - "ProviderMetaData", - "ToolCallsResult", - "AssistantMessageSegment", "ToolCallMessageSegment", - "LLMResponse", + "ToolCallsResult", ] diff --git a/astrbot/core/provider/entities.py b/astrbot/core/provider/entities.py index 85687c41..dc188f14 100644 --- a/astrbot/core/provider/entities.py +++ b/astrbot/core/provider/entities.py @@ -1,20 +1,24 @@ -import enum import base64 +import enum import json -from astrbot.core.utils.io import download_image_by_url -from astrbot import logger from dataclasses import dataclass, field -from typing import List, Dict, Type, Any -from astrbot.core.agent.tool import ToolSet -from openai.types.chat.chat_completion import ChatCompletion +from typing import Any + +from anthropic.types import Message as AnthropicMessage from google.genai.types import GenerateContentResponse -from anthropic.types import Message -from openai.types.chat.chat_completion_message_tool_call import ( - ChatCompletionMessageToolCall, +from openai.types.chat.chat_completion import ChatCompletion + +import astrbot.core.message.components as Comp +from astrbot import logger +from astrbot.core.agent.message import ( + AssistantMessageSegment, + ToolCall, + ToolCallMessageSegment, ) +from astrbot.core.agent.tool import ToolSet from astrbot.core.db.po import Conversation from astrbot.core.message.message_event_result import MessageChain -import astrbot.core.message.components as Comp +from astrbot.core.utils.io import download_image_by_url class ProviderType(enum.Enum): @@ -26,56 +30,31 @@ class ProviderType(enum.Enum): @dataclass -class ProviderMetaData: +class ProviderMeta: + """The basic metadata of a provider instance.""" + + id: str + """the unique id of the provider instance that user configured""" + model: str | None + """the model name of the provider instance currently used""" type: str - """提供商适配器名称，如 openai, ollama""" - desc: str = "" - """提供商适配器描述.""" + """the name of the provider adapter, such as openai, ollama""" provider_type: ProviderType = ProviderType.CHAT_COMPLETION - cls_type: Type | None = None + """the capability type of the provider adapter""" + +@dataclass +class ProviderMetaData(ProviderMeta): + """The metadata of a provider adapter for registration.""" + + desc: str = "" + """the short description of the provider adapter""" + cls_type: Any = None + """the class type of the provider adapter""" default_config_tmpl: dict | None = None - """平台的默认配置模板""" + """the default configuration template of the provider adapter""" provider_display_name: str | None = None - """显示在 WebUI 配置页中的提供商名称，如空则是 type""" - - -@dataclass -class ToolCallMessageSegment: - """OpenAI 格式的上下文中 role 为 tool 的消息段。参考: https://platform.openai.com/docs/guides/function-calling""" - - tool_call_id: str - content: str - role: str = "tool" - - def to_dict(self): - return { - "tool_call_id": self.tool_call_id, - "content": self.content, - "role": self.role, - } - - -@dataclass -class AssistantMessageSegment: - """OpenAI 格式的上下文中 role 为 assistant 的消息段。参考: https://platform.openai.com/docs/guides/function-calling""" - - content: str | None = None - tool_calls: List[ChatCompletionMessageToolCall | Dict] = field(default_factory=list) - role: str = "assistant" - - def to_dict(self): - ret: dict[str, str | list[dict]] = { - "role": self.role, - } - if self.content: - ret["content"] = self.content - if self.tool_calls: - tool_calls_dict = [ - tc if isinstance(tc, dict) else tc.to_dict() for tc in self.tool_calls - ] - ret["tool_calls"] = tool_calls_dict - return ret + """the display name of the provider shown in the WebUI configuration page; if empty, the type is used""" @dataclass @@ -84,38 +63,46 @@ class ToolCallsResult: tool_calls_info: AssistantMessageSegment """函数调用的信息""" - tool_calls_result: List[ToolCallMessageSegment] + tool_calls_result: list[ToolCallMessageSegment] """函数调用的结果""" - def to_openai_messages(self) -> List[Dict]: + def to_openai_messages(self) -> list[dict]: ret = [ - self.tool_calls_info.to_dict(), - *[item.to_dict() for item in self.tool_calls_result], + self.tool_calls_info.model_dump(), + *[item.model_dump() for item in self.tool_calls_result], ] return ret + def to_openai_messages_model( + self, + ) -> list[AssistantMessageSegment | ToolCallMessageSegment]: + return [ + self.tool_calls_info, + *self.tool_calls_result, + ] + @dataclass class ProviderRequest: - prompt: str + prompt: str | None = None """提示词""" - session_id: str = "" + session_id: str | None = "" """会话 ID""" image_urls: list[str] = field(default_factory=list) """图片 URL 列表""" func_tool: ToolSet | None = None """可用的函数工具""" contexts: list[dict] = field(default_factory=list) - """上下文。格式与 openai 的上下文格式一致： + """ + OpenAI 格式上下文列表。 参考 https://platform.openai.com/docs/api-reference/chat/create#chat-create-messages """ system_prompt: str = "" """系统提示词""" conversation: Conversation | None = None - + """关联的对话对象""" tool_calls_result: list[ToolCallsResult] | ToolCallsResult | None = None """附加的上次请求后工具调用的结果。参考: https://platform.openai.com/docs/guides/function-calling#handling-function-calls""" - model: str | None = None """模型名称，为 None 时使用提供商的默认模型""" @@ -175,13 +162,13 @@ class ProviderRequest: return result_parts - async def assemble_context(self) -> Dict: + async def assemble_context(self) -> dict: """将请求(prompt 和 image_urls)包装成 OpenAI 的消息格式。""" if self.image_urls: user_content = { "role": "user", "content": [ - {"type": "text", "text": self.prompt if self.prompt else "[图片]"} + {"type": "text", "text": self.prompt if self.prompt else "[图片]"}, ], } for image_url in self.image_urls: @@ -197,11 +184,10 @@ class ProviderRequest: logger.warning(f"图片 {image_url} 得到的结果为空，将忽略。") continue user_content["content"].append( - {"type": "image_url", "image_url": {"url": image_data}} + {"type": "image_url", "image_url": {"url": image_data}}, ) return user_content - else: - return {"role": "user", "content": self.prompt} + return {"role": "user", "content": self.prompt} async def _encode_image_bs64(self, image_url: str) -> str: """将图片转换为 base64""" @@ -216,34 +202,44 @@ class ProviderRequest: @dataclass class LLMResponse: role: str - """角色, assistant, tool, err""" + """The role of the message, e.g., assistant, tool, err""" result_chain: MessageChain | None = None - """返回的消息链""" - tools_call_args: List[Dict[str, Any]] = field(default_factory=list) - """工具调用参数""" - tools_call_name: List[str] = field(default_factory=list) - """工具调用名称""" - tools_call_ids: List[str] = field(default_factory=list) - """工具调用 ID""" + """A chain of message components representing the text completion from LLM.""" + tools_call_args: list[dict[str, Any]] = field(default_factory=list) + """Tool call arguments.""" + tools_call_name: list[str] = field(default_factory=list) + """Tool call names.""" + tools_call_ids: list[str] = field(default_factory=list) + """Tool call IDs.""" + tools_call_extra_content: dict[str, dict[str, Any]] = field(default_factory=dict) + """Tool call extra content. tool_call_id -> extra_content dict""" + reasoning_content: str = "" + """The reasoning content extracted from the LLM, if any.""" - raw_completion: ChatCompletion | GenerateContentResponse | Message | None = None - _new_record: Dict[str, Any] | None = None + raw_completion: ( + ChatCompletion | GenerateContentResponse | AnthropicMessage | None + ) = None + """The raw completion response from the LLM provider.""" _completion_text: str = "" + """The plain text of the completion.""" is_chunk: bool = False - """是否是流式输出的单个 Chunk""" + """Indicates if the response is a chunked response.""" def __init__( self, role: str, completion_text: str = "", result_chain: MessageChain | None = None, - tools_call_args: List[Dict[str, Any]] | None = None, - tools_call_name: List[str] | None = None, - tools_call_ids: List[str] | None = None, - raw_completion: ChatCompletion | None = None, - _new_record: Dict[str, Any] | None = None, + tools_call_args: list[dict[str, Any]] | None = None, + tools_call_name: list[str] | None = None, + tools_call_ids: list[str] | None = None, + tools_call_extra_content: dict[str, dict[str, Any]] | None = None, + raw_completion: ChatCompletion + | GenerateContentResponse + | AnthropicMessage + | None = None, is_chunk: bool = False, ): """初始化 LLMResponse @@ -255,6 +251,7 @@ class LLMResponse: tools_call_args (List[Dict[str, any]], optional): 工具调用参数. Defaults to None. tools_call_name (List[str], optional): 工具调用名称. Defaults to None. raw_completion (ChatCompletion, optional): 原始响应, OpenAI 格式. Defaults to None. + """ if tools_call_args is None: tools_call_args = [] @@ -262,6 +259,8 @@ class LLMResponse: tools_call_name = [] if tools_call_ids is None: tools_call_ids = [] + if tools_call_extra_content is None: + tools_call_extra_content = {} self.role = role self.completion_text = completion_text @@ -269,8 +268,8 @@ class LLMResponse: self.tools_call_args = tools_call_args self.tools_call_name = tools_call_name self.tools_call_ids = tools_call_ids + self.tools_call_extra_content = tools_call_extra_content self.raw_completion = raw_completion - self._new_record = _new_record self.is_chunk = is_chunk @property @@ -291,19 +290,41 @@ class LLMResponse: else: self._completion_text = value - def to_openai_tool_calls(self) -> List[Dict]: - """将工具调用信息转换为 OpenAI 格式""" + def to_openai_tool_calls(self) -> list[dict]: + """Convert to OpenAI tool calls format. Deprecated, use to_openai_to_calls_model instead.""" + ret = [] + for idx, tool_call_arg in enumerate(self.tools_call_args): + payload = { + "id": self.tools_call_ids[idx], + "function": { + "name": self.tools_call_name[idx], + "arguments": json.dumps(tool_call_arg), + }, + "type": "function", + } + if self.tools_call_extra_content.get(self.tools_call_ids[idx]): + payload["extra_content"] = self.tools_call_extra_content[ + self.tools_call_ids[idx] + ] + ret.append(payload) + return ret + + def to_openai_to_calls_model(self) -> list[ToolCall]: + """The same as to_openai_tool_calls but return pydantic model.""" ret = [] for idx, tool_call_arg in enumerate(self.tools_call_args): ret.append( - { - "id": self.tools_call_ids[idx], - "function": { - "name": self.tools_call_name[idx], - "arguments": json.dumps(tool_call_arg), - }, - "type": "function", - } + ToolCall( + id=self.tools_call_ids[idx], + function=ToolCall.FunctionBody( + name=self.tools_call_name[idx], + arguments=json.dumps(tool_call_arg), + ), + # the extra_content will not serialize if it's None when calling ToolCall.model_dump() + extra_content=self.tools_call_extra_content.get( + self.tools_call_ids[idx] + ), + ), ) return ret diff --git a/astrbot/core/provider/func_tool_manager.py b/astrbot/core/provider/func_tool_manager.py index 51cde0eb..8e04423e 100644 --- a/astrbot/core/provider/func_tool_manager.py +++ b/astrbot/core/provider/func_tool_manager.py @@ -1,17 +1,19 @@ from __future__ import annotations + +import asyncio +import copy import json import os -import asyncio +from collections.abc import Awaitable, Callable +from typing import Any + import aiohttp -from typing import Dict, List, Awaitable, Callable, Any from astrbot import logger from astrbot.core import sp - +from astrbot.core.agent.mcp_client import MCPClient, MCPTool +from astrbot.core.agent.tool import FunctionTool, ToolSet from astrbot.core.utils.astrbot_path import get_astrbot_data_path -from astrbot.core.agent.mcp_client import MCPClient -from astrbot.core.agent.tool import ToolSet, FunctionTool - DEFAULT_MCP_CONFIG = {"mcpServers": {}} @@ -23,14 +25,23 @@ SUPPORTED_TYPES = [ "boolean", ] # json schema 支持的数据类型 - +PY_TO_JSON_TYPE = { + "int": "number", + "float": "number", + "bool": "boolean", + "str": "string", + "dict": "object", + "list": "array", + "tuple": "array", + "set": "array", +} # alias FuncTool = FunctionTool def _prepare_config(config: dict) -> dict: """准备配置，处理嵌套格式""" - if "mcpServers" in config and config["mcpServers"]: + if config.get("mcpServers"): first_key = next(iter(config["mcpServers"])) config = config["mcpServers"][first_key] config.pop("active", None) @@ -72,8 +83,7 @@ async def _quick_test_mcp_connection(config: dict) -> tuple[bool, str]: ) as response: if response.status == 200: return True, "" - else: - return False, f"HTTP {response.status}: {response.reason}" + return False, f"HTTP {response.status}: {response.reason}" else: async with session.get( url, @@ -85,8 +95,7 @@ async def _quick_test_mcp_connection(config: dict) -> tuple[bool, str]: ) as response: if response.status == 200: return True, "" - else: - return False, f"HTTP {response.status}: {response.reason}" + return False, f"HTTP {response.status}: {response.reason}" except asyncio.TimeoutError: return False, f"连接超时: {timeout}秒" @@ -96,10 +105,10 @@ async def _quick_test_mcp_connection(config: dict) -> tuple[bool, str]: class FunctionToolManager: def __init__(self) -> None: - self.func_list: List[FuncTool] = [] - self.mcp_client_dict: Dict[str, MCPClient] = {} + self.func_list: list[FuncTool] = [] + self.mcp_client_dict: dict[str, MCPClient] = {} """MCP 服务列表""" - self.mcp_client_event: Dict[str, asyncio.Event] = {} + self.mcp_client_event: dict[str, asyncio.Event] = {} def empty(self) -> bool: return len(self.func_list) == 0 @@ -107,7 +116,7 @@ class FunctionToolManager: def spec_to_func( self, name: str, - func_args: list, + func_args: list[dict], desc: str, handler: Callable[..., Awaitable[Any]], ) -> FuncTool: @@ -116,10 +125,9 @@ class FunctionToolManager: "properties": {}, } for param in func_args: - params["properties"][param["name"]] = { - "type": param["type"], - "description": param["description"], - } + p = copy.deepcopy(param) + p.pop("name", None) + params["properties"][param["name"]] = p return FuncTool( name=name, parameters=params, @@ -150,14 +158,12 @@ class FunctionToolManager: func_args=func_args, desc=desc, handler=handler, - ) + ), ) logger.info(f"添加函数调用工具: {name}") def remove_func(self, name: str) -> None: - """ - 删除一个函数调用工具。 - """ + """删除一个函数调用工具。""" for i, f in enumerate(self.func_list): if f.name == name: self.func_list.pop(i) @@ -202,16 +208,16 @@ class FunctionToolManager: logger.info(f"未找到 MCP 服务配置文件，已创建默认配置文件 {mcp_json_file}") return - mcp_server_json_obj: Dict[str, Dict] = json.load( - open(mcp_json_file, "r", encoding="utf-8") + mcp_server_json_obj: dict[str, dict] = json.load( + open(mcp_json_file, encoding="utf-8"), )["mcpServers"] - for name in mcp_server_json_obj.keys(): + for name in mcp_server_json_obj: cfg = mcp_server_json_obj[name] if cfg.get("active", True): event = asyncio.Event() asyncio.create_task( - self._init_mcp_client_task_wrapper(name, cfg, event) + self._init_mcp_client_task_wrapper(name, cfg, event), ) self.mcp_client_event[name] = event @@ -257,18 +263,15 @@ class FunctionToolManager: self.func_list = [ f for f in self.func_list - if not (f.origin == "mcp" and f.mcp_server_name == name) + if not (isinstance(f, MCPTool) and f.mcp_server_name == name) ] # 将 MCP 工具转换为 FuncTool 并添加到 func_list for tool in mcp_client.tools: - func_tool = FuncTool( - name=tool.name, - parameters=tool.inputSchema, - description=tool.description, - origin="mcp", - mcp_server_name=name, + func_tool = MCPTool( + mcp_tool=tool, mcp_client=mcp_client, + mcp_server_name=name, ) self.func_list.append(func_tool) @@ -277,19 +280,22 @@ class FunctionToolManager: async def _terminate_mcp_client(self, name: str) -> None: """关闭并清理MCP客户端""" if name in self.mcp_client_dict: + client = self.mcp_client_dict[name] try: # 关闭MCP连接 - await self.mcp_client_dict[name].cleanup() - self.mcp_client_dict.pop(name) + await client.cleanup() except Exception as e: logger.error(f"清空 MCP 客户端资源 {name}: {e}。") - # 移除关联的FuncTool - self.func_list = [ - f - for f in self.func_list - if not (f.origin == "mcp" and f.mcp_server_name == name) - ] - logger.info(f"已关闭 MCP 服务 {name}") + finally: + # Remove client from dict after cleanup attempt (successful or not) + self.mcp_client_dict.pop(name, None) + # 移除关联的FuncTool + self.func_list = [ + f + for f in self.func_list + if not (isinstance(f, MCPTool) and f.mcp_server_name == name) + ] + logger.info(f"已关闭 MCP 服务 {name}") @staticmethod async def test_mcp_server_connection(config: dict) -> list[str]: @@ -325,9 +331,11 @@ class FunctionToolManager: event (asyncio.Event): Event to signal when the MCP client is ready. ready_future (asyncio.Future): Future to signal when the MCP client is ready. timeout (int): Timeout for the initialization. + Raises: TimeoutError: If the initialization does not complete within the specified timeout. Exception: If there is an error during initialization. + """ if not event: event = asyncio.Event() @@ -336,7 +344,7 @@ class FunctionToolManager: if name in self.mcp_client_dict: return asyncio.create_task( - self._init_mcp_client_task_wrapper(name, config, event, ready_future) + self._init_mcp_client_task_wrapper(name, config, event, ready_future), ) try: await asyncio.wait_for(ready_future, timeout=timeout) @@ -349,13 +357,16 @@ class FunctionToolManager: raise exc async def disable_mcp_server( - self, name: str | None = None, timeout: float = 10 + self, + name: str | None = None, + timeout: float = 10, ) -> None: """Disable an MCP server by its name. Args: name (str): The name of the MCP server to disable. If None, ALL MCP servers will be disabled. timeout (int): Timeout. + """ if name: if name not in self.mcp_client_event: @@ -372,7 +383,7 @@ class FunctionToolManager: self.func_list = [ f for f in self.func_list - if f.origin != "mcp" or f.mcp_server_name != name + if not (isinstance(f, MCPTool) and f.mcp_server_name == name) ] else: running_events = [ @@ -386,30 +397,26 @@ class FunctionToolManager: finally: self.mcp_client_event.clear() self.mcp_client_dict.clear() - self.func_list = [f for f in self.func_list if f.origin != "mcp"] + self.func_list = [ + f for f in self.func_list if not isinstance(f, MCPTool) + ] def get_func_desc_openai_style(self, omit_empty_parameter_field=False) -> list: - """ - 获得 OpenAI API 风格的**已经激活**的工具描述 - """ + """获得 OpenAI API 风格的**已经激活**的工具描述""" tools = [f for f in self.func_list if f.active] toolset = ToolSet(tools) return toolset.openai_schema( - omit_empty_parameter_field=omit_empty_parameter_field + omit_empty_parameter_field=omit_empty_parameter_field, ) def get_func_desc_anthropic_style(self) -> list: - """ - 获得 Anthropic API 风格的**已经激活**的工具描述 - """ + """获得 Anthropic API 风格的**已经激活**的工具描述""" tools = [f for f in self.func_list if f.active] toolset = ToolSet(tools) return toolset.anthropic_schema() def get_func_desc_google_genai_style(self) -> dict: - """ - 获得 Google GenAI API 风格的**已经激活**的工具描述 - """ + """获得 Google GenAI API 风格的**已经激活**的工具描述""" tools = [f for f in self.func_list if f.active] toolset = ToolSet(tools) return toolset.google_schema() @@ -418,13 +425,18 @@ class FunctionToolManager: """停用一个已经注册的函数调用工具。 Returns: - 如果没找到，会返回 False""" + 如果没找到，会返回 False + + """ func_tool = self.get_func(name) if func_tool is not None: func_tool.active = False inactivated_llm_tools: list = sp.get( - "inactivated_llm_tools", [], scope="global", scope_id="global" + "inactivated_llm_tools", + [], + scope="global", + scope_id="global", ) if name not in inactivated_llm_tools: inactivated_llm_tools.append(name) @@ -445,13 +457,16 @@ class FunctionToolManager: if func_tool.handler_module_path in star_map: if not star_map[func_tool.handler_module_path].activated: raise ValueError( - f"此函数调用工具所属的插件 {star_map[func_tool.handler_module_path].name} 已被禁用，请先在管理面板启用再激活此工具。" + f"此函数调用工具所属的插件 {star_map[func_tool.handler_module_path].name} 已被禁用，请先在管理面板启用再激活此工具。", ) func_tool.active = True inactivated_llm_tools: list = sp.get( - "inactivated_llm_tools", [], scope="global", scope_id="global" + "inactivated_llm_tools", + [], + scope="global", + scope_id="global", ) if name in inactivated_llm_tools: inactivated_llm_tools.remove(name) @@ -479,7 +494,7 @@ class FunctionToolManager: return DEFAULT_MCP_CONFIG try: - with open(self.mcp_config_path, "r", encoding="utf-8") as f: + with open(self.mcp_config_path, encoding="utf-8") as f: return json.load(f) except Exception as e: logger.error(f"加载 MCP 配置失败: {e}") @@ -509,7 +524,8 @@ class FunctionToolManager: if response.status == 200: data = await response.json() mcp_server_list = data.get("data", {}).get( - "mcp_server_list", [] + "mcp_server_list", + [], ) local_mcp_config = self.load_mcp_config() @@ -541,23 +557,23 @@ class FunctionToolManager: self.enable_mcp_server( name=name, config=local_mcp_config["mcpServers"][name], - ) + ), ) await asyncio.gather(*tasks) logger.info( - f"从 ModelScope 同步了 {synced_count} 个 MCP 服务器" + f"从 ModelScope 同步了 {synced_count} 个 MCP 服务器", ) else: logger.warning("没有找到可用的 ModelScope MCP 服务器") else: raise Exception( - f"ModelScope API 请求失败: HTTP {response.status}" + f"ModelScope API 请求失败: HTTP {response.status}", ) except aiohttp.ClientError as e: - raise Exception(f"网络连接错误: {str(e)}") + raise Exception(f"网络连接错误: {e!s}") except Exception as e: - raise Exception(f"同步 ModelScope MCP 服务器时发生错误: {str(e)}") + raise Exception(f"同步 ModelScope MCP 服务器时发生错误: {e!s}") def __str__(self): return str(self.func_list) diff --git a/astrbot/core/provider/manager.py b/astrbot/core/provider/manager.py index 6666b33e..ec255041 100644 --- a/astrbot/core/provider/manager.py +++ b/astrbot/core/provider/manager.py @@ -1,21 +1,20 @@ import asyncio import traceback -from typing import List from astrbot.core import logger, sp from astrbot.core.astrbot_config_mgr import AstrBotConfigManager from astrbot.core.db import BaseDatabase +from ..persona_mgr import PersonaManager from .entities import ProviderType from .provider import ( + EmbeddingProvider, Provider, + RerankProvider, STTProvider, TTSProvider, - EmbeddingProvider, - RerankProvider, ) from .register import llm_tools, provider_cls_map -from ..persona_mgr import PersonaManager class ProviderManager: @@ -28,7 +27,7 @@ class ProviderManager: self.persona_mgr = persona_mgr self.acm = acm config = acm.confs["default"] - self.providers_config: List = config["provider"] + self.providers_config: list = config["provider"] self.provider_settings: dict = config["provider_settings"] self.provider_stt_settings: dict = config.get("provider_stt_settings", {}) self.provider_tts_settings: dict = config.get("provider_tts_settings", {}) @@ -36,15 +35,15 @@ class ProviderManager: # 人格相关属性，v4.0.0 版本后被废弃，推荐使用 PersonaManager self.default_persona_name = persona_mgr.default_persona - self.provider_insts: List[Provider] = [] + self.provider_insts: list[Provider] = [] """加载的 Provider 的实例""" - self.stt_provider_insts: List[STTProvider] = [] + self.stt_provider_insts: list[STTProvider] = [] """加载的 Speech To Text Provider 的实例""" - self.tts_provider_insts: List[TTSProvider] = [] + self.tts_provider_insts: list[TTSProvider] = [] """加载的 Text To Speech Provider 的实例""" - self.embedding_provider_insts: List[EmbeddingProvider] = [] + self.embedding_provider_insts: list[EmbeddingProvider] = [] """加载的 Embedding Provider 的实例""" - self.rerank_provider_insts: List[RerankProvider] = [] + self.rerank_provider_insts: list[RerankProvider] = [] """加载的 Rerank Provider 的实例""" self.inst_map: dict[ str, @@ -77,7 +76,10 @@ class ProviderManager: return self.persona_mgr.selected_default_persona_v3 async def set_provider( - self, provider_id: str, provider_type: ProviderType, umo: str | None = None + self, + provider_id: str, + provider_type: ProviderType, + umo: str | None = None, ): """设置提供商。 @@ -87,6 +89,7 @@ class ProviderManager: umo (str, optional): 用户会话 ID，用于提供商会话隔离。 Version 4.0.0: 这个版本下已经默认隔离提供商 + """ if provider_id not in self.inst_map: raise ValueError(f"提供商 {provider_id} 不存在，无法设置。") @@ -101,17 +104,20 @@ class ProviderManager: prov = self.inst_map[provider_id] if provider_type == ProviderType.TEXT_TO_SPEECH and isinstance( - prov, TTSProvider + prov, + TTSProvider, ): self.curr_tts_provider_inst = prov sp.put("curr_provider_tts", provider_id, scope="global", scope_id="global") elif provider_type == ProviderType.SPEECH_TO_TEXT and isinstance( - prov, STTProvider + prov, + STTProvider, ): self.curr_stt_provider_inst = prov sp.put("curr_provider_stt", provider_id, scope="global", scope_id="global") elif provider_type == ProviderType.CHAT_COMPLETION and isinstance( - prov, Provider + prov, + Provider, ): self.curr_provider_inst = prov sp.put("curr_provider", provider_id, scope="global", scope_id="global") @@ -121,7 +127,9 @@ class ProviderManager: return self.inst_map.get(provider_id) def get_using_provider( - self, provider_type: ProviderType, umo=None + self, + provider_type: ProviderType, + umo=None, ) -> Provider | STTProvider | TTSProvider | None: """获取正在使用的提供商实例。 @@ -131,6 +139,7 @@ class ProviderManager: Returns: Provider: 正在使用的提供商实例。 + """ provider = None if umo: @@ -175,7 +184,11 @@ class ProviderManager: async def initialize(self): # 逐个初始化提供商 for provider_config in self.providers_config: - await self.load_provider(provider_config) + try: + await self.load_provider(provider_config) + except Exception as e: + logger.error(traceback.format_exc()) + logger.error(e) # 设置默认提供商 selected_provider_id = sp.get( @@ -216,7 +229,7 @@ class ProviderManager: return logger.info( - f"载入 {provider_config['type']}({provider_config['id']}) 服务提供商 ..." + f"载入 {provider_config['type']}({provider_config['id']}) 服务提供商 ...", ) # 动态导入 @@ -228,6 +241,8 @@ class ProviderManager: ) case "zhipu_chat_completion": from .sources.zhipu_source import ProviderZhipu as ProviderZhipu + case "groq_chat_completion": + from .sources.groq_source import ProviderGroq as ProviderGroq case "anthropic_chat_completion": from .sources.anthropic_source import ( ProviderAnthropic as ProviderAnthropic, @@ -256,6 +271,10 @@ class ProviderManager: from .sources.whisper_selfhosted_source import ( ProviderOpenAIWhisperSelfHost as ProviderOpenAIWhisperSelfHost, ) + case "xinference_stt": + from .sources.xinference_stt_provider import ( + ProviderXinferenceSTT as ProviderXinferenceSTT, + ) case "openai_tts_api": from .sources.openai_tts_api_source import ( ProviderOpenAITTSAPI as ProviderOpenAITTSAPI, @@ -308,20 +327,28 @@ class ProviderManager: from .sources.vllm_rerank_source import ( VLLMRerankProvider as VLLMRerankProvider, ) + case "xinference_rerank": + from .sources.xinference_rerank_source import ( + XinferenceRerankProvider as XinferenceRerankProvider, + ) + case "bailian_rerank": + from .sources.bailian_rerank_source import ( + BailianRerankProvider as BailianRerankProvider, + ) except (ImportError, ModuleNotFoundError) as e: logger.critical( - f"加载 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}。可能是因为有未安装的依赖。" + f"加载 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}。可能是因为有未安装的依赖。", ) return except Exception as e: logger.critical( - f"加载 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}。未知原因" + f"加载 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}。未知原因", ) return if provider_config["type"] not in provider_cls_map: logger.error( - f"未找到适用于 {provider_config['type']}({provider_config['id']}) 的提供商适配器，请检查是否已经安装或者名称填写错误。已跳过。" + f"未找到适用于 {provider_config['type']}({provider_config['id']}) 的提供商适配器，请检查是否已经安装或者名称填写错误。已跳过。", ) return @@ -333,6 +360,8 @@ class ProviderManager: logger.error(f"无法找到 {provider_metadata.type} 的类") return + provider_metadata.id = provider_config["id"] + if provider_metadata.provider_type == ProviderType.SPEECH_TO_TEXT: # STT 任务 inst = cls_type(provider_config, self.provider_settings) @@ -347,7 +376,7 @@ class ProviderManager: ): self.curr_stt_provider_inst = inst logger.info( - f"已选择 {provider_config['type']}({provider_config['id']}) 作为当前语音转文本提供商适配器。" + f"已选择 {provider_config['type']}({provider_config['id']}) 作为当前语音转文本提供商适配器。", ) if not self.curr_stt_provider_inst: self.curr_stt_provider_inst = inst @@ -363,7 +392,7 @@ class ProviderManager: if self.provider_settings.get("provider_id") == provider_config["id"]: self.curr_tts_provider_inst = inst logger.info( - f"已选择 {provider_config['type']}({provider_config['id']}) 作为当前文本转语音提供商适配器。" + f"已选择 {provider_config['type']}({provider_config['id']}) 作为当前文本转语音提供商适配器。", ) if not self.curr_tts_provider_inst: self.curr_tts_provider_inst = inst @@ -373,7 +402,6 @@ class ProviderManager: inst = cls_type( provider_config, self.provider_settings, - self.selected_default_persona, ) if getattr(inst, "initialize", None): @@ -386,7 +414,7 @@ class ProviderManager: ): self.curr_provider_inst = inst logger.info( - f"已选择 {provider_config['type']}({provider_config['id']}) 作为当前提供商适配器。" + f"已选择 {provider_config['type']}({provider_config['id']}) 作为当前提供商适配器。", ) if not self.curr_provider_inst: self.curr_provider_inst = inst @@ -404,9 +432,11 @@ class ProviderManager: self.inst_map[provider_config["id"]] = inst except Exception as e: - logger.error(traceback.format_exc()) logger.error( - f"实例化 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}" + f"实例化 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}", + ) + raise Exception( + f"实例化 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}", ) async def reload(self, provider_config: dict): @@ -426,7 +456,7 @@ class ProviderManager: elif self.curr_provider_inst is None and len(self.provider_insts) > 0: self.curr_provider_inst = self.provider_insts[0] logger.info( - f"自动选择 {self.curr_provider_inst.meta().id} 作为当前提供商适配器。" + f"自动选择 {self.curr_provider_inst.meta().id} 作为当前提供商适配器。", ) if len(self.stt_provider_insts) == 0: @@ -434,7 +464,7 @@ class ProviderManager: elif self.curr_stt_provider_inst is None and len(self.stt_provider_insts) > 0: self.curr_stt_provider_inst = self.stt_provider_insts[0] logger.info( - f"自动选择 {self.curr_stt_provider_inst.meta().id} 作为当前语音转文本提供商适配器。" + f"自动选择 {self.curr_stt_provider_inst.meta().id} 作为当前语音转文本提供商适配器。", ) if len(self.tts_provider_insts) == 0: @@ -442,7 +472,7 @@ class ProviderManager: elif self.curr_tts_provider_inst is None and len(self.tts_provider_insts) > 0: self.curr_tts_provider_inst = self.tts_provider_insts[0] logger.info( - f"自动选择 {self.curr_tts_provider_inst.meta().id} 作为当前文本转语音提供商适配器。" + f"自动选择 {self.curr_tts_provider_inst.meta().id} 作为当前文本转语音提供商适配器。", ) def get_insts(self): @@ -451,7 +481,7 @@ class ProviderManager: async def terminate_provider(self, provider_id: str): if provider_id in self.inst_map: logger.info( - f"终止 {provider_id} 提供商适配器({len(self.provider_insts)}, {len(self.stt_provider_insts)}, {len(self.tts_provider_insts)}) ..." + f"终止 {provider_id} 提供商适配器({len(self.provider_insts)}, {len(self.stt_provider_insts)}, {len(self.tts_provider_insts)}) ...", ) if self.inst_map[provider_id] in self.provider_insts: @@ -478,7 +508,7 @@ class ProviderManager: await self.inst_map[provider_id].terminate() # type: ignore logger.info( - f"{provider_id} 提供商适配器已终止({len(self.provider_insts)}, {len(self.stt_provider_insts)}, {len(self.tts_provider_insts)})" + f"{provider_id} 提供商适配器已终止({len(self.provider_insts)}, {len(self.stt_provider_insts)}, {len(self.tts_provider_insts)})", ) del self.inst_map[provider_id] diff --git a/astrbot/core/provider/provider.py b/astrbot/core/provider/provider.py index 901036b1..3ae5f2bd 100644 --- a/astrbot/core/provider/provider.py +++ b/astrbot/core/provider/provider.py @@ -1,162 +1,170 @@ import abc -from typing import List -from typing import AsyncGenerator +import asyncio +from collections.abc import AsyncGenerator + +from astrbot.core.agent.message import Message from astrbot.core.agent.tool import ToolSet from astrbot.core.provider.entities import ( LLMResponse, - ToolCallsResult, - ProviderType, + ProviderMeta, RerankResult, + ToolCallsResult, ) from astrbot.core.provider.register import provider_cls_map -from astrbot.core.db.po import Personality -from dataclasses import dataclass - - -@dataclass -class ProviderMeta: - id: str - model: str - type: str - provider_type: ProviderType class AbstractProvider(abc.ABC): + """Provider Abstract Class""" + def __init__(self, provider_config: dict) -> None: super().__init__() self.model_name = "" self.provider_config = provider_config def set_model(self, model_name: str): - """设置当前使用的模型名称""" + """Set the current model name""" self.model_name = model_name def get_model(self) -> str: - """获得当前使用的模型名称""" + """Get the current model name""" return self.model_name def meta(self) -> ProviderMeta: - """获取 Provider 的元数据""" + """Get the provider metadata""" provider_type_name = self.provider_config["type"] meta_data = provider_cls_map.get(provider_type_name) - provider_type = meta_data.provider_type if meta_data else None - return ProviderMeta( - id=self.provider_config["id"], + if not meta_data: + raise ValueError(f"Provider type {provider_type_name} not registered") + meta = ProviderMeta( + id=self.provider_config.get("id", "default"), model=self.get_model(), type=provider_type_name, - provider_type=provider_type, + provider_type=meta_data.provider_type, ) + return meta class Provider(AbstractProvider): + """Chat Provider""" + def __init__( self, provider_config: dict, provider_settings: dict, - default_persona: Personality | None = None, ) -> None: super().__init__(provider_config) - self.provider_settings = provider_settings - self.curr_personality = default_persona - """维护了当前的使用的 persona，即人格。可能为 None""" - @abc.abstractmethod def get_current_key(self) -> str: - raise NotImplementedError() + raise NotImplementedError - def get_keys(self) -> List[str]: + def get_keys(self) -> list[str]: """获得提供商 Key""" keys = self.provider_config.get("key", [""]) return keys or [""] @abc.abstractmethod def set_key(self, key: str): - raise NotImplementedError() + raise NotImplementedError @abc.abstractmethod - async def get_models(self) -> List[str]: + async def get_models(self) -> list[str]: """获得支持的模型列表""" - raise NotImplementedError() + raise NotImplementedError @abc.abstractmethod async def text_chat( self, - prompt: str, - session_id: str = None, - image_urls: list[str] = None, - func_tool: ToolSet = None, - contexts: list = None, - system_prompt: str = None, - tool_calls_result: ToolCallsResult | list[ToolCallsResult] = None, + prompt: str | None = None, + session_id: str | None = None, + image_urls: list[str] | None = None, + func_tool: ToolSet | None = None, + contexts: list[Message] | list[dict] | None = None, + system_prompt: str | None = None, + tool_calls_result: ToolCallsResult | list[ToolCallsResult] | None = None, model: str | None = None, **kwargs, ) -> LLMResponse: """获得 LLM 的文本对话结果。会使用当前的模型进行对话。 Args: - prompt: 提示词 + prompt: 提示词，和 contexts 二选一使用，如果都指定，则会将 prompt（以及可能的 image_urls） 作为最新的一条记录添加到 contexts 中 session_id: 会话 ID(此属性已经被废弃) image_urls: 图片 URL 列表 - tools: Function-calling 工具 - contexts: 上下文 + tools: tool set + contexts: 上下文，和 prompt 二选一使用 tool_calls_result: 回传给 LLM 的工具调用结果。参考: https://platform.openai.com/docs/guides/function-calling kwargs: 其他参数 Notes: - 如果传入了 image_urls，将会在对话时附上图片。如果模型不支持图片输入，将会抛出错误。 - 如果传入了 tools，将会使用 tools 进行 Function-calling。如果模型不支持 Function-calling，将会抛出错误。 + """ ... async def text_chat_stream( self, - prompt: str, - session_id: str = None, - image_urls: list[str] = None, - func_tool: ToolSet = None, - contexts: list = None, - system_prompt: str = None, - tool_calls_result: ToolCallsResult | list[ToolCallsResult] = None, + prompt: str | None = None, + session_id: str | None = None, + image_urls: list[str] | None = None, + func_tool: ToolSet | None = None, + contexts: list[Message] | list[dict] | None = None, + system_prompt: str | None = None, + tool_calls_result: ToolCallsResult | list[ToolCallsResult] | None = None, model: str | None = None, **kwargs, ) -> AsyncGenerator[LLMResponse, None]: """获得 LLM 的流式文本对话结果。会使用当前的模型进行对话。在生成的最后会返回一次完整的结果。 Args: - prompt: 提示词 + prompt: 提示词，和 contexts 二选一使用，如果都指定，则会将 prompt（以及可能的 image_urls） 作为最新的一条记录添加到 contexts 中 session_id: 会话 ID(此属性已经被废弃) image_urls: 图片 URL 列表 - tools: Function-calling 工具 - contexts: 上下文 + tools: tool set + contexts: 上下文，和 prompt 二选一使用 tool_calls_result: 回传给 LLM 的工具调用结果。参考: https://platform.openai.com/docs/guides/function-calling kwargs: 其他参数 Notes: - 如果传入了 image_urls，将会在对话时附上图片。如果模型不支持图片输入，将会抛出错误。 - 如果传入了 tools，将会使用 tools 进行 Function-calling。如果模型不支持 Function-calling，将会抛出错误。 + """ ... - async def pop_record(self, context: List): - """ - 弹出 context 第一条非系统提示词对话记录 - """ + async def pop_record(self, context: list): + """弹出 context 第一条非系统提示词对话记录""" poped = 0 indexs_to_pop = [] for idx, record in enumerate(context): if record["role"] == "system": continue - else: - indexs_to_pop.append(idx) - poped += 1 - if poped == 2: - break + indexs_to_pop.append(idx) + poped += 1 + if poped == 2: + break for idx in reversed(indexs_to_pop): context.pop(idx) + def _ensure_message_to_dicts( + self, + messages: list[dict] | list[Message] | None, + ) -> list[dict]: + """Convert a list of Message objects to a list of dictionaries.""" + if not messages: + return [] + dicts: list[dict] = [] + for message in messages: + if isinstance(message, Message): + dicts.append(message.model_dump()) + else: + dicts.append(message) + + return dicts + class STTProvider(AbstractProvider): def __init__(self, provider_config: dict, provider_settings: dict) -> None: @@ -167,7 +175,7 @@ class STTProvider(AbstractProvider): @abc.abstractmethod async def get_text(self, audio_url: str) -> str: """获取音频的文本""" - raise NotImplementedError() + raise NotImplementedError class TTSProvider(AbstractProvider): @@ -179,7 +187,7 @@ class TTSProvider(AbstractProvider): @abc.abstractmethod async def get_audio(self, text: str) -> str: """获取文本的音频，返回音频文件路径""" - raise NotImplementedError() + raise NotImplementedError class EmbeddingProvider(AbstractProvider): @@ -203,6 +211,73 @@ class EmbeddingProvider(AbstractProvider): """获取向量的维度""" ... + async def get_embeddings_batch( + self, + texts: list[str], + batch_size: int = 16, + tasks_limit: int = 3, + max_retries: int = 3, + progress_callback=None, + ) -> list[list[float]]: + """批量获取文本的向量，分批处理以节省内存 + + Args: + texts: 文本列表 + batch_size: 每批处理的文本数量 + tasks_limit: 并发任务数量限制 + max_retries: 失败时的最大重试次数 + progress_callback: 进度回调函数，接收参数 (current, total) + + Returns: + 向量列表 + + """ + semaphore = asyncio.Semaphore(tasks_limit) + all_embeddings: list[list[float]] = [] + failed_batches: list[tuple[int, list[str]]] = [] + completed_count = 0 + total_count = len(texts) + + async def process_batch(batch_idx: int, batch_texts: list[str]): + nonlocal completed_count + async with semaphore: + for attempt in range(max_retries): + try: + batch_embeddings = await self.get_embeddings(batch_texts) + all_embeddings.extend(batch_embeddings) + completed_count += len(batch_texts) + if progress_callback: + await progress_callback(completed_count, total_count) + return + except Exception as e: + if attempt == max_retries - 1: + # 最后一次重试失败，记录失败的批次 + failed_batches.append((batch_idx, batch_texts)) + raise Exception( + f"批次 {batch_idx} 处理失败，已重试 {max_retries} 次: {e!s}", + ) + # 等待一段时间后重试，使用指数退避 + await asyncio.sleep(2**attempt) + + tasks = [] + for i in range(0, len(texts), batch_size): + batch_texts = texts[i : i + batch_size] + batch_idx = i // batch_size + tasks.append(process_batch(batch_idx, batch_texts)) + + # 收集所有任务的结果，包括失败的任务 + results = await asyncio.gather(*tasks, return_exceptions=True) + + # 检查是否有失败的任务 + errors = [r for r in results if isinstance(r, Exception)] + if errors: + error_msg = ( + f"有 {len(errors)} 个批次处理失败: {'; '.join(str(e) for e in errors)}" + ) + raise Exception(error_msg) + + return all_embeddings + class RerankProvider(AbstractProvider): def __init__(self, provider_config: dict, provider_settings: dict) -> None: @@ -212,7 +287,10 @@ class RerankProvider(AbstractProvider): @abc.abstractmethod async def rerank( - self, query: str, documents: list[str], top_n: int | None = None + self, + query: str, + documents: list[str], + top_n: int | None = None, ) -> list[RerankResult]: """获取查询和文档的重排序分数""" ... diff --git a/astrbot/core/provider/register.py b/astrbot/core/provider/register.py index 02d7934d..3ad83784 100644 --- a/astrbot/core/provider/register.py +++ b/astrbot/core/provider/register.py @@ -1,11 +1,11 @@ -from typing import List, Dict -from .entities import ProviderMetaData, ProviderType from astrbot.core import logger + +from .entities import ProviderMetaData, ProviderType from .func_tool_manager import FuncCall -provider_registry: List[ProviderMetaData] = [] +provider_registry: list[ProviderMetaData] = [] """维护了通过装饰器注册的 Provider""" -provider_cls_map: Dict[str, ProviderMetaData] = {} +provider_cls_map: dict[str, ProviderMetaData] = {} """维护了 Provider 类型名称和 ProviderMetadata 的映射""" llm_tools = FuncCall() @@ -15,15 +15,15 @@ def register_provider_adapter( provider_type_name: str, desc: str, provider_type: ProviderType = ProviderType.CHAT_COMPLETION, - default_config_tmpl: dict = None, - provider_display_name: str = None, + default_config_tmpl: dict | None = None, + provider_display_name: str | None = None, ): """用于注册平台适配器的带参装饰器""" def decorator(cls): if provider_type_name in provider_cls_map: raise ValueError( - f"检测到大模型提供商适配器 {provider_type_name} 已经注册，可能发生了大模型提供商适配器类型命名冲突。" + f"检测到大模型提供商适配器 {provider_type_name} 已经注册，可能发生了大模型提供商适配器类型命名冲突。", ) # 添加必备选项 @@ -36,6 +36,8 @@ def register_provider_adapter( default_config_tmpl["id"] = provider_type_name pm = ProviderMetaData( + id="default", # will be replaced when instantiated + model=None, type=provider_type_name, desc=desc, provider_type=provider_type, diff --git a/astrbot/core/provider/sources/anthropic_source.py b/astrbot/core/provider/sources/anthropic_source.py index 57bffdc8..f05d205c 100644 --- a/astrbot/core/provider/sources/anthropic_source.py +++ b/astrbot/core/provider/sources/anthropic_source.py @@ -1,39 +1,38 @@ -import json -import anthropic import base64 -from typing import List +import json +from collections.abc import AsyncGenerator from mimetypes import guess_type +import anthropic from anthropic import AsyncAnthropic from anthropic.types import Message -from astrbot.core.utils.io import download_image_by_url -from astrbot.api.provider import Provider from astrbot import logger -from astrbot.core.provider.func_tool_manager import FuncCall -from ..register import register_provider_adapter +from astrbot.api.provider import Provider from astrbot.core.provider.entities import LLMResponse -from typing import AsyncGenerator +from astrbot.core.provider.func_tool_manager import ToolSet +from astrbot.core.utils.io import download_image_by_url + +from ..register import register_provider_adapter @register_provider_adapter( - "anthropic_chat_completion", "Anthropic Claude API 提供商适配器" + "anthropic_chat_completion", + "Anthropic Claude API 提供商适配器", ) class ProviderAnthropic(Provider): def __init__( self, provider_config, provider_settings, - default_persona=None, ) -> None: super().__init__( provider_config, provider_settings, - default_persona, ) self.chosen_api_key: str = "" - self.api_keys: List = super().get_keys() + self.api_keys: list = super().get_keys() self.chosen_api_key = self.api_keys[0] if len(self.api_keys) > 0 else "" self.base_url = provider_config.get("api_base", "https://api.anthropic.com") self.timeout = provider_config.get("timeout", 120) @@ -41,7 +40,9 @@ class ProviderAnthropic(Provider): self.timeout = int(self.timeout) self.client = AsyncAnthropic( - api_key=self.chosen_api_key, timeout=self.timeout, base_url=self.base_url + api_key=self.chosen_api_key, + timeout=self.timeout, + base_url=self.base_url, ) self.set_model(provider_config["model_config"]["model"]) @@ -54,6 +55,7 @@ class ProviderAnthropic(Provider): Returns: system_prompt: 系统提示内容 new_messages: 处理后的消息列表，去除系统提示 + """ system_prompt = "" new_messages = [] @@ -73,18 +75,19 @@ class ProviderAnthropic(Provider): "input": ( json.loads(tool_call["function"]["arguments"]) if isinstance( - tool_call["function"]["arguments"], str + tool_call["function"]["arguments"], + str, ) else tool_call["function"]["arguments"] ), "id": tool_call["id"], - } + }, ) new_messages.append( { "role": "assistant", "content": blocks, - } + }, ) elif message["role"] == "tool": new_messages.append( @@ -95,16 +98,16 @@ class ProviderAnthropic(Provider): "type": "tool_result", "tool_use_id": message["tool_call_id"], "content": message["content"], - } + }, ], - } + }, ) else: new_messages.append(message) return system_prompt, new_messages - async def _query(self, payloads: dict, tools: FuncCall) -> LLMResponse: + async def _query(self, payloads: dict, tools: ToolSet | None) -> LLMResponse: if tools: if tool_list := tools.get_func_desc_anthropic_style(): payloads["tools"] = tool_list @@ -135,7 +138,9 @@ class ProviderAnthropic(Provider): return llm_response async def _query_stream( - self, payloads: dict, tools: FuncCall + self, + payloads: dict, + tools: ToolSet | None, ) -> AsyncGenerator[LLMResponse, None]: if tools: if tool_list := tools.get_func_desc_anthropic_style(): @@ -154,7 +159,9 @@ class ProviderAnthropic(Provider): if event.content_block.type == "text": # 文本块开始 yield LLMResponse( - role="assistant", completion_text="", is_chunk=True + role="assistant", + completion_text="", + is_chunk=True, ) elif event.content_block.type == "tool_use": # 工具使用块开始，初始化缓冲区 @@ -198,7 +205,7 @@ class ProviderAnthropic(Provider): "id": tool_info["id"], "name": tool_info["name"], "input": tool_info["input"], - } + }, ) yield LLMResponse( @@ -218,7 +225,9 @@ class ProviderAnthropic(Provider): # 返回最终的完整结果 final_response = LLMResponse( - role="assistant", completion_text=final_text, is_chunk=False + role="assistant", + completion_text=final_text, + is_chunk=False, ) if final_tool_calls: @@ -232,7 +241,7 @@ class ProviderAnthropic(Provider): async def text_chat( self, - prompt, + prompt=None, session_id=None, image_urls=None, func_tool=None, @@ -244,8 +253,13 @@ class ProviderAnthropic(Provider): ) -> LLMResponse: if contexts is None: contexts = [] - new_record = await self.assemble_context(prompt, image_urls) - context_query = [*contexts, new_record] + new_record = None + if prompt is not None: + new_record = await self.assemble_context(prompt, image_urls) + context_query = self._ensure_message_to_dicts(contexts) + if new_record: + context_query.append(new_record) + if system_prompt: context_query.insert(0, {"role": "system", "content": system_prompt}) @@ -295,8 +309,12 @@ class ProviderAnthropic(Provider): ): if contexts is None: contexts = [] - new_record = await self.assemble_context(prompt, image_urls) - context_query = [*contexts, new_record] + new_record = None + if prompt is not None: + new_record = await self.assemble_context(prompt, image_urls) + context_query = self._ensure_message_to_dicts(contexts) + if new_record: + context_query.append(new_record) if system_prompt: context_query.insert(0, {"role": "system", "content": system_prompt}) @@ -326,7 +344,7 @@ class ProviderAnthropic(Provider): async for llm_response in self._query_stream(payloads, func_tool): yield llm_response - async def assemble_context(self, text: str, image_urls: List[str] = None): + async def assemble_context(self, text: str, image_urls: list[str] | None = None): """组装上下文，支持文本和图片""" if not image_urls: return {"role": "user", "content": text} @@ -365,15 +383,13 @@ class ProviderAnthropic(Provider): else image_data ), }, - } + }, ) return {"role": "user", "content": content} async def encode_image_bs64(self, image_url: str) -> str: - """ - 将图片转换为 base64 - """ + """将图片转换为 base64""" if image_url.startswith("base64://"): return image_url.replace("base64://", "data:image/jpeg;base64,") with open(image_url, "rb") as f: @@ -384,7 +400,7 @@ class ProviderAnthropic(Provider): def get_current_key(self) -> str: return self.chosen_api_key - async def get_models(self) -> List[str]: + async def get_models(self) -> list[str]: models_str = [] models = await self.client.models.list() models = sorted(models.data, key=lambda x: x.id) diff --git a/astrbot/core/provider/sources/azure_tts_source.py b/astrbot/core/provider/sources/azure_tts_source.py index 6ddf452d..e85d9179 100644 --- a/astrbot/core/provider/sources/azure_tts_source.py +++ b/astrbot/core/provider/sources/azure_tts_source.py @@ -1,15 +1,15 @@ -import uuid -import time +import asyncio +import hashlib import json import re -import hashlib -import random -import asyncio +import secrets +import time +import uuid from pathlib import Path -from typing import Dict from xml.sax.saxutils import escape from httpx import AsyncClient, Timeout + from astrbot.core.config.default import VERSION from ..entities import ProviderType @@ -21,7 +21,7 @@ TEMP_DIR.mkdir(parents=True, exist_ok=True) class OTTSProvider: - def __init__(self, config: Dict): + def __init__(self, config: dict): self.skey = config["OTTS_SKEY"] self.api_url = config["OTTS_URL"] self.auth_time_url = config["OTTS_AUTH_TIME"] @@ -54,11 +54,13 @@ class OTTSProvider: async def _generate_signature(self) -> str: await self._sync_time() timestamp = int(time.time()) + self.time_offset - nonce = "".join(random.choices("abcdefghijklmnopqrstuvwxyz0123456789", k=10)) + nonce = "".join( + secrets.choice("abcdefghijklmnopqrstuvwxyz0123456789") for _ in range(10) + ) path = re.sub(r"^https?://[^/]+", "", self.api_url) or "/" return f"{timestamp}-{nonce}-0-{hashlib.md5(f'{path}-{timestamp}-{nonce}-0-{self.skey}'.encode()).hexdigest()}" - async def get_audio(self, text: str, voice_params: Dict) -> str: + async def get_audio(self, text: str, voice_params: dict) -> str: file_path = TEMP_DIR / f"otts-{uuid.uuid4()}.wav" signature = await self._generate_signature() for attempt in range(self.retry_count): @@ -86,7 +88,7 @@ class OTTSProvider: return str(file_path.resolve()) except Exception as e: if attempt == self.retry_count - 1: - raise RuntimeError(f"OTTS请求失败: {str(e)}") from e + raise RuntimeError(f"OTTS请求失败: {e!s}") from e await asyncio.sleep(0.5 * (attempt + 1)) @@ -94,7 +96,8 @@ class AzureNativeProvider(TTSProvider): def __init__(self, provider_config: dict, provider_settings: dict): super().__init__(provider_config, provider_settings) self.subscription_key = provider_config.get( - "azure_tts_subscription_key", "" + "azure_tts_subscription_key", + "", ).strip() if not re.fullmatch(r"^[a-zA-Z0-9]{32}$", self.subscription_key): raise ValueError("无效的Azure订阅密钥") @@ -119,7 +122,7 @@ class AzureNativeProvider(TTSProvider): "User-Agent": f"AstrBot/{VERSION}", "Content-Type": "application/ssml+xml", "X-Microsoft-OutputFormat": "riff-48khz-16bit-mono-pcm", - } + }, ) return self @@ -132,7 +135,8 @@ class AzureNativeProvider(TTSProvider): f"https://{self.region}.api.cognitive.microsoft.com/sts/v1.0/issuetoken" ) response = await self.client.post( - token_url, headers={"Ocp-Apim-Subscription-Key": self.subscription_key} + token_url, + headers={"Ocp-Apim-Subscription-Key": self.subscription_key}, ) response.raise_for_status() self.token = response.text diff --git a/astrbot/core/provider/sources/bailian_rerank_source.py b/astrbot/core/provider/sources/bailian_rerank_source.py new file mode 100644 index 00000000..e6f6f1a4 --- /dev/null +++ b/astrbot/core/provider/sources/bailian_rerank_source.py @@ -0,0 +1,236 @@ +import os + +import aiohttp + +from astrbot import logger + +from ..entities import ProviderType, RerankResult +from ..provider import RerankProvider +from ..register import register_provider_adapter + + +class BailianRerankError(Exception): + """百炼重排序服务异常基类""" + + pass + + +class BailianAPIError(BailianRerankError): + """百炼API返回错误""" + + pass + + +class BailianNetworkError(BailianRerankError): + """百炼网络请求错误""" + + pass + + +@register_provider_adapter( + "bailian_rerank", "阿里云百炼文本排序适配器", provider_type=ProviderType.RERANK +) +class BailianRerankProvider(RerankProvider): + """阿里云百炼文本重排序适配器.""" + + def __init__(self, provider_config: dict, provider_settings: dict) -> None: + super().__init__(provider_config, provider_settings) + self.provider_config = provider_config + self.provider_settings = provider_settings + + # API配置 + self.api_key = provider_config.get("rerank_api_key") or os.getenv( + "DASHSCOPE_API_KEY", "" + ) + if not self.api_key: + raise ValueError("阿里云百炼 API Key 不能为空。") + + self.model = provider_config.get("rerank_model", "qwen3-rerank") + self.timeout = provider_config.get("timeout", 30) + self.return_documents = provider_config.get("return_documents", False) + self.instruct = provider_config.get("instruct", "") + + self.base_url = provider_config.get( + "rerank_api_base", + "https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank", + ) + + # 设置HTTP客户端 + headers = { + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + } + + self.client = aiohttp.ClientSession( + headers=headers, timeout=aiohttp.ClientTimeout(total=self.timeout) + ) + + # 设置模型名称 + self.set_model(self.model) + + logger.info(f"AstrBot 百炼 Rerank 初始化完成。模型: {self.model}") + + def _build_payload( + self, query: str, documents: list[str], top_n: int | None + ) -> dict: + """构建请求载荷 + + Args: + query: 查询文本 + documents: 文档列表 + top_n: 返回前N个结果，如果为None则返回所有结果 + + Returns: + 请求载荷字典 + """ + base = {"model": self.model, "input": {"query": query, "documents": documents}} + + params = { + k: v + for k, v in [ + ("top_n", top_n if top_n is not None and top_n > 0 else None), + ("return_documents", True if self.return_documents else None), + ( + "instruct", + self.instruct + if self.instruct and self.model == "qwen3-rerank" + else None, + ), + ] + if v is not None + } + + if params: + base["parameters"] = params + + return base + + def _parse_results(self, data: dict) -> list[RerankResult]: + """解析API响应结果 + + Args: + data: API响应数据 + + Returns: + 重排序结果列表 + + Raises: + BailianAPIError: API返回错误 + KeyError: 结果缺少必要字段 + """ + # 检查响应状态 + if data.get("code", "200") != "200": + raise BailianAPIError( + f"百炼 API 错误: {data.get('code')} – {data.get('message', '')}" + ) + + results = data.get("output", {}).get("results", []) + if not results: + logger.warning(f"百炼 Rerank 返回空结果: {data}") + return [] + + # 转换为RerankResult对象，使用.get()避免KeyError + rerank_results = [] + for idx, result in enumerate(results): + try: + index = result.get("index", idx) + relevance_score = result.get("relevance_score", 0.0) + + if relevance_score is None: + logger.warning(f"结果 {idx} 缺少 relevance_score，使用默认值 0.0") + relevance_score = 0.0 + + rerank_result = RerankResult( + index=index, relevance_score=relevance_score + ) + rerank_results.append(rerank_result) + except Exception as e: + logger.warning(f"解析结果 {idx} 时出错: {e}, result={result}") + continue + + return rerank_results + + def _log_usage(self, data: dict) -> None: + """记录使用量信息 + + Args: + data: API响应数据 + """ + tokens = data.get("usage", {}).get("total_tokens", 0) + if tokens > 0: + logger.debug(f"百炼 Rerank 消耗 Token: {tokens}") + + async def rerank( + self, + query: str, + documents: list[str], + top_n: int | None = None, + ) -> list[RerankResult]: + """ + 对文档进行重排序 + + Args: + query: 查询文本 + documents: 待排序的文档列表 + top_n: 返回前N个结果，如果为None则使用配置中的默认值 + + Returns: + 重排序结果列表 + """ + if not documents: + logger.warning("文档列表为空，返回空结果") + return [] + + if not query.strip(): + logger.warning("查询文本为空，返回空结果") + return [] + + # 检查限制 + if len(documents) > 500: + logger.warning( + f"文档数量({len(documents)})超过限制(500)，将截断前500个文档" + ) + documents = documents[:500] + + try: + # 构建请求载荷，如果top_n为None则返回所有重排序结果 + payload = self._build_payload(query, documents, top_n) + + logger.debug( + f"百炼 Rerank 请求: query='{query[:50]}...', 文档数量={len(documents)}" + ) + + # 发送请求 + async with self.client.post(self.base_url, json=payload) as response: + response.raise_for_status() + response_data = await response.json() + + # 解析结果并记录使用量 + results = self._parse_results(response_data) + self._log_usage(response_data) + + logger.debug(f"百炼 Rerank 成功返回 {len(results)} 个结果") + + return results + + except aiohttp.ClientError as e: + error_msg = f"网络请求失败: {e}" + logger.error(f"百炼 Rerank 网络请求失败: {e}") + raise BailianNetworkError(error_msg) from e + except BailianRerankError: + raise + except Exception as e: + error_msg = f"重排序失败: {e}" + logger.error(f"百炼 Rerank 处理失败: {e}") + raise BailianRerankError(error_msg) from e + + async def terminate(self) -> None: + """关闭HTTP客户端会话.""" + if self.client: + logger.info("关闭 百炼 Rerank 客户端会话") + try: + await self.client.close() + except Exception as e: + logger.error(f"关闭 百炼 Rerank 客户端时出错: {e}") + finally: + self.client = None diff --git a/astrbot/core/provider/sources/coze_api_client.py b/astrbot/core/provider/sources/coze_api_client.py index a768979c..e8f3a1e2 100644 --- a/astrbot/core/provider/sources/coze_api_client.py +++ b/astrbot/core/provider/sources/coze_api_client.py @@ -1,8 +1,11 @@ -import json import asyncio -import aiohttp import io -from typing import Dict, List, Any, AsyncGenerator +import json +from collections.abc import AsyncGenerator +from typing import Any + +import aiohttp + from astrbot.core import logger @@ -32,7 +35,9 @@ class CozeAPIClient: "Accept": "text/event-stream", } self.session = aiohttp.ClientSession( - headers=headers, timeout=timeout, connector=connector + headers=headers, + timeout=timeout, + connector=connector, ) return self.session @@ -46,6 +51,7 @@ class CozeAPIClient: file_data (bytes): 文件的二进制数据 Returns: str: 上传成功后返回的 file_id + """ session = await self._ensure_session() url = f"{self.api_base}/v1/files/upload" @@ -64,12 +70,12 @@ class CozeAPIClient: response_text = await response.text() logger.debug( - f"文件上传响应状态: {response.status}, 内容: {response_text}" + f"文件上传响应状态: {response.status}, 内容: {response_text}", ) if response.status != 200: raise Exception( - f"文件上传失败，状态码: {response.status}, 响应: {response_text}" + f"文件上传失败，状态码: {response.status}, 响应: {response_text}", ) try: @@ -88,8 +94,8 @@ class CozeAPIClient: logger.error("文件上传超时") raise Exception("文件上传超时") except Exception as e: - logger.error(f"文件上传失败: {str(e)}") - raise Exception(f"文件上传失败: {str(e)}") + logger.error(f"文件上传失败: {e!s}") + raise Exception(f"文件上传失败: {e!s}") async def download_image(self, image_url: str) -> bytes: """下载图片并返回字节数据 @@ -98,6 +104,7 @@ class CozeAPIClient: image_url (str): 图片的URL Returns: bytes: 图片的二进制数据 + """ session = await self._ensure_session() @@ -110,19 +117,19 @@ class CozeAPIClient: return image_data except Exception as e: - logger.error(f"下载图片失败 {image_url}: {str(e)}") - raise Exception(f"下载图片失败: {str(e)}") + logger.error(f"下载图片失败 {image_url}: {e!s}") + raise Exception(f"下载图片失败: {e!s}") async def chat_messages( self, bot_id: str, user_id: str, - additional_messages: List[Dict] | None = None, + additional_messages: list[dict] | None = None, conversation_id: str | None = None, auto_save_history: bool = True, stream: bool = True, timeout: float = 120, - ) -> AsyncGenerator[Dict[str, Any], None]: + ) -> AsyncGenerator[dict[str, Any], None]: """发送聊天消息并返回流式响应 Args: @@ -133,6 +140,7 @@ class CozeAPIClient: auto_save_history: 是否自动保存历史 stream: 是否流式响应 timeout: 超时时间 + """ session = await self._ensure_session() url = f"{self.api_base}/v3/chat" @@ -198,7 +206,7 @@ class CozeAPIClient: except asyncio.TimeoutError: raise Exception(f"Coze API 流式请求超时 ({timeout}秒)") except Exception as e: - raise Exception(f"Coze API 流式请求失败: {str(e)}") + raise Exception(f"Coze API 流式请求失败: {e!s}") async def clear_context(self, conversation_id: str): """清空会话上下文 @@ -207,6 +215,7 @@ class CozeAPIClient: conversation_id: 会话ID Returns: dict: API响应结果 + """ session = await self._ensure_session() url = f"{self.api_base}/v3/conversation/message/clear_context" @@ -230,7 +239,7 @@ class CozeAPIClient: except asyncio.TimeoutError: raise Exception("Coze API 请求超时") except aiohttp.ClientError as e: - raise Exception(f"Coze API 请求失败: {str(e)}") + raise Exception(f"Coze API 请求失败: {e!s}") async def get_message_list( self, @@ -248,6 +257,7 @@ class CozeAPIClient: offset: 偏移量 Returns: dict: API响应结果 + """ session = await self._ensure_session() url = f"{self.api_base}/v3/conversation/message/list" @@ -264,8 +274,8 @@ class CozeAPIClient: return await response.json() except Exception as e: - logger.error(f"获取Coze消息列表失败: {str(e)}") - raise Exception(f"获取Coze消息列表失败: {str(e)}") + logger.error(f"获取Coze消息列表失败: {e!s}") + raise Exception(f"获取Coze消息列表失败: {e!s}") async def close(self): """关闭会话""" @@ -275,8 +285,8 @@ class CozeAPIClient: if __name__ == "__main__": - import os import asyncio + import os async def test_coze_api_client(): api_key = os.getenv("COZE_API_KEY", "") diff --git a/astrbot/core/provider/sources/coze_source.py b/astrbot/core/provider/sources/coze_source.py index 639af081..6f1355bf 100644 --- a/astrbot/core/provider/sources/coze_source.py +++ b/astrbot/core/provider/sources/coze_source.py @@ -1,13 +1,15 @@ -import json -import os import base64 import hashlib -from typing import AsyncGenerator, Dict -from astrbot.core.message.message_event_result import MessageChain +import json +import os +from collections.abc import AsyncGenerator + import astrbot.core.message.components as Comp -from astrbot.api.provider import Provider from astrbot import logger +from astrbot.api.provider import Provider +from astrbot.core.message.message_event_result import MessageChain from astrbot.core.provider.entities import LLMResponse + from ..register import register_provider_adapter from .coze_api_client import CozeAPIClient @@ -18,12 +20,10 @@ class ProviderCoze(Provider): self, provider_config, provider_settings, - default_persona=None, ) -> None: super().__init__( provider_config, provider_settings, - default_persona, ) self.api_key = provider_config.get("coze_api_key", "") if not self.api_key: @@ -34,18 +34,18 @@ class ProviderCoze(Provider): self.api_base: str = provider_config.get("coze_api_base", "https://api.coze.cn") if not isinstance(self.api_base, str) or not self.api_base.startswith( - ("http://", "https://") + ("http://", "https://"), ): raise Exception( - "Coze API Base URL 格式不正确，必须以 http:// 或 https:// 开头。" + "Coze API Base URL 格式不正确，必须以 http:// 或 https:// 开头。", ) self.timeout = provider_config.get("timeout", 120) if isinstance(self.timeout, str): self.timeout = int(self.timeout) self.auto_save_history = provider_config.get("auto_save_history", True) - self.conversation_ids: Dict[str, str] = {} - self.file_id_cache: Dict[str, Dict[str, str]] = {} + self.conversation_ids: dict[str, str] = {} + self.file_id_cache: dict[str, dict[str, str]] = {} # 创建 API 客户端 self.api_client = CozeAPIClient(api_key=self.api_key, api_base=self.api_base) @@ -59,8 +59,8 @@ class ProviderCoze(Provider): Returns: str: 缓存键 - """ + """ try: if is_base64 and data.startswith("data:image/"): try: @@ -71,26 +71,24 @@ class ProviderCoze(Provider): except Exception: cache_key = hashlib.md5(encoded.encode("utf-8")).hexdigest() return cache_key + elif data.startswith(("http://", "https://")): + # URL图片，使用URL作为缓存键 + cache_key = hashlib.md5(data.encode("utf-8")).hexdigest() + return cache_key else: - if data.startswith(("http://", "https://")): - # URL图片，使用URL作为缓存键 - cache_key = hashlib.md5(data.encode("utf-8")).hexdigest() - return cache_key - else: - clean_path = ( - data.split("_")[0] - if "_" in data and len(data.split("_")) >= 3 - else data - ) + clean_path = ( + data.split("_")[0] + if "_" in data and len(data.split("_")) >= 3 + else data + ) - if os.path.exists(clean_path): - with open(clean_path, "rb") as f: - file_content = f.read() - cache_key = hashlib.md5(file_content).hexdigest() - return cache_key - else: - cache_key = hashlib.md5(clean_path.encode("utf-8")).hexdigest() - return cache_key + if os.path.exists(clean_path): + with open(clean_path, "rb") as f: + file_content = f.read() + cache_key = hashlib.md5(file_content).hexdigest() + return cache_key + cache_key = hashlib.md5(clean_path.encode("utf-8")).hexdigest() + return cache_key except Exception as e: cache_key = hashlib.md5(data.encode("utf-8")).hexdigest() @@ -117,7 +115,9 @@ class ProviderCoze(Provider): return file_id async def _download_and_upload_image( - self, image_url: str, session_id: str | None = None + self, + image_url: str, + session_id: str | None = None, ) -> str: """下载图片并上传到 Coze，返回 file_id""" # 计算哈希实现缓存 @@ -142,14 +142,15 @@ class ProviderCoze(Provider): return file_id except Exception as e: - logger.error(f"处理图片失败 {image_url}: {str(e)}") - raise Exception(f"处理图片失败: {str(e)}") + logger.error(f"处理图片失败 {image_url}: {e!s}") + raise Exception(f"处理图片失败: {e!s}") async def _process_context_images( - self, content: str | list, session_id: str + self, + content: str | list, + session_id: str, ) -> str: """处理上下文中的图片内容，将 base64 图片上传并替换为 file_id""" - try: if isinstance(content, str): return content @@ -184,14 +185,15 @@ class ProviderCoze(Provider): continue # 计算哈希用于缓存 cache_key = self._generate_cache_key( - image_data, is_base64=image_data.startswith("data:image/") + image_data, + is_base64=image_data.startswith("data:image/"), ) # 检查缓存 if cache_key in self.file_id_cache[session_id]: file_id = self.file_id_cache[session_id][cache_key] processed_content.append( - {"type": "image", "file_id": file_id} + {"type": "image", "file_id": file_id}, ) else: # 上传图片并缓存 @@ -207,7 +209,8 @@ class ProviderCoze(Provider): elif image_data.startswith(("http://", "https://")): # URL 图片 file_id = await self._download_and_upload_image( - image_data, session_id + image_data, + session_id, ) # 为URL图片也添加缓存 self.file_id_cache[session_id][cache_key] = file_id @@ -222,22 +225,21 @@ class ProviderCoze(Provider): ) else: logger.warning( - f"无法处理的图片格式: {image_data[:50]}..." + f"无法处理的图片格式: {image_data[:50]}...", ) continue processed_content.append( - {"type": "image", "file_id": file_id} + {"type": "image", "file_id": file_id}, ) result = json.dumps(processed_content, ensure_ascii=False) return result except Exception as e: - logger.error(f"处理上下文图片失败: {str(e)}") + logger.error(f"处理上下文图片失败: {e!s}") if isinstance(content, str): return content - else: - return json.dumps(content, ensure_ascii=False) + return json.dumps(content, ensure_ascii=False) async def text_chat( self, @@ -262,8 +264,10 @@ class ProviderCoze(Provider): system_prompt (str): 系统提示语 tool_calls_result (ToolCallsResult | List[ToolCallsResult]): 工具调用结果(不支持) model (str): 模型名称(不支持) + Returns: LLMResponse: LLM响应对象 + """ accumulated_content = "" final_response = None @@ -291,8 +295,7 @@ class ProviderCoze(Provider): if accumulated_content: chain = MessageChain(chain=[Comp.Plain(accumulated_content)]) return LLMResponse(role="assistant", result_chain=chain) - else: - return LLMResponse(role="assistant", completion_text="") + return LLMResponse(role="assistant", completion_text="") async def text_chat_stream( self, @@ -319,9 +322,14 @@ class ProviderCoze(Provider): if system_prompt: if not self.auto_save_history or not conversation_id: additional_messages.append( - {"role": "system", "content": system_prompt, "content_type": "text"} + { + "role": "system", + "content": system_prompt, + "content_type": "text", + }, ) + contexts = self._ensure_message_to_dicts(contexts) if not self.auto_save_history and contexts: # 如果关闭了自动保存历史，传入上下文 for ctx in contexts: @@ -343,14 +351,15 @@ class ProviderCoze(Provider): ) ): processed_content = await self._process_context_images( - content, user_id + content, + user_id, ) additional_messages.append( { "role": ctx["role"], "content": processed_content, "content_type": "object_string", - } + }, ) else: # 纯文本 @@ -363,7 +372,7 @@ class ProviderCoze(Provider): else json.dumps(content, ensure_ascii=False) ), "content_type": "text", - } + }, ) else: logger.info(f"[Coze] 跳过格式不正确的上下文: {ctx}") @@ -380,7 +389,8 @@ class ProviderCoze(Provider): if url.startswith(("http://", "https://")): # 网络图片 file_id = await self._download_and_upload_image( - url, user_id + url, + user_id, ) else: # 本地文件或 base64 @@ -389,37 +399,41 @@ class ProviderCoze(Provider): _, encoded = url.split(",", 1) image_data = base64.b64decode(encoded) cache_key = self._generate_cache_key( - url, is_base64=True + url, + is_base64=True, ) file_id = await self._upload_file( - image_data, user_id, cache_key + image_data, + user_id, + cache_key, + ) + # 本地文件 + elif os.path.exists(url): + with open(url, "rb") as f: + image_data = f.read() + # 用文件路径和修改时间来缓存 + file_stat = os.stat(url) + cache_key = self._generate_cache_key( + f"{url}_{file_stat.st_mtime}_{file_stat.st_size}", + is_base64=False, + ) + file_id = await self._upload_file( + image_data, + user_id, + cache_key, ) else: - # 本地文件 - if os.path.exists(url): - with open(url, "rb") as f: - image_data = f.read() - # 用文件路径和修改时间来缓存 - file_stat = os.stat(url) - cache_key = self._generate_cache_key( - f"{url}_{file_stat.st_mtime}_{file_stat.st_size}", - is_base64=False, - ) - file_id = await self._upload_file( - image_data, user_id, cache_key - ) - else: - logger.warning(f"图片文件不存在: {url}") - continue + logger.warning(f"图片文件不存在: {url}") + continue object_string_content.append( { "type": "image", "file_id": file_id, - } + }, ) except Exception as e: - logger.error(f"处理图片失败 {url}: {str(e)}") + logger.error(f"处理图片失败 {url}: {e!s}") continue if object_string_content: @@ -429,18 +443,17 @@ class ProviderCoze(Provider): "role": "user", "content": content, "content_type": "object_string", - } - ) - else: - # 纯文本 - if prompt: - additional_messages.append( - { - "role": "user", - "content": prompt, - "content_type": "text", - } + }, ) + # 纯文本 + elif prompt: + additional_messages.append( + { + "role": "user", + "content": prompt, + "content_type": "text", + }, + ) try: accumulated_content = "" @@ -534,10 +547,10 @@ class ProviderCoze(Provider): ) except Exception as e: - logger.error(f"Coze 流式请求失败: {str(e)}") + logger.error(f"Coze 流式请求失败: {e!s}") yield LLMResponse( role="err", - completion_text=f"Coze 流式请求失败: {str(e)}", + completion_text=f"Coze 流式请求失败: {e!s}", is_chunk=False, ) @@ -558,12 +571,11 @@ class ProviderCoze(Provider): if "code" in response and response["code"] == 0: self.conversation_ids.pop(user_id, None) return True - else: - logger.warning(f"清空 Coze 会话上下文失败: {response}") - return False + logger.warning(f"清空 Coze 会话上下文失败: {response}") + return False except Exception as e: - logger.error(f"清空 Coze 会话失败: {str(e)}") + logger.error(f"清空 Coze 会话失败: {e!s}") return False async def get_current_key(self): @@ -590,7 +602,10 @@ class ProviderCoze(Provider): self.bot_id = model async def get_human_readable_context( - self, session_id: str, page: int = 1, page_size: int = 10 + self, + session_id: str, + page: int = 1, + page_size: int = 10, ): """获取人类可读的上下文历史""" user_id = session_id @@ -627,7 +642,7 @@ class ProviderCoze(Provider): return readable_history except Exception as e: - logger.error(f"获取 Coze 消息历史失败: {str(e)}") + logger.error(f"获取 Coze 消息历史失败: {e!s}") return [] async def terminate(self): diff --git a/astrbot/core/provider/sources/dashscope_source.py b/astrbot/core/provider/sources/dashscope_source.py index 4e14d20d..7c690e04 100644 --- a/astrbot/core/provider/sources/dashscope_source.py +++ b/astrbot/core/provider/sources/dashscope_source.py @@ -1,15 +1,17 @@ -import re import asyncio import functools -from typing import List -from .. import Provider, Personality -from ..entities import LLMResponse -from ..func_tool_manager import FuncCall -from ..register import register_provider_adapter -from astrbot.core.message.message_event_result import MessageChain -from .openai_source import ProviderOpenAIOfficial -from astrbot.core import logger, sp +import re + from dashscope import Application +from dashscope.app.application_response import ApplicationResponse + +from astrbot.core import logger, sp +from astrbot.core.message.message_event_result import MessageChain + +from .. import Provider +from ..entities import LLMResponse +from ..register import register_provider_adapter +from .openai_source import ProviderOpenAIOfficial @register_provider_adapter("dashscope", "Dashscope APP 适配器。") @@ -18,13 +20,11 @@ class ProviderDashscope(ProviderOpenAIOfficial): self, provider_config: dict, provider_settings: dict, - default_persona: Personality | None = None, ) -> None: Provider.__init__( self, provider_config, provider_settings, - default_persona, ) self.api_key = provider_config.get("dashscope_api_key", "") if not self.api_key: @@ -51,6 +51,7 @@ class ProviderDashscope(ProviderOpenAIOfficial): Returns: bool: 是否有 RAG 选项 + """ if self.rag_options and ( len(self.rag_options.get("pipeline_ids", [])) > 0 @@ -62,14 +63,16 @@ class ProviderDashscope(ProviderOpenAIOfficial): async def text_chat( self, prompt: str, - session_id: str = None, - image_urls: List[str] = [], - func_tool: FuncCall = None, - contexts: List = None, - system_prompt: str = None, + session_id=None, + image_urls=None, + func_tool=None, + contexts=None, + system_prompt=None, model=None, **kwargs, ) -> LLMResponse: + if image_urls is None: + image_urls = [] if contexts is None: contexts = [] # 获得会话变量 @@ -122,31 +125,34 @@ class ProviderDashscope(ProviderOpenAIOfficial): ) response = await asyncio.get_event_loop().run_in_executor(None, partial) + assert isinstance(response, ApplicationResponse) + logger.debug(f"dashscope resp: {response}") if response.status_code != 200: logger.error( - f"阿里云百炼请求失败: request_id={response.request_id}, code={response.status_code}, message={response.message}, 请参考文档：https://help.aliyun.com/zh/model-studio/developer-reference/error-code" + f"阿里云百炼请求失败: request_id={response.request_id}, code={response.status_code}, message={response.message}, 请参考文档：https://help.aliyun.com/zh/model-studio/developer-reference/error-code", ) return LLMResponse( role="err", result_chain=MessageChain().message( - f"阿里云百炼请求失败: message={response.message} code={response.status_code}" + f"阿里云百炼请求失败: message={response.message} code={response.status_code}", ), ) - output_text = response.output.get("text", "") + output_text = response.output.get("text", "") or "" # RAG 引用脚标格式化 output_text = re.sub(r"\[(\d+)\]", r"[\1]", output_text) if self.output_reference and response.output.get("doc_references", None): - ref_str = "" - for ref in response.output.get("doc_references", []): + ref_parts = [] + for ref in response.output.get("doc_references", []) or []: ref_title = ( ref.get("title", "") if ref.get("title") else ref.get("doc_name", "") ) - ref_str += f"{ref['index_id']}. {ref_title}\n" + ref_parts.append(f"{ref['index_id']}. {ref_title}\n") + ref_str = "".join(ref_parts) output_text += f"\n\n回答来源:\n{ref_str}" llm_response = LLMResponse("assistant") diff --git a/astrbot/core/provider/sources/dashscope_tts.py b/astrbot/core/provider/sources/dashscope_tts.py index efda31ca..44e9965c 100644 --- a/astrbot/core/provider/sources/dashscope_tts.py +++ b/astrbot/core/provider/sources/dashscope_tts.py @@ -3,7 +3,7 @@ import base64 import logging import os import uuid -from typing import Optional, Tuple + import aiohttp import dashscope from dashscope.audio.tts_v2 import AudioFormat, SpeechSynthesizer @@ -15,14 +15,17 @@ except ( ): # pragma: no cover - older dashscope versions without Qwen TTS support MultiModalConversation = None +from astrbot.core.utils.astrbot_path import get_astrbot_data_path + from ..entities import ProviderType from ..provider import TTSProvider from ..register import register_provider_adapter -from astrbot.core.utils.astrbot_path import get_astrbot_data_path @register_provider_adapter( - "dashscope_tts", "Dashscope TTS API", provider_type=ProviderType.TEXT_TO_SPEECH + "dashscope_tts", + "Dashscope TTS API", + provider_type=ProviderType.TEXT_TO_SPEECH, ) class ProviderDashscopeTTSAPI(TTSProvider): def __init__( @@ -33,7 +36,7 @@ class ProviderDashscopeTTSAPI(TTSProvider): super().__init__(provider_config, provider_settings) self.chosen_api_key: str = provider_config.get("api_key", "") self.voice: str = provider_config.get("dashscope_tts_voice", "loongstella") - self.set_model(provider_config.get("model", None)) + self.set_model(provider_config.get("model")) self.timeout_ms = float(provider_config.get("timeout", 20)) * 1000 dashscope.api_key = self.chosen_api_key @@ -52,7 +55,7 @@ class ProviderDashscopeTTSAPI(TTSProvider): if not audio_bytes: raise RuntimeError( - "Audio synthesis failed, returned empty content. The model may not be supported or the service is unavailable." + "Audio synthesis failed, returned empty content. The model may not be supported or the service is unavailable.", ) path = os.path.join(temp_dir, f"dashscope_tts_{uuid.uuid4()}{ext}") @@ -63,7 +66,7 @@ class ProviderDashscopeTTSAPI(TTSProvider): def _call_qwen_tts(self, model: str, text: str): if MultiModalConversation is None: raise RuntimeError( - "dashscope SDK missing MultiModalConversation. Please upgrade the dashscope package to use Qwen TTS models." + "dashscope SDK missing MultiModalConversation. Please upgrade the dashscope package to use Qwen TTS models.", ) kwargs = { @@ -74,24 +77,26 @@ class ProviderDashscopeTTSAPI(TTSProvider): } if not self.voice: logging.warning( - "No voice specified for Qwen TTS model, using default 'Cherry'." + "No voice specified for Qwen TTS model, using default 'Cherry'.", ) return MultiModalConversation.call(**kwargs) async def _synthesize_with_qwen_tts( - self, model: str, text: str - ) -> Tuple[Optional[bytes], str]: + self, + model: str, + text: str, + ) -> tuple[bytes | None, str]: loop = asyncio.get_event_loop() response = await loop.run_in_executor(None, self._call_qwen_tts, model, text) audio_bytes = await self._extract_audio_from_response(response) if not audio_bytes: raise RuntimeError( - f"Audio synthesis failed for model '{model}'. {response}" + f"Audio synthesis failed for model '{model}'. {response}", ) ext = ".wav" return audio_bytes, ext - async def _extract_audio_from_response(self, response) -> Optional[bytes]: + async def _extract_audio_from_response(self, response) -> bytes | None: output = getattr(response, "output", None) audio_obj = getattr(output, "audio", None) if output is not None else None if not audio_obj: @@ -102,7 +107,7 @@ class ProviderDashscopeTTSAPI(TTSProvider): try: return base64.b64decode(data_b64) except (ValueError, TypeError): - logging.error("Failed to decode base64 audio data.") + logging.exception("Failed to decode base64 audio data.") return None url = getattr(audio_obj, "url", None) @@ -110,23 +115,28 @@ class ProviderDashscopeTTSAPI(TTSProvider): return await self._download_audio_from_url(url) return None - async def _download_audio_from_url(self, url: str) -> Optional[bytes]: + async def _download_audio_from_url(self, url: str) -> bytes | None: if not url: return None timeout = max(self.timeout_ms / 1000, 1) if self.timeout_ms else 20 try: - async with aiohttp.ClientSession() as session: - async with session.get( - url, timeout=aiohttp.ClientTimeout(total=timeout) - ) as response: - return await response.read() + async with ( + aiohttp.ClientSession() as session, + session.get( + url, + timeout=aiohttp.ClientTimeout(total=timeout), + ) as response, + ): + return await response.read() except (aiohttp.ClientError, asyncio.TimeoutError, OSError) as e: - logging.error(f"Failed to download audio from URL {url}: {e}") + logging.exception(f"Failed to download audio from URL {url}: {e}") return None async def _synthesize_with_cosyvoice( - self, model: str, text: str - ) -> Tuple[Optional[bytes], str]: + self, + model: str, + text: str, + ) -> tuple[bytes | None, str]: synthesizer = SpeechSynthesizer( model=model, voice=self.voice, @@ -134,13 +144,16 @@ class ProviderDashscopeTTSAPI(TTSProvider): ) loop = asyncio.get_event_loop() audio_bytes = await loop.run_in_executor( - None, synthesizer.call, text, self.timeout_ms + None, + synthesizer.call, + text, + self.timeout_ms, ) if not audio_bytes: resp = synthesizer.get_response() if resp and isinstance(resp, dict): raise RuntimeError( - f"Audio synthesis failed for model '{model}'. {resp}".strip() + f"Audio synthesis failed for model '{model}'. {resp}".strip(), ) return audio_bytes, ".wav" diff --git a/astrbot/core/provider/sources/dify_source.py b/astrbot/core/provider/sources/dify_source.py index e19e912a..7850a982 100644 --- a/astrbot/core/provider/sources/dify_source.py +++ b/astrbot/core/provider/sources/dify_source.py @@ -1,15 +1,15 @@ -import astrbot.core.message.components as Comp import os -from typing import List -from .. import Provider -from ..entities import LLMResponse -from ..func_tool_manager import FuncCall -from ..register import register_provider_adapter -from astrbot.core.utils.dify_api_client import DifyAPIClient -from astrbot.core.utils.io import download_image_by_url, download_file + +import astrbot.core.message.components as Comp from astrbot.core import logger, sp from astrbot.core.message.message_event_result import MessageChain from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.dify_api_client import DifyAPIClient +from astrbot.core.utils.io import download_file, download_image_by_url + +from .. import Provider +from ..entities import LLMResponse +from ..register import register_provider_adapter @register_provider_adapter("dify", "Dify APP 适配器。") @@ -18,12 +18,10 @@ class ProviderDify(Provider): self, provider_config, provider_settings, - default_persona=None, ) -> None: super().__init__( provider_config, provider_settings, - default_persona, ) self.api_key = provider_config.get("dify_api_key", "") if not self.api_key: @@ -34,10 +32,12 @@ class ProviderDify(Provider): raise Exception("Dify API 类型不能为空。") self.model_name = "dify" self.workflow_output_key = provider_config.get( - "dify_workflow_output_key", "astrbot_wf_output" + "dify_workflow_output_key", + "astrbot_wf_output", ) self.dify_query_input_key = provider_config.get( - "dify_query_input_key", "astrbot_text_query" + "dify_query_input_key", + "astrbot_text_query", ) if not self.dify_query_input_key: self.dify_query_input_key = "astrbot_text_query" @@ -55,11 +55,11 @@ class ProviderDify(Provider): async def text_chat( self, prompt: str, - session_id: str = None, - image_urls: List[str] = None, - func_tool: FuncCall = None, - contexts: List = None, - system_prompt: str = None, + session_id=None, + image_urls=None, + func_tool=None, + contexts=None, + system_prompt=None, tool_calls_result=None, model=None, **kwargs, @@ -78,12 +78,13 @@ class ProviderDify(Provider): else image_url ) file_response = await self.api_client.file_upload( - image_path, user=session_id + image_path, + user=session_id, ) logger.debug(f"Dify 上传图片响应：{file_response}") if "id" not in file_response: logger.warning( - f"上传图片后得到未知的 Dify 响应：{file_response}，图片将忽略。" + f"上传图片后得到未知的 Dify 响应：{file_response}，图片将忽略。", ) continue files_payload.append( @@ -91,7 +92,7 @@ class ProviderDify(Provider): "type": "image", "transfer_method": "local_file", "upload_file_id": file_response["id"], - } + }, ) # 获得会话变量 @@ -134,7 +135,7 @@ class ProviderDify(Provider): elif chunk["event"] == "error": logger.error(f"Dify 出现错误：{chunk}") raise Exception( - f"Dify 出现错误 status: {chunk['status']} message: {chunk['message']}" + f"Dify 出现错误 status: {chunk['status']} message: {chunk['message']}", ) case "workflow": @@ -151,37 +152,37 @@ class ProviderDify(Provider): match chunk["event"]: case "workflow_started": logger.info( - f"Dify 工作流(ID: {chunk['workflow_run_id']})开始运行。" + f"Dify 工作流(ID: {chunk['workflow_run_id']})开始运行。", ) case "node_finished": logger.debug( - f"Dify 工作流节点(ID: {chunk['data']['node_id']} Title: {chunk['data'].get('title', '')})运行结束。" + f"Dify 工作流节点(ID: {chunk['data']['node_id']} Title: {chunk['data'].get('title', '')})运行结束。", ) case "workflow_finished": logger.info( - f"Dify 工作流(ID: {chunk['workflow_run_id']})运行结束" + f"Dify 工作流(ID: {chunk['workflow_run_id']})运行结束", ) logger.debug(f"Dify 工作流结果：{chunk}") if chunk["data"]["error"]: logger.error( - f"Dify 工作流出现错误：{chunk['data']['error']}" + f"Dify 工作流出现错误：{chunk['data']['error']}", ) raise Exception( - f"Dify 工作流出现错误：{chunk['data']['error']}" + f"Dify 工作流出现错误：{chunk['data']['error']}", ) if ( self.workflow_output_key not in chunk["data"]["outputs"] ): raise Exception( - f"Dify 工作流的输出不包含指定的键名：{self.workflow_output_key}" + f"Dify 工作流的输出不包含指定的键名：{self.workflow_output_key}", ) result = chunk case _: raise Exception(f"未知的 Dify API 类型：{self.api_type}") except Exception as e: - logger.error(f"Dify 请求失败：{str(e)}") - return LLMResponse(role="err", completion_text=f"Dify 请求失败：{str(e)}") + logger.error(f"Dify 请求失败：{e!s}") + return LLMResponse(role="err", completion_text=f"Dify 请求失败：{e!s}") if not result: logger.warning("Dify 请求结果为空，请查看 Debug 日志。") @@ -223,7 +224,7 @@ class ProviderDify(Provider): # Chat return MessageChain(chain=[Comp.Plain(chunk)]) - async def parse_file(item: dict) -> Comp: + async def parse_file(item: dict): match item["type"]: case "image": return Comp.Image(file=item["url"], url=item["url"]) diff --git a/astrbot/core/provider/sources/edge_tts_source.py b/astrbot/core/provider/sources/edge_tts_source.py index 44c2d175..8bbf6232 100644 --- a/astrbot/core/provider/sources/edge_tts_source.py +++ b/astrbot/core/provider/sources/edge_tts_source.py @@ -1,14 +1,17 @@ -import uuid -import os -import edge_tts -import subprocess import asyncio -from ..provider import TTSProvider -from ..entities import ProviderType -from ..register import register_provider_adapter +import os +import subprocess +import uuid + +import edge_tts + from astrbot.core import logger from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from ..entities import ProviderType +from ..provider import TTSProvider +from ..register import register_provider_adapter + """ edge_tts 方式，能够免费、快速生成语音，使用需要先安装edge-tts库 ``` @@ -19,7 +22,9 @@ Windows 如果提示找不到指定文件，以管理员身份运行命令行窗 @register_provider_adapter( - "edge_tts", "Microsoft Edge TTS", provider_type=ProviderType.TEXT_TO_SPEECH + "edge_tts", + "Microsoft Edge TTS", + provider_type=ProviderType.TEXT_TO_SPEECH, ) class ProviderEdgeTTS(TTSProvider): def __init__( @@ -31,9 +36,9 @@ class ProviderEdgeTTS(TTSProvider): # 设置默认语音，如果没有指定则使用中文小萱 self.voice = provider_config.get("edge-tts-voice", "zh-CN-XiaoxiaoNeural") - self.rate = provider_config.get("rate", None) - self.volume = provider_config.get("volume", None) - self.pitch = provider_config.get("pitch", None) + self.rate = provider_config.get("rate") + self.volume = provider_config.get("volume") + self.pitch = provider_config.get("pitch") self.timeout = provider_config.get("timeout", 30) self.proxy = os.getenv("https_proxy", None) @@ -97,26 +102,25 @@ class ProviderEdgeTTS(TTSProvider): os.remove(mp3_path) if os.path.exists(wav_path) and os.path.getsize(wav_path) > 0: return wav_path - else: - logger.error("生成的WAV文件不存在或为空") - raise RuntimeError("生成的WAV文件不存在或为空") + logger.error("生成的WAV文件不存在或为空") + raise RuntimeError("生成的WAV文件不存在或为空") except subprocess.CalledProcessError as e: logger.error( - f"FFmpeg 转换失败: {e.stderr.decode() if e.stderr else str(e)}" + f"FFmpeg 转换失败: {e.stderr.decode() if e.stderr else str(e)}", ) try: if os.path.exists(mp3_path): os.remove(mp3_path) except Exception: pass - raise RuntimeError(f"FFmpeg 转换失败: {str(e)}") + raise RuntimeError(f"FFmpeg 转换失败: {e!s}") except Exception as e: - logger.error(f"音频生成失败: {str(e)}") + logger.error(f"音频生成失败: {e!s}") try: if os.path.exists(mp3_path): os.remove(mp3_path) except Exception: pass - raise RuntimeError(f"音频生成失败: {str(e)}") + raise RuntimeError(f"音频生成失败: {e!s}") diff --git a/astrbot/core/provider/sources/fishaudio_tts_api_source.py b/astrbot/core/provider/sources/fishaudio_tts_api_source.py index 49c78239..ca571c3e 100644 --- a/astrbot/core/provider/sources/fishaudio_tts_api_source.py +++ b/astrbot/core/provider/sources/fishaudio_tts_api_source.py @@ -1,15 +1,18 @@ import os -import uuid import re -import ormsgpack -from pydantic import BaseModel, conint -from httpx import AsyncClient +import uuid from typing import Annotated, Literal -from ..provider import TTSProvider -from ..entities import ProviderType -from ..register import register_provider_adapter + +import ormsgpack +from httpx import AsyncClient +from pydantic import BaseModel, conint + from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from ..entities import ProviderType +from ..provider import TTSProvider +from ..register import register_provider_adapter + class ServeReferenceAudio(BaseModel): audio: bytes @@ -35,7 +38,9 @@ class ServeTTSRequest(BaseModel): @register_provider_adapter( - "fishaudio_tts_api", "FishAudio TTS API", provider_type=ProviderType.TEXT_TO_SPEECH + "fishaudio_tts_api", + "FishAudio TTS API", + provider_type=ProviderType.TEXT_TO_SPEECH, ) class ProviderFishAudioTTSAPI(TTSProvider): def __init__( @@ -48,16 +53,16 @@ class ProviderFishAudioTTSAPI(TTSProvider): self.reference_id: str = provider_config.get("fishaudio-tts-reference-id", "") self.character: str = provider_config.get("fishaudio-tts-character", "可莉") self.api_base: str = provider_config.get( - "api_base", "https://api.fish-audio.cn/v1" + "api_base", + "https://api.fish-audio.cn/v1", ) self.headers = { "Authorization": f"Bearer {self.chosen_api_key}", } - self.set_model(provider_config.get("model", None)) + self.set_model(provider_config.get("model")) async def _get_reference_id_by_character(self, character: str) -> str: - """ - 获取角色的reference_id + """获取角色的reference_id Args: character: 角色名称 @@ -67,13 +72,16 @@ class ProviderFishAudioTTSAPI(TTSProvider): exception: APIException: 获取语音角色列表为空 + """ sort_options = ["score", "task_count", "created_at"] async with AsyncClient(base_url=self.api_base.replace("/v1", "")) as client: for sort_by in sort_options: params = {"title": character, "sort_by": sort_by} response = await client.get( - "/model", params=params, headers=self.headers + "/model", + params=params, + headers=self.headers, ) resp_data = response.json() if resp_data["total"] == 0: @@ -84,14 +92,14 @@ class ProviderFishAudioTTSAPI(TTSProvider): return None def _validate_reference_id(self, reference_id: str) -> bool: - """ - 验证reference_id格式是否有效 + """验证reference_id格式是否有效 Args: reference_id: 参考模型ID Returns: bool: ID是否有效 + """ if not reference_id or not reference_id.strip(): return False @@ -109,7 +117,7 @@ class ProviderFishAudioTTSAPI(TTSProvider): raise ValueError( f"无效的FishAudio参考模型ID: '{self.reference_id}'. " f"请确保ID是32位十六进制字符串（例如: 626bb6d3f3364c9cbc3aa6a67300a664）。" - f"您可以从 https://fish.audio/zh-CN/discovery 获取有效的模型ID。" + f"您可以从 https://fish.audio/zh-CN/discovery 获取有效的模型ID。", ) reference_id = self.reference_id.strip() else: diff --git a/astrbot/core/provider/sources/gemini_embedding_source.py b/astrbot/core/provider/sources/gemini_embedding_source.py index baccf52a..8d11cce5 100644 --- a/astrbot/core/provider/sources/gemini_embedding_source.py +++ b/astrbot/core/provider/sources/gemini_embedding_source.py @@ -1,9 +1,10 @@ from google import genai from google.genai import types from google.genai.errors import APIError + +from ..entities import ProviderType from ..provider import EmbeddingProvider from ..register import register_provider_adapter -from ..entities import ProviderType @register_provider_adapter( @@ -18,41 +19,38 @@ class GeminiEmbeddingProvider(EmbeddingProvider): self.provider_settings = provider_settings api_key: str = provider_config.get("embedding_api_key") - api_base: str = provider_config.get("embedding_api_base", None) + api_base: str = provider_config.get("embedding_api_base") timeout: int = int(provider_config.get("timeout", 20)) http_options = types.HttpOptions(timeout=timeout * 1000) if api_base: - if api_base.endswith("/"): - api_base = api_base[:-1] + api_base = api_base.removesuffix("/") http_options.base_url = api_base self.client = genai.Client(api_key=api_key, http_options=http_options).aio self.model = provider_config.get( - "embedding_model", "gemini-embedding-exp-03-07" + "embedding_model", + "gemini-embedding-exp-03-07", ) - self.dimension = provider_config.get("embedding_dimensions", 768) async def get_embedding(self, text: str) -> list[float]: - """ - 获取文本的嵌入 - """ + """获取文本的嵌入""" try: result = await self.client.models.embed_content( - model=self.model, contents=text + model=self.model, + contents=text, ) return result.embeddings[0].values except APIError as e: raise Exception(f"Gemini Embedding API请求失败: {e.message}") async def get_embeddings(self, texts: list[str]) -> list[list[float]]: - """ - 批量获取文本的嵌入 - """ + """批量获取文本的嵌入""" try: result = await self.client.models.embed_content( - model=self.model, contents=texts + model=self.model, + contents=texts, ) return [embedding.values for embedding in result.embeddings] except APIError as e: @@ -60,4 +58,4 @@ class GeminiEmbeddingProvider(EmbeddingProvider): def get_dim(self) -> int: """获取向量的维度""" - return self.dimension + return self.provider_config.get("embedding_dimensions", 768) diff --git a/astrbot/core/provider/sources/gemini_source.py b/astrbot/core/provider/sources/gemini_source.py index b14a9bdc..fec1fb80 100644 --- a/astrbot/core/provider/sources/gemini_source.py +++ b/astrbot/core/provider/sources/gemini_source.py @@ -3,7 +3,6 @@ import base64 import json import logging import random -from typing import Optional, List from collections.abc import AsyncGenerator from google import genai @@ -32,7 +31,8 @@ logging.getLogger("google_genai.types").addFilter(SuppressNonTextPartsWarning()) @register_provider_adapter( - "googlegenai_chat_completion", "Google Gemini Chat Completion 提供商适配器" + "googlegenai_chat_completion", + "Google Gemini Chat Completion 提供商适配器", ) class ProviderGoogleGenAI(Provider): CATEGORY_MAPPING = { @@ -53,18 +53,16 @@ class ProviderGoogleGenAI(Provider): self, provider_config, provider_settings, - default_persona=None, ) -> None: super().__init__( provider_config, provider_settings, - default_persona, ) - self.api_keys: List = super().get_keys() + self.api_keys: list = super().get_keys() self.chosen_api_key: str = self.api_keys[0] if len(self.api_keys) > 0 else "" self.timeout: int = int(provider_config.get("timeout", 180)) - self.api_base: Optional[str] = provider_config.get("api_base", None) + self.api_base: str | None = provider_config.get("api_base", None) if self.api_base and self.api_base.endswith("/"): self.api_base = self.api_base[:-1] @@ -87,7 +85,8 @@ class ProviderGoogleGenAI(Provider): user_safety_config = self.provider_config.get("gm_safety_settings", {}) self.safety_settings = [ types.SafetySetting( - category=harm_category, threshold=self.THRESHOLD_MAPPING[threshold_str] + category=harm_category, + threshold=self.THRESHOLD_MAPPING[threshold_str], ) for config_key, harm_category in self.CATEGORY_MAPPING.items() if (threshold_str := user_safety_config.get(config_key)) @@ -104,27 +103,25 @@ class ProviderGoogleGenAI(Provider): if len(keys) > 0: self.set_key(random.choice(keys)) logger.info( - f"检测到 Key 异常({e.message})，正在尝试更换 API Key 重试... 当前 Key: {self.chosen_api_key[:12]}..." + f"检测到 Key 异常({e.message})，正在尝试更换 API Key 重试... 当前 Key: {self.chosen_api_key[:12]}...", ) await asyncio.sleep(1) return True - else: - logger.error( - f"检测到 Key 异常({e.message})，且已没有可用的 Key。 当前 Key: {self.chosen_api_key[:12]}..." - ) - raise Exception("达到了 Gemini 速率限制, 请稍后再试...") - else: logger.error( - f"发生了错误(gemini_source)。Provider 配置如下: {self.provider_config}" + f"检测到 Key 异常({e.message})，且已没有可用的 Key。 当前 Key: {self.chosen_api_key[:12]}...", ) - raise e + raise Exception("达到了 Gemini 速率限制, 请稍后再试...") + logger.error( + f"发生了错误(gemini_source)。Provider 配置如下: {self.provider_config}", + ) + raise e async def _prepare_query_config( self, payloads: dict, - tools: Optional[ToolSet] = None, - system_instruction: Optional[str] = None, - modalities: Optional[list[str]] = None, + tools: ToolSet | None = None, + system_instruction: str | None = None, + modalities: list[str] | None = None, temperature: float = 0.7, ) -> types.GenerateContentConfig: """准备查询配置""" @@ -152,7 +149,7 @@ class ProviderGoogleGenAI(Provider): logger.warning("代码执行工具与搜索工具互斥，已忽略搜索工具") if url_context: logger.warning( - "代码执行工具与URL上下文工具互斥，已忽略URL上下文工具" + "代码执行工具与URL上下文工具互斥，已忽略URL上下文工具", ) else: if native_search: @@ -163,13 +160,13 @@ class ProviderGoogleGenAI(Provider): tool_list.append(types.Tool(url_context=types.UrlContext())) else: logger.warning( - "当前 SDK 版本不支持 URL 上下文工具，已忽略该设置，请升级 google-genai 包" + "当前 SDK 版本不支持 URL 上下文工具，已忽略该设置，请升级 google-genai 包", ) elif "gemini-2.0-lite" in model_name: if native_coderunner or native_search or url_context: logger.warning( - "gemini-2.0-lite 不支持代码执行、搜索工具和URL上下文，将忽略这些设置" + "gemini-2.0-lite 不支持代码执行、搜索工具和URL上下文，将忽略这些设置", ) tool_list = None @@ -186,7 +183,7 @@ class ProviderGoogleGenAI(Provider): tool_list.append(types.Tool(url_context=types.UrlContext())) else: logger.warning( - "当前 SDK 版本不支持 URL 上下文工具，已忽略该设置，请升级 google-genai 包" + "当前 SDK 版本不支持 URL 上下文工具，已忽略该设置，请升级 google-genai 包", ) if not tool_list: @@ -196,7 +193,7 @@ class ProviderGoogleGenAI(Provider): logger.warning("已启用原生工具，函数工具将被忽略") elif tools and (func_desc := tools.get_func_desc_google_genai_style()): tool_list = [ - types.Tool(function_declarations=func_desc["function_declarations"]) + types.Tool(function_declarations=func_desc["function_declarations"]), ] return types.GenerateContentConfig( @@ -223,8 +220,9 @@ class ProviderGoogleGenAI(Provider): thinking_budget=min( int( self.provider_config.get("gm_thinking_config", {}).get( - "budget", 0 - ) + "budget", + 0, + ), ), 24576, ), @@ -234,7 +232,7 @@ class ProviderGoogleGenAI(Provider): else None ), automatic_function_calling=types.AutomaticFunctionCallingConfig( - disable=True + disable=True, ), ) @@ -268,7 +266,7 @@ class ProviderGoogleGenAI(Provider): [ self.provider_config.get("gm_native_coderunner", False), self.provider_config.get("gm_native_search", False), - ] + ], ) for message in payloads["messages"]: role, content = message["role"], message.get("content") @@ -292,19 +290,30 @@ class ProviderGoogleGenAI(Provider): parts = [types.Part.from_text(text=content)] append_or_extend(gemini_contents, parts, types.ModelContent) elif not native_tool_enabled and "tool_calls" in message: - parts = [ - types.Part.from_function_call( + parts = [] + for tool in message["tool_calls"]: + part = types.Part.from_function_call( name=tool["function"]["name"], args=json.loads(tool["function"]["arguments"]), ) - for tool in message["tool_calls"] - ] + # we should set thought_signature back to part if exists + # for more info about thought_signature, see: + # https://ai.google.dev/gemini-api/docs/thought-signatures + if "extra_content" in tool and tool["extra_content"]: + ts_bs64 = ( + tool["extra_content"] + .get("google", {}) + .get("thought_signature") + ) + if ts_bs64: + part.thought_signature = base64.b64decode(ts_bs64) + parts.append(part) append_or_extend(gemini_contents, parts, types.ModelContent) else: logger.warning("assistant 角色的消息内容为空，已添加空格占位") if native_tool_enabled and "tool_calls" in message: logger.warning( - "检测到启用Gemini原生工具，且上下文中存在函数调用，建议使用 /reset 重置上下文" + "检测到启用Gemini原生工具，且上下文中存在函数调用，建议使用 /reset 重置上下文", ) parts = [types.Part.from_text(text=" ")] append_or_extend(gemini_contents, parts, types.ModelContent) @@ -317,7 +326,7 @@ class ProviderGoogleGenAI(Provider): "name": message["tool_call_id"], "content": message["content"], }, - ) + ), ] append_or_extend(gemini_contents, parts, types.UserContent) @@ -326,9 +335,20 @@ class ProviderGoogleGenAI(Provider): return gemini_contents - @staticmethod + def _extract_reasoning_content(self, candidate: types.Candidate) -> str: + """Extract reasoning content from candidate parts""" + if not candidate.content or not candidate.content.parts: + return "" + + thought_buf: list[str] = [ + (p.text or "") for p in candidate.content.parts if p.thought + ] + return "".join(thought_buf).strip() + def _process_content_parts( - candidate: types.Candidate, llm_response: LLMResponse + self, + candidate: types.Candidate, + llm_response: LLMResponse, ) -> MessageChain: """处理内容部分并构建消息链""" if not candidate.content: @@ -357,6 +377,11 @@ class ProviderGoogleGenAI(Provider): logger.warning(f"收到的 candidate.content.parts 为空: {candidate}") raise Exception("API 返回的 candidate.content.parts 为空。") + # 提取 reasoning content + reasoning = self._extract_reasoning_content(candidate) + if reasoning: + llm_response.reasoning_content = reasoning + chain = [] part: types.Part @@ -379,10 +404,15 @@ class ProviderGoogleGenAI(Provider): llm_response.role = "tool" llm_response.tools_call_name.append(part.function_call.name) llm_response.tools_call_args.append(part.function_call.args) - # gemini 返回的 function_call.id 可能为 None - llm_response.tools_call_ids.append( - part.function_call.id or part.function_call.name - ) + # function_call.id might be None, use name as fallback + tool_call_id = part.function_call.id or part.function_call.name + llm_response.tools_call_ids.append(tool_call_id) + # extra_content + if part.thought_signature: + ts_bs64 = base64.b64encode(part.thought_signature).decode("utf-8") + llm_response.tools_call_extra_content[tool_call_id] = { + "google": {"thought_signature": ts_bs64} + } elif ( part.inline_data and part.inline_data.mime_type @@ -406,17 +436,22 @@ class ProviderGoogleGenAI(Provider): conversation = self._prepare_conversation(payloads) temperature = payloads.get("temperature", 0.7) - result: Optional[types.GenerateContentResponse] = None + result: types.GenerateContentResponse | None = None while True: try: config = await self._prepare_query_config( - payloads, tools, system_instruction, modalities, temperature + payloads, + tools, + system_instruction, + modalities, + temperature, ) result = await self.client.models.generate_content( model=self.get_model(), contents=conversation, config=config, ) + logger.debug(f"genai result: {result}") if not result.candidates: logger.error(f"请求失败, 返回的 candidates 为空: {result}") @@ -427,7 +462,7 @@ class ProviderGoogleGenAI(Provider): raise Exception("温度参数已超过最大值2，仍然发生recitation") temperature += 0.2 logger.warning( - f"发生了recitation，正在提高温度至{temperature:.1f}重试..." + f"发生了recitation，正在提高温度至{temperature:.1f}重试...", ) continue @@ -438,7 +473,7 @@ class ProviderGoogleGenAI(Provider): e.message = "" if "Developer instruction is not enabled" in e.message: logger.warning( - f"{self.get_model()} 不支持 system prompt，已自动去除(影响人格设置)" + f"{self.get_model()} 不支持 system prompt，已自动去除(影响人格设置)", ) system_instruction = None elif "Function calling is not enabled" in e.message: @@ -451,7 +486,7 @@ class ProviderGoogleGenAI(Provider): or "only supports text output" in e.message ): logger.warning( - f"{self.get_model()} 不支持多模态输出，降级为文本模态" + f"{self.get_model()} 不支持多模态输出，降级为文本模态", ) modalities = ["Text"] else: @@ -461,12 +496,15 @@ class ProviderGoogleGenAI(Provider): llm_response = LLMResponse("assistant") llm_response.raw_completion = result llm_response.result_chain = self._process_content_parts( - result.candidates[0], llm_response + result.candidates[0], + llm_response, ) return llm_response async def _query_stream( - self, payloads: dict, tools: ToolSet | None + self, + payloads: dict, + tools: ToolSet | None, ) -> AsyncGenerator[LLMResponse, None]: """流式请求 Gemini API""" system_instruction = next( @@ -480,7 +518,9 @@ class ProviderGoogleGenAI(Provider): while True: try: config = await self._prepare_query_config( - payloads, tools, system_instruction + payloads, + tools, + system_instruction, ) result = await self.client.models.generate_content_stream( model=self.get_model(), @@ -493,7 +533,7 @@ class ProviderGoogleGenAI(Provider): e.message = "" if "Developer instruction is not enabled" in e.message: logger.warning( - f"{self.get_model()} 不支持 system prompt，已自动去除(影响人格设置)" + f"{self.get_model()} 不支持 system prompt，已自动去除(影响人格设置)", ) system_instruction = None elif "Function calling is not enabled" in e.message: @@ -505,6 +545,7 @@ class ProviderGoogleGenAI(Provider): # Accumulate the complete response text for the final response accumulated_text = "" + accumulated_reasoning = "" final_response = None async for chunk in result: @@ -523,14 +564,25 @@ class ProviderGoogleGenAI(Provider): llm_response = LLMResponse("assistant", is_chunk=False) llm_response.raw_completion = chunk llm_response.result_chain = self._process_content_parts( - chunk.candidates[0], llm_response + chunk.candidates[0], + llm_response, ) yield llm_response return + _f = False + + # 提取 reasoning content + reasoning = self._extract_reasoning_content(chunk.candidates[0]) + if reasoning: + _f = True + accumulated_reasoning += reasoning + llm_response.reasoning_content = reasoning if chunk.text: + _f = True accumulated_text += chunk.text llm_response.result_chain = MessageChain(chain=[Comp.Plain(chunk.text)]) + if _f: yield llm_response if chunk.candidates[0].finish_reason: @@ -539,7 +591,8 @@ class ProviderGoogleGenAI(Provider): final_response = LLMResponse("assistant", is_chunk=False) final_response.raw_completion = chunk final_response.result_chain = self._process_content_parts( - chunk.candidates[0], final_response + chunk.candidates[0], + final_response, ) break @@ -547,10 +600,14 @@ class ProviderGoogleGenAI(Provider): if not final_response: final_response = LLMResponse("assistant", is_chunk=False) + # Set the complete accumulated reasoning in the final response + if accumulated_reasoning: + final_response.reasoning_content = accumulated_reasoning + # Set the complete accumulated text in the final response if accumulated_text: final_response.result_chain = MessageChain( - chain=[Comp.Plain(accumulated_text)] + chain=[Comp.Plain(accumulated_text)], ) elif not final_response.result_chain: # If no text was accumulated and no final response was set, provide empty space @@ -560,7 +617,7 @@ class ProviderGoogleGenAI(Provider): async def text_chat( self, - prompt: str, + prompt=None, session_id=None, image_urls=None, func_tool=None, @@ -572,8 +629,12 @@ class ProviderGoogleGenAI(Provider): ) -> LLMResponse: if contexts is None: contexts = [] - new_record = await self.assemble_context(prompt, image_urls) - context_query = [*contexts, new_record] + new_record = None + if prompt is not None: + new_record = await self.assemble_context(prompt, image_urls) + context_query = self._ensure_message_to_dicts(contexts) + if new_record: + context_query.append(new_record) if system_prompt: context_query.insert(0, {"role": "system", "content": system_prompt}) @@ -609,7 +670,7 @@ class ProviderGoogleGenAI(Provider): async def text_chat_stream( self, - prompt, + prompt=None, session_id=None, image_urls=None, func_tool=None, @@ -621,8 +682,12 @@ class ProviderGoogleGenAI(Provider): ) -> AsyncGenerator[LLMResponse, None]: if contexts is None: contexts = [] - new_record = await self.assemble_context(prompt, image_urls) - context_query = [*contexts, new_record] + new_record = None + if prompt is not None: + new_record = await self.assemble_context(prompt, image_urls) + context_query = self._ensure_message_to_dicts(contexts) + if new_record: + context_query.append(new_record) if system_prompt: context_query.insert(0, {"role": "system", "content": system_prompt}) @@ -680,9 +745,7 @@ class ProviderGoogleGenAI(Provider): self._init_client() async def assemble_context(self, text: str, image_urls: list[str] | None = None): - """ - 组装上下文。 - """ + """组装上下文。""" if image_urls: user_content = { "role": "user", @@ -704,22 +767,18 @@ class ProviderGoogleGenAI(Provider): { "type": "image_url", "image_url": {"url": image_data}, - } + }, ) return user_content - else: - return {"role": "user", "content": text} + return {"role": "user", "content": text} async def encode_image_bs64(self, image_url: str) -> str: - """ - 将图片转换为 base64 - """ + """将图片转换为 base64""" if image_url.startswith("base64://"): return image_url.replace("base64://", "data:image/jpeg;base64,") with open(image_url, "rb") as f: image_bs64 = base64.b64encode(f.read()).decode("utf-8") return "data:image/jpeg;base64," + image_bs64 - return "" async def terminate(self): logger.info("Google GenAI 适配器已终止。") diff --git a/astrbot/core/provider/sources/gemini_tts_source.py b/astrbot/core/provider/sources/gemini_tts_source.py index 48cb4833..0bf92b32 100644 --- a/astrbot/core/provider/sources/gemini_tts_source.py +++ b/astrbot/core/provider/sources/gemini_tts_source.py @@ -13,7 +13,9 @@ from ..register import register_provider_adapter @register_provider_adapter( - "gemini_tts", "Gemini TTS API", provider_type=ProviderType.TEXT_TO_SPEECH + "gemini_tts", + "Gemini TTS API", + provider_type=ProviderType.TEXT_TO_SPEECH, ) class ProviderGeminiTTSAPI(TTSProvider): def __init__( @@ -28,13 +30,13 @@ class ProviderGeminiTTSAPI(TTSProvider): http_options = types.HttpOptions(timeout=timeout * 1000) if api_base: - if api_base.endswith("/"): - api_base = api_base[:-1] + api_base = api_base.removesuffix("/") http_options.base_url = api_base self.client = genai.Client(api_key=api_key, http_options=http_options).aio self.model: str = provider_config.get( - "gemini_tts_model", "gemini-2.5-flash-preview-tts" + "gemini_tts_model", + "gemini-2.5-flash-preview-tts", ) self.prefix: str | None = provider_config.get( "gemini_tts_prefix", @@ -54,8 +56,8 @@ class ProviderGeminiTTSAPI(TTSProvider): voice_config=types.VoiceConfig( prebuilt_voice_config=types.PrebuiltVoiceConfig( voice_name=self.voice_name, - ) - ) + ), + ), ), ), ) diff --git a/astrbot/core/provider/sources/groq_source.py b/astrbot/core/provider/sources/groq_source.py new file mode 100644 index 00000000..fcc8f238 --- /dev/null +++ b/astrbot/core/provider/sources/groq_source.py @@ -0,0 +1,15 @@ +from ..register import register_provider_adapter +from .openai_source import ProviderOpenAIOfficial + + +@register_provider_adapter( + "groq_chat_completion", "Groq Chat Completion Provider Adapter" +) +class ProviderGroq(ProviderOpenAIOfficial): + def __init__( + self, + provider_config: dict, + provider_settings: dict, + ) -> None: + super().__init__(provider_config, provider_settings) + self.reasoning_key = "reasoning" diff --git a/astrbot/core/provider/sources/gsv_selfhosted_source.py b/astrbot/core/provider/sources/gsv_selfhosted_source.py index 6c4d872a..7f8d39ea 100644 --- a/astrbot/core/provider/sources/gsv_selfhosted_source.py +++ b/astrbot/core/provider/sources/gsv_selfhosted_source.py @@ -3,12 +3,14 @@ import os import uuid import aiohttp -from ..provider import TTSProvider -from ..entities import ProviderType -from ..register import register_provider_adapter + from astrbot import logger from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from ..entities import ProviderType +from ..provider import TTSProvider +from ..register import register_provider_adapter + @register_provider_adapter( provider_type_name="gsv_tts_selfhost", @@ -24,7 +26,7 @@ class ProviderGSVTTS(TTSProvider): super().__init__(provider_config, provider_settings) self.api_base = provider_config.get("api_base", "http://127.0.0.1:9880").rstrip( - "/" + "/", ) self.gpt_weights_path: str = provider_config.get("gpt_weights_path", "") self.sovits_weights_path: str = provider_config.get("sovits_weights_path", "") @@ -40,7 +42,7 @@ class ProviderGSVTTS(TTSProvider): async def initialize(self): """异步初始化：在 ProviderManager 中被调用""" self._session = aiohttp.ClientSession( - timeout=aiohttp.ClientTimeout(total=self.timeout) + timeout=aiohttp.ClientTimeout(total=self.timeout), ) try: await self._set_model_weights() @@ -52,12 +54,15 @@ class ProviderGSVTTS(TTSProvider): def get_session(self) -> aiohttp.ClientSession: if not self._session or self._session.closed: raise RuntimeError( - "[GSV TTS] Provider HTTP session is not ready or closed." + "[GSV TTS] Provider HTTP session is not ready or closed.", ) return self._session async def _make_request( - self, endpoint: str, params=None, retries: int = 3 + self, + endpoint: str, + params=None, + retries: int = 3, ) -> bytes | None: """发起请求""" for attempt in range(retries): @@ -67,13 +72,13 @@ class ProviderGSVTTS(TTSProvider): if response.status != 200: error_text = await response.text() raise Exception( - f"[GSV TTS] Request to {endpoint} failed with status {response.status}: {error_text}" + f"[GSV TTS] Request to {endpoint} failed with status {response.status}: {error_text}", ) return await response.read() except Exception as e: if attempt < retries - 1: logger.warning( - f"[GSV TTS] 请求 {endpoint} 第 {attempt + 1} 次失败：{e}，重试中..." + f"[GSV TTS] 请求 {endpoint} 第 {attempt + 1} 次失败：{e}，重试中...", ) await asyncio.sleep(1) else: @@ -98,7 +103,7 @@ class ProviderGSVTTS(TTSProvider): {"weights_path": self.sovits_weights_path}, ) logger.info( - f"[GSV TTS] 成功设置 SoVITS 模型路径：{self.sovits_weights_path}" + f"[GSV TTS] 成功设置 SoVITS 模型路径：{self.sovits_weights_path}", ) else: logger.info("[GSV TTS] SoVITS 模型路径未配置，将使用内置 SoVITS 模型") @@ -127,12 +132,10 @@ class ProviderGSVTTS(TTSProvider): with open(path, "wb") as f: f.write(result) return path - else: - raise Exception(f"[GSV TTS] 合成失败，输入文本：{text}，错误信息：{result}") + raise Exception(f"[GSV TTS] 合成失败，输入文本：{text}，错误信息：{result}") def build_synthesis_params(self, text: str) -> dict: - """ - 构建语音合成所需的参数字典。 + """构建语音合成所需的参数字典。 当前仅包含默认参数 + 文本，未来可在此基础上动态添加如情绪、角色等语义控制字段。 """ diff --git a/astrbot/core/provider/sources/gsvi_tts_source.py b/astrbot/core/provider/sources/gsvi_tts_source.py index c2444819..d8b17171 100644 --- a/astrbot/core/provider/sources/gsvi_tts_source.py +++ b/astrbot/core/provider/sources/gsvi_tts_source.py @@ -1,15 +1,20 @@ import os -import uuid -import aiohttp import urllib.parse -from ..provider import TTSProvider -from ..entities import ProviderType -from ..register import register_provider_adapter +import uuid + +import aiohttp + from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from ..entities import ProviderType +from ..provider import TTSProvider +from ..register import register_provider_adapter + @register_provider_adapter( - "gsvi_tts_api", "GSVI TTS API", provider_type=ProviderType.TEXT_TO_SPEECH + "gsvi_tts_api", + "GSVI TTS API", + provider_type=ProviderType.TEXT_TO_SPEECH, ) class ProviderGSVITTS(TTSProvider): def __init__( @@ -19,8 +24,7 @@ class ProviderGSVITTS(TTSProvider): ) -> None: super().__init__(provider_config, provider_settings) self.api_base = provider_config.get("api_base", "http://127.0.0.1:5000") - if self.api_base.endswith("/"): - self.api_base = self.api_base[:-1] + self.api_base = self.api_base.removesuffix("/") self.character = provider_config.get("character") self.emotion = provider_config.get("emotion") @@ -49,7 +53,7 @@ class ProviderGSVITTS(TTSProvider): else: error_text = await response.text() raise Exception( - f"GSVI TTS API 请求失败，状态码: {response.status}，错误: {error_text}" + f"GSVI TTS API 请求失败，状态码: {response.status}，错误: {error_text}", ) return path diff --git a/astrbot/core/provider/sources/minimax_tts_api_source.py b/astrbot/core/provider/sources/minimax_tts_api_source.py index 5b210835..5ffc7cc6 100644 --- a/astrbot/core/provider/sources/minimax_tts_api_source.py +++ b/astrbot/core/provider/sources/minimax_tts_api_source.py @@ -1,17 +1,22 @@ import json import os import uuid +from collections.abc import AsyncIterator + import aiohttp -from typing import Dict, List, Union, AsyncIterator -from astrbot.core.utils.astrbot_path import get_astrbot_data_path + from astrbot.api import logger +from astrbot.core.utils.astrbot_path import get_astrbot_data_path + from ..entities import ProviderType from ..provider import TTSProvider from ..register import register_provider_adapter @register_provider_adapter( - "minimax_tts_api", "MiniMax TTS API", provider_type=ProviderType.TEXT_TO_SPEECH + "minimax_tts_api", + "MiniMax TTS API", + provider_type=ProviderType.TEXT_TO_SPEECH, ) class ProviderMiniMaxTTSAPI(TTSProvider): def __init__( @@ -22,19 +27,21 @@ class ProviderMiniMaxTTSAPI(TTSProvider): super().__init__(provider_config, provider_settings) self.chosen_api_key: str = provider_config.get("api_key", "") self.api_base: str = provider_config.get( - "api_base", "https://api.minimax.chat/v1/t2a_v2" + "api_base", + "https://api.minimax.chat/v1/t2a_v2", ) self.group_id: str = provider_config.get("minimax-group-id", "") self.set_model(provider_config.get("model", "")) self.lang_boost: str = provider_config.get("minimax-langboost", "auto") self.is_timber_weight: bool = provider_config.get( - "minimax-is-timber-weight", False + "minimax-is-timber-weight", + False, ) - self.timber_weight: List[Dict[str, Union[str, int]]] = json.loads( + self.timber_weight: list[dict[str, str | int]] = json.loads( provider_config.get( "minimax-timber-weight", '[{"voice_id": "Chinese (Mandarin)_Warm_Girl", "weight": 1}]', - ) + ), ) self.voice_setting: dict = { @@ -47,7 +54,8 @@ class ProviderMiniMaxTTSAPI(TTSProvider): "emotion": provider_config.get("minimax-voice-emotion", "neutral"), "latex_read": provider_config.get("minimax-voice-latex", False), "english_normalization": provider_config.get( - "minimax-voice-english-normalization", False + "minimax-voice-english-normalization", + False, ), } @@ -66,7 +74,7 @@ class ProviderMiniMaxTTSAPI(TTSProvider): def _build_tts_stream_body(self, text: str): """构建流式请求体""" - dict_body: Dict[str, object] = { + dict_body: dict[str, object] = { "model": self.model_name, "text": text, "stream": True, @@ -82,44 +90,46 @@ class ProviderMiniMaxTTSAPI(TTSProvider): async def _call_tts_stream(self, text: str) -> AsyncIterator[bytes]: """进行流式请求""" try: - async with aiohttp.ClientSession() as session: - async with session.post( + async with ( + aiohttp.ClientSession() as session, + session.post( self.concat_base_url, headers=self.headers, data=self._build_tts_stream_body(text), timeout=aiohttp.ClientTimeout(total=60), - ) as response: - response.raise_for_status() + ) as response, + ): + response.raise_for_status() - buffer = b"" - while True: - chunk = await response.content.read(8192) - if not chunk: - break + buffer = b"" + while True: + chunk = await response.content.read(8192) + if not chunk: + break - buffer += chunk + buffer += chunk - while b"\n\n" in buffer: - try: - message, buffer = buffer.split(b"\n\n", 1) - if message.startswith(b"data: "): - try: - data = json.loads(message[6:]) - if "extra_info" in data: - continue - audio = data.get("data", {}).get("audio") - if audio is not None: - yield audio - except json.JSONDecodeError: - logger.warning( - "Failed to parse JSON data from SSE message" - ) + while b"\n\n" in buffer: + try: + message, buffer = buffer.split(b"\n\n", 1) + if message.startswith(b"data: "): + try: + data = json.loads(message[6:]) + if "extra_info" in data: continue - except ValueError: - buffer = buffer[-1024:] + audio = data.get("data", {}).get("audio") + if audio is not None: + yield audio + except json.JSONDecodeError: + logger.warning( + "Failed to parse JSON data from SSE message", + ) + continue + except ValueError: + buffer = buffer[-1024:] except aiohttp.ClientError as e: - raise Exception(f"MiniMax TTS API请求失败: {str(e)}") + raise Exception(f"MiniMax TTS API请求失败: {e!s}") async def _audio_play(self, audio_stream: AsyncIterator[str]) -> bytes: """解码数据流到 audio 比特流""" diff --git a/astrbot/core/provider/sources/openai_embedding_source.py b/astrbot/core/provider/sources/openai_embedding_source.py index 79b2e83b..368e610e 100644 --- a/astrbot/core/provider/sources/openai_embedding_source.py +++ b/astrbot/core/provider/sources/openai_embedding_source.py @@ -1,7 +1,8 @@ from openai import AsyncOpenAI + +from ..entities import ProviderType from ..provider import EmbeddingProvider from ..register import register_provider_adapter -from ..entities import ProviderType @register_provider_adapter( @@ -17,27 +18,23 @@ class OpenAIEmbeddingProvider(EmbeddingProvider): self.client = AsyncOpenAI( api_key=provider_config.get("embedding_api_key"), base_url=provider_config.get( - "embedding_api_base", "https://api.openai.com/v1" + "embedding_api_base", + "https://api.openai.com/v1", ), timeout=int(provider_config.get("timeout", 20)), ) self.model = provider_config.get("embedding_model", "text-embedding-3-small") - self.dimension = provider_config.get("embedding_dimensions", 1024) async def get_embedding(self, text: str) -> list[float]: - """ - 获取文本的嵌入 - """ + """获取文本的嵌入""" embedding = await self.client.embeddings.create(input=text, model=self.model) return embedding.data[0].embedding async def get_embeddings(self, texts: list[str]) -> list[list[float]]: - """ - 批量获取文本的嵌入 - """ + """批量获取文本的嵌入""" embeddings = await self.client.embeddings.create(input=texts, model=self.model) return [item.embedding for item in embeddings.data] def get_dim(self) -> int: """获取向量的维度""" - return self.dimension + return self.provider_config.get("embedding_dimensions", 1024) diff --git a/astrbot/core/provider/sources/openai_source.py b/astrbot/core/provider/sources/openai_source.py index 81342ad5..3f1d283c 100644 --- a/astrbot/core/provider/sources/openai_source.py +++ b/astrbot/core/provider/sources/openai_source.py @@ -1,73 +1,101 @@ +import asyncio import base64 +import inspect import json import os -import inspect import random -import asyncio -import astrbot.core.message.components as Comp +import re +from collections.abc import AsyncGenerator -from openai import AsyncOpenAI, AsyncAzureOpenAI -from openai.types.chat.chat_completion import ChatCompletion - -from openai._exceptions import NotFoundError, UnprocessableEntityError +from openai import AsyncAzureOpenAI, AsyncOpenAI +from openai._exceptions import NotFoundError from openai.lib.streaming.chat._completions import ChatCompletionStreamState -from astrbot.core.utils.io import download_image_by_url -from astrbot.core.message.message_event_result import MessageChain +from openai.types.chat.chat_completion import ChatCompletion +from openai.types.chat.chat_completion_chunk import ChatCompletionChunk -from astrbot.api.provider import Provider +import astrbot.core.message.components as Comp from astrbot import logger -from astrbot.core.provider.func_tool_manager import FuncCall -from typing import List, AsyncGenerator -from ..register import register_provider_adapter +from astrbot.api.provider import Provider +from astrbot.core.agent.message import Message +from astrbot.core.agent.tool import ToolSet +from astrbot.core.message.message_event_result import MessageChain from astrbot.core.provider.entities import LLMResponse, ToolCallsResult +from astrbot.core.utils.io import download_image_by_url + +from ..register import register_provider_adapter @register_provider_adapter( - "openai_chat_completion", "OpenAI API Chat Completion 提供商适配器" + "openai_chat_completion", + "OpenAI API Chat Completion 提供商适配器", ) class ProviderOpenAIOfficial(Provider): - def __init__( - self, - provider_config, - provider_settings, - default_persona=None, - ) -> None: - super().__init__( - provider_config, - provider_settings, - default_persona, - ) + def __init__(self, provider_config, provider_settings) -> None: + super().__init__(provider_config, provider_settings) self.chosen_api_key = None - self.api_keys: List = super().get_keys() + self.api_keys: list = super().get_keys() self.chosen_api_key = self.api_keys[0] if len(self.api_keys) > 0 else None self.timeout = provider_config.get("timeout", 120) + self.custom_headers = provider_config.get("custom_headers", {}) if isinstance(self.timeout, str): self.timeout = int(self.timeout) - # 适配 azure openai #332 + + if not isinstance(self.custom_headers, dict) or not self.custom_headers: + self.custom_headers = None + else: + for key in self.custom_headers: + self.custom_headers[key] = str(self.custom_headers[key]) + if "api_version" in provider_config: - # 使用 azure api + # Using Azure OpenAI API self.client = AsyncAzureOpenAI( api_key=self.chosen_api_key, api_version=provider_config.get("api_version", None), - base_url=provider_config.get("api_base", None), + default_headers=self.custom_headers, + base_url=provider_config.get("api_base", ""), timeout=self.timeout, ) else: - # 使用 openai api + # Using OpenAI Official API self.client = AsyncOpenAI( api_key=self.chosen_api_key, base_url=provider_config.get("api_base", None), + default_headers=self.custom_headers, timeout=self.timeout, ) self.default_params = inspect.signature( - self.client.chat.completions.create + self.client.chat.completions.create, ).parameters.keys() model_config = provider_config.get("model_config", {}) model = model_config.get("model", "unknown") self.set_model(model) + self.reasoning_key = "reasoning_content" + + def _maybe_inject_xai_search(self, payloads: dict, **kwargs): + """当开启 xAI 原生搜索时，向请求体注入 Live Search 参数。 + + - 仅在 provider_config.xai_native_search 为 True 时生效 + - 默认注入 {"mode": "auto"} + - 允许通过 kwargs 使用 xai_search_mode 覆盖（on/auto/off） + """ + if not bool(self.provider_config.get("xai_native_search", False)): + return + + mode = kwargs.get("xai_search_mode", "auto") + mode = str(mode).lower() + if mode not in ("auto", "on", "off"): + mode = "auto" + + # off 时不注入，保持与未开启一致 + if mode == "off": + return + + # OpenAI SDK 不识别的字段会在 _query/_query_stream 中放入 extra_body + payloads["search_parameters"] = {"mode": mode} + async def get_models(self): try: models_str = [] @@ -79,12 +107,12 @@ class ProviderOpenAIOfficial(Provider): except NotFoundError as e: raise Exception(f"获取模型列表失败：{e}") - async def _query(self, payloads: dict, tools: FuncCall) -> LLMResponse: + async def _query(self, payloads: dict, tools: ToolSet | None) -> LLMResponse: if tools: model = payloads.get("model", "").lower() omit_empty_param_field = "gemini" in model tool_list = tools.get_func_desc_openai_style( - omit_empty_parameter_field=omit_empty_param_field + omit_empty_parameter_field=omit_empty_param_field, ) if tool_list: payloads["tools"] = tool_list @@ -92,7 +120,7 @@ class ProviderOpenAIOfficial(Provider): # 不在默认参数中的参数放在 extra_body 中 extra_body = {} to_del = [] - for key in payloads.keys(): + for key in payloads: if key not in self.default_params: extra_body[key] = payloads[key] to_del.append(key) @@ -111,29 +139,33 @@ class ProviderOpenAIOfficial(Provider): del payloads["tools"] completion = await self.client.chat.completions.create( - **payloads, stream=False, extra_body=extra_body + **payloads, + stream=False, + extra_body=extra_body, ) if not isinstance(completion, ChatCompletion): raise Exception( - f"API 返回的 completion 类型错误：{type(completion)}: {completion}。" + f"API 返回的 completion 类型错误：{type(completion)}: {completion}。", ) logger.debug(f"completion: {completion}") - llm_response = await self.parse_openai_completion(completion, tools) + llm_response = await self._parse_openai_completion(completion, tools) return llm_response async def _query_stream( - self, payloads: dict, tools: FuncCall + self, + payloads: dict, + tools: ToolSet | None, ) -> AsyncGenerator[LLMResponse, None]: """流式查询API，逐步返回结果""" if tools: model = payloads.get("model", "").lower() omit_empty_param_field = "gemini" in model tool_list = tools.get_func_desc_openai_style( - omit_empty_parameter_field=omit_empty_param_field + omit_empty_parameter_field=omit_empty_param_field, ) if tool_list: payloads["tools"] = tool_list @@ -147,7 +179,7 @@ class ProviderOpenAIOfficial(Provider): extra_body.update(custom_extra_body) to_del = [] - for key in payloads.keys(): + for key in payloads: if key not in self.default_params: extra_body[key] = payloads[key] to_del.append(key) @@ -155,7 +187,9 @@ class ProviderOpenAIOfficial(Provider): del payloads[key] stream = await self.client.chat.completions.create( - **payloads, stream=True, extra_body=extra_body + **payloads, + stream=True, + extra_body=extra_body, ) llm_response = LLMResponse("assistant", is_chunk=True) @@ -170,45 +204,91 @@ class ProviderOpenAIOfficial(Provider): if len(chunk.choices) == 0: continue delta = chunk.choices[0].delta - # 处理文本内容 + # logger.debug(f"chunk delta: {delta}") + # handle the content delta + reasoning = self._extract_reasoning_content(chunk) + _y = False + if reasoning: + llm_response.reasoning_content = reasoning + _y = True if delta.content: completion_text = delta.content llm_response.result_chain = MessageChain( - chain=[Comp.Plain(completion_text)] + chain=[Comp.Plain(completion_text)], ) + _y = True + if _y: yield llm_response final_completion = state.get_final_completion() - llm_response = await self.parse_openai_completion(final_completion, tools) + llm_response = await self._parse_openai_completion(final_completion, tools) yield llm_response - async def parse_openai_completion( - self, completion: ChatCompletion, tools: FuncCall - ): - """解析 OpenAI 的 ChatCompletion 响应""" + def _extract_reasoning_content( + self, + completion: ChatCompletion | ChatCompletionChunk, + ) -> str: + """Extract reasoning content from OpenAI ChatCompletion if available.""" + reasoning_text = "" + if len(completion.choices) == 0: + return reasoning_text + if isinstance(completion, ChatCompletion): + choice = completion.choices[0] + reasoning_attr = getattr(choice.message, self.reasoning_key, None) + if reasoning_attr: + reasoning_text = str(reasoning_attr) + elif isinstance(completion, ChatCompletionChunk): + delta = completion.choices[0].delta + reasoning_attr = getattr(delta, self.reasoning_key, None) + if reasoning_attr: + reasoning_text = str(reasoning_attr) + return reasoning_text + + async def _parse_openai_completion( + self, completion: ChatCompletion, tools: ToolSet | None + ) -> LLMResponse: + """Parse OpenAI ChatCompletion into LLMResponse""" llm_response = LLMResponse("assistant") if len(completion.choices) == 0: raise Exception("API 返回的 completion 为空。") choice = completion.choices[0] + # parse the text completion if choice.message.content is not None: # text completion completion_text = str(choice.message.content).strip() + # specially, some providers may set tags around reasoning content in the completion text, + # we use regex to remove them, and store then in reasoning_content field + reasoning_pattern = re.compile(r"(.*?)", re.DOTALL) + matches = reasoning_pattern.findall(completion_text) + if matches: + llm_response.reasoning_content = "\n".join( + [match.strip() for match in matches], + ) + completion_text = reasoning_pattern.sub("", completion_text).strip() llm_response.result_chain = MessageChain().message(completion_text) - if choice.message.tool_calls: - # tools call (function calling) + # parse the reasoning content if any + # the priority is higher than the tag extraction + llm_response.reasoning_content = self._extract_reasoning_content(completion) + + # parse tool calls if any + if choice.message.tool_calls and tools is not None: args_ls = [] func_name_ls = [] tool_call_ids = [] + tool_call_extra_content_dict = {} for tool_call in choice.message.tool_calls: if isinstance(tool_call, str): # workaround for #1359 tool_call = json.loads(tool_call) for tool in tools.func_list: - if tool.name == tool_call.function.name: + if ( + tool_call.type == "function" + and tool.name == tool_call.function.name + ): # workaround for #1454 if isinstance(tool_call.function.arguments, str): args = json.loads(tool_call.function.arguments) @@ -217,16 +297,21 @@ class ProviderOpenAIOfficial(Provider): args_ls.append(args) func_name_ls.append(tool_call.function.name) tool_call_ids.append(tool_call.id) + + # gemini-2.5 / gemini-3 series extra_content handling + extra_content = getattr(tool_call, "extra_content", None) + if extra_content is not None: + tool_call_extra_content_dict[tool_call.id] = extra_content llm_response.role = "tool" llm_response.tools_call_args = args_ls llm_response.tools_call_name = func_name_ls llm_response.tools_call_ids = tool_call_ids - + llm_response.tools_call_extra_content = tool_call_extra_content_dict + # specially handle finish reason if choice.finish_reason == "content_filter": raise Exception( - "API 返回的 completion 由于内容安全过滤被拒绝(非 AstrBot)。" + "API 返回的 completion 由于内容安全过滤被拒绝(非 AstrBot)。", ) - if llm_response.completion_text is None and not llm_response.tools_call_args: logger.error(f"API 返回的 completion 无法解析：{completion}。") raise Exception(f"API 返回的 completion 无法解析：{completion}。") @@ -237,9 +322,9 @@ class ProviderOpenAIOfficial(Provider): async def _prepare_chat_payload( self, - prompt: str, + prompt: str | None, image_urls: list[str] | None = None, - contexts: list | None = None, + contexts: list[dict] | list[Message] | None = None, system_prompt: str | None = None, tool_calls_result: ToolCallsResult | list[ToolCallsResult] | None = None, model: str | None = None, @@ -248,8 +333,12 @@ class ProviderOpenAIOfficial(Provider): """准备聊天所需的有效载荷和上下文""" if contexts is None: contexts = [] - new_record = await self.assemble_context(prompt, image_urls) - context_query = [*contexts, new_record] + new_record = None + if prompt is not None: + new_record = await self.assemble_context(prompt, image_urls) + context_query = self._ensure_message_to_dicts(contexts) + if new_record: + context_query.append(new_record) if system_prompt: context_query.insert(0, {"role": "system", "content": system_prompt}) @@ -270,6 +359,9 @@ class ProviderOpenAIOfficial(Provider): payloads = {"messages": context_query, **model_config} + # xAI origin search tool inject + self._maybe_inject_xai_search(payloads, **kwargs) + return payloads, context_query async def _handle_api_error( @@ -277,16 +369,16 @@ class ProviderOpenAIOfficial(Provider): e: Exception, payloads: dict, context_query: list, - func_tool: FuncCall, + func_tool: ToolSet | None, chosen_key: str, - available_api_keys: List[str], + available_api_keys: list[str], retry_cnt: int, max_retries: int, ) -> tuple: """处理API错误并尝试恢复""" if "429" in str(e): logger.warning( - f"API 调用过于频繁，尝试使用其他 Key 重试。当前 Key: {chosen_key[:12]}" + f"API 调用过于频繁，尝试使用其他 Key 重试。当前 Key: {chosen_key[:12]}", ) # 最后一次不等待 if retry_cnt < max_retries - 1: @@ -302,11 +394,10 @@ class ProviderOpenAIOfficial(Provider): context_query, func_tool, ) - else: - raise e - elif "maximum context length" in str(e): + raise e + if "maximum context length" in str(e): logger.warning( - f"上下文长度超过限制。尝试弹出最早的记录然后重试。当前记录条数: {len(context_query)}" + f"上下文长度超过限制。尝试弹出最早的记录然后重试。当前记录条数: {len(context_query)}", ) await self.pop_record(context_query) payloads["messages"] = context_query @@ -318,7 +409,7 @@ class ProviderOpenAIOfficial(Provider): context_query, func_tool, ) - elif "The model is not a VLM" in str(e): # siliconcloud + if "The model is not a VLM" in str(e): # siliconcloud # 尝试删除所有 image new_contexts = await self._remove_image_from_context(context_query) payloads["messages"] = new_contexts @@ -331,36 +422,34 @@ class ProviderOpenAIOfficial(Provider): context_query, func_tool, ) - elif ( + if ( "Function calling is not enabled" in str(e) or ("tool" in str(e).lower() and "support" in str(e).lower()) or ("function" in str(e).lower() and "support" in str(e).lower()) ): # openai, ollama, gemini openai, siliconcloud 的错误提示与 code 不统一，只能通过字符串匹配 logger.info( - f"{self.get_model()} 不支持函数工具调用，已自动去除，不影响使用。" + f"{self.get_model()} 不支持函数工具调用，已自动去除，不影响使用。", ) - if "tools" in payloads: - del payloads["tools"] + payloads.pop("tools", None) return False, chosen_key, available_api_keys, payloads, context_query, None - else: - logger.error(f"发生了错误。Provider 配置如下: {self.provider_config}") + logger.error(f"发生了错误。Provider 配置如下: {self.provider_config}") - if "tool" in str(e).lower() and "support" in str(e).lower(): - logger.error("疑似该模型不支持函数调用工具调用。请输入 /tool off_all") + if "tool" in str(e).lower() and "support" in str(e).lower(): + logger.error("疑似该模型不支持函数调用工具调用。请输入 /tool off_all") - if "Connection error." in str(e): - proxy = os.environ.get("http_proxy", None) - if proxy: - logger.error( - f"可能为代理原因，请检查代理是否正常。当前代理: {proxy}" - ) + if "Connection error." in str(e): + proxy = os.environ.get("http_proxy", None) + if proxy: + logger.error( + f"可能为代理原因，请检查代理是否正常。当前代理: {proxy}", + ) - raise e + raise e async def text_chat( self, - prompt, + prompt=None, session_id=None, image_urls=None, func_tool=None, @@ -392,12 +481,6 @@ class ProviderOpenAIOfficial(Provider): self.client.api_key = chosen_key llm_response = await self._query(payloads, func_tool) break - except UnprocessableEntityError as e: - logger.warning(f"不可处理的实体错误：{e}，尝试删除图片。") - # 尝试删除所有 image - new_contexts = await self._remove_image_from_context(context_query) - payloads["messages"] = new_contexts - context_query = new_contexts except Exception as e: last_exception = e ( @@ -420,7 +503,7 @@ class ProviderOpenAIOfficial(Provider): if success: break - if retry_cnt == max_retries - 1: + if retry_cnt == max_retries - 1 or llm_response is None: logger.error(f"API 调用失败，重试 {max_retries} 次仍然失败。") if last_exception is None: raise Exception("未知错误") @@ -429,11 +512,11 @@ class ProviderOpenAIOfficial(Provider): async def text_chat_stream( self, - prompt: str, - session_id: str = None, - image_urls: List[str] = [], - func_tool: FuncCall = None, - contexts=[], + prompt=None, + session_id=None, + image_urls=None, + func_tool=None, + contexts=None, system_prompt=None, tool_calls_result=None, model=None, @@ -462,12 +545,6 @@ class ProviderOpenAIOfficial(Provider): async for response in self._query_stream(payloads, func_tool): yield response break - except UnprocessableEntityError as e: - logger.warning(f"不可处理的实体错误：{e}，尝试删除图片。") - # 尝试删除所有 image - new_contexts = await self._remove_image_from_context(context_query) - payloads["messages"] = new_contexts - context_query = new_contexts except Exception as e: last_exception = e ( @@ -496,10 +573,8 @@ class ProviderOpenAIOfficial(Provider): raise Exception("未知错误") raise last_exception - async def _remove_image_from_context(self, contexts: List): - """ - 从上下文中删除所有带有 image 的记录 - """ + async def _remove_image_from_context(self, contexts: list): + """从上下文中删除所有带有 image 的记录""" new_contexts = [] for context in contexts: @@ -520,13 +595,17 @@ class ProviderOpenAIOfficial(Provider): def get_current_key(self) -> str: return self.client.api_key - def get_keys(self) -> List[str]: + def get_keys(self) -> list[str]: return self.api_keys def set_key(self, key): self.client.api_key = key - async def assemble_context(self, text: str, image_urls: List[str] = None) -> dict: + async def assemble_context( + self, + text: str, + image_urls: list[str] | None = None, + ) -> dict: """组装成符合 OpenAI 格式的 role 为 user 的消息段""" if image_urls: user_content = { @@ -549,19 +628,15 @@ class ProviderOpenAIOfficial(Provider): { "type": "image_url", "image_url": {"url": image_data}, - } + }, ) return user_content - else: - return {"role": "user", "content": text} + return {"role": "user", "content": text} async def encode_image_bs64(self, image_url: str) -> str: - """ - 将图片转换为 base64 - """ + """将图片转换为 base64""" if image_url.startswith("base64://"): return image_url.replace("base64://", "data:image/jpeg;base64,") with open(image_url, "rb") as f: image_bs64 = base64.b64encode(f.read()).decode("utf-8") return "data:image/jpeg;base64," + image_bs64 - return "" diff --git a/astrbot/core/provider/sources/openai_tts_api_source.py b/astrbot/core/provider/sources/openai_tts_api_source.py index c188a9fa..d71e9811 100644 --- a/astrbot/core/provider/sources/openai_tts_api_source.py +++ b/astrbot/core/provider/sources/openai_tts_api_source.py @@ -1,14 +1,19 @@ import os import uuid -from openai import AsyncOpenAI, NOT_GIVEN -from ..provider import TTSProvider -from ..entities import ProviderType -from ..register import register_provider_adapter + +from openai import NOT_GIVEN, AsyncOpenAI + from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from ..entities import ProviderType +from ..provider import TTSProvider +from ..register import register_provider_adapter + @register_provider_adapter( - "openai_tts_api", "OpenAI TTS API", provider_type=ProviderType.TEXT_TO_SPEECH + "openai_tts_api", + "OpenAI TTS API", + provider_type=ProviderType.TEXT_TO_SPEECH, ) class ProviderOpenAITTSAPI(TTSProvider): def __init__( @@ -26,17 +31,20 @@ class ProviderOpenAITTSAPI(TTSProvider): self.client = AsyncOpenAI( api_key=self.chosen_api_key, - base_url=provider_config.get("api_base", None), + base_url=provider_config.get("api_base"), timeout=timeout, ) - self.set_model(provider_config.get("model", None)) + self.set_model(provider_config.get("model", "")) async def get_audio(self, text: str) -> str: temp_dir = os.path.join(get_astrbot_data_path(), "temp") path = os.path.join(temp_dir, f"openai_tts_api_{uuid.uuid4()}.wav") async with self.client.audio.speech.with_streaming_response.create( - model=self.model_name, voice=self.voice, response_format="wav", input=text + model=self.model_name, + voice=self.voice, + response_format="wav", + input=text, ) as response: with open(path, "wb") as f: async for chunk in response.iter_bytes(chunk_size=1024): diff --git a/astrbot/core/provider/sources/sensevoice_selfhosted_source.py b/astrbot/core/provider/sources/sensevoice_selfhosted_source.py index b6e3331f..67947c68 100644 --- a/astrbot/core/provider/sources/sensevoice_selfhosted_source.py +++ b/astrbot/core/provider/sources/sensevoice_selfhosted_source.py @@ -1,22 +1,24 @@ -""" -Author: diudiu62 +"""Author: diudiu62 Date: 2025-02-24 18:04:18 LastEditTime: 2025-02-25 14:06:30 """ import asyncio -from datetime import datetime import os import re +from datetime import datetime + from funasr_onnx import SenseVoiceSmall from funasr_onnx.utils.postprocess_utils import rich_transcription_postprocess -from ..provider import STTProvider -from ..entities import ProviderType -from astrbot.core.utils.io import download_file -from ..register import register_provider_adapter + from astrbot.core import logger +from astrbot.core.utils.io import download_file from astrbot.core.utils.tencent_record_helper import tencent_silk_to_wav +from ..entities import ProviderType +from ..provider import STTProvider +from ..register import register_provider_adapter + @register_provider_adapter( "sensevoice_stt_selfhost", @@ -30,7 +32,7 @@ class ProviderSenseVoiceSTTSelfHost(STTProvider): provider_settings: dict, ) -> None: super().__init__(provider_config, provider_settings) - self.set_model(provider_config.get("stt_model", None)) + self.set_model(provider_config.get("stt_model")) self.model = None self.is_emotion = provider_config.get("is_emotion", False) @@ -39,7 +41,8 @@ class ProviderSenseVoiceSTTSelfHost(STTProvider): # 将模型加载放到线程池中执行 self.model = await asyncio.get_event_loop().run_in_executor( - None, lambda: SenseVoiceSmall(self.model_name, quantize=True, batch_size=16) + None, + lambda: SenseVoiceSmall(self.model_name, quantize=True, batch_size=16), ) logger.info("SenseVoice 模型加载完成。") @@ -55,8 +58,7 @@ class ProviderSenseVoiceSTTSelfHost(STTProvider): if silk_header in file_header: return True - else: - return False + return False async def get_text(self, audio_url: str) -> str: try: diff --git a/astrbot/core/provider/sources/vllm_rerank_source.py b/astrbot/core/provider/sources/vllm_rerank_source.py index 2620e345..3e6f3d33 100644 --- a/astrbot/core/provider/sources/vllm_rerank_source.py +++ b/astrbot/core/provider/sources/vllm_rerank_source.py @@ -1,8 +1,10 @@ import aiohttp + from astrbot import logger + +from ..entities import ProviderType, RerankResult from ..provider import RerankProvider from ..register import register_provider_adapter -from ..entities import ProviderType, RerankResult @register_provider_adapter( @@ -30,7 +32,10 @@ class VLLMRerankProvider(RerankProvider): ) async def rerank( - self, query: str, documents: list[str], top_n: int | None = None + self, + query: str, + documents: list[str], + top_n: int | None = None, ) -> list[RerankResult]: payload = { "query": query, @@ -40,14 +45,15 @@ class VLLMRerankProvider(RerankProvider): if top_n is not None: payload["top_n"] = top_n async with self.client.post( - f"{self.base_url}/v1/rerank", json=payload + f"{self.base_url}/v1/rerank", + json=payload, ) as response: response_data = await response.json() results = response_data.get("results", []) if not results: logger.warning( - f"Rerank API 返回了空的列表数据。原始响应: {response_data}" + f"Rerank API 返回了空的列表数据。原始响应: {response_data}", ) return [ diff --git a/astrbot/core/provider/sources/volcengine_tts.py b/astrbot/core/provider/sources/volcengine_tts.py index 12e7ed9c..f5d758f5 100644 --- a/astrbot/core/provider/sources/volcengine_tts.py +++ b/astrbot/core/provider/sources/volcengine_tts.py @@ -1,18 +1,23 @@ -import uuid +import asyncio import base64 import json import os import traceback -import asyncio +import uuid + import aiohttp -from ..provider import TTSProvider -from ..entities import ProviderType -from ..register import register_provider_adapter + from astrbot import logger +from ..entities import ProviderType +from ..provider import TTSProvider +from ..register import register_provider_adapter + @register_provider_adapter( - "volcengine_tts", "火山引擎 TTS", provider_type=ProviderType.TEXT_TO_SPEECH + "volcengine_tts", + "火山引擎 TTS", + provider_type=ProviderType.TEXT_TO_SPEECH, ) class ProviderVolcengineTTS(TTSProvider): def __init__(self, provider_config: dict, provider_settings: dict) -> None: @@ -23,7 +28,8 @@ class ProviderVolcengineTTS(TTSProvider): self.voice_type = provider_config.get("volcengine_voice_type", "") self.speed_ratio = provider_config.get("volcengine_speed_ratio", 1.0) self.api_base = provider_config.get( - "api_base", "https://openspeech.bytedance.com/api/v1/tts" + "api_base", + "https://openspeech.bytedance.com/api/v1/tts", ) self.timeout = provider_config.get("timeout", 20) @@ -66,43 +72,44 @@ class ProviderVolcengineTTS(TTSProvider): logger.debug(f"请求体: {json.dumps(payload, ensure_ascii=False)[:100]}...") try: - async with aiohttp.ClientSession() as session: - async with session.post( + async with ( + aiohttp.ClientSession() as session, + session.post( self.api_base, data=json.dumps(payload), headers=headers, timeout=self.timeout, - ) as response: - logger.debug(f"响应状态码: {response.status}") + ) as response, + ): + logger.debug(f"响应状态码: {response.status}") - response_text = await response.text() - logger.debug(f"响应内容: {response_text[:200]}...") + response_text = await response.text() + logger.debug(f"响应内容: {response_text[:200]}...") - if response.status == 200: - resp_data = json.loads(response_text) + if response.status == 200: + resp_data = json.loads(response_text) - if "data" in resp_data: - audio_data = base64.b64decode(resp_data["data"]) + if "data" in resp_data: + audio_data = base64.b64decode(resp_data["data"]) - os.makedirs("data/temp", exist_ok=True) + os.makedirs("data/temp", exist_ok=True) - file_path = f"data/temp/volcengine_tts_{uuid.uuid4()}.mp3" + file_path = f"data/temp/volcengine_tts_{uuid.uuid4()}.mp3" - loop = asyncio.get_running_loop() - await loop.run_in_executor( - None, lambda: open(file_path, "wb").write(audio_data) - ) - - return file_path - else: - error_msg = resp_data.get("message", "未知错误") - raise Exception(f"火山引擎 TTS API 返回错误: {error_msg}") - else: - raise Exception( - f"火山引擎 TTS API 请求失败: {response.status}, {response_text}" + loop = asyncio.get_running_loop() + await loop.run_in_executor( + None, + lambda: open(file_path, "wb").write(audio_data), ) + return file_path + error_msg = resp_data.get("message", "未知错误") + raise Exception(f"火山引擎 TTS API 返回错误: {error_msg}") + raise Exception( + f"火山引擎 TTS API 请求失败: {response.status}, {response_text}", + ) + except Exception as e: error_details = traceback.format_exc() logger.debug(f"火山引擎 TTS 异常详情: {error_details}") - raise Exception(f"火山引擎 TTS 异常: {str(e)}") + raise Exception(f"火山引擎 TTS 异常: {e!s}") diff --git a/astrbot/core/provider/sources/whisper_api_source.py b/astrbot/core/provider/sources/whisper_api_source.py index dfe28697..8f6d9e29 100644 --- a/astrbot/core/provider/sources/whisper_api_source.py +++ b/astrbot/core/provider/sources/whisper_api_source.py @@ -1,13 +1,16 @@ -import uuid import os -from openai import AsyncOpenAI, NOT_GIVEN -from ..provider import STTProvider -from ..entities import ProviderType -from astrbot.core.utils.io import download_file -from ..register import register_provider_adapter +import uuid + +from openai import NOT_GIVEN, AsyncOpenAI + from astrbot.core import logger -from astrbot.core.utils.tencent_record_helper import tencent_silk_to_wav from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.io import download_file +from astrbot.core.utils.tencent_record_helper import tencent_silk_to_wav + +from ..entities import ProviderType +from ..provider import STTProvider +from ..register import register_provider_adapter @register_provider_adapter( @@ -26,11 +29,11 @@ class ProviderOpenAIWhisperAPI(STTProvider): self.client = AsyncOpenAI( api_key=self.chosen_api_key, - base_url=provider_config.get("api_base", None), + base_url=provider_config.get("api_base"), timeout=provider_config.get("timeout", NOT_GIVEN), ) - self.set_model(provider_config.get("model", None)) + self.set_model(provider_config.get("model")) async def _is_silk_file(self, file_path): silk_header = b"SILK" @@ -39,11 +42,10 @@ class ProviderOpenAIWhisperAPI(STTProvider): if silk_header in file_header: return True - else: - return False + return False async def get_text(self, audio_url: str) -> str: - """only supports mp3, mp4, mpeg, m4a, wav, webm""" + """Only supports mp3, mp4, mpeg, m4a, wav, webm""" is_tencent = False if audio_url.startswith("http"): diff --git a/astrbot/core/provider/sources/whisper_selfhosted_source.py b/astrbot/core/provider/sources/whisper_selfhosted_source.py index 7cb76cc4..fbdc7d62 100644 --- a/astrbot/core/provider/sources/whisper_selfhosted_source.py +++ b/astrbot/core/provider/sources/whisper_selfhosted_source.py @@ -1,14 +1,17 @@ -import uuid -import os import asyncio +import os +import uuid + import whisper -from ..provider import STTProvider -from ..entities import ProviderType -from astrbot.core.utils.io import download_file -from ..register import register_provider_adapter + from astrbot.core import logger -from astrbot.core.utils.tencent_record_helper import tencent_silk_to_wav from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.io import download_file +from astrbot.core.utils.tencent_record_helper import tencent_silk_to_wav + +from ..entities import ProviderType +from ..provider import STTProvider +from ..register import register_provider_adapter @register_provider_adapter( @@ -23,14 +26,16 @@ class ProviderOpenAIWhisperSelfHost(STTProvider): provider_settings: dict, ) -> None: super().__init__(provider_config, provider_settings) - self.set_model(provider_config.get("model", None)) + self.set_model(provider_config.get("model")) self.model = None async def initialize(self): loop = asyncio.get_event_loop() logger.info("下载或者加载 Whisper 模型中，这可能需要一些时间 ...") self.model = await loop.run_in_executor( - None, whisper.load_model, self.model_name + None, + whisper.load_model, + self.model_name, ) logger.info("Whisper 模型加载完成。") @@ -41,8 +46,7 @@ class ProviderOpenAIWhisperSelfHost(STTProvider): if silk_header in file_header: return True - else: - return False + return False async def get_text(self, audio_url: str) -> str: loop = asyncio.get_event_loop() diff --git a/astrbot/core/provider/sources/xinference_rerank_source.py b/astrbot/core/provider/sources/xinference_rerank_source.py new file mode 100644 index 00000000..29f3ab09 --- /dev/null +++ b/astrbot/core/provider/sources/xinference_rerank_source.py @@ -0,0 +1,116 @@ +from xinference_client.client.restful.async_restful_client import ( + AsyncClient as Client, +) + +from astrbot import logger + +from ..entities import ProviderType, RerankResult +from ..provider import RerankProvider +from ..register import register_provider_adapter + + +@register_provider_adapter( + "xinference_rerank", + "Xinference Rerank 适配器", + provider_type=ProviderType.RERANK, +) +class XinferenceRerankProvider(RerankProvider): + def __init__(self, provider_config: dict, provider_settings: dict) -> None: + super().__init__(provider_config, provider_settings) + self.provider_config = provider_config + self.provider_settings = provider_settings + self.base_url = provider_config.get("rerank_api_base", "http://127.0.0.1:8000") + self.base_url = self.base_url.rstrip("/") + self.timeout = provider_config.get("timeout", 20) + self.model_name = provider_config.get("rerank_model", "BAAI/bge-reranker-base") + self.api_key = provider_config.get("rerank_api_key") + self.launch_model_if_not_running = provider_config.get( + "launch_model_if_not_running", + False, + ) + self.client = None + self.model = None + self.model_uid = None + + async def initialize(self): + if self.api_key: + logger.info("Xinference Rerank: Using API key for authentication.") + self.client = Client(self.base_url, api_key=self.api_key) + else: + logger.info("Xinference Rerank: No API key provided.") + self.client = Client(self.base_url) + + try: + running_models = await self.client.list_models() + for uid, model_spec in running_models.items(): + if model_spec.get("model_name") == self.model_name: + logger.info( + f"Model '{self.model_name}' is already running with UID: {uid}", + ) + self.model_uid = uid + break + + if self.model_uid is None: + if self.launch_model_if_not_running: + logger.info(f"Launching {self.model_name} model...") + self.model_uid = await self.client.launch_model( + model_name=self.model_name, + model_type="rerank", + ) + logger.info("Model launched.") + else: + logger.warning( + f"Model '{self.model_name}' is not running and auto-launch is disabled. Provider will not be available.", + ) + return + + if self.model_uid: + self.model = await self.client.get_model(self.model_uid) + + except Exception as e: + logger.error(f"Failed to initialize Xinference model: {e}") + logger.debug( + f"Xinference initialization failed with exception: {e}", + exc_info=True, + ) + self.model = None + + async def rerank( + self, + query: str, + documents: list[str], + top_n: int | None = None, + ) -> list[RerankResult]: + if not self.model: + logger.error("Xinference rerank model is not initialized.") + return [] + try: + response = await self.model.rerank(documents, query, top_n) + results = response.get("results", []) + logger.debug(f"Rerank API response: {response}") + + if not results: + logger.warning( + f"Rerank API returned an empty list. Original response: {response}", + ) + + return [ + RerankResult( + index=result["index"], + relevance_score=result["relevance_score"], + ) + for result in results + ] + except Exception as e: + logger.error(f"Xinference rerank failed: {e}") + logger.debug(f"Xinference rerank failed with exception: {e}", exc_info=True) + return [] + + async def terminate(self) -> None: + """关闭客户端会话""" + if self.client: + logger.info("Closing Xinference rerank client...") + try: + await self.client.close() + except Exception as e: + logger.error(f"Failed to close Xinference client: {e}", exc_info=True) diff --git a/astrbot/core/provider/sources/xinference_stt_provider.py b/astrbot/core/provider/sources/xinference_stt_provider.py new file mode 100644 index 00000000..9c69a003 --- /dev/null +++ b/astrbot/core/provider/sources/xinference_stt_provider.py @@ -0,0 +1,197 @@ +import os +import uuid + +import aiohttp +from xinference_client.client.restful.async_restful_client import ( + AsyncClient as Client, +) + +from astrbot.core import logger +from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.tencent_record_helper import tencent_silk_to_wav + +from ..entities import ProviderType +from ..provider import STTProvider +from ..register import register_provider_adapter + + +@register_provider_adapter( + "xinference_stt", + "Xinference STT", + provider_type=ProviderType.SPEECH_TO_TEXT, +) +class ProviderXinferenceSTT(STTProvider): + def __init__(self, provider_config: dict, provider_settings: dict) -> None: + super().__init__(provider_config, provider_settings) + self.provider_config = provider_config + self.provider_settings = provider_settings + self.base_url = provider_config.get("api_base", "http://127.0.0.1:9997") + self.base_url = self.base_url.rstrip("/") + self.timeout = provider_config.get("timeout", 180) + self.model_name = provider_config.get("model", "whisper-large-v3") + self.api_key = provider_config.get("api_key") + self.launch_model_if_not_running = provider_config.get( + "launch_model_if_not_running", + False, + ) + self.client = None + self.model_uid = None + + async def initialize(self): + if self.api_key: + logger.info("Xinference STT: Using API key for authentication.") + self.client = Client(self.base_url, api_key=self.api_key) + else: + logger.info("Xinference STT: No API key provided.") + self.client = Client(self.base_url) + + try: + running_models = await self.client.list_models() + for uid, model_spec in running_models.items(): + if model_spec.get("model_name") == self.model_name: + logger.info( + f"Model '{self.model_name}' is already running with UID: {uid}", + ) + self.model_uid = uid + break + + if self.model_uid is None: + if self.launch_model_if_not_running: + logger.info(f"Launching {self.model_name} model...") + self.model_uid = await self.client.launch_model( + model_name=self.model_name, + model_type="audio", + ) + logger.info("Model launched.") + else: + logger.warning( + f"Model '{self.model_name}' is not running and auto-launch is disabled. Provider will not be available.", + ) + return + + except Exception as e: + logger.error(f"Failed to initialize Xinference model: {e}") + logger.debug( + f"Xinference initialization failed with exception: {e}", + exc_info=True, + ) + + async def get_text(self, audio_url: str) -> str: + if not self.model_uid or self.client is None or self.client.session is None: + logger.error("Xinference STT model is not initialized.") + return "" + + audio_bytes = None + temp_files = [] + is_tencent = False + + try: + # 1. Get audio bytes + if audio_url.startswith("http"): + if "multimedia.nt.qq.com.cn" in audio_url: + is_tencent = True + async with aiohttp.ClientSession() as session: + async with session.get(audio_url, timeout=self.timeout) as resp: + if resp.status == 200: + audio_bytes = await resp.read() + else: + logger.error( + f"Failed to download audio from {audio_url}, status: {resp.status}", + ) + return "" + elif os.path.exists(audio_url): + with open(audio_url, "rb") as f: + audio_bytes = f.read() + else: + logger.error(f"File not found: {audio_url}") + return "" + + if not audio_bytes: + logger.error("Audio bytes are empty.") + return "" + + # 2. Check for conversion + needs_conversion = False + if ( + audio_url.endswith((".amr", ".silk")) + or is_tencent + or b"SILK" in audio_bytes[:8] + ): + needs_conversion = True + + # 3. Perform conversion if needed + if needs_conversion: + logger.info("Audio requires conversion, using temporary files...") + temp_dir = os.path.join(get_astrbot_data_path(), "temp") + os.makedirs(temp_dir, exist_ok=True) + + input_path = os.path.join(temp_dir, str(uuid.uuid4())) + output_path = os.path.join(temp_dir, str(uuid.uuid4()) + ".wav") + temp_files.extend([input_path, output_path]) + + with open(input_path, "wb") as f: + f.write(audio_bytes) + + logger.info("Converting silk/amr file to wav ...") + await tencent_silk_to_wav(input_path, output_path) + + with open(output_path, "rb") as f: + audio_bytes = f.read() + + # 4. Transcribe + # 官方asyncCLient的客户端似乎实现有点问题，这里直接用aiohttp实现openai标准兼容请求，提交issue等待官方修复后再改回来 + url = f"{self.base_url}/v1/audio/transcriptions" + headers = { + "accept": "application/json", + } + if self.client and self.client._headers: + headers.update(self.client._headers) + + data = aiohttp.FormData() + data.add_field("model", self.model_uid) + data.add_field( + "file", + audio_bytes, + filename="audio.wav", + content_type="audio/wav", + ) + + async with self.client.session.post( + url, + data=data, + headers=headers, + timeout=self.timeout, + ) as resp: + if resp.status == 200: + result = await resp.json() + text = result.get("text", "") + logger.debug(f"Xinference STT result: {text}") + return text + error_text = await resp.text() + logger.error( + f"Xinference STT transcription failed with status {resp.status}: {error_text}", + ) + return "" + + except Exception as e: + logger.error(f"Xinference STT failed: {e}") + logger.debug(f"Xinference STT failed with exception: {e}", exc_info=True) + return "" + finally: + # 5. Cleanup + for temp_file in temp_files: + try: + if os.path.exists(temp_file): + os.remove(temp_file) + logger.debug(f"Removed temporary file: {temp_file}") + except Exception as e: + logger.error(f"Failed to remove temporary file {temp_file}: {e}") + + async def terminate(self) -> None: + """关闭客户端会话""" + if self.client: + logger.info("Closing Xinference STT client...") + try: + await self.client.close() + except Exception as e: + logger.error(f"Failed to close Xinference client: {e}", exc_info=True) diff --git a/astrbot/core/provider/sources/zhipu_source.py b/astrbot/core/provider/sources/zhipu_source.py index e7b6ee4f..ed4bc0bf 100644 --- a/astrbot/core/provider/sources/zhipu_source.py +++ b/astrbot/core/provider/sources/zhipu_source.py @@ -12,10 +12,5 @@ class ProviderZhipu(ProviderOpenAIOfficial): self, provider_config: dict, provider_settings: dict, - default_persona=None, ) -> None: - super().__init__( - provider_config, - provider_settings, - default_persona, - ) + super().__init__(provider_config, provider_settings) diff --git a/astrbot/core/star/__init__.py b/astrbot/core/star/__init__.py index 70e06d0d..e27db740 100644 --- a/astrbot/core/star/__init__.py +++ b/astrbot/core/star/__init__.py @@ -1,10 +1,11 @@ +from astrbot.core import html_renderer +from astrbot.core.provider import Provider +from astrbot.core.star.star_tools import StarTools +from astrbot.core.utils.command_parser import CommandParserMixin + +from .context import Context from .star import StarMetadata, star_map, star_registry from .star_manager import PluginManager -from .context import Context -from astrbot.core.provider import Provider -from astrbot.core.utils.command_parser import CommandParserMixin -from astrbot.core import html_renderer -from astrbot.core.star.star_tools import StarTools class Star(CommandParserMixin): @@ -36,24 +37,28 @@ class Star(CommandParserMixin): ) async def html_render( - self, tmpl: str, data: dict, return_url=True, options: dict | None = None + self, + tmpl: str, + data: dict, + return_url=True, + options: dict | None = None, ) -> str: """渲染 HTML""" return await html_renderer.render_custom_template( - tmpl, data, return_url=return_url, options=options + tmpl, + data, + return_url=return_url, + options=options, ) async def initialize(self): """当插件被激活时会调用这个方法""" - pass async def terminate(self): """当插件被禁用、重载插件时会调用这个方法""" - pass def __del__(self): """[Deprecated] 当插件被禁用、重载插件时会调用这个方法""" - pass -__all__ = ["Star", "StarMetadata", "PluginManager", "Context", "Provider", "StarTools"] +__all__ = ["Context", "PluginManager", "Provider", "Star", "StarMetadata", "StarTools"] diff --git a/astrbot/core/star/config.py b/astrbot/core/star/config.py index 23a522dc..a9af974c 100644 --- a/astrbot/core/star/config.py +++ b/astrbot/core/star/config.py @@ -1,23 +1,20 @@ -""" -此功能已过时，参考 https://astrbot.app/dev/plugin.html#%E6%B3%A8%E5%86%8C%E6%8F%92%E4%BB%B6%E9%85%8D%E7%BD%AE-beta -""" +"""此功能已过时，参考 https://astrbot.app/dev/plugin.html#%E6%B3%A8%E5%86%8C%E6%8F%92%E4%BB%B6%E9%85%8D%E7%BD%AE-beta""" -from typing import Union -import os import json +import os + from astrbot.core.utils.astrbot_path import get_astrbot_data_path -def load_config(namespace: str) -> Union[dict, bool]: - """ - 从配置文件中加载配置。 +def load_config(namespace: str) -> dict | bool: + """从配置文件中加载配置。 namespace: str, 配置的唯一识别符，也就是配置文件的名字。 返回值: 当配置文件存在时，返回 namespace 对应配置文件的内容dict，否则返回 False。 """ path = os.path.join(get_astrbot_data_path(), "config", f"{namespace}.json") if not os.path.exists(path): return False - with open(path, "r", encoding="utf-8-sig") as f: + with open(path, encoding="utf-8-sig") as f: ret = {} data = json.load(f) for k in data: @@ -26,8 +23,7 @@ def load_config(namespace: str) -> Union[dict, bool]: def put_config(namespace: str, name: str, key: str, value, description: str): - """ - 将配置项写入以namespace为名字的配置文件，如果key不存在于目标配置文件中。当前 value 仅支持 str, int, float, bool, list 类型（暂不支持 dict）。 + """将配置项写入以namespace为名字的配置文件，如果key不存在于目标配置文件中。当前 value 仅支持 str, int, float, bool, list 类型（暂不支持 dict）。 namespace: str, 配置的唯一识别符，也就是配置文件的名字。 name: str, 配置项的显示名字。 key: str, 配置项的键。 @@ -51,7 +47,7 @@ def put_config(namespace: str, name: str, key: str, value, description: str): if not os.path.exists(path): with open(path, "w", encoding="utf-8-sig") as f: f.write("{}") - with open(path, "r", encoding="utf-8-sig") as f: + with open(path, encoding="utf-8-sig") as f: d = json.load(f) assert isinstance(d, dict) if key not in d: @@ -69,8 +65,7 @@ def put_config(namespace: str, name: str, key: str, value, description: str): def update_config(namespace: str, key: str, value): - """ - 更新配置文件中的配置项。 + """更新配置文件中的配置项。 namespace: str, 配置的唯一识别符，也就是配置文件的名字。 key: str, 配置项的键。 value: str, int, float, bool, list, 配置项的值。 @@ -78,7 +73,7 @@ def update_config(namespace: str, key: str, value): path = os.path.join(get_astrbot_data_path(), "config", f"{namespace}.json") if not os.path.exists(path): raise FileNotFoundError(f"配置文件 {namespace}.json 不存在。") - with open(path, "r", encoding="utf-8-sig") as f: + with open(path, encoding="utf-8-sig") as f: d = json.load(f) assert isinstance(d, dict) if key not in d: diff --git a/astrbot/core/star/context.py b/astrbot/core/star/context.py index 31616e7d..21c1ad8f 100644 --- a/astrbot/core/star/context.py +++ b/astrbot/core/star/context.py @@ -1,47 +1,56 @@ +import logging from asyncio import Queue -from typing import List, Union +from collections.abc import Awaitable, Callable +from typing import Any -from astrbot.core.provider.provider import ( - Provider, - TTSProvider, - STTProvider, - EmbeddingProvider, - RerankProvider, -) -from astrbot.core.provider.entities import ProviderType -from astrbot.core.db import BaseDatabase +from deprecated import deprecated + +from astrbot.core.agent.hooks import BaseAgentRunHooks +from astrbot.core.agent.message import Message +from astrbot.core.agent.runners.tool_loop_agent_runner import ToolLoopAgentRunner +from astrbot.core.agent.tool import ToolSet +from astrbot.core.astrbot_config_mgr import AstrBotConfigManager from astrbot.core.config.astrbot_config import AstrBotConfig -from astrbot.core.provider.func_tool_manager import FunctionToolManager -from astrbot.core.platform.astr_message_event import MessageSesion +from astrbot.core.conversation_mgr import ConversationManager +from astrbot.core.db import BaseDatabase +from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager from astrbot.core.message.message_event_result import MessageChain -from astrbot.core.provider.manager import ProviderManager +from astrbot.core.persona_mgr import PersonaManager from astrbot.core.platform import Platform +from astrbot.core.platform.astr_message_event import AstrMessageEvent, MessageSesion from astrbot.core.platform.manager import PlatformManager from astrbot.core.platform_message_history_mgr import PlatformMessageHistoryManager -from astrbot.core.astrbot_config_mgr import AstrBotConfigManager -from astrbot.core.persona_mgr import PersonaManager -from .star import star_registry, StarMetadata, star_map -from .star_handler import star_handlers_registry, StarHandlerMetadata, EventType +from astrbot.core.provider.entities import LLMResponse, ProviderRequest, ProviderType +from astrbot.core.provider.func_tool_manager import FunctionTool, FunctionToolManager +from astrbot.core.provider.manager import ProviderManager +from astrbot.core.provider.provider import ( + EmbeddingProvider, + Provider, + RerankProvider, + STTProvider, + TTSProvider, +) +from astrbot.core.star.filter.platform_adapter_type import ( + ADAPTER_NAME_2_TYPE, + PlatformAdapterType, +) + +from ..exceptions import ProviderNotFoundError from .filter.command import CommandFilter from .filter.regex import RegexFilter -from typing import Awaitable, Any, Callable -from astrbot.core.conversation_mgr import ConversationManager -from astrbot.core.star.filter.platform_adapter_type import ( - PlatformAdapterType, - ADAPTER_NAME_2_TYPE, -) -from deprecated import deprecated +from .star import StarMetadata, star_map, star_registry +from .star_handler import EventType, StarHandlerMetadata, star_handlers_registry + +logger = logging.getLogger("astrbot") class Context: - """ - 暴露给插件的接口上下文。 - """ + """暴露给插件的接口上下文。""" registered_web_apis: list = [] # back compatibility - _register_tasks: List[Awaitable] = [] + _register_tasks: list[Awaitable] = [] _star_manager = None def __init__( @@ -55,6 +64,7 @@ class Context: message_history_manager: PlatformMessageHistoryManager, persona_manager: PersonaManager, astrbot_config_mgr: AstrBotConfigManager, + knowledge_base_manager: KnowledgeBaseManager, ): self._event_queue = event_queue """事件队列。消息平台通过事件队列传递消息事件。""" @@ -68,6 +78,154 @@ class Context: self.message_history_manager = message_history_manager self.persona_manager = persona_manager self.astrbot_config_mgr = astrbot_config_mgr + self.kb_manager = knowledge_base_manager + + async def llm_generate( + self, + *, + chat_provider_id: str, + prompt: str | None = None, + image_urls: list[str] | None = None, + tools: ToolSet | None = None, + system_prompt: str | None = None, + contexts: list[Message] | None = None, + **kwargs: Any, + ) -> LLMResponse: + """Call the LLM to generate a response. The method will not automatically execute tool calls. If you want to use tool calls, please use `tool_loop_agent()`. + + .. versionadded:: 4.5.7 (sdk) + + Args: + chat_provider_id: The chat provider ID to use. + prompt: The prompt to send to the LLM, if `contexts` and `prompt` are both provided, `prompt` will be appended as the last user message + image_urls: List of image URLs to include in the prompt, if `contexts` and `prompt` are both provided, `image_urls` will be appended to the last user message + tools: ToolSet of tools available to the LLM + system_prompt: System prompt to guide the LLM's behavior, if provided, it will always insert as the first system message in the context + contexts: context messages for the LLM + **kwargs: Additional keyword arguments for LLM generation, OpenAI compatible + + Raises: + ChatProviderNotFoundError: If the specified chat provider ID is not found + Exception: For other errors during LLM generation + """ + prov = await self.provider_manager.get_provider_by_id(chat_provider_id) + if not prov or not isinstance(prov, Provider): + raise ProviderNotFoundError(f"Provider {chat_provider_id} not found") + llm_resp = await prov.text_chat( + prompt=prompt, + image_urls=image_urls, + func_tool=tools, + contexts=contexts, + system_prompt=system_prompt, + **kwargs, + ) + return llm_resp + + async def tool_loop_agent( + self, + *, + event: AstrMessageEvent, + chat_provider_id: str, + prompt: str | None = None, + image_urls: list[str] | None = None, + tools: ToolSet | None = None, + system_prompt: str | None = None, + contexts: list[Message] | None = None, + max_steps: int = 30, + tool_call_timeout: int = 60, + **kwargs: Any, + ) -> LLMResponse: + """Run an agent loop that allows the LLM to call tools iteratively until a final answer is produced. + If you do not pass the agent_context parameter, the method will recreate a new agent context. + + .. versionadded:: 4.5.7 (sdk) + + Args: + chat_provider_id: The chat provider ID to use. + prompt: The prompt to send to the LLM, if `contexts` and `prompt` are both provided, `prompt` will be appended as the last user message + image_urls: List of image URLs to include in the prompt, if `contexts` and `prompt` are both provided, `image_urls` will be appended to the last user message + tools: ToolSet of tools available to the LLM + system_prompt: System prompt to guide the LLM's behavior, if provided, it will always insert as the first system message in the context + contexts: context messages for the LLM + max_steps: Maximum number of tool calls before stopping the loop + **kwargs: Additional keyword arguments. The kwargs will not be passed to the LLM directly for now, but can include: + agent_hooks: BaseAgentRunHooks[AstrAgentContext] - hooks to run during agent execution + agent_context: AstrAgentContext - context to use for the agent + + Returns: + The final LLMResponse after tool calls are completed. + + Raises: + ChatProviderNotFoundError: If the specified chat provider ID is not found + Exception: For other errors during LLM generation + """ + # Import here to avoid circular imports + from astrbot.core.astr_agent_context import ( + AgentContextWrapper, + AstrAgentContext, + ) + from astrbot.core.astr_agent_tool_exec import FunctionToolExecutor + + prov = await self.provider_manager.get_provider_by_id(chat_provider_id) + if not prov or not isinstance(prov, Provider): + raise ProviderNotFoundError(f"Provider {chat_provider_id} not found") + + agent_hooks = kwargs.get("agent_hooks") or BaseAgentRunHooks[AstrAgentContext]() + agent_context = kwargs.get("agent_context") + + context_ = [] + for msg in contexts or []: + if isinstance(msg, Message): + context_.append(msg.model_dump()) + else: + context_.append(msg) + + request = ProviderRequest( + prompt=prompt, + image_urls=image_urls or [], + func_tool=tools, + contexts=context_, + system_prompt=system_prompt or "", + ) + if agent_context is None: + agent_context = AstrAgentContext( + context=self, + event=event, + ) + agent_runner = ToolLoopAgentRunner() + tool_executor = FunctionToolExecutor() + await agent_runner.reset( + provider=prov, + request=request, + run_context=AgentContextWrapper( + context=agent_context, + tool_call_timeout=tool_call_timeout, + ), + tool_executor=tool_executor, + agent_hooks=agent_hooks, + streaming=kwargs.get("stream", False), + ) + async for _ in agent_runner.step_until_done(max_steps): + pass + llm_resp = agent_runner.get_final_llm_resp() + if not llm_resp: + raise Exception("Agent did not produce a final LLM response") + return llm_resp + + async def get_current_chat_provider_id(self, umo: str) -> str: + """Get the ID of the currently used chat provider. + + Args: + umo(str): unified_message_origin value, if provided and user has enabled provider session isolation, the provider preferred by that session will be used. + + Raises: + ProviderNotFoundError: If the specified chat provider is not found + + """ + prov = self.get_using_provider(umo) + if not prov: + raise ProviderNotFoundError("Provider not found") + return prov.meta().id def get_registered_star(self, star_name: str) -> StarMetadata | None: """根据插件名获取插件的 Metadata""" @@ -75,7 +233,7 @@ class Context: if star.name == star_name: return star - def get_all_stars(self) -> List[StarMetadata]: + def get_all_stars(self) -> list[StarMetadata]: """获取当前载入的所有插件 Metadata 的列表""" return star_registry @@ -88,6 +246,7 @@ class Context: Returns: 如果没找到，会返回 False + """ return self.provider_manager.llm_tools.activate_llm_tool(name, star_map) @@ -95,17 +254,14 @@ class Context: """停用一个已经注册的函数调用工具。 Returns: - 如果没找到，会返回 False""" + 如果没找到，会返回 False + + """ return self.provider_manager.llm_tools.deactivate_llm_tool(name) - def register_provider(self, provider: Provider): - """ - 注册一个 LLM Provider(Chat_Completion 类型)。 - """ - self.provider_manager.provider_insts.append(provider) - def get_provider_by_id( - self, provider_id: str + self, + provider_id: str, ) -> ( Provider | TTSProvider | STTProvider | EmbeddingProvider | RerankProvider | None ): @@ -113,28 +269,28 @@ class Context: prov = self.provider_manager.inst_map.get(provider_id) return prov - def get_all_providers(self) -> List[Provider]: + def get_all_providers(self) -> list[Provider]: """获取所有用于文本生成任务的 LLM Provider(Chat_Completion 类型)。""" return self.provider_manager.provider_insts - def get_all_tts_providers(self) -> List[TTSProvider]: + def get_all_tts_providers(self) -> list[TTSProvider]: """获取所有用于 TTS 任务的 Provider。""" return self.provider_manager.tts_provider_insts - def get_all_stt_providers(self) -> List[STTProvider]: + def get_all_stt_providers(self) -> list[STTProvider]: """获取所有用于 STT 任务的 Provider。""" return self.provider_manager.stt_provider_insts - def get_all_embedding_providers(self) -> List[EmbeddingProvider]: + def get_all_embedding_providers(self) -> list[EmbeddingProvider]: """获取所有用于 Embedding 任务的 Provider。""" return self.provider_manager.embedding_provider_insts def get_using_provider(self, umo: str | None = None) -> Provider | None: - """ - 获取当前使用的用于文本生成任务的 LLM Provider(Chat_Completion 类型)。通过 /provider 指令切换。 + """获取当前使用的用于文本生成任务的 LLM Provider(Chat_Completion 类型)。通过 /provider 指令切换。 Args: umo(str): unified_message_origin 值，如果传入并且用户启用了提供商会话隔离，则使用该会话偏好的提供商。 + """ prov = self.provider_manager.get_using_provider( provider_type=ProviderType.CHAT_COMPLETION, @@ -145,11 +301,11 @@ class Context: return prov def get_using_tts_provider(self, umo: str | None = None) -> TTSProvider | None: - """ - 获取当前使用的用于 TTS 任务的 Provider。 + """获取当前使用的用于 TTS 任务的 Provider。 Args: umo(str): unified_message_origin 值，如果传入，则使用该会话偏好的提供商。 + """ prov = self.provider_manager.get_using_provider( provider_type=ProviderType.TEXT_TO_SPEECH, @@ -160,11 +316,11 @@ class Context: return prov def get_using_stt_provider(self, umo: str | None = None) -> STTProvider | None: - """ - 获取当前使用的用于 STT 任务的 Provider。 + """获取当前使用的用于 STT 任务的 Provider。 Args: umo(str): unified_message_origin 值，如果传入，则使用该会话偏好的提供商。 + """ prov = self.provider_manager.get_using_provider( provider_type=ProviderType.SPEECH_TO_TEXT, @@ -179,59 +335,14 @@ class Context: if not umo: # using default config return self._config - else: - return self.astrbot_config_mgr.get_conf(umo) - - def get_db(self) -> BaseDatabase: - """获取 AstrBot 数据库。""" - return self._db - - def get_event_queue(self) -> Queue: - """ - 获取事件队列。 - """ - return self._event_queue - - @deprecated(version="4.0.0", reason="Use get_platform_inst instead") - def get_platform( - self, platform_type: Union[PlatformAdapterType, str] - ) -> Platform | None: - """ - 获取指定类型的平台适配器。 - - 该方法已经过时，请使用 get_platform_inst 方法。(>= AstrBot v4.0.0) - """ - for platform in self.platform_manager.platform_insts: - name = platform.meta().name - if isinstance(platform_type, str): - if name == platform_type: - return platform - else: - if ( - name in ADAPTER_NAME_2_TYPE - and ADAPTER_NAME_2_TYPE[name] & platform_type - ): - return platform - - def get_platform_inst(self, platform_id: str) -> Platform | None: - """ - 获取指定 ID 的平台适配器实例。 - - Args: - platform_id (str): 平台适配器的唯一标识符。你可以通过 event.get_platform_id() 获取。 - - Returns: - Platform: 平台适配器实例，如果未找到则返回 None。 - """ - for platform in self.platform_manager.platform_insts: - if platform.meta().id == platform_id: - return platform + return self.astrbot_config_mgr.get_conf(umo) async def send_message( - self, session: Union[str, MessageSesion], message_chain: MessageChain + self, + session: str | MessageSesion, + message_chain: MessageChain, ) -> bool: - """ - 根据 session(unified_msg_origin) 主动发送消息。 + """根据 session(unified_msg_origin) 主动发送消息。 @param session: 消息会话。通过 event.session 或者 event.unified_msg_origin 获取。 @param message_chain: 消息链。 @@ -242,7 +353,6 @@ class Context: NOTE: qq_official(QQ 官方 API 平台) 不支持此方法 """ - if isinstance(session, str): try: session = MessageSesion.from_str(session) @@ -255,10 +365,93 @@ class Context: return True return False + def add_llm_tools(self, *tools: FunctionTool) -> None: + """添加 LLM 工具。""" + tool_name = {tool.name for tool in self.provider_manager.llm_tools.func_list} + module_path = "" + for tool in tools: + if not module_path: + _parts = [] + module_part = tool.__module__.split(".") + flags = ["packages", "plugins"] + for i, part in enumerate(module_part): + _parts.append(part) + if part in flags and i + 1 < len(module_part): + _parts.append(module_part[i + 1]) + break + tool.handler_module_path = ".".join(_parts) + module_path = tool.handler_module_path + else: + tool.handler_module_path = module_path + logger.info( + f"plugin(module_path {module_path}) added LLM tool: {tool.name}" + ) + + if tool.name in tool_name: + logger.warning("替换已存在的 LLM 工具: " + tool.name) + self.provider_manager.llm_tools.remove_func(tool.name) + self.provider_manager.llm_tools.func_list.append(tool) + + def register_web_api( + self, + route: str, + view_handler: Awaitable, + methods: list, + desc: str, + ): + for idx, api in enumerate(self.registered_web_apis): + if api[0] == route and methods == api[2]: + self.registered_web_apis[idx] = (route, view_handler, methods, desc) + return + self.registered_web_apis.append((route, view_handler, methods, desc)) + """ 以下的方法已经不推荐使用。请从 AstrBot 文档查看更好的注册方式。 """ + def get_event_queue(self) -> Queue: + """获取事件队列。""" + return self._event_queue + + @deprecated(version="4.0.0", reason="Use get_platform_inst instead") + def get_platform(self, platform_type: PlatformAdapterType | str) -> Platform | None: + """获取指定类型的平台适配器。 + + 该方法已经过时，请使用 get_platform_inst 方法。(>= AstrBot v4.0.0) + """ + for platform in self.platform_manager.platform_insts: + name = platform.meta().name + if isinstance(platform_type, str): + if name == platform_type: + return platform + elif ( + name in ADAPTER_NAME_2_TYPE + and ADAPTER_NAME_2_TYPE[name] & platform_type + ): + return platform + + def get_platform_inst(self, platform_id: str) -> Platform | None: + """获取指定 ID 的平台适配器实例。 + + Args: + platform_id (str): 平台适配器的唯一标识符。你可以通过 event.get_platform_id() 获取。 + + Returns: + Platform: 平台适配器实例，如果未找到则返回 None。 + + """ + for platform in self.platform_manager.platform_insts: + if platform.meta().id == platform_id: + return platform + + def get_db(self) -> BaseDatabase: + """获取 AstrBot 数据库。""" + return self._db + + def register_provider(self, provider: Provider): + """注册一个 LLM Provider(Chat_Completion 类型)。""" + self.provider_manager.provider_insts.append(provider) + def register_llm_tool( self, name: str, @@ -266,8 +459,7 @@ class Context: desc: str, func_obj: Callable[..., Awaitable[Any]], ) -> None: - """ - 为函数调用（function-calling / tools-use）添加工具。 + """[DEPRECATED]为函数调用（function-calling / tools-use）添加工具。 @param name: 函数名 @param func_args: 函数参数列表，格式为 [{"type": "string", "name": "arg_name", "description": "arg_description"}, ...] @@ -289,7 +481,7 @@ class Context: self.provider_manager.llm_tools.add_func(name, func_args, desc, func_obj) def unregister_llm_tool(self, name: str) -> None: - """删除一个函数调用工具。如果再要启用，需要重新注册。""" + """[DEPRECATED]删除一个函数调用工具。如果再要启用，需要重新注册。""" self.provider_manager.llm_tools.remove_func(name) def register_commands( @@ -302,8 +494,7 @@ class Context: use_regex=False, ignore_prefix=False, ): - """ - 注册一个命令。 + """注册一个命令。 [Deprecated] 推荐使用装饰器注册指令。该方法将在未来的版本中被移除。 @@ -327,21 +518,10 @@ class Context: md.event_filters.append(RegexFilter(regex=command_name)) else: md.event_filters.append( - CommandFilter(command_name=command_name, handler_md=md) + CommandFilter(command_name=command_name, handler_md=md), ) star_handlers_registry.append(md) def register_task(self, task: Awaitable, desc: str): - """ - 注册一个异步任务。 - """ + """[DEPRECATED]注册一个异步任务。""" self._register_tasks.append(task) - - def register_web_api( - self, route: str, view_handler: Awaitable, methods: list, desc: str - ): - for idx, api in enumerate(self.registered_web_apis): - if api[0] == route and methods == api[2]: - self.registered_web_apis[idx] = (route, view_handler, methods, desc) - return - self.registered_web_apis.append((route, view_handler, methods, desc)) diff --git a/astrbot/core/star/filter/__init__.py b/astrbot/core/star/filter/__init__.py index c2f78e27..e550017a 100644 --- a/astrbot/core/star/filter/__init__.py +++ b/astrbot/core/star/filter/__init__.py @@ -1,7 +1,8 @@ import abc -from astrbot.core.platform.message_type import MessageType -from astrbot.core.platform.astr_message_event import AstrMessageEvent + from astrbot.core.config import AstrBotConfig +from astrbot.core.platform.astr_message_event import AstrMessageEvent +from astrbot.core.platform.message_type import MessageType class HandlerFilter(abc.ABC): @@ -11,4 +12,4 @@ class HandlerFilter(abc.ABC): raise NotImplementedError -__all__ = ["HandlerFilter", "MessageType", "AstrMessageEvent", "AstrBotConfig"] +__all__ = ["AstrBotConfig", "AstrMessageEvent", "HandlerFilter", "MessageType"] diff --git a/astrbot/core/star/filter/command.py b/astrbot/core/star/filter/command.py index 3d67cb75..2a9868fd 100755 --- a/astrbot/core/star/filter/command.py +++ b/astrbot/core/star/filter/command.py @@ -1,20 +1,20 @@ -import re import inspect +import re import types import typing -from typing import List, Any, Type, Dict -from . import HandlerFilter -from astrbot.core.platform.astr_message_event import AstrMessageEvent +from typing import Any + from astrbot.core.config import AstrBotConfig -from .custom_filter import CustomFilter +from astrbot.core.platform.astr_message_event import AstrMessageEvent + from ..star_handler import StarHandlerMetadata +from . import HandlerFilter +from .custom_filter import CustomFilter class GreedyStr(str): """标记指令完成其他参数接收后的所有剩余文本。""" - pass - def unwrap_optional(annotation) -> tuple: """去掉 Optional[T] / Union[T, None] / T|None，返回 T""" @@ -22,10 +22,9 @@ def unwrap_optional(annotation) -> tuple: non_none_args = [a for a in args if a is not type(None)] if len(non_none_args) == 1: return (non_none_args[0],) - elif len(non_none_args) > 1: + if len(non_none_args) > 1: return tuple(non_none_args) - else: - return () + return () # 标准指令受到 wake_prefix 的制约。 @@ -37,28 +36,30 @@ class CommandFilter(HandlerFilter): command_name: str, alias: set | None = None, handler_md: StarHandlerMetadata | None = None, - parent_command_names: List[str] = [""], + parent_command_names: list[str] | None = None, ): self.command_name = command_name self.alias = alias if alias else set() - self.parent_command_names = parent_command_names + self.parent_command_names = ( + parent_command_names if parent_command_names is not None else [""] + ) if handler_md: self.init_handler_md(handler_md) - self.custom_filter_list: List[CustomFilter] = [] + self.custom_filter_list: list[CustomFilter] = [] # Cache for complete command names list self._cmpl_cmd_names: list | None = None def print_types(self): - result = "" + parts = [] for k, v in self.handler_params.items(): if isinstance(v, type): - result += f"{k}({v.__name__})," + parts.append(f"{k}({v.__name__}),") elif isinstance(v, types.UnionType) or typing.get_origin(v) is typing.Union: - result += f"{k}({v})," + parts.append(f"{k}({v}),") else: - result += f"{k}({type(v).__name__})={v}," - result = result.rstrip(",") + parts.append(f"{k}({type(v).__name__})={v},") + result = "".join(parts).rstrip(",") return result def init_handler_md(self, handle_md: StarHandlerMetadata): @@ -89,8 +90,10 @@ class CommandFilter(HandlerFilter): return True def validate_and_convert_params( - self, params: List[Any], param_type: Dict[str, Type] - ) -> Dict[str, Any]: + self, + params: list[Any], + param_type: dict[str, type], + ) -> dict[str, Any]: """将参数列表 params 根据 param_type 转换为参数字典。""" result = {} param_items = list(param_type.items()) @@ -101,7 +104,7 @@ class CommandFilter(HandlerFilter): # GreedyStr 必须是最后一个参数 if i != len(param_items) - 1: raise ValueError( - f"参数 '{param_name}' (GreedyStr) 必须是最后一个参数。" + f"参数 '{param_name}' (GreedyStr) 必须是最后一个参数。", ) # 将剩余的所有部分合并成一个字符串 @@ -111,17 +114,16 @@ class CommandFilter(HandlerFilter): # 没有 GreedyStr 的情况 if i >= len(params): if ( - isinstance(param_type_or_default_val, (Type, types.UnionType)) + isinstance(param_type_or_default_val, (type, types.UnionType)) or typing.get_origin(param_type_or_default_val) is typing.Union or param_type_or_default_val is inspect.Parameter.empty ): # 是类型 raise ValueError( - f"必要参数缺失。该指令完整参数: {self.print_types()}" + f"必要参数缺失。该指令完整参数: {self.print_types()}", ) - else: - # 是默认值 - result[param_name] = param_type_or_default_val + # 是默认值 + result[param_name] = param_type_or_default_val else: # 尝试强制转换 try: @@ -142,7 +144,7 @@ class CommandFilter(HandlerFilter): result[param_name] = False else: raise ValueError( - f"参数 {param_name} 必须是布尔值（true/false, yes/no, 1/0）。" + f"参数 {param_name} 必须是布尔值（true/false, yes/no, 1/0）。", ) elif isinstance(param_type_or_default_val, int): result[param_name] = int(params[i]) @@ -165,7 +167,7 @@ class CommandFilter(HandlerFilter): result[param_name] = param_type_or_default_val(params[i]) except ValueError: raise ValueError( - f"参数 {param_name} 类型错误。完整参数: {self.print_types()}" + f"参数 {param_name} 类型错误。完整参数: {self.print_types()}", ) return result diff --git a/astrbot/core/star/filter/command_group.py b/astrbot/core/star/filter/command_group.py index e01fa2c5..e1c2efb2 100755 --- a/astrbot/core/star/filter/command_group.py +++ b/astrbot/core/star/filter/command_group.py @@ -1,10 +1,10 @@ from __future__ import annotations -from typing import List, Union +from astrbot.core.config import AstrBotConfig +from astrbot.core.platform.astr_message_event import AstrMessageEvent + from . import HandlerFilter from .command import CommandFilter -from astrbot.core.platform.astr_message_event import AstrMessageEvent -from astrbot.core.config import AstrBotConfig from .custom_filter import CustomFilter @@ -18,25 +18,27 @@ class CommandGroupFilter(HandlerFilter): ): self.group_name = group_name self.alias = alias if alias else set() - self.sub_command_filters: List[Union[CommandFilter, CommandGroupFilter]] = [] - self.custom_filter_list: List[CustomFilter] = [] + self.sub_command_filters: list[CommandFilter | CommandGroupFilter] = [] + self.custom_filter_list: list[CustomFilter] = [] self.parent_group = parent_group # Cache for complete command names list self._cmpl_cmd_names: list | None = None def add_sub_command_filter( - self, sub_command_filter: Union[CommandFilter, CommandGroupFilter] + self, + sub_command_filter: CommandFilter | CommandGroupFilter, ): self.sub_command_filters.append(sub_command_filter) def add_custom_filter(self, custom_filter: CustomFilter): self.custom_filter_list.append(custom_filter) - def get_complete_command_names(self) -> List[str]: + def get_complete_command_names(self) -> list[str]: """遍历父节点获取完整的指令名。 - 新版本 v3.4.29 采用预编译指令，不再从指令组递归遍历子指令，因此这个方法是返回包括别名在内的整个指令名列表。""" + 新版本 v3.4.29 采用预编译指令，不再从指令组递归遍历子指令，因此这个方法是返回包括别名在内的整个指令名列表。 + """ if self._cmpl_cmd_names is not None: return self._cmpl_cmd_names @@ -59,12 +61,12 @@ class CommandGroupFilter(HandlerFilter): # 以树的形式打印出来 def print_cmd_tree( self, - sub_command_filters: List[Union[CommandFilter, CommandGroupFilter]], + sub_command_filters: list[CommandFilter | CommandGroupFilter], prefix: str = "", event: AstrMessageEvent | None = None, cfg: AstrBotConfig | None = None, ) -> str: - result = "" + parts = [] for sub_filter in sub_command_filters: if isinstance(sub_filter, CommandFilter): custom_filter_pass = True @@ -72,31 +74,32 @@ class CommandGroupFilter(HandlerFilter): custom_filter_pass = sub_filter.custom_filter_ok(event, cfg) if custom_filter_pass: cmd_th = sub_filter.print_types() - result += f"{prefix}├── {sub_filter.command_name}" + line = f"{prefix}├── {sub_filter.command_name}" if cmd_th: - result += f" ({cmd_th})" + line += f" ({cmd_th})" else: - result += " (无参数指令)" + line += " (无参数指令)" if sub_filter.handler_md and sub_filter.handler_md.desc: - result += f": {sub_filter.handler_md.desc}" + line += f": {sub_filter.handler_md.desc}" - result += "\n" + parts.append(line + "\n") elif isinstance(sub_filter, CommandGroupFilter): custom_filter_pass = True if event and cfg: custom_filter_pass = sub_filter.custom_filter_ok(event, cfg) if custom_filter_pass: - result += f"{prefix}├── {sub_filter.group_name}" - result += "\n" - result += sub_filter.print_cmd_tree( - sub_filter.sub_command_filters, - prefix + "│ ", - event=event, - cfg=cfg, + parts.append(f"{prefix}├── {sub_filter.group_name}\n") + parts.append( + sub_filter.print_cmd_tree( + sub_filter.sub_command_filters, + prefix + "│ ", + event=event, + cfg=cfg, + ) ) - return result + return "".join(parts) def custom_filter_ok(self, event: AstrMessageEvent, cfg: AstrBotConfig) -> bool: for custom_filter in self.custom_filter_list: @@ -125,7 +128,7 @@ class CommandGroupFilter(HandlerFilter): + self.print_cmd_tree(self.sub_command_filters, event=event, cfg=cfg) ) raise ValueError( - f"参数不足。{self.group_name} 指令组下有如下指令，请参考：\n" + tree + f"参数不足。{self.group_name} 指令组下有如下指令，请参考：\n" + tree, ) return self.startswith(event.message_str) diff --git a/astrbot/core/star/filter/custom_filter.py b/astrbot/core/star/filter/custom_filter.py index 9a76b74f..d57b5cac 100644 --- a/astrbot/core/star/filter/custom_filter.py +++ b/astrbot/core/star/filter/custom_filter.py @@ -1,8 +1,9 @@ -from abc import abstractmethod, ABCMeta +from abc import ABCMeta, abstractmethod + +from astrbot.core.config import AstrBotConfig +from astrbot.core.platform.astr_message_event import AstrMessageEvent from . import HandlerFilter -from astrbot.core.platform.astr_message_event import AstrMessageEvent -from astrbot.core.config import AstrBotConfig class CustomFilterMeta(ABCMeta): @@ -38,7 +39,7 @@ class CustomFilterOr(CustomFilter): super().__init__() if not isinstance(filter1, (CustomFilter, CustomFilterAnd, CustomFilterOr)): raise ValueError( - "CustomFilter lass can only operate with other CustomFilter." + "CustomFilter lass can only operate with other CustomFilter.", ) self.filter1 = filter1 self.filter2 = filter2 @@ -52,7 +53,7 @@ class CustomFilterAnd(CustomFilter): super().__init__() if not isinstance(filter1, (CustomFilter, CustomFilterAnd, CustomFilterOr)): raise ValueError( - "CustomFilter lass can only operate with other CustomFilter." + "CustomFilter lass can only operate with other CustomFilter.", ) self.filter1 = filter1 self.filter2 = filter2 diff --git a/astrbot/core/star/filter/event_message_type.py b/astrbot/core/star/filter/event_message_type.py index ce36ec9e..7f350bd3 100644 --- a/astrbot/core/star/filter/event_message_type.py +++ b/astrbot/core/star/filter/event_message_type.py @@ -1,9 +1,11 @@ import enum -from . import HandlerFilter -from astrbot.core.platform.astr_message_event import AstrMessageEvent + from astrbot.core.config import AstrBotConfig +from astrbot.core.platform.astr_message_event import AstrMessageEvent from astrbot.core.platform.message_type import MessageType +from . import HandlerFilter + class EventMessageType(enum.Flag): GROUP_MESSAGE = enum.auto() diff --git a/astrbot/core/star/filter/permission.py b/astrbot/core/star/filter/permission.py index 307b492a..3374544c 100644 --- a/astrbot/core/star/filter/permission.py +++ b/astrbot/core/star/filter/permission.py @@ -1,7 +1,9 @@ import enum -from . import HandlerFilter -from astrbot.core.platform.astr_message_event import AstrMessageEvent + from astrbot.core.config import AstrBotConfig +from astrbot.core.platform.astr_message_event import AstrMessageEvent + +from . import HandlerFilter class PermissionType(enum.Flag): diff --git a/astrbot/core/star/filter/platform_adapter_type.py b/astrbot/core/star/filter/platform_adapter_type.py index 4c551078..1182ff9b 100644 --- a/astrbot/core/star/filter/platform_adapter_type.py +++ b/astrbot/core/star/filter/platform_adapter_type.py @@ -1,7 +1,9 @@ import enum -from . import HandlerFilter -from astrbot.core.platform.astr_message_event import AstrMessageEvent + from astrbot.core.config import AstrBotConfig +from astrbot.core.platform.astr_message_event import AstrMessageEvent + +from . import HandlerFilter class PlatformAdapterType(enum.Flag): diff --git a/astrbot/core/star/filter/regex.py b/astrbot/core/star/filter/regex.py index af9cb3a5..cd5bebdb 100644 --- a/astrbot/core/star/filter/regex.py +++ b/astrbot/core/star/filter/regex.py @@ -1,7 +1,9 @@ import re -from . import HandlerFilter -from astrbot.core.platform.astr_message_event import AstrMessageEvent + from astrbot.core.config import AstrBotConfig +from astrbot.core.platform.astr_message_event import AstrMessageEvent + +from . import HandlerFilter # 正则表达式过滤器不会受到 wake_prefix 的制约。 diff --git a/astrbot/core/star/register/__init__.py b/astrbot/core/star/register/__init__.py index 0519e8ca..15fe1e9c 100644 --- a/astrbot/core/star/register/__init__.py +++ b/astrbot/core/star/register/__init__.py @@ -1,37 +1,37 @@ from .star import register_star from .star_handler import ( + register_after_message_sent, + register_agent, register_command, register_command_group, - register_event_message_type, - register_platform_adapter_type, - register_regex, - register_permission_type, register_custom_filter, + register_event_message_type, + register_llm_tool, register_on_astrbot_loaded, - register_on_platform_loaded, + register_on_decorating_result, register_on_llm_request, register_on_llm_response, - register_llm_tool, - register_agent, - register_on_decorating_result, - register_after_message_sent, + register_on_platform_loaded, + register_permission_type, + register_platform_adapter_type, + register_regex, ) __all__ = [ - "register_star", + "register_after_message_sent", + "register_agent", "register_command", "register_command_group", - "register_event_message_type", - "register_platform_adapter_type", - "register_regex", - "register_permission_type", "register_custom_filter", + "register_event_message_type", + "register_llm_tool", "register_on_astrbot_loaded", - "register_on_platform_loaded", + "register_on_decorating_result", "register_on_llm_request", "register_on_llm_response", - "register_llm_tool", - "register_agent", - "register_on_decorating_result", - "register_after_message_sent", + "register_on_platform_loaded", + "register_permission_type", + "register_platform_adapter_type", + "register_regex", + "register_star", ] diff --git a/astrbot/core/star/register/star.py b/astrbot/core/star/register/star.py index a5190dd5..617cd5ff 100644 --- a/astrbot/core/star/register/star.py +++ b/astrbot/core/star/register/star.py @@ -6,7 +6,11 @@ _warned_register_star = False def register_star( - name: str, author: str, desc: str, version: str, repo: str | None = None + name: str, + author: str, + desc: str, + version: str, + repo: str | None = None, ): """注册一个插件(Star)。 @@ -29,8 +33,8 @@ def register_star( ... 帮助信息会被自动提取。使用 `/plugin <插件名> 可以查看帮助信息。` - """ + """ global _warned_register_star if not _warned_register_star: _warned_register_star = True diff --git a/astrbot/core/star/register/star_handler.py b/astrbot/core/star/register/star_handler.py index d1c5a6dc..ee3c0968 100644 --- a/astrbot/core/star/register/star_handler.py +++ b/astrbot/core/star/register/star_handler.py @@ -1,26 +1,31 @@ from __future__ import annotations + +import re +from collections.abc import Awaitable, Callable +from typing import Any + import docstring_parser -from ..star_handler import star_handlers_registry, StarHandlerMetadata, EventType -from ..filter.command import CommandFilter -from ..filter.command_group import CommandGroupFilter -from ..filter.event_message_type import EventMessageTypeFilter, EventMessageType -from ..filter.platform_adapter_type import ( - PlatformAdapterTypeFilter, - PlatformAdapterType, -) -from ..filter.permission import PermissionTypeFilter, PermissionType -from ..filter.custom_filter import CustomFilterAnd, CustomFilterOr -from ..filter.regex import RegexFilter -from typing import Awaitable, Any, Callable -from astrbot.core.provider.func_tool_manager import SUPPORTED_TYPES -from astrbot.core.provider.register import llm_tools +from astrbot.core import logger from astrbot.core.agent.agent import Agent -from astrbot.core.agent.tool import FunctionTool from astrbot.core.agent.handoff import HandoffTool from astrbot.core.agent.hooks import BaseAgentRunHooks +from astrbot.core.agent.tool import FunctionTool from astrbot.core.astr_agent_context import AstrAgentContext -from astrbot.core import logger +from astrbot.core.provider.func_tool_manager import PY_TO_JSON_TYPE, SUPPORTED_TYPES +from astrbot.core.provider.register import llm_tools + +from ..filter.command import CommandFilter +from ..filter.command_group import CommandGroupFilter +from ..filter.custom_filter import CustomFilterAnd, CustomFilterOr +from ..filter.event_message_type import EventMessageType, EventMessageTypeFilter +from ..filter.permission import PermissionType, PermissionTypeFilter +from ..filter.platform_adapter_type import ( + PlatformAdapterType, + PlatformAdapterTypeFilter, +) +from ..filter.regex import RegexFilter +from ..star_handler import EventType, StarHandlerMetadata, star_handlers_registry def get_handler_full_name(awaitable: Callable[..., Awaitable[Any]]) -> str: @@ -39,27 +44,26 @@ def get_handler_or_create( md = star_handlers_registry.get_handler_by_full_name(handler_full_name) if md: return md - else: - md = StarHandlerMetadata( - event_type=event_type, - handler_full_name=handler_full_name, - handler_name=handler.__name__, - handler_module_path=handler.__module__, - handler=handler, - event_filters=[], - ) + md = StarHandlerMetadata( + event_type=event_type, + handler_full_name=handler_full_name, + handler_name=handler.__name__, + handler_module_path=handler.__module__, + handler=handler, + event_filters=[], + ) - # 插件handler的附加额外信息 - if handler.__doc__: - md.desc = handler.__doc__.strip() - if "desc" in kwargs: - md.desc = kwargs["desc"] - del kwargs["desc"] - md.extras_configs = kwargs + # 插件handler的附加额外信息 + if handler.__doc__: + md.desc = handler.__doc__.strip() + if "desc" in kwargs: + md.desc = kwargs["desc"] + del kwargs["desc"] + md.extras_configs = kwargs - if not dont_add: - star_handlers_registry.append(md) - return md + if not dont_add: + star_handlers_registry.append(md) + return md def register_command( @@ -78,20 +82,22 @@ def register_command( command_name.parent_group.get_complete_command_names() ) new_command = CommandFilter( - sub_command, alias, None, parent_command_names=parent_command_names + sub_command, + alias, + None, + parent_command_names=parent_command_names, ) command_name.parent_group.add_sub_command_filter(new_command) else: logger.warning( - f"注册指令{command_name} 的子指令时未提供 sub_command 参数。" + f"注册指令{command_name} 的子指令时未提供 sub_command 参数。", ) + # 裸指令 + elif command_name is None: + logger.warning("注册裸指令时未提供 command_name 参数。") else: - # 裸指令 - if command_name is None: - logger.warning("注册裸指令时未提供 command_name 参数。") - else: - new_command = CommandFilter(command_name, alias, None) - add_to_event_filters = True + new_command = CommandFilter(command_name, alias, None) + add_to_event_filters = True def decorator(awaitable): if not add_to_event_filters: @@ -99,7 +105,9 @@ def register_command( True # 打一个标记，表示这是一个子指令，再 wakingstage 阶段这个 handler 将会直接被跳过（其父指令会接管） ) handler_md = get_handler_or_create( - awaitable, EventType.AdapterMessageEvent, **kwargs + awaitable, + EventType.AdapterMessageEvent, + **kwargs, ) if new_command: new_command.init_handler_md(handler_md) @@ -116,6 +124,7 @@ def register_custom_filter(custom_type_filter, *args, **kwargs): custom_type_filter: 在裸指令时为CustomFilter对象 在指令组时为父指令的RegisteringCommandable对象，即self或者command_group的返回 raise_error: 如果没有权限，是否抛出错误到消息平台，并且停止事件传播。默认为 True + """ add_to_event_filters = False raise_error = True @@ -140,19 +149,20 @@ def register_custom_filter(custom_type_filter, *args, **kwargs): def decorator(awaitable): # 裸指令，子指令与指令组的区分，指令组会因为标记跳过wake。 if ( - not add_to_event_filters - and isinstance(awaitable, RegisteringCommandable) - or (add_to_event_filters and isinstance(awaitable, RegisteringCommandable)) - ): + not add_to_event_filters and isinstance(awaitable, RegisteringCommandable) + ) or (add_to_event_filters and isinstance(awaitable, RegisteringCommandable)): # 指令组 与 根指令组，添加到本层的grouphandle中一起判断 awaitable.parent_group.add_custom_filter(custom_filter) else: handler_md = get_handler_or_create( - awaitable, EventType.AdapterMessageEvent, **kwargs + awaitable, + EventType.AdapterMessageEvent, + **kwargs, ) if not add_to_event_filters and not isinstance( - awaitable, RegisteringCommandable + awaitable, + RegisteringCommandable, ): # 底层子指令 handle_full_name = get_handler_full_name(awaitable) @@ -171,7 +181,9 @@ def register_custom_filter(custom_type_filter, *args, **kwargs): else: # 裸指令 handler_md = get_handler_or_create( - awaitable, EventType.AdapterMessageEvent, **kwargs + awaitable, + EventType.AdapterMessageEvent, + **kwargs, ) handler_md.event_filters.append(custom_filter) @@ -194,20 +206,23 @@ def register_command_group( logger.warning(f"{command_group_name} 指令组的子指令组 sub_command 未指定") else: new_group = CommandGroupFilter( - sub_command, alias, parent_group=command_group_name.parent_group + sub_command, + alias, + parent_group=command_group_name.parent_group, ) command_group_name.parent_group.add_sub_command_filter(new_group) + # 根指令组 + elif command_group_name is None: + logger.warning("根指令组的名称未指定") else: - # 根指令组 - if command_group_name is None: - logger.warning("根指令组的名称未指定") - else: - new_group = CommandGroupFilter(command_group_name, alias) + new_group = CommandGroupFilter(command_group_name, alias) def decorator(obj): if new_group: handler_md = get_handler_or_create( - obj, EventType.AdapterMessageEvent, **kwargs + obj, + EventType.AdapterMessageEvent, + **kwargs, ) handler_md.event_filters.append(new_group) @@ -220,9 +235,7 @@ def register_command_group( class RegisteringCommandable: """用于指令组级联注册""" - group: Callable[..., Callable[..., "RegisteringCommandable"]] = ( - register_command_group - ) + group: Callable[..., Callable[..., RegisteringCommandable]] = register_command_group command: Callable[..., Callable[..., None]] = register_command custom_filter: Callable[..., Callable[..., None]] = register_custom_filter @@ -235,7 +248,9 @@ def register_event_message_type(event_message_type: EventMessageType, **kwargs): def decorator(awaitable): handler_md = get_handler_or_create( - awaitable, EventType.AdapterMessageEvent, **kwargs + awaitable, + EventType.AdapterMessageEvent, + **kwargs, ) handler_md.event_filters.append(EventMessageTypeFilter(event_message_type)) return awaitable @@ -244,14 +259,15 @@ def register_event_message_type(event_message_type: EventMessageType, **kwargs): def register_platform_adapter_type( - platform_adapter_type: PlatformAdapterType, **kwargs + platform_adapter_type: PlatformAdapterType, + **kwargs, ): """注册一个 PlatformAdapterType""" def decorator(awaitable): handler_md = get_handler_or_create(awaitable, EventType.AdapterMessageEvent) handler_md.event_filters.append( - PlatformAdapterTypeFilter(platform_adapter_type) + PlatformAdapterTypeFilter(platform_adapter_type), ) return awaitable @@ -263,7 +279,9 @@ def register_regex(regex: str, **kwargs): def decorator(awaitable): handler_md = get_handler_or_create( - awaitable, EventType.AdapterMessageEvent, **kwargs + awaitable, + EventType.AdapterMessageEvent, + **kwargs, ) handler_md.event_filters.append(RegexFilter(regex)) return awaitable @@ -277,12 +295,13 @@ def register_permission_type(permission_type: PermissionType, raise_error: bool Args: permission_type: PermissionType raise_error: 如果没有权限，是否抛出错误到消息平台，并且停止事件传播。默认为 True + """ def decorator(awaitable): handler_md = get_handler_or_create(awaitable, EventType.AdapterMessageEvent) handler_md.event_filters.append( - PermissionTypeFilter(permission_type, raise_error) + PermissionTypeFilter(permission_type, raise_error), ) return awaitable @@ -300,9 +319,7 @@ def register_on_astrbot_loaded(**kwargs): def register_on_platform_loaded(**kwargs): - """ - 当平台加载完成时 - """ + """当平台加载完成时""" def decorator(awaitable): _ = get_handler_or_create(awaitable, EventType.OnPlatformLoadedEvent, **kwargs) @@ -324,6 +341,7 @@ def register_on_llm_request(**kwargs): ``` 请务必接收两个参数：event, request + """ def decorator(awaitable): @@ -346,6 +364,7 @@ def register_on_llm_response(**kwargs): ``` 请务必接收两个参数：event, request + """ def decorator(awaitable): @@ -365,7 +384,7 @@ def register_llm_tool(name: str | None = None, **kwargs): async def get_weather(event: AstrMessageEvent, location: str): \'\'\'获取天气信息。 - Args: + Args: location(string): 地点 \'\'\' # 处理逻辑 @@ -386,8 +405,8 @@ def register_llm_tool(name: str | None = None, **kwargs): event.stop_event() yield ``` - """ + """ name_ = name registering_agent = None if kwargs.get("registering_agent"): @@ -399,18 +418,37 @@ def register_llm_tool(name: str | None = None, **kwargs): docstring = docstring_parser.parse(func_doc) args = [] for arg in docstring.params: - if arg.type_name not in SUPPORTED_TYPES: + sub_type_name = None + type_name = arg.type_name + if not type_name: raise ValueError( - f"LLM 函数工具 {awaitable.__module__}_{llm_tool_name} 不支持的参数类型：{arg.type_name}" + f"LLM 函数工具 {awaitable.__module__}_{llm_tool_name} 的参数 {arg.arg_name} 缺少类型注释。", ) - args.append( - { - "type": arg.type_name, - "name": arg.arg_name, - "description": arg.description, - } - ) - # print(llm_tool_name, registering_agent) + # parse type_name to handle cases like "list[string]" + match = re.match(r"(\w+)\[(\w+)\]", type_name) + if match: + type_name = match.group(1) + sub_type_name = match.group(2) + type_name = PY_TO_JSON_TYPE.get(type_name, type_name) + if sub_type_name: + sub_type_name = PY_TO_JSON_TYPE.get(sub_type_name, sub_type_name) + if type_name not in SUPPORTED_TYPES or ( + sub_type_name and sub_type_name not in SUPPORTED_TYPES + ): + raise ValueError( + f"LLM 函数工具 {awaitable.__module__}_{llm_tool_name} 不支持的参数类型：{arg.type_name}", + ) + + arg_json_schema = { + "type": type_name, + "name": arg.arg_name, + "description": arg.description, + } + if sub_type_name: + if type_name == "array": + arg_json_schema["items"] = {"type": sub_type_name} + args.append(arg_json_schema) + if not registering_agent: doc_desc = docstring.description.strip() if docstring.description else "" md = get_handler_or_create(awaitable, EventType.OnCallingFuncToolEvent) @@ -454,6 +492,7 @@ def register_agent( instruction: Agent 的指令 tools: Agent 使用的工具列表 run_hooks: Agent 运行时的钩子函数 + """ tools_ = tools or [] @@ -478,7 +517,9 @@ def register_on_decorating_result(**kwargs): def decorator(awaitable): _ = get_handler_or_create( - awaitable, EventType.OnDecoratingResultEvent, **kwargs + awaitable, + EventType.OnDecoratingResultEvent, + **kwargs, ) return awaitable @@ -490,7 +531,9 @@ def register_after_message_sent(**kwargs): def decorator(awaitable): _ = get_handler_or_create( - awaitable, EventType.OnAfterMessageSentEvent, **kwargs + awaitable, + EventType.OnAfterMessageSentEvent, + **kwargs, ) return awaitable diff --git a/astrbot/core/star/session_llm_manager.py b/astrbot/core/star/session_llm_manager.py index 8fb88c6b..8c40f25c 100644 --- a/astrbot/core/star/session_llm_manager.py +++ b/astrbot/core/star/session_llm_manager.py @@ -1,6 +1,4 @@ -""" -会话服务管理器 - 负责管理每个会话的LLM、TTS等服务的启停状态 -""" +"""会话服务管理器 - 负责管理每个会话的LLM、TTS等服务的启停状态""" from astrbot.core import logger, sp from astrbot.core.platform.astr_message_event import AstrMessageEvent @@ -22,10 +20,14 @@ class SessionServiceManager: Returns: bool: True表示启用，False表示禁用 + """ # 获取会话服务配置 session_services = sp.get( - "session_service_config", {}, scope="umo", scope_id=session_id + "session_service_config", + {}, + scope="umo", + scope_id=session_id, ) # 如果配置了该会话的LLM状态，返回该状态 @@ -43,13 +45,17 @@ class SessionServiceManager: Args: session_id: 会话ID (unified_msg_origin) enabled: True表示启用，False表示禁用 + """ session_config = ( sp.get("session_service_config", {}, scope="umo", scope_id=session_id) or {} ) session_config["llm_enabled"] = enabled sp.put( - "session_service_config", session_config, scope="umo", scope_id=session_id + "session_service_config", + session_config, + scope="umo", + scope_id=session_id, ) @staticmethod @@ -61,6 +67,7 @@ class SessionServiceManager: Returns: bool: True表示应该处理，False表示跳过 + """ session_id = event.unified_msg_origin return SessionServiceManager.is_llm_enabled_for_session(session_id) @@ -78,10 +85,14 @@ class SessionServiceManager: Returns: bool: True表示启用，False表示禁用 + """ # 获取会话服务配置 session_services = sp.get( - "session_service_config", {}, scope="umo", scope_id=session_id + "session_service_config", + {}, + scope="umo", + scope_id=session_id, ) # 如果配置了该会话的TTS状态，返回该状态 @@ -99,17 +110,21 @@ class SessionServiceManager: Args: session_id: 会话ID (unified_msg_origin) enabled: True表示启用，False表示禁用 + """ session_config = ( sp.get("session_service_config", {}, scope="umo", scope_id=session_id) or {} ) session_config["tts_enabled"] = enabled sp.put( - "session_service_config", session_config, scope="umo", scope_id=session_id + "session_service_config", + session_config, + scope="umo", + scope_id=session_id, ) logger.info( - f"会话 {session_id} 的TTS状态已更新为: {'启用' if enabled else '禁用'}" + f"会话 {session_id} 的TTS状态已更新为: {'启用' if enabled else '禁用'}", ) @staticmethod @@ -121,6 +136,7 @@ class SessionServiceManager: Returns: bool: True表示应该处理，False表示跳过 + """ session_id = event.unified_msg_origin return SessionServiceManager.is_tts_enabled_for_session(session_id) @@ -138,10 +154,14 @@ class SessionServiceManager: Returns: bool: True表示启用，False表示禁用 + """ # 获取会话服务配置 session_services = sp.get( - "session_service_config", {}, scope="umo", scope_id=session_id + "session_service_config", + {}, + scope="umo", + scope_id=session_id, ) # 如果配置了该会话的整体状态，返回该状态 @@ -159,17 +179,21 @@ class SessionServiceManager: Args: session_id: 会话ID (unified_msg_origin) enabled: True表示启用，False表示禁用 + """ session_config = ( sp.get("session_service_config", {}, scope="umo", scope_id=session_id) or {} ) session_config["session_enabled"] = enabled sp.put( - "session_service_config", session_config, scope="umo", scope_id=session_id + "session_service_config", + session_config, + scope="umo", + scope_id=session_id, ) logger.info( - f"会话 {session_id} 的整体状态已更新为: {'启用' if enabled else '禁用'}" + f"会话 {session_id} 的整体状态已更新为: {'启用' if enabled else '禁用'}", ) @staticmethod @@ -181,6 +205,7 @@ class SessionServiceManager: Returns: bool: True表示应该处理，False表示跳过 + """ session_id = event.unified_msg_origin return SessionServiceManager.is_session_enabled(session_id) @@ -198,9 +223,13 @@ class SessionServiceManager: Returns: str: 自定义名称，如果没有设置则返回None + """ session_services = sp.get( - "session_service_config", {}, scope="umo", scope_id=session_id + "session_service_config", + {}, + scope="umo", + scope_id=session_id, ) return session_services.get("custom_name") @@ -211,6 +240,7 @@ class SessionServiceManager: Args: session_id: 会话ID (unified_msg_origin) custom_name: 自定义名称，可以为空字符串来清除名称 + """ session_config = ( sp.get("session_service_config", {}, scope="umo", scope_id=session_id) or {} @@ -221,11 +251,14 @@ class SessionServiceManager: # 如果传入空名称，则删除自定义名称 session_config.pop("custom_name", None) sp.put( - "session_service_config", session_config, scope="umo", scope_id=session_id + "session_service_config", + session_config, + scope="umo", + scope_id=session_id, ) logger.info( - f"会话 {session_id} 的自定义名称已更新为: {custom_name.strip() if custom_name and custom_name.strip() else '已清除'}" + f"会话 {session_id} 的自定义名称已更新为: {custom_name.strip() if custom_name and custom_name.strip() else '已清除'}", ) @staticmethod @@ -237,6 +270,7 @@ class SessionServiceManager: Returns: str: 显示名称 + """ custom_name = SessionServiceManager.get_session_custom_name(session_id) if custom_name: diff --git a/astrbot/core/star/session_plugin_manager.py b/astrbot/core/star/session_plugin_manager.py index 94a0c8a4..c74546fe 100644 --- a/astrbot/core/star/session_plugin_manager.py +++ b/astrbot/core/star/session_plugin_manager.py @@ -1,9 +1,6 @@ -""" -会话插件管理器 - 负责管理每个会话的插件启停状态 -""" +"""会话插件管理器 - 负责管理每个会话的插件启停状态""" -from astrbot.core import sp, logger -from typing import Dict, List +from astrbot.core import logger, sp from astrbot.core.platform.astr_message_event import AstrMessageEvent @@ -20,10 +17,14 @@ class SessionPluginManager: Returns: bool: True表示启用，False表示禁用 + """ # 获取会话插件配置 session_plugin_config = sp.get( - "session_plugin_config", {}, scope="umo", scope_id=session_id + "session_plugin_config", + {}, + scope="umo", + scope_id=session_id, ) session_config = session_plugin_config.get(session_id, {}) @@ -43,7 +44,9 @@ class SessionPluginManager: @staticmethod def set_plugin_status_for_session( - session_id: str, plugin_name: str, enabled: bool + session_id: str, + plugin_name: str, + enabled: bool, ) -> None: """设置插件在指定会话中的启停状态 @@ -51,10 +54,14 @@ class SessionPluginManager: session_id: 会话ID (unified_msg_origin) plugin_name: 插件名称 enabled: True表示启用，False表示禁用 + """ # 获取当前配置 session_plugin_config = sp.get( - "session_plugin_config", {}, scope="umo", scope_id=session_id + "session_plugin_config", + {}, + scope="umo", + scope_id=session_id, ) if session_id not in session_plugin_config: session_plugin_config[session_id] = { @@ -91,11 +98,11 @@ class SessionPluginManager: ) logger.info( - f"会话 {session_id} 的插件 {plugin_name} 状态已更新为: {'启用' if enabled else '禁用'}" + f"会话 {session_id} 的插件 {plugin_name} 状态已更新为: {'启用' if enabled else '禁用'}", ) @staticmethod - def get_session_plugin_config(session_id: str) -> Dict[str, List[str]]: + def get_session_plugin_config(session_id: str) -> dict[str, list[str]]: """获取指定会话的插件配置 Args: @@ -103,16 +110,21 @@ class SessionPluginManager: Returns: Dict[str, List[str]]: 包含enabled_plugins和disabled_plugins的字典 + """ session_plugin_config = sp.get( - "session_plugin_config", {}, scope="umo", scope_id=session_id + "session_plugin_config", + {}, + scope="umo", + scope_id=session_id, ) return session_plugin_config.get( - session_id, {"enabled_plugins": [], "disabled_plugins": []} + session_id, + {"enabled_plugins": [], "disabled_plugins": []}, ) @staticmethod - def filter_handlers_by_session(event: AstrMessageEvent, handlers: List) -> List: + def filter_handlers_by_session(event: AstrMessageEvent, handlers: list) -> list: """根据会话配置过滤处理器列表 Args: @@ -121,6 +133,7 @@ class SessionPluginManager: Returns: List: 过滤后的处理器列表 + """ from astrbot.core.star.star import star_map @@ -145,12 +158,13 @@ class SessionPluginManager: # 检查插件是否在当前会话中启用 if SessionPluginManager.is_plugin_enabled_for_session( - session_id, plugin.name + session_id, + plugin.name, ): filtered_handlers.append(handler) else: logger.debug( - f"插件 {plugin.name} 在会话 {session_id} 中被禁用，跳过处理器 {handler.handler_name}" + f"插件 {plugin.name} 在会话 {session_id} 中被禁用，跳过处理器 {handler.handler_name}", ) return filtered_handlers diff --git a/astrbot/core/star/star.py b/astrbot/core/star/star.py index 0563e8cc..c5b7b124 100644 --- a/astrbot/core/star/star.py +++ b/astrbot/core/star/star.py @@ -16,8 +16,7 @@ if TYPE_CHECKING: @dataclass class StarMetadata: - """ - 插件的元数据。 + """插件的元数据。 当 activated 为 False 时，star_cls 可能为 None，请不要在插件未激活时调用 star_cls 的方法。 """ @@ -56,6 +55,12 @@ class StarMetadata: star_handler_full_names: list[str] = field(default_factory=list) """注册的 Handler 的全名列表""" + display_name: str | None = None + """用于展示的插件名称""" + + logo_path: str | None = None + """插件 Logo 的路径""" + def __str__(self) -> str: return f"Plugin {self.name} ({self.version}) by {self.author}: {self.desc}" diff --git a/astrbot/core/star/star_handler.py b/astrbot/core/star/star_handler.py index 80b5adb6..141f9180 100644 --- a/astrbot/core/star/star_handler.py +++ b/astrbot/core/star/star_handler.py @@ -1,7 +1,10 @@ from __future__ import annotations + import enum +from collections.abc import Awaitable, Callable from dataclasses import dataclass, field -from typing import Callable, Awaitable, Any, List, Dict, TypeVar, Generic +from typing import Any, Generic, TypeVar + from .filter import HandlerFilter from .star import star_map @@ -10,8 +13,8 @@ T = TypeVar("T", bound="StarHandlerMetadata") class StarHandlerRegistry(Generic[T]): def __init__(self): - self.star_handlers_map: Dict[str, StarHandlerMetadata] = {} - self._handlers: List[StarHandlerMetadata] = [] + self.star_handlers_map: dict[str, StarHandlerMetadata] = {} + self._handlers: list[StarHandlerMetadata] = [] def append(self, handler: StarHandlerMetadata): """添加一个 Handler，并保持按优先级有序""" @@ -31,7 +34,7 @@ class StarHandlerRegistry(Generic[T]): event_type: EventType, only_activated=True, plugins_name: list[str] | None = None, - ) -> List[StarHandlerMetadata]: + ) -> list[StarHandlerMetadata]: handlers = [] for handler in self._handlers: # 过滤事件类型 @@ -64,8 +67,9 @@ class StarHandlerRegistry(Generic[T]): return self.star_handlers_map.get(full_name, None) def get_handlers_by_module_name( - self, module_name: str - ) -> List[StarHandlerMetadata]: + self, + module_name: str, + ) -> list[StarHandlerMetadata]: return [ handler for handler in self._handlers @@ -126,7 +130,7 @@ class StarHandlerMetadata: handler: Callable[..., Awaitable[Any]] """Handler 的函数对象，应当是一个异步函数""" - event_filters: List[HandlerFilter] + event_filters: list[HandlerFilter] """一个适配器消息事件过滤器，用于描述这个 Handler 能够处理、应该处理的适配器消息事件""" desc: str = "" @@ -138,5 +142,6 @@ class StarHandlerMetadata: def __lt__(self, other: StarHandlerMetadata): """定义小于运算符以支持优先队列""" return self.extras_configs.get("priority", 0) < other.extras_configs.get( - "priority", 0 + "priority", + 0, ) diff --git a/astrbot/core/star/star_manager.py b/astrbot/core/star/star_manager.py index 41700673..abdedc24 100644 --- a/astrbot/core/star/star_manager.py +++ b/astrbot/core/star/star_manager.py @@ -1,6 +1,4 @@ -""" -插件的重载、启停、安装、卸载等操作。 -""" +"""插件的重载、启停、安装、卸载等操作。""" import asyncio import functools @@ -15,6 +13,7 @@ from types import ModuleType import yaml from astrbot.core import logger, pip_installer, sp +from astrbot.core.agent.handoff import FunctionTool, HandoffTool from astrbot.core.config.astrbot_config import AstrBotConfig from astrbot.core.provider.register import llm_tools from astrbot.core.utils.astrbot_path import ( @@ -22,7 +21,6 @@ from astrbot.core.utils.astrbot_path import ( get_astrbot_plugin_path, ) from astrbot.core.utils.io import remove_dir -from astrbot.core.agent.handoff import HandoffTool, FunctionTool from . import StarMetadata from .context import Context @@ -52,11 +50,13 @@ class PluginManager: """存储插件配置的路径。data/config""" self.reserved_plugin_path = os.path.abspath( os.path.join( - os.path.dirname(os.path.abspath(__file__)), "../../../packages" - ) + os.path.dirname(os.path.abspath(__file__)), + "../../../packages", + ), ) """保留插件的路径。在 packages 目录下""" self.conf_schema_fname = "_conf_schema.json" + self.logo_fname = "logo.png" """插件配置 Schema 文件名""" self._pm_lock = asyncio.Lock() """StarManager操作互斥锁""" @@ -79,7 +79,7 @@ class PluginManager: except asyncio.CancelledError: pass except Exception as e: - logger.error(f"插件热重载监视任务异常: {str(e)}") + logger.error(f"插件热重载监视任务异常: {e!s}") logger.error(traceback.format_exc()) async def _handle_file_changes(self, changes): @@ -94,11 +94,13 @@ class PluginManager: continue if star.reserved: plugin_dir_path = os.path.join( - self.reserved_plugin_path, star.root_dir_name + self.reserved_plugin_path, + star.root_dir_name, ) else: plugin_dir_path = os.path.join( - self.plugin_store_path, star.root_dir_name + self.plugin_store_path, + star.root_dir_name, ) plugins_to_check.append((plugin_dir_path, star.name)) reloaded_plugins = set() @@ -142,14 +144,14 @@ class PluginManager: logger.info(f"插件 {d} 未找到 main.py 或者 {d}.py，跳过。") continue if os.path.exists(os.path.join(path, d, "main.py")) or os.path.exists( - os.path.join(path, d, d + ".py") + os.path.join(path, d, d + ".py"), ): modules.append( { "pname": d, "module": module_str, "module_path": os.path.join(path, d, module_str), - } + }, ) return modules @@ -185,7 +187,7 @@ class PluginManager: try: await pip_installer.install(requirements_path=pth) except Exception as e: - logger.error(f"更新插件 {p} 的依赖失败。Code: {str(e)}") + logger.error(f"更新插件 {p} 的依赖失败。Code: {e!s}") @staticmethod def _load_plugin_metadata(plugin_path: str, plugin_obj=None) -> StarMetadata | None: @@ -200,7 +202,8 @@ class PluginManager: if os.path.exists(os.path.join(plugin_path, "metadata.yaml")): with open( - os.path.join(plugin_path, "metadata.yaml"), "r", encoding="utf-8" + os.path.join(plugin_path, "metadata.yaml"), + encoding="utf-8", ) as f: metadata = yaml.safe_load(f) elif plugin_obj and hasattr(plugin_obj, "info"): @@ -218,7 +221,7 @@ class PluginManager: or "author" not in metadata ): raise Exception( - "插件元数据信息不完整。name, desc, version, author 是必须的字段。" + "插件元数据信息不完整。name, desc, version, author 是必须的字段。", ) metadata = StarMetadata( name=metadata["name"], @@ -226,13 +229,15 @@ class PluginManager: desc=metadata["desc"], version=metadata["version"], repo=metadata["repo"] if "repo" in metadata else None, + display_name=metadata.get("display_name", None), ) return metadata @staticmethod def _get_plugin_related_modules( - plugin_root_dir: str, is_reserved: bool = False + plugin_root_dir: str, + is_reserved: bool = False, ) -> list[str]: """获取与指定插件相关的所有已加载模块名 @@ -244,6 +249,7 @@ class PluginManager: Returns: list[str]: 与该插件相关的模块名列表 + """ prefix = "packages." if is_reserved else "data.plugins." return [ @@ -266,6 +272,7 @@ class PluginManager: module_patterns: 要移除的模块名模式列表（例如 ["data.plugins", "packages"]） root_dir_name: 插件根目录名，用于移除与该插件相关的所有模块 is_reserved: 插件是否为保留插件（影响模块路径前缀） + """ if module_patterns: for pattern in module_patterns: @@ -276,7 +283,8 @@ class PluginManager: if root_dir_name: for module_name in self._get_plugin_related_modules( - root_dir_name, is_reserved + root_dir_name, + is_reserved, ): try: del sys.modules[module_name] @@ -295,6 +303,7 @@ class PluginManager: tuple: 返回 load() 方法的结果，包含 (success, error_message) - success (bool): 重载是否成功 - error_message (str|None): 错误信息，成功时为 None + """ async with self._pm_lock: specified_module_path = None @@ -313,7 +322,7 @@ class PluginManager: except Exception as e: logger.warning(traceback.format_exc()) logger.warning( - f"插件 {smd.name} 未被正常终止: {str(e)}, 可能会导致该插件运行不正常。" + f"插件 {smd.name} 未被正常终止: {e!s}, 可能会导致该插件运行不正常。", ) if smd.name and smd.module_path: await self._unbind_plugin(smd.name, smd.module_path) @@ -330,7 +339,7 @@ class PluginManager: except Exception as e: logger.warning(traceback.format_exc()) logger.warning( - f"插件 {smd.name} 未被正常终止: {str(e)}, 可能会导致该插件运行不正常。" + f"插件 {smd.name} 未被正常终止: {e!s}, 可能会导致该插件运行不正常。", ) if smd.name: await self._unbind_plugin(smd.name, specified_module_path) @@ -351,6 +360,7 @@ class PluginManager: tuple: (success, error_message) - success (bool): 是否全部加载成功 - error_message (str|None): 错误信息，成功时为 None + """ inactivated_plugins = await sp.global_get("inactivated_plugins", []) inactivated_llm_tools = await sp.global_get("inactivated_llm_tools", []) @@ -369,7 +379,8 @@ class PluginManager: # module_path = plugin_module['module_path'] root_dir_name = plugin_module["pname"] # 插件的目录名 reserved = plugin_module.get( - "reserved", False + "reserved", + False, ) # 是否是保留插件。目前在 packages/ 目录下的都是保留插件。保留插件不可以卸载。 path = "data.plugins." if not reserved else "packages." @@ -392,7 +403,7 @@ class PluginManager: module = __import__(path, fromlist=[module_str]) except Exception as e: logger.error(traceback.format_exc()) - logger.error(f"插件 {root_dir_name} 导入失败。原因：{str(e)}") + logger.error(f"插件 {root_dir_name} 导入失败。原因：{e!s}") continue # 检查 _conf_schema.json @@ -403,17 +414,20 @@ class PluginManager: else os.path.join(self.reserved_plugin_path, root_dir_name) ) plugin_schema_path = os.path.join( - plugin_dir_path, self.conf_schema_fname + plugin_dir_path, + self.conf_schema_fname, ) if os.path.exists(plugin_schema_path): # 加载插件配置 - with open(plugin_schema_path, "r", encoding="utf-8") as f: + with open(plugin_schema_path, encoding="utf-8") as f: plugin_config = AstrBotConfig( config_path=os.path.join( - self.plugin_config_path, f"{root_dir_name}_config.json" + self.plugin_config_path, + f"{root_dir_name}_config.json", ), schema=json.loads(f.read()), ) + logo_path = os.path.join(plugin_dir_path, self.logo_fname) if path in star_map: # 通过 __init__subclass__ 注册插件 @@ -422,7 +436,7 @@ class PluginManager: try: # yaml 文件的元数据优先 metadata_yaml = self._load_plugin_metadata( - plugin_path=plugin_dir_path + plugin_path=plugin_dir_path, ) if metadata_yaml: metadata.name = metadata_yaml.name @@ -430,9 +444,10 @@ class PluginManager: metadata.desc = metadata_yaml.desc metadata.version = metadata_yaml.version metadata.repo = metadata_yaml.repo + metadata.display_name = metadata_yaml.display_name except Exception as e: logger.warning( - f"插件 {root_dir_name} 元数据载入失败: {str(e)}。使用默认元数据。" + f"插件 {root_dir_name} 元数据载入失败: {e!s}。使用默认元数据。", ) logger.info(metadata) metadata.config = plugin_config @@ -441,15 +456,16 @@ class PluginManager: if plugin_config and metadata.star_cls_type: try: metadata.star_cls = metadata.star_cls_type( - context=self.context, config=plugin_config + context=self.context, + config=plugin_config, ) except TypeError as _: metadata.star_cls = metadata.star_cls_type( - context=self.context + context=self.context, ) elif metadata.star_cls_type: metadata.star_cls = metadata.star_cls_type( - context=self.context + context=self.context, ) else: logger.info(f"插件 {metadata.name} 已被禁用。") @@ -465,7 +481,7 @@ class PluginManager: # 绑定 handler related_handlers = ( star_handlers_registry.get_handlers_by_module_name( - metadata.module_path + metadata.module_path, ) ) for handler in related_handlers: @@ -501,7 +517,7 @@ class PluginManager: else: # v3.4.0 以前的方式注册插件 logger.debug( - f"插件 {path} 未通过装饰器注册。尝试通过旧版本方式载入。" + f"插件 {path} 未通过装饰器注册。尝试通过旧版本方式载入。", ) classes = self._get_classes(module) @@ -510,19 +526,21 @@ class PluginManager: if plugin_config: try: obj = getattr(module, classes[0])( - context=self.context, config=plugin_config + context=self.context, + config=plugin_config, ) # 实例化插件类 except TypeError as _: obj = getattr(module, classes[0])( - context=self.context + context=self.context, ) # 实例化插件类 else: obj = getattr(module, classes[0])( - context=self.context + context=self.context, ) # 实例化插件类 metadata = self._load_plugin_metadata( - plugin_path=plugin_dir_path, plugin_obj=obj + plugin_path=plugin_dir_path, + plugin_obj=obj, ) if not metadata: raise Exception(f"无法找到插件 {plugin_dir_path} 的元数据。") @@ -540,13 +558,15 @@ class PluginManager: if metadata.module_path in inactivated_plugins: metadata.activated = False - assert metadata.module_path is not None, ( - f"插件 {metadata.name} 的模块路径为空。" - ) + # Plugin logo path + if os.path.exists(logo_path): + metadata.logo_path = logo_path + + assert metadata.module_path, f"插件 {metadata.name} 模块路径为空" full_names = [] for handler in star_handlers_registry.get_handlers_by_module_name( - metadata.module_path + metadata.module_path, ): full_names.append(handler.handler_full_name) @@ -556,7 +576,8 @@ class PluginManager: and handler.handler_name in alter_cmd[metadata.name] ): cmd_type = alter_cmd[metadata.name][handler.handler_name].get( - "permission", "member" + "permission", + "member", ) found_permission_filter = False for filter_ in handler.event_filters: @@ -572,12 +593,12 @@ class PluginManager: PermissionTypeFilter( PermissionType.ADMIN if cmd_type == "admin" - else PermissionType.MEMBER - ) + else PermissionType.MEMBER, + ), ) logger.debug( - f"插入权限过滤器 {cmd_type} 到 {metadata.name} 的 {handler.handler_name} 方法。" + f"插入权限过滤器 {cmd_type} 到 {metadata.name} 的 {handler.handler_name} 方法。", ) metadata.star_handler_full_names = full_names @@ -592,7 +613,7 @@ class PluginManager: for line in errors.split("\n"): logger.error(f"| {line}") logger.error("----------------------------------") - fail_rec += f"加载 {root_dir_name} 插件时出现问题，原因 {str(e)}。\n" + fail_rec += f"加载 {root_dir_name} 插件时出现问题，原因 {e!s}。\n" # 清除 pip.main 导致的多余的 logging handlers for handler in logging.root.handlers[:]: @@ -600,9 +621,8 @@ class PluginManager: if not fail_rec: return True, None - else: - self.failed_plugin_info = fail_rec - return False, fail_rec + self.failed_plugin_info = fail_rec + return False, fail_rec async def install_plugin(self, repo_url: str, proxy=""): """从仓库 URL 安装插件 @@ -618,6 +638,7 @@ class PluginManager: - repo: 插件的仓库 URL - readme: README.md 文件的内容(如果存在) 如果找不到插件元数据则返回 None。 + """ async with self._pm_lock: plugin_path = await self.updator.install(repo_url, proxy) @@ -642,11 +663,11 @@ class PluginManager: if os.path.exists(readme_path): try: - with open(readme_path, "r", encoding="utf-8") as f: + with open(readme_path, encoding="utf-8") as f: readme_content = f.read() except Exception as e: logger.warning( - f"读取插件 {dir_name} 的 README.md 文件失败: {str(e)}" + f"读取插件 {dir_name} 的 README.md 文件失败: {e!s}", ) plugin_info = None @@ -659,14 +680,22 @@ class PluginManager: return plugin_info - async def uninstall_plugin(self, plugin_name: str): + async def uninstall_plugin( + self, + plugin_name: str, + delete_config: bool = False, + delete_data: bool = False, + ): """卸载指定的插件。 Args: plugin_name (str): 要卸载的插件名称 + delete_config (bool): 是否删除插件配置文件，默认为 False + delete_data (bool): 是否删除插件数据，默认为 False Raises: Exception: 当插件不存在、是保留插件时，或删除插件文件夹失败时抛出异常 + """ async with self._pm_lock: plugin = self.context.get_registered_star(plugin_name) @@ -683,7 +712,7 @@ class PluginManager: except Exception as e: logger.warning(traceback.format_exc()) logger.warning( - f"插件 {plugin_name} 未被正常终止 {str(e)}, 可能会导致资源泄露等问题。" + f"插件 {plugin_name} 未被正常终止 {e!s}, 可能会导致资源泄露等问题。", ) # 从 star_registry 和 star_map 中删除 @@ -692,19 +721,66 @@ class PluginManager: await self._unbind_plugin(plugin_name, plugin.module_path) + # 删除插件文件夹 try: remove_dir(os.path.join(ppath, root_dir_name)) except Exception as e: raise Exception( - f"移除插件成功，但是删除插件文件夹失败: {str(e)}。您可以手动删除该文件夹，位于 addons/plugins/ 下。" + f"移除插件成功，但是删除插件文件夹失败: {e!s}。您可以手动删除该文件夹，位于 addons/plugins/ 下。", ) + # 删除插件配置文件 + if delete_config and root_dir_name: + config_file = os.path.join( + self.plugin_config_path, + f"{root_dir_name}_config.json", + ) + if os.path.exists(config_file): + try: + os.remove(config_file) + logger.info(f"已删除插件 {plugin_name} 的配置文件") + except Exception as e: + logger.warning(f"删除插件配置文件失败: {e!s}") + + # 删除插件持久化数据 + # 注意：需要检查两个可能的目录名（plugin_data 和 plugins_data） + # data/temp 目录可能被多个插件共享，不自动删除以防误删 + if delete_data and root_dir_name: + data_base_dir = os.path.dirname(ppath) # data/ + + # 删除 data/plugin_data 下的插件持久化数据（单数形式，新版本） + plugin_data_dir = os.path.join( + data_base_dir, "plugin_data", root_dir_name + ) + if os.path.exists(plugin_data_dir): + try: + remove_dir(plugin_data_dir) + logger.info( + f"已删除插件 {plugin_name} 的持久化数据 (plugin_data)" + ) + except Exception as e: + logger.warning(f"删除插件持久化数据失败 (plugin_data): {e!s}") + + # 删除 data/plugins_data 下的插件持久化数据（复数形式，旧版本兼容） + plugins_data_dir = os.path.join( + data_base_dir, "plugins_data", root_dir_name + ) + if os.path.exists(plugins_data_dir): + try: + remove_dir(plugins_data_dir) + logger.info( + f"已删除插件 {plugin_name} 的持久化数据 (plugins_data)" + ) + except Exception as e: + logger.warning(f"删除插件持久化数据失败 (plugins_data): {e!s}") + async def _unbind_plugin(self, plugin_name: str, plugin_module_path: str): """解绑并移除一个插件。 Args: plugin_name: 要解绑的插件名称 plugin_module_path: 插件的完整模块路径 + """ plugin = None del star_map[plugin_module_path] @@ -714,10 +790,10 @@ class PluginManager: del star_registry[i] break for handler in star_handlers_registry.get_handlers_by_module_name( - plugin_module_path + plugin_module_path, ): logger.info( - f"移除了插件 {plugin_name} 的处理函数 {handler.handler_name} ({len(star_handlers_registry)})" + f"移除了插件 {plugin_name} 的处理函数 {handler.handler_name} ({len(star_handlers_registry)})", ) star_handlers_registry.remove(handler) @@ -728,11 +804,25 @@ class PluginManager: ]: del star_handlers_registry.star_handlers_map[k] + # llm_tools 中移除该插件的工具函数绑定 + to_remove = [] + for func_tool in llm_tools.func_list: + mp = func_tool.handler_module_path + if ( + mp + and mp.startswith(plugin_module_path) + and not mp.endswith(("packages", "data.plugins")) + ): + to_remove.append(func_tool) + for func_tool in to_remove: + llm_tools.func_list.remove(func_tool) + if plugin is None: return self._purge_modules( - root_dir_name=plugin.root_dir_name, is_reserved=plugin.reserved + root_dir_name=plugin.root_dir_name, + is_reserved=plugin.reserved, ) async def update_plugin(self, plugin_name: str, proxy=""): @@ -747,8 +837,7 @@ class PluginManager: await self.reload(plugin_name) async def turn_off_plugin(self, plugin_name: str): - """ - 禁用一个插件。 + """禁用一个插件。 调用插件的 terminate() 方法， 将插件的 module_path 加入到 data/shared_preferences.json 的 inactivated_plugins 列表中。 并且同时将插件启用的 llm_tool 禁用。 @@ -767,12 +856,18 @@ class PluginManager: inactivated_plugins.append(plugin.module_path) inactivated_llm_tools: list = list( - set(await sp.global_get("inactivated_llm_tools", [])) + set(await sp.global_get("inactivated_llm_tools", [])), ) # 后向兼容 # 禁用插件启用的 llm_tool for func_tool in llm_tools.func_list: - if func_tool.handler_module_path == plugin.module_path: + mp = func_tool.handler_module_path + if ( + plugin.module_path + and mp + and plugin.module_path.startswith(mp) + and not mp.endswith(("packages", "data.plugins")) + ): func_tool.active = False if func_tool.name not in inactivated_llm_tools: inactivated_llm_tools.append(func_tool.name) @@ -797,7 +892,8 @@ class PluginManager: if "__del__" in star_metadata.star_cls_type.__dict__: asyncio.get_event_loop().run_in_executor( - None, star_metadata.star_cls.__del__ + None, + star_metadata.star_cls.__del__, ) elif "terminate" in star_metadata.star_cls_type.__dict__: await star_metadata.star_cls.terminate() @@ -814,8 +910,12 @@ class PluginManager: # 启用插件启用的 llm_tool for func_tool in llm_tools.func_list: + mp = func_tool.handler_module_path if ( - func_tool.handler_module_path == plugin.module_path + plugin.module_path + and mp + and plugin.module_path.startswith(mp) + and not mp.endswith(("packages", "data.plugins")) and func_tool.name in inactivated_llm_tools ): inactivated_llm_tools.remove(func_tool.name) @@ -824,8 +924,6 @@ class PluginManager: await self.reload(plugin_name) - # plugin.activated = True - async def install_plugin_from_file(self, zip_file_path: str): dir_name = os.path.basename(zip_file_path).replace(".zip", "") dir_name = dir_name.removesuffix("-master").removesuffix("-main").lower() @@ -836,7 +934,7 @@ class PluginManager: try: os.remove(zip_file_path) except BaseException as e: - logger.warning(f"删除插件压缩包失败: {str(e)}") + logger.warning(f"删除插件压缩包失败: {e!s}") # await self.reload() await self.load(specified_dir_name=dir_name) @@ -857,10 +955,10 @@ class PluginManager: if os.path.exists(readme_path): try: - with open(readme_path, "r", encoding="utf-8") as f: + with open(readme_path, encoding="utf-8") as f: readme_content = f.read() except Exception as e: - logger.warning(f"读取插件 {dir_name} 的 README.md 文件失败: {str(e)}") + logger.warning(f"读取插件 {dir_name} 的 README.md 文件失败: {e!s}") plugin_info = None if plugin: diff --git a/astrbot/core/star/star_tools.py b/astrbot/core/star/star_tools.py index 6f9dfe2f..7a66449b 100644 --- a/astrbot/core/star/star_tools.py +++ b/astrbot/core/star/star_tools.py @@ -1,5 +1,4 @@ -""" -插件开发工具集 +"""插件开发工具集 封装了许多常用的操作，方便插件开发者使用 说明: @@ -21,47 +20,49 @@ import inspect import os import uuid +from collections.abc import Awaitable, Callable from pathlib import Path -from typing import Union, Awaitable, Callable, Any, List, Optional, ClassVar +from typing import Any, ClassVar + +from astrbot.api.platform import AstrBotMessage, MessageMember, MessageType from astrbot.core.message.components import BaseMessageComponent from astrbot.core.message.message_event_result import MessageChain -from astrbot.api.platform import MessageMember, AstrBotMessage, MessageType from astrbot.core.platform.astr_message_event import MessageSesion -from astrbot.core.star.context import Context -from astrbot.core.star.star import star_map -from astrbot.core.utils.astrbot_path import get_astrbot_data_path from astrbot.core.platform.sources.aiocqhttp.aiocqhttp_message_event import ( AiocqhttpMessageEvent, ) from astrbot.core.platform.sources.aiocqhttp.aiocqhttp_platform_adapter import ( AiocqhttpAdapter, ) +from astrbot.core.star.context import Context +from astrbot.core.star.star import star_map +from astrbot.core.utils.astrbot_path import get_astrbot_data_path class StarTools: - """ - 提供给插件使用的便捷工具函数集合 + """提供给插件使用的便捷工具函数集合 这些方法封装了一些常用操作，使插件开发更加简单便捷! """ - _context: ClassVar[Optional[Context]] = None + _context: ClassVar[Context | None] = None @classmethod def initialize(cls, context: Context) -> None: - """ - 初始化StarTools，设置context引用 + """初始化StarTools，设置context引用 Args: context: 暴露给插件的上下文 + """ cls._context = context @classmethod async def send_message( - cls, session: Union[str, MessageSesion], message_chain: MessageChain + cls, + session: str | MessageSesion, + message_chain: MessageChain, ) -> bool: - """ - 根据session(unified_msg_origin)主动发送消息 + """根据session(unified_msg_origin)主动发送消息 Args: session: 消息会话。通过event.session或者event.unified_msg_origin获取 @@ -75,6 +76,7 @@ class StarTools: Note: qq_official(QQ官方API平台)不支持此方法 + """ if cls._context is None: raise ValueError("StarTools not initialized") @@ -88,21 +90,22 @@ class StarTools: message_chain: MessageChain, platform: str = "aiocqhttp", ): - """ - 根据 id(例如qq号, 群号等) 直接, 主动地发送消息 + """根据 id(例如qq号, 群号等) 直接, 主动地发送消息 Args: type (str): 消息类型, 可选: PrivateMessage, GroupMessage id (str): 目标ID, 例如QQ号, 群号等 message_chain (MessageChain): 消息链 platform (str): 可选的平台名称，默认平台(aiocqhttp), 目前只支持 aiocqhttp + """ if cls._context is None: raise ValueError("StarTools not initialized") platforms = cls._context.platform_manager.get_insts() if platform == "aiocqhttp": adapter = next( - (p for p in platforms if isinstance(p, AiocqhttpAdapter)), None + (p for p in platforms if isinstance(p, AiocqhttpAdapter)), + None, ) if adapter is None: raise ValueError("未找到适配器: AiocqhttpAdapter") @@ -122,14 +125,13 @@ class StarTools: self_id: str, session_id: str, sender: MessageMember, - message: List[BaseMessageComponent], + message: list[BaseMessageComponent], message_str: str, message_id: str = "", raw_message: object = None, group_id: str = "", ) -> AstrBotMessage: - """ - 创建一个AstrBot消息对象 + """创建一个AstrBot消息对象 Args: type (str): 消息类型, 例如 "GroupMessage" "FriendMessage" "OtherMessage" @@ -145,6 +147,7 @@ class StarTools: Returns: AstrBotMessage: 创建的消息对象 + """ abm = AstrBotMessage() abm.type = MessageType(type) @@ -162,23 +165,27 @@ class StarTools: @classmethod async def create_event( - cls, abm: AstrBotMessage, platform: str = "aiocqhttp", is_wake: bool = True + cls, + abm: AstrBotMessage, + platform: str = "aiocqhttp", + is_wake: bool = True, ) -> None: - """ - 创建并提交事件到指定平台 + """创建并提交事件到指定平台 当有需要创建一个事件, 触发某些处理流程时, 使用该方法 Args: abm (AstrBotMessage): 要提交的消息对象, 请先使用 create_message 创建 platform (str): 可选的平台名称，默认平台(aiocqhttp), 目前只支持 aiocqhttp is_wake (bool): 是否标记为唤醒事件, 默认为 True, 只有唤醒事件才会被 llm 响应 + """ if cls._context is None: raise ValueError("StarTools not initialized") platforms = cls._context.platform_manager.get_insts() if platform == "aiocqhttp": adapter = next( - (p for p in platforms if isinstance(p, AiocqhttpAdapter)), None + (p for p in platforms if isinstance(p, AiocqhttpAdapter)), + None, ) if adapter is None: raise ValueError("未找到适配器: AiocqhttpAdapter") @@ -196,12 +203,12 @@ class StarTools: @classmethod def activate_llm_tool(cls, name: str) -> bool: - """ - 激活一个已经注册的函数调用工具 + """激活一个已经注册的函数调用工具 注册的工具默认是激活状态 Args: name (str): 工具名称 + """ if cls._context is None: raise ValueError("StarTools not initialized") @@ -209,11 +216,11 @@ class StarTools: @classmethod def deactivate_llm_tool(cls, name: str) -> bool: - """ - 停用一个已经注册的函数调用工具 + """停用一个已经注册的函数调用工具 Args: name (str): 工具名称 + """ if cls._context is None: raise ValueError("StarTools not initialized") @@ -227,14 +234,14 @@ class StarTools: desc: str, func_obj: Callable[..., Awaitable[Any]], ) -> None: - """ - 为函数调用（function-calling/tools-use）添加工具 + """为函数调用（function-calling/tools-use）添加工具 Args: name (str): 工具名称 func_args (list): 函数参数列表 desc (str): 工具描述 func_obj (Awaitable): 函数对象，必须是异步函数 + """ if cls._context is None: raise ValueError("StarTools not initialized") @@ -242,21 +249,20 @@ class StarTools: @classmethod def unregister_llm_tool(cls, name: str) -> None: - """ - 删除一个函数调用工具 + """删除一个函数调用工具 如果再要启用，需要重新注册 Args: name (str): 工具名称 + """ if cls._context is None: raise ValueError("StarTools not initialized") cls._context.unregister_llm_tool(name) @classmethod - def get_data_dir(cls, plugin_name: Optional[str] = None) -> Path: - """ - 返回插件数据目录的绝对路径。 + def get_data_dir(cls, plugin_name: str | None = None) -> Path: + """返回插件数据目录的绝对路径。 此方法会在 data/plugin_data 目录下为插件创建一个专属的数据目录。如果未提供插件名称， 会自动从调用栈中获取插件信息。 @@ -272,6 +278,7 @@ class StarTools: - 无法获取调用者模块信息 - 无法获取模块的元数据信息 - 创建目录失败（权限不足或其他IO错误） + """ if not plugin_name: frame = inspect.currentframe() @@ -294,7 +301,7 @@ class StarTools: raise ValueError("无法获取插件名称") data_dir = Path( - os.path.join(get_astrbot_data_path(), "plugin_data", plugin_name) + os.path.join(get_astrbot_data_path(), "plugin_data", plugin_name), ) try: diff --git a/astrbot/core/star/updator.py b/astrbot/core/star/updator.py index a2245537..8793ad50 100644 --- a/astrbot/core/star/updator.py +++ b/astrbot/core/star/updator.py @@ -1,12 +1,13 @@ import os -import zipfile import shutil +import zipfile -from ..updator import RepoZipUpdator -from astrbot.core.utils.io import remove_dir, on_error -from ..star.star import StarMetadata from astrbot.core import logger from astrbot.core.utils.astrbot_path import get_astrbot_plugin_path +from astrbot.core.utils.io import on_error, remove_dir + +from ..star.star import StarMetadata +from ..updator import RepoZipUpdator class PluginUpdator(RepoZipUpdator): @@ -44,7 +45,7 @@ class PluginUpdator(RepoZipUpdator): remove_dir(plugin_path) except BaseException as e: logger.error( - f"删除旧版本插件 {plugin_path} 文件夹失败: {str(e)}，使用覆盖安装。" + f"删除旧版本插件 {plugin_path} 文件夹失败: {e!s}，使用覆盖安装。", ) self.unzip_file(plugin_path + ".zip", plugin_path) @@ -64,18 +65,17 @@ class PluginUpdator(RepoZipUpdator): if os.path.isdir(os.path.join(target_dir, update_dir, f)): if os.path.exists(os.path.join(target_dir, f)): shutil.rmtree(os.path.join(target_dir, f), onerror=on_error) - else: - if os.path.exists(os.path.join(target_dir, f)): - os.remove(os.path.join(target_dir, f)) + elif os.path.exists(os.path.join(target_dir, f)): + os.remove(os.path.join(target_dir, f)) shutil.move(os.path.join(target_dir, update_dir, f), target_dir) try: logger.info( - f"删除临时文件: {zip_path} 和 {os.path.join(target_dir, update_dir)}" + f"删除临时文件: {zip_path} 和 {os.path.join(target_dir, update_dir)}", ) shutil.rmtree(os.path.join(target_dir, update_dir), onerror=on_error) os.remove(zip_path) except BaseException: logger.warning( - f"删除更新文件失败，可以手动删除 {zip_path} 和 {os.path.join(target_dir, update_dir)}" + f"删除更新文件失败，可以手动删除 {zip_path} 和 {os.path.join(target_dir, update_dir)}", ) diff --git a/astrbot/core/umop_config_router.py b/astrbot/core/umop_config_router.py index dd2063e5..07858da5 100644 --- a/astrbot/core/umop_config_router.py +++ b/astrbot/core/umop_config_router.py @@ -15,7 +15,10 @@ class UmopConfigRouter: """加载路由表""" # 从 SharedPreferences 中加载 umop_to_conf_id 映射 sp_data = self.sp.get( - "umop_config_routing", {}, scope="global", scope_id="global" + "umop_config_routing", + {}, + scope="global", + scope_id="global", ) self.umop_to_conf_id = sp_data @@ -37,6 +40,7 @@ class UmopConfigRouter: Returns: str | None: 配置文件 ID，如果没有找到则返回 None + """ for pattern, conf_id in self.umop_to_conf_id.items(): if self._is_umo_match(pattern, umo): @@ -52,11 +56,12 @@ class UmopConfigRouter: Raises: ValueError: 如果 new_routing 中的 key 格式不正确 + """ - for part in new_routing.keys(): + for part in new_routing: if not isinstance(part, str) or len(part.split(":")) != 3: raise ValueError( - "umop keys must be strings in the format [platform_id]:[message_type]:[session_id], with optional wildcards * or empty for all" + "umop keys must be strings in the format [platform_id]:[message_type]:[session_id], with optional wildcards * or empty for all", ) self.umop_to_conf_id = new_routing @@ -71,10 +76,11 @@ class UmopConfigRouter: Raises: ValueError: 如果 umo 格式不正确 + """ if not isinstance(umo, str) or len(umo.split(":")) != 3: raise ValueError( - "umop must be a string in the format [platform_id]:[message_type]:[session_id], with optional wildcards * or empty for all" + "umop must be a string in the format [platform_id]:[message_type]:[session_id], with optional wildcards * or empty for all", ) self.umop_to_conf_id[umo] = conf_id diff --git a/astrbot/core/updator.py b/astrbot/core/updator.py index 439c05a5..d13bab68 100644 --- a/astrbot/core/updator.py +++ b/astrbot/core/updator.py @@ -1,12 +1,15 @@ import os -import psutil import sys import time -from .zip_updator import ReleaseInfo, RepoZipUpdator + +import psutil + from astrbot.core import logger from astrbot.core.config.default import VERSION -from astrbot.core.utils.io import download_file from astrbot.core.utils.astrbot_path import get_astrbot_path +from astrbot.core.utils.io import download_file + +from .zip_updator import ReleaseInfo, RepoZipUpdator class AstrBotUpdator(RepoZipUpdator): @@ -67,11 +70,16 @@ class AstrBotUpdator(RepoZipUpdator): raise e async def check_update( - self, url: str, current_version: str, consider_prerelease: bool = True + self, + url: str, + current_version: str, + consider_prerelease: bool = True, ) -> ReleaseInfo: """检查更新""" return await super().check_update( - self.ASTRBOT_RELEASE_API, VERSION, consider_prerelease + self.ASTRBOT_RELEASE_API, + VERSION, + consider_prerelease, ) async def get_releases(self) -> list: @@ -99,7 +107,7 @@ class AstrBotUpdator(RepoZipUpdator): else: if len(str(version)) != 40: raise Exception("commit hash 长度不正确，应为 40") - file_url = f"https://github.com/Soulter/AstrBot/archive/{version}.zip" + file_url = f"https://github.com/AstrBotDevs/AstrBot/archive/{version}.zip" logger.info(f"准备更新至指定版本的 AstrBot Core: {version}") if proxy: diff --git a/astrbot/core/utils/astrbot_path.py b/astrbot/core/utils/astrbot_path.py index 64ed9229..e13379b9 100644 --- a/astrbot/core/utils/astrbot_path.py +++ b/astrbot/core/utils/astrbot_path.py @@ -1,5 +1,4 @@ -""" -Astrbot统一路径获取 +"""Astrbot统一路径获取 项目路径：固定为源码所在路径 根目录路径：默认为当前工作目录，可通过环境变量 ASTRBOT_ROOT 指定 @@ -14,7 +13,7 @@ import os def get_astrbot_path() -> str: """获取Astrbot项目路径""" return os.path.realpath( - os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../../") + os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../../"), ) @@ -22,8 +21,7 @@ def get_astrbot_root() -> str: """获取Astrbot根目录路径""" if path := os.environ.get("ASTRBOT_ROOT"): return os.path.realpath(path) - else: - return os.path.realpath(os.getcwd()) + return os.path.realpath(os.getcwd()) def get_astrbot_data_path() -> str: diff --git a/astrbot/core/utils/dify_api_client.py b/astrbot/core/utils/dify_api_client.py index 15a6b71f..ea8ff9df 100644 --- a/astrbot/core/utils/dify_api_client.py +++ b/astrbot/core/utils/dify_api_client.py @@ -1,8 +1,11 @@ import codecs import json +from collections.abc import AsyncGenerator +from typing import Any + +from aiohttp import ClientResponse, ClientSession + from astrbot.core import logger -from aiohttp import ClientSession, ClientResponse -from typing import Dict, List, Any, AsyncGenerator async def _stream_sse(resp: ClientResponse) -> AsyncGenerator[dict, None]: @@ -25,7 +28,6 @@ async def _stream_sse(resp: ClientResponse) -> AsyncGenerator[dict, None]: yield json.loads(buffer[5:]) except json.JSONDecodeError: logger.warning(f"Drop invalid dify json data: {buffer[5:]}") - pass class DifyAPIClient: @@ -39,50 +41,60 @@ class DifyAPIClient: async def chat_messages( self, - inputs: Dict, + inputs: dict, query: str, user: str, response_mode: str = "streaming", conversation_id: str = "", - files: List[Dict[str, Any]] = [], + files: list[dict[str, Any]] | None = None, timeout: float = 60, - ) -> AsyncGenerator[Dict[str, Any], None]: + ) -> AsyncGenerator[dict[str, Any], None]: + if files is None: + files = [] url = f"{self.api_base}/chat-messages" payload = locals() payload.pop("self") payload.pop("timeout") logger.info(f"chat_messages payload: {payload}") async with self.session.post( - url, json=payload, headers=self.headers, timeout=timeout + url, + json=payload, + headers=self.headers, + timeout=timeout, ) as resp: if resp.status != 200: text = await resp.text() raise Exception( - f"Dify /chat-messages 接口请求失败：{resp.status}. {text}" + f"Dify /chat-messages 接口请求失败：{resp.status}. {text}", ) async for event in _stream_sse(resp): yield event async def workflow_run( self, - inputs: Dict, + inputs: dict, user: str, response_mode: str = "streaming", - files: List[Dict[str, Any]] = [], + files: list[dict[str, Any]] | None = None, timeout: float = 60, ): + if files is None: + files = [] url = f"{self.api_base}/workflows/run" payload = locals() payload.pop("self") payload.pop("timeout") logger.info(f"workflow_run payload: {payload}") async with self.session.post( - url, json=payload, headers=self.headers, timeout=timeout + url, + json=payload, + headers=self.headers, + timeout=timeout, ) as resp: if resp.status != 200: text = await resp.text() raise Exception( - f"Dify /workflows/run 接口请求失败：{resp.status}. {text}" + f"Dify /workflows/run 接口请求失败：{resp.status}. {text}", ) async for event in _stream_sse(resp): yield event @@ -91,7 +103,7 @@ class DifyAPIClient: self, file_path: str, user: str, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: url = f"{self.api_base}/files/upload" with open(file_path, "rb") as f: payload = { @@ -99,7 +111,9 @@ class DifyAPIClient: "file": f, } async with self.session.post( - url, data=payload, headers=self.headers + url, + data=payload, + headers=self.headers, ) as resp: return await resp.json() # {"id": "xxx", ...} @@ -126,7 +140,11 @@ class DifyAPIClient: return await resp.json() async def rename( - self, conversation_id: str, name: str, user: str, auto_generate: bool = False + self, + conversation_id: str, + name: str, + user: str, + auto_generate: bool = False, ): # /conversations/:conversation_id/name url = f"{self.api_base}/conversations/{conversation_id}/name" diff --git a/astrbot/core/utils/io.py b/astrbot/core/utils/io.py index 723953f8..073c0493 100644 --- a/astrbot/core/utils/io.py +++ b/astrbot/core/utils/io.py @@ -1,29 +1,26 @@ +import base64 +import logging import os -import ssl import shutil import socket +import ssl import time -import aiohttp -import base64 -import zipfile import uuid -import psutil -import logging +import zipfile +from pathlib import Path +import aiohttp import certifi - -from typing import Union - +import psutil from PIL import Image + from .astrbot_path import get_astrbot_data_path logger = logging.getLogger("astrbot") def on_error(func, path, exc_info): - """ - a callback of the rmtree function. - """ + """A callback of the rmtree function.""" import stat if not os.access(path, os.W_OK): @@ -52,7 +49,7 @@ def port_checker(port: int, host: str = "localhost"): return False -def save_temp_img(img: Union[Image.Image, str]) -> str: +def save_temp_img(img: Image.Image | str) -> str: temp_dir = os.path.join(get_astrbot_data_path(), "temp") # 获得文件创建时间，清除超过 12 小时的 try: @@ -78,61 +75,75 @@ def save_temp_img(img: Union[Image.Image, str]) -> str: async def download_image_by_url( - url: str, post: bool = False, post_data: dict = None, path=None + url: str, + post: bool = False, + post_data: dict | None = None, + path: str | None = None, ) -> str: - """ - 下载图片, 返回 path - """ + """下载图片, 返回 path""" try: ssl_context = ssl.create_default_context( - cafile=certifi.where() + cafile=certifi.where(), ) # 使用 certifi 提供的 CA 证书 connector = aiohttp.TCPConnector(ssl=ssl_context) # 使用 certifi 的根证书 async with aiohttp.ClientSession( - trust_env=True, connector=connector + trust_env=True, + connector=connector, ) as session: if post: async with session.post(url, json=post_data) as resp: if not path: return save_temp_img(await resp.read()) - else: - with open(path, "wb") as f: - f.write(await resp.read()) - return path + with open(path, "wb") as f: + f.write(await resp.read()) + return path else: async with session.get(url) as resp: if not path: return save_temp_img(await resp.read()) - else: - with open(path, "wb") as f: - f.write(await resp.read()) - return path + with open(path, "wb") as f: + f.write(await resp.read()) + return path except (aiohttp.ClientConnectorSSLError, aiohttp.ClientConnectorCertificateError): - # 关闭SSL验证 + # 关闭SSL验证（仅在证书验证失败时作为fallback） + logger.warning( + f"SSL certificate verification failed for {url}. " + "Disabling SSL verification (CERT_NONE) as a fallback. " + "This is insecure and exposes the application to man-in-the-middle attacks. " + "Please investigate and resolve certificate issues." + ) ssl_context = ssl.create_default_context() - ssl_context.set_ciphers("DEFAULT") + ssl_context.check_hostname = False + ssl_context.verify_mode = ssl.CERT_NONE async with aiohttp.ClientSession() as session: if post: - async with session.get(url, ssl=ssl_context) as resp: - return save_temp_img(await resp.read()) + async with session.post(url, json=post_data, ssl=ssl_context) as resp: + if not path: + return save_temp_img(await resp.read()) + with open(path, "wb") as f: + f.write(await resp.read()) + return path else: async with session.get(url, ssl=ssl_context) as resp: - return save_temp_img(await resp.read()) + if not path: + return save_temp_img(await resp.read()) + with open(path, "wb") as f: + f.write(await resp.read()) + return path except Exception as e: raise e async def download_file(url: str, path: str, show_progress: bool = False): - """ - 从指定 url 下载文件到指定路径 path - """ + """从指定 url 下载文件到指定路径 path""" try: ssl_context = ssl.create_default_context( - cafile=certifi.where() + cafile=certifi.where(), ) # 使用 certifi 提供的 CA 证书 connector = aiohttp.TCPConnector(ssl=ssl_context) async with aiohttp.ClientSession( - trust_env=True, connector=connector + trust_env=True, + connector=connector, ) as session: async with session.get(url, timeout=1800) as resp: if resp.status != 200: @@ -150,16 +161,30 @@ async def download_file(url: str, path: str, show_progress: bool = False): f.write(chunk) downloaded_size += len(chunk) if show_progress: - elapsed_time = time.time() - start_time + elapsed_time = ( + time.time() - start_time + if time.time() - start_time > 0 + else 1 + ) speed = downloaded_size / 1024 / elapsed_time # KB/s print( f"\r下载进度: {downloaded_size / total_size:.2%} 速度: {speed:.2f} KB/s", end="", ) except (aiohttp.ClientConnectorSSLError, aiohttp.ClientConnectorCertificateError): - # 关闭SSL验证 + # 关闭SSL验证（仅在证书验证失败时作为fallback） + logger.warning( + "SSL 证书验证失败，已关闭 SSL 验证（不安全，仅用于临时下载）。请检查目标服务器的证书配置。" + ) + logger.warning( + f"SSL certificate verification failed for {url}. " + "Falling back to unverified connection (CERT_NONE). " + "This is insecure and exposes the application to man-in-the-middle attacks. " + "Please investigate certificate issues with the remote server." + ) ssl_context = ssl.create_default_context() - ssl_context.set_ciphers("DEFAULT") + ssl_context.check_hostname = False + ssl_context.verify_mode = ssl.CERT_NONE async with aiohttp.ClientSession() as session: async with session.get(url, ssl=ssl_context, timeout=120) as resp: total_size = int(resp.headers.get("content-length", 0)) @@ -209,7 +234,7 @@ async def get_dashboard_version(): if os.path.exists(dist_dir): version_file = os.path.join(dist_dir, "assets", "version") if os.path.exists(version_file): - with open(version_file, "r") as f: + with open(version_file, encoding="utf-8") as f: v = f.read().strip() return v return None @@ -221,32 +246,42 @@ async def download_dashboard( latest: bool = True, version: str | None = None, proxy: str | None = None, -): +) -> None: """下载管理面板文件""" if path is None: - path = os.path.join(get_astrbot_data_path(), "dashboard.zip") + zip_path = Path(get_astrbot_data_path()).absolute() / "dashboard.zip" + else: + zip_path = Path(path).absolute() if latest or len(str(version)) != 40: ver_name = "latest" if latest else version dashboard_release_url = f"https://astrbot-registry.soulter.top/download/astrbot-dashboard/{ver_name}/dist.zip" logger.info( - f"准备下载指定发行版本的 AstrBot WebUI 文件: {dashboard_release_url}" + f"准备下载指定发行版本的 AstrBot WebUI 文件: {dashboard_release_url}", ) try: - await download_file(dashboard_release_url, path, show_progress=True) + await download_file( + dashboard_release_url, + str(zip_path), + show_progress=True, + ) except BaseException as _: if latest: - dashboard_release_url = "https://github.com/Soulter/AstrBot/releases/latest/download/dist.zip" + dashboard_release_url = "https://github.com/AstrBotDevs/AstrBot/releases/latest/download/dist.zip" else: - dashboard_release_url = f"https://github.com/Soulter/AstrBot/releases/download/{version}/dist.zip" + dashboard_release_url = f"https://github.com/AstrBotDevs/AstrBot/releases/download/{version}/dist.zip" if proxy: dashboard_release_url = f"{proxy}/{dashboard_release_url}" - await download_file(dashboard_release_url, path, show_progress=True) + await download_file( + dashboard_release_url, + str(zip_path), + show_progress=True, + ) else: url = f"https://github.com/AstrBotDevs/astrbot-release-harbour/releases/download/release-{version}/dist.zip" logger.info(f"准备下载指定版本的 AstrBot WebUI: {url}") if proxy: url = f"{proxy}/{url}" - await download_file(url, path, show_progress=True) - with zipfile.ZipFile(path, "r") as z: + await download_file(url, str(zip_path), show_progress=True) + with zipfile.ZipFile(zip_path, "r") as z: z.extractall(extract_path) diff --git a/astrbot/core/utils/log_pipe.py b/astrbot/core/utils/log_pipe.py index bf5402f1..2e931dd8 100644 --- a/astrbot/core/utils/log_pipe.py +++ b/astrbot/core/utils/log_pipe.py @@ -1,5 +1,5 @@ -import threading import os +import threading from logging import Logger diff --git a/astrbot/core/utils/metrics.py b/astrbot/core/utils/metrics.py index 7fe9bde0..f12019e3 100644 --- a/astrbot/core/utils/metrics.py +++ b/astrbot/core/utils/metrics.py @@ -1,10 +1,12 @@ -import aiohttp -import sys import os import socket +import sys import uuid -from astrbot.core.config import VERSION + +import aiohttp + from astrbot.core import db_helper, logger +from astrbot.core.config import VERSION class Metric: @@ -21,7 +23,7 @@ class Metric: if os.path.exists(id_file): try: - with open(id_file, "r") as f: + with open(id_file) as f: Metric._iid_cache = f.read().strip() return Metric._iid_cache except Exception: @@ -39,8 +41,7 @@ class Metric: @staticmethod async def upload(**kwargs): - """ - 上传相关非敏感的指标以更好地了解 AstrBot 的使用情况。上传的指标不会包含任何有关消息文本、用户信息等敏感信息。 + """上传相关非敏感的指标以更好地了解 AstrBot 的使用情况。上传的指标不会包含任何有关消息文本、用户信息等敏感信息。 Powered by TickStats. """ @@ -64,7 +65,6 @@ class Metric: ) except Exception as e: logger.error(f"保存指标到数据库失败: {e}") - pass try: async with aiohttp.ClientSession(trust_env=True) as session: diff --git a/astrbot/core/utils/path_util.py b/astrbot/core/utils/path_util.py index 0d8511f0..9520d481 100644 --- a/astrbot/core/utils/path_util.py +++ b/astrbot/core/utils/path_util.py @@ -19,24 +19,23 @@ def path_Mapping(mappings, srcPath: str) -> str: # 切割后大于4个项目，或者只有1个项目，那肯定是错误的，只能是2，3，4个项目 logger.warning(f"路径映射规则错误: {mapping}") continue - else: - # rule.len == 3 or 4 - if os.path.exists(rule[0] + ":" + rule[1]): - # 前面两个项目合并路径存在，说明是本地Window路径。后面一个或两个项目组成的路径本地大概率无法解析，直接拼接 - from_ = rule[0] + ":" + rule[1] - if len(rule) == 3: - to_ = rule[2] - else: - to_ = rule[2] + ":" + rule[3] + # rule.len == 3 or 4 + elif os.path.exists(rule[0] + ":" + rule[1]): + # 前面两个项目合并路径存在，说明是本地Window路径。后面一个或两个项目组成的路径本地大概率无法解析，直接拼接 + from_ = rule[0] + ":" + rule[1] + if len(rule) == 3: + to_ = rule[2] else: - # 前面两个项目合并路径不存在，说明第一个项目是本地Linux路径，后面一个或两个项目直接拼接。 - from_ = rule[0] - if len(rule) == 3: - to_ = rule[1] + ":" + rule[2] - else: - # 这种情况下存在四个项目，说明规则也是错误的 - logger.warning(f"路径映射规则错误: {mapping}") - continue + to_ = rule[2] + ":" + rule[3] + else: + # 前面两个项目合并路径不存在，说明第一个项目是本地Linux路径，后面一个或两个项目直接拼接。 + from_ = rule[0] + if len(rule) == 3: + to_ = rule[1] + ":" + rule[2] + else: + # 这种情况下存在四个项目，说明规则也是错误的 + logger.warning(f"路径映射规则错误: {mapping}") + continue from_ = from_.removesuffix("/") from_ = from_.removesuffix("\\") diff --git a/astrbot/core/utils/pip_installer.py b/astrbot/core/utils/pip_installer.py index 88cc2130..6076a114 100644 --- a/astrbot/core/utils/pip_installer.py +++ b/astrbot/core/utils/pip_installer.py @@ -1,20 +1,20 @@ -import logging import asyncio +import logging import sys logger = logging.getLogger("astrbot") class PipInstaller: - def __init__(self, pip_install_arg: str, pypi_index_url: str = None): + def __init__(self, pip_install_arg: str, pypi_index_url: str | None = None): self.pip_install_arg = pip_install_arg self.pypi_index_url = pypi_index_url async def install( self, - package_name: str = None, - requirements_path: str = None, - mirror: str = None, + package_name: str | None = None, + requirements_path: str | None = None, + mirror: str | None = None, ): args = ["install"] if package_name: diff --git a/astrbot/core/utils/session_waiter.py b/astrbot/core/utils/session_waiter.py index c27a5411..33b7cb17 100644 --- a/astrbot/core/utils/session_waiter.py +++ b/astrbot/core/utils/session_waiter.py @@ -1,24 +1,22 @@ -""" -会话控制 -""" +"""会话控制""" import abc import asyncio -import time -import functools import copy +import functools +import time +from collections.abc import Awaitable, Callable +from typing import Any + import astrbot.core.message.components as Comp -from typing import Dict, Any, Callable, Awaitable, List from astrbot.core.platform import AstrMessageEvent -USER_SESSIONS: Dict[str, "SessionWaiter"] = {} # 存储 SessionWaiter 实例 -FILTERS: List["SessionFilter"] = [] # 存储 SessionFilter 实例 +USER_SESSIONS: dict[str, "SessionWaiter"] = {} # 存储 SessionWaiter 实例 +FILTERS: list["SessionFilter"] = [] # 存储 SessionFilter 实例 class SessionController: - """ - 控制一个 Session 是否已经结束 - """ + """控制一个 Session 是否已经结束""" def __init__(self): self.future = asyncio.Future() @@ -29,7 +27,7 @@ class SessionController: self.timeout: float | int = None """上次保持(keep)开始时的超时时间""" - self.history_chains: List[List[Comp.BaseMessageComponent]] = [] + self.history_chains: list[list[Comp.BaseMessageComponent]] = [] def stop(self, error: Exception = None): """立即结束这个会话""" @@ -39,13 +37,14 @@ class SessionController: else: self.future.set_result(None) - def keep(self, timeout: float | int = 0, reset_timeout=False): + def keep(self, timeout: float = 0, reset_timeout=False): """保持这个会话 Args: timeout (float): 必填。会话超时时间。 当 reset_timeout 设置为 True 时, 代表重置超时时间, timeout 必须 > 0, 如果 <= 0 则立即结束会话。 当 reset_timeout 设置为 False 时, 代表继续维持原来的超时时间, 新 timeout = 原来剩余的timeout + timeout (可以 < 0) + """ new_ts = time.time() @@ -81,7 +80,7 @@ class SessionController: pass # 避免报错 # finally: - def get_history_chains(self) -> List[List[Comp.BaseMessageComponent]]: + def get_history_chains(self) -> list[list[Comp.BaseMessageComponent]]: """获取历史消息链""" return self.history_chains @@ -92,7 +91,6 @@ class SessionFilter: @abc.abstractmethod def filter(self, event: AstrMessageEvent) -> str: """根据事件返回一个会话标识符""" - pass class DefaultSessionFilter(SessionFilter): @@ -120,7 +118,9 @@ class SessionWaiter: """需要保证一个 session 同时只有一个 trigger""" async def register_wait( - self, handler: Callable[[str], Awaitable[Any]], timeout: int = 30 + self, + handler: Callable[[str], Awaitable[Any]], + timeout: int = 30, ) -> Any: """等待外部输入并处理""" self.handler = handler @@ -149,7 +149,7 @@ class SessionWaiter: @classmethod async def trigger(cls, session_id: str, event: AstrMessageEvent): """外部输入触发会话处理""" - session = USER_SESSIONS.get(session_id, None) + session = USER_SESSIONS.get(session_id) if not session or session.session_controller.future.done(): return @@ -157,7 +157,7 @@ class SessionWaiter: if not session.session_controller.future.done(): if session.record_history_chains: session.session_controller.history_chains.append( - [copy.deepcopy(comp) for comp in event.get_messages()] + [copy.deepcopy(comp) for comp in event.get_messages()], ) try: # TODO: 这里使用 create_task，跟踪 task，防止超时后这里 handler 仍然在执行 @@ -167,8 +167,7 @@ class SessionWaiter: def session_waiter(timeout: int = 30, record_history_chains: bool = False): - """ - 装饰器：自动将函数注册为 SessionWaiter 处理函数，并等待外部输入触发执行。 + """装饰器：自动将函数注册为 SessionWaiter 处理函数，并等待外部输入触发执行。 :param timeout: 超时时间（秒） :param record_history_chain: 是否自动记录历史消息链。可以通过 controller.get_history_chains() 获取。深拷贝。 diff --git a/astrbot/core/utils/shared_preferences.py b/astrbot/core/utils/shared_preferences.py index c1368f18..c6b4c5ed 100644 --- a/astrbot/core/utils/shared_preferences.py +++ b/astrbot/core/utils/shared_preferences.py @@ -1,11 +1,12 @@ -from astrbot.core.db import BaseDatabase -from astrbot.core.db.po import Preference -import threading import asyncio import os -from typing import TypeVar, Any, overload -from .astrbot_path import get_astrbot_data_path +import threading +from typing import Any, TypeVar, overload +from astrbot.core.db import BaseDatabase +from astrbot.core.db.po import Preference + +from .astrbot_path import get_astrbot_data_path _VT = TypeVar("_VT") @@ -14,7 +15,8 @@ class SharedPreferences: def __init__(self, db_helper: BaseDatabase, json_storage_path=None): if json_storage_path is None: json_storage_path = os.path.join( - get_astrbot_data_path(), "shared_preferences.json" + get_astrbot_data_path(), + "shared_preferences.json", ) self.path = json_storage_path self.db_helper = db_helper @@ -38,13 +40,15 @@ class SharedPreferences: else: ret = default return ret - else: - raise ValueError( - "scope_id and key cannot be None when getting a specific preference." - ) + raise ValueError( + "scope_id and key cannot be None when getting a specific preference.", + ) async def range_get_async( - self, scope: str, scope_id: str | None = None, key: str | None = None + self, + scope: str, + scope_id: str | None = None, + key: str | None = None, ) -> list[Preference]: """获取指定范围的偏好设置 Note: 返回 Preference 列表，其中的 value 属性是一个 dict，value["val"] 为值。scope_id 和 key 可以为 None，这时返回该范围下所有的偏好设置。 @@ -54,21 +58,33 @@ class SharedPreferences: @overload async def session_get( - self, umo: None, key: str, default: Any = None + self, + umo: None, + key: str, + default: Any = None, ) -> list[Preference]: ... @overload async def session_get( - self, umo: str, key: None, default: Any = None + self, + umo: str, + key: None, + default: Any = None, ) -> list[Preference]: ... @overload async def session_get( - self, umo: None, key: None, default: Any = None + self, + umo: None, + key: None, + default: Any = None, ) -> list[Preference]: ... async def session_get( - self, umo: str | None, key: str | None = None, default: _VT = None + self, + umo: str | None, + key: str | None = None, + default: _VT = None, ) -> _VT | list[Preference]: """获取会话范围的偏好设置 @@ -85,7 +101,9 @@ class SharedPreferences: async def global_get(self, key: str, default: _VT = None) -> _VT: ... async def global_get( - self, key: str | None, default: _VT = None + self, + key: str | None, + default: _VT = None, ) -> _VT | list[Preference]: """获取全局范围的偏好设置 @@ -98,7 +116,10 @@ class SharedPreferences: async def put_async(self, scope: str, scope_id: str, key: str, value: Any): """设置指定范围和键的偏好设置""" await self.db_helper.insert_preference_or_update( - scope, scope_id, key, {"val": value} + scope, + scope_id, + key, + {"val": value}, ) async def session_put(self, umo: str, key: str, value: Any): @@ -139,7 +160,7 @@ class SharedPreferences: if scope_id is None or key is None: # result = asyncio.run(self.range_get_async(scope, scope_id, key)) raise ValueError( - "scope_id and key cannot be None when getting a specific preference." + "scope_id and key cannot be None when getting a specific preference.", ) result = asyncio.run_coroutine_threadsafe( self.get_async(scope or "unknown", scope_id or "unknown", key, default), @@ -149,11 +170,15 @@ class SharedPreferences: return result if result is not None else default def range_get( - self, scope: str, scope_id: str | None = None, key: str | None = None + self, + scope: str, + scope_id: str | None = None, + key: str | None = None, ) -> list[Preference]: """获取指定范围的偏好设置（已弃用）""" result = asyncio.run_coroutine_threadsafe( - self.range_get_async(scope, scope_id, key), self._sync_loop + self.range_get_async(scope, scope_id, key), + self._sync_loop, ).result() return result diff --git a/astrbot/core/utils/t2i/__init__.py b/astrbot/core/utils/t2i/__init__.py index 8ce209ad..5038a46f 100644 --- a/astrbot/core/utils/t2i/__init__.py +++ b/astrbot/core/utils/t2i/__init__.py @@ -8,6 +8,9 @@ class RenderStrategy(ABC): @abstractmethod def render_custom_template( - self, tmpl_str: str, tmpl_data: dict, return_url: bool + self, + tmpl_str: str, + tmpl_data: dict, + return_url: bool, ) -> str: pass diff --git a/astrbot/core/utils/t2i/network_strategy.py b/astrbot/core/utils/t2i/network_strategy.py index c43f9ed2..7ebba566 100644 --- a/astrbot/core/utils/t2i/network_strategy.py +++ b/astrbot/core/utils/t2i/network_strategy.py @@ -1,14 +1,17 @@ -import aiohttp import asyncio -import ssl -import certifi import logging import random -from . import RenderStrategy +import ssl + +import aiohttp +import certifi + from astrbot.core.config import VERSION from astrbot.core.utils.io import download_image_by_url from astrbot.core.utils.t2i.template_manager import TemplateManager +from . import RenderStrategy + ASTRBOT_T2I_DEFAULT_ENDPOINT = "https://t2i.soulter.top/text2img" logger = logging.getLogger("astrbot") @@ -38,7 +41,7 @@ class NetworkRenderStrategy(RenderStrategy): try: async with aiohttp.ClientSession() as session: async with session.get( - "https://api.soulter.top/astrbot/t2i-endpoints" + "https://api.soulter.top/astrbot/t2i-endpoints", ) as resp: if resp.status == 200: data = await resp.json() @@ -49,14 +52,13 @@ class NetworkRenderStrategy(RenderStrategy): if ep.get("active") and ep.get("url") ] logger.info( - f"Successfully got {len(self.endpoints)} official T2I endpoints." + f"Successfully got {len(self.endpoints)} official T2I endpoints.", ) except Exception as e: logger.error(f"Failed to get official endpoints: {e}") def _clean_url(self, url: str): - if url.endswith("/"): - url = url[:-1] + url = url.removesuffix("/") if not url.endswith("text2img"): url += "/text2img" return url @@ -69,7 +71,6 @@ class NetworkRenderStrategy(RenderStrategy): options: dict | None = None, ) -> str: """使用自定义文转图模板""" - default_options = {"full_page": True, "type": "jpeg", "quality": 40} if options: default_options |= options @@ -89,21 +90,26 @@ class NetworkRenderStrategy(RenderStrategy): if return_url: ssl_context = ssl.create_default_context(cafile=certifi.where()) connector = aiohttp.TCPConnector(ssl=ssl_context) - async with aiohttp.ClientSession( - trust_env=True, connector=connector - ) as session: - async with session.post( - f"{endpoint}/generate", json=post_data - ) as resp: - if resp.status == 200: - ret = await resp.json() - return f"{endpoint}/{ret['data']['id']}" - else: - raise Exception(f"HTTP {resp.status}") + async with ( + aiohttp.ClientSession( + trust_env=True, + connector=connector, + ) as session, + session.post( + f"{endpoint}/generate", + json=post_data, + ) as resp, + ): + if resp.status == 200: + ret = await resp.json() + return f"{endpoint}/{ret['data']['id']}" + raise Exception(f"HTTP {resp.status}") else: # download_image_by_url 失败时抛异常 return await download_image_by_url( - f"{endpoint}/generate", post=True, post_data=post_data + f"{endpoint}/generate", + post=True, + post_data=post_data, ) except Exception as e: last_exception = e @@ -114,15 +120,18 @@ class NetworkRenderStrategy(RenderStrategy): raise RuntimeError(f"All endpoints failed: {last_exception}") async def render( - self, text: str, return_url: bool = False, template_name: str | None = "base" + self, + text: str, + return_url: bool = False, + template_name: str | None = "base", ) -> str: - """ - 返回图像的文件路径 - """ + """返回图像的文件路径""" if not template_name: template_name = "base" tmpl_str = await self.get_template(name=template_name) text = text.replace("`", "\\`") return await self.render_custom_template( - tmpl_str, {"text": text, "version": f"v{VERSION}"}, return_url + tmpl_str, + {"text": text, "version": f"v{VERSION}"}, + return_url, ) diff --git a/astrbot/core/utils/t2i/renderer.py b/astrbot/core/utils/t2i/renderer.py index 122189f9..2ce7a5eb 100644 --- a/astrbot/core/utils/t2i/renderer.py +++ b/astrbot/core/utils/t2i/renderer.py @@ -1,7 +1,8 @@ -from .network_strategy import NetworkRenderStrategy -from .local_strategy import LocalRenderStrategy from astrbot.core.log import LogManager +from .local_strategy import LocalRenderStrategy +from .network_strategy import NetworkRenderStrategy + logger = LogManager.GetLogger(log_name="astrbot") @@ -30,7 +31,10 @@ class HtmlRenderer: @example: 参见 https://astrbot.app 插件开发部分。 """ return await self.network_strategy.render_custom_template( - tmpl_str, tmpl_data, return_url, options + tmpl_str, + tmpl_data, + return_url, + options, ) async def render_t2i( @@ -44,11 +48,13 @@ class HtmlRenderer: if use_network: try: return await self.network_strategy.render( - text, return_url=return_url, template_name=template_name + text, + return_url=return_url, + template_name=template_name, ) except BaseException as e: logger.error( - f"Failed to render image via AstrBot API: {e}. Falling back to local rendering." + f"Failed to render image via AstrBot API: {e}. Falling back to local rendering.", ) return await self.local_strategy.render(text) else: diff --git a/astrbot/core/utils/t2i/template_manager.py b/astrbot/core/utils/t2i/template_manager.py index b441a908..6d44f735 100644 --- a/astrbot/core/utils/t2i/template_manager.py +++ b/astrbot/core/utils/t2i/template_manager.py @@ -2,12 +2,12 @@ import os import shutil + from astrbot.core.utils.astrbot_path import get_astrbot_data_path, get_astrbot_path class TemplateManager: - """ - 负责管理 t2i HTML 模板的 CRUD 和重置操作。 + """负责管理 t2i HTML 模板的 CRUD 和重置操作。 采用“用户覆盖内置”策略：用户模板存储在 data 目录中，并优先于内置模板加载。 所有创建、更新、删除操作仅影响用户目录，以确保更新框架时用户数据安全。 """ @@ -16,7 +16,12 @@ class TemplateManager: def __init__(self): self.builtin_template_dir = os.path.join( - get_astrbot_path(), "astrbot", "core", "utils", "t2i", "template" + get_astrbot_path(), + "astrbot", + "core", + "utils", + "t2i", + "template", ) self.user_template_dir = os.path.join(get_astrbot_data_path(), "t2i_templates") @@ -43,12 +48,11 @@ class TemplateManager: def _read_file(self, path: str) -> str: """读取文件内容。""" - with open(path, "r", encoding="utf-8") as f: + with open(path, encoding="utf-8") as f: return f.read() def list_templates(self) -> list[dict]: - """ - 列出所有可用模板。 + """列出所有可用模板。 该列表是内置模板和用户模板的合并视图，用户模板将覆盖同名的内置模板。 """ dirs_to_scan = [self.builtin_template_dir, self.user_template_dir] @@ -63,8 +67,7 @@ class TemplateManager: ] def get_template(self, name: str) -> str: - """ - 获取指定模板的内容。 + """获取指定模板的内容。 优先从用户目录加载，如果不存在则回退到内置目录。 """ user_path = self._get_user_template_path(name) @@ -86,8 +89,7 @@ class TemplateManager: f.write(content) def update_template(self, name: str, content: str): - """ - 更新一个模板。此操作始终写入用户目录。 + """更新一个模板。此操作始终写入用户目录。 如果更新的是一个内置模板，此操作实际上会在用户目录中创建一个修改后的副本， 从而实现对内置模板的“覆盖”。 """ @@ -96,8 +98,7 @@ class TemplateManager: f.write(content) def delete_template(self, name: str): - """ - 仅删除用户目录中的模板文件。 + """仅删除用户目录中的模板文件。 如果删除的是一个覆盖了内置模板的用户模板，这将有效地“恢复”到内置版本。 """ path = self._get_user_template_path(name) @@ -106,7 +107,5 @@ class TemplateManager: os.remove(path) def reset_default_template(self): - """ - 将核心模板从内置目录强制重置到用户目录。 - """ + """将核心模板从内置目录强制重置到用户目录。""" self._copy_core_templates(overwrite=True) diff --git a/astrbot/core/utils/tencent_record_helper.py b/astrbot/core/utils/tencent_record_helper.py index 2c97a01e..9cc36571 100644 --- a/astrbot/core/utils/tencent_record_helper.py +++ b/astrbot/core/utils/tencent_record_helper.py @@ -1,10 +1,11 @@ +import asyncio import base64 -import wave import os import subprocess -from io import BytesIO -import asyncio import tempfile +import wave +from io import BytesIO + from astrbot.core import logger from astrbot.core.utils.astrbot_path import get_astrbot_data_path @@ -35,7 +36,7 @@ async def wav_to_tencent_silk(wav_path: str, output_path: str) -> int: import pilk except (ImportError, ModuleNotFoundError) as _: raise Exception( - "pilk 模块未安装，请前往管理面板->控制台->安装pip库 安装 pilk 这个库" + "pilk 模块未安装，请前往管理面板->控制台->安装pip库 安装 pilk 这个库", ) # with wave.open(wav_path, 'rb') as wav: # wav_data = wav.readframes(wav.getnframes()) @@ -60,8 +61,7 @@ async def wav_to_tencent_silk(wav_path: str, output_path: str) -> int: async def convert_to_pcm_wav(input_path: str, output_path: str) -> str: - """ - 将 MP3 或其他音频格式转换为 PCM 16bit WAV，采样率24000Hz，单声道。 + """将 MP3 或其他音频格式转换为 PCM 16bit WAV，采样率24000Hz，单声道。 若转换失败则抛出异常。 """ try: @@ -99,13 +99,11 @@ async def convert_to_pcm_wav(input_path: str, output_path: str) -> str: if os.path.exists(output_path) and os.path.getsize(output_path) > 0: return output_path - else: - raise RuntimeError("生成的WAV文件不存在或为空") + raise RuntimeError("生成的WAV文件不存在或为空") async def audio_to_tencent_silk_base64(audio_path: str) -> tuple[str, float]: - """ - 将 MP3/WAV 文件转为 Tencent Silk 并返回 base64 编码与时长（秒）。 + """将 MP3/WAV 文件转为 Tencent Silk 并返回 base64 编码与时长（秒）。 参数: - audio_path: 输入音频文件路径（.mp3 或 .wav） @@ -125,7 +123,9 @@ async def audio_to_tencent_silk_base64(audio_path: str) -> tuple[str, float]: # 是否需要转换为 WAV ext = os.path.splitext(audio_path)[1].lower() temp_wav = tempfile.NamedTemporaryFile( - suffix=".wav", delete=False, dir=temp_dir + suffix=".wav", + delete=False, + dir=temp_dir, ).name if ext != ".wav": @@ -140,12 +140,18 @@ async def audio_to_tencent_silk_base64(audio_path: str) -> tuple[str, float]: rate = wav_file.getframerate() silk_path = tempfile.NamedTemporaryFile( - suffix=".silk", delete=False, dir=temp_dir + suffix=".silk", + delete=False, + dir=temp_dir, ).name try: duration = await asyncio.to_thread( - pilk.encode, wav_path, silk_path, pcm_rate=rate, tencent=True + pilk.encode, + wav_path, + silk_path, + pcm_rate=rate, + tencent=True, ) with open(silk_path, "rb") as f: diff --git a/astrbot/core/utils/version_comparator.py b/astrbot/core/utils/version_comparator.py index f7ad65fc..e3bf7495 100644 --- a/astrbot/core/utils/version_comparator.py +++ b/astrbot/core/utils/version_comparator.py @@ -38,15 +38,15 @@ class VersionComparator: for i in range(length): if v1_parts[i] > v2_parts[i]: return 1 - elif v1_parts[i] < v2_parts[i]: + if v1_parts[i] < v2_parts[i]: return -1 # 比较预发布标签 if v1_prerelease is None and v2_prerelease is not None: return 1 # 没有预发布标签的版本高于有预发布标签的版本 - elif v1_prerelease is not None and v2_prerelease is None: + if v1_prerelease is not None and v2_prerelease is None: return -1 # 有预发布标签的版本低于没有预发布标签的版本 - elif v1_prerelease is not None and v2_prerelease is not None: + if v1_prerelease is not None and v2_prerelease is not None: len_pre = max(len(v1_prerelease), len(v2_prerelease)) for i in range(len_pre): p1 = v1_prerelease[i] if i < len(v1_prerelease) else None @@ -54,21 +54,18 @@ class VersionComparator: if p1 is None and p2 is not None: return -1 - elif p1 is not None and p2 is None: + if p1 is not None and p2 is None: return 1 - elif isinstance(p1, int) and isinstance(p2, str): + if isinstance(p1, int) and isinstance(p2, str): return -1 - elif isinstance(p1, str) and isinstance(p2, int): + if isinstance(p1, str) and isinstance(p2, int): return 1 - elif isinstance(p1, int) and isinstance(p2, int): + if (isinstance(p1, int) and isinstance(p2, int)) or ( + isinstance(p1, str) and isinstance(p2, str) + ): if p1 > p2: return 1 - elif p1 < p2: - return -1 - elif isinstance(p1, str) and isinstance(p2, str): - if p1 > p2: - return 1 - elif p1 < p2: + if p1 < p2: return -1 return 0 # 预发布标签完全相同 diff --git a/astrbot/core/zip_updator.py b/astrbot/core/zip_updator.py index 7e5f3bfb..728dfdab 100644 --- a/astrbot/core/zip_updator.py +++ b/astrbot/core/zip_updator.py @@ -1,14 +1,14 @@ -import aiohttp import os import re -import zipfile import shutil - import ssl +import zipfile + +import aiohttp import certifi -from astrbot.core.utils.io import on_error, download_file from astrbot.core import logger +from astrbot.core.utils.io import download_file, on_error from astrbot.core.utils.version_comparator import VersionComparator @@ -18,7 +18,10 @@ class ReleaseInfo: body: str def __init__( - self, version: str = "", published_at: str = "", body: str = "" + self, + version: str = "", + published_at: str = "", + body: str = "", ) -> None: self.version = version self.published_at = published_at @@ -34,29 +37,31 @@ class RepoZipUpdator: self.rm_on_error = on_error async def fetch_release_info(self, url: str, latest: bool = True) -> list: - """ - 请求版本信息。 + """请求版本信息。 返回一个列表，每个元素是一个字典，包含版本号、发布时间、更新内容、commit hash等信息。 """ try: ssl_context = ssl.create_default_context( - cafile=certifi.where() + cafile=certifi.where(), ) # 新增：创建基于 certifi 的 SSL 上下文 connector = aiohttp.TCPConnector( - ssl=ssl_context + ssl=ssl_context, ) # 新增：使用 TCPConnector 指定 SSL 上下文 - async with aiohttp.ClientSession( - trust_env=True, connector=connector - ) as session: - async with session.get(url) as response: - # 检查 HTTP 状态码 - if response.status != 200: - text = await response.text() - logger.error( - f"请求 {url} 失败，状态码: {response.status}, 内容: {text}" - ) - raise Exception(f"请求失败，状态码: {response.status}") - result = await response.json() + async with ( + aiohttp.ClientSession( + trust_env=True, + connector=connector, + ) as session, + session.get(url) as response, + ): + # 检查 HTTP 状态码 + if response.status != 200: + text = await response.text() + logger.error( + f"请求 {url} 失败，状态码: {response.status}, 内容: {text}", + ) + raise Exception(f"请求失败，状态码: {response.status}") + result = await response.json() if not result: return [] # if latest: @@ -72,7 +77,7 @@ class RepoZipUpdator: "body": release["body"], "tag_name": release["tag_name"], "zipball_url": release["zipball_url"], - } + }, ) except Exception as e: logger.error(f"解析版本信息时发生异常: {e}") @@ -80,8 +85,7 @@ class RepoZipUpdator: return ret def github_api_release_parser(self, releases: list) -> list: - """ - 解析 GitHub API 返回的 releases 信息。 + """解析 GitHub API 返回的 releases 信息。 返回一个列表，每个元素是一个字典，包含版本号、发布时间、更新内容、commit hash等信息。 """ ret = [] @@ -93,22 +97,25 @@ class RepoZipUpdator: "body": release["body"], "tag_name": release["tag_name"], "zipball_url": release["zipball_url"], - } + }, ) return ret def unzip(self): - raise NotImplementedError() + raise NotImplementedError async def update(self): - raise NotImplementedError() + raise NotImplementedError def compare_version(self, v1: str, v2: str) -> int: """Semver 版本比较""" return VersionComparator.compare_version(v1, v2) async def check_update( - self, url: str, current_version: str, consider_prerelease: bool = True + self, + url: str, + current_version: str, + consider_prerelease: bool = True, ) -> ReleaseInfo | None: update_data = await self.fetch_release_info(url) @@ -157,7 +164,7 @@ class RepoZipUpdator: releases = await self.fetch_release_info(url=release_url) except Exception as e: logger.warning( - f"获取 {author}/{repo} 的 GitHub Releases 失败: {e}，将尝试下载默认分支" + f"获取 {author}/{repo} 的 GitHub Releases 失败: {e}，将尝试下载默认分支", ) releases = [] if not releases: @@ -173,7 +180,7 @@ class RepoZipUpdator: proxy = proxy.rstrip("/") release_url = f"{proxy}/{release_url}" logger.info( - f"检查到设置了镜像站，将使用镜像站下载 {author}/{repo} 仓库源码: {release_url}" + f"检查到设置了镜像站，将使用镜像站下载 {author}/{repo} 仓库源码: {release_url}", ) await download_file(release_url, target_path + ".zip") @@ -194,13 +201,10 @@ class RepoZipUpdator: repo = match.group(2) branch = match.group(4) return author, repo, branch - else: - raise ValueError("无效的 GitHub URL") + raise ValueError("无效的 GitHub URL") def unzip_file(self, zip_path: str, target_dir: str): - """ - 解压缩文件, 并将压缩包内**第一个**文件夹内的文件移动到 target_dir - """ + """解压缩文件, 并将压缩包内**第一个**文件夹内的文件移动到 target_dir""" os.makedirs(target_dir, exist_ok=True) update_dir = "" with zipfile.ZipFile(zip_path, "r") as z: @@ -213,20 +217,19 @@ class RepoZipUpdator: if os.path.isdir(os.path.join(target_dir, update_dir, f)): if os.path.exists(os.path.join(target_dir, f)): shutil.rmtree(os.path.join(target_dir, f), onerror=on_error) - else: - if os.path.exists(os.path.join(target_dir, f)): - os.remove(os.path.join(target_dir, f)) + elif os.path.exists(os.path.join(target_dir, f)): + os.remove(os.path.join(target_dir, f)) shutil.move(os.path.join(target_dir, update_dir, f), target_dir) try: logger.debug( - f"删除临时更新文件: {zip_path} 和 {os.path.join(target_dir, update_dir)}" + f"删除临时更新文件: {zip_path} 和 {os.path.join(target_dir, update_dir)}", ) shutil.rmtree(os.path.join(target_dir, update_dir), onerror=on_error) os.remove(zip_path) except BaseException: logger.warning( - f"删除更新文件失败，可以手动删除 {zip_path} 和 {os.path.join(target_dir, update_dir)}" + f"删除更新文件失败，可以手动删除 {zip_path} 和 {os.path.join(target_dir, update_dir)}", ) def format_name(self, name: str) -> str: diff --git a/astrbot/dashboard/routes/__init__.py b/astrbot/dashboard/routes/__init__.py index ef2fa3e8..b7997cf8 100644 --- a/astrbot/dashboard/routes/__init__.py +++ b/astrbot/dashboard/routes/__init__.py @@ -1,29 +1,31 @@ from .auth import AuthRoute -from .plugin import PluginRoute -from .config import ConfigRoute -from .update import UpdateRoute -from .stat import StatRoute -from .log import LogRoute -from .static_file import StaticFileRoute from .chat import ChatRoute -from .tools import ToolsRoute +from .config import ConfigRoute from .conversation import ConversationRoute from .file import FileRoute -from .session_management import SessionManagementRoute +from .knowledge_base import KnowledgeBaseRoute +from .log import LogRoute from .persona import PersonaRoute +from .plugin import PluginRoute +from .session_management import SessionManagementRoute +from .stat import StatRoute +from .static_file import StaticFileRoute +from .tools import ToolsRoute +from .update import UpdateRoute __all__ = [ "AuthRoute", - "PluginRoute", - "ConfigRoute", - "UpdateRoute", - "StatRoute", - "LogRoute", - "StaticFileRoute", "ChatRoute", - "ToolsRoute", + "ConfigRoute", "ConversationRoute", "FileRoute", - "SessionManagementRoute", + "KnowledgeBaseRoute", + "LogRoute", "PersonaRoute", + "PluginRoute", + "SessionManagementRoute", + "StatRoute", + "StaticFileRoute", + "ToolsRoute", + "UpdateRoute", ] diff --git a/astrbot/dashboard/routes/auth.py b/astrbot/dashboard/routes/auth.py index 87af4b61..4ee0d57d 100644 --- a/astrbot/dashboard/routes/auth.py +++ b/astrbot/dashboard/routes/auth.py @@ -1,10 +1,13 @@ -import jwt -import datetime import asyncio -from .route import Route, Response, RouteContext +import datetime + +import jwt from quart import request -from astrbot.core import DEMO_MODE + from astrbot import logger +from astrbot.core import DEMO_MODE + +from .route import Response, Route, RouteContext class AuthRoute(Route): @@ -37,13 +40,12 @@ class AuthRoute(Route): "token": self.generate_jwt(username), "username": username, "change_pwd_hint": change_pwd_hint, - } + }, ) .__dict__ ) - else: - await asyncio.sleep(3) - return Response().error("用户名或密码错误").__dict__ + await asyncio.sleep(3) + return Response().error("用户名或密码错误").__dict__ async def edit_account(self): if DEMO_MODE: diff --git a/astrbot/dashboard/routes/chat.py b/astrbot/dashboard/routes/chat.py index 71fd3472..1ad78956 100644 --- a/astrbot/dashboard/routes/chat.py +++ b/astrbot/dashboard/routes/chat.py @@ -1,16 +1,19 @@ -import uuid +import asyncio import json import os -import asyncio +import uuid from contextlib import asynccontextmanager -from .route import Route, Response, RouteContext -from astrbot.core.platform.sources.webchat.webchat_queue_mgr import webchat_queue_mgr -from quart import request, Response as QuartResponse, g, make_response -from astrbot.core.db import BaseDatabase + +from quart import Response as QuartResponse +from quart import g, make_response, request + from astrbot.core import logger from astrbot.core.core_lifecycle import AstrBotCoreLifecycle +from astrbot.core.db import BaseDatabase +from astrbot.core.platform.sources.webchat.webchat_queue_mgr import webchat_queue_mgr from astrbot.core.utils.astrbot_path import get_astrbot_data_path -from astrbot.core.platform.astr_message_event import MessageSession + +from .route import Response, Route, RouteContext @asynccontextmanager @@ -32,11 +35,14 @@ class ChatRoute(Route): super().__init__(context) self.routes = { "/chat/send": ("POST", self.chat), - "/chat/new_conversation": ("GET", self.new_conversation), - "/chat/conversations": ("GET", self.get_conversations), - "/chat/get_conversation": ("GET", self.get_conversation), - "/chat/delete_conversation": ("GET", self.delete_conversation), - "/chat/rename_conversation": ("POST", self.rename_conversation), + "/chat/new_session": ("GET", self.new_session), + "/chat/sessions": ("GET", self.get_sessions), + "/chat/get_session": ("GET", self.get_session), + "/chat/delete_session": ("GET", self.delete_webchat_session), + "/chat/update_session_display_name": ( + "POST", + self.update_session_display_name, + ), "/chat/get_file": ("GET", self.get_file), "/chat/post_image": ("POST", self.post_image), "/chat/post_file": ("POST", self.post_file), @@ -49,6 +55,7 @@ class ChatRoute(Route): self.supported_imgs = ["jpg", "jpeg", "png", "gif", "webp"] self.conv_mgr = core_lifecycle.conversation_manager self.platform_history_mgr = core_lifecycle.platform_message_history_manager + self.db = db self.running_convs: dict[str, bool] = {} @@ -70,10 +77,9 @@ class ChatRoute(Route): if filename_ext == ".wav": return QuartResponse(f.read(), mimetype="audio/wav") - elif filename_ext[1:] in self.supported_imgs: + if filename_ext[1:] in self.supported_imgs: return QuartResponse(f.read(), mimetype="image/jpeg") - else: - return QuartResponse(f.read()) + return QuartResponse(f.read()) except (FileNotFoundError, OSError): return Response().error("File access error").__dict__ @@ -96,7 +102,7 @@ class ChatRoute(Route): return Response().error("Missing key: file").__dict__ file = post_data["file"] - filename = f"{str(uuid.uuid4())}" + filename = f"{uuid.uuid4()!s}" # 通过文件格式判断文件类型 if file.content_type.startswith("audio"): filename += ".wav" @@ -113,28 +119,33 @@ class ChatRoute(Route): if "message" not in post_data and "image_url" not in post_data: return Response().error("Missing key: message or image_url").__dict__ - if "conversation_id" not in post_data: - return Response().error("Missing key: conversation_id").__dict__ + if "session_id" not in post_data and "conversation_id" not in post_data: + return ( + Response().error("Missing key: session_id or conversation_id").__dict__ + ) message = post_data["message"] - conversation_id = post_data["conversation_id"] + # conversation_id = post_data["conversation_id"] + session_id = post_data.get("session_id", post_data.get("conversation_id")) image_url = post_data.get("image_url") audio_url = post_data.get("audio_url") selected_provider = post_data.get("selected_provider") selected_model = post_data.get("selected_model") + enable_streaming = post_data.get("enable_streaming", True) # 默认为 True + if not message and not image_url and not audio_url: return ( Response() .error("Message and image_url and audio_url are empty") .__dict__ ) - if not conversation_id: - return Response().error("conversation_id is empty").__dict__ + if not session_id: + return Response().error("session_id is empty").__dict__ - # append user message - webchat_conv_id = await self._get_webchat_conv_id_from_conv_id(conversation_id) + # 追加用户消息 + webchat_conv_id = session_id - # Get conversation-specific queues + # 获取会话特定的队列 back_queue = webchat_queue_mgr.get_or_create_back_queue(webchat_conv_id) new_his = {"type": "user", "message": message} @@ -179,7 +190,7 @@ class ChatRoute(Route): except Exception as e: if not client_disconnected: logger.debug( - f"[WebChat] 用户 {username} 断开聊天长连接。 {e}" + f"[WebChat] 用户 {username} 断开聊天长连接。 {e}", ) client_disconnected = True @@ -197,8 +208,10 @@ class ChatRoute(Route): or not streaming or type == "break" ): - # append bot message + # 追加机器人消息 new_his = {"type": "bot", "message": result_text} + if "reasoning" in result: + new_his["reasoning"] = result["reasoning"] await self.platform_history_mgr.insert( platform_id="webchat", user_id=webchat_conv_id, @@ -209,7 +222,7 @@ class ChatRoute(Route): except BaseException as e: logger.exception(f"WebChat stream unexpected error: {e}", exc_info=True) - # Put message to conversation-specific queue + # 将消息放入会话特定的队列 chat_queue = webchat_queue_mgr.get_or_create_queue(webchat_conv_id) await chat_queue.put( ( @@ -221,8 +234,9 @@ class ChatRoute(Route): "audio_url": audio_url, "selected_provider": selected_provider, "selected_model": selected_model, + "enable_streaming": enable_streaming, }, - ) + ), ) response = await make_response( @@ -237,84 +251,112 @@ class ChatRoute(Route): response.timeout = None # fix SSE auto disconnect issue return response - async def _get_webchat_conv_id_from_conv_id(self, conversation_id: str) -> str: - """从对话 ID 中提取 WebChat 会话 ID - - NOTE: 关于这里为什么要单独做一个 WebChat 的 Conversation ID 出来，这个是为了向前兼容。 - """ - conversation = await self.conv_mgr.get_conversation( - unified_msg_origin="webchat", conversation_id=conversation_id - ) - if not conversation: - raise ValueError(f"Conversation with ID {conversation_id} not found.") - conv_user_id = conversation.user_id - webchat_session_id = MessageSession.from_str(conv_user_id).session_id - if "!" not in webchat_session_id: - raise ValueError(f"Invalid conv user ID: {conv_user_id}") - return webchat_session_id.split("!")[-1] - - async def delete_conversation(self): - conversation_id = request.args.get("conversation_id") - if not conversation_id: - return Response().error("Missing key: conversation_id").__dict__ + async def delete_webchat_session(self): + """Delete a Platform session and all its related data.""" + session_id = request.args.get("session_id") + if not session_id: + return Response().error("Missing key: session_id").__dict__ username = g.get("username", "guest") - # Clean up queues when deleting conversation - webchat_queue_mgr.remove_queues(conversation_id) - webchat_conv_id = await self._get_webchat_conv_id_from_conv_id(conversation_id) - await self.conv_mgr.delete_conversation( - unified_msg_origin=f"webchat:FriendMessage:webchat!{username}!{webchat_conv_id}", - conversation_id=conversation_id, - ) + # 验证会话是否存在且属于当前用户 + session = await self.db.get_platform_session_by_id(session_id) + if not session: + return Response().error(f"Session {session_id} not found").__dict__ + if session.creator != username: + return Response().error("Permission denied").__dict__ + + # 删除该会话下的所有对话 + unified_msg_origin = f"{session.platform_id}:FriendMessage:{session.platform_id}!{username}!{session_id}" + await self.conv_mgr.delete_conversations_by_user_id(unified_msg_origin) + + # 删除消息历史 await self.platform_history_mgr.delete( - platform_id="webchat", user_id=webchat_conv_id, offset_sec=99999999 + platform_id=session.platform_id, + user_id=session_id, + offset_sec=99999999, ) + + # 清理队列（仅对 webchat） + if session.platform_id == "webchat": + webchat_queue_mgr.remove_queues(session_id) + + # 删除会话 + await self.db.delete_platform_session(session_id) + return Response().ok().__dict__ - async def new_conversation(self): + async def new_session(self): + """Create a new Platform session (default: webchat).""" username = g.get("username", "guest") - webchat_conv_id = str(uuid.uuid4()) - conv_id = await self.conv_mgr.new_conversation( - unified_msg_origin=f"webchat:FriendMessage:webchat!{username}!{webchat_conv_id}", - platform_id="webchat", - content=[], + + # 获取可选的 platform_id 参数，默认为 webchat + platform_id = request.args.get("platform_id", "webchat") + + # 创建新会话 + session = await self.db.create_platform_session( + creator=username, + platform_id=platform_id, + is_group=0, ) - return Response().ok(data={"conversation_id": conv_id}).__dict__ - async def rename_conversation(self): - post_data = await request.json - if "conversation_id" not in post_data or "title" not in post_data: - return Response().error("Missing key: conversation_id or title").__dict__ - - conversation_id = post_data["conversation_id"] - title = post_data["title"] - - await self.conv_mgr.update_conversation( - unified_msg_origin="webchat", # fake - conversation_id=conversation_id, - title=title, + return ( + Response() + .ok( + data={ + "session_id": session.session_id, + "platform_id": session.platform_id, + } + ) + .__dict__ ) - return Response().ok(message="重命名成功！").__dict__ - async def get_conversations(self): - conversations = await self.conv_mgr.get_conversations(platform_id="webchat") - # remove content - conversations_ = [] - for conv in conversations: - conv.history = None - conversations_.append(conv) - return Response().ok(data=conversations_).__dict__ + async def get_sessions(self): + """Get all Platform sessions for the current user.""" + username = g.get("username", "guest") - async def get_conversation(self): - conversation_id = request.args.get("conversation_id") - if not conversation_id: - return Response().error("Missing key: conversation_id").__dict__ + # 获取可选的 platform_id 参数 + platform_id = request.args.get("platform_id") - webchat_conv_id = await self._get_webchat_conv_id_from_conv_id(conversation_id) + sessions = await self.db.get_platform_sessions_by_creator( + creator=username, + platform_id=platform_id, + page=1, + page_size=100, # 暂时返回前100个 + ) - # Get platform message history + # 转换为字典格式，并添加额外信息 + sessions_data = [] + for session in sessions: + sessions_data.append( + { + "session_id": session.session_id, + "platform_id": session.platform_id, + "creator": session.creator, + "display_name": session.display_name, + "is_group": session.is_group, + "created_at": session.created_at.astimezone().isoformat(), + "updated_at": session.updated_at.astimezone().isoformat(), + } + ) + + return Response().ok(data=sessions_data).__dict__ + + async def get_session(self): + """Get session information and message history by session_id.""" + session_id = request.args.get("session_id") + if not session_id: + return Response().error("Missing key: session_id").__dict__ + + # 获取会话信息以确定 platform_id + session = await self.db.get_platform_session_by_id(session_id) + platform_id = session.platform_id if session else "webchat" + + # Get platform message history using session_id history_ls = await self.platform_history_mgr.get( - platform_id="webchat", user_id=webchat_conv_id, page=1, page_size=1000 + platform_id=platform_id, + user_id=session_id, + page=1, + page_size=1000, ) history_res = [history.model_dump() for history in history_ls] @@ -324,8 +366,37 @@ class ChatRoute(Route): .ok( data={ "history": history_res, - "is_running": self.running_convs.get(webchat_conv_id, False), - } + "is_running": self.running_convs.get(session_id, False), + }, ) .__dict__ ) + + async def update_session_display_name(self): + """Update a Platform session's display name.""" + post_data = await request.json + + session_id = post_data.get("session_id") + display_name = post_data.get("display_name") + + if not session_id: + return Response().error("Missing key: session_id").__dict__ + if display_name is None: + return Response().error("Missing key: display_name").__dict__ + + username = g.get("username", "guest") + + # 验证会话是否存在且属于当前用户 + session = await self.db.get_platform_session_by_id(session_id) + if not session: + return Response().error(f"Session {session_id} not found").__dict__ + if session.creator != username: + return Response().error("Permission denied").__dict__ + + # 更新 display_name + await self.db.update_platform_session( + session_id=session_id, + display_name=display_name, + ) + + return Response().ok().__dict__ diff --git a/astrbot/dashboard/routes/config.py b/astrbot/dashboard/routes/config.py index 01b9b243..b947d26f 100644 --- a/astrbot/dashboard/routes/config.py +++ b/astrbot/dashboard/routes/config.py @@ -1,27 +1,29 @@ -import typing -import traceback -import os +import asyncio import inspect -from .route import Route, Response, RouteContext -from astrbot.core.provider.entities import ProviderType +import os +import traceback + from quart import request + +from astrbot.core import file_token_service, logger +from astrbot.core.config.astrbot_config import AstrBotConfig from astrbot.core.config.default import ( - DEFAULT_CONFIG, CONFIG_METADATA_2, - DEFAULT_VALUE_MAP, CONFIG_METADATA_3, CONFIG_METADATA_3_SYSTEM, + DEFAULT_CONFIG, + DEFAULT_VALUE_MAP, ) -from astrbot.core.utils.astrbot_path import get_astrbot_path -from astrbot.core.config.astrbot_config import AstrBotConfig from astrbot.core.core_lifecycle import AstrBotCoreLifecycle -from astrbot.core.platform.register import platform_registry, platform_cls_map +from astrbot.core.platform.register import platform_cls_map, platform_registry +from astrbot.core.provider import Provider +from astrbot.core.provider.entities import ProviderType +from astrbot.core.provider.provider import RerankProvider from astrbot.core.provider.register import provider_registry from astrbot.core.star.star import star_registry -from astrbot.core import logger, file_token_service -from astrbot.core.provider import Provider -from astrbot.core.provider.provider import RerankProvider -import asyncio +from astrbot.core.utils.astrbot_path import get_astrbot_path + +from .route import Response, Route, RouteContext def try_cast(value: str, type_: str): @@ -34,9 +36,7 @@ def try_cast(value: str, type_: str): type_ == "float" and isinstance(value, str) and value.replace(".", "", 1).isdigit() - ): - return float(value) - elif type_ == "float" and isinstance(value, int): + ) or (type_ == "float" and isinstance(value, int)): return float(value) elif type_ == "float": try: @@ -45,9 +45,7 @@ def try_cast(value: str, type_: str): return None -def validate_config( - data, schema: dict, is_core: bool -) -> typing.Tuple[typing.List[str], typing.Dict]: +def validate_config(data, schema: dict, is_core: bool) -> tuple[list[str], dict]: errors = [] def validate(data: dict, metadata: dict = schema, path=""): @@ -64,7 +62,7 @@ def validate_config( continue if meta["type"] == "list" and not isinstance(value, list): errors.append( - f"错误的类型 {path}{key}: 期望是 list, 得到了 {type(value).__name__}" + f"错误的类型 {path}{key}: 期望是 list, 得到了 {type(value).__name__}", ) elif ( meta["type"] == "list" @@ -83,31 +81,31 @@ def validate_config( casted = try_cast(value, "int") if casted is None: errors.append( - f"错误的类型 {path}{key}: 期望是 int, 得到了 {type(value).__name__}" + f"错误的类型 {path}{key}: 期望是 int, 得到了 {type(value).__name__}", ) data[key] = casted elif meta["type"] == "float" and not isinstance(value, float): casted = try_cast(value, "float") if casted is None: errors.append( - f"错误的类型 {path}{key}: 期望是 float, 得到了 {type(value).__name__}" + f"错误的类型 {path}{key}: 期望是 float, 得到了 {type(value).__name__}", ) data[key] = casted elif meta["type"] == "bool" and not isinstance(value, bool): errors.append( - f"错误的类型 {path}{key}: 期望是 bool, 得到了 {type(value).__name__}" + f"错误的类型 {path}{key}: 期望是 bool, 得到了 {type(value).__name__}", ) elif meta["type"] in ["string", "text"] and not isinstance(value, str): errors.append( - f"错误的类型 {path}{key}: 期望是 string, 得到了 {type(value).__name__}" + f"错误的类型 {path}{key}: 期望是 string, 得到了 {type(value).__name__}", ) elif meta["type"] == "list" and not isinstance(value, list): errors.append( - f"错误的类型 {path}{key}: 期望是 list, 得到了 {type(value).__name__}" + f"错误的类型 {path}{key}: 期望是 list, 得到了 {type(value).__name__}", ) elif meta["type"] == "object" and not isinstance(value, dict): errors.append( - f"错误的类型 {path}{key}: 期望是 dict, 得到了 {type(value).__name__}" + f"错误的类型 {path}{key}: 期望是 dict, 得到了 {type(value).__name__}", ) if is_core: @@ -130,7 +128,9 @@ def save_config(post_config: dict, config: AstrBotConfig, is_core: bool = False) try: if is_core: errors, post_config = validate_config( - post_config, CONFIG_METADATA_2, is_core + post_config, + CONFIG_METADATA_2, + is_core, ) else: errors, post_config = validate_config(post_config, config.schema, is_core) @@ -146,7 +146,9 @@ def save_config(post_config: dict, config: AstrBotConfig, is_core: bool = False) class ConfigRoute(Route): def __init__( - self, context: RouteContext, core_lifecycle: AstrBotCoreLifecycle + self, + context: RouteContext, + core_lifecycle: AstrBotCoreLifecycle, ) -> None: super().__init__(context) self.core_lifecycle = core_lifecycle @@ -178,6 +180,7 @@ class ConfigRoute(Route): "/config/provider/check_one": ("GET", self.check_one_provider_status), "/config/provider/list": ("GET", self.get_provider_config_list), "/config/provider/model_list": ("GET", self.get_provider_model_list), + "/config/provider/get_embedding_dim": ("POST", self.get_embedding_dim), } self.register_routes() @@ -201,7 +204,7 @@ class ConfigRoute(Route): return Response().ok(message="更新成功").__dict__ except Exception as e: logger.error(traceback.format_exc()) - return Response().error(f"更新路由表失败: {str(e)}").__dict__ + return Response().error(f"更新路由表失败: {e!s}").__dict__ async def update_ucr(self): """更新 UMOP 配置路由表""" @@ -220,7 +223,7 @@ class ConfigRoute(Route): return Response().ok(message="更新成功").__dict__ except Exception as e: logger.error(traceback.format_exc()) - return Response().error(f"更新路由表失败: {str(e)}").__dict__ + return Response().error(f"更新路由表失败: {e!s}").__dict__ async def delete_ucr(self): """删除 UMOP 配置路由表中的一项""" @@ -240,7 +243,7 @@ class ConfigRoute(Route): return Response().ok(message="删除成功").__dict__ except Exception as e: logger.error(traceback.format_exc()) - return Response().error(f"删除路由表项失败: {str(e)}").__dict__ + return Response().error(f"删除路由表项失败: {e!s}").__dict__ async def get_default_config(self): """获取默认配置文件""" @@ -307,13 +310,12 @@ class ConfigRoute(Route): success = self.acm.delete_conf(conf_id) if success: return Response().ok(message="删除成功").__dict__ - else: - return Response().error("删除失败").__dict__ + return Response().error("删除失败").__dict__ except ValueError as e: return Response().error(str(e)).__dict__ except Exception as e: logger.error(traceback.format_exc()) - return Response().error(f"删除配置文件失败: {str(e)}").__dict__ + return Response().error(f"删除配置文件失败: {e!s}").__dict__ async def update_abconf(self): """更新指定 AstrBot 配置文件信息""" @@ -331,13 +333,12 @@ class ConfigRoute(Route): success = self.acm.update_conf_info(conf_id, name=name) if success: return Response().ok(message="更新成功").__dict__ - else: - return Response().error("更新失败").__dict__ + return Response().error("更新失败").__dict__ except ValueError as e: return Response().error(str(e)).__dict__ except Exception as e: logger.error(traceback.format_exc()) - return Response().error(f"更新配置文件失败: {str(e)}").__dict__ + return Response().error(f"更新配置文件失败: {e!s}").__dict__ async def _test_single_provider(self, provider): """辅助函数：测试单个 provider 的可用性""" @@ -354,17 +355,18 @@ class ConfigRoute(Route): "error": None, } logger.debug( - f"Attempting to check provider: {status_info['name']} (ID: {status_info['id']}, Type: {status_info['type']}, Model: {status_info['model']})" + f"Attempting to check provider: {status_info['name']} (ID: {status_info['id']}, Type: {status_info['type']}, Model: {status_info['model']})", ) if provider_capability_type == ProviderType.CHAT_COMPLETION: try: logger.debug(f"Sending 'Ping' to provider: {status_info['name']}") response = await asyncio.wait_for( - provider.text_chat(prompt="REPLY `PONG` ONLY"), timeout=45.0 + provider.text_chat(prompt="REPLY `PONG` ONLY"), + timeout=45.0, ) logger.debug( - f"Received response from {status_info['name']}: {response}" + f"Received response from {status_info['name']}: {response}", ) if response is not None: status_info["status"] = "available" @@ -388,14 +390,14 @@ class ConfigRoute(Route): except Exception as _: pass logger.info( - f"Provider {status_info['name']} (ID: {status_info['id']}) is available. Response snippet: '{response_text_snippet}'" + f"Provider {status_info['name']} (ID: {status_info['id']}) is available. Response snippet: '{response_text_snippet}'", ) else: status_info["error"] = ( "Test call returned None, but expected an LLMResponse object." ) logger.warning( - f"Provider {status_info['name']} (ID: {status_info['id']}) test call returned None." + f"Provider {status_info['name']} (ID: {status_info['id']}) test call returned None.", ) except asyncio.TimeoutError: @@ -403,16 +405,16 @@ class ConfigRoute(Route): "Connection timed out after 45 seconds during test call." ) logger.warning( - f"Provider {status_info['name']} (ID: {status_info['id']}) timed out." + f"Provider {status_info['name']} (ID: {status_info['id']}) timed out.", ) except Exception as e: error_message = str(e) status_info["error"] = error_message logger.warning( - f"Provider {status_info['name']} (ID: {status_info['id']}) is unavailable. Error: {error_message}" + f"Provider {status_info['name']} (ID: {status_info['id']}) is unavailable. Error: {error_message}", ) logger.debug( - f"Traceback for {status_info['name']}:\n{traceback.format_exc()}" + f"Traceback for {status_info['name']}:\n{traceback.format_exc()}", ) elif provider_capability_type == ProviderType.EMBEDDING: @@ -434,7 +436,7 @@ class ConfigRoute(Route): exc_info=True, ) status_info["status"] = "unavailable" - status_info["error"] = f"Embedding test failed: {str(e)}" + status_info["error"] = f"Embedding test failed: {e!s}" elif provider_capability_type == ProviderType.TEXT_TO_SPEECH: try: @@ -449,17 +451,20 @@ class ConfigRoute(Route): ) except Exception as e: logger.error( - f"Error testing TTS provider {provider_name}: {e}", exc_info=True + f"Error testing TTS provider {provider_name}: {e}", + exc_info=True, ) status_info["status"] = "unavailable" - status_info["error"] = f"TTS test failed: {str(e)}" + status_info["error"] = f"TTS test failed: {e!s}" elif provider_capability_type == ProviderType.SPEECH_TO_TEXT: try: logger.debug( - f"Sending health check audio to provider: {status_info['name']}" + f"Sending health check audio to provider: {status_info['name']}", ) sample_audio_path = os.path.join( - get_astrbot_path(), "samples", "stt_health_check.wav" + get_astrbot_path(), + "samples", + "stt_health_check.wav", ) if not os.path.exists(sample_audio_path): status_info["status"] = "unavailable" @@ -467,7 +472,7 @@ class ConfigRoute(Route): "STT test failed: sample audio file not found." ) logger.warning( - f"STT test for {status_info['name']} failed: sample audio file not found at {sample_audio_path}" + f"STT test for {status_info['name']} failed: sample audio file not found at {sample_audio_path}", ) else: text_result = await provider.get_text(sample_audio_path) @@ -479,7 +484,7 @@ class ConfigRoute(Route): else text_result ) logger.info( - f"Provider {status_info['name']} (ID: {status_info['id']}) is available. Response snippet: '{snippet}'" + f"Provider {status_info['name']} (ID: {status_info['id']}) is available. Response snippet: '{snippet}'", ) else: status_info["status"] = "unavailable" @@ -487,14 +492,15 @@ class ConfigRoute(Route): f"STT test failed: unexpected result type {type(text_result)}" ) logger.warning( - f"STT test for {status_info['name']} failed: unexpected result type {type(text_result)}" + f"STT test for {status_info['name']} failed: unexpected result type {type(text_result)}", ) except Exception as e: logger.error( - f"Error testing STT provider {provider_name}: {e}", exc_info=True + f"Error testing STT provider {provider_name}: {e}", + exc_info=True, ) status_info["status"] = "unavailable" - status_info["error"] = f"STT test failed: {str(e)}" + status_info["error"] = f"STT test failed: {e!s}" elif provider_capability_type == ProviderType.RERANK: try: assert isinstance(provider, RerankProvider) @@ -506,11 +512,11 @@ class ConfigRoute(Route): exc_info=True, ) status_info["status"] = "unavailable" - status_info["error"] = f"Rerank test failed: {str(e)}" + status_info["error"] = f"Rerank test failed: {e!s}" else: logger.debug( - f"Provider {provider_name} is not a Chat Completion or Embedding provider. Marking as available without test. Meta: {meta}" + f"Provider {provider_name} is not a Chat Completion or Embedding provider. Marking as available without test. Meta: {meta}", ) status_info["status"] = "available" status_info["error"] = ( @@ -520,7 +526,10 @@ class ConfigRoute(Route): return status_info def _error_response( - self, message: str, status_code: int = 500, log_fn=logger.error + self, + message: str, + status_code: int = 500, + log_fn=logger.error, ): log_fn(message) # 记录更详细的traceback信息，但只在是严重错误时 @@ -533,7 +542,9 @@ class ConfigRoute(Route): provider_id = request.args.get("id") if not provider_id: return self._error_response( - "Missing provider_id parameter", 400, logger.warning + "Missing provider_id parameter", + 400, + logger.warning, ) logger.info(f"API call: /config/provider/check_one id={provider_id}") @@ -543,7 +554,7 @@ class ConfigRoute(Route): if not target: logger.warning( - f"Provider with id '{provider_id}' not found in provider_manager." + f"Provider with id '{provider_id}' not found in provider_manager.", ) return ( Response() @@ -556,7 +567,8 @@ class ConfigRoute(Route): except Exception as e: return self._error_response( - f"Critical error checking provider {provider_id}: {e}", 500 + f"Critical error checking provider {provider_id}: {e}", + 500, ) async def get_configs(self): @@ -601,6 +613,61 @@ class ConfigRoute(Route): logger.error(traceback.format_exc()) return Response().error(str(e)).__dict__ + async def get_embedding_dim(self): + """获取嵌入模型的维度""" + post_data = await request.json + provider_config = post_data.get("provider_config", None) + if not provider_config: + return Response().error("缺少参数 provider_config").__dict__ + + try: + # 动态导入 EmbeddingProvider + from astrbot.core.provider.provider import EmbeddingProvider + from astrbot.core.provider.register import provider_cls_map + + # 获取 provider 类型 + provider_type = provider_config.get("type", None) + if not provider_type: + return Response().error("provider_config 缺少 type 字段").__dict__ + + # 获取对应的 provider 类 + if provider_type not in provider_cls_map: + return ( + Response() + .error(f"未找到适用于 {provider_type} 的提供商适配器") + .__dict__ + ) + + provider_metadata = provider_cls_map[provider_type] + cls_type = provider_metadata.cls_type + + if not cls_type: + return Response().error(f"无法找到 {provider_type} 的类").__dict__ + + # 实例化 provider + inst = cls_type(provider_config, {}) + + # 检查是否是 EmbeddingProvider + if not isinstance(inst, EmbeddingProvider): + return Response().error("提供商不是 EmbeddingProvider 类型").__dict__ + + # 初始化 + if getattr(inst, "initialize", None): + await inst.initialize() + + # 获取嵌入向量维度 + vec = await inst.get_embedding("echo") + dim = len(vec) + + logger.info( + f"检测到 {provider_config.get('id', 'unknown')} 的嵌入向量维度为 {dim}", + ) + + return Response().ok({"embedding_dimensions": dim}).__dict__ + except Exception as e: + logger.error(traceback.format_exc()) + return Response().error(f"获取嵌入维度失败: {e!s}").__dict__ + async def get_platform_list(self): """获取所有平台的列表""" platform_list = [] @@ -640,7 +707,7 @@ class ConfigRoute(Route): try: save_config(self.config, self.config, is_core=True) await self.core_lifecycle.platform_manager.load_platform( - new_platform_config + new_platform_config, ) except Exception as e: return Response().error(str(e)).__dict__ @@ -652,7 +719,7 @@ class ConfigRoute(Route): try: save_config(self.config, self.config, is_core=True) await self.core_lifecycle.provider_manager.load_provider( - new_provider_config + new_provider_config, ) except Exception as e: return Response().error(str(e)).__dict__ @@ -749,9 +816,9 @@ class ConfigRoute(Route): if cache_key in self._logo_token_cache: cached_token = self._logo_token_cache[cache_key] # 确保platform_default_tmpl[platform.name]存在且为字典 - if platform.name not in platform_default_tmpl: - platform_default_tmpl[platform.name] = {} - elif not isinstance(platform_default_tmpl[platform.name], dict): + if platform.name not in platform_default_tmpl or not isinstance( + platform_default_tmpl[platform.name], dict + ): platform_default_tmpl[platform.name] = {} platform_default_tmpl[platform.name]["logo_token"] = cached_token logger.debug(f"Using cached logo token for platform {platform.name}") @@ -773,13 +840,14 @@ class ConfigRoute(Route): # 检查文件是否存在并注册令牌 if os.path.exists(logo_file_path): logo_token = await file_token_service.register_file( - logo_file_path, timeout=3600 + logo_file_path, + timeout=3600, ) # 确保platform_default_tmpl[platform.name]存在且为字典 - if platform.name not in platform_default_tmpl: - platform_default_tmpl[platform.name] = {} - elif not isinstance(platform_default_tmpl[platform.name], dict): + if platform.name not in platform_default_tmpl or not isinstance( + platform_default_tmpl[platform.name], dict + ): platform_default_tmpl[platform.name] = {} platform_default_tmpl[platform.name]["logo_token"] = logo_token @@ -790,18 +858,18 @@ class ConfigRoute(Route): logger.debug(f"Logo token registered for platform {platform.name}") else: logger.warning( - f"Platform {platform.name} logo file not found: {logo_file_path}" + f"Platform {platform.name} logo file not found: {logo_file_path}", ) except (ImportError, AttributeError) as e: logger.warning( - f"Failed to import required modules for platform {platform.name}: {e}" + f"Failed to import required modules for platform {platform.name}: {e}", ) - except (OSError, IOError) as e: + except OSError as e: logger.warning(f"File system error for platform {platform.name} logo: {e}") except Exception as e: logger.warning( - f"Unexpected error registering logo for platform {platform.name}: {e}" + f"Unexpected error registering logo for platform {platform.name}: {e}", ) async def _get_astrbot_config(self): @@ -820,7 +888,7 @@ class ConfigRoute(Route): # 收集logo注册任务 if platform.logo_path: logo_registration_tasks.append( - self._register_platform_logo(platform, platform_default_tmpl) + self._register_platform_logo(platform, platform_default_tmpl), ) # 并行执行logo注册 @@ -852,13 +920,15 @@ class ConfigRoute(Route): "description": f"{plugin_name} 配置", "type": "object", "items": plugin_md.config.schema, # 初始化时通过 __setattr__ 存入了 schema - } + }, } break return ret - async def _save_astrbot_configs(self, post_configs: dict, conf_id: str = None): + async def _save_astrbot_configs( + self, post_configs: dict, conf_id: str | None = None + ): try: if conf_id not in self.acm.confs: raise ValueError(f"配置文件 {conf_id} 不存在") diff --git a/astrbot/dashboard/routes/conversation.py b/astrbot/dashboard/routes/conversation.py index 56f892f2..d19fdf79 100644 --- a/astrbot/dashboard/routes/conversation.py +++ b/astrbot/dashboard/routes/conversation.py @@ -1,10 +1,13 @@ -import traceback import json -from .route import Route, Response, RouteContext -from astrbot.core import logger +import traceback + from quart import request -from astrbot.core.db import BaseDatabase + +from astrbot.core import logger from astrbot.core.core_lifecycle import AstrBotCoreLifecycle +from astrbot.core.db import BaseDatabase + +from .route import Response, Route, RouteContext class ConversationRoute(Route): @@ -55,12 +58,10 @@ class ConversationRoute(Route): exclude_platforms.split(",") if exclude_platforms else [] ) - if page < 1: - page = 1 + page = max(page, 1) if page_size < 1: page_size = 20 - if page_size > 100: - page_size = 100 + page_size = min(page_size, 100) try: ( @@ -76,8 +77,8 @@ class ConversationRoute(Route): exclude_platforms=exclude_platform_list, ) except Exception as e: - logger.error(f"数据库查询出错: {str(e)}\n{traceback.format_exc()}") - return Response().error(f"数据库查询出错: {str(e)}").__dict__ + logger.error(f"数据库查询出错: {e!s}\n{traceback.format_exc()}") + return Response().error(f"数据库查询出错: {e!s}").__dict__ # 计算总页数 total_pages = ( @@ -96,9 +97,9 @@ class ConversationRoute(Route): return Response().ok(result).__dict__ except Exception as e: - error_msg = f"获取对话列表失败: {str(e)}\n{traceback.format_exc()}" + error_msg = f"获取对话列表失败: {e!s}\n{traceback.format_exc()}" logger.error(error_msg) - return Response().error(f"获取对话列表失败: {str(e)}").__dict__ + return Response().error(f"获取对话列表失败: {e!s}").__dict__ async def get_conv_detail(self): """获取指定对话详情（通过POST请求）""" @@ -111,7 +112,8 @@ class ConversationRoute(Route): return Response().error("缺少必要参数: user_id 和 cid").__dict__ conversation = await self.conv_mgr.get_conversation( - unified_msg_origin=user_id, conversation_id=cid + unified_msg_origin=user_id, + conversation_id=cid, ) if not conversation: return Response().error("对话不存在").__dict__ @@ -127,14 +129,14 @@ class ConversationRoute(Route): "history": conversation.history, "created_at": conversation.created_at, "updated_at": conversation.updated_at, - } + }, ) .__dict__ ) except Exception as e: - logger.error(f"获取对话详情失败: {str(e)}\n{traceback.format_exc()}") - return Response().error(f"获取对话详情失败: {str(e)}").__dict__ + logger.error(f"获取对话详情失败: {e!s}\n{traceback.format_exc()}") + return Response().error(f"获取对话详情失败: {e!s}").__dict__ async def upd_conv(self): """更新对话信息(标题和角色ID)""" @@ -148,7 +150,8 @@ class ConversationRoute(Route): if not user_id or not cid: return Response().error("缺少必要参数: user_id 和 cid").__dict__ conversation = await self.conv_mgr.get_conversation( - unified_msg_origin=user_id, conversation_id=cid + unified_msg_origin=user_id, + conversation_id=cid, ) if not conversation: return Response().error("对话不存在").__dict__ @@ -162,8 +165,8 @@ class ConversationRoute(Route): return Response().ok({"message": "对话信息更新成功"}).__dict__ except Exception as e: - logger.error(f"更新对话信息失败: {str(e)}\n{traceback.format_exc()}") - return Response().error(f"更新对话信息失败: {str(e)}").__dict__ + logger.error(f"更新对话信息失败: {e!s}\n{traceback.format_exc()}") + return Response().error(f"更新对话信息失败: {e!s}").__dict__ async def del_conv(self): """删除对话""" @@ -188,17 +191,18 @@ class ConversationRoute(Route): if not user_id or not cid: failed_items.append( - f"user_id:{user_id}, cid:{cid} - 缺少必要参数" + f"user_id:{user_id}, cid:{cid} - 缺少必要参数", ) continue try: await self.core_lifecycle.conversation_manager.delete_conversation( - unified_msg_origin=user_id, conversation_id=cid + unified_msg_origin=user_id, + conversation_id=cid, ) deleted_count += 1 except Exception as e: - failed_items.append(f"user_id:{user_id}, cid:{cid} - {str(e)}") + failed_items.append(f"user_id:{user_id}, cid:{cid} - {e!s}") message = f"成功删除 {deleted_count} 个对话" if failed_items: @@ -212,26 +216,26 @@ class ConversationRoute(Route): "deleted_count": deleted_count, "failed_count": len(failed_items), "failed_items": failed_items, - } + }, ) .__dict__ ) - else: - # 单个删除 - user_id = data.get("user_id") - cid = data.get("cid") + # 单个删除 + user_id = data.get("user_id") + cid = data.get("cid") - if not user_id or not cid: - return Response().error("缺少必要参数: user_id 和 cid").__dict__ + if not user_id or not cid: + return Response().error("缺少必要参数: user_id 和 cid").__dict__ - await self.core_lifecycle.conversation_manager.delete_conversation( - unified_msg_origin=user_id, conversation_id=cid - ) - return Response().ok({"message": "对话删除成功"}).__dict__ + await self.core_lifecycle.conversation_manager.delete_conversation( + unified_msg_origin=user_id, + conversation_id=cid, + ) + return Response().ok({"message": "对话删除成功"}).__dict__ except Exception as e: - logger.error(f"删除对话失败: {str(e)}\n{traceback.format_exc()}") - return Response().error(f"删除对话失败: {str(e)}").__dict__ + logger.error(f"删除对话失败: {e!s}\n{traceback.format_exc()}") + return Response().error(f"删除对话失败: {e!s}").__dict__ async def update_history(self): """更新对话历史内容""" @@ -260,7 +264,8 @@ class ConversationRoute(Route): ) conversation = await self.conv_mgr.get_conversation( - unified_msg_origin=user_id, conversation_id=cid + unified_msg_origin=user_id, + conversation_id=cid, ) if not conversation: return Response().error("对话不存在").__dict__ @@ -268,11 +273,13 @@ class ConversationRoute(Route): history = json.loads(history) if isinstance(history, str) else history await self.conv_mgr.update_conversation( - unified_msg_origin=user_id, conversation_id=cid, history=history + unified_msg_origin=user_id, + conversation_id=cid, + history=history, ) return Response().ok({"message": "对话历史更新成功"}).__dict__ except Exception as e: - logger.error(f"更新对话历史失败: {str(e)}\n{traceback.format_exc()}") - return Response().error(f"更新对话历史失败: {str(e)}").__dict__ + logger.error(f"更新对话历史失败: {e!s}\n{traceback.format_exc()}") + return Response().error(f"更新对话历史失败: {e!s}").__dict__ diff --git a/astrbot/dashboard/routes/file.py b/astrbot/dashboard/routes/file.py index 8ea73d08..71d867fe 100644 --- a/astrbot/dashboard/routes/file.py +++ b/astrbot/dashboard/routes/file.py @@ -1,8 +1,10 @@ -from .route import Route, RouteContext -from astrbot import logger from quart import abort, send_file + +from astrbot import logger from astrbot.core import file_token_service +from .route import Route, RouteContext + class FileRoute(Route): def __init__( diff --git a/astrbot/dashboard/routes/knowledge_base.py b/astrbot/dashboard/routes/knowledge_base.py new file mode 100644 index 00000000..b692f23b --- /dev/null +++ b/astrbot/dashboard/routes/knowledge_base.py @@ -0,0 +1,1244 @@ +"""知识库管理 API 路由""" + +import asyncio +import os +import traceback +import uuid + +import aiofiles +from quart import request + +from astrbot.core import logger +from astrbot.core.core_lifecycle import AstrBotCoreLifecycle +from astrbot.core.provider.provider import EmbeddingProvider, RerankProvider + +from ..utils import generate_tsne_visualization +from .route import Response, Route, RouteContext + + +class KnowledgeBaseRoute(Route): + """知识库管理路由 + + 提供知识库、文档、检索、会话配置等 API 接口 + """ + + def __init__( + self, + context: RouteContext, + core_lifecycle: AstrBotCoreLifecycle, + ) -> None: + super().__init__(context) + self.core_lifecycle = core_lifecycle + self.kb_manager = None # 延迟初始化 + self.kb_db = None + self.session_config_db = None # 会话配置数据库 + self.retrieval_manager = None + self.upload_progress = {} # 存储上传进度 {task_id: {status, file_index, file_total, stage, current, total}} + self.upload_tasks = {} # 存储后台上传任务 {task_id: {"status", "result", "error"}} + + # 注册路由 + self.routes = { + # 知识库管理 + "/kb/list": ("GET", self.list_kbs), + "/kb/create": ("POST", self.create_kb), + "/kb/get": ("GET", self.get_kb), + "/kb/update": ("POST", self.update_kb), + "/kb/delete": ("POST", self.delete_kb), + "/kb/stats": ("GET", self.get_kb_stats), + # 文档管理 + "/kb/document/list": ("GET", self.list_documents), + "/kb/document/upload": ("POST", self.upload_document), + "/kb/document/upload/url": ("POST", self.upload_document_from_url), + "/kb/document/upload/progress": ("GET", self.get_upload_progress), + "/kb/document/get": ("GET", self.get_document), + "/kb/document/delete": ("POST", self.delete_document), + # # 块管理 + "/kb/chunk/list": ("GET", self.list_chunks), + "/kb/chunk/delete": ("POST", self.delete_chunk), + # # 多媒体管理 + # "/kb/media/list": ("GET", self.list_media), + # "/kb/media/delete": ("POST", self.delete_media), + # 检索 + "/kb/retrieve": ("POST", self.retrieve), + # 会话知识库配置 + "/kb/session/config/get": ("GET", self.get_session_kb_config), + "/kb/session/config/set": ("POST", self.set_session_kb_config), + "/kb/session/config/delete": ("POST", self.delete_session_kb_config), + } + self.register_routes() + + def _get_kb_manager(self): + return self.core_lifecycle.kb_manager + + async def _background_upload_task( + self, + task_id: str, + kb_helper, + files_to_upload: list, + chunk_size: int, + chunk_overlap: int, + batch_size: int, + tasks_limit: int, + max_retries: int, + ): + """后台上传任务""" + try: + # 初始化任务状态 + self.upload_tasks[task_id] = { + "status": "processing", + "result": None, + "error": None, + } + self.upload_progress[task_id] = { + "status": "processing", + "file_index": 0, + "file_total": len(files_to_upload), + "stage": "waiting", + "current": 0, + "total": 100, + } + + uploaded_docs = [] + failed_docs = [] + + for file_idx, file_info in enumerate(files_to_upload): + try: + # 更新整体进度 + self.upload_progress[task_id].update( + { + "status": "processing", + "file_index": file_idx, + "file_name": file_info["file_name"], + "stage": "parsing", + "current": 0, + "total": 100, + }, + ) + + # 创建进度回调函数 + async def progress_callback(stage, current, total): + if task_id in self.upload_progress: + self.upload_progress[task_id].update( + { + "status": "processing", + "file_index": file_idx, + "file_name": file_info["file_name"], + "stage": stage, + "current": current, + "total": total, + }, + ) + + doc = await kb_helper.upload_document( + file_name=file_info["file_name"], + file_content=file_info["file_content"], + file_type=file_info["file_type"], + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + progress_callback=progress_callback, + ) + + uploaded_docs.append(doc.model_dump()) + except Exception as e: + logger.error(f"上传文档 {file_info['file_name']} 失败: {e}") + failed_docs.append( + {"file_name": file_info["file_name"], "error": str(e)}, + ) + + # 更新任务完成状态 + result = { + "task_id": task_id, + "uploaded": uploaded_docs, + "failed": failed_docs, + "total": len(files_to_upload), + "success_count": len(uploaded_docs), + "failed_count": len(failed_docs), + } + + self.upload_tasks[task_id] = { + "status": "completed", + "result": result, + "error": None, + } + self.upload_progress[task_id]["status"] = "completed" + + except Exception as e: + logger.error(f"后台上传任务 {task_id} 失败: {e}") + logger.error(traceback.format_exc()) + self.upload_tasks[task_id] = { + "status": "failed", + "result": None, + "error": str(e), + } + if task_id in self.upload_progress: + self.upload_progress[task_id]["status"] = "failed" + + async def list_kbs(self): + """获取知识库列表 + + Query 参数: + - page: 页码 (默认 1) + - page_size: 每页数量 (默认 20) + - refresh_stats: 是否刷新统计信息 (默认 false，首次加载时可设为 true) + """ + try: + kb_manager = self._get_kb_manager() + page = request.args.get("page", 1, type=int) + page_size = request.args.get("page_size", 20, type=int) + + kbs = await kb_manager.list_kbs() + + # 转换为字典列表 + kb_list = [] + for kb in kbs: + kb_list.append(kb.model_dump()) + + return ( + Response() + .ok({"items": kb_list, "page": page, "page_size": page_size}) + .__dict__ + ) + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"获取知识库列表失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"获取知识库列表失败: {e!s}").__dict__ + + async def create_kb(self): + """创建知识库 + + Body: + - kb_name: 知识库名称 (必填) + - description: 描述 (可选) + - emoji: 图标 (可选) + - embedding_provider_id: 嵌入模型提供商ID (可选) + - rerank_provider_id: 重排序模型提供商ID (可选) + - chunk_size: 分块大小 (可选, 默认512) + - chunk_overlap: 块重叠大小 (可选, 默认50) + - top_k_dense: 密集检索数量 (可选, 默认50) + - top_k_sparse: 稀疏检索数量 (可选, 默认50) + - top_m_final: 最终返回数量 (可选, 默认5) + """ + try: + kb_manager = self._get_kb_manager() + data = await request.json + kb_name = data.get("kb_name") + if not kb_name: + return Response().error("知识库名称不能为空").__dict__ + + description = data.get("description") + emoji = data.get("emoji") + embedding_provider_id = data.get("embedding_provider_id") + rerank_provider_id = data.get("rerank_provider_id") + chunk_size = data.get("chunk_size") + chunk_overlap = data.get("chunk_overlap") + top_k_dense = data.get("top_k_dense") + top_k_sparse = data.get("top_k_sparse") + top_m_final = data.get("top_m_final") + + # pre-check embedding dim + if not embedding_provider_id: + return Response().error("缺少参数 embedding_provider_id").__dict__ + prv = await kb_manager.provider_manager.get_provider_by_id( + embedding_provider_id, + ) # type: ignore + if not prv or not isinstance(prv, EmbeddingProvider): + return ( + Response().error(f"嵌入模型不存在或类型错误({type(prv)})").__dict__ + ) + try: + vec = await prv.get_embedding("astrbot") + if len(vec) != prv.get_dim(): + raise ValueError( + f"嵌入向量维度不匹配，实际是 {len(vec)}，然而配置是 {prv.get_dim()}", + ) + except Exception as e: + return Response().error(f"测试嵌入模型失败: {e!s}").__dict__ + # pre-check rerank + if rerank_provider_id: + rerank_prv: RerankProvider = ( + await kb_manager.provider_manager.get_provider_by_id( + rerank_provider_id, + ) + ) # type: ignore + if not rerank_prv: + return Response().error("重排序模型不存在").__dict__ + # 检查重排序模型可用性 + try: + res = await rerank_prv.rerank( + query="astrbot", + documents=["astrbot knowledge base"], + ) + if not res: + raise ValueError("重排序模型返回结果异常") + except Exception as e: + return ( + Response() + .error(f"测试重排序模型失败: {e!s}，请检查控制台日志输出。") + .__dict__ + ) + + kb_helper = await kb_manager.create_kb( + kb_name=kb_name, + description=description, + emoji=emoji, + embedding_provider_id=embedding_provider_id, + rerank_provider_id=rerank_provider_id, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + top_k_dense=top_k_dense, + top_k_sparse=top_k_sparse, + top_m_final=top_m_final, + ) + kb = kb_helper.kb + + return Response().ok(kb.model_dump(), "创建知识库成功").__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"创建知识库失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"创建知识库失败: {e!s}").__dict__ + + async def get_kb(self): + """获取知识库详情 + + Query 参数: + - kb_id: 知识库 ID (必填) + """ + try: + kb_manager = self._get_kb_manager() + kb_id = request.args.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + + kb_helper = await kb_manager.get_kb(kb_id) + if not kb_helper: + return Response().error("知识库不存在").__dict__ + kb = kb_helper.kb + + return Response().ok(kb.model_dump()).__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"获取知识库详情失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"获取知识库详情失败: {e!s}").__dict__ + + async def update_kb(self): + """更新知识库 + + Body: + - kb_id: 知识库 ID (必填) + - kb_name: 新的知识库名称 (可选) + - description: 新的描述 (可选) + - emoji: 新的图标 (可选) + - embedding_provider_id: 新的嵌入模型提供商ID (可选) + - rerank_provider_id: 新的重排序模型提供商ID (可选) + - chunk_size: 分块大小 (可选) + - chunk_overlap: 块重叠大小 (可选) + - top_k_dense: 密集检索数量 (可选) + - top_k_sparse: 稀疏检索数量 (可选) + - top_m_final: 最终返回数量 (可选) + """ + try: + kb_manager = self._get_kb_manager() + data = await request.json + + kb_id = data.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + + kb_name = data.get("kb_name") + description = data.get("description") + emoji = data.get("emoji") + embedding_provider_id = data.get("embedding_provider_id") + rerank_provider_id = data.get("rerank_provider_id") + chunk_size = data.get("chunk_size") + chunk_overlap = data.get("chunk_overlap") + top_k_dense = data.get("top_k_dense") + top_k_sparse = data.get("top_k_sparse") + top_m_final = data.get("top_m_final") + + # 检查是否至少提供了一个更新字段 + if all( + v is None + for v in [ + kb_name, + description, + emoji, + embedding_provider_id, + rerank_provider_id, + chunk_size, + chunk_overlap, + top_k_dense, + top_k_sparse, + top_m_final, + ] + ): + return Response().error("至少需要提供一个更新字段").__dict__ + + kb_helper = await kb_manager.update_kb( + kb_id=kb_id, + kb_name=kb_name, + description=description, + emoji=emoji, + embedding_provider_id=embedding_provider_id, + rerank_provider_id=rerank_provider_id, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + top_k_dense=top_k_dense, + top_k_sparse=top_k_sparse, + top_m_final=top_m_final, + ) + + if not kb_helper: + return Response().error("知识库不存在").__dict__ + + kb = kb_helper.kb + return Response().ok(kb.model_dump(), "更新知识库成功").__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"更新知识库失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"更新知识库失败: {e!s}").__dict__ + + async def delete_kb(self): + """删除知识库 + + Body: + - kb_id: 知识库 ID (必填) + """ + try: + kb_manager = self._get_kb_manager() + data = await request.json + + kb_id = data.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + + success = await kb_manager.delete_kb(kb_id) + if not success: + return Response().error("知识库不存在").__dict__ + + return Response().ok(message="删除知识库成功").__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"删除知识库失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"删除知识库失败: {e!s}").__dict__ + + async def get_kb_stats(self): + """获取知识库统计信息 + + Query 参数: + - kb_id: 知识库 ID (必填) + """ + try: + kb_manager = self._get_kb_manager() + kb_id = request.args.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + + kb_helper = await kb_manager.get_kb(kb_id) + if not kb_helper: + return Response().error("知识库不存在").__dict__ + kb = kb_helper.kb + + stats = { + "kb_id": kb.kb_id, + "kb_name": kb.kb_name, + "doc_count": kb.doc_count, + "chunk_count": kb.chunk_count, + "created_at": kb.created_at.isoformat(), + "updated_at": kb.updated_at.isoformat(), + } + + return Response().ok(stats).__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"获取知识库统计失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"获取知识库统计失败: {e!s}").__dict__ + + # ===== 文档管理 API ===== + + async def list_documents(self): + """获取文档列表 + + Query 参数: + - kb_id: 知识库 ID (必填) + - page: 页码 (默认 1) + - page_size: 每页数量 (默认 20) + """ + try: + kb_manager = self._get_kb_manager() + kb_id = request.args.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + kb_helper = await kb_manager.get_kb(kb_id) + if not kb_helper: + return Response().error("知识库不存在").__dict__ + + page = request.args.get("page", 1, type=int) + page_size = request.args.get("page_size", 100, type=int) + + offset = (page - 1) * page_size + limit = page_size + + doc_list = await kb_helper.list_documents(offset=offset, limit=limit) + + doc_list = [doc.model_dump() for doc in doc_list] + + return ( + Response() + .ok({"items": doc_list, "page": page, "page_size": page_size}) + .__dict__ + ) + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"获取文档列表失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"获取文档列表失败: {e!s}").__dict__ + + async def upload_document(self): + """上传文档 + + 支持两种方式: + 1. multipart/form-data 文件上传（支持多文件，最多10个） + 2. JSON 格式 base64 编码上传（支持多文件，最多10个） + + Form Data (multipart/form-data): + - kb_id: 知识库 ID (必填) + - file: 文件对象 (必填，可多个，字段名为 file, file1, file2, ... 或 files[]) + + JSON Body (application/json): + - kb_id: 知识库 ID (必填) + - files: 文件数组 (必填) + - file_name: 文件名 (必填) + - file_content: base64 编码的文件内容 (必填) + + 返回: + - task_id: 任务ID，用于查询上传进度和结果 + """ + try: + kb_manager = self._get_kb_manager() + + # 检查 Content-Type + content_type = request.content_type + kb_id = None + chunk_size = None + chunk_overlap = None + batch_size = 32 + tasks_limit = 3 + max_retries = 3 + files_to_upload = [] # 存储待上传的文件信息列表 + + if content_type and "multipart/form-data" not in content_type: + return ( + Response().error("Content-Type 须为 multipart/form-data").__dict__ + ) + form_data = await request.form + files = await request.files + + kb_id = form_data.get("kb_id") + chunk_size = int(form_data.get("chunk_size", 512)) + chunk_overlap = int(form_data.get("chunk_overlap", 50)) + batch_size = int(form_data.get("batch_size", 32)) + tasks_limit = int(form_data.get("tasks_limit", 3)) + max_retries = int(form_data.get("max_retries", 3)) + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + + # 收集所有文件 + file_list = [] + # 支持 file, file1, file2, ... 或 files[] 格式 + for key in files.keys(): + if key == "file" or key.startswith("file") or key == "files[]": + file_items = files.getlist(key) + file_list.extend(file_items) + + if not file_list: + return Response().error("缺少文件").__dict__ + + # 限制文件数量 + if len(file_list) > 10: + return Response().error("最多只能上传10个文件").__dict__ + + # 处理每个文件 + for file in file_list: + file_name = file.filename + + # 保存到临时文件 + temp_file_path = f"data/temp/{uuid.uuid4()}_{file_name}" + await file.save(temp_file_path) + + try: + # 异步读取文件内容 + async with aiofiles.open(temp_file_path, "rb") as f: + file_content = await f.read() + + # 提取文件类型 + file_type = ( + file_name.rsplit(".", 1)[-1].lower() if "." in file_name else "" + ) + + files_to_upload.append( + { + "file_name": file_name, + "file_content": file_content, + "file_type": file_type, + }, + ) + finally: + # 清理临时文件 + if os.path.exists(temp_file_path): + os.remove(temp_file_path) + + # 获取知识库 + kb_helper = await kb_manager.get_kb(kb_id) + if not kb_helper: + return Response().error("知识库不存在").__dict__ + + # 生成任务ID + task_id = str(uuid.uuid4()) + + # 初始化任务状态 + self.upload_tasks[task_id] = { + "status": "pending", + "result": None, + "error": None, + } + + # 启动后台任务 + asyncio.create_task( + self._background_upload_task( + task_id=task_id, + kb_helper=kb_helper, + files_to_upload=files_to_upload, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + ), + ) + + return ( + Response() + .ok( + { + "task_id": task_id, + "file_count": len(files_to_upload), + "message": "task created, processing in background", + }, + ) + .__dict__ + ) + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"上传文档失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"上传文档失败: {e!s}").__dict__ + + async def get_upload_progress(self): + """获取上传进度和结果 + + Query 参数: + - task_id: 任务 ID (必填) + + 返回状态: + - pending: 任务待处理 + - processing: 任务处理中 + - completed: 任务完成 + - failed: 任务失败 + """ + try: + task_id = request.args.get("task_id") + if not task_id: + return Response().error("缺少参数 task_id").__dict__ + + # 检查任务是否存在 + if task_id not in self.upload_tasks: + return Response().error("找不到该任务").__dict__ + + task_info = self.upload_tasks[task_id] + status = task_info["status"] + + # 构建返回数据 + response_data = { + "task_id": task_id, + "status": status, + } + + # 如果任务正在处理，返回进度信息 + if status == "processing" and task_id in self.upload_progress: + response_data["progress"] = self.upload_progress[task_id] + + # 如果任务完成，返回结果 + if status == "completed": + response_data["result"] = task_info["result"] + # 清理已完成的任务 + # del self.upload_tasks[task_id] + # if task_id in self.upload_progress: + # del self.upload_progress[task_id] + + # 如果任务失败，返回错误信息 + if status == "failed": + response_data["error"] = task_info["error"] + + return Response().ok(response_data).__dict__ + + except Exception as e: + logger.error(f"获取上传进度失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"获取上传进度失败: {e!s}").__dict__ + + async def get_document(self): + """获取文档详情 + + Query 参数: + - doc_id: 文档 ID (必填) + """ + try: + kb_manager = self._get_kb_manager() + kb_id = request.args.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + doc_id = request.args.get("doc_id") + if not doc_id: + return Response().error("缺少参数 doc_id").__dict__ + kb_helper = await kb_manager.get_kb(kb_id) + if not kb_helper: + return Response().error("知识库不存在").__dict__ + + doc = await kb_helper.get_document(doc_id) + if not doc: + return Response().error("文档不存在").__dict__ + + return Response().ok(doc.model_dump()).__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"获取文档详情失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"获取文档详情失败: {e!s}").__dict__ + + async def delete_document(self): + """删除文档 + + Body: + - kb_id: 知识库 ID (必填) + - doc_id: 文档 ID (必填) + """ + try: + kb_manager = self._get_kb_manager() + data = await request.json + + kb_id = data.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + doc_id = data.get("doc_id") + if not doc_id: + return Response().error("缺少参数 doc_id").__dict__ + + kb_helper = await kb_manager.get_kb(kb_id) + if not kb_helper: + return Response().error("知识库不存在").__dict__ + + await kb_helper.delete_document(doc_id) + return Response().ok(message="删除文档成功").__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"删除文档失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"删除文档失败: {e!s}").__dict__ + + async def delete_chunk(self): + """删除文本块 + + Body: + - kb_id: 知识库 ID (必填) + - chunk_id: 块 ID (必填) + """ + try: + kb_manager = self._get_kb_manager() + data = await request.json + + kb_id = data.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + chunk_id = data.get("chunk_id") + if not chunk_id: + return Response().error("缺少参数 chunk_id").__dict__ + doc_id = data.get("doc_id") + if not doc_id: + return Response().error("缺少参数 doc_id").__dict__ + + kb_helper = await kb_manager.get_kb(kb_id) + if not kb_helper: + return Response().error("知识库不存在").__dict__ + + await kb_helper.delete_chunk(chunk_id, doc_id) + return Response().ok(message="删除文本块成功").__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"删除文本块失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"删除文本块失败: {e!s}").__dict__ + + async def list_chunks(self): + """获取块列表 + + Query 参数: + - kb_id: 知识库 ID (必填) + - page: 页码 (默认 1) + - page_size: 每页数量 (默认 20) + """ + try: + kb_manager = self._get_kb_manager() + kb_id = request.args.get("kb_id") + doc_id = request.args.get("doc_id") + page = request.args.get("page", 1, type=int) + page_size = request.args.get("page_size", 100, type=int) + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + if not doc_id: + return Response().error("缺少参数 doc_id").__dict__ + kb_helper = await kb_manager.get_kb(kb_id) + offset = (page - 1) * page_size + limit = page_size + if not kb_helper: + return Response().error("知识库不存在").__dict__ + chunk_list = await kb_helper.get_chunks_by_doc_id( + doc_id=doc_id, + offset=offset, + limit=limit, + ) + return ( + Response() + .ok( + data={ + "items": chunk_list, + "page": page, + "page_size": page_size, + "total": await kb_helper.get_chunk_count_by_doc_id(doc_id), + }, + ) + .__dict__ + ) + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"获取块列表失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"获取块列表失败: {e!s}").__dict__ + + # ===== 检索 API ===== + + async def retrieve(self): + """检索知识库 + + Body: + - query: 查询文本 (必填) + - kb_ids: 知识库 ID 列表 (必填) + - top_k: 返回结果数量 (可选, 默认 5) + - debug: 是否启用调试模式，返回 t-SNE 可视化图片 (可选, 默认 False) + """ + try: + kb_manager = self._get_kb_manager() + data = await request.json + + query = data.get("query") + kb_names = data.get("kb_names") + debug = data.get("debug", False) + + if not query: + return Response().error("缺少参数 query").__dict__ + if not kb_names or not isinstance(kb_names, list): + return Response().error("缺少参数 kb_names 或格式错误").__dict__ + + top_k = data.get("top_k", 5) + + results = await kb_manager.retrieve( + query=query, + kb_names=kb_names, + top_m_final=top_k, + ) + result_list = [] + if results: + result_list = results["results"] + + response_data = { + "results": result_list, + "total": len(result_list), + "query": query, + } + + # Debug 模式：生成 t-SNE 可视化 + if debug: + try: + img_base64 = await generate_tsne_visualization( + query, + kb_names, + kb_manager, + ) + if img_base64: + response_data["visualization"] = img_base64 + except Exception as e: + logger.error(f"生成 t-SNE 可视化失败: {e}") + logger.error(traceback.format_exc()) + response_data["visualization_error"] = str(e) + + return Response().ok(response_data).__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"检索失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"检索失败: {e!s}").__dict__ + + # ===== 会话知识库配置 API ===== + + async def get_session_kb_config(self): + """获取会话的知识库配置 + + Query 参数: + - session_id: 会话 ID (必填) + + 返回: + - kb_ids: 知识库 ID 列表 + - top_k: 返回结果数量 + - enable_rerank: 是否启用重排序 + """ + try: + from astrbot.core import sp + + session_id = request.args.get("session_id") + + if not session_id: + return Response().error("缺少参数 session_id").__dict__ + + # 从 SharedPreferences 获取配置 + config = await sp.session_get(session_id, "kb_config", default={}) + + logger.debug(f"[KB配置] 读取到配置: session_id={session_id}") + + # 如果没有配置，返回默认值 + if not config: + config = {"kb_ids": [], "top_k": 5, "enable_rerank": True} + + return Response().ok(config).__dict__ + + except Exception as e: + logger.error(f"[KB配置] 获取配置时出错: {e}", exc_info=True) + return Response().error(f"获取会话知识库配置失败: {e!s}").__dict__ + + async def set_session_kb_config(self): + """设置会话的知识库配置 + + Body: + - scope: 配置范围 (目前只支持 "session") + - scope_id: 会话 ID (必填) + - kb_ids: 知识库 ID 列表 (必填) + - top_k: 返回结果数量 (可选, 默认 5) + - enable_rerank: 是否启用重排序 (可选, 默认 true) + """ + try: + from astrbot.core import sp + + data = await request.json + + scope = data.get("scope") + scope_id = data.get("scope_id") + kb_ids = data.get("kb_ids", []) + top_k = data.get("top_k", 5) + enable_rerank = data.get("enable_rerank", True) + + # 验证参数 + if scope != "session": + return Response().error("目前仅支持 session 范围的配置").__dict__ + + if not scope_id: + return Response().error("缺少参数 scope_id").__dict__ + + if not isinstance(kb_ids, list): + return Response().error("kb_ids 必须是列表").__dict__ + + # 验证知识库是否存在 + kb_mgr = self._get_kb_manager() + invalid_ids = [] + valid_ids = [] + for kb_id in kb_ids: + kb_helper = await kb_mgr.get_kb(kb_id) + if kb_helper: + valid_ids.append(kb_id) + else: + invalid_ids.append(kb_id) + logger.warning(f"[KB配置] 知识库不存在: {kb_id}") + + if invalid_ids: + logger.warning(f"[KB配置] 以下知识库ID无效: {invalid_ids}") + + # 允许保存空列表，表示明确不使用任何知识库 + if kb_ids and not valid_ids: + # 只有当用户提供了 kb_ids 但全部无效时才报错 + return Response().error(f"所有提供的知识库ID都无效: {kb_ids}").__dict__ + + # 如果 kb_ids 为空列表，表示用户想清空配置 + if not kb_ids: + valid_ids = [] + + # 构建配置对象（只保存有效的ID） + config = { + "kb_ids": valid_ids, + "top_k": top_k, + "enable_rerank": enable_rerank, + } + + # 保存到 SharedPreferences + await sp.session_put(scope_id, "kb_config", config) + + # 立即验证是否保存成功 + verify_config = await sp.session_get(scope_id, "kb_config", default={}) + + if verify_config == config: + return ( + Response() + .ok( + {"valid_ids": valid_ids, "invalid_ids": invalid_ids}, + "保存知识库配置成功", + ) + .__dict__ + ) + logger.error("[KB配置] 配置保存失败，验证不匹配") + return Response().error("配置保存失败").__dict__ + + except Exception as e: + logger.error(f"[KB配置] 设置配置时出错: {e}", exc_info=True) + return Response().error(f"设置会话知识库配置失败: {e!s}").__dict__ + + async def delete_session_kb_config(self): + """删除会话的知识库配置 + + Body: + - scope: 配置范围 (目前只支持 "session") + - scope_id: 会话 ID (必填) + """ + try: + from astrbot.core import sp + + data = await request.json + + scope = data.get("scope") + scope_id = data.get("scope_id") + + # 验证参数 + if scope != "session": + return Response().error("目前仅支持 session 范围的配置").__dict__ + + if not scope_id: + return Response().error("缺少参数 scope_id").__dict__ + + # 从 SharedPreferences 删除配置 + await sp.session_remove(scope_id, "kb_config") + + return Response().ok(message="删除知识库配置成功").__dict__ + + except Exception as e: + logger.error(f"删除会话知识库配置失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"删除会话知识库配置失败: {e!s}").__dict__ + + async def upload_document_from_url(self): + """从 URL 上传文档 + + Body: + - kb_id: 知识库 ID (必填) + - url: 要提取内容的网页 URL (必填) + - chunk_size: 分块大小 (可选, 默认512) + - chunk_overlap: 块重叠大小 (可选, 默认50) + - batch_size: 批处理大小 (可选, 默认32) + - tasks_limit: 并发任务限制 (可选, 默认3) + - max_retries: 最大重试次数 (可选, 默认3) + + 返回: + - task_id: 任务ID，用于查询上传进度和结果 + """ + try: + kb_manager = self._get_kb_manager() + data = await request.json + + kb_id = data.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + + url = data.get("url") + if not url: + return Response().error("缺少参数 url").__dict__ + + chunk_size = data.get("chunk_size", 512) + chunk_overlap = data.get("chunk_overlap", 50) + batch_size = data.get("batch_size", 32) + tasks_limit = data.get("tasks_limit", 3) + max_retries = data.get("max_retries", 3) + enable_cleaning = data.get("enable_cleaning", False) + cleaning_provider_id = data.get("cleaning_provider_id") + + # 获取知识库 + kb_helper = await kb_manager.get_kb(kb_id) + if not kb_helper: + return Response().error("知识库不存在").__dict__ + + # 生成任务ID + task_id = str(uuid.uuid4()) + + # 初始化任务状态 + self.upload_tasks[task_id] = { + "status": "pending", + "result": None, + "error": None, + } + + # 启动后台任务 + asyncio.create_task( + self._background_upload_from_url_task( + task_id=task_id, + kb_helper=kb_helper, + url=url, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + enable_cleaning=enable_cleaning, + cleaning_provider_id=cleaning_provider_id, + ), + ) + + return ( + Response() + .ok( + { + "task_id": task_id, + "url": url, + "message": "URL upload task created, processing in background", + }, + ) + .__dict__ + ) + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"从URL上传文档失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"从URL上传文档失败: {e!s}").__dict__ + + async def _background_upload_from_url_task( + self, + task_id: str, + kb_helper, + url: str, + chunk_size: int, + chunk_overlap: int, + batch_size: int, + tasks_limit: int, + max_retries: int, + enable_cleaning: bool, + cleaning_provider_id: str | None, + ): + """后台上传URL任务""" + try: + # 初始化任务状态 + self.upload_tasks[task_id] = { + "status": "processing", + "result": None, + "error": None, + } + self.upload_progress[task_id] = { + "status": "processing", + "file_index": 0, + "file_total": 1, + "file_name": f"URL: {url}", + "stage": "extracting", + "current": 0, + "total": 100, + } + + # 创建进度回调函数 + async def progress_callback(stage, current, total): + if task_id in self.upload_progress: + self.upload_progress[task_id].update( + { + "status": "processing", + "file_index": 0, + "file_name": f"URL: {url}", + "stage": stage, + "current": current, + "total": total, + }, + ) + + # 上传文档 + doc = await kb_helper.upload_from_url( + url=url, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + progress_callback=progress_callback, + enable_cleaning=enable_cleaning, + cleaning_provider_id=cleaning_provider_id, + ) + + # 更新任务完成状态 + result = { + "task_id": task_id, + "uploaded": [doc.model_dump()], + "failed": [], + "total": 1, + "success_count": 1, + "failed_count": 0, + } + + self.upload_tasks[task_id] = { + "status": "completed", + "result": result, + "error": None, + } + self.upload_progress[task_id]["status"] = "completed" + + except Exception as e: + logger.error(f"后台上传URL任务 {task_id} 失败: {e}") + logger.error(traceback.format_exc()) + self.upload_tasks[task_id] = { + "status": "failed", + "result": None, + "error": str(e), + } + if task_id in self.upload_progress: + self.upload_progress[task_id]["status"] = "failed" diff --git a/astrbot/dashboard/routes/log.py b/astrbot/dashboard/routes/log.py index e47f9d77..eb02fdf4 100644 --- a/astrbot/dashboard/routes/log.py +++ b/astrbot/dashboard/routes/log.py @@ -1,8 +1,11 @@ import asyncio import json + from quart import make_response -from astrbot.core import logger, LogBroker -from .route import Route, RouteContext, Response + +from astrbot.core import LogBroker, logger + +from .route import Response, Route, RouteContext class LogRoute(Route): @@ -11,7 +14,9 @@ class LogRoute(Route): self.log_broker = log_broker self.app.add_url_rule("/api/live-log", view_func=self.log, methods=["GET"]) self.app.add_url_rule( - "/api/log-history", view_func=self.log_history, methods=["GET"] + "/api/log-history", + view_func=self.log_history, + methods=["GET"], ) async def log(self): @@ -55,7 +60,7 @@ class LogRoute(Route): .ok( data={ "logs": logs, - } + }, ) .__dict__ ) diff --git a/astrbot/dashboard/routes/persona.py b/astrbot/dashboard/routes/persona.py index 032471ee..7ddb75f1 100644 --- a/astrbot/dashboard/routes/persona.py +++ b/astrbot/dashboard/routes/persona.py @@ -1,9 +1,12 @@ import traceback -from .route import Route, Response, RouteContext -from astrbot.core import logger + from quart import request -from astrbot.core.db import BaseDatabase + +from astrbot.core import logger from astrbot.core.core_lifecycle import AstrBotCoreLifecycle +from astrbot.core.db import BaseDatabase + +from .route import Response, Route, RouteContext class PersonaRoute(Route): @@ -46,13 +49,13 @@ class PersonaRoute(Route): else None, } for persona in personas - ] + ], ) .__dict__ ) except Exception as e: - logger.error(f"获取人格列表失败: {str(e)}\n{traceback.format_exc()}") - return Response().error(f"获取人格列表失败: {str(e)}").__dict__ + logger.error(f"获取人格列表失败: {e!s}\n{traceback.format_exc()}") + return Response().error(f"获取人格列表失败: {e!s}").__dict__ async def get_persona_detail(self): """获取指定人格的详细信息""" @@ -81,13 +84,13 @@ class PersonaRoute(Route): "updated_at": persona.updated_at.isoformat() if persona.updated_at else None, - } + }, ) .__dict__ ) except Exception as e: - logger.error(f"获取人格详情失败: {str(e)}\n{traceback.format_exc()}") - return Response().error(f"获取人格详情失败: {str(e)}").__dict__ + logger.error(f"获取人格详情失败: {e!s}\n{traceback.format_exc()}") + return Response().error(f"获取人格详情失败: {e!s}").__dict__ async def create_persona(self): """创建新人格""" @@ -136,15 +139,15 @@ class PersonaRoute(Route): if persona.updated_at else None, }, - } + }, ) .__dict__ ) except ValueError as e: return Response().error(str(e)).__dict__ except Exception as e: - logger.error(f"创建人格失败: {str(e)}\n{traceback.format_exc()}") - return Response().error(f"创建人格失败: {str(e)}").__dict__ + logger.error(f"创建人格失败: {e!s}\n{traceback.format_exc()}") + return Response().error(f"创建人格失败: {e!s}").__dict__ async def update_persona(self): """更新人格信息""" @@ -177,8 +180,8 @@ class PersonaRoute(Route): except ValueError as e: return Response().error(str(e)).__dict__ except Exception as e: - logger.error(f"更新人格失败: {str(e)}\n{traceback.format_exc()}") - return Response().error(f"更新人格失败: {str(e)}").__dict__ + logger.error(f"更新人格失败: {e!s}\n{traceback.format_exc()}") + return Response().error(f"更新人格失败: {e!s}").__dict__ async def delete_persona(self): """删除人格""" @@ -195,5 +198,5 @@ class PersonaRoute(Route): except ValueError as e: return Response().error(str(e)).__dict__ except Exception as e: - logger.error(f"删除人格失败: {str(e)}\n{traceback.format_exc()}") - return Response().error(f"删除人格失败: {str(e)}").__dict__ + logger.error(f"删除人格失败: {e!s}\n{traceback.format_exc()}") + return Response().error(f"删除人格失败: {e!s}").__dict__ diff --git a/astrbot/dashboard/routes/plugin.py b/astrbot/dashboard/routes/plugin.py index 84933969..597a245d 100644 --- a/astrbot/dashboard/routes/plugin.py +++ b/astrbot/dashboard/routes/plugin.py @@ -1,24 +1,23 @@ -import traceback -import aiohttp -import os import json +import os +import ssl +import traceback from datetime import datetime -import ssl +import aiohttp import certifi - -from .route import Route, Response, RouteContext -from astrbot.core import logger from quart import request -from astrbot.core.star.star_manager import PluginManager + +from astrbot.core import DEMO_MODE, file_token_service, logger from astrbot.core.core_lifecycle import AstrBotCoreLifecycle -from astrbot.core.star.star_handler import star_handlers_registry from astrbot.core.star.filter.command import CommandFilter from astrbot.core.star.filter.command_group import CommandGroupFilter from astrbot.core.star.filter.permission import PermissionTypeFilter from astrbot.core.star.filter.regex import RegexFilter -from astrbot.core.star.star_handler import EventType -from astrbot.core import DEMO_MODE +from astrbot.core.star.star_handler import EventType, star_handlers_registry +from astrbot.core.star.star_manager import PluginManager + +from .route import Response, Route, RouteContext class PluginRoute(Route): @@ -54,6 +53,8 @@ class PluginRoute(Route): EventType.OnAfterMessageSentEvent: "发送消息后", } + self._logo_cache = {} + async def reload_plugins(self): if DEMO_MODE: return ( @@ -104,29 +105,33 @@ class PluginRoute(Route): for url in urls: try: - async with aiohttp.ClientSession( - trust_env=True, connector=connector - ) as session: - async with session.get(url) as response: - if response.status == 200: - remote_data = await response.json() + async with ( + aiohttp.ClientSession( + trust_env=True, + connector=connector, + ) as session, + session.get(url) as response, + ): + if response.status == 200: + remote_data = await response.json() - # 检查远程数据是否为空 - if not remote_data or ( - isinstance(remote_data, dict) and len(remote_data) == 0 - ): - logger.warning(f"远程插件市场数据为空: {url}") - continue # 继续尝试其他URL或使用缓存 + # 检查远程数据是否为空 + if not remote_data or ( + isinstance(remote_data, dict) and len(remote_data) == 0 + ): + logger.warning(f"远程插件市场数据为空: {url}") + continue # 继续尝试其他URL或使用缓存 - logger.info("成功获取远程插件市场数据") - # 获取最新的MD5并保存到缓存 - current_md5 = await self._get_remote_md5() - self._save_plugin_cache( - cache_file, remote_data, current_md5 - ) - return Response().ok(remote_data).__dict__ - else: - logger.error(f"请求 {url} 失败，状态码：{response.status}") + logger.info("成功获取远程插件市场数据") + # 获取最新的MD5并保存到缓存 + current_md5 = await self._get_remote_md5() + self._save_plugin_cache( + cache_file, + remote_data, + current_md5, + ) + return Response().ok(remote_data).__dict__ + logger.error(f"请求 {url} 失败，状态码：{response.status}") except Exception as e: logger.error(f"请求 {url} 失败，错误：{e}") @@ -147,7 +152,7 @@ class PluginRoute(Route): return False # 加载缓存文件 - with open(cache_file, "r", encoding="utf-8") as f: + with open(cache_file, encoding="utf-8") as f: cache_data = json.load(f) cached_md5 = cache_data.get("md5") @@ -163,7 +168,7 @@ class PluginRoute(Route): is_valid = cached_md5 == remote_md5 logger.debug( - f"插件数据MD5: 本地={cached_md5}, 远程={remote_md5}, 有效={is_valid}" + f"插件数据MD5: 本地={cached_md5}, 远程={remote_md5}, 有效={is_valid}", ) return is_valid @@ -177,18 +182,20 @@ class PluginRoute(Route): ssl_context = ssl.create_default_context(cafile=certifi.where()) connector = aiohttp.TCPConnector(ssl=ssl_context) - async with aiohttp.ClientSession( - trust_env=True, connector=connector - ) as session: - async with session.get( - "https://api.soulter.top/astrbot/plugins-md5" - ) as response: - if response.status == 200: - data = await response.json() - return data.get("md5", "") - else: - logger.error(f"获取MD5失败，状态码：{response.status}") - return "" + async with ( + aiohttp.ClientSession( + trust_env=True, + connector=connector, + ) as session, + session.get( + "https://api.soulter.top/astrbot/plugins-md5", + ) as response, + ): + if response.status == 200: + data = await response.json() + return data.get("md5", "") + logger.error(f"获取MD5失败，状态码：{response.status}") + return "" except Exception as e: logger.error(f"获取远程MD5失败: {e}") return "" @@ -197,19 +204,19 @@ class PluginRoute(Route): """加载本地缓存的插件市场数据""" try: if os.path.exists(cache_file): - with open(cache_file, "r", encoding="utf-8") as f: + with open(cache_file, encoding="utf-8") as f: cache_data = json.load(f) # 检查缓存是否有效 if "data" in cache_data and "timestamp" in cache_data: logger.debug( - f"加载缓存文件: {cache_file}, 缓存时间: {cache_data['timestamp']}" + f"加载缓存文件: {cache_file}, 缓存时间: {cache_data['timestamp']}", ) return cache_data["data"] except Exception as e: logger.warning(f"加载插件市场缓存失败: {e}") return None - def _save_plugin_cache(self, cache_file: str, data, md5: str = None): + def _save_plugin_cache(self, cache_file: str, data, md5: str | None = None): """保存插件市场数据到本地缓存""" try: # 确保目录存在 @@ -227,12 +234,27 @@ class PluginRoute(Route): except Exception as e: logger.warning(f"保存插件市场缓存失败: {e}") + async def get_plugin_logo_token(self, logo_path: str): + try: + if token := self._logo_cache.get(logo_path): + if not await file_token_service.check_token_expired(token): + return self._logo_cache[logo_path] + token = await file_token_service.register_file(logo_path, timeout=300) + self._logo_cache[logo_path] = token + return token + except Exception as e: + logger.warning(f"获取插件 Logo 失败: {e}") + return None + async def get_plugins(self): _plugin_resp = [] plugin_name = request.args.get("name") for plugin in self.plugin_manager.context.get_all_stars(): if plugin_name and plugin.name != plugin_name: continue + logo_url = None + if plugin.logo_path: + logo_url = await self.get_plugin_logo_token(plugin.logo_path) _t = { "name": plugin.name, "repo": "" if plugin.repo is None else plugin.repo, @@ -243,8 +265,10 @@ class PluginRoute(Route): "activated": plugin.activated, "online_vesion": "", "handlers": await self.get_plugin_handlers_info( - plugin.star_handler_full_names + plugin.star_handler_full_names, ), + "display_name": plugin.display_name, + "logo": f"/api/file/{logo_url}" if logo_url else None, } _plugin_resp.append(_t) return ( @@ -260,13 +284,15 @@ class PluginRoute(Route): for handler_full_name in handler_full_names: info = {} handler = star_handlers_registry.star_handlers_map.get( - handler_full_name, None + handler_full_name, + None, ) if handler is None: continue info["event_type"] = handler.event_type.name info["event_type_h"] = self.translated_event_type.get( - handler.event_type, handler.event_type.name + handler.event_type, + handler.event_type.name, ) info["handler_full_name"] = handler.handler_full_name info["desc"] = handler.desc @@ -289,7 +315,7 @@ class PluginRoute(Route): info["cmd"] = filter.get_complete_command_names()[0] info["cmd"] = info["cmd"].strip() info["sub_command"] = filter.print_cmd_tree( - filter.sub_command_filters + filter.sub_command_filters, ) elif isinstance(filter, RegexFilter): info["type"] = "正则匹配" @@ -369,9 +395,15 @@ class PluginRoute(Route): post_data = await request.json plugin_name = post_data["name"] + delete_config = post_data.get("delete_config", False) + delete_data = post_data.get("delete_data", False) try: logger.info(f"正在卸载插件 {plugin_name}") - await self.plugin_manager.uninstall_plugin(plugin_name) + await self.plugin_manager.uninstall_plugin( + plugin_name, + delete_config=delete_config, + delete_data=delete_data, + ) logger.info(f"卸载插件 {plugin_name} 成功") return Response().ok(None, "卸载成功").__dict__ except Exception as e: @@ -455,7 +487,8 @@ class PluginRoute(Route): return Response().error(f"插件 {plugin_name} 不存在").__dict__ plugin_dir = os.path.join( - self.plugin_manager.plugin_store_path, plugin_obj.root_dir_name + self.plugin_manager.plugin_store_path, + plugin_obj.root_dir_name, ) if not os.path.isdir(plugin_dir): @@ -469,7 +502,7 @@ class PluginRoute(Route): return Response().error(f"插件 {plugin_name} 没有README文件").__dict__ try: - with open(readme_path, "r", encoding="utf-8") as f: + with open(readme_path, encoding="utf-8") as f: readme_content = f.read() return ( @@ -479,4 +512,4 @@ class PluginRoute(Route): ) except Exception as e: logger.error(f"/api/plugin/readme: {traceback.format_exc()}") - return Response().error(f"读取README文件失败: {str(e)}").__dict__ + return Response().error(f"读取README文件失败: {e!s}").__dict__ diff --git a/astrbot/dashboard/routes/route.py b/astrbot/dashboard/routes/route.py index ec455ce3..1105b69a 100644 --- a/astrbot/dashboard/routes/route.py +++ b/astrbot/dashboard/routes/route.py @@ -1,7 +1,9 @@ -from astrbot.core.config.astrbot_config import AstrBotConfig from dataclasses import dataclass + from quart import Quart +from astrbot.core.config.astrbot_config import AstrBotConfig + @dataclass class RouteContext: diff --git a/astrbot/dashboard/routes/session_management.py b/astrbot/dashboard/routes/session_management.py index 1d632171..0b16c094 100644 --- a/astrbot/dashboard/routes/session_management.py +++ b/astrbot/dashboard/routes/session_management.py @@ -47,7 +47,10 @@ class SessionManagementRoute(Route): # 获取活跃的会话数据（处于对话内的会话） sessions_data, total = await self.db_helper.get_session_conversations( - page, page_size, search_query, platform + page, + page_size, + search_query, + platform, ) provider_manager = self.core_lifecycle.provider_manager @@ -80,13 +83,13 @@ class SessionManagementRoute(Route): "stt_provider_id": None, "tts_provider_id": None, "session_enabled": SessionServiceManager.is_session_enabled( - session_id + session_id, ), "llm_enabled": SessionServiceManager.is_llm_enabled_for_session( - session_id + session_id, ), "tts_enabled": SessionServiceManager.is_tts_enabled_for_session( - session_id + session_id, ), "platform": session_id.split(":")[0] if ":" in session_id @@ -95,7 +98,7 @@ class SessionManagementRoute(Route): if session_id.count(":") >= 1 else "unknown", "session_name": SessionServiceManager.get_session_display_name( - session_id + session_id, ), "session_raw_name": session_id.split(":")[2] if session_id.count(":") >= 2 @@ -105,13 +108,16 @@ class SessionManagementRoute(Route): # 获取 provider 信息 chat_provider = provider_manager.get_using_provider( - provider_type=ProviderType.CHAT_COMPLETION, umo=session_id + provider_type=ProviderType.CHAT_COMPLETION, + umo=session_id, ) tts_provider = provider_manager.get_using_provider( - provider_type=ProviderType.TEXT_TO_SPEECH, umo=session_id + provider_type=ProviderType.TEXT_TO_SPEECH, + umo=session_id, ) stt_provider = provider_manager.get_using_provider( - provider_type=ProviderType.SPEECH_TO_TEXT, umo=session_id + provider_type=ProviderType.SPEECH_TO_TEXT, + umo=session_id, ) if chat_provider: meta = chat_provider.meta() @@ -139,7 +145,7 @@ class SessionManagementRoute(Route): "name": meta.id, "model": meta.model, "type": meta.type, - } + }, ) available_stt_providers = [] @@ -151,7 +157,7 @@ class SessionManagementRoute(Route): "name": meta.id, "model": meta.model, "type": meta.type, - } + }, ) available_tts_providers = [] @@ -163,7 +169,7 @@ class SessionManagementRoute(Route): "name": meta.id, "model": meta.model, "type": meta.type, - } + }, ) result = { @@ -185,15 +191,15 @@ class SessionManagementRoute(Route): return Response().ok(result).__dict__ except Exception as e: - error_msg = f"获取会话列表失败: {str(e)}\n{traceback.format_exc()}" + error_msg = f"获取会话列表失败: {e!s}\n{traceback.format_exc()}" logger.error(error_msg) - return Response().error(f"获取会话列表失败: {str(e)}").__dict__ + return Response().error(f"获取会话列表失败: {e!s}").__dict__ async def _update_single_session_persona(self, session_id: str, persona_name: str): """更新单个会话的 persona 的内部方法""" conversation_manager = self.core_lifecycle.star_context.conversation_manager conversation_id = await conversation_manager.get_curr_conversation_id( - session_id + session_id, ) conv = None @@ -207,11 +213,16 @@ class SessionManagementRoute(Route): # 更新 persona await conversation_manager.update_conversation_persona_id( - session_id, persona_name + session_id, + persona_name, ) async def _handle_batch_operation( - self, session_ids: list, operation_func, operation_name: str, **kwargs + self, + session_ids: list, + operation_func, + operation_name: str, + **kwargs, ): """通用的批量操作处理方法""" success_count = 0 @@ -222,7 +233,7 @@ class SessionManagementRoute(Route): await operation_func(session_id, **kwargs) success_count += 1 except Exception as e: - logger.error(f"批量{operation_name} 会话 {session_id} 失败: {str(e)}") + logger.error(f"批量{operation_name} 会话 {session_id} 失败: {e!s}") error_sessions.append(session_id) if error_sessions: @@ -234,21 +245,20 @@ class SessionManagementRoute(Route): "success_count": success_count, "error_count": len(error_sessions), "error_sessions": error_sessions, - } + }, ) .__dict__ ) - else: - return ( - Response() - .ok( - { - "message": f"成功批量{operation_name} {success_count} 个会话", - "success_count": success_count, - } - ) - .__dict__ + return ( + Response() + .ok( + { + "message": f"成功批量{operation_name} {success_count} 个会话", + "success_count": success_count, + }, ) + .__dict__ + ) async def update_session_persona(self): """更新指定会话的 persona，支持批量操作""" @@ -271,29 +281,31 @@ class SessionManagementRoute(Route): "更新人格", persona_name=persona_name, ) - else: - session_id = data.get("session_id") - if not session_id: - return Response().error("缺少必要参数: session_id").__dict__ + session_id = data.get("session_id") + if not session_id: + return Response().error("缺少必要参数: session_id").__dict__ - await self._update_single_session_persona(session_id, persona_name) - return ( - Response() - .ok( - { - "message": f"成功更新会话 {session_id} 的人格为 {persona_name}" - } - ) - .__dict__ + await self._update_single_session_persona(session_id, persona_name) + return ( + Response() + .ok( + { + "message": f"成功更新会话 {session_id} 的人格为 {persona_name}", + }, ) + .__dict__ + ) except Exception as e: - error_msg = f"更新会话人格失败: {str(e)}\n{traceback.format_exc()}" + error_msg = f"更新会话人格失败: {e!s}\n{traceback.format_exc()}" logger.error(error_msg) - return Response().error(f"更新会话人格失败: {str(e)}").__dict__ + return Response().error(f"更新会话人格失败: {e!s}").__dict__ async def _update_single_session_provider( - self, session_id: str, provider_id: str, provider_type_enum + self, + session_id: str, + provider_id: str, + provider_type_enum, ): """更新单个会话的 provider 的内部方法""" provider_manager = self.core_lifecycle.star_context.provider_manager @@ -344,28 +356,29 @@ class SessionManagementRoute(Route): provider_id=provider_id, provider_type_enum=provider_type_enum, ) - else: - session_id = data.get("session_id") - if not session_id: - return Response().error("缺少必要参数: session_id").__dict__ + session_id = data.get("session_id") + if not session_id: + return Response().error("缺少必要参数: session_id").__dict__ - await self._update_single_session_provider( - session_id, provider_id, provider_type_enum - ) - return ( - Response() - .ok( - { - "message": f"成功更新会话 {session_id} 的 {provider_type} 提供商为 {provider_id}" - } - ) - .__dict__ + await self._update_single_session_provider( + session_id, + provider_id, + provider_type_enum, + ) + return ( + Response() + .ok( + { + "message": f"成功更新会话 {session_id} 的 {provider_type} 提供商为 {provider_id}", + }, ) + .__dict__ + ) except Exception as e: - error_msg = f"更新会话提供商失败: {str(e)}\n{traceback.format_exc()}" + error_msg = f"更新会话提供商失败: {e!s}\n{traceback.format_exc()}" logger.error(error_msg) - return Response().error(f"更新会话提供商失败: {str(e)}").__dict__ + return Response().error(f"更新会话提供商失败: {e!s}").__dict__ async def get_session_plugins(self): """获取指定会话的插件配置信息""" @@ -384,7 +397,8 @@ class SessionManagementRoute(Route): if plugin.activated and not plugin.reserved: plugin_name = plugin.name or "" plugin_enabled = SessionPluginManager.is_plugin_enabled_for_session( - session_id, plugin_name + session_id, + plugin_name, ) all_plugins.append( @@ -393,7 +407,7 @@ class SessionManagementRoute(Route): "author": plugin.author, "desc": plugin.desc, "enabled": plugin_enabled, - } + }, ) return ( @@ -402,15 +416,15 @@ class SessionManagementRoute(Route): { "session_id": session_id, "plugins": all_plugins, - } + }, ) .__dict__ ) except Exception as e: - error_msg = f"获取会话插件配置失败: {str(e)}\n{traceback.format_exc()}" + error_msg = f"获取会话插件配置失败: {e!s}\n{traceback.format_exc()}" logger.error(error_msg) - return Response().error(f"获取会话插件配置失败: {str(e)}").__dict__ + return Response().error(f"获取会话插件配置失败: {e!s}").__dict__ async def update_session_plugin(self): """更新指定会话的插件启停状态""" @@ -448,7 +462,9 @@ class SessionManagementRoute(Route): # 使用 SessionPluginManager 更新插件状态 SessionPluginManager.set_plugin_status_for_session( - session_id, plugin_name, enabled + session_id, + plugin_name, + enabled, ) return ( @@ -459,15 +475,15 @@ class SessionManagementRoute(Route): "session_id": session_id, "plugin_name": plugin_name, "enabled": enabled, - } + }, ) .__dict__ ) except Exception as e: - error_msg = f"更新会话插件状态失败: {str(e)}\n{traceback.format_exc()}" + error_msg = f"更新会话插件状态失败: {e!s}\n{traceback.format_exc()}" logger.error(error_msg) - return Response().error(f"更新会话插件状态失败: {str(e)}").__dict__ + return Response().error(f"更新会话插件状态失败: {e!s}").__dict__ async def _update_single_session_llm(self, session_id: str, enabled: bool): """更新单个会话的LLM状态的内部方法""" @@ -495,28 +511,27 @@ class SessionManagementRoute(Route): enabled=enabled, ) return result - else: - session_id = data.get("session_id") - if not session_id: - return Response().error("缺少必要参数: session_id").__dict__ + session_id = data.get("session_id") + if not session_id: + return Response().error("缺少必要参数: session_id").__dict__ - await self._update_single_session_llm(session_id, enabled) - return ( - Response() - .ok( - { - "message": f"LLM已{'启用' if enabled else '禁用'}", - "session_id": session_id, - "llm_enabled": enabled, - } - ) - .__dict__ + await self._update_single_session_llm(session_id, enabled) + return ( + Response() + .ok( + { + "message": f"LLM已{'启用' if enabled else '禁用'}", + "session_id": session_id, + "llm_enabled": enabled, + }, ) + .__dict__ + ) except Exception as e: - error_msg = f"更新会话LLM状态失败: {str(e)}\n{traceback.format_exc()}" + error_msg = f"更新会话LLM状态失败: {e!s}\n{traceback.format_exc()}" logger.error(error_msg) - return Response().error(f"更新会话LLM状态失败: {str(e)}").__dict__ + return Response().error(f"更新会话LLM状态失败: {e!s}").__dict__ async def _update_single_session_tts(self, session_id: str, enabled: bool): """更新单个会话的TTS状态的内部方法""" @@ -544,28 +559,27 @@ class SessionManagementRoute(Route): enabled=enabled, ) return result - else: - session_id = data.get("session_id") - if not session_id: - return Response().error("缺少必要参数: session_id").__dict__ + session_id = data.get("session_id") + if not session_id: + return Response().error("缺少必要参数: session_id").__dict__ - await self._update_single_session_tts(session_id, enabled) - return ( - Response() - .ok( - { - "message": f"TTS已{'启用' if enabled else '禁用'}", - "session_id": session_id, - "tts_enabled": enabled, - } - ) - .__dict__ + await self._update_single_session_tts(session_id, enabled) + return ( + Response() + .ok( + { + "message": f"TTS已{'启用' if enabled else '禁用'}", + "session_id": session_id, + "tts_enabled": enabled, + }, ) + .__dict__ + ) except Exception as e: - error_msg = f"更新会话TTS状态失败: {str(e)}\n{traceback.format_exc()}" + error_msg = f"更新会话TTS状态失败: {e!s}\n{traceback.format_exc()}" logger.error(error_msg) - return Response().error(f"更新会话TTS状态失败: {str(e)}").__dict__ + return Response().error(f"更新会话TTS状态失败: {e!s}").__dict__ async def update_session_name(self): """更新指定会话的自定义名称""" @@ -588,17 +602,17 @@ class SessionManagementRoute(Route): "session_id": session_id, "custom_name": custom_name, "display_name": SessionServiceManager.get_session_display_name( - session_id + session_id, ), - } + }, ) .__dict__ ) except Exception as e: - error_msg = f"更新会话名称失败: {str(e)}\n{traceback.format_exc()}" + error_msg = f"更新会话名称失败: {e!s}\n{traceback.format_exc()}" logger.error(error_msg) - return Response().error(f"更新会话名称失败: {str(e)}").__dict__ + return Response().error(f"更新会话名称失败: {e!s}").__dict__ async def update_session_status(self): """更新指定会话的整体启停状态""" @@ -623,15 +637,15 @@ class SessionManagementRoute(Route): "message": f"会话整体状态已更新为: {'启用' if session_enabled else '禁用'}", "session_id": session_id, "session_enabled": session_enabled, - } + }, ) .__dict__ ) except Exception as e: - error_msg = f"更新会话整体状态失败: {str(e)}\n{traceback.format_exc()}" + error_msg = f"更新会话整体状态失败: {e!s}\n{traceback.format_exc()}" logger.error(error_msg) - return Response().error(f"更新会话整体状态失败: {str(e)}").__dict__ + return Response().error(f"更新会话整体状态失败: {e!s}").__dict__ async def delete_session(self): """删除指定会话及其所有相关数据""" @@ -649,13 +663,13 @@ class SessionManagementRoute(Route): try: await conversation_manager.delete_conversations_by_user_id(session_id) except Exception as e: - logger.warning(f"删除会话 {session_id} 的对话失败: {str(e)}") + logger.warning(f"删除会话 {session_id} 的对话失败: {e!s}") # 2. 清除会话的偏好设置数据（清空该会话的所有配置） try: await sp.clear_async("umo", session_id) except Exception as e: - logger.warning(f"清除会话 {session_id} 的偏好设置失败: {str(e)}") + logger.warning(f"清除会话 {session_id} 的偏好设置失败: {e!s}") return ( Response() @@ -663,12 +677,12 @@ class SessionManagementRoute(Route): { "message": f"会话 {session_id} 及其相关所有对话数据已成功删除", "session_id": session_id, - } + }, ) .__dict__ ) except Exception as e: - error_msg = f"删除会话失败: {str(e)}\n{traceback.format_exc()}" + error_msg = f"删除会话失败: {e!s}\n{traceback.format_exc()}" logger.error(error_msg) - return Response().error(f"删除会话失败: {str(e)}").__dict__ + return Response().error(f"删除会话失败: {e!s}").__dict__ diff --git a/astrbot/dashboard/routes/stat.py b/astrbot/dashboard/routes/stat.py index d13eb802..8df690cc 100644 --- a/astrbot/dashboard/routes/stat.py +++ b/astrbot/dashboard/routes/stat.py @@ -1,17 +1,19 @@ -import traceback -import psutil -import time import threading +import time +import traceback + import aiohttp -from .route import Route, Response, RouteContext -from astrbot.core import logger +import psutil from quart import request + +from astrbot.core import DEMO_MODE, logger +from astrbot.core.config import VERSION from astrbot.core.core_lifecycle import AstrBotCoreLifecycle from astrbot.core.db import BaseDatabase -from astrbot.core.config import VERSION -from astrbot.core.utils.io import get_dashboard_version -from astrbot.core import DEMO_MODE from astrbot.core.db.migration.helper import check_migration_needed_v4 +from astrbot.core.utils.io import get_dashboard_version + +from .route import Response, Route, RouteContext class StatRoute(Route): @@ -70,7 +72,7 @@ class StatRoute(Route): "dashboard_version": await get_dashboard_version(), "change_pwd_hint": self.is_default_cred(), "need_migration": need_migration, - } + }, ) .__dict__ ) @@ -116,17 +118,17 @@ class StatRoute(Route): # 计算运行时长组件 running_time = self._get_running_time_components( - int(time.time()) - self.core_lifecycle.start_time + int(time.time()) - self.core_lifecycle.start_time, ) stat_dict.update( { "platform": self.db_helper.get_grouped_base_stats( - offset_sec + offset_sec, ).platform, "message_count": self.db_helper.get_total_message_count() or 0, "platform_count": len( - self.core_lifecycle.platform_manager.get_insts() + self.core_lifecycle.platform_manager.get_insts(), ), "plugin_count": len(plugins), "plugins": plugin_info, @@ -139,7 +141,7 @@ class StatRoute(Route): "cpu_percent": round(cpu_percent, 1), "thread_count": thread_count, "start_time": self.core_lifecycle.start_time, - } + }, ) return Response().ok(stat_dict).__dict__ @@ -148,9 +150,7 @@ class StatRoute(Route): return Response().error(e.__str__()).__dict__ async def test_ghproxy_connection(self): - """ - 测试 GitHub 代理连接是否可用。 - """ + """测试 GitHub 代理连接是否可用。""" try: data = await request.get_json() proxy_url: str = data.get("proxy_url") @@ -163,23 +163,23 @@ class StatRoute(Route): test_url = f"{proxy_url}/https://github.com/AstrBotDevs/AstrBot/raw/refs/heads/master/.python-version" start_time = time.time() - async with aiohttp.ClientSession() as session: - async with session.get( - test_url, timeout=aiohttp.ClientTimeout(total=10) - ) as response: - if response.status == 200: - end_time = time.time() - _ = await response.text() - ret = { - "latency": round((end_time - start_time) * 1000, 2), - } - return Response().ok(data=ret).__dict__ - else: - return ( - Response() - .error(f"Failed. Status code: {response.status}") - .__dict__ - ) + async with ( + aiohttp.ClientSession() as session, + session.get( + test_url, + timeout=aiohttp.ClientTimeout(total=10), + ) as response, + ): + if response.status == 200: + end_time = time.time() + _ = await response.text() + ret = { + "latency": round((end_time - start_time) * 1000, 2), + } + return Response().ok(data=ret).__dict__ + return ( + Response().error(f"Failed. Status code: {response.status}").__dict__ + ) except Exception as e: logger.error(traceback.format_exc()) - return Response().error(f"Error: {str(e)}").__dict__ + return Response().error(f"Error: {e!s}").__dict__ diff --git a/astrbot/dashboard/routes/t2i.py b/astrbot/dashboard/routes/t2i.py index 04f87bc9..db70a882 100644 --- a/astrbot/dashboard/routes/t2i.py +++ b/astrbot/dashboard/routes/t2i.py @@ -1,11 +1,13 @@ # astrbot/dashboard/routes/t2i.py from dataclasses import asdict + from quart import jsonify, request from astrbot.core import logger from astrbot.core.core_lifecycle import AstrBotCoreLifecycle from astrbot.core.utils.t2i.template_manager import TemplateManager + from .route import Response, Route, RouteContext @@ -49,7 +51,7 @@ class T2iRoute(Route): try: active_template = self.config.get("t2i_active_template", "base") return jsonify( - asdict(Response().ok(data={"active_template": active_template})) + asdict(Response().ok(data={"active_template": active_template})), ) except Exception as e: logger.error("Error in get_active_template", exc_info=True) @@ -62,7 +64,7 @@ class T2iRoute(Route): try: content = self.manager.get_template(name) return jsonify( - asdict(Response().ok(data={"name": name, "content": content})) + asdict(Response().ok(data={"name": name, "content": content})), ) except FileNotFoundError: response = jsonify(asdict(Response().error("Template not found"))) @@ -81,7 +83,7 @@ class T2iRoute(Route): content = data.get("content") if not name or not content: response = jsonify( - asdict(Response().error("Name and content are required.")) + asdict(Response().error("Name and content are required.")), ) response.status_code = 400 return response @@ -91,15 +93,16 @@ class T2iRoute(Route): response = jsonify( asdict( Response().ok( - data={"name": name}, message="Template created successfully." - ) - ) + data={"name": name}, + message="Template created successfully.", + ), + ), ) response.status_code = 201 return response except FileExistsError: response = jsonify( - asdict(Response().error("Template with this name already exists.")) + asdict(Response().error("Template with this name already exists.")), ) response.status_code = 409 return response @@ -149,7 +152,7 @@ class T2iRoute(Route): name = name.strip() self.manager.delete_template(name) return jsonify( - asdict(Response().ok(message="Template deleted successfully.")) + asdict(Response().ok(message="Template deleted successfully.")), ) except FileNotFoundError: response = jsonify(asdict(Response().error("Template not found."))) @@ -189,7 +192,7 @@ class T2iRoute(Route): except FileNotFoundError: response = jsonify( - asdict(Response().error(f"模板 '{name}' 不存在，无法应用。")) + asdict(Response().error(f"模板 '{name}' 不存在，无法应用。")), ) response.status_code = 404 return response @@ -215,9 +218,9 @@ class T2iRoute(Route): return jsonify( asdict( Response().ok( - message="Default template has been reset and activated." - ) - ) + message="Default template has been reset and activated.", + ), + ), ) except FileNotFoundError as e: response = jsonify(asdict(Response().error(str(e)))) diff --git a/astrbot/dashboard/routes/tools.py b/astrbot/dashboard/routes/tools.py index 8fd89919..64cd78ca 100644 --- a/astrbot/dashboard/routes/tools.py +++ b/astrbot/dashboard/routes/tools.py @@ -13,7 +13,9 @@ DEFAULT_MCP_CONFIG = {"mcpServers": {}} class ToolsRoute(Route): def __init__( - self, context: RouteContext, core_lifecycle: AstrBotCoreLifecycle + self, + context: RouteContext, + core_lifecycle: AstrBotCoreLifecycle, ) -> None: super().__init__(context) self.core_lifecycle = core_lifecycle @@ -64,7 +66,7 @@ class ToolsRoute(Route): return Response().ok(servers).__dict__ except Exception as e: logger.error(traceback.format_exc()) - return Response().error(f"获取 MCP 服务器列表失败: {str(e)}").__dict__ + return Response().error(f"获取 MCP 服务器列表失败: {e!s}").__dict__ async def add_mcp_server(self): try: @@ -105,23 +107,22 @@ class ToolsRoute(Route): if self.tool_mgr.save_mcp_config(config): try: await self.tool_mgr.enable_mcp_server( - name, server_config, timeout=30 + name, + server_config, + timeout=30, ) except TimeoutError: return Response().error(f"启用 MCP 服务器 {name} 超时。").__dict__ except Exception as e: logger.error(traceback.format_exc()) return ( - Response() - .error(f"启用 MCP 服务器 {name} 失败: {str(e)}") - .__dict__ + Response().error(f"启用 MCP 服务器 {name} 失败: {e!s}").__dict__ ) return Response().ok(None, f"成功添加 MCP 服务器 {name}").__dict__ - else: - return Response().error("保存配置失败").__dict__ + return Response().error("保存配置失败").__dict__ except Exception as e: logger.error(traceback.format_exc()) - return Response().error(f"添加 MCP 服务器失败: {str(e)}").__dict__ + return Response().error(f"添加 MCP 服务器失败: {e!s}").__dict__ async def update_mcp_server(self): try: @@ -139,7 +140,8 @@ class ToolsRoute(Route): # 获取活动状态 active = server_data.get( - "active", config["mcpServers"][name].get("active", True) + "active", + config["mcpServers"][name].get("active", True), ) # 创建新的配置对象 @@ -177,19 +179,21 @@ class ToolsRoute(Route): except TimeoutError as e: return ( Response() - .error(f"启用前停用 MCP 服务器时 {name} 超时: {str(e)}") + .error(f"启用前停用 MCP 服务器时 {name} 超时: {e!s}") .__dict__ ) except Exception as e: logger.error(traceback.format_exc()) return ( Response() - .error(f"启用前停用 MCP 服务器时 {name} 失败: {str(e)}") + .error(f"启用前停用 MCP 服务器时 {name} 失败: {e!s}") .__dict__ ) try: await self.tool_mgr.enable_mcp_server( - name, config["mcpServers"][name], timeout=30 + name, + config["mcpServers"][name], + timeout=30, ) except TimeoutError: return ( @@ -199,34 +203,30 @@ class ToolsRoute(Route): logger.error(traceback.format_exc()) return ( Response() - .error(f"启用 MCP 服务器 {name} 失败: {str(e)}") + .error(f"启用 MCP 服务器 {name} 失败: {e!s}") + .__dict__ + ) + # 如果要停用服务器 + elif name in self.tool_mgr.mcp_client_dict: + try: + await self.tool_mgr.disable_mcp_server(name, timeout=10) + except TimeoutError: + return ( + Response().error(f"停用 MCP 服务器 {name} 超时。").__dict__ + ) + except Exception as e: + logger.error(traceback.format_exc()) + return ( + Response() + .error(f"停用 MCP 服务器 {name} 失败: {e!s}") .__dict__ ) - else: - # 如果要停用服务器 - if name in self.tool_mgr.mcp_client_dict: - try: - await self.tool_mgr.disable_mcp_server(name, timeout=10) - except TimeoutError: - return ( - Response() - .error(f"停用 MCP 服务器 {name} 超时。") - .__dict__ - ) - except Exception as e: - logger.error(traceback.format_exc()) - return ( - Response() - .error(f"停用 MCP 服务器 {name} 失败: {str(e)}") - .__dict__ - ) return Response().ok(None, f"成功更新 MCP 服务器 {name}").__dict__ - else: - return Response().error("保存配置失败").__dict__ + return Response().error("保存配置失败").__dict__ except Exception as e: logger.error(traceback.format_exc()) - return Response().error(f"更新 MCP 服务器失败: {str(e)}").__dict__ + return Response().error(f"更新 MCP 服务器失败: {e!s}").__dict__ async def delete_mcp_server(self): try: @@ -255,20 +255,17 @@ class ToolsRoute(Route): logger.error(traceback.format_exc()) return ( Response() - .error(f"停用 MCP 服务器 {name} 失败: {str(e)}") + .error(f"停用 MCP 服务器 {name} 失败: {e!s}") .__dict__ ) return Response().ok(None, f"成功删除 MCP 服务器 {name}").__dict__ - else: - return Response().error("保存配置失败").__dict__ + return Response().error("保存配置失败").__dict__ except Exception as e: logger.error(traceback.format_exc()) - return Response().error(f"删除 MCP 服务器失败: {str(e)}").__dict__ + return Response().error(f"删除 MCP 服务器失败: {e!s}").__dict__ async def test_mcp_connection(self): - """ - 测试 MCP 服务器连接 - """ + """测试 MCP 服务器连接""" try: server_data = await request.json config = server_data.get("mcp_server_config", None) @@ -283,9 +280,8 @@ class ToolsRoute(Route): if len(keys) > 1: return Response().error("一次只能配置一个 MCP 服务器配置").__dict__ config = config["mcpServers"][keys[0]] - else: - if not config: - return Response().error("MCP 服务器配置不能为空").__dict__ + elif not config: + return Response().error("MCP 服务器配置不能为空").__dict__ tools_name = await self.tool_mgr.test_mcp_server_connection(config) return ( @@ -294,17 +290,25 @@ class ToolsRoute(Route): except Exception as e: logger.error(traceback.format_exc()) - return Response().error(f"测试 MCP 连接失败: {str(e)}").__dict__ + return Response().error(f"测试 MCP 连接失败: {e!s}").__dict__ async def get_tool_list(self): """获取所有注册的工具列表""" try: tools = self.tool_mgr.func_list - tools_dict = [tool.__dict__() for tool in tools] + tools_dict = [ + { + "name": tool.name, + "description": tool.description, + "parameters": tool.parameters, + "active": tool.active, + } + for tool in tools + ] return Response().ok(data=tools_dict).__dict__ except Exception as e: logger.error(traceback.format_exc()) - return Response().error(f"获取工具列表失败: {str(e)}").__dict__ + return Response().error(f"获取工具列表失败: {e!s}").__dict__ async def toggle_tool(self): """启用或停用指定的工具""" @@ -320,18 +324,17 @@ class ToolsRoute(Route): try: ok = self.tool_mgr.activate_llm_tool(tool_name, star_map=star_map) except ValueError as e: - return Response().error(f"启用工具失败: {str(e)}").__dict__ + return Response().error(f"启用工具失败: {e!s}").__dict__ else: ok = self.tool_mgr.deactivate_llm_tool(tool_name) if ok: return Response().ok(None, "操作成功。").__dict__ - else: - return Response().error(f"工具 {tool_name} 不存在或操作失败。").__dict__ + return Response().error(f"工具 {tool_name} 不存在或操作失败。").__dict__ except Exception as e: logger.error(traceback.format_exc()) - return Response().error(f"操作工具失败: {str(e)}").__dict__ + return Response().error(f"操作工具失败: {e!s}").__dict__ async def sync_provider(self): """同步 MCP 提供者配置""" @@ -348,4 +351,4 @@ class ToolsRoute(Route): return Response().ok(message="同步成功").__dict__ except Exception as e: logger.error(traceback.format_exc()) - return Response().error(f"同步失败: {str(e)}").__dict__ + return Response().error(f"同步失败: {e!s}").__dict__ diff --git a/astrbot/dashboard/routes/update.py b/astrbot/dashboard/routes/update.py index 24184354..b0520c31 100644 --- a/astrbot/dashboard/routes/update.py +++ b/astrbot/dashboard/routes/update.py @@ -1,13 +1,15 @@ import traceback -from .route import Route, Response, RouteContext + from quart import request -from astrbot.core.core_lifecycle import AstrBotCoreLifecycle -from astrbot.core.updator import AstrBotUpdator -from astrbot.core import logger, pip_installer -from astrbot.core.utils.io import download_dashboard, get_dashboard_version + +from astrbot.core import DEMO_MODE, logger, pip_installer from astrbot.core.config.default import VERSION -from astrbot.core import DEMO_MODE -from astrbot.core.db.migration.helper import do_migration_v4, check_migration_needed_v4 +from astrbot.core.core_lifecycle import AstrBotCoreLifecycle +from astrbot.core.db.migration.helper import check_migration_needed_v4, do_migration_v4 +from astrbot.core.updator import AstrBotUpdator +from astrbot.core.utils.io import download_dashboard, get_dashboard_version + +from .route import Response, Route, RouteContext CLEAR_SITE_DATA_HEADERS = {"Clear-Site-Data": '"cache"'} @@ -40,46 +42,39 @@ class UpdateRoute(Route): data = await request.json pim = data.get("platform_id_map", {}) await do_migration_v4( - self.core_lifecycle.db, pim, self.core_lifecycle.astrbot_config + self.core_lifecycle.db, + pim, + self.core_lifecycle.astrbot_config, ) return Response().ok(None, "迁移成功。").__dict__ except Exception as e: logger.error(f"迁移失败: {traceback.format_exc()}") - return Response().error(f"迁移失败: {str(e)}").__dict__ + return Response().error(f"迁移失败: {e!s}").__dict__ async def check_update(self): type_ = request.args.get("type", None) try: dv = await get_dashboard_version() - # WebUI 版本独立于核心版本：不再用 dv 与 v{VERSION} 比较，避免误报 if type_ == "dashboard": return ( Response() - .ok( - { - "has_new_version": False, - "current_version": dv, - "installed": bool(dv), - } - ) + .ok({"has_new_version": dv != f"v{VERSION}", "current_version": dv}) .__dict__ ) - else: - ret = await self.astrbot_updator.check_update(None, None, False) - return Response( - status="success", - message=str(ret) if ret is not None else "已经是最新版本了。", - data={ - "version": f"v{VERSION}", - "has_new_version": ret is not None, - "dashboard_version": dv, - # dv正常获取则不会提示需要更新 - "dashboard_has_new_version": not bool(dv), - }, - ).__dict__ + ret = await self.astrbot_updator.check_update(None, None, False) + return Response( + status="success", + message=str(ret) if ret is not None else "已经是最新版本了。", + data={ + "version": f"v{VERSION}", + "has_new_version": ret is not None, + "dashboard_version": dv, + "dashboard_has_new_version": bool(dv and dv != f"v{VERSION}"), + }, + ).__dict__ except Exception as e: - logger.warning(f"检查更新失败: {str(e)} (不影响除项目更新外的正常使用)") + logger.warning(f"检查更新失败: {e!s} (不影响除项目更新外的正常使用)") return Response().error(e.__str__()).__dict__ async def get_releases(self): @@ -106,7 +101,9 @@ class UpdateRoute(Route): try: await self.astrbot_updator.update( - latest=latest, version=version, proxy=proxy + latest=latest, + version=version, + proxy=proxy, ) try: @@ -129,13 +126,12 @@ class UpdateRoute(Route): .__dict__ ) return ret, 200, CLEAR_SITE_DATA_HEADERS - else: - ret = ( - Response() - .ok(None, "更新成功，AstrBot 将在下次启动时应用新的代码。") - .__dict__ - ) - return ret, 200, CLEAR_SITE_DATA_HEADERS + ret = ( + Response() + .ok(None, "更新成功，AstrBot 将在下次启动时应用新的代码。") + .__dict__ + ) + return ret, 200, CLEAR_SITE_DATA_HEADERS except Exception as e: logger.error(f"/api/update_project: {traceback.format_exc()}") return Response().error(e.__str__()).__dict__ diff --git a/astrbot/dashboard/server.py b/astrbot/dashboard/server.py index 9b9588d4..84976f2b 100644 --- a/astrbot/dashboard/server.py +++ b/astrbot/dashboard/server.py @@ -39,7 +39,7 @@ class AstrBotDashboard: self.data_path = os.path.abspath(webui_dir) else: self.data_path = os.path.abspath( - os.path.join(get_astrbot_data_path(), "dist") + os.path.join(get_astrbot_data_path(), "dist"), ) self.app = Quart("dashboard", static_folder=self.data_path, static_url_path="/") @@ -53,11 +53,15 @@ class AstrBotDashboard: logging.getLogger(self.app.name).removeHandler(default_handler) self.context = RouteContext(self.config, self.app) self.ur = UpdateRoute( - self.context, core_lifecycle.astrbot_updator, core_lifecycle + self.context, + core_lifecycle.astrbot_updator, + core_lifecycle, ) self.sr = StatRoute(self.context, db, core_lifecycle) self.pr = PluginRoute( - self.context, core_lifecycle, core_lifecycle.plugin_manager + self.context, + core_lifecycle, + core_lifecycle.plugin_manager, ) self.cr = ConfigRoute(self.context, core_lifecycle) self.lr = LogRoute(self.context, core_lifecycle.log_broker) @@ -68,10 +72,13 @@ class AstrBotDashboard: self.conversation_route = ConversationRoute(self.context, db, core_lifecycle) self.file_route = FileRoute(self.context) self.session_management_route = SessionManagementRoute( - self.context, db, core_lifecycle + self.context, + db, + core_lifecycle, ) self.persona_route = PersonaRoute(self.context, db, core_lifecycle) self.t2i_route = T2iRoute(self.context, core_lifecycle) + self.kb_route = KnowledgeBaseRoute(self.context, core_lifecycle) self.app.add_url_rule( "/api/plug/", @@ -84,9 +91,7 @@ class AstrBotDashboard: self._init_jwt_secret() async def srv_plug_route(self, subpath, *args, **kwargs): - """ - 插件路由 - """ + """插件路由""" registered_web_apis = self.core_lifecycle.star_context.registered_web_apis for api in registered_web_apis: route, view_handler, methods, _ = api @@ -96,18 +101,17 @@ class AstrBotDashboard: async def auth_middleware(self): if not request.path.startswith("/api"): - return + return None allowed_endpoints = ["/api/auth/login", "/api/file"] if any(request.path.startswith(prefix) for prefix in allowed_endpoints): - return - # claim jwt + return None + # 声明 JWT token = request.headers.get("Authorization") if not token: r = jsonify(Response().error("未授权").__dict__) r.status_code = 401 return r - if token.startswith("Bearer "): - token = token[7:] + token = token.removeprefix("Bearer ") try: payload = jwt.decode(token, self._jwt_secret, algorithms=["HS256"]) g.username = payload["username"] @@ -121,9 +125,7 @@ class AstrBotDashboard: return r def check_port_in_use(self, port: int) -> bool: - """ - 跨平台检测端口是否被占用 - """ + """跨平台检测端口是否被占用""" try: # 创建 IPv4 TCP Socket sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) @@ -135,7 +137,7 @@ class AstrBotDashboard: # result 为 0 表示端口被占用 return result == 0 except Exception as e: - logger.warning(f"检查端口 {port} 时发生错误: {str(e)}") + logger.warning(f"检查端口 {port} 时发生错误: {e!s}") # 如果出现异常，保守起见认为端口可能被占用 return True @@ -156,10 +158,10 @@ class AstrBotDashboard: ] return "\n ".join(proc_info) except (psutil.NoSuchProcess, psutil.AccessDenied) as e: - return f"无法获取进程详细信息(可能需要管理员权限): {str(e)}" + return f"无法获取进程详细信息(可能需要管理员权限): {e!s}" return "未找到占用进程" except Exception as e: - return f"获取进程信息失败: {str(e)}" + return f"获取进程信息失败: {e!s}" def _init_jwt_secret(self): if not self.config.get("dashboard", {}).get("jwt_secret", None): @@ -177,12 +179,17 @@ class AstrBotDashboard: else: port = self.core_lifecycle.astrbot_config["dashboard"].get("port", 6185) host = self.core_lifecycle.astrbot_config["dashboard"].get("host", "0.0.0.0") + enable = self.core_lifecycle.astrbot_config["dashboard"].get("enable", True) + + if not enable: + logger.info("WebUI 已被禁用") + return None logger.info(f"正在启动 WebUI, 监听地址: http://{host}:{port}") if host == "0.0.0.0": logger.info( - "提示: WebUI 将监听所有网络接口，请注意安全。（可在 data/cmd_config.json 中配置 dashboard.host 以修改 host）" + "提示: WebUI 将监听所有网络接口，请注意安全。（可在 data/cmd_config.json 中配置 dashboard.host 以修改 host）", ) if host not in ["localhost", "127.0.0.1"]: @@ -201,16 +208,17 @@ class AstrBotDashboard: f"请确保：\n" f"1. 没有其他 AstrBot 实例正在运行\n" f"2. 端口 {port} 没有被其他程序占用\n" - f"3. 如需使用其他端口，请修改配置文件" + f"3. 如需使用其他端口，请修改配置文件", ) raise Exception(f"端口 {port} 已被占用") - display = f"\n ✨✨✨\n AstrBot v{VERSION} WebUI 已启动，可访问\n\n" - display += f" ➜ 本地: http://localhost:{port}\n" + parts = [f"\n ✨✨✨\n AstrBot v{VERSION} WebUI 已启动，可访问\n\n"] + parts.append(f" ➜ 本地: http://localhost:{port}\n") for ip in ip_addr: - display += f" ➜ 网络: http://{ip}:{port}\n" - display += " ➜ 默认用户名和密码: astrbot\n ✨✨✨\n" + parts.append(f" ➜ 网络: http://{ip}:{port}\n") + parts.append(" ➜ 默认用户名和密码: astrbot\n ✨✨✨\n") + display = "".join(parts) if not ip_addr: display += ( @@ -220,7 +228,9 @@ class AstrBotDashboard: logger.info(display) return self.app.run_task( - host=host, port=port, shutdown_trigger=self.shutdown_trigger + host=host, + port=port, + shutdown_trigger=self.shutdown_trigger, ) async def shutdown_trigger(self): diff --git a/astrbot/dashboard/utils.py b/astrbot/dashboard/utils.py new file mode 100644 index 00000000..b81faad0 --- /dev/null +++ b/astrbot/dashboard/utils.py @@ -0,0 +1,165 @@ +import base64 +import os +import traceback +from io import BytesIO + +from astrbot.api import logger +from astrbot.core.db.vec_db.faiss_impl import FaissVecDB +from astrbot.core.knowledge_base.kb_helper import KBHelper +from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager + + +async def generate_tsne_visualization( + query: str, + kb_names: list[str], + kb_manager: KnowledgeBaseManager, +) -> str | None: + """生成 t-SNE 可视化图片 + + Args: + query: 查询文本 + kb_names: 知识库名称列表 + kb_manager: 知识库管理器 + + Returns: + 图片路径或 None + + """ + try: + import faiss + import matplotlib + import numpy as np + + matplotlib.use("Agg") # 使用非交互式后端 + import matplotlib.pyplot as plt + from sklearn.manifold import TSNE + except ImportError as e: + raise Exception( + "缺少必要的库以生成 t-SNE 可视化。请安装 matplotlib 和 scikit-learn: {e}", + ) from e + + try: + # 获取第一个知识库的向量数据 + kb_helper: KBHelper | None = None + for kb_name in kb_names: + kb_helper = await kb_manager.get_kb_by_name(kb_name) + if kb_helper: + break + + if not kb_helper: + logger.warning("未找到知识库") + return None + + kb = kb_helper.kb + index_path = f"data/knowledge_base/{kb.kb_id}/index.faiss" + + # 读取 FAISS 索引 + if not os.path.exists(index_path): + logger.warning(f"FAISS 索引不存在: {index_path}") + return None + + index = faiss.read_index(index_path) + + if index.ntotal == 0: + logger.warning("索引为空") + return None + + # 提取所有向量 + logger.info(f"提取 {index.ntotal} 个向量用于可视化...") + if isinstance(index, faiss.IndexIDMap): + base_index = faiss.downcast_index(index.index) + if hasattr(base_index, "reconstruct_n"): + vectors = base_index.reconstruct_n(0, index.ntotal) + else: + vectors = np.zeros((index.ntotal, index.d), dtype=np.float32) + for i in range(index.ntotal): + base_index.reconstruct(i, vectors[i]) + elif hasattr(index, "reconstruct_n"): + vectors = index.reconstruct_n(0, index.ntotal) + else: + vectors = np.zeros((index.ntotal, index.d), dtype=np.float32) + for i in range(index.ntotal): + index.reconstruct(i, vectors[i]) + + # 获取查询向量 + vec_db: FaissVecDB = kb_helper.vec_db # type: ignore + embedding_provider = vec_db.embedding_provider + query_embedding = await embedding_provider.get_embedding(query) + query_vector = np.array([query_embedding], dtype=np.float32) + + # 合并所有向量和查询向量 + all_vectors = np.vstack([vectors, query_vector]) + + # t-SNE 降维 + logger.info("开始 t-SNE 降维...") + perplexity = min(30, all_vectors.shape[0] - 1) + tsne = TSNE(n_components=2, random_state=42, perplexity=perplexity) + vectors_2d = tsne.fit_transform(all_vectors) + + # 分离知识库向量和查询向量 + kb_vectors_2d = vectors_2d[:-1] + query_vector_2d = vectors_2d[-1] + + # 可视化 + logger.info("生成可视化图表...") + plt.figure(figsize=(14, 10)) + + # 绘制知识库向量 + scatter = plt.scatter( + kb_vectors_2d[:, 0], + kb_vectors_2d[:, 1], + alpha=0.5, + s=40, + c=range(len(kb_vectors_2d)), + cmap="viridis", + label="Knowledge Base Vectors", + ) + + # 绘制查询向量（红色 X） + plt.scatter( + query_vector_2d[0], + query_vector_2d[1], + c="red", + s=300, + marker="X", + edgecolors="black", + linewidths=2, + label="Query", + zorder=5, + ) + + # 添加查询文本标注 + plt.annotate( + "Query", + (query_vector_2d[0], query_vector_2d[1]), + xytext=(10, 10), + textcoords="offset points", + fontsize=10, + bbox={"boxstyle": "round,pad=0.5", "fc": "yellow", "alpha": 0.7}, + arrowprops={"arrowstyle": "->", "connectionstyle": "arc3,rad=0"}, + ) + + plt.colorbar(scatter, label="Vector Index") + plt.title( + f"t-SNE Visualization: Query in Knowledge Base\n" + f"({index.ntotal} vectors, {index.d} dimensions, KB: {kb.kb_name})", + fontsize=14, + pad=20, + ) + plt.xlabel("t-SNE Dimension 1", fontsize=12) + plt.ylabel("t-SNE Dimension 2", fontsize=12) + plt.grid(True, alpha=0.3) + plt.legend(fontsize=10, loc="upper right") + + # base64 编码图片返回 + buffer = BytesIO() + plt.savefig(buffer, format="png", dpi=150, bbox_inches="tight") + plt.close() + buffer.seek(0) + img_base64 = base64.b64encode(buffer.read()).decode("utf-8") + return img_base64 + + except Exception as e: + logger.error(f"生成 t-SNE 可视化时出错: {e}") + logger.error(traceback.format_exc()) + return None diff --git a/changelogs/v3.4.0.md b/changelogs/v3.4.0.md index 6b5a52c0..ee3f0f2d 100644 --- a/changelogs/v3.4.0.md +++ b/changelogs/v3.4.0.md @@ -6,7 +6,7 @@ 4. 解耦了 Provider 的相关处理逻辑。 5. 解耦了 Platform 相关处理逻辑。 6. aiocqhttp 适配器支持设置群聊白名单、私聊白名单； -7. aiocqhttp 适配器将图片转换成 base64 格式上报，而不需要先上传到图床；https://github.com/Soulter/AstrBot/issues/219 +7. aiocqhttp 适配器将图片转换成 base64 格式上报，而不需要先上传到图床；https://github.com/AstrBotDevs/AstrBot/issues/219 8. qq_official 适配器在群聊/ C2C 场景下以 base64 格式直接上传到 QQ 服务器，而不需要先上传到图床； 9. 移除了对 nakuru 适配器的支持； 10. 移除了 update, reboot 等指令； diff --git a/changelogs/v3.4.16.md b/changelogs/v3.4.16.md index 15587666..96767f79 100644 --- a/changelogs/v3.4.16.md +++ b/changelogs/v3.4.16.md @@ -1,6 +1,6 @@ # What's Changed -- [gewechat] [修复每次启动astrbot都需要扫码的问题](https://github.com/Soulter/AstrBot/commit/fd5d7dd37a6d74f81a148bbebef8516aa0cb5540) -- [core] [Provider 重复时不直接报错闪退](https://github.com/Soulter/AstrBot/commit/b61f9be18db9a6b8b3c5b6b36553f66dd2b79375) https://github.com/Soulter/AstrBot/issues/265 -- [core] [弱化更新报错](https://github.com/Soulter/AstrBot/commit/0ba0150fd8ff2062dbe83889163888ba3e33bd49) https://github.com/Soulter/AstrBot/issues/267 +- [gewechat] [修复每次启动astrbot都需要扫码的问题](https://github.com/AstrBotDevs/AstrBot/commit/fd5d7dd37a6d74f81a148bbebef8516aa0cb5540) +- [core] [Provider 重复时不直接报错闪退](https://github.com/AstrBotDevs/AstrBot/commit/b61f9be18db9a6b8b3c5b6b36553f66dd2b79375) https://github.com/AstrBotDevs/AstrBot/issues/265 +- [core] [弱化更新报错](https://github.com/AstrBotDevs/AstrBot/commit/0ba0150fd8ff2062dbe83889163888ba3e33bd49) https://github.com/AstrBotDevs/AstrBot/issues/267 - 修复 webui 无法从本地上传插件的问题 \ No newline at end of file diff --git a/changelogs/v3.4.33.md b/changelogs/v3.4.33.md index f77a8854..f53e928d 100644 --- a/changelogs/v3.4.33.md +++ b/changelogs/v3.4.33.md @@ -1,12 +1,12 @@ # What's Changed 1. ✨ 新增: add English README by @CAICAIIs -2. ✨ 新增: perf: 优化网页录音 [#283](https://github.com/Soulter/AstrBot/issues/283) by @Fridemn -3. ✨ 新增: 添加对于 Edge-TTS 的支持 [#471](https://github.com/Soulter/AstrBot/issues/471) by @Fridemn +2. ✨ 新增: perf: 优化网页录音 [#283](https://github.com/AstrBotDevs/AstrBot/issues/283) by @Fridemn +3. ✨ 新增: 添加对于 Edge-TTS 的支持 [#471](https://github.com/AstrBotDevs/AstrBot/issues/471) by @Fridemn 4. ⚡ 优化: 为防止输入一大堆 k，改 k 键为 Ctrl 键；改为长按录音，松手结束；为防止误触改为只有点击输入框之后才会生效 by @Fridemn -5. ⚡ 优化: 插件市场非列表视图能够正常搜索 [#640](https://github.com/Soulter/AstrBot/issues/640) by @Fridemn +5. ⚡ 优化: 插件市场非列表视图能够正常搜索 [#640](https://github.com/AstrBotDevs/AstrBot/issues/640) by @Fridemn 6. ⚡ 优化: 插件市场帮助按钮 tooltip 移入时会消失无法点击其中链接，更改为按钮触发 by @Quirrel-zh -7. ‼️‼️ 🐛 修复: v3.4.32 无法记忆历史的会话 [#630](https://github.com/Soulter/AstrBot/issues/630) +7. ‼️‼️ 🐛 修复: v3.4.32 无法记忆历史的会话 [#630](https://github.com/AstrBotDevs/AstrBot/issues/630) 8. ‼️🐛 修复: 钩子函数无法终止事件传播的问题；修复某些情况下终止事件传播后仍然会请求 LLM 的问题 9. ‼️🐛 修复: OneBot V11 通知类事件某些情况无法回复问题 by @CAICAIIs 10. 🐛 修复: Correct STT model path and improve logging in provider manager and pip installer diff --git a/changelogs/v3.4.37.md b/changelogs/v3.4.37.md index 6899a480..e6b730cf 100644 --- a/changelogs/v3.4.37.md +++ b/changelogs/v3.4.37.md @@ -1,6 +1,6 @@ # What's Changed 1. ✨ 新增: 支持接入钉钉 #643 -2. ✨ 新增: 支持设置私聊是否需要唤醒前缀唤醒 [#735](https://github.com/Soulter/AstrBot/issues/735) +2. ✨ 新增: 支持设置私聊是否需要唤醒前缀唤醒 [#735](https://github.com/AstrBotDevs/AstrBot/issues/735) 3. 🐛 修复: 无法正常保存插件的 list 类型配置 #737 4. 🐛 修复: 部分情况下使用 aiocqhttp 报错 int 不能与 str 进行 '+' 操作的问题 \ No newline at end of file diff --git a/changelogs/v3.5.0.md b/changelogs/v3.5.0.md index e0b68bd4..9121f338 100644 --- a/changelogs/v3.5.0.md +++ b/changelogs/v3.5.0.md @@ -6,11 +6,11 @@ ## ✨ 新增的功能 1. ‼️ 新增支持接入 MCP 服务器 @Soulter @AraragiEro -1. ‼️ 新增支持本地渲染 Markdown，并支持自定义字体，详见 -> [#957](https://github.com/Soulter/AstrBot/issues/957#issuecomment-2749981802) +1. ‼️ 新增支持本地渲染 Markdown，并支持自定义字体，详见 -> [#957](https://github.com/AstrBotDevs/AstrBot/issues/957#issuecomment-2749981802) 2. 新增支持在 WebUI 管理所有与大模型的对话 -3. 适配完整的 function-calling 流程。[#804](https://github.com/Soulter/AstrBot/issues/804) [#566](https://github.com/Soulter/AstrBot/issues/566) +3. 适配完整的 function-calling 流程。[#804](https://github.com/AstrBotDevs/AstrBot/issues/804) [#566](https://github.com/AstrBotDevs/AstrBot/issues/566) 4. 新增支持消息平台热重载，不再需要重启 AstrBot -5. 新增支持阿里云百炼应用的 RAG 应用 [#878](https://github.com/Soulter/AstrBot/issues/878) +5. 新增支持阿里云百炼应用的 RAG 应用 [#878](https://github.com/AstrBotDevs/AstrBot/issues/878) 6. 新增 `/plugin get` OP 指令下载插件。如 `/plugin get Raven95676/astrbot_plugin_wordle` 7. 新增 `/newgroup` OP 指令，支持私聊 bot 给指定群聊创建新的对话。by @LunarMeal 8. Gewechat 下支持 `添加好友`, `接收/发送视频`, `获取群信息`, `接收/发送表情包` by @Moyuyanli @Soulter @XuYingJie-cmd @NiceAir @@ -20,10 +20,10 @@ ## 🎈 功能性优化 0. 更加美观的 WebUI 设计，降低疲劳程度。 -1. 微信下，忽略 `微信团队` 的消息 [#859](https://github.com/Soulter/AstrBot/issues/859) -2. 完善 Dify 的图片输入功能 [#893](https://github.com/Soulter/AstrBot/issues/893) +1. 微信下，忽略 `微信团队` 的消息 [#859](https://github.com/AstrBotDevs/AstrBot/issues/859) +2. 完善 Dify 的图片输入功能 [#893](https://github.com/AstrBotDevs/AstrBot/issues/893) 3. 消息平台和配置提供商配置页中，自动更新旧的配置项 -4. 优化钉钉在配置错误之后堵塞整个线程的问题 [#885](https://github.com/Soulter/AstrBot/issues/885) +4. 优化钉钉在配置错误之后堵塞整个线程的问题 [#885](https://github.com/AstrBotDevs/AstrBot/issues/885) 5. WebUI 删除插件时提供二次确认避免误删 @zhx8702 6. WebUI 优化新版本时的信息显示 7. 发送消息失败时的报错回显优化 @@ -33,13 +33,13 @@ ## 🐛 修复的 Bug 1. 修复图片没有被存储到聊天上下文历史记录 -2. 修复 Telegram 下无法识别图片描述(Caption) [#910](https://github.com/Soulter/AstrBot/issues/910) -3. 修复 Telegram Topic 群组下引用消息来源错误的问题 [#908](https://github.com/Soulter/AstrBot/issues/908) -4. 修复 Telegram 下 `/start` 指令的一些问题 [#751](https://github.com/Soulter/AstrBot/issues/751) -5. WebUI 插件市场卡片显示风格的过滤问题。[#927](https://github.com/Soulter/AstrBot/issues/927) -6. 统一 SSL 证书验证逻辑，修复 `SSLCertVerificationError` 的问题。by @IGCrystal [#950](https://github.com/Soulter/AstrBot/issues/950) +2. 修复 Telegram 下无法识别图片描述(Caption) [#910](https://github.com/AstrBotDevs/AstrBot/issues/910) +3. 修复 Telegram Topic 群组下引用消息来源错误的问题 [#908](https://github.com/AstrBotDevs/AstrBot/issues/908) +4. 修复 Telegram 下 `/start` 指令的一些问题 [#751](https://github.com/AstrBotDevs/AstrBot/issues/751) +5. WebUI 插件市场卡片显示风格的过滤问题。[#927](https://github.com/AstrBotDevs/AstrBot/issues/927) +6. 统一 SSL 证书验证逻辑，修复 `SSLCertVerificationError` 的问题。by @IGCrystal [#950](https://github.com/AstrBotDevs/AstrBot/issues/950) 7. 修复可能形成 SQL 注入的风险 -8. 修复本地上传插件时无法重载插件的问题 [#995](https://github.com/Soulter/AstrBot/issues/995) by @zhx8702 +8. 修复本地上传插件时无法重载插件的问题 [#995](https://github.com/AstrBotDevs/AstrBot/issues/995) by @zhx8702 ## 🧩 新增的插件 diff --git a/changelogs/v3.5.1.md b/changelogs/v3.5.1.md index 80ef10b4..fa97f4e8 100644 --- a/changelogs/v3.5.1.md +++ b/changelogs/v3.5.1.md @@ -4,16 +4,16 @@ ## ✨ 新增的功能 -1. 适配 `gemini-2.0-flash-exp-image-generation` 对图片模态的输入 [#1017](https://github.com/Soulter/AstrBot/issues/1017) +1. 适配 `gemini-2.0-flash-exp-image-generation` 对图片模态的输入 [#1017](https://github.com/AstrBotDevs/AstrBot/issues/1017) 2. 在 MessageChain 类中添加 at 和 at_all 方法，用于快速添加 At 消息 @left666 3. Gewechat Client 增加获取通讯录列表接口 -4. 支持 /llm 指令快捷启停 LLM 功能 [#296](https://github.com/Soulter/AstrBot/issues/296) +4. 支持 /llm 指令快捷启停 LLM 功能 [#296](https://github.com/AstrBotDevs/AstrBot/issues/296) ## 🎈 功能性优化 1. Edge TTS 支持使用代理 2. 在 Lifecycle 新增插件资源清理逻辑 @Raven95676 -3. Docker 镜像提供内置 FFmpeg [#979](https://github.com/Soulter/AstrBot/issues/979) +3. Docker 镜像提供内置 FFmpeg [#979](https://github.com/AstrBotDevs/AstrBot/issues/979) 4. 优化无对话情况下设置人格的反馈 @Raven95676 5. 若禁用提供商，自动切换到另一个可用的提供商 @Raven95676 6. openai_source 同步支持随机请求均衡，同时优化 LLM 请求逻辑的异常处理 diff --git a/changelogs/v3.5.12.md b/changelogs/v3.5.12.md index bdcf2941..b241ce84 100644 --- a/changelogs/v3.5.12.md +++ b/changelogs/v3.5.12.md @@ -1,7 +1,7 @@ # What's Changed -1. 新增：支持 MCP 的 Streamable HTTP 传输方式。详见 [#1637](https://github.com/Soulter/AstrBot/issues/1637) -2. 新增：支持 MCP 的 SSE 传输方式的自定义请求头。详见 [#1659](https://github.com/Soulter/AstrBot/issues/1659) +1. 新增：支持 MCP 的 Streamable HTTP 传输方式。详见 [#1637](https://github.com/AstrBotDevs/AstrBot/issues/1637) +2. 新增：支持 MCP 的 SSE 传输方式的自定义请求头。详见 [#1659](https://github.com/AstrBotDevs/AstrBot/issues/1659) 3. 优化：将 /llm 和 /model 和 /provider 指令设置为管理员指令 4. 修复：修复插件的 priority 部分失效的问题 5. 修复：修复 QQ 下合并转发消息内无法发送文件等问题，尽可能修复了各种文件、语音、视频、图片无法发送的问题 diff --git a/changelogs/v3.5.25.md b/changelogs/v3.5.25.md index 5576efaa..4d43ca7d 100644 --- a/changelogs/v3.5.25.md +++ b/changelogs/v3.5.25.md @@ -1,12 +1,12 @@ # What's Changed -1. 修复: 修复插件可能存在的无法正常禁用的问题 ([#2352](https://github.com/Soulter/AstrBot/issues/2352)) -2. ❗修复：当返回文本为空并且存在函数调用时错误地被终止事件，导致函数调用结果未被正常返回 ([#2491](https://github.com/Soulter/AstrBot/issues/2491)) -3. 修复：修复无法清空 AstrBot 配置下的 http_proxy 代理的问题 ([#2434](https://github.com/Soulter/AstrBot/issues/2434)) -4. ❗修复：Gemini 下开启流式输出时，持久化的消息结果不完整 ([#2424](https://github.com/Soulter/AstrBot/issues/2424)) -5. 修复：注册文件时由于 file:/// 前缀，导致文件被误判为不存在的问题 ([#2325](https://github.com/Soulter/AstrBot/issues/2325)) -6. 优化: 为部分类型供应商添加默认的温度选项 ([#2321](https://github.com/Soulter/AstrBot/issues/2321)) -7. 优化: 适配 Qwen3 模型非流式输出下需要传入 enable_think 参数（否则报错） ([#2424](https://github.com/Soulter/AstrBot/issues/2424)) +1. 修复: 修复插件可能存在的无法正常禁用的问题 ([#2352](https://github.com/AstrBotDevs/AstrBot/issues/2352)) +2. ❗修复：当返回文本为空并且存在函数调用时错误地被终止事件，导致函数调用结果未被正常返回 ([#2491](https://github.com/AstrBotDevs/AstrBot/issues/2491)) +3. 修复：修复无法清空 AstrBot 配置下的 http_proxy 代理的问题 ([#2434](https://github.com/AstrBotDevs/AstrBot/issues/2434)) +4. ❗修复：Gemini 下开启流式输出时，持久化的消息结果不完整 ([#2424](https://github.com/AstrBotDevs/AstrBot/issues/2424)) +5. 修复：注册文件时由于 file:/// 前缀，导致文件被误判为不存在的问题 ([#2325](https://github.com/AstrBotDevs/AstrBot/issues/2325)) +6. 优化: 为部分类型供应商添加默认的温度选项 ([#2321](https://github.com/AstrBotDevs/AstrBot/issues/2321)) +7. 优化: 适配 Qwen3 模型非流式输出下需要传入 enable_think 参数（否则报错） ([#2424](https://github.com/AstrBotDevs/AstrBot/issues/2424)) 8. 优化：支持配置工具调用轮数上限，默认 30 9. 新增: 添加 WebUI 语义化预发布版本提醒和检测功能 diff --git a/changelogs/v4.0.0-beta.5.md b/changelogs/v4.0.0-beta.5.md index 291f6e61..88f3815d 100644 --- a/changelogs/v4.0.0-beta.5.md +++ b/changelogs/v4.0.0-beta.5.md @@ -12,4 +12,4 @@ 6. 优化：增加 abconf_data 缓存，优化性能 7. 优化: 适配 qwen3 的 thinking 类模型 8. 优化: 完善对 rerank model 的可用性检测 -9. 新增: 给添加 edge_tts 新增 rate, volume, pitch 参数 ([#2625](https://github.com/Soulter/AstrBot/issues/2625)) +9. 新增: 给添加 edge_tts 新增 rate, volume, pitch 参数 ([#2625](https://github.com/AstrBotDevs/AstrBot/issues/2625)) diff --git a/changelogs/v4.1.0.md b/changelogs/v4.1.0.md index 9c60ca4a..3c1c8f2a 100644 --- a/changelogs/v4.1.0.md +++ b/changelogs/v4.1.0.md @@ -2,14 +2,14 @@ > 如果已经使用自定义文转图模板，此次升级之后将会被覆盖，请提前备份。路径在 `astrbot/core/utils/t2i/template` 目录下。 -0. ‼️‼️‼️ 修复 LLM 仍会调用已禁用的工具的问题 ([#2729](https://github.com/Soulter/AstrBot/issues/2729)) +0. ‼️‼️‼️ 修复 LLM 仍会调用已禁用的工具的问题 ([#2729](https://github.com/AstrBotDevs/AstrBot/issues/2729)) 1. ‼️ 修复 WebChat 下，Agent 长时任务时，SSE 连接自动断开的问题 -2. ‼️ 修复自定义文转图模板更新版本后会被覆盖的问题 ([#2677](https://github.com/Soulter/AstrBot/issues/2677)) -3. 修复 Satori 适配器教程链接 ([#2668](https://github.com/Soulter/AstrBot/issues/2668)) -4. 修复插件页表格视图中，点击状态字段表头排序不起作用的问题 ([#2714](https://github.com/Soulter/AstrBot/issues/2714)) -5. 修复工具调用时的 content 内容在重新加载后没有显示在 webchat 的问题 ([#2727](https://github.com/Soulter/AstrBot/issues/2727)) -6. 允许添加多个 tavily API Key 进行轮询 ([#2725](https://github.com/Soulter/AstrBot/issues/2725)) -7. 添加 --webui-dir 启动参数以支持指定 WebUI 构建文件目录 ([#2680](https://github.com/Soulter/AstrBot/issues/2680)) -8. 兼容指令名和第一个参数之间没有空格的情况 ([#2650](https://github.com/Soulter/AstrBot/issues/2650)) -9. 支持在 WebUI 自定义 OpenAI API extra_body 参数 ([#2719](https://github.com/Soulter/AstrBot/issues/2719)) -10. 增加 on_platform_loaded 钩子以在消息平台适配器实例化完成后触发 ([#2651](https://github.com/Soulter/AstrBot/issues/2651)) +2. ‼️ 修复自定义文转图模板更新版本后会被覆盖的问题 ([#2677](https://github.com/AstrBotDevs/AstrBot/issues/2677)) +3. 修复 Satori 适配器教程链接 ([#2668](https://github.com/AstrBotDevs/AstrBot/issues/2668)) +4. 修复插件页表格视图中，点击状态字段表头排序不起作用的问题 ([#2714](https://github.com/AstrBotDevs/AstrBot/issues/2714)) +5. 修复工具调用时的 content 内容在重新加载后没有显示在 webchat 的问题 ([#2727](https://github.com/AstrBotDevs/AstrBot/issues/2727)) +6. 允许添加多个 tavily API Key 进行轮询 ([#2725](https://github.com/AstrBotDevs/AstrBot/issues/2725)) +7. 添加 --webui-dir 启动参数以支持指定 WebUI 构建文件目录 ([#2680](https://github.com/AstrBotDevs/AstrBot/issues/2680)) +8. 兼容指令名和第一个参数之间没有空格的情况 ([#2650](https://github.com/AstrBotDevs/AstrBot/issues/2650)) +9. 支持在 WebUI 自定义 OpenAI API extra_body 参数 ([#2719](https://github.com/AstrBotDevs/AstrBot/issues/2719)) +10. 增加 on_platform_loaded 钩子以在消息平台适配器实例化完成后触发 ([#2651](https://github.com/AstrBotDevs/AstrBot/issues/2651)) diff --git a/changelogs/v4.1.3.md b/changelogs/v4.1.3.md index 5d78125e..267b43c7 100644 --- a/changelogs/v4.1.3.md +++ b/changelogs/v4.1.3.md @@ -1,8 +1,8 @@ # What's Changed -0. ‼️ fix: 修复 4.0.0 版本之后，配置默认 TTS 或者 STT 模型之后仍无法生效的问题 ([#2758](https://github.com/Soulter/AstrBot/issues/2758)) -1. ‼️ fix: 修复分段回复时，引用消息单独发送导致第一条消息内容为空的问题 ([#2757](https://github.com/Soulter/AstrBot/issues/2757)) -2. feat: 支持在 WebUI 复制提供商配置以简化操作 ([#2767](https://github.com/Soulter/AstrBot/issues/2767)) -3. fix: handle image value correctly for mcp BlobResourceContents ([#2753](https://github.com/Soulter/AstrBot/issues/2753)) -4. feat: 增加 QQ 群名称识别到 system prompt, 并提供相应的配置 ([#2770](https://github.com/Soulter/AstrBot/issues/2770)) +0. ‼️ fix: 修复 4.0.0 版本之后，配置默认 TTS 或者 STT 模型之后仍无法生效的问题 ([#2758](https://github.com/AstrBotDevs/AstrBot/issues/2758)) +1. ‼️ fix: 修复分段回复时，引用消息单独发送导致第一条消息内容为空的问题 ([#2757](https://github.com/AstrBotDevs/AstrBot/issues/2757)) +2. feat: 支持在 WebUI 复制提供商配置以简化操作 ([#2767](https://github.com/AstrBotDevs/AstrBot/issues/2767)) +3. fix: handle image value correctly for mcp BlobResourceContents ([#2753](https://github.com/AstrBotDevs/AstrBot/issues/2753)) +4. feat: 增加 QQ 群名称识别到 system prompt, 并提供相应的配置 ([#2770](https://github.com/AstrBotDevs/AstrBot/issues/2770)) 5. fix: parameter type/default handling in CommandFilter diff --git a/changelogs/v4.1.4.md b/changelogs/v4.1.4.md index 85bb7705..08a0b49a 100644 --- a/changelogs/v4.1.4.md +++ b/changelogs/v4.1.4.md @@ -1,10 +1,10 @@ # What's Changed -0. ‼️ fix: 修复 4.0.0 版本之后，配置默认 TTS 或者 STT 模型之后仍无法生效的问题 ([#2758](https://github.com/Soulter/AstrBot/issues/2758)) -1. ‼️ fix: 修复分段回复时，引用消息单独发送导致第一条消息内容为空的问题 ([#2757](https://github.com/Soulter/AstrBot/issues/2757)) -2. feat: 支持在 WebUI 复制提供商配置以简化操作 ([#2767](https://github.com/Soulter/AstrBot/issues/2767)) -3. fix: handle image value correctly for mcp BlobResourceContents ([#2753](https://github.com/Soulter/AstrBot/issues/2753)) -4. feat: 增加 QQ 群名称识别到 system prompt, 并提供相应的配置 ([#2770](https://github.com/Soulter/AstrBot/issues/2770)) +0. ‼️ fix: 修复 4.0.0 版本之后，配置默认 TTS 或者 STT 模型之后仍无法生效的问题 ([#2758](https://github.com/AstrBotDevs/AstrBot/issues/2758)) +1. ‼️ fix: 修复分段回复时，引用消息单独发送导致第一条消息内容为空的问题 ([#2757](https://github.com/AstrBotDevs/AstrBot/issues/2757)) +2. feat: 支持在 WebUI 复制提供商配置以简化操作 ([#2767](https://github.com/AstrBotDevs/AstrBot/issues/2767)) +3. fix: handle image value correctly for mcp BlobResourceContents ([#2753](https://github.com/AstrBotDevs/AstrBot/issues/2753)) +4. feat: 增加 QQ 群名称识别到 system prompt, 并提供相应的配置 ([#2770](https://github.com/AstrBotDevs/AstrBot/issues/2770)) 5. fix: 修复 4.1.3 的异常问题 **总之上个版本有很严重的 bug 赶快更新!** diff --git a/changelogs/v4.3.3.md b/changelogs/v4.3.3.md index 2299187a..53e403fc 100644 --- a/changelogs/v4.3.3.md +++ b/changelogs/v4.3.3.md @@ -1,12 +1,12 @@ # What's Changed -1. fix: 修复了代码执行器插件不能正确获得发送来文件的问题 ([#2970](https://github.com/Soulter/AstrBot/issues/2970)) -2. fix: 修改的 DeepSeek 默认 modalities，避免默认勾选图像导致的报错。 ([#2963](https://github.com/Soulter/AstrBot/issues/2963)) -3. fix: 事件钩子终止事件传播后不继续执行 ([#2989](https://github.com/Soulter/AstrBot/issues/2989)) +1. fix: 修复了代码执行器插件不能正确获得发送来文件的问题 ([#2970](https://github.com/AstrBotDevs/AstrBot/issues/2970)) +2. fix: 修改的 DeepSeek 默认 modalities，避免默认勾选图像导致的报错。 ([#2963](https://github.com/AstrBotDevs/AstrBot/issues/2963)) +3. fix: 事件钩子终止事件传播后不继续执行 ([#2989](https://github.com/AstrBotDevs/AstrBot/issues/2989)) 4. fix: 启动了 TTS 但未配置 TTS 模型时，At 和 Reply 发送人无效 -5. fix: 修复 session-management 中人格错误的显示为默认人格的问题 ([#3000](https://github.com/Soulter/AstrBot/issues/3000)) -6. fix: 修复了删除对话时，聊天增强中的记录未被清除，导致新对话中仍然出现之前的聊天记录。 ([#3002](https://github.com/Soulter/AstrBot/issues/3002)) -7. fix: 修复阿里云百炼平台 TTS 下接入 CosyVoice V2, Qwen TTS 生成报错的问题 ([#2964](https://github.com/Soulter/AstrBot/issues/2964)) -8. perf: 优化 SQLite 参数配置，对话和会话管理增加输入防抖机制 ([#2969](https://github.com/Soulter/AstrBot/issues/2969)) -9. feat: 在新对话中重用先前的对话人格设置 ([#3005](https://github.com/Soulter/AstrBot/issues/3005)) -10. feat: 从 WebUI 更新后清除浏览器缓存 ([#2958](https://github.com/Soulter/AstrBot/issues/2958)) +5. fix: 修复 session-management 中人格错误的显示为默认人格的问题 ([#3000](https://github.com/AstrBotDevs/AstrBot/issues/3000)) +6. fix: 修复了删除对话时，聊天增强中的记录未被清除，导致新对话中仍然出现之前的聊天记录。 ([#3002](https://github.com/AstrBotDevs/AstrBot/issues/3002)) +7. fix: 修复阿里云百炼平台 TTS 下接入 CosyVoice V2, Qwen TTS 生成报错的问题 ([#2964](https://github.com/AstrBotDevs/AstrBot/issues/2964)) +8. perf: 优化 SQLite 参数配置，对话和会话管理增加输入防抖机制 ([#2969](https://github.com/AstrBotDevs/AstrBot/issues/2969)) +9. feat: 在新对话中重用先前的对话人格设置 ([#3005](https://github.com/AstrBotDevs/AstrBot/issues/3005)) +10. feat: 从 WebUI 更新后清除浏览器缓存 ([#2958](https://github.com/AstrBotDevs/AstrBot/issues/2958)) diff --git a/changelogs/v4.5.0.md b/changelogs/v4.5.0.md new file mode 100644 index 00000000..a40fe6ca --- /dev/null +++ b/changelogs/v4.5.0.md @@ -0,0 +1,14 @@ +## What's Changed + +1. 修复：部分情况下，MCP、配置文件的代码编辑器一直显示 `loading...` 的问题(bump monaco-editor version to 0.54.0)。 +2. 新增：重构创建消息平台时的流程及一些 UI 优化 ([#3102](https://github.com/AstrBotDevs/AstrBot/issues/3102)) +3. 新增：全新的自带知识库功能。 +4. 新增：插件支持显示可读名称和 Logo。 +5. 修复： dashboard.enable 配置未生效。 +6. 新增：Misskey 适配器支持文件上传、投票内容感知功能和重构部分代码 ([#2986](https://github.com/AstrBotDevs/AstrBot/issues/2986)) +7. 新增：优化 Misskey 适配器的通知和聊天消息处理，改进 @用户提及逻辑 ([#3075](https://github.com/AstrBotDevs/AstrBot/issues/3075)) +8. 新增：QQ 官方机器人增加沙盒模式选项，让本地部署能跳过 IP 白名单验证 ([#3087](https://github.com/AstrBotDevs/AstrBot/issues/3087)) +9. 新增：Satori 添加对合并转发消息功能的支持 ([#3050](https://github.com/AstrBotDevs/AstrBot/issues/3050)) +10. 修复：人格预设对话的重复注入 ([#3088](https://github.com/AstrBotDevs/AstrBot/issues/3088)) +11. 新增：适配第三方 Gemini 思考片段过滤 ([#3139](https://github.com/AstrBotDevs/AstrBot/issues/3139)) +12. 重构：从主模块和依赖项中移除 Google 搜索引擎集成 ([#3154](https://github.com/AstrBotDevs/AstrBot/issues/3154)) diff --git a/changelogs/v4.5.1.md b/changelogs/v4.5.1.md new file mode 100644 index 00000000..94462b80 --- /dev/null +++ b/changelogs/v4.5.1.md @@ -0,0 +1,7 @@ +## What's Changed + +1. 修复：第一次启动时不再错误地弹出迁移提醒 +2. 新增：Xinference Rerank Provider, STT Provider +3. 新增: xAI Grok Live Search +4. 优化: 插件卡片左下角恢复 文档 按钮并新增 插件配置 按钮。 +5. 优化: 更好地适配 Class 方式注册 LLM Tool。 diff --git a/changelogs/v4.5.2.md b/changelogs/v4.5.2.md new file mode 100644 index 00000000..5945e548 --- /dev/null +++ b/changelogs/v4.5.2.md @@ -0,0 +1,8 @@ +## What's Changed + +1. 修复：>= Python 3.12 版本下可能导致 LLM Tool 注册错误的问题。 +2. 优化：更好地适配 Class 方式注册 LLM Tool 的场景。引入 `call` 方法。 +3. 新增：`ConversationManager` 类支持 `add_message_pair` 方法，简化对话消息的添加操作。 +4. 新增：增加对 Tool Parameters 的参数验证，确保工具参数符合 JSON Schema 标准。 +5. 新增：增加 LLM Message Schema 定义，提升消息结构的规范性和一致性。 +6. 新增：支持对 WebUI 的侧边栏模块进行自定义配置(入口在侧边栏下方的设置页中)。 diff --git a/changelogs/v4.5.3.md b/changelogs/v4.5.3.md new file mode 100644 index 00000000..1e15510a --- /dev/null +++ b/changelogs/v4.5.3.md @@ -0,0 +1,5 @@ +## What's Changed + +> hotfix version of 4.5.2 + +1. 修复：修正 `get_tool_list` 方法中工具字典推导式的错误导致的 WebUI MCP 页面工具列表无法显示的问题。 diff --git a/changelogs/v4.5.4.md b/changelogs/v4.5.4.md new file mode 100644 index 00000000..42e149cb --- /dev/null +++ b/changelogs/v4.5.4.md @@ -0,0 +1,5 @@ +## What's Changed + +1. 修复：Docker 镜像部分依赖问题导致某些情况下无法启动容器的问题； +2. 优化：插件卡片样式 +3. 修复：部分情况下 Windows 一键启动部署时，更新 / 部署失败的问题； diff --git a/changelogs/v4.5.5.md b/changelogs/v4.5.5.md new file mode 100644 index 00000000..9fda2b0e --- /dev/null +++ b/changelogs/v4.5.5.md @@ -0,0 +1,3 @@ +## What's Changed + +1. 修复：部署失败 diff --git a/changelogs/v4.5.6.md b/changelogs/v4.5.6.md new file mode 100644 index 00000000..51cbde60 --- /dev/null +++ b/changelogs/v4.5.6.md @@ -0,0 +1,3 @@ +## What's Changed + +1. 修复：构建失败 diff --git a/changelogs/v4.5.7.md b/changelogs/v4.5.7.md new file mode 100644 index 00000000..06317f57 --- /dev/null +++ b/changelogs/v4.5.7.md @@ -0,0 +1,12 @@ +## What's Changed + +1. 新增：支持为 OpenAI API 提供商自定义请求头 ([#3581](https://github.com/AstrBotDevs/AstrBot/issues/3581)) +2. 新增：为 WebChat 为 Thinking 模型添加思考过程展示功能；支持快捷切换流式输出 / 非流式输出。([#3632](https://github.com/AstrBotDevs/AstrBot/issues/3632)) +3. 新增：优化插件调用 LLM 和 Agent 的路径，为 Context 类引入多个调用 LLM 和 Agent 的便捷方法 ([#3636](https://github.com/AstrBotDevs/AstrBot/issues/3636)) +4. 优化：改善不支持流式输出的消息平台的回退策略 ([#3547](https://github.com/AstrBotDevs/AstrBot/issues/3547)) +5. 优化：当同一个会话（umo）下同时有多个请求时，执行排队处理，避免并发请求导致的上下文混乱问题 ([#3607](https://github.com/AstrBotDevs/AstrBot/issues/3607)) +6. 优化：优化 WebUI 的登录界面和 Changelog 页面的显示效果 +7. 修复：修复在知识库名字过长的情况下，“选择知识库”按钮显示异常的问题 ([#3582](https://github.com/AstrBotDevs/AstrBot/issues/3582)) +8. 修复：修复部分情况下，分段消息发送时导致的死锁问题（由 PR #3607 引入） +9. 修复：钉钉适配器使用部分指令无法生效的问题 ([#3634](https://github.com/AstrBotDevs/AstrBot/issues/3634)) +10. 其他：为部分适配器添加缺失的 send_streaming 方法 ([#3545](https://github.com/AstrBotDevs/AstrBot/issues/3545)) diff --git a/changelogs/v4.5.8.md b/changelogs/v4.5.8.md new file mode 100644 index 00000000..2f236462 --- /dev/null +++ b/changelogs/v4.5.8.md @@ -0,0 +1,5 @@ +## What's Changed + +hot fix of 4.5.7 + +fix: 无法正常发送图片，报错 `pydantic_core._pydantic_core.ValidationError` diff --git a/changelogs/v4.6.0.md b/changelogs/v4.6.0.md new file mode 100644 index 00000000..ca543990 --- /dev/null +++ b/changelogs/v4.6.0.md @@ -0,0 +1,23 @@ +## What's Changed + +1. 新增: 支持 gemini-3 系列的 thought signature ([#3698](https://github.com/AstrBotDevs/AstrBot/issues/3698)) +2. 新增: 支持知识库的 Agentic 检索功能 ([#3667](https://github.com/AstrBotDevs/AstrBot/issues/3667)) +3. 新增: 为知识库添加 URL 文档解析器 ([#3622](https://github.com/AstrBotDevs/AstrBot/issues/3622)) +4. 修复(core.platform): 修复启用多个企业微信智能机器人适配器时消息混乱的问题 ([#3693](https://github.com/AstrBotDevs/AstrBot/issues/3693)) +5. 修复: MCP Server 连接成功一段时间后，调用 mcp 工具时可能出现 `anyio.ClosedResourceError` 错误 ([#3700](https://github.com/AstrBotDevs/AstrBot/issues/3700)) +6. 新增(chat): 重构聊天组件结构并添加新功能 ([#3701](https://github.com/AstrBotDevs/AstrBot/issues/3701)) +7. 修复(dashboard.i18n): 完善缺失的英文国际化键值 ([#3699](https://github.com/AstrBotDevs/AstrBot/issues/3699)) +8. 重构: 实现 WebChat 会话管理及从版本 4.6 迁移到 4.7 +9. 持续集成(docker-build): 每日构建 Nightly 版本 Docker 镜像 ([#3120](https://github.com/AstrBotDevs/AstrBot/issues/3120)) + +--- + +1. feat: add supports for gemini-3 series thought signature ([#3698](https://github.com/AstrBotDevs/AstrBot/issues/3698)) +2. feat: supports knowledge base agentic search ([#3667](https://github.com/AstrBotDevs/AstrBot/issues/3667)) +3. feat: Add URL document parser for knowledge base ([#3622](https://github.com/AstrBotDevs/AstrBot/issues/3622)) +4. fix(core.platform): fix message mix-up issue when enabling multiple WeCom AI Bot adapters ([#3693](https://github.com/AstrBotDevs/AstrBot/issues/3693)) +5. fix: fix `anyio.ClosedResourceError` that may occur when calling mcp tools after a period of successful connection to MCP Server ([#3700](https://github.com/AstrBotDevs/AstrBot/issues/3700)) +6. feat(chat): refactor chat component structure and add new features ([#3701](https://github.com/AstrBotDevs/AstrBot/issues/3701)) +7. fix(dashboard.i18n): complete the missing i18n keys for en([#3699](https://github.com/AstrBotDevs/AstrBot/issues/3699)) +8. refactor: Implement WebChat session management and migration from version 4.6 to 4.7 +9. ci(docker-build): build nightly image everyday ([#3120](https://github.com/AstrBotDevs/AstrBot/issues/3120)) diff --git a/changelogs/v4.6.1.md b/changelogs/v4.6.1.md new file mode 100644 index 00000000..97c6f8a3 --- /dev/null +++ b/changelogs/v4.6.1.md @@ -0,0 +1,29 @@ +## What's Changed + +**hot fix of v4.6.0** + +fix(core.db): 修复升级后 webchat 相关对话数据未正确迁移的问题 ([#3745](https://github.com/AstrBotDevs/AstrBot/issues/3745)) + +--- + +1. 新增: 支持 gemini-3 系列的 thought signature ([#3698](https://github.com/AstrBotDevs/AstrBot/issues/3698)) +2. 新增: 支持知识库的 Agentic 检索功能 ([#3667](https://github.com/AstrBotDevs/AstrBot/issues/3667)) +3. 新增: 为知识库添加 URL 文档解析器 ([#3622](https://github.com/AstrBotDevs/AstrBot/issues/3622)) +4. 修复(core.platform): 修复启用多个企业微信智能机器人适配器时消息混乱的问题 ([#3693](https://github.com/AstrBotDevs/AstrBot/issues/3693)) +5. 修复: MCP Server 连接成功一段时间后，调用 mcp 工具时可能出现 `anyio.ClosedResourceError` 错误 ([#3700](https://github.com/AstrBotDevs/AstrBot/issues/3700)) +6. 新增(chat): 重构聊天组件结构并添加新功能 ([#3701](https://github.com/AstrBotDevs/AstrBot/issues/3701)) +7. 修复(dashboard.i18n): 完善缺失的英文国际化键值 ([#3699](https://github.com/AstrBotDevs/AstrBot/issues/3699)) +8. 重构: 实现 WebChat 会话管理及从版本 4.6 迁移到 4.7 +9. 持续集成(docker-build): 每日构建 Nightly 版本 Docker 镜像 ([#3120](https://github.com/AstrBotDevs/AstrBot/issues/3120)) + +--- + +1. feat: add supports for gemini-3 series thought signature ([#3698](https://github.com/AstrBotDevs/AstrBot/issues/3698)) +2. feat: supports knowledge base agentic search ([#3667](https://github.com/AstrBotDevs/AstrBot/issues/3667)) +3. feat: Add URL document parser for knowledge base ([#3622](https://github.com/AstrBotDevs/AstrBot/issues/3622)) +4. fix(core.platform): fix message mix-up issue when enabling multiple WeCom AI Bot adapters ([#3693](https://github.com/AstrBotDevs/AstrBot/issues/3693)) +5. fix: fix `anyio.ClosedResourceError` that may occur when calling mcp tools after a period of successful connection to MCP Server ([#3700](https://github.com/AstrBotDevs/AstrBot/issues/3700)) +6. feat(chat): refactor chat component structure and add new features ([#3701](https://github.com/AstrBotDevs/AstrBot/issues/3701)) +7. fix(dashboard.i18n): complete the missing i18n keys for en([#3699](https://github.com/AstrBotDevs/AstrBot/issues/3699)) +8. refactor: Implement WebChat session management and migration from version 4.6 to 4.7 +9. ci(docker-build): build nightly image everyday ([#3120](https://github.com/AstrBotDevs/AstrBot/issues/3120)) diff --git a/compose.yml b/compose.yml index 491c0cd7..2b318530 100644 --- a/compose.yml +++ b/compose.yml @@ -7,7 +7,7 @@ services: image: soulter/astrbot:latest container_name: astrbot restart: always - ports: # mappings description: https://github.com/Soulter/AstrBot/issues/497 + ports: # mappings description: https://github.com/AstrBotDevs/AstrBot/issues/497 - "6185:6185" # 必选，AstrBot WebUI 端口 - "6195:6195" # 可选, 企业微信 Webhook 端口 - "6199:6199" # 可选, QQ 个人号 WebSocket 端口 diff --git a/dashboard/.gitignore b/dashboard/.gitignore index 12ac6472..6e03962a 100644 --- a/dashboard/.gitignore +++ b/dashboard/.gitignore @@ -1,2 +1,3 @@ node_modules/ -.DS_Store \ No newline at end of file +.DS_Store +dist/ \ No newline at end of file diff --git a/dashboard/src/assets/images/icon-no-shadow.svg b/dashboard/src/assets/images/icon-no-shadow.svg new file mode 100644 index 00000000..4268e03e --- /dev/null +++ b/dashboard/src/assets/images/icon-no-shadow.svg @@ -0,0 +1 @@ +
 \ No newline at end of file diff --git a/dashboard/src/components/chat/Chat.vue b/dashboard/src/components/chat/Chat.vue index 5a63e5f4..09acd1b7 100644 --- a/dashboard/src/components/chat/Chat.vue +++ b/dashboard/src/components/chat/Chat.vue @@ -1,95 +1,46 @@ - - \ No newline at end of file + diff --git a/dashboard/src/components/chat/ChatInput.vue b/dashboard/src/components/chat/ChatInput.vue new file mode 100644 index 00000000..7ca0ec94 --- /dev/null +++ b/dashboard/src/components/chat/ChatInput.vue @@ -0,0 +1,283 @@ + + + + + diff --git a/dashboard/src/components/chat/ConversationSidebar.vue b/dashboard/src/components/chat/ConversationSidebar.vue new file mode 100644 index 00000000..5abc1bed --- /dev/null +++ b/dashboard/src/components/chat/ConversationSidebar.vue @@ -0,0 +1,296 @@ + + + + + + diff --git a/dashboard/src/components/chat/MessageList.vue b/dashboard/src/components/chat/MessageList.vue index 5832deae..15f3c1d3 100644 --- a/dashboard/src/components/chat/MessageList.vue +++ b/dashboard/src/components/chat/MessageList.vue @@ -33,33 +33,53 @@ - ✨ + mdi-star-four-points-small
- -
- - -
-
- + +
+ {{ tm('message.loading') }} +
+ +
-
+
@@ -125,7 +145,8 @@ export default { copiedMessages: new Set(), isUserNearBottom: true, scrollThreshold: 1, - scrollTimer: null + scrollTimer: null, + expandedReasoning: new Set(), // Track which reasoning blocks are expanded }; }, mounted() { @@ -142,6 +163,22 @@ export default { } }, methods: { + // Toggle reasoning expansion state + toggleReasoning(messageIndex) { + if (this.expandedReasoning.has(messageIndex)) { + this.expandedReasoning.delete(messageIndex); + } else { + this.expandedReasoning.add(messageIndex); + } + // Force reactivity + this.expandedReasoning = new Set(this.expandedReasoning); + }, + + // Check if reasoning is expanded + isReasoningExpanded(messageIndex) { + return this.expandedReasoning.has(messageIndex); + }, + // 复制代码到剪贴板 copyCodeToClipboard(code) { navigator.clipboard.writeText(code).then(() => { @@ -348,7 +385,7 @@ export default { @keyframes fadeIn { from { opacity: 0; - transform: translateY(10px); + transform: translateY(0); } to { @@ -539,6 +576,69 @@ export default { .fade-in { animation: fadeIn 0.3s ease-in-out; } + +/* Reasoning 区块样式 */ +.reasoning-container { + margin-bottom: 12px; + margin-top: 6px; + border: 1px solid var(--v-theme-border); + border-radius: 8px; + overflow: hidden; + width: fit-content; +} + +.v-theme--dark .reasoning-container { + background-color: rgba(103, 58, 183, 0.08); +} + +.reasoning-header { + display: inline-flex; + align-items: center; + padding: 8px 8px; + cursor: pointer; + user-select: none; + transition: background-color 0.2s ease; + border-radius: 8px; +} + +.reasoning-header:hover { + background-color: rgba(103, 58, 183, 0.08); +} + +.v-theme--dark .reasoning-header:hover { + background-color: rgba(103, 58, 183, 0.15); +} + +.reasoning-icon { + margin-right: 6px; + color: var(--v-theme-secondary); + transition: transform 0.2s ease; +} + +.reasoning-label { + font-size: 13px; + font-weight: 500; + color: var(--v-theme-secondary); + letter-spacing: 0.3px; +} + +.reasoning-content { + padding: 0px 12px; + border-top: 1px solid var(--v-theme-border); + color: gray; + animation: fadeIn 0.2s ease-in-out; + font-style: italic; +} + +.reasoning-text { + font-size: 14px; + line-height: 1.6; + color: var(--v-theme-secondaryText); +} + +.v-theme--dark .reasoning-text { + opacity: 0.85; +} \ No newline at end of file diff --git a/dashboard/src/components/shared/AstrBotConfig.vue b/dashboard/src/components/shared/AstrBotConfig.vue index e45327eb..d6c6fee9 100644 --- a/dashboard/src/components/shared/AstrBotConfig.vue +++ b/dashboard/src/components/shared/AstrBotConfig.vue @@ -7,6 +7,8 @@ import ProviderSelector from './ProviderSelector.vue' import PersonaSelector from './PersonaSelector.vue' import KnowledgeBaseSelector from './KnowledgeBaseSelector.vue' import { useI18n } from '@/i18n/composables' +import axios from 'axios' +import { useToast } from '@/utils/toast' const props = defineProps({ metadata: { @@ -40,6 +42,7 @@ const currentEditingKey = ref('') const currentEditingLanguage = ref('json') const currentEditingTheme = ref('vs-light') let currentEditingKeyIterable = null +const loadingEmbeddingDim = ref(false) function openEditorDialog(key, value, theme, language) { currentEditingKey.value = key @@ -49,10 +52,34 @@ function openEditorDialog(key, value, theme, language) { dialog.value = true } + function saveEditedContent() { dialog.value = false } +async function getEmbeddingDimensions(providerConfig) { + if (loadingEmbeddingDim.value) return + + loadingEmbeddingDim.value = true + try { + const response = await axios.post('/api/config/provider/get_embedding_dim', { + provider_config: providerConfig + }) + + if (response.data.status != "error" && response.data.data?.embedding_dimensions) { + console.log(response.data.data.embedding_dimensions) + providerConfig.embedding_dimensions = response.data.data.embedding_dimensions + useToast().success("获取成功: " + response.data.data.embedding_dimensions) + } else { + useToast().error(response.data.message) + } + } catch (error) { + console.error('Error getting embedding dimensions:', error) + } finally { + loadingEmbeddingDim.value = false + } +} + function getValueBySelector(obj, selector) { const keys = selector.split('.') let current = obj @@ -184,6 +211,29 @@ function hasVisibleItemsAfter(items, currentIndex) { v-model="iterable[key]" />
+ +
+ + + 自动检测 + +
+
diff --git a/dashboard/src/components/shared/AstrBotConfigV4.vue b/dashboard/src/components/shared/AstrBotConfigV4.vue index 52f793b9..77fe39eb 100644 --- a/dashboard/src/components/shared/AstrBotConfigV4.vue +++ b/dashboard/src/components/shared/AstrBotConfigV4.vue @@ -130,19 +130,6 @@ function hasVisibleItemsAfter(items, currentIndex) { return false } -// 将 options 和 labels 转换为 v-select 的 items 格式 -function getSelectItems(itemMeta) { - if (!itemMeta?.options) { - return [] - } - if (itemMeta?.labels && itemMeta.labels.length === itemMeta.options.length) { - return itemMeta.options.map((value, index) => ({ - title: itemMeta.labels[index], - value: value - })) - } - return itemMeta.options -}