diff --git a/.gitignore b/.gitignore index 4ddd0641f..35c92ee93 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ __pycache__/ *.pyd *.so .pytest_cache/ +pytest-cache-files-*/ .mypy_cache/ .ruff_cache/ .coverage @@ -22,6 +23,13 @@ wheels/ .eggs/ pip-wheel-metadata/ +# +fork-docs/ +tmp/ +openspec/ +scripts/ +docs/zh/reference + # Virtual environments .venv/ venv/ diff --git a/src-new/astrbot_sdk/clients/db.py b/src-new/astrbot_sdk/clients/db.py index 7c270d23d..161c0c8bc 100644 --- a/src-new/astrbot_sdk/clients/db.py +++ b/src-new/astrbot_sdk/clients/db.py @@ -18,14 +18,11 @@ - 数据永久存储，除非用户显式删除 - 值类型支持任意 JSON 数据 - 支持前缀查询键列表 - -TODO: - - 缺少批量操作支持 (set_many, get_many) - - 缺少数据变更事件通知 """ from __future__ import annotations +from collections.abc import AsyncIterator, Mapping, Sequence from typing import Any from ._proxy import CapabilityProxy @@ -108,3 +105,67 @@ class DBClient: if not isinstance(keys, (list, tuple)): return [] return [str(item) for item in keys] + + async def get_many(self, keys: Sequence[str]) -> dict[str, Any | None]: + """批量获取多个键的值。 + + Args: + keys: 要读取的键列表 + + Returns: + 一个 dict，key 为键名，value 为对应值（不存在则为 None） + + 示例: + values = await ctx.db.get_many(["user:1", "user:2"]) + if values["user:1"] is None: + print("user:1 missing") + """ + output = await self._proxy.call("db.get_many", {"keys": list(keys)}) + items = output.get("items") + if not isinstance(items, (list, tuple)): + return {} + result: dict[str, Any | None] = {} + for item in items: + if not isinstance(item, dict): + continue + key = item.get("key") + if not isinstance(key, str): + continue + result[key] = item.get("value") + return result + + async def set_many( + self, items: Mapping[str, Any] | Sequence[tuple[str, Any]] + ) -> None: + """批量写入多个键值对。 + + Args: + items: 键值对集合（dict 或二元组序列） + + 示例: + await ctx.db.set_many({"user:1": {"name": "a"}, "user:2": {"name": "b"}}) + """ + if isinstance(items, Mapping): + pairs = list(items.items()) + else: + pairs = list(items) + + payload_items: list[dict[str, Any]] = [ + {"key": str(key), "value": value} for key, value in pairs + ] + await self._proxy.call("db.set_many", {"items": payload_items}) + + def watch(self, prefix: str | None = None) -> AsyncIterator[dict[str, Any]]: + """订阅 KV 变更事件（流式）。 + + Args: + prefix: 键前缀过滤；None 表示订阅所有键 + + Yields: + 变更事件 dict：{"op": "set"|"delete", "key": str, "value": Any|None} + + 示例: + async for event in ctx.db.watch("user:"): + print(event["op"], event["key"]) + """ + return self._proxy.stream("db.watch", {"prefix": prefix}) diff --git a/src-new/astrbot_sdk/protocol/descriptors.py b/src-new/astrbot_sdk/protocol/descriptors.py index 7df6017a1..bdeca089c 100644 --- a/src-new/astrbot_sdk/protocol/descriptors.py +++ b/src-new/astrbot_sdk/protocol/descriptors.py @@ -127,6 +127,37 @@ DB_LIST_OUTPUT_SCHEMA = _object_schema( required=("keys",), keys={"type": "array", "items": {"type": "string"}}, ) +DB_GET_MANY_INPUT_SCHEMA = _object_schema( + required=("keys",), + keys={"type": "array", "items": {"type": "string"}}, +) +DB_GET_MANY_OUTPUT_SCHEMA = _object_schema( + required=("items",), + items={ + "type": "array", + "items": _object_schema( + required=("key", "value"), + key={"type": "string"}, + value=_nullable({}), + ), + }, +) +DB_SET_MANY_INPUT_SCHEMA = _object_schema( + required=("items",), + items={ + "type": "array", + "items": _object_schema( + required=("key", "value"), + key={"type": "string"}, + value={}, + ), + }, +) +DB_SET_MANY_OUTPUT_SCHEMA = _object_schema() +DB_WATCH_INPUT_SCHEMA = _object_schema( + prefix=_nullable({"type": "string"}), +) +DB_WATCH_OUTPUT_SCHEMA = _object_schema() SESSION_REF_SCHEMA = _object_schema( required=("conversation_id",), conversation_id={"type": "string"}, @@ -218,6 +249,18 @@ BUILTIN_CAPABILITY_SCHEMAS: dict[str, dict[str, JSONSchema]] = { "input": DB_LIST_INPUT_SCHEMA, "output": DB_LIST_OUTPUT_SCHEMA, }, + "db.get_many": { + "input": DB_GET_MANY_INPUT_SCHEMA, + "output": DB_GET_MANY_OUTPUT_SCHEMA, + }, + "db.set_many": { + "input": DB_SET_MANY_INPUT_SCHEMA, + "output": DB_SET_MANY_OUTPUT_SCHEMA, + }, + "db.watch": { + "input": DB_WATCH_INPUT_SCHEMA, + "output": DB_WATCH_OUTPUT_SCHEMA, + }, "platform.send": { "input": PLATFORM_SEND_INPUT_SCHEMA, "output": PLATFORM_SEND_OUTPUT_SCHEMA, @@ -491,10 +534,16 @@ __all__ = [ "DB_DELETE_OUTPUT_SCHEMA", "DB_GET_INPUT_SCHEMA", "DB_GET_OUTPUT_SCHEMA", + "DB_GET_MANY_INPUT_SCHEMA", + "DB_GET_MANY_OUTPUT_SCHEMA", "DB_LIST_INPUT_SCHEMA", "DB_LIST_OUTPUT_SCHEMA", "DB_SET_INPUT_SCHEMA", "DB_SET_OUTPUT_SCHEMA", + "DB_SET_MANY_INPUT_SCHEMA", + "DB_SET_MANY_OUTPUT_SCHEMA", + "DB_WATCH_INPUT_SCHEMA", + "DB_WATCH_OUTPUT_SCHEMA", "EventTrigger", "HandlerDescriptor", "JSONSchema", diff --git a/src-new/astrbot_sdk/runtime/capability_router.py b/src-new/astrbot_sdk/runtime/capability_router.py index 32028d9dd..de33aa696 100644 --- a/src-new/astrbot_sdk/runtime/capability_router.py +++ b/src-new/astrbot_sdk/runtime/capability_router.py @@ -109,6 +109,7 @@ CAPABILITY_NAME_PATTERN = re.compile(r"^[a-z][a-z0-9_]*\.[a-z][a-z0-9_]*$") class StreamExecution: iterator: AsyncIterator[dict[str, Any]] finalize: FinalizeHandler + collect_chunks: bool = True StreamHandler = Callable[ @@ -134,8 +135,18 @@ class CapabilityRouter: self.db_store: dict[str, Any] = {} self.memory_store: dict[str, dict[str, Any]] = {} self.sent_messages: list[dict[str, Any]] = [] + self._db_watch_subscriptions: dict[ + str, tuple[str | None, asyncio.Queue[dict[str, Any]]] + ] = {} self._register_builtin_capabilities() + def _emit_db_change(self, *, op: str, key: str, value: Any | None) -> None: + event = {"op": op, "key": key, "value": value} + for prefix, queue in list(self._db_watch_subscriptions.values()): + if prefix is not None and not key.startswith(prefix): + continue + queue.put_nowait(event) + def descriptors(self) -> list[CapabilityDescriptor]: return [ entry.descriptor for entry in self._registrations.values() if entry.exposed @@ -227,6 +238,7 @@ class CapabilityRouter: return StreamExecution( iterator=execution.iterator, finalize=validated_finalize, + collect_chunks=execution.collect_chunks, ) def _register_builtin_capabilities(self) -> None: @@ -329,13 +341,17 @@ class CapabilityRouter: _request_id: str, payload: dict[str, Any], _token ) -> dict[str, Any]: key = str(payload.get("key", "")) - self.db_store[key] = payload.get("value") + value = payload.get("value") + self.db_store[key] = value + self._emit_db_change(op="set", key=key, value=value) return {} async def db_delete( _request_id: str, payload: dict[str, Any], _token ) -> dict[str, Any]: - self.db_store.pop(str(payload.get("key", "")), None) + key = str(payload.get("key", "")) + self.db_store.pop(key, None) + self._emit_db_change(op="delete", key=key, value=None) return {} async def db_list( @@ -347,6 +363,63 @@ class CapabilityRouter: keys = [item for item in keys if item.startswith(prefix)] return {"keys": keys} + async def db_get_many( + _request_id: str, payload: dict[str, Any], _token + ) -> dict[str, Any]: + keys_payload = payload.get("keys") + if not isinstance(keys_payload, (list, tuple)): + raise AstrBotError.invalid_input("db.get_many 的 keys 必须是数组") + keys = [str(item) for item in keys_payload] + items = [{"key": key, "value": self.db_store.get(key)} for key in keys] + return {"items": items} + + async def db_set_many( + _request_id: str, payload: dict[str, Any], _token + ) -> dict[str, Any]: + items_payload = payload.get("items") + if not isinstance(items_payload, (list, tuple)): + raise AstrBotError.invalid_input("db.set_many 的 items 必须是数组") + for entry in items_payload: + if not isinstance(entry, dict): + raise AstrBotError.invalid_input( + "db.set_many 的 items 必须是 object 数组" + ) + key = str(entry.get("key", "")) + value = entry.get("value") + self.db_store[key] = value + self._emit_db_change(op="set", key=key, value=value) + return {} + + async def db_watch( + request_id: str, payload: dict[str, Any], _token + ) -> StreamExecution: + prefix = payload.get("prefix") + prefix_value: str | None + if isinstance(prefix, str): + prefix_value = prefix + elif prefix is None: + prefix_value = None + else: + raise AstrBotError.invalid_input( + "db.watch 的 prefix 必须是 string 或 null" + ) + + queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue() + self._db_watch_subscriptions[request_id] = (prefix_value, queue) + + async def iterator() -> AsyncIterator[dict[str, Any]]: + try: + while True: + yield await queue.get() + finally: + self._db_watch_subscriptions.pop(request_id, None) + + return StreamExecution( + iterator=iterator(), + finalize=lambda _chunks: {}, + collect_chunks=False, + ) + async def platform_send( _request_id: str, payload: dict[str, Any], _token ) -> dict[str, Any]: @@ -460,6 +533,23 @@ class CapabilityRouter: builtin_descriptor("db.list", "列出 KV"), call_handler=db_list, ) + self.register( + builtin_descriptor("db.get_many", "批量读取 KV"), + call_handler=db_get_many, + ) + self.register( + builtin_descriptor("db.set_many", "批量写入 KV"), + call_handler=db_set_many, + ) + self.register( + builtin_descriptor( + "db.watch", + "订阅 KV 变更", + supports_stream=True, + cancelable=True, + ), + stream_handler=db_watch, + ) self.register( builtin_descriptor("platform.send", "发送消息"), call_handler=platform_send, diff --git a/src-new/astrbot_sdk/runtime/peer.py b/src-new/astrbot_sdk/runtime/peer.py index d737c9f57..f501c420d 100644 --- a/src-new/astrbot_sdk/runtime/peer.py +++ b/src-new/astrbot_sdk/runtime/peer.py @@ -507,9 +507,11 @@ class Peer: "stream=true 必须返回 StreamExecution" ) await self._send(EventMessage(id=message.id, phase="started")) + collect_chunks = execution.collect_chunks chunks: list[dict[str, Any]] = [] async for chunk in execution.iterator: - chunks.append(chunk) + if collect_chunks: + chunks.append(chunk) await self._send( EventMessage(id=message.id, phase="delta", data=chunk) ) diff --git a/tests_v4/test_capability_router.py b/tests_v4/test_capability_router.py index 759010aa8..de8a9118c 100644 --- a/tests_v4/test_capability_router.py +++ b/tests_v4/test_capability_router.py @@ -164,6 +164,9 @@ class TestCapabilityRouterInit: assert "db.set" in capability_names assert "db.delete" in capability_names assert "db.list" in capability_names + assert "db.get_many" in capability_names + assert "db.set_many" in capability_names + assert "db.watch" in capability_names # Platform capabilities assert "platform.send" in capability_names @@ -182,6 +185,14 @@ class TestCapabilityRouterInit: assert descriptors[name].input_schema == schema["input"] assert descriptors[name].output_schema == schema["output"] + def test_db_watch_descriptor_supports_stream(self): + router = CapabilityRouter() + descriptors = { + descriptor.name: descriptor for descriptor in router.descriptors() + } + + assert descriptors["db.watch"].supports_stream is True + class TestCapabilityRouterRegister: """Tests for CapabilityRouter.register method.""" @@ -350,6 +361,85 @@ class TestCapabilityRouterExecute: request_id="req-1", ) + +class TestCapabilityRouterDBWatch: + """Router-level tests for db.watch behavior.""" + + @pytest.mark.asyncio + async def test_db_watch_receives_set_and_delete_events(self): + router = CapabilityRouter() + token = CancelToken() + + execution = await router.execute( + "db.watch", + {"prefix": None}, + stream=True, + cancel_token=token, + request_id="watch-1", + ) + assert isinstance(execution, StreamExecution) + assert execution.collect_chunks is False + + await router.execute( + "db.set", + {"key": "a", "value": 1}, + stream=False, + cancel_token=token, + request_id="set-1", + ) + await router.execute( + "db.delete", + {"key": "a"}, + stream=False, + cancel_token=token, + request_id="del-1", + ) + + event1 = await anext(execution.iterator) + event2 = await anext(execution.iterator) + assert event1 == {"op": "set", "key": "a", "value": 1} + assert event2 == {"op": "delete", "key": "a", "value": None} + + close = getattr(execution.iterator, "aclose", None) + if close is not None: + await close() + + @pytest.mark.asyncio + async def test_db_watch_prefix_filters_events(self): + router = CapabilityRouter() + token = CancelToken() + + execution = await router.execute( + "db.watch", + {"prefix": "user:"}, + stream=True, + cancel_token=token, + request_id="watch-2", + ) + assert isinstance(execution, StreamExecution) + + await router.execute( + "db.set", + {"key": "sys:1", "value": 1}, + stream=False, + cancel_token=token, + request_id="set-2", + ) + await router.execute( + "db.set", + {"key": "user:1", "value": {"ok": True}}, + stream=False, + cancel_token=token, + request_id="set-3", + ) + + event = await anext(execution.iterator) + assert event == {"op": "set", "key": "user:1", "value": {"ok": True}} + + close = getattr(execution.iterator, "aclose", None) + if close is not None: + await close() + @pytest.mark.asyncio async def test_execute_missing_capability_raises(self): """execute should raise for unknown capability.""" diff --git a/tests_v4/test_db_client.py b/tests_v4/test_db_client.py index 8023b53b8..8d15c1310 100644 --- a/tests_v4/test_db_client.py +++ b/tests_v4/test_db_client.py @@ -237,3 +237,103 @@ class TestDBClientList: proxy.call.assert_called_once_with("db.list", {"prefix": None}) assert result == [] + + +class TestDBClientGetMany: + """Tests for DBClient.get_many() method.""" + + @pytest.mark.asyncio + async def test_get_many_returns_mapping(self): + proxy = MagicMock(spec=CapabilityProxy) + proxy.call = AsyncMock( + return_value={ + "items": [ + {"key": "a", "value": 1}, + {"key": "b", "value": None}, + ] + } + ) + client = DBClient(proxy) + + result = await client.get_many(["a", "b"]) + + proxy.call.assert_called_once_with("db.get_many", {"keys": ["a", "b"]}) + assert result == {"a": 1, "b": None} + + @pytest.mark.asyncio + async def test_get_many_returns_empty_dict_for_malformed_items(self): + proxy = MagicMock(spec=CapabilityProxy) + proxy.call = AsyncMock(return_value={"items": "not-a-list"}) + client = DBClient(proxy) + + result = await client.get_many(["a"]) + + assert result == {} + + +class TestDBClientSetMany: + """Tests for DBClient.set_many() method.""" + + @pytest.mark.asyncio + async def test_set_many_accepts_mapping(self): + proxy = MagicMock(spec=CapabilityProxy) + proxy.call = AsyncMock(return_value={}) + client = DBClient(proxy) + + await client.set_many({"a": 1, "b": 2}) + + proxy.call.assert_called_once_with( + "db.set_many", + {"items": [{"key": "a", "value": 1}, {"key": "b", "value": 2}]}, + ) + + @pytest.mark.asyncio + async def test_set_many_accepts_sequence_pairs(self): + proxy = MagicMock(spec=CapabilityProxy) + proxy.call = AsyncMock(return_value={}) + client = DBClient(proxy) + + await client.set_many([("a", True), ("b", {"x": 1})]) + + proxy.call.assert_called_once_with( + "db.set_many", + {"items": [{"key": "a", "value": True}, {"key": "b", "value": {"x": 1}}]}, + ) + + +class TestDBClientWatch: + """Tests for DBClient.watch() method.""" + + @pytest.mark.asyncio + async def test_watch_calls_proxy_stream_and_yields_events(self): + async def gen(): + yield {"op": "set", "key": "a", "value": 1} + yield {"op": "delete", "key": "a", "value": None} + + proxy = MagicMock(spec=CapabilityProxy) + proxy.stream = MagicMock(return_value=gen()) + client = DBClient(proxy) + + iterator = client.watch() + + proxy.stream.assert_called_once_with("db.watch", {"prefix": None}) + events = [event async for event in iterator] + assert events == [ + {"op": "set", "key": "a", "value": 1}, + {"op": "delete", "key": "a", "value": None}, + ] + + @pytest.mark.asyncio + async def test_watch_with_prefix(self): + async def gen(): + yield {"op": "set", "key": "user:1", "value": {"ok": True}} + + proxy = MagicMock(spec=CapabilityProxy) + proxy.stream = MagicMock(return_value=gen()) + client = DBClient(proxy) + + iterator = client.watch(prefix="user:") + + proxy.stream.assert_called_once_with("db.watch", {"prefix": "user:"}) + events = [event async for event in iterator] + assert events == [{"op": "set", "key": "user:1", "value": {"ok": True}}]