98 lines
3.1 KiB
Python
98 lines
3.1 KiB
Python
"""Working Memory - 基于 Redis 的短期任务记忆"""
|
||
|
||
import json
|
||
import logging
|
||
from datetime import datetime, timezone
|
||
from typing import Any
|
||
|
||
import redis.asyncio as aioredis
|
||
|
||
from agentkit.memory.base import Memory, MemoryItem
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class WorkingMemory(Memory):
|
||
"""Working Memory - 当前任务的上下文和中间状态
|
||
|
||
基于 Redis 实现，支持自动过期（TTL）。
|
||
生命周期：单次任务。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
redis: aioredis.Redis,
|
||
key_prefix: str = "agentkit:working",
|
||
default_ttl: int = 3600,
|
||
):
|
||
self._redis = redis
|
||
self._key_prefix = key_prefix
|
||
self._default_ttl = default_ttl
|
||
|
||
def _make_key(self, key: str) -> str:
|
||
return f"{self._key_prefix}:{key}"
|
||
|
||
async def store(self, key: str, value: Any, metadata: dict[str, Any] | None = None) -> None:
|
||
redis_key = self._make_key(key)
|
||
item = MemoryItem(
|
||
key=key,
|
||
value=value,
|
||
metadata=metadata or {},
|
||
created_at=datetime.now(timezone.utc),
|
||
)
|
||
await self._redis.setex(
|
||
redis_key,
|
||
self._default_ttl,
|
||
json.dumps(item.to_dict(), default=str),
|
||
)
|
||
|
||
async def retrieve(self, key: str) -> MemoryItem | None:
|
||
redis_key = self._make_key(key)
|
||
data = await self._redis.get(redis_key)
|
||
if data is None:
|
||
return None
|
||
item_dict = json.loads(data)
|
||
return MemoryItem(
|
||
key=item_dict["key"],
|
||
value=item_dict["value"],
|
||
metadata=item_dict.get("metadata", {}),
|
||
score=item_dict.get("score", 1.0),
|
||
created_at=datetime.fromisoformat(item_dict["created_at"]) if item_dict.get("created_at") else datetime.now(timezone.utc),
|
||
)
|
||
|
||
async def search(self, query: str, top_k: int = 5, filters: dict[str, Any] | None = None) -> list[MemoryItem]:
|
||
"""Working Memory 不支持语义检索，按 key 前缀匹配"""
|
||
pattern = self._make_key(f"{query}*")
|
||
keys = []
|
||
async for key in self._redis.scan_iter(match=pattern, count=top_k):
|
||
keys.append(key)
|
||
if len(keys) >= top_k:
|
||
break
|
||
|
||
items = []
|
||
for key in keys:
|
||
data = await self._redis.get(key)
|
||
if data:
|
||
item_dict = json.loads(data)
|
||
items.append(MemoryItem(
|
||
key=item_dict["key"],
|
||
value=item_dict["value"],
|
||
metadata=item_dict.get("metadata", {}),
|
||
score=1.0,
|
||
created_at=datetime.now(timezone.utc),
|
||
))
|
||
return items
|
||
|
||
async def delete(self, key: str) -> bool:
|
||
redis_key = self._make_key(key)
|
||
return bool(await self._redis.delete(redis_key))
|
||
|
||
async def clear(self, prefix: str = "") -> int:
|
||
"""清除指定前缀的所有 Working Memory"""
|
||
pattern = self._make_key(f"{prefix}*")
|
||
count = 0
|
||
async for key in self._redis.scan_iter(match=pattern):
|
||
await self._redis.delete(key)
|
||
count += 1
|
||
return count
|