feat: 重构 agentscope 缓存架构，新增消息和附件缓存

2026-03-25 17:41:55 +08:00
parent d22ded21f8
commit 599c597e69
25 changed files with 1509 additions and 78 deletions
@@ -0,0 +1,228 @@
+from __future__ import annotations
+
+from datetime import datetime, timedelta, timezone
+from typing import Any, cast
+
+import pytest
+
+from core.agentscope.caches.context_messages_cache import ContextMessagesCache
+from schemas.domain.automation import ContextWindowMode, MessageContextConfig
+
+
+class _FakeCacheStore:
+    def __init__(self) -> None:
+        self.hash_store: dict[str, dict[str, str]] = {}
+        self.set_store: dict[str, set[str]] = {}
+
+    async def hgetall(self, key: str) -> dict[str, str]:
+        return dict(self.hash_store.get(key, {}))
+
+    async def hset(self, key: str, mapping: dict[str, str]) -> int:
+        self.hash_store[key] = dict(mapping)
+        return 1
+
+    async def hincrby(self, key: str, field: str, amount: int = 1) -> int:
+        del key, field, amount
+        return 0
+
+    async def expire(self, key: str, ttl_seconds: int) -> int:
+        del key, ttl_seconds
+        return 1
+
+    async def delete(self, *keys: str) -> int:
+        for key in keys:
+            self.hash_store.pop(key, None)
+            self.set_store.pop(key, None)
+        return len(keys)
+
+    async def sadd(self, key: str, *members: str) -> int:
+        values = self.set_store.setdefault(key, set())
+        before = len(values)
+        for member in members:
+            values.add(member)
+        return len(values) - before
+
+    async def smembers(self, key: str) -> set[str]:
+        return set(self.set_store.get(key, set()))
+
+
+@pytest.mark.asyncio
+async def test_context_messages_cache_set_get_roundtrip() -> None:
+    store = _FakeCacheStore()
+    cache = ContextMessagesCache(
+        client=store,
+        key_prefix="agent:context-messages",
+        ttl_seconds=600,
+    )
+    config = MessageContextConfig(window_mode=ContextWindowMode.DAY, window_count=2)
+    messages: list[dict[str, object]] = [
+        {"role": "user", "content": "hello", "timestamp": "2026-03-25T08:00:00+00:00"}
+    ]
+
+    await cache.set(
+        thread_id="thread-1",
+        runtime_mode="chat",
+        context_config=config,
+        messages=messages,
+    )
+    loaded = await cache.get(
+        thread_id="thread-1",
+        runtime_mode="chat",
+        context_config=config,
+    )
+
+    assert loaded is not None
+    assert loaded[0]["content"] == "hello"
+
+
+@pytest.mark.asyncio
+async def test_context_messages_cache_append_skips_when_not_visible() -> None:
+    store = _FakeCacheStore()
+    cache = ContextMessagesCache(
+        client=store,
+        key_prefix="agent:context-messages",
+        ttl_seconds=600,
+    )
+    config = MessageContextConfig(
+        window_mode=ContextWindowMode.NUMBER,
+        window_count=1,
+    )
+    await cache.set(
+        thread_id="thread-1",
+        runtime_mode="chat",
+        context_config=config,
+        messages=cast(
+            list[dict[str, Any]],
+            [
+                {
+                    "role": "user",
+                    "content": "q1",
+                    "timestamp": "2026-03-25T08:00:00+00:00",
+                }
+            ],
+        ),
+    )
+
+    await cache.append_message(
+        thread_id="thread-1",
+        runtime_mode="chat",
+        visibility_mask=1,
+        message={"role": "assistant", "content": "a1"},
+    )
+
+    loaded = await cache.get(
+        thread_id="thread-1",
+        runtime_mode="chat",
+        context_config=config,
+    )
+    assert loaded is not None
+    assert len(loaded) == 1
+
+
+@pytest.mark.asyncio
+async def test_context_messages_cache_append_trims_number_window() -> None:
+    store = _FakeCacheStore()
+    cache = ContextMessagesCache(
+        client=store,
+        key_prefix="agent:context-messages",
+        ttl_seconds=600,
+    )
+    config = MessageContextConfig(
+        window_mode=ContextWindowMode.NUMBER,
+        window_count=1,
+    )
+    await cache.set(
+        thread_id="thread-1",
+        runtime_mode="chat",
+        context_config=config,
+        messages=cast(
+            list[dict[str, Any]],
+            [
+                {
+                    "role": "user",
+                    "content": "q0",
+                    "timestamp": "2026-03-25T08:00:00+00:00",
+                },
+                {
+                    "role": "assistant",
+                    "content": "a0",
+                    "timestamp": "2026-03-25T08:01:00+00:00",
+                },
+                {
+                    "role": "user",
+                    "content": "q1",
+                    "timestamp": "2026-03-25T08:02:00+00:00",
+                },
+            ],
+        ),
+    )
+
+    await cache.append_message(
+        thread_id="thread-1",
+        runtime_mode="chat",
+        visibility_mask=2,
+        message={"role": "assistant", "content": "a1"},
+    )
+
+    loaded = await cache.get(
+        thread_id="thread-1",
+        runtime_mode="chat",
+        context_config=config,
+    )
+    assert loaded is not None
+    assert [str(item["content"]) for item in loaded] == ["q1", "a1"]
+
+
+@pytest.mark.asyncio
+async def test_context_messages_cache_append_trims_day_window() -> None:
+    store = _FakeCacheStore()
+    cache = ContextMessagesCache(
+        client=store,
+        key_prefix="agent:context-messages",
+        ttl_seconds=600,
+    )
+    config = MessageContextConfig(window_mode=ContextWindowMode.DAY, window_count=2)
+
+    now = datetime(2026, 3, 25, 10, 0, tzinfo=timezone.utc)
+    yesterday = now - timedelta(days=1)
+    two_days_ago = now - timedelta(days=2)
+
+    await cache.set(
+        thread_id="thread-1",
+        runtime_mode="chat",
+        context_config=config,
+        messages=cast(
+            list[dict[str, Any]],
+            [
+                {
+                    "role": "assistant",
+                    "content": "d-2",
+                    "timestamp": two_days_ago.isoformat(),
+                },
+                {
+                    "role": "assistant",
+                    "content": "d-1",
+                    "timestamp": yesterday.isoformat(),
+                },
+            ],
+        ),
+    )
+
+    await cache.append_message(
+        thread_id="thread-1",
+        runtime_mode="chat",
+        visibility_mask=2,
+        message={
+            "role": "assistant",
+            "content": "d0",
+            "timestamp": now.isoformat(),
+        },
+    )
+
+    loaded = await cache.get(
+        thread_id="thread-1",
+        runtime_mode="chat",
+        context_config=config,
+    )
+    assert loaded is not None
+    assert [str(item["content"]) for item in loaded] == ["d-1", "d0"]