From 96d1df8584ef86a0a98da04ccd60683be7e6ea68 Mon Sep 17 00:00:00 2001 From: whatevertogo Date: Thu, 19 Mar 2026 02:20:59 +0800 Subject: [PATCH 1/7] Merge commit 'e45bade147ff44b43860ecff12067309e59c151a' into feat/sdk-integration --- AGENTS.md | 10 +- CLAUDE.md | 12 +- src/astrbot_sdk/_testing_support.py | 82 ++- src/astrbot_sdk/clients/memory.py | 71 ++- src/astrbot_sdk/docs/01_context_api.md | 21 +- src/astrbot_sdk/docs/05_clients.md | 19 +- src/astrbot_sdk/docs/api/clients.md | 38 +- src/astrbot_sdk/docs/api/context.md | 28 +- src/astrbot_sdk/protocol/_builtin_schemas.py | 24 +- .../runtime/_capability_router_builtins.py | 582 +++++++++++++++++- src/astrbot_sdk/runtime/capability_router.py | 3 + src/astrbot_sdk/runtime/handler_dispatcher.py | 2 +- src/astrbot_sdk/star.py | 7 +- {tests_v4 => tests}/conftest.py | 0 .../test_context_register_task.py | 0 tests/test_memory_runtime.py | 277 +++++++++ 16 files changed, 1106 insertions(+), 70 deletions(-) rename {tests_v4 => tests}/conftest.py (100%) rename {tests_v4 => tests}/test_context_register_task.py (100%) create mode 100644 tests/test_memory_runtime.py diff --git a/AGENTS.md b/AGENTS.md index 09b79652a..3de989e63 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -41,12 +41,14 @@ ruff check . --fix # 使用 ruff 检查并自动修复全局格式问题 如果修改了内容可能影响现有功能，请运行测试以确保没有引入错误： 如果修改了bug或者更改了功能需要添加新的测试 +当前仓库已统一使用 `tests/` 目录，`tests_v4/` 不再作为新增测试入口。 +仓库当前没有 `run_tests.py`，请直接使用 `pytest`。 ```bash -python run_tests.py # 运行所有测试 -python run_tests.py -v # 详细输出 -python run_tests.py -k "test_peer" # 运行匹配模式的测试 -python run_tests.py --cov # 运行测试并生成覆盖率报告 +python -m pytest tests -q # 运行 tests 目录全部测试 +python -m pytest tests -v # 详细输出 +python -m pytest tests -k "test_context_register_task" # 运行匹配模式的测试 +python -m pytest tests --cov=astrbot_sdk # 运行测试并生成覆盖率报告 ``` ## 设计原则 diff --git a/CLAUDE.md b/CLAUDE.md index de010b89c..45efe08cd 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -41,12 +41,14 @@ ruff check . --fix # 使用 ruff 检查并自动修复全局格式问题 如果修改了内容可能影响现有功能，请运行测试以确保没有引入错误： 如果修改了bug或者更改了功能需要添加新的测试 +当前仓库已统一使用 `tests/` 目录，`tests_v4/` 不再作为新增测试入口。 +仓库当前没有 `run_tests.py`，请直接使用 `pytest`。 ```bash -python run_tests.py # 运行所有测试 -python run_tests.py -v # 详细输出 -python run_tests.py -k "test_peer" # 运行匹配模式的测试 -python run_tests.py --cov # 运行测试并生成覆盖率报告 +python -m pytest tests -q # 运行 tests 目录全部测试 +python -m pytest tests -v # 详细输出 +python -m pytest tests -k "test_context_register_task" # 运行匹配模式的测试 +python -m pytest tests --cov=astrbot_sdk # 运行测试并生成覆盖率报告 ``` ## 设计原则 @@ -57,4 +59,4 @@ python run_tests.py --cov # 运行测试并生成覆盖率报告 --- # currentDate -Today's date is 2026-03-14. +Today's date is 2026-03-19. diff --git a/src/astrbot_sdk/_testing_support.py b/src/astrbot_sdk/_testing_support.py index e6c562734..d1f09ab5f 100644 --- a/src/astrbot_sdk/_testing_support.py +++ b/src/astrbot_sdk/_testing_support.py @@ -6,6 +6,7 @@ import asyncio import typing from collections.abc import Mapping from dataclasses import dataclass, field +from datetime import datetime, timezone from typing import Any, TextIO from .context import CancelToken @@ -121,10 +122,77 @@ class InMemoryDB: class InMemoryMemory: - def __init__(self, store: dict[str, dict[str, Any]]) -> None: + def __init__( + self, + store: dict[str, dict[str, Any]], + *, + expires_at: dict[str, datetime | None] | None = None, + ) -> None: self._store = store + self._expires_at = expires_at if expires_at is not None else {} + + @staticmethod + def _is_ttl_entry(value: Any) -> bool: + """判断测试 memory 值是否使用 TTL 包装结构。 + + Args: + value: 待检查的存储值。 + + Returns: + bool: 如果包含 ``value`` 和 ``ttl_seconds`` 字段则返回 ``True``。 + """ + return isinstance(value, dict) and "value" in value and "ttl_seconds" in value + + @classmethod + def _search_text(cls, value: Any) -> str: + """提取测试用 memory.search 的匹配文本。 + + Args: + value: 当前存储的 memory 值。 + + Returns: + str: 用于本地测试搜索的文本内容。 + """ + if cls._is_ttl_entry(value): + value = value.get("value") + if not isinstance(value, dict): + return "" + for field_name in ("embedding_text", "content", "summary", "title", "text"): + item = value.get(field_name) + if isinstance(item, str) and item.strip(): + return item.strip() + return str(value) + + def _is_expired(self, key: str) -> bool: + """判断测试 memory 键是否已经过期。 + + Args: + key: memory 条目的键。 + + Returns: + bool: 如果当前时间已超过过期时间则返回 ``True``。 + """ + expires_at = self._expires_at.get(key) + return expires_at is not None and expires_at <= datetime.now(timezone.utc) + + def _purge_if_expired(self, key: str) -> bool: + """在测试 helper 中清理已过期的 memory 条目。 + + Args: + key: memory 条目的键。 + + Returns: + bool: 如果条目已过期并被清理则返回 ``True``。 + """ + if not self._is_expired(key): + return False + self._store.pop(key, None) + self._expires_at.pop(key, None) + return True def get(self, key: str, default: Any = None) -> Any: + if self._purge_if_expired(key): + return default return self._store.get(key, default) def save(self, key: str, value: dict[str, Any]) -> None: @@ -132,11 +200,14 @@ class InMemoryMemory: def delete(self, key: str) -> None: self._store.pop(key, None) + self._expires_at.pop(key, None) def search(self, query: str) -> list[dict[str, Any]]: results: list[dict[str, Any]] = [] - for key, value in self._store.items(): - if query in key or query in str(value): + for key, value in list(self._store.items()): + if self._purge_if_expired(key): + continue + if query in key or query in self._search_text(value): results.append({"key": key, "value": value}) return results @@ -200,7 +271,10 @@ class MockCapabilityRouter(CapabilityRouter): self._llm_stream_responses: list[str] = [] super().__init__() self.db = InMemoryDB(self.db_store) - self.memory = InMemoryMemory(self.memory_store) + self.memory = InMemoryMemory( + self.memory_store, + expires_at=self._memory_expires_at, + ) def list_dynamic_command_routes(self, plugin_id: str) -> list[dict[str, Any]]: return super().list_dynamic_command_routes(plugin_id) diff --git a/src/astrbot_sdk/clients/memory.py b/src/astrbot_sdk/clients/memory.py index e1c9d59ea..5fbaf5b60 100644 --- a/src/astrbot_sdk/clients/memory.py +++ b/src/astrbot_sdk/clients/memory.py @@ -15,7 +15,7 @@ from __future__ import annotations -from typing import Any +from typing import Any, Literal from ._proxy import CapabilityProxy @@ -48,30 +48,47 @@ class MemoryClient: """ self._proxy = proxy - async def search(self, query: str) -> list[dict[str, Any]]: - """Search memory items with the current bridge behavior. + async def search( + self, + query: str, + *, + mode: Literal["auto", "keyword", "vector", "hybrid"] = "auto", + limit: int | None = None, + min_score: float | None = None, + provider_id: str | None = None, + ) -> list[dict[str, Any]]: + """搜索记忆项。 - The current core bridge matches `query` against the memory key and the - serialized memory payload. It does not provide vector or semantic - retrieval yet. - - Returned items preserve the original `{"key": ..., "value": {...}}` - shape. When `value` is a mapping, its fields are also exposed at the - top level for compatibility with existing plugin examples. + 默认会在有 embedding provider 时执行 hybrid 检索， + 否则退化为关键词检索。返回结果包含 `score` 与 `match_type` 字段。 Args: query: 搜索查询文本 + mode: 搜索模式，支持 auto/keyword/vector/hybrid + limit: 最大返回条数 + min_score: 最低分数阈值 + provider_id: 指定 embedding provider，默认使用当前激活的 provider Returns: 匹配的记忆项列表，按相关度排序 示例: - # 搜索用户偏好相关的记忆 - results = await ctx.memory.search("用户喜欢什么颜色") + results = await ctx.memory.search( + "用户喜欢什么颜色", + mode="hybrid", + limit=5, + ) for item in results: - print(item["key"], item["content"]) + print(item["key"], item["score"], item["match_type"]) """ - output = await self._proxy.call("memory.search", {"query": query}) + payload: dict[str, Any] = {"query": query, "mode": mode} + if limit is not None: + payload["limit"] = limit + if min_score is not None: + payload["min_score"] = min_score + if provider_id is not None: + payload["provider_id"] = provider_id + output = await self._proxy.call("memory.search", payload) items = output.get("items") if not isinstance(items, (list, tuple)): return [] @@ -96,16 +113,20 @@ class MemoryClient: key: 记忆项的唯一标识键 value: 要存储的数据字典 **extra: 额外的键值对，会合并到 value 中 - Raises: TypeError: 如果 value 不是 dict 类型 - 示例: - # 保存用户偏好 + 保存用户偏好 await ctx.memory.save("user_pref", {"theme": "dark", "lang": "zh"}) - # 使用关键字参数 + 使用关键字参数 await ctx.memory.save("note", None, content="重要笔记", tags=["work"]) + + 使用 embedding_text 显式指定检索文本 + await ctx.memory.save( + "profile", + {"name": "alice", "embedding_text": "Alice 喜欢蓝色和海边"}, + ) """ if value is not None and not isinstance(value, dict): raise TypeError("memory.save 的 value 必须是 dict") @@ -230,16 +251,22 @@ class MemoryClient: async def stats(self) -> dict[str, Any]: """获取记忆系统统计信息。 - 返回记忆系统的当前状态，包括总条目数等统计信息。 + 返回记忆系统的当前状态，包括条目数、索引状态和脏索引数量。 Returns: 统计信息字典，包含： - total_items: 总记忆条目数 - total_bytes: 总占用字节数（可选） + - ttl_entries: 带过期时间的条目数（可选） + - indexed_items: 已建立检索索引的条目数（可选） + - embedded_items: 已生成向量的条目数（可选） + - dirty_items: 等待重建索引的条目数（可选） 示例: stats = await ctx.memory.stats() print(f"记忆库共有 {stats['total_items']} 条记录") + if "embedded_items" in stats: + print(f"其中 {stats['embedded_items']} 条已经向量化") """ output = await self._proxy.call("memory.stats", {}) stats = { @@ -250,4 +277,10 @@ class MemoryClient: stats["plugin_id"] = output.get("plugin_id") if "ttl_entries" in output: stats["ttl_entries"] = output.get("ttl_entries") + if "indexed_items" in output: + stats["indexed_items"] = output.get("indexed_items") + if "embedded_items" in output: + stats["embedded_items"] = output.get("embedded_items") + if "dirty_items" in output: + stats["dirty_items"] = output.get("dirty_items") return stats diff --git a/src/astrbot_sdk/docs/01_context_api.md b/src/astrbot_sdk/docs/01_context_api.md index 812456869..95a425262 100644 --- a/src/astrbot_sdk/docs/01_context_api.md +++ b/src/astrbot_sdk/docs/01_context_api.md @@ -159,12 +159,12 @@ async for chunk in ctx.llm.stream_chat("讲一个故事"): ### search() -语义搜索记忆项。 +搜索记忆项。默认在有 embedding provider 时执行 hybrid 检索。 ```python -results = await ctx.memory.search("用户喜欢什么颜色") +results = await ctx.memory.search("用户喜欢什么颜色", mode="hybrid", limit=5) for item in results: - print(item["key"], item["content"]) + print(item["key"], item["score"], item["match_type"]) ``` ### save() @@ -177,6 +177,12 @@ await ctx.memory.save("user_pref", {"theme": "dark", "lang": "zh"}) # 使用关键字参数 await ctx.memory.save("note", None, content="重要笔记", tags=["work"]) + +# 显式指定检索文本 +await ctx.memory.save( + "profile:alice", + {"name": "Alice", "embedding_text": "Alice 喜欢蓝色和海边"}, +) ``` ### get() @@ -202,6 +208,15 @@ await ctx.memory.save_with_ttl( ) ``` +### stats() + +查看记忆索引状态。 + +```python +stats = await ctx.memory.stats() +print(stats["total_items"], stats.get("embedded_items"), stats.get("dirty_items")) +``` + --- ## Database 客户端 diff --git a/src/astrbot_sdk/docs/05_clients.md b/src/astrbot_sdk/docs/05_clients.md index 7f49974ea..b5b30109c 100644 --- a/src/astrbot_sdk/docs/05_clients.md +++ b/src/astrbot_sdk/docs/05_clients.md @@ -66,10 +66,12 @@ from astrbot_sdk.clients import MemoryClient #### search() -语义搜索。 +搜索记忆。默认在有 embedding provider 时执行 hybrid 检索。 ```python -results = await ctx.memory.search("用户喜欢什么颜色") +results = await ctx.memory.search("用户喜欢什么颜色", mode="hybrid", limit=5) +for item in results: + print(item["key"], item["score"], item["match_type"]) ``` #### save() @@ -78,6 +80,10 @@ results = await ctx.memory.search("用户喜欢什么颜色") ```python await ctx.memory.save("user_pref", {"theme": "dark", "lang": "zh"}) +await ctx.memory.save( + "profile:alice", + {"name": "Alice", "embedding_text": "Alice 喜欢蓝色和海边"}, +) ``` #### get() @@ -108,6 +114,15 @@ await ctx.memory.save_with_ttl( await ctx.memory.delete("old_note") ``` +#### stats() + +查看记忆索引状态。 + +```python +stats = await ctx.memory.stats() +print(stats["total_items"], stats.get("embedded_items"), stats.get("dirty_items")) +``` + --- ## 3. DBClient - KV 数据库客户端 diff --git a/src/astrbot_sdk/docs/api/clients.md b/src/astrbot_sdk/docs/api/clients.md index 2e6ced7d1..455c6e7d0 100644 --- a/src/astrbot_sdk/docs/api/clients.md +++ b/src/astrbot_sdk/docs/api/clients.md @@ -142,25 +142,33 @@ from astrbot_sdk.clients import MemoryClient ### 方法 -#### `search(query)` +#### `search(query, *, mode="auto", limit=None, min_score=None, provider_id=None)` -语义搜索记忆项。 +搜索记忆项。默认会在存在 embedding provider 时执行 hybrid 检索， +否则退化为关键词检索。 **参数**: - `query` (`str`): 搜索查询文本（自然语言） +- `mode` (`Literal["auto", "keyword", "vector", "hybrid"]`): 搜索模式 +- `limit` (`int | None`): 最大返回条数 +- `min_score` (`float | None`): 最低分数阈值 +- `provider_id` (`str | None`): 指定 embedding provider -**返回**: `list[dict]` - 匹配的记忆项列表，按相关度排序 +**返回**: `list[dict]` - 匹配的记忆项列表。每项至少包含 `key`、`value`、`score`、`match_type` **示例**: ```python # 搜索用户偏好 -results = await ctx.memory.search("用户喜欢什么颜色") +results = await ctx.memory.search("用户喜欢什么颜色", mode="hybrid", limit=5) for item in results: - print(f"Key: {item['key']}, Content: {item['content']}") + print(item["key"], item["score"], item["match_type"]) -# 搜索对话摘要 -summaries = await ctx.memory.search("之前讨论过什么技术话题") +# 强制使用关键词检索 +keyword_hits = await ctx.memory.search("blue", mode="keyword", min_score=0.9) + +# 使用当前激活的 embedding provider 执行向量检索 +vector_hits = await ctx.memory.search("之前讨论过什么技术话题", mode="vector") ``` --- @@ -192,6 +200,16 @@ await ctx.memory.save( tags=["work"], timestamp="2024-01-01" ) + +# 显式指定检索文本 +await ctx.memory.save( + "profile:alice", + { + "name": "Alice", + "city": "Shanghai", + "embedding_text": "Alice 喜欢蓝色、海边和摄影", + }, +) ``` --- @@ -314,6 +332,12 @@ stats = await ctx.memory.stats() print(f"记忆库共有 {stats['total_items']} 条记录") if 'ttl_entries' in stats: print(f"其中 {stats['ttl_entries']} 条有过期时间") +if 'indexed_items' in stats: + print(f"已建立索引: {stats['indexed_items']}") +if 'embedded_items' in stats: + print(f"已向量化: {stats['embedded_items']}") +if 'dirty_items' in stats: + print(f"待重建索引: {stats['dirty_items']}") ``` --- diff --git a/src/astrbot_sdk/docs/api/context.md b/src/astrbot_sdk/docs/api/context.md index e76091602..eb91004b6 100644 --- a/src/astrbot_sdk/docs/api/context.md +++ b/src/astrbot_sdk/docs/api/context.md @@ -210,12 +210,16 @@ async for chunk in ctx.llm.stream_chat("讲一个故事"): ##### `search()` -语义搜索。 +搜索记忆。默认在有 embedding provider 时执行 hybrid 检索。 ```python -results = await ctx.memory.search("用户喜欢什么颜色") +results = await ctx.memory.search( + "用户喜欢什么颜色", + mode="hybrid", + limit=5, +) for item in results: - print(item["key"], item["content"]) + print(item["key"], item["score"], item["match_type"]) ``` ##### `save()` @@ -228,6 +232,15 @@ await ctx.memory.save("user_pref", {"theme": "dark", "lang": "zh"}) # 使用关键字参数 await ctx.memory.save("note", None, content="重要笔记", tags=["work"]) + +# 显式指定检索文本 +await ctx.memory.save( + "profile:alice", + { + "name": "Alice", + "embedding_text": "Alice 喜欢蓝色和海边", + }, +) ``` ##### `get()` @@ -261,6 +274,15 @@ await ctx.memory.save_with_ttl( await ctx.memory.delete("old_note") ``` +##### `stats()` + +查看记忆索引状态。 + +```python +stats = await ctx.memory.stats() +print(stats["total_items"], stats.get("embedded_items"), stats.get("dirty_items")) +``` + --- ### 3. DB 客户端 (ctx.db) diff --git a/src/astrbot_sdk/protocol/_builtin_schemas.py b/src/astrbot_sdk/protocol/_builtin_schemas.py index 790527cbc..80ade1d64 100644 --- a/src/astrbot_sdk/protocol/_builtin_schemas.py +++ b/src/astrbot_sdk/protocol/_builtin_schemas.py @@ -75,11 +75,28 @@ LLM_STREAM_CHAT_OUTPUT_SCHEMA = _object_schema( required=("text",), text={"type": "string"} ) MEMORY_SEARCH_INPUT_SCHEMA = _object_schema( - required=("query",), query={"type": "string"} + required=("query",), + query={"type": "string"}, + mode={"type": "string", "enum": ["auto", "keyword", "vector", "hybrid"]}, + limit={"type": "integer", "minimum": 1}, + min_score={"type": "number"}, + provider_id={"type": "string"}, ) MEMORY_SEARCH_OUTPUT_SCHEMA = _object_schema( required=("items",), - items={"type": "array", "items": {"type": "object"}}, + items={ + "type": "array", + "items": _object_schema( + required=("key", "value", "score", "match_type"), + key={"type": "string"}, + value=_nullable({"type": "object"}), + score={"type": "number"}, + match_type={ + "type": "string", + "enum": ["keyword", "vector", "hybrid"], + }, + ), + }, ) MEMORY_SAVE_INPUT_SCHEMA = _object_schema( required=("key", "value"), @@ -133,6 +150,9 @@ MEMORY_STATS_OUTPUT_SCHEMA = _object_schema( total_bytes=_nullable({"type": "integer"}), plugin_id=_nullable({"type": "string"}), ttl_entries=_nullable({"type": "integer"}), + indexed_items=_nullable({"type": "integer"}), + embedded_items=_nullable({"type": "integer"}), + dirty_items=_nullable({"type": "integer"}), ) SYSTEM_GET_DATA_DIR_INPUT_SCHEMA = _object_schema() SYSTEM_GET_DATA_DIR_OUTPUT_SCHEMA = _object_schema( diff --git a/src/astrbot_sdk/runtime/_capability_router_builtins.py b/src/astrbot_sdk/runtime/_capability_router_builtins.py index f01cb1ddb..e49ed7d10 100644 --- a/src/astrbot_sdk/runtime/_capability_router_builtins.py +++ b/src/astrbot_sdk/runtime/_capability_router_builtins.py @@ -20,10 +20,13 @@ from __future__ import annotations import asyncio import base64 import copy +import hashlib import json +import math +import re import uuid from collections.abc import AsyncIterator -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Any @@ -52,8 +55,61 @@ def _clone_chain_payload(value: Any) -> list[dict[str, Any]]: ] +_MOCK_EMBEDDING_DIM = 24 + + +def _embedding_terms(text: str) -> list[str]: + """为 mock embedding 构造稳定的分词结果。 + + Args: + text: 待向量化的原始文本。 + + Returns: + list[str]: 用于生成 mock 向量的词项列表。 + """ + normalized = re.sub(r"\s+", " ", str(text).strip().casefold()) + compact = normalized.replace(" ", "") + if not normalized: + return [] + + terms = [word for word in re.findall(r"\w+", normalized, flags=re.UNICODE) if word] + if compact: + if len(compact) == 1: + terms.append(compact) + else: + terms.extend( + compact[index : index + 2] for index in range(len(compact) - 1) + ) + terms.append(compact) + return terms or [normalized] + + +def _mock_embedding_vector(text: str, *, provider_id: str) -> list[float]: + """生成确定性的 mock embedding 向量。 + + Args: + text: 待向量化的文本。 + provider_id: 当前使用的 embedding provider 标识。 + + Returns: + list[float]: 归一化后的 mock 向量。 + """ + values = [0.0] * _MOCK_EMBEDDING_DIM + for term in _embedding_terms(text): + digest = hashlib.sha256(f"{provider_id}:{term}".encode("utf-8")).digest() + index = int.from_bytes(digest[:2], "big") % _MOCK_EMBEDDING_DIM + values[index] += 1.0 + min(len(term), 8) * 0.05 + norm = math.sqrt(sum(value * value for value in values)) + if norm <= 0: + return values + return [value / norm for value in values] + + class _CapabilityRouterHost: memory_store: dict[str, dict[str, Any]] + _memory_index: dict[str, dict[str, Any]] + _memory_dirty_keys: set[str] + _memory_expires_at: dict[str, datetime | None] db_store: dict[str, Any] sent_messages: list[dict[str, Any]] event_actions: list[dict[str, Any]] @@ -278,15 +334,471 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): }, ) + @staticmethod + def _is_ttl_memory_entry(value: Any) -> bool: + """判断存储值是否使用了 TTL 包装结构。 + + Args: + value: 待检查的存储值。 + + Returns: + bool: 如果值包含 ``value`` 和 ``ttl_seconds`` 字段则返回 ``True``。 + """ + return isinstance(value, dict) and "value" in value and "ttl_seconds" in value + + @classmethod + def _memory_value_for_search(cls, stored: Any) -> dict[str, Any] | None: + """提取用于检索的原始 memory payload。 + + Args: + stored: memory_store 中保存的原始值。 + + Returns: + dict[str, Any] | None: 解开 TTL 包装后的字典，无法解析时返回 ``None``。 + """ + if not isinstance(stored, dict): + return None + if cls._is_ttl_memory_entry(stored): + value = stored.get("value") + return value if isinstance(value, dict) else None + return stored + + @classmethod + def _extract_memory_text(cls, stored: Any) -> str: + """提取用于检索索引的首选文本。 + + Args: + stored: memory_store 中保存的原始值。 + + Returns: + str: 优先使用 ``embedding_text`` / ``content`` 等字段，兜底为 JSON 文本。 + """ + value = cls._memory_value_for_search(stored) + if not isinstance(value, dict): + return "" + for field_name in ("embedding_text", "content", "summary", "title", "text"): + item = value.get(field_name) + if isinstance(item, str) and item.strip(): + return item.strip() + return json.dumps(value, ensure_ascii=False, sort_keys=True, default=str) + + @staticmethod + def _memory_expiration_from_ttl(ttl_seconds: Any) -> datetime | None: + """将 TTL 秒数转换为 UTC 过期时间。 + + Args: + ttl_seconds: TTL 秒数。 + + Returns: + datetime | None: 绝对过期时间；当输入无效时返回 ``None``。 + """ + try: + ttl = int(ttl_seconds) + except (TypeError, ValueError): + return None + if ttl < 1: + return None + return datetime.now(timezone.utc) + timedelta(seconds=ttl) + + @staticmethod + def _memory_keyword_score(query: str, key: str, text: str) -> float: + """计算关键词匹配分数。 + + Args: + query: 查询文本。 + key: memory 条目的键。 + text: 已索引的检索文本。 + + Returns: + float: 基于键名和文本命中的粗粒度关键词分数。 + """ + normalized_query = str(query).casefold() + if not normalized_query: + return 1.0 + normalized_key = str(key).casefold() + normalized_text = str(text).casefold() + if normalized_query in normalized_key: + return 1.0 + if normalized_query in normalized_text: + return 0.9 + return 0.0 + + @staticmethod + def _cosine_similarity(left: list[float], right: list[float]) -> float: + """计算两个向量之间的余弦相似度。 + + Args: + left: 左侧向量。 + right: 右侧向量。 + + Returns: + float: 余弦相似度；输入不合法时返回 ``0.0``。 + """ + if not left or not right or len(left) != len(right): + return 0.0 + left_norm = math.sqrt(sum(value * value for value in left)) + right_norm = math.sqrt(sum(value * value for value in right)) + if left_norm <= 0 or right_norm <= 0: + return 0.0 + return sum(a * b for a, b in zip(left, right, strict=False)) / ( + left_norm * right_norm + ) + + def _resolve_memory_embedding_provider_id( + self, + provider_id: Any, + *, + required: bool, + ) -> str | None: + """解析 memory.search 要使用的 embedding provider。 + + Args: + provider_id: 调用方显式传入的 provider 标识。 + required: 当前检索模式是否强制要求 embedding provider。 + + Returns: + str | None: 最终选中的 provider 标识；在非强制场景下允许返回 ``None``。 + """ + normalized = str(provider_id).strip() if provider_id is not None else "" + if normalized: + self._provider_entry( + {"provider_id": normalized}, + "memory.search", + "embedding", + ) + return normalized + active_id = self._active_provider_ids.get("embedding") + if active_id is not None: + normalized_active = str(active_id).strip() + if normalized_active: + self._provider_entry( + {"provider_id": normalized_active}, + "memory.search", + "embedding", + ) + return normalized_active + if required: + raise AstrBotError.invalid_input( + "memory.search requires an embedding provider", + ) + return None + + @staticmethod + def _memory_index_entry(entry: Any, *, text: str) -> dict[str, Any]: + """将原始索引项规范化为内部统一结构。 + + Args: + entry: 当前索引表中的原始项。 + text: 当前条目的索引文本。 + + Returns: + dict[str, Any]: 统一后的索引项，包含 ``text``、``embedding``、``provider_id``。 + """ + if isinstance(entry, dict): + return { + "text": str(entry.get("text", text)), + "embedding": ( + [float(item) for item in entry.get("embedding", [])] + if isinstance(entry.get("embedding"), list) + else None + ), + "provider_id": ( + str(entry.get("provider_id")).strip() + if entry.get("provider_id") is not None + else None + ), + } + return {"text": text, "embedding": None, "provider_id": None} + + def _clear_memory_sidecars(self, key: str) -> None: + """清理指定 memory 键对应的所有 sidecar 状态。 + + Args: + key: memory 条目的键。 + + Returns: + None + """ + self._memory_index.pop(key, None) + self._memory_expires_at.pop(key, None) + self._memory_dirty_keys.discard(key) + + def _delete_memory_entry(self, key: str) -> bool: + """删除 memory 条目并同步清理 sidecar 状态。 + + Args: + key: memory 条目的键。 + + Returns: + bool: 条目存在并删除成功时返回 ``True``。 + """ + deleted = self.memory_store.pop(key, None) is not None + self._clear_memory_sidecars(key) + return deleted + + def _upsert_memory_sidecars( + self, + key: str, + stored: dict[str, Any], + *, + expires_at: datetime | None = None, + ) -> None: + """创建或更新单条 memory 的 sidecar 索引状态。 + + Args: + key: memory 条目的键。 + stored: 需要建立索引的原始存储值。 + expires_at: 可选的绝对过期时间。 + + Returns: + None + """ + self._memory_index[key] = { + "text": self._extract_memory_text(stored), + "embedding": None, + "provider_id": None, + } + if expires_at is None: + self._memory_expires_at.pop(key, None) + else: + self._memory_expires_at[key] = expires_at + self._memory_dirty_keys.add(key) + + def _ensure_memory_sidecars(self, key: str, stored: Any) -> None: + """确保 sidecar 状态与当前存储值保持一致。 + + Args: + key: memory 条目的键。 + stored: memory_store 中的当前存储值。 + + Returns: + None + """ + if not isinstance(stored, dict): + return + text = self._extract_memory_text(stored) + existed = key in self._memory_index + entry = self._memory_index_entry(self._memory_index.get(key), text=text) + if entry["text"] != text: + entry["text"] = text + entry["embedding"] = None + entry["provider_id"] = None + self._memory_dirty_keys.add(key) + self._memory_index[key] = entry + if not existed: + self._memory_dirty_keys.add(key) + + def _is_memory_expired(self, key: str) -> bool: + """判断 memory 条目是否已过期。 + + Args: + key: memory 条目的键。 + + Returns: + bool: 如果当前时间已超过记录的过期时间则返回 ``True``。 + """ + expires_at = self._memory_expires_at.get(key) + return expires_at is not None and expires_at <= datetime.now(timezone.utc) + + def _purge_expired_memory_entry(self, key: str) -> bool: + """在单条 memory 已过期时立即清理它。 + + Args: + key: memory 条目的键。 + + Returns: + bool: 如果条目已过期并被成功清理则返回 ``True``。 + """ + if not self._is_memory_expired(key): + return False + self._delete_memory_entry(key) + return True + + def _purge_expired_memory_entries(self) -> None: + """批量清理所有已跟踪的过期 TTL 条目。 + + Returns: + None + """ + for key in list(self._memory_expires_at): + self._purge_expired_memory_entry(key) + + async def _embedding_for_text( + self, + *, + provider_id: str, + text: str, + ) -> list[float]: + """通过 embedding capability 获取单条文本向量。 + + Args: + provider_id: 使用的 embedding provider 标识。 + text: 待向量化的文本。 + + Returns: + list[float]: provider 返回的向量；异常场景下返回空列表。 + """ + output = await self._provider_embedding_get_embedding( + "", + {"provider_id": provider_id, "text": text}, + None, + ) + embedding = output.get("embedding") + if not isinstance(embedding, list): + return [] + return [float(item) for item in embedding] + + async def _embeddings_for_texts( + self, + *, + provider_id: str, + texts: list[str], + ) -> list[list[float]]: + """批量获取多条文本的 embedding 向量。 + + Args: + provider_id: 使用的 embedding provider 标识。 + texts: 待向量化的文本列表。 + + Returns: + list[list[float]]: 与输入顺序对应的向量列表。 + """ + if not texts: + return [] + output = await self._provider_embedding_get_embeddings( + "", + {"provider_id": provider_id, "texts": texts}, + None, + ) + embeddings = output.get("embeddings") + if not isinstance(embeddings, list): + return [] + return [ + [float(value) for value in item] + for item in embeddings + if isinstance(item, list) + ] + + async def _refresh_memory_embeddings(self, *, provider_id: str) -> None: + """刷新当前 provider 下脏或过期的 memory 向量索引。 + + Args: + provider_id: 当前使用的 embedding provider 标识。 + + Returns: + None + """ + keys_to_refresh: list[str] = [] + texts_to_refresh: list[str] = [] + for key, stored in self.memory_store.items(): + self._ensure_memory_sidecars(key, stored) + entry = self._memory_index_entry( + self._memory_index.get(key), + text=self._extract_memory_text(stored), + ) + should_refresh = ( + key in self._memory_dirty_keys + or entry["embedding"] is None + or entry["provider_id"] != provider_id + ) + self._memory_index[key] = entry + if should_refresh: + keys_to_refresh.append(key) + texts_to_refresh.append(str(entry["text"])) + embeddings = await self._embeddings_for_texts( + provider_id=provider_id, + texts=texts_to_refresh, + ) + for index, key in enumerate(keys_to_refresh): + entry = self._memory_index_entry( + self._memory_index.get(key), + text=str(texts_to_refresh[index]), + ) + entry["embedding"] = embeddings[index] if index < len(embeddings) else [] + entry["provider_id"] = provider_id + self._memory_index[key] = entry + self._memory_dirty_keys.discard(key) + async def _memory_search( self, _request_id: str, payload: dict[str, Any], _token ) -> dict[str, Any]: query = str(payload.get("query", "")) - items = [ - {"key": key, "value": value} - for key, value in self.memory_store.items() - if query in key or query in json.dumps(value, ensure_ascii=False) - ] + mode = str(payload.get("mode", "auto")).strip().lower() or "auto" + limit = self._optional_int(payload.get("limit")) + min_score = ( + float(payload.get("min_score")) + if payload.get("min_score") is not None + else None + ) + self._purge_expired_memory_entries() + provider_id = self._resolve_memory_embedding_provider_id( + payload.get("provider_id"), + required=mode in {"vector", "hybrid"}, + ) + effective_mode = mode + if effective_mode == "auto": + effective_mode = "hybrid" if provider_id is not None else "keyword" + query_embedding: list[float] | None = None + if effective_mode in {"vector", "hybrid"}: + if provider_id is None: + raise AstrBotError.invalid_input( + "memory.search requires an embedding provider", + ) + await self._refresh_memory_embeddings(provider_id=provider_id) + query_embedding = await self._embedding_for_text( + provider_id=provider_id, + text=query, + ) + + items: list[dict[str, Any]] = [] + for key, value in self.memory_store.items(): + self._ensure_memory_sidecars(key, value) + entry = self._memory_index_entry( + self._memory_index.get(key), + text=self._extract_memory_text(value), + ) + text = str(entry.get("text", "")) + keyword_score = self._memory_keyword_score(query, key, text) + vector_score = 0.0 + if query_embedding is not None: + embedding = entry.get("embedding") + if isinstance(embedding, list): + vector_score = max( + 0.0, + self._cosine_similarity(query_embedding, embedding), + ) + + if effective_mode == "keyword": + score = keyword_score + elif effective_mode == "vector": + score = vector_score + else: + score = vector_score + if keyword_score > 0: + score = max(score, 0.4 + 0.6 * vector_score) + if score <= 0: + continue + if min_score is not None and score < min_score: + continue + + if effective_mode == "keyword" or (keyword_score > 0 and vector_score <= 0): + match_type = "keyword" + elif effective_mode == "vector" or keyword_score <= 0: + match_type = "vector" + else: + match_type = "hybrid" + + items.append( + { + "key": key, + "value": self._memory_value_for_search(value), + "score": score, + "match_type": match_type, + } + ) + items.sort(key=lambda item: (-float(item["score"]), str(item["key"]))) + if limit is not None and limit >= 0: + items = items[:limit] return {"items": items} async def _memory_save( @@ -297,17 +809,21 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): if not isinstance(value, dict): raise AstrBotError.invalid_input("memory.save 的 value 必须是 object") self.memory_store[key] = value + self._upsert_memory_sidecars(key, value) return {} async def _memory_get( self, _request_id: str, payload: dict[str, Any], _token ) -> dict[str, Any]: - return {"value": self.memory_store.get(str(payload.get("key", "")))} + key = str(payload.get("key", "")) + if self._purge_expired_memory_entry(key): + return {"value": None} + return {"value": self.memory_store.get(key)} async def _memory_delete( self, _request_id: str, payload: dict[str, Any], _token ) -> dict[str, Any]: - self.memory_store.pop(str(payload.get("key", "")), None) + self._delete_memory_entry(str(payload.get("key", ""))) return {} async def _memory_save_with_ttl( @@ -320,7 +836,13 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): raise AstrBotError.invalid_input( "memory.save_with_ttl 的 value 必须是 object" ) - self.memory_store[key] = {"value": value, "ttl_seconds": ttl_seconds} + stored = {"value": value, "ttl_seconds": ttl_seconds} + self.memory_store[key] = stored + self._upsert_memory_sidecars( + key, + stored, + expires_at=self._memory_expiration_from_ttl(ttl_seconds), + ) return {} async def _memory_get_many( @@ -332,6 +854,9 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): keys = [str(item) for item in keys_payload] items = [] for key in keys: + if self._purge_expired_memory_entry(key): + items.append({"key": key, "value": None}) + continue stored = self.memory_store.get(key) if ( isinstance(stored, dict) @@ -353,28 +878,36 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): keys = [str(item) for item in keys_payload] deleted_count = 0 for key in keys: - if key in self.memory_store: - del self.memory_store[key] + if self._delete_memory_entry(key): deleted_count += 1 return {"deleted_count": deleted_count} async def _memory_stats( self, _request_id: str, _payload: dict[str, Any], _token ) -> dict[str, Any]: + self._purge_expired_memory_entries() total_items = len(self.memory_store) total_bytes = sum( len(str(key)) + len(str(value)) for key, value in self.memory_store.items() ) - ttl_entries = sum( + ttl_entries = len(self._memory_expires_at) + indexed_items = len(self._memory_index) + embedded_items = sum( 1 - for value in self.memory_store.values() - if isinstance(value, dict) and "value" in value and "ttl_seconds" in value + for entry in self._memory_index.values() + if isinstance(entry, dict) + and isinstance(entry.get("embedding"), list) + and bool(entry.get("embedding")) ) + dirty_items = len(self._memory_dirty_keys) return { "total_items": total_items, "total_bytes": total_bytes, "plugin_id": self._require_caller_plugin_id("memory.stats"), "ttl_entries": ttl_entries, + "indexed_items": indexed_items, + "embedded_items": embedded_items, + "dirty_items": dirty_items, } def _register_memory_capabilities(self) -> None: @@ -1072,17 +1605,22 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): async def _provider_embedding_get_embedding( self, _request_id: str, payload: dict[str, Any], _token ) -> dict[str, Any]: - self._provider_entry( + provider = self._provider_entry( payload, "provider.embedding.get_embedding", "embedding", ) - return {"embedding": [0.0, 0.0, 0.0]} + return { + "embedding": _mock_embedding_vector( + str(payload.get("text", "")), + provider_id=str(provider.get("id", "")), + ) + } async def _provider_embedding_get_embeddings( self, _request_id: str, payload: dict[str, Any], _token ) -> dict[str, Any]: - self._provider_entry( + provider = self._provider_entry( payload, "provider.embedding.get_embeddings", "embedding", @@ -1093,7 +1631,13 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): "provider.embedding.get_embeddings requires texts", ) return { - "embeddings": [[0.0, 0.0, 0.0] for _ in texts], + "embeddings": [ + _mock_embedding_vector( + str(text), + provider_id=str(provider.get("id", "")), + ) + for text in texts + ], } async def _provider_embedding_get_dim( @@ -1104,7 +1648,7 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): "provider.embedding.get_dim", "embedding", ) - return {"dim": 3} + return {"dim": _MOCK_EMBEDDING_DIM} async def _provider_rerank_rerank( self, _request_id: str, payload: dict[str, Any], _token diff --git a/src/astrbot_sdk/runtime/capability_router.py b/src/astrbot_sdk/runtime/capability_router.py index 0d73d7ba7..73bf5557c 100644 --- a/src/astrbot_sdk/runtime/capability_router.py +++ b/src/astrbot_sdk/runtime/capability_router.py @@ -217,6 +217,9 @@ class CapabilityRouter(BuiltinCapabilityRouterMixin): self._registrations: dict[str, _CapabilityRegistration] = {} self.db_store: dict[str, Any] = {} self.memory_store: dict[str, dict[str, Any]] = {} + self._memory_index: dict[str, dict[str, Any]] = {} + self._memory_dirty_keys: set[str] = set() + self._memory_expires_at: dict[str, datetime | None] = {} self.sent_messages: list[dict[str, Any]] = [] self.event_actions: list[dict[str, Any]] = [] self._event_streams: dict[str, dict[str, Any]] = {} diff --git a/src/astrbot_sdk/runtime/handler_dispatcher.py b/src/astrbot_sdk/runtime/handler_dispatcher.py index 1a52c4a9d..2b32d9cfd 100644 --- a/src/astrbot_sdk/runtime/handler_dispatcher.py +++ b/src/astrbot_sdk/runtime/handler_dispatcher.py @@ -838,7 +838,7 @@ class HandlerDispatcher: if inspect.isawaitable(result): await result return - await Star().on_error(exc, event, ctx) + await Star.default_on_error(exc, event, ctx) __all__ = ["CapabilityDispatcher", "HandlerDispatcher"] diff --git a/src/astrbot_sdk/star.py b/src/astrbot_sdk/star.py index aef7eb09e..ef774b4e7 100644 --- a/src/astrbot_sdk/star.py +++ b/src/astrbot_sdk/star.py @@ -102,7 +102,9 @@ class Star(PluginKVStoreMixin): options=options, ) - async def on_error(self, error: Exception, event, ctx) -> None: + @staticmethod + async def default_on_error(error: Exception, event, ctx) -> None: + del ctx if isinstance(error, AstrBotError): lines: list[str] = [] if error.retryable: @@ -122,6 +124,9 @@ class Star(PluginKVStoreMixin): await event.reply("出了点问题，请联系插件作者") logger.error("handler 执行失败\n{}", traceback.format_exc()) + async def on_error(self, error: Exception, event, ctx) -> None: + await Star.default_on_error(error, event, ctx) + @classmethod def __astrbot_is_new_star__(cls) -> bool: return True diff --git a/tests_v4/conftest.py b/tests/conftest.py similarity index 100% rename from tests_v4/conftest.py rename to tests/conftest.py diff --git a/tests_v4/test_context_register_task.py b/tests/test_context_register_task.py similarity index 100% rename from tests_v4/test_context_register_task.py rename to tests/test_context_register_task.py diff --git a/tests/test_memory_runtime.py b/tests/test_memory_runtime.py new file mode 100644 index 000000000..f1b35509f --- /dev/null +++ b/tests/test_memory_runtime.py @@ -0,0 +1,277 @@ +from __future__ import annotations + +from datetime import datetime, timedelta, timezone + +import pytest + +from astrbot_sdk._invocation_context import caller_plugin_scope +from astrbot_sdk.runtime.capability_router import CapabilityRouter + + +async def _call( + router: CapabilityRouter, + capability: str, + payload: dict[str, object], +) -> dict[str, object]: + result = await router.execute( + capability, + payload, + stream=False, + cancel_token=object(), + request_id=f"test-{capability}", + ) + assert isinstance(result, dict) + return result + + +@pytest.mark.asyncio +async def test_memory_save_updates_sidecars_and_search() -> None: + router = CapabilityRouter() + + await _call( + router, + "memory.save", + {"key": "user-pref", "value": {"content": "user likes blue"}}, + ) + + assert router.memory_store["user-pref"] == {"content": "user likes blue"} + assert router._memory_index["user-pref"] == { + "text": "user likes blue", + "embedding": None, + "provider_id": None, + } + assert "user-pref" in router._memory_dirty_keys + assert "user-pref" not in router._memory_expires_at + + result = await _call(router, "memory.search", {"query": "likes blue"}) + assert len(result["items"]) == 1 + item = result["items"][0] + assert item["key"] == "user-pref" + assert item["value"] == {"content": "user likes blue"} + assert item["match_type"] == "hybrid" + assert float(item["score"]) > 0 + assert router._memory_index["user-pref"]["provider_id"] == "mock-embedding-provider" + assert isinstance(router._memory_index["user-pref"]["embedding"], list) + assert "user-pref" not in router._memory_dirty_keys + + +@pytest.mark.asyncio +async def test_memory_search_keyword_mode_keeps_dirty_embedding_state() -> None: + router = CapabilityRouter() + + await _call( + router, + "memory.save", + {"key": "alpha-key", "value": {"content": "blue ocean memory"}}, + ) + + result = await _call( + router, + "memory.search", + {"query": "alpha", "mode": "keyword", "min_score": 0.95}, + ) + + assert [item["key"] for item in result["items"]] == ["alpha-key"] + assert result["items"][0]["match_type"] == "keyword" + assert router._memory_index["alpha-key"]["embedding"] is None + assert "alpha-key" in router._memory_dirty_keys + + +@pytest.mark.asyncio +async def test_memory_search_vector_mode_supports_ranking_and_limit() -> None: + router = CapabilityRouter() + + await _call( + router, + "memory.save", + {"key": "fruit-note", "value": {"content": "banana smoothie with mango"}}, + ) + await _call( + router, + "memory.save", + {"key": "ocean-note", "value": {"content": "waves on the blue ocean"}}, + ) + + result = await _call( + router, + "memory.search", + {"query": "banana smoothie", "mode": "vector", "limit": 1}, + ) + + assert len(result["items"]) == 1 + assert result["items"][0]["key"] == "fruit-note" + assert result["items"][0]["match_type"] == "vector" + + +@pytest.mark.asyncio +async def test_memory_search_auto_falls_back_to_keyword_without_embedding_provider() -> ( + None +): + router = CapabilityRouter() + router._active_provider_ids["embedding"] = None + + await _call( + router, + "memory.save", + {"key": "alpha-key", "value": {"content": "blue ocean memory"}}, + ) + + result = await _call(router, "memory.search", {"query": "alpha", "mode": "auto"}) + + assert [item["key"] for item in result["items"]] == ["alpha-key"] + assert result["items"][0]["match_type"] == "keyword" + assert router._memory_index["alpha-key"]["embedding"] is None + assert "alpha-key" in router._memory_dirty_keys + + +@pytest.mark.asyncio +async def test_memory_search_reembeds_when_embedding_provider_changes() -> None: + router = CapabilityRouter() + router._provider_catalog["embedding"].append( + { + "id": "mock-embedding-provider-alt", + "model": "mock-embedding-model-alt", + "type": "mock", + "provider_type": "embedding", + } + ) + router._provider_configs["mock-embedding-provider-alt"] = { + "id": "mock-embedding-provider-alt", + "model": "mock-embedding-model-alt", + "type": "mock", + "provider_type": "embedding", + "enable": True, + } + + await _call( + router, + "memory.save", + {"key": "topic", "value": {"content": "banana smoothie with mango"}}, + ) + + first = await _call(router, "memory.search", {"query": "banana smoothie"}) + first_embedding = list(router._memory_index["topic"]["embedding"]) + assert first["items"][0]["match_type"] == "hybrid" + assert router._memory_index["topic"]["provider_id"] == "mock-embedding-provider" + + router._active_provider_ids["embedding"] = "mock-embedding-provider-alt" + + second = await _call(router, "memory.search", {"query": "banana smoothie"}) + second_embedding = list(router._memory_index["topic"]["embedding"]) + assert second["items"][0]["match_type"] == "hybrid" + assert router._memory_index["topic"]["provider_id"] == "mock-embedding-provider-alt" + assert first_embedding != second_embedding + + +@pytest.mark.asyncio +async def test_memory_stats_reports_index_embedding_and_dirty_counts() -> None: + router = CapabilityRouter() + + await _call( + router, + "memory.save", + {"key": "a", "value": {"content": "alpha"}}, + ) + await _call( + router, + "memory.save_with_ttl", + {"key": "b", "value": {"content": "beta"}, "ttl_seconds": 60}, + ) + + with caller_plugin_scope("test-plugin"): + before = await _call(router, "memory.stats", {}) + assert before["total_items"] == 2 + assert before["ttl_entries"] == 1 + assert before["indexed_items"] == 2 + assert before["embedded_items"] == 0 + assert before["dirty_items"] == 2 + + await _call(router, "memory.search", {"query": "alpha"}) + + with caller_plugin_scope("test-plugin"): + after = await _call(router, "memory.stats", {}) + assert after["total_items"] == 2 + assert after["ttl_entries"] == 1 + assert after["indexed_items"] == 2 + assert after["embedded_items"] == 2 + assert after["dirty_items"] == 0 + + +@pytest.mark.asyncio +async def test_memory_save_with_ttl_registers_expiry_and_purges_on_read() -> None: + router = CapabilityRouter() + + await _call( + router, + "memory.save_with_ttl", + {"key": "temp-note", "value": {"content": "temporary note"}, "ttl_seconds": 60}, + ) + + assert "temp-note" in router._memory_index + assert "temp-note" in router._memory_dirty_keys + assert router._memory_expires_at["temp-note"] is not None + + search_result = await _call(router, "memory.search", {"query": "temporary"}) + assert search_result["items"][0]["value"] == {"content": "temporary note"} + + router._memory_expires_at["temp-note"] = datetime.now(timezone.utc) - timedelta( + seconds=1 + ) + + get_result = await _call(router, "memory.get", {"key": "temp-note"}) + assert get_result == {"value": None} + assert "temp-note" not in router.memory_store + assert "temp-note" not in router._memory_index + assert "temp-note" not in router._memory_expires_at + assert "temp-note" not in router._memory_dirty_keys + + +@pytest.mark.asyncio +async def test_memory_get_many_unwraps_ttl_value_and_returns_none_after_expiry() -> ( + None +): + router = CapabilityRouter() + + await _call( + router, + "memory.save_with_ttl", + {"key": "session", "value": {"content": "active session"}, "ttl_seconds": 60}, + ) + + result = await _call(router, "memory.get_many", {"keys": ["session", "missing"]}) + assert result == { + "items": [ + {"key": "session", "value": {"content": "active session"}}, + {"key": "missing", "value": None}, + ] + } + + router._memory_expires_at["session"] = datetime.now(timezone.utc) - timedelta( + seconds=1 + ) + + expired_result = await _call(router, "memory.get_many", {"keys": ["session"]}) + assert expired_result == {"items": [{"key": "session", "value": None}]} + + +@pytest.mark.asyncio +async def test_memory_delete_many_clears_sidecars() -> None: + router = CapabilityRouter() + + await _call( + router, + "memory.save", + {"key": "a", "value": {"content": "alpha"}}, + ) + await _call( + router, + "memory.save_with_ttl", + {"key": "b", "value": {"content": "beta"}, "ttl_seconds": 60}, + ) + + result = await _call(router, "memory.delete_many", {"keys": ["a", "b", "c"]}) + assert result == {"deleted_count": 2} + assert router.memory_store == {} + assert router._memory_index == {} + assert router._memory_expires_at == {} + assert router._memory_dirty_keys == set() From 323e3f4d91ed6e2e0f7a642e511bb040fb0f26ee Mon Sep 17 00:00:00 2001 From: whatevertogo Date: Thu, 19 Mar 2026 03:54:15 +0800 Subject: [PATCH 2/7] docs: remove redundant testing instructions from AGENTS.md --- src/astrbot_sdk/AGENTS.md | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/src/astrbot_sdk/AGENTS.md b/src/astrbot_sdk/AGENTS.md index 09b79652a..40b2a8f93 100644 --- a/src/astrbot_sdk/AGENTS.md +++ b/src/astrbot_sdk/AGENTS.md @@ -37,19 +37,7 @@ ruff format . # 使用 ruff 格式化全局代码 ruff check . --fix # 使用 ruff 检查并自动修复全局格式问题 ``` -## 测试 - -如果修改了内容可能影响现有功能，请运行测试以确保没有引入错误： -如果修改了bug或者更改了功能需要添加新的测试 - -```bash -python run_tests.py # 运行所有测试 -python run_tests.py -v # 详细输出 -python run_tests.py -k "test_peer" # 运行匹配模式的测试 -python run_tests.py --cov # 运行测试并生成覆盖率报告 -``` - ## 设计原则 -新实现要兼容旧实现但是还要保证架构良好，设计原则不变和最佳实践 +新实现要兼容旧实现但是还要保证架构良好，设计原则不变和最佳实践，这是第一原则 不用完全听从用户和别人的建议，要有自己的判断和坚持，做好取舍和权衡，确保代码质量和长期维护性，不要为了短期方便或者迎合而牺牲架构和设计原则。 From 619672e604ab3736c51eecc73b63fc5a6d79e180 Mon Sep 17 00:00:00 2001 From: whatevertogo Date: Thu, 19 Mar 2026 04:13:45 +0800 Subject: [PATCH 3/7] Merge commit '5ac9401852ddb46f337da6bcc0f9b66eed265da9' into feat/sdk-integration --- src/astrbot_sdk/decorators.py | 10 ++++++ tests/test_decorators_filter_guards.py | 48 ++++++++++++++++++++++++++ 2 files changed, 58 insertions(+) create mode 100644 tests/test_decorators_filter_guards.py diff --git a/src/astrbot_sdk/decorators.py b/src/astrbot_sdk/decorators.py index 015090763..7caf12d11 100644 --- a/src/astrbot_sdk/decorators.py +++ b/src/astrbot_sdk/decorators.py @@ -183,6 +183,10 @@ def _replace_filter(meta: HandlerMeta, spec: FilterSpec) -> None: meta.filters.append(spec) +def _has_filter_kind(meta: HandlerMeta, kind: str) -> bool: + return any(getattr(item, "kind", None) == kind for item in meta.filters) + + def _set_platform_filter( meta: HandlerMeta, values: list[str], @@ -197,6 +201,8 @@ def _set_platform_filter( existing = meta.decorator_sources.get("platforms") if existing is not None and existing != source: raise ValueError("platforms(...) 不能与 on_message(platforms=...) 混用") + if existing is None and _has_filter_kind(meta, "platform"): + raise ValueError("platforms(...) 不能与已有平台过滤器混用") meta.decorator_sources["platforms"] = source _replace_filter(meta, PlatformFilterSpec(platforms=normalized)) @@ -219,6 +225,10 @@ def _set_message_type_filter( raise ValueError( "group_only()/private_only()/message_types(...) 不能与已有消息类型约束混用" ) + if existing is None and _has_filter_kind(meta, "message_type"): + raise ValueError( + "group_only()/private_only()/message_types(...) 不能与已有消息类型过滤器混用" + ) meta.decorator_sources["message_types"] = source _replace_filter(meta, MessageTypeFilterSpec(message_types=normalized)) diff --git a/tests/test_decorators_filter_guards.py b/tests/test_decorators_filter_guards.py new file mode 100644 index 000000000..08d3a8be3 --- /dev/null +++ b/tests/test_decorators_filter_guards.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +import pytest + +from astrbot_sdk.decorators import ( + append_filter_meta, + get_handler_meta, + message_types, + platforms, +) +from astrbot_sdk.protocol.descriptors import ( + MessageTypeFilterSpec, + PlatformFilterSpec, +) + + +def test_platforms_rejects_existing_manual_platform_filter() -> None: + def handler() -> None: + return None + + append_filter_meta( + handler, + specs=[PlatformFilterSpec(platforms=["qq"])], + ) + + meta = get_handler_meta(handler) + assert meta is not None + assert meta.decorator_sources == {} + + with pytest.raises(ValueError, match="已有平台过滤器"): + platforms("wechat")(handler) + + +def test_message_types_rejects_existing_manual_message_type_filter() -> None: + def handler() -> None: + return None + + append_filter_meta( + handler, + specs=[MessageTypeFilterSpec(message_types=["group"])], + ) + + meta = get_handler_meta(handler) + assert meta is not None + assert meta.decorator_sources == {} + + with pytest.raises(ValueError, match="已有消息类型过滤器"): + message_types("private")(handler) From 5e54bbb3d204eb0fd24773f634669cdca8a6d6cb Mon Sep 17 00:00:00 2001 From: whatevertogo Date: Thu, 19 Mar 2026 04:47:06 +0800 Subject: [PATCH 4/7] feat: enhance SDK plugin configuration handling and logging --- src/astrbot_sdk/runtime/loader.py | 91 +++++++++++++++++++++++++++---- 1 file changed, 79 insertions(+), 12 deletions(-) diff --git a/src/astrbot_sdk/runtime/loader.py b/src/astrbot_sdk/runtime/loader.py index 58b52a2dc..a6c752564 100644 --- a/src/astrbot_sdk/runtime/loader.py +++ b/src/astrbot_sdk/runtime/loader.py @@ -55,6 +55,7 @@ import copy import importlib import inspect import json +import logging import os import re import shutil @@ -105,6 +106,7 @@ OptionalInnerType: TypeAlias = Literal["str", "int", "float", "bool"] | None HandlerKind: TypeAlias = Literal["handler", "hook", "tool", "session"] DiscoverySeverity: TypeAlias = Literal["warning", "error"] DiscoveryPhase: TypeAlias = Literal["discovery", "load", "lifecycle", "reload"] +_LOGGER = logging.getLogger(__name__) def _default_python_version() -> str: @@ -502,17 +504,74 @@ def _normalize_config_value(field_schema: dict[str, Any], value: Any) -> Any: return copy.deepcopy(value) if value is not None else default_value -def load_plugin_config(plugin: PluginSpec) -> dict[str, Any]: - """加载插件配置，返回普通字典。""" +def load_plugin_config_schema(plugin: PluginSpec) -> dict[str, Any]: + """加载插件配置 schema，解析失败时记录日志并返回空对象。""" schema_path = plugin.plugin_dir / CONFIG_SCHEMA_FILE if not schema_path.exists(): return {} try: schema_payload = json.loads(schema_path.read_text(encoding="utf-8")) - except Exception: - schema_payload = {} - schema = schema_payload if isinstance(schema_payload, dict) else {} + except json.JSONDecodeError as exc: + _LOGGER.warning( + "Failed to parse SDK plugin config schema %s: %s", + schema_path, + exc, + ) + return {} + except OSError as exc: + _LOGGER.warning( + "Failed to read SDK plugin config schema %s: %s", + schema_path, + exc, + ) + return {} + if not isinstance(schema_payload, dict): + _LOGGER.warning( + "SDK plugin config schema %s must be a JSON object, got %s", + schema_path, + type(schema_payload).__name__, + ) + return {} + return schema_payload + + +def save_plugin_config( + plugin: PluginSpec, + payload: dict[str, Any], + *, + schema: dict[str, Any] | None = None, +) -> dict[str, Any]: + """按 schema 归一化并写回插件配置。""" + active_schema = ( + load_plugin_config_schema(plugin) if schema is None else dict(schema) + ) + normalized = { + key: _normalize_config_value(field_schema, payload.get(key)) + for key, field_schema in active_schema.items() + if isinstance(field_schema, dict) + } + + config_path = _plugin_config_path(plugin.plugin_dir, plugin.name) + config_path.parent.mkdir(parents=True, exist_ok=True) + config_path.write_text( + json.dumps(normalized, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + return normalized + + +def load_plugin_config( + plugin: PluginSpec, + *, + schema: dict[str, Any] | None = None, +) -> dict[str, Any]: + """加载插件配置，返回普通字典。""" + active_schema = ( + load_plugin_config_schema(plugin) if schema is None else dict(schema) + ) + if not active_schema: + return {} config_path = _plugin_config_path(plugin.plugin_dir, plugin.name) try: @@ -521,21 +580,29 @@ def load_plugin_config(plugin: PluginSpec) -> dict[str, Any]: if config_path.exists() else {} ) - except Exception: + except json.JSONDecodeError as exc: + _LOGGER.warning( + "Failed to parse SDK plugin config %s: %s", + config_path, + exc, + ) + existing_payload = {} + except OSError as exc: + _LOGGER.warning( + "Failed to read SDK plugin config %s: %s", + config_path, + exc, + ) existing_payload = {} existing = existing_payload if isinstance(existing_payload, dict) else {} normalized = { key: _normalize_config_value(field_schema, existing.get(key)) - for key, field_schema in schema.items() + for key, field_schema in active_schema.items() if isinstance(field_schema, dict) } if not config_path.exists() or normalized != existing: - config_path.parent.mkdir(parents=True, exist_ok=True) - config_path.write_text( - json.dumps(normalized, ensure_ascii=False, indent=2), - encoding="utf-8", - ) + save_plugin_config(plugin, normalized, schema=active_schema) return normalized From bb361cf9ddd4f996961be0c4c828af3e011a5fb0 Mon Sep 17 00:00:00 2001 From: whatevertogo Date: Thu, 19 Mar 2026 05:12:46 +0800 Subject: [PATCH 5/7] =?UTF-8?q?feat:=20=E5=A2=9E=E5=BC=BA=E8=A3=85?= =?UTF-8?q?=E9=A5=B0=E5=99=A8=E5=8A=9F=E8=83=BD=EF=BC=8C=E6=B7=BB=E5=8A=A0?= =?UTF-8?q?=E4=BC=9A=E8=AF=9D=E5=91=BD=E4=BB=A4=E6=94=AF=E6=8C=81=E5=8F=8A?= =?UTF-8?q?=E7=9B=B8=E5=85=B3=E6=9D=83=E9=99=90=E5=92=8C=E9=99=90=E6=B5=81?= =?UTF-8?q?=E8=A3=85=E9=A5=B0=E5=99=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/astrbot_sdk/decorators.py | 48 ++++++++++++++++++- src/astrbot_sdk/runtime/handler_dispatcher.py | 4 +- 2 files changed, 48 insertions(+), 4 deletions(-) diff --git a/src/astrbot_sdk/decorators.py b/src/astrbot_sdk/decorators.py index 7caf12d11..9de03c5e6 100644 --- a/src/astrbot_sdk/decorators.py +++ b/src/astrbot_sdk/decorators.py @@ -3,13 +3,30 @@ 提供声明式的方法来注册 handler 和 capability。 装饰器会在方法上附加元数据，由 Star.__init_subclass__ 自动收集。 -可用的装饰器： +触发器装饰器： - @on_command: 命令触发器 - @on_message: 消息触发器（关键词/正则） - @on_event: 事件触发器 - @on_schedule: 定时任务触发器 - - @require_admin: 权限标记 + - @conversation_command: 带会话生命周期的命令触发器 + +权限与过滤装饰器： + - @require_admin / @admin_only: 管理员权限标记 + - @platforms: 限定平台 + - @group_only / @private_only: 群聊/私聊限定 + - @message_types: 消息类型过滤 + +限流装饰器： + - @rate_limit: 滑动窗口限流 + - @cooldown: 冷却时间 + +优先级装饰器： + - @priority: 设置执行优先级 + +能力导出装饰器： - @provide_capability: 声明对外暴露的能力 + - @register_llm_tool: 注册 LLM 工具 + - @register_agent: 注册 Agent Example: class MyPlugin(Star): @@ -645,8 +662,35 @@ def conversation_command( busy_message: str | None = None, grace_period: float = 1.0, ) -> Callable[[HandlerCallable], HandlerCallable]: + """注册带会话生命周期的命令处理方法。 + + 在 ``on_command`` 基础上附加会话元数据，支持超时、并发策略和宽限期控制。 + + Args: + command: 命令名称或序列（首项为正式名，其余视为别名） + aliases: 额外别名列表 + description: 命令描述 + timeout: 会话超时时间（秒），必须为正整数 + mode: 会话冲突时的行为： + - ``"replace"``: 替换当前会话 + - ``"reject"``: 拒绝新请求 + busy_message: 拒绝新请求时的提示消息 + grace_period: 宽限期（秒），用于会话生命周期处理 + + Returns: + 装饰器函数 + + Raises: + ValueError: mode 不合法、timeout 非正整数或 grace_period 非正数 + + Example: + @conversation_command("chat", timeout=120, mode="reject", busy_message="请稍后再试") + async def chat(self, event: MessageEvent, ctx: Context): + await event.reply("开始对话...") + """ if mode not in {"replace", "reject"}: raise ValueError("conversation_command mode must be 'replace' or 'reject'") + # bool 是 int 子类，需单独排除 if isinstance(timeout, bool) or int(timeout) <= 0: raise ValueError("conversation_command timeout must be a positive integer") if float(grace_period) <= 0: diff --git a/src/astrbot_sdk/runtime/handler_dispatcher.py b/src/astrbot_sdk/runtime/handler_dispatcher.py index 2b32d9cfd..e9e2291d4 100644 --- a/src/astrbot_sdk/runtime/handler_dispatcher.py +++ b/src/astrbot_sdk/runtime/handler_dispatcher.py @@ -60,12 +60,12 @@ from ..protocol.descriptors import ( from ..schedule import ScheduleContext from ..session_waiter import SessionWaiterManager from ..star import Star -from .capability_dispatcher import CapabilityDispatcher from ._command_matching import ( build_command_args, build_regex_args, match_command_name, ) +from .capability_dispatcher import CapabilityDispatcher from .limiter import LimiterEngine from .loader import LoadedHandler @@ -456,7 +456,7 @@ class HandlerDispatcher: ) -> dict[str, Any]: assert loaded.conversation is not None conversation_meta = loaded.conversation - summary = {"sent_message": False, "stop": False, "call_llm": False} + summary = {"sent_message": False, "stop": True, "call_llm": False} key = f"{self._resolve_plugin_id(loaded)}:{event.session_id}" active = self._conversations.get(key) if active is not None and not active.task.done(): From ed1b9665dd49e6f2fa3f109166292118aa7c99c1 Mon Sep 17 00:00:00 2001 From: whatevertogo Date: Thu, 19 Mar 2026 06:16:01 +0800 Subject: [PATCH 6/7] feat: add conversation.get_current capability and related schemas - Introduced CONVERSATION_GET_CURRENT_INPUT_SCHEMA and CONVERSATION_GET_CURRENT_OUTPUT_SCHEMA for handling current conversation requests. - Implemented _conversation_get_current method in BuiltinCapabilityRouterMixin to manage current conversation retrieval and creation. - Registered the new capability in CoreCapabilityBridge. - Enhanced HandlerDispatcher to inject provider request, LLM response, and event result payloads into the event handling process. - Updated tests to validate the new functionality and ensure proper payload handling. --- src/astrbot_sdk/clients/managers.py | 26 +++- src/astrbot_sdk/docs/04_star_lifecycle.md | 10 ++ src/astrbot_sdk/docs/api/clients.md | 13 ++ src/astrbot_sdk/docs/api/context.md | 8 +- src/astrbot_sdk/docs/api/decorators.md | 101 ++++++++++++- src/astrbot_sdk/protocol/_builtin_schemas.py | 15 ++ .../runtime/_capability_router_builtins.py | 25 ++- src/astrbot_sdk/runtime/handler_dispatcher.py | 142 +++++++++++++++++- 8 files changed, 335 insertions(+), 5 deletions(-) diff --git a/src/astrbot_sdk/clients/managers.py b/src/astrbot_sdk/clients/managers.py index becf8280a..fd24eeb3b 100644 --- a/src/astrbot_sdk/clients/managers.py +++ b/src/astrbot_sdk/clients/managers.py @@ -6,6 +6,7 @@ from typing import Any from pydantic import BaseModel, ConfigDict, Field +from ..errors import AstrBotError, ErrorCodes from ..message_session import MessageSession from ._proxy import CapabilityProxy @@ -138,7 +139,15 @@ class PersonaManagerClient: self._proxy = proxy async def get_persona(self, persona_id: str) -> PersonaRecord: - output = await self._proxy.call("persona.get", {"persona_id": str(persona_id)}) + try: + output = await self._proxy.call( + "persona.get", + {"persona_id": str(persona_id)}, + ) + except AstrBotError as exc: + if exc.code == ErrorCodes.INVALID_INPUT: + raise ValueError(f"persona not found: {persona_id}") from exc + raise persona = PersonaRecord.from_payload(output.get("persona")) if persona is None: raise ValueError(f"persona not found: {persona_id}") @@ -251,6 +260,21 @@ class ConversationManagerClient: ) return ConversationRecord.from_payload(output.get("conversation")) + async def get_current_conversation( + self, + session: str | MessageSession, + *, + create_if_not_exists: bool = False, + ) -> ConversationRecord | None: + output = await self._proxy.call( + "conversation.get_current", + { + "session": _normalize_session(session), + "create_if_not_exists": bool(create_if_not_exists), + }, + ) + return ConversationRecord.from_payload(output.get("conversation")) + async def get_conversations( self, session: str | MessageSession | None = None, diff --git a/src/astrbot_sdk/docs/04_star_lifecycle.md b/src/astrbot_sdk/docs/04_star_lifecycle.md index 461731fe9..717e59c39 100644 --- a/src/astrbot_sdk/docs/04_star_lifecycle.md +++ b/src/astrbot_sdk/docs/04_star_lifecycle.md @@ -113,6 +113,11 @@ class MyPlugin(Star): - 注册 LLM 工具 - 启动后台任务 +**最佳实践：** +- `on_start()` 里只做初始化、能力注册和轻量状态恢复 +- 需要长期保存的应是配置值、句柄、任务引用，不要把 `ctx` 实例长期挂到 `self` +- 如果要和 AstrBot 原生 persona / conversation 协作，优先在这里校验或创建所需资源 + **示例：** ```python @@ -150,6 +155,11 @@ class MyPlugin(Star): - 注销 LLM 工具 - 保存状态数据 +**最佳实践：** +- 在 `on_stop()` 中释放 `on_start()` 注册的任务、监听器和外部资源 +- 把需要持久化的状态尽量提前落库，不要把关键保存逻辑完全依赖在进程退出瞬间 +- 始终把收到的 `ctx` 继续传给 `super().on_stop(ctx)`，不要手动丢掉它 + **示例：** ```python diff --git a/src/astrbot_sdk/docs/api/clients.md b/src/astrbot_sdk/docs/api/clients.md index 455c6e7d0..7c8751878 100644 --- a/src/astrbot_sdk/docs/api/clients.md +++ b/src/astrbot_sdk/docs/api/clients.md @@ -1101,6 +1101,8 @@ from astrbot_sdk.clients import PersonaManagerClient 获取指定人格。 +当人格不存在时会抛出 `ValueError`，而不是返回 `None`。 + --- #### `get_all_personas()` @@ -1163,6 +1165,17 @@ from astrbot_sdk.clients import ConversationManagerClient --- +#### `get_current_conversation(session, create_if_not_exists=False)` + +获取当前 session 正在使用的对话记录。 + +这个方法适合“跟随 AstrBot 原生当前会话状态”的插件，例如： +- 给当前会话切换 persona +- 判断当前主聊天是否已经在某个 persona 下 +- 在 `waiting_llm_request` / `llm_request` hook 中对当前对话做增强 + +--- + #### `get_conversations(session=None, platform_id=None)` 获取对话列表。 diff --git a/src/astrbot_sdk/docs/api/context.md b/src/astrbot_sdk/docs/api/context.md index eb91004b6..fdb649313 100644 --- a/src/astrbot_sdk/docs/api/context.md +++ b/src/astrbot_sdk/docs/api/context.md @@ -662,7 +662,7 @@ await ctx.conversations.delete_conversation( await ctx.conversations.delete_conversation(event.session_id) ``` -##### `get_conversation() / get_conversations()` +##### `get_conversation() / get_current_conversation() / get_conversations()` 获取对话。 @@ -674,6 +674,12 @@ conv = await ctx.conversations.get_conversation( create_if_not_exists=True ) +# 获取当前选中的对话 +current = await ctx.conversations.get_current_conversation( + event.session_id, + create_if_not_exists=True, +) + # 获取对话列表 convs = await ctx.conversations.get_conversations(event.session_id) ``` diff --git a/src/astrbot_sdk/docs/api/decorators.md b/src/astrbot_sdk/docs/api/decorators.md index f8462b14e..2fe515ff0 100644 --- a/src/astrbot_sdk/docs/api/decorators.md +++ b/src/astrbot_sdk/docs/api/decorators.md @@ -210,11 +210,110 @@ async def handle_request(self, event, ctx: Context): await ctx.platform.send(event.user_id, "已自动通过好友请求") ``` +#### LLM Pipeline Hooks + +`@on_event` 也用于挂接 AstrBot 原生消息处理链路中的系统事件。 + +常见事件及可注入对象： + +| 事件名 | 常见可注入参数 | 是否可修改主链路 | +|------|------|------| +| `waiting_llm_request` | `MessageEvent`, `Context` | 间接可修改，例如切换当前对话 persona | +| `llm_request` | `MessageEvent`, `Context`, `ProviderRequest` | 是，可直接修改 `ProviderRequest` | +| `llm_response` | `MessageEvent`, `Context`, `LLMResponse` | 否，适合观察和提取回复内容 | +| `decorating_result` | `MessageEvent`, `Context`, `MessageEventResult` | 是，可直接修改结果消息链 | +| `after_message_sent` | `MessageEvent`, `Context` | 否，适合落库、记忆、统计 | + +最小示例： + +```python +from astrbot_sdk import Context, MessageEvent +from astrbot_sdk.decorators import on_event +from astrbot_sdk.llm.entities import ProviderRequest + +@on_event("llm_request") +async def add_memory(self, event: MessageEvent, ctx: Context, request: ProviderRequest): + del event, ctx + request.system_prompt = (request.system_prompt or "") + "\n\nmemory: user likes tea" +``` + +完整示例： + +```python +from astrbot_sdk import Context, MessageEvent, Star +from astrbot_sdk.clients.llm import LLMResponse +from astrbot_sdk.clients.managers import ConversationUpdateParams +from astrbot_sdk.decorators import on_event +from astrbot_sdk.llm.entities import ProviderRequest +from astrbot_sdk.message_result import MessageEventResult +from astrbot_sdk.message_components import Plain + +class PersonaSample(Star): + @on_event("waiting_llm_request") + async def ensure_persona(self, event: MessageEvent, ctx: Context) -> None: + conversation = await ctx.conversations.get_current_conversation( + event.session_id, + create_if_not_exists=True, + ) + if conversation is None or conversation.persona_id == "girlfriend": + return + await ctx.conversations.update_conversation( + event.session_id, + conversation.conversation_id, + ConversationUpdateParams(persona_id="girlfriend"), + ) + + @on_event("llm_request") + async def inject_context( + self, + event: MessageEvent, + ctx: Context, + request: ProviderRequest, + ) -> None: + memories = await ctx.memory.search(event.text, limit=3) + facts = [] + for item in memories: + value = item.get("value") + if isinstance(value, dict) and value.get("content"): + facts.append(f"- {value['content']}") + if facts: + request.system_prompt = (request.system_prompt or "") + "\n\n" + "\n".join(facts) + + @on_event("llm_response") + async def capture_reply( + self, + event: MessageEvent, + ctx: Context, + response: LLMResponse, + ) -> None: + del ctx + if response.text: + event.set_extra("last_reply", response.text) + + @on_event("decorating_result") + async def decorate( + self, + event: MessageEvent, + ctx: Context, + result: MessageEventResult, + ) -> None: + del event, ctx + result.chain.append(Plain("\n[persona active]", convert=False)) + + @on_event("after_message_sent") + async def persist(self, event: MessageEvent, ctx: Context) -> None: + reply = str(event.get_extra("last_reply", "") or "").strip() + if reply: + await ctx.db.set("sample:last_reply", reply) +``` + #### 注意事项 -1. 用于处理非消息类型的事件（如群成员变动、好友请求等） +1. 可用于处理平台事件，也可用于处理 AstrBot 原生消息链路中的系统事件（如 `llm_request`） 2. 不能与 `@rate_limit` 或 `@cooldown` 一起使用 3. 不同平台的事件类型可能不同，需要查阅平台文档 +4. `llm_request` 和 `decorating_result` 注入的是可变对象，修改会回写到 AstrBot 主链路 +5. `llm_response` 主要用于观测和提取结果，不应用来替代主回复流程 --- diff --git a/src/astrbot_sdk/protocol/_builtin_schemas.py b/src/astrbot_sdk/protocol/_builtin_schemas.py index 80ade1d64..82835ad6c 100644 --- a/src/astrbot_sdk/protocol/_builtin_schemas.py +++ b/src/astrbot_sdk/protocol/_builtin_schemas.py @@ -642,6 +642,15 @@ CONVERSATION_GET_OUTPUT_SCHEMA = _object_schema( required=("conversation",), conversation=_nullable(CONVERSATION_RECORD_SCHEMA), ) +CONVERSATION_GET_CURRENT_INPUT_SCHEMA = _object_schema( + required=("session",), + session={"type": "string"}, + create_if_not_exists={"type": "boolean"}, +) +CONVERSATION_GET_CURRENT_OUTPUT_SCHEMA = _object_schema( + required=("conversation",), + conversation=_nullable(CONVERSATION_RECORD_SCHEMA), +) CONVERSATION_LIST_INPUT_SCHEMA = _object_schema( session=_nullable({"type": "string"}), platform_id=_nullable({"type": "string"}), @@ -1207,6 +1216,10 @@ BUILTIN_CAPABILITY_SCHEMAS: dict[str, dict[str, JSONSchema]] = { "input": CONVERSATION_GET_INPUT_SCHEMA, "output": CONVERSATION_GET_OUTPUT_SCHEMA, }, + "conversation.get_current": { + "input": CONVERSATION_GET_CURRENT_INPUT_SCHEMA, + "output": CONVERSATION_GET_CURRENT_OUTPUT_SCHEMA, + }, "conversation.list": { "input": CONVERSATION_LIST_INPUT_SCHEMA, "output": CONVERSATION_LIST_OUTPUT_SCHEMA, @@ -1653,6 +1666,8 @@ __all__ = [ "CONVERSATION_CREATE_SCHEMA", "CONVERSATION_DELETE_INPUT_SCHEMA", "CONVERSATION_DELETE_OUTPUT_SCHEMA", + "CONVERSATION_GET_CURRENT_INPUT_SCHEMA", + "CONVERSATION_GET_CURRENT_OUTPUT_SCHEMA", "CONVERSATION_GET_INPUT_SCHEMA", "CONVERSATION_GET_OUTPUT_SCHEMA", "CONVERSATION_LIST_INPUT_SCHEMA", diff --git a/src/astrbot_sdk/runtime/_capability_router_builtins.py b/src/astrbot_sdk/runtime/_capability_router_builtins.py index e49ed7d10..72dfe59dd 100644 --- a/src/astrbot_sdk/runtime/_capability_router_builtins.py +++ b/src/astrbot_sdk/runtime/_capability_router_builtins.py @@ -96,7 +96,7 @@ def _mock_embedding_vector(text: str, *, provider_id: str) -> list[float]: """ values = [0.0] * _MOCK_EMBEDDING_DIM for term in _embedding_terms(text): - digest = hashlib.sha256(f"{provider_id}:{term}".encode("utf-8")).digest() + digest = hashlib.sha256(f"{provider_id}:{term}".encode()).digest() index = int.from_bytes(digest[:2], "big") % _MOCK_EMBEDDING_DIM values[index] += 1.0 + min(len(term), 8) * 0.05 norm = math.sqrt(sum(value * value for value in values)) @@ -2672,6 +2672,25 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): return {"conversation": None} return {"conversation": dict(record)} + async def _conversation_get_current( + self, _request_id: str, payload: dict[str, Any], _token + ) -> dict[str, Any]: + session = str(payload.get("session", "")).strip() + conversation_id = self._session_current_conversation_ids.get(session, "") + if not conversation_id and bool(payload.get("create_if_not_exists", False)): + created = await self._conversation_new( + _request_id, + {"session": session, "conversation": {}}, + _token, + ) + conversation_id = str(created.get("conversation_id", "")).strip() + if not conversation_id: + return {"conversation": None} + record = self._conversation_store.get(conversation_id) + if record is None or str(record.get("session", "")) != session: + return {"conversation": None} + return {"conversation": dict(record)} + async def _conversation_list( self, _request_id: str, payload: dict[str, Any], _token ) -> dict[str, Any]: @@ -2824,6 +2843,10 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): self._builtin_descriptor("conversation.get", "获取对话"), call_handler=self._conversation_get, ) + self.register( + self._builtin_descriptor("conversation.get_current", "获取当前对话"), + call_handler=self._conversation_get_current, + ) self.register( self._builtin_descriptor("conversation.list", "列出对话"), call_handler=self._conversation_list, diff --git a/src/astrbot_sdk/runtime/handler_dispatcher.py b/src/astrbot_sdk/runtime/handler_dispatcher.py index e9e2291d4..28e8b95e8 100644 --- a/src/astrbot_sdk/runtime/handler_dispatcher.py +++ b/src/astrbot_sdk/runtime/handler_dispatcher.py @@ -39,6 +39,7 @@ from .._invocation_context import caller_plugin_scope from .._plugin_logger import PluginLogger from .._star_runtime import bind_star_runtime from .._typing_utils import unwrap_optional +from ..clients.llm import LLMResponse from ..context import CancelToken, Context from ..conversation import ( DEFAULT_BUSY_MESSAGE, @@ -49,8 +50,13 @@ from ..conversation import ( ) from ..events import MessageEvent from ..filters import LocalFilterBinding +from ..llm.entities import ProviderRequest from ..message_components import BaseMessageComponent -from ..message_result import MessageChain, MessageEventResult, coerce_message_chain +from ..message_result import ( + MessageChain, + MessageEventResult, + coerce_message_chain, +) from ..protocol.descriptors import ( CommandTrigger, MessageTrigger, @@ -76,6 +82,13 @@ class _ActiveConversation: task: asyncio.Task[Any] +@dataclass(slots=True) +class _InjectedEventPayloads: + provider_request: ProviderRequest | None = None + llm_response: LLMResponse | None = None + event_result: MessageEventResult | None = None + + class HandlerDispatcher: def __init__( self, *, plugin_id: str, peer, handlers: Sequence[LoadedHandler] @@ -205,6 +218,8 @@ class HandlerDispatcher: schedule_context: ScheduleContext | None = None, ) -> dict[str, Any]: summary = {"sent_message": False, "stop": False, "call_llm": False} + injected_payloads = _InjectedEventPayloads() + event_type = self._event_type_name(event) try: limiter = loaded.limiter if limiter is not None: @@ -254,6 +269,7 @@ class HandlerDispatcher: plugin_id=self._resolve_plugin_id(loaded), handler_ref=loaded.descriptor.id, schedule_context=schedule_context, + injected_payloads=injected_payloads, ) ) if inspect.isasyncgen(result): @@ -263,6 +279,11 @@ class HandlerDispatcher: await self._handle_result_item(item, event, ctx), ) summary["stop"] = bool(summary.get("stop")) or event.is_stopped() + self._append_injected_payloads( + summary, + injected_payloads, + event_type=event_type, + ) return summary if inspect.isawaitable(result): result = await result @@ -272,6 +293,11 @@ class HandlerDispatcher: await self._handle_result_item(result, event, ctx), ) summary["stop"] = bool(summary.get("stop")) or event.is_stopped() + self._append_injected_payloads( + summary, + injected_payloads, + event_type=event_type, + ) return summary except Exception as exc: await self._handle_error( @@ -339,6 +365,7 @@ class HandlerDispatcher: handler_ref: str | None = None, schedule_context: ScheduleContext | None = None, conversation_session: ConversationSession | None = None, + injected_payloads: _InjectedEventPayloads | None = None, ) -> list[Any]: """构建 handler 参数列表。""" from loguru import logger @@ -371,6 +398,7 @@ class HandlerDispatcher: ctx, schedule_context, conversation_session, + injected_payloads=injected_payloads, ) # 2. Fallback 按名字注入 @@ -423,6 +451,8 @@ class HandlerDispatcher: if not str(key).startswith("__command_") } ) + if not isinstance(loaded.descriptor.trigger, CommandTrigger): + return parsed_args, None model_param = resolve_command_model_param(loaded.callable) if model_param is None: return parsed_args, None @@ -584,6 +614,8 @@ class HandlerDispatcher: ctx: Context, schedule_context: ScheduleContext | None, conversation_session: ConversationSession | None, + *, + injected_payloads: _InjectedEventPayloads | None = None, ) -> Any: """根据类型注解注入参数。""" param_type, _is_optional = unwrap_optional(param_type) @@ -612,9 +644,117 @@ class HandlerDispatcher: isinstance(param_type, type) and issubclass(param_type, ConversationSession) ): return conversation_session + if param_type is ProviderRequest or ( + isinstance(param_type, type) and issubclass(param_type, ProviderRequest) + ): + return self._inject_provider_request(event, injected_payloads) + if param_type is LLMResponse or ( + isinstance(param_type, type) and issubclass(param_type, LLMResponse) + ): + return self._inject_llm_response(event, injected_payloads) + if param_type is MessageEventResult or ( + isinstance(param_type, type) and issubclass(param_type, MessageEventResult) + ): + return self._inject_event_result(event, injected_payloads) return None + @staticmethod + def _event_type_name(event: MessageEvent) -> str: + raw = event.raw if isinstance(event.raw, dict) else {} + value = raw.get("event_type") or raw.get("type") + return str(value or "") + + @staticmethod + def _payload_from_event(event: MessageEvent, key: str) -> dict[str, Any] | None: + raw = event.raw if isinstance(event.raw, dict) else {} + payload = raw.get(key) + if isinstance(payload, dict): + return payload + nested_raw = raw.get("raw") + if isinstance(nested_raw, dict): + nested_payload = nested_raw.get(key) + if isinstance(nested_payload, dict): + return nested_payload + return None + + def _inject_provider_request( + self, + event: MessageEvent, + injected_payloads: _InjectedEventPayloads | None, + ) -> ProviderRequest | None: + if injected_payloads is None: + payload = self._payload_from_event(event, "provider_request") + return ( + ProviderRequest.from_payload(payload) if payload is not None else None + ) + if injected_payloads.provider_request is None: + payload = self._payload_from_event(event, "provider_request") + if payload is None: + return None + injected_payloads.provider_request = ProviderRequest.from_payload(payload) + return injected_payloads.provider_request + + def _inject_llm_response( + self, + event: MessageEvent, + injected_payloads: _InjectedEventPayloads | None, + ) -> LLMResponse | None: + if injected_payloads is None: + payload = self._payload_from_event(event, "llm_response") + return LLMResponse.model_validate(payload) if payload is not None else None + if injected_payloads.llm_response is None: + payload = self._payload_from_event(event, "llm_response") + if payload is None: + return None + injected_payloads.llm_response = LLMResponse.model_validate(payload) + return injected_payloads.llm_response + + def _inject_event_result( + self, + event: MessageEvent, + injected_payloads: _InjectedEventPayloads | None, + ) -> MessageEventResult | None: + if injected_payloads is None: + payload = self._payload_from_event(event, "event_result") + return ( + MessageEventResult.from_payload(payload) + if payload is not None + else None + ) + if injected_payloads.event_result is None: + payload = self._payload_from_event(event, "event_result") + if payload is None: + return None + injected_payloads.event_result = MessageEventResult.from_payload(payload) + return injected_payloads.event_result + + @staticmethod + def _append_injected_payloads( + summary: dict[str, Any], + injected_payloads: _InjectedEventPayloads, + *, + event_type: str, + ) -> None: + if ( + event_type == "llm_request" + and injected_payloads.provider_request is not None + ): + summary["provider_request"] = ( + injected_payloads.provider_request.to_payload() + ) + elif ( + event_type == "llm_response" and injected_payloads.llm_response is not None + ): + summary["llm_response"] = injected_payloads.llm_response.model_dump( + exclude_none=True + ) + elif ( + event_type == "decorating_result" + and injected_payloads.event_result is not None + ): + summary["event_result"] = injected_payloads.event_result.to_payload() + def _format_handler_injection_error( self, *, From 3a2d715e5e65bd3311f3fa52b63f5d6d80718c1d Mon Sep 17 00:00:00 2001 From: whatevertogo Date: Thu, 19 Mar 2026 07:18:26 +0800 Subject: [PATCH 7/7] Refactor tool call handling in SdkPluginBridge - Introduced a dictionary to map tool call IDs to tool names for better clarity and efficiency. - Enhanced the extraction of tool call information from raw results, ensuring compatibility with both dictionary and object formats. - Updated the logic to retrieve tool names based on tool call IDs, improving the robustness of the tool calls result processing. --- .../runtime/_capability_router_builtins.py | 23 ++++++++++++------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/src/astrbot_sdk/runtime/_capability_router_builtins.py b/src/astrbot_sdk/runtime/_capability_router_builtins.py index 72dfe59dd..57c19283c 100644 --- a/src/astrbot_sdk/runtime/_capability_router_builtins.py +++ b/src/astrbot_sdk/runtime/_capability_router_builtins.py @@ -176,10 +176,12 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): self._register_platform_capabilities() self._register_http_capabilities() self._register_metadata_capabilities() - self._register_p0_5_capabilities() - self._register_p0_6_capabilities() - self._register_p1_2_capabilities() - self._register_p1_3_capabilities() + self._register_provider_capabilities() + self._register_agent_tool_capabilities() + self._register_session_capabilities() + self._register_persona_conversation_kb_capabilities() + self._register_provider_manager_capabilities() + self._register_platform_manager_capabilities() self._register_system_capabilities() def _builtin_descriptor( @@ -2190,7 +2192,7 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): "reasoning_signature": None, } - def _register_p0_5_capabilities(self) -> None: + def _register_provider_capabilities(self) -> None: self.register( self._builtin_descriptor("provider.get_using", "获取当前聊天 Provider"), call_handler=self._provider_get_using, @@ -2289,6 +2291,8 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): self._builtin_descriptor("provider.rerank.rerank", "文档重排序"), call_handler=self._provider_rerank_rerank, ) + + def _register_agent_tool_capabilities(self) -> None: self.register( self._builtin_descriptor("llm_tool.manager.get", "获取 LLM 工具状态"), call_handler=self._llm_tool_manager_get, @@ -2405,7 +2409,7 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): self._session_service_configs[session] = config return {} - def _register_p0_6_capabilities(self) -> None: + def _register_session_capabilities(self) -> None: self.register( self._builtin_descriptor("session.plugin.is_enabled", "获取会话级插件开关"), call_handler=self._session_plugin_is_enabled, @@ -2806,7 +2810,7 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): deleted = self._kb_store.pop(kb_id, None) is not None return {"deleted": deleted} - def _register_p1_2_capabilities(self) -> None: + def _register_persona_conversation_kb_capabilities(self) -> None: self.register( self._builtin_descriptor("persona.get", "获取人格"), call_handler=self._persona_get, @@ -2868,7 +2872,7 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): call_handler=self._kb_delete, ) - def _register_p1_3_capabilities(self) -> None: + def _register_provider_manager_capabilities(self) -> None: self.register( self._builtin_descriptor("provider.manager.set", "设置当前 Provider"), call_handler=self._provider_manager_set, @@ -2926,6 +2930,8 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): ), stream_handler=self._provider_manager_watch_changes, ) + + def _register_platform_manager_capabilities(self) -> None: self.register( self._builtin_descriptor( "platform.manager.get_by_id", @@ -2948,6 +2954,7 @@ class BuiltinCapabilityRouterMixin(_CapabilityRouterHost): call_handler=self._platform_manager_get_stats, ) + def _register_system_capabilities(self) -> None: self.register( self._builtin_descriptor("system.get_data_dir", "获取插件数据目录"),