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"]
@@ -145,6 +145,38 @@ async def test_store_persists_tool_output_with_summary_as_content(
    assert append_kwargs["visibility_mask"] == (1 << 0)


+@pytest.mark.asyncio
+async def test_store_sets_history_only_visibility_for_automation_worker_output(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    captured: dict[str, object] = {}
+    fake_chat_session = SimpleNamespace(state_snapshot={}, message_count=3)
+    _patch_repositories(monkeypatch, captured, fake_chat_session)
+
+    store = store_module.SqlAlchemyEventStore(session_factory=lambda: _FakeSessionCtx())
+    await store.persist(
+        {
+            "type": "TEXT_MESSAGE_END",
+            "threadId": "00000000-0000-0000-0000-000000000001",
+            "runId": "run-auto-1",
+            "messageId": "assistant-auto-1",
+            "role": "assistant",
+            "stage": "worker",
+            "runtime_mode": "automation",
+            "status": "success",
+            "answer": "automation-result",
+            "key_points": [],
+            "result_type": "summary",
+            "suggested_actions": [],
+            "error": None,
+        }
+    )
+
+    append_kwargs = cast(dict[str, Any], captured["append_kwargs"])
+    assert append_kwargs["content"] == "automation-result"
+    assert append_kwargs["visibility_mask"] == (1 << 0)
+
+
@pytest.mark.asyncio
 async def test_store_persists_router_step_output_for_cost_tracking(
    monkeypatch: pytest.MonkeyPatch,
@@ -5,7 +5,7 @@ from uuid import uuid4

 import pytest

-from core.agentscope.persistence.user_context_cache import UserContextCache
+from core.agentscope.caches.user_context_cache import UserContextCache
 from schemas.shared.user import (
    UserContext,
    parse_profile_settings,
@@ -28,6 +28,7 @@ async def test_tool_result_event_uses_runtime_tool_call_id() -> None:
        session_id="thread-1",
        run_id="run-1",
        stage="worker",
+        runtime_mode="chat",
        emit_text_events=False,
        emit_tool_events=True,
    )
@@ -65,3 +66,4 @@ async def test_tool_result_event_uses_runtime_tool_call_id() -> None:
    result_events = [e for e in pipeline.events if e.get("type") == "TOOL_CALL_RESULT"]
    assert len(result_events) == 1
    assert result_events[0]["tool_call_id"] == "runtime-call-123"
+    assert result_events[0]["runtime_mode"] == "chat"
@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from schemas.agent.runtime_models import RouterAgentOutput
+
+
+def test_router_agent_output_coerces_key_entity_value_to_string() -> None:
+    payload = {
+        "normalized_task_input": {
+            "user_text": "test",
+            "multimodal_summary": [],
+            "context_summary": "",
+        },
+        "key_entities": [
+            {
+                "name": "priority",
+                "type": "number",
+                "value": 8,
+            }
+        ],
+        "constraints": [],
+        "task_typing": {
+            "primary": "planning",
+            "secondary": [],
+        },
+        "execution_mode": "onestep",
+        "result_typing": {
+            "primary": "summary",
+            "secondary": [],
+        },
+        "ui": {
+            "ui_mode": "none",
+            "ui_decision_reason": "test",
+        },
+    }
+
+    model = RouterAgentOutput.model_validate(payload)
+
+    assert model.key_entities[0].value == "8"