agentengine-sdk-python 0.6.2__py3-none-any.whl → 0.6.9__py3-none-any.whl

ksadk/memory/backends/memory.py

@@ -0,0 +1,112 @@
+"""In-Memory Backend - 内存存储 (开发/测试用)"""
+
+import json
+from datetime import datetime, timedelta
+from typing import Any, Optional, List, Dict
+from collections import defaultdict
+
+from ksadk.memory.backends.base import BaseMemoryBackend
+
+
+class InMemoryBackend(BaseMemoryBackend):
+    """内存存储后端
+    
+    适用于开发和测试，数据在进程退出后丢失。
+    """
+    
+    def __init__(self):
+        # {session_id: {key: (value, expires_at)}}
+        self._data: Dict[str, Dict[str, tuple]] = defaultdict(dict)
+        # {session_id: [messages]}
+        self._messages: Dict[str, List[Dict]] = defaultdict(list)
+    
+    def _make_key(self, key: str, session_id: Optional[str]) -> tuple:
+        """生成复合键"""
+        sid = session_id or "__global__"
+        return sid, key
+    
+    def _is_expired(self, expires_at: Optional[datetime]) -> bool:
+        if expires_at is None:
+            return False
+        return datetime.utcnow() > expires_at
+    
+    def get(self, key: str, session_id: Optional[str] = None) -> Optional[Any]:
+        sid, k = self._make_key(key, session_id)
+        
+        if sid not in self._data or k not in self._data[sid]:
+            return None
+        
+        value, expires_at = self._data[sid][k]
+        
+        if self._is_expired(expires_at):
+            del self._data[sid][k]
+            return None
+        
+        return value
+    
+    def set(
+        self,
+        key: str,
+        value: Any,
+        session_id: Optional[str] = None,
+        ttl: Optional[timedelta] = None,
+    ) -> bool:
+        sid, k = self._make_key(key, session_id)
+        
+        expires_at = None
+        if ttl:
+            expires_at = datetime.utcnow() + ttl
+        
+        self._data[sid][k] = (value, expires_at)
+        return True
+    
+    def delete(self, key: str, session_id: Optional[str] = None) -> bool:
+        sid, k = self._make_key(key, session_id)
+        
+        if sid in self._data and k in self._data[sid]:
+            del self._data[sid][k]
+            return True
+        return False
+    
+    def exists(self, key: str, session_id: Optional[str] = None) -> bool:
+        return self.get(key, session_id) is not None
+    
+    def clear_session(self, session_id: str) -> bool:
+        if session_id in self._data:
+            del self._data[session_id]
+        if session_id in self._messages:
+            del self._messages[session_id]
+        return True
+    
+    def add_message(
+        self,
+        session_id: str,
+        role: str,
+        content: str,
+        metadata: Optional[Dict] = None,
+    ) -> bool:
+        msg = {
+            "role": role,
+            "content": content,
+            "timestamp": datetime.utcnow().isoformat() + "Z",
+        }
+        if metadata:
+            msg["metadata"] = metadata
+        
+        self._messages[session_id].append(msg)
+        return True
+    
+    def get_messages(
+        self,
+        session_id: str,
+        limit: Optional[int] = None,
+    ) -> List[Dict]:
+        messages = self._messages.get(session_id, [])
+        
+        if limit:
+            return messages[-limit:]
+        return messages.copy()
+    
+    def close(self) -> None:
+        self._data.clear()
+        self._messages.clear()

ksadk/memory/backends/redis.py

@@ -0,0 +1,146 @@
+"""Redis Backend - Redis 存储 (短期记忆)"""
+
+import json
+from datetime import datetime, timedelta
+from typing import Any, Optional, List, Dict
+
+from ksadk.memory.backends.base import BaseMemoryBackend
+
+
+class RedisBackend(BaseMemoryBackend):
+    """Redis 存储后端
+    
+    适用于短期记忆和会话状态存储。
+    
+    使用示例:
+        backend = RedisBackend(url="redis://localhost:6379/0")
+    """
+    
+    def __init__(
+        self,
+        url: str = "redis://localhost:6379/0",
+        prefix: str = "ksadk:memory:",
+        default_ttl: Optional[timedelta] = None,
+    ):
+        self.url = url
+        self.prefix = prefix
+        self.default_ttl = default_ttl
+        self._client = None
+    
+    @property
+    def client(self):
+        """懒加载 Redis 客户端"""
+        if self._client is None:
+            try:
+                import redis
+                self._client = redis.from_url(self.url, decode_responses=True)
+            except ImportError:
+                raise ImportError("redis package required. Install with: pip install redis")
+        return self._client
+    
+    def _make_key(self, key: str, session_id: Optional[str]) -> str:
+        """生成 Redis 键"""
+        if session_id:
+            return f"{self.prefix}session:{session_id}:{key}"
+        return f"{self.prefix}global:{key}"
+    
+    def _messages_key(self, session_id: str) -> str:
+        return f"{self.prefix}messages:{session_id}"
+    
+    def get(self, key: str, session_id: Optional[str] = None) -> Optional[Any]:
+        redis_key = self._make_key(key, session_id)
+        value = self.client.get(redis_key)
+        
+        if value is None:
+            return None
+        
+        try:
+            return json.loads(value)
+        except json.JSONDecodeError:
+            return value
+    
+    def set(
+        self,
+        key: str,
+        value: Any,
+        session_id: Optional[str] = None,
+        ttl: Optional[timedelta] = None,
+    ) -> bool:
+        redis_key = self._make_key(key, session_id)
+        
+        # 序列化
+        if not isinstance(value, str):
+            value = json.dumps(value, ensure_ascii=False)
+        
+        # 设置过期时间
+        ex = None
+        if ttl:
+            ex = int(ttl.total_seconds())
+        elif self.default_ttl:
+            ex = int(self.default_ttl.total_seconds())
+        
+        self.client.set(redis_key, value, ex=ex)
+        return True
+    
+    def delete(self, key: str, session_id: Optional[str] = None) -> bool:
+        redis_key = self._make_key(key, session_id)
+        return self.client.delete(redis_key) > 0
+    
+    def exists(self, key: str, session_id: Optional[str] = None) -> bool:
+        redis_key = self._make_key(key, session_id)
+        return self.client.exists(redis_key) > 0
+    
+    def clear_session(self, session_id: str) -> bool:
+        pattern = f"{self.prefix}session:{session_id}:*"
+        keys = list(self.client.scan_iter(match=pattern))
+        
+        # 也删除消息
+        keys.append(self._messages_key(session_id))
+        
+        if keys:
+            self.client.delete(*keys)
+        return True
+    
+    def add_message(
+        self,
+        session_id: str,
+        role: str,
+        content: str,
+        metadata: Optional[Dict] = None,
+    ) -> bool:
+        msg = {
+            "role": role,
+            "content": content,
+            "timestamp": datetime.utcnow().isoformat() + "Z",
+        }
+        if metadata:
+            msg["metadata"] = metadata
+        
+        redis_key = self._messages_key(session_id)
+        self.client.rpush(redis_key, json.dumps(msg, ensure_ascii=False))
+        
+        # 可以设置消息历史的过期时间
+        if self.default_ttl:
+            self.client.expire(redis_key, int(self.default_ttl.total_seconds()))
+        
+        return True
+    
+    def get_messages(
+        self,
+        session_id: str,
+        limit: Optional[int] = None,
+    ) -> List[Dict]:
+        redis_key = self._messages_key(session_id)
+        
+        if limit:
+            # 获取最后 N 条
+            raw = self.client.lrange(redis_key, -limit, -1)
+        else:
+            raw = self.client.lrange(redis_key, 0, -1)
+        
+        return [json.loads(m) for m in raw]
+    
+    def close(self) -> None:
+        if self._client:
+            self._client.close()
+            self._client = None

ksadk/memory/langchain_tool.py

@@ -0,0 +1,39 @@
+"""LangChain / LangGraph long-term memory tool wrappers."""
+
+from __future__ import annotations
+
+import logging
+
+from ksadk.memory.tool import load_memory as _load_memory
+from ksadk.memory.tool import save_memory as _save_memory
+
+logger = logging.getLogger(__name__)
+
+
+def create_langchain_tools():
+    try:
+        from langchain_core.tools import tool
+
+        @tool
+        def load_memory_tool(query: str) -> str:
+            """检索当前用户的长期记忆。"""
+
+            return _load_memory(query)
+
+        @tool
+        def save_memory_tool(content: str) -> dict:
+            """保存一条长期记忆。"""
+
+            return _save_memory(content)
+
+        load_memory_tool.name = "load_memory"
+        save_memory_tool.name = "save_memory"
+        return load_memory_tool, save_memory_tool
+    except ImportError:
+        logger.warning(
+            "langchain-core not installed, returning raw functions. Install with: pip install langchain-core"
+        )
+        return _load_memory, _save_memory
+
+
+load_memory, save_memory = create_langchain_tools()

ksadk/memory/ltm_backend_factory.py

@@ -0,0 +1,24 @@
+"""Shared long-term memory backend resolution."""
+
+from __future__ import annotations
+
+
+def get_long_term_memory_backend_cls(backend: str) -> type:
+    if backend == "local":
+        from ksadk.memory.adk.backends.inmemory_ltm_backend import InMemoryLTMBackend
+
+        return InMemoryLTMBackend
+
+    if backend == "http":
+        from ksadk.memory.adk.backends.http_ltm_backend import HttpLTMBackend
+
+        return HttpLTMBackend
+
+    if backend == "sdk":
+        from ksadk.memory.adk.backends.sdk_ltm_backend import SdkLTMBackend
+
+        return SdkLTMBackend
+
+    raise ValueError(
+        f"Unsupported long term memory backend: {backend}. Available: local, http, sdk"
+    )

ksadk/memory/manager.py

@@ -0,0 +1,215 @@
+"""Memory Manager - 统一记忆管理接口"""
+
+import os
+import logging
+from datetime import timedelta
+from typing import Any, Optional, List, Dict, Type
+
+from ksadk.memory.backends.base import BaseMemoryBackend
+from ksadk.memory.backends.memory import InMemoryBackend
+
+logger = logging.getLogger(__name__)
+
+
+# 后端注册表
+_BACKENDS: Dict[str, Type[BaseMemoryBackend]] = {
+    "memory": InMemoryBackend,
+}
+
+
+def register_backend(name: str, backend_class: Type[BaseMemoryBackend]):
+    """注册自定义后端"""
+    _BACKENDS[name] = backend_class
+
+
+class MemoryManager:
+    """统一记忆管理器
+    
+    支持可插拔后端，自动从环境变量读取配置。
+    
+    使用示例:
+        # 从环境变量自动配置
+        memory = MemoryManager.from_env()
+        
+        # 手动配置
+        memory = MemoryManager(backend="redis", url="redis://localhost:6379")
+        
+        # 存取数据
+        memory.set("key", {"data": "value"}, session_id="sess-123")
+        data = memory.get("key", session_id="sess-123")
+        
+        # 消息历史
+        memory.add_message("sess-123", "user", "你好")
+        memory.add_message("sess-123", "assistant", "你好！有什么可以帮助你的？")
+        messages = memory.get_messages("sess-123")
+    """
+    
+    def __init__(
+        self,
+        backend: str = "memory",
+        url: Optional[str] = None,
+        prefix: str = "ksadk:memory:",
+        default_ttl: Optional[timedelta] = None,
+        **kwargs,
+    ):
+        """初始化 Memory Manager
+        
+        Args:
+            backend: 后端类型 ("memory", "redis")
+            url: 后端连接 URL
+            prefix: 键名前缀
+            default_ttl: 默认过期时间
+            **kwargs: 传递给后端的额外参数
+        """
+        self.backend_name = backend
+        self._backend = self._create_backend(backend, url, prefix, default_ttl, **kwargs)
+    
+    def _create_backend(
+        self,
+        backend: str,
+        url: Optional[str],
+        prefix: str,
+        default_ttl: Optional[timedelta],
+        **kwargs,
+    ) -> BaseMemoryBackend:
+        """创建后端实例"""
+        # 延迟导入 Redis backend
+        if backend == "redis":
+            from ksadk.memory.backends.redis import RedisBackend
+            _BACKENDS["redis"] = RedisBackend
+        
+        if backend not in _BACKENDS:
+            raise ValueError(f"Unknown backend: {backend}. Available: {list(_BACKENDS.keys())}")
+        
+        backend_class = _BACKENDS[backend]
+        
+        # 根据后端类型传递参数
+        if backend == "memory":
+            return backend_class()
+        elif backend == "redis":
+            return backend_class(url=url, prefix=prefix, default_ttl=default_ttl, **kwargs)
+        else:
+            return backend_class(**kwargs)
+    
+    @classmethod
+    def from_env(cls) -> "MemoryManager":
+        """从环境变量创建 MemoryManager
+        
+        环境变量:
+            KSADK_MEMORY_BACKEND: 后端类型 (默认 "memory")
+            KSADK_MEMORY_URL: 连接 URL (如 redis://localhost:6379)
+            KSADK_MEMORY_PREFIX: 键名前缀 (默认 "ksadk:memory:")
+            KSADK_MEMORY_TTL: 默认 TTL 秒数 (可选)
+        """
+        backend = os.environ.get("KSADK_MEMORY_BACKEND", "memory")
+        url = os.environ.get("KSADK_MEMORY_URL", "")
+        prefix = os.environ.get("KSADK_MEMORY_PREFIX", "ksadk:memory:")
+        
+        ttl = None
+        ttl_str = os.environ.get("KSADK_MEMORY_TTL", "")
+        if ttl_str:
+            ttl = timedelta(seconds=int(ttl_str))
+        
+        logger.debug(f"MemoryManager.from_env: backend={backend}, url={url[:20]}...")
+        
+        return cls(
+            backend=backend,
+            url=url or None,
+            prefix=prefix,
+            default_ttl=ttl,
+        )
+    
+    # ===== 键值操作 =====
+    
+    def get(self, key: str, session_id: Optional[str] = None) -> Optional[Any]:
+        """获取值"""
+        return self._backend.get(key, session_id)
+    
+    def set(
+        self,
+        key: str,
+        value: Any,
+        session_id: Optional[str] = None,
+        ttl: Optional[timedelta] = None,
+    ) -> bool:
+        """设置值"""
+        return self._backend.set(key, value, session_id, ttl)
+    
+    def delete(self, key: str, session_id: Optional[str] = None) -> bool:
+        """删除键"""
+        return self._backend.delete(key, session_id)
+    
+    def exists(self, key: str, session_id: Optional[str] = None) -> bool:
+        """检查键是否存在"""
+        return self._backend.exists(key, session_id)
+    
+    # ===== 会话操作 =====
+    
+    def clear_session(self, session_id: str) -> bool:
+        """清除会话所有数据"""
+        return self._backend.clear_session(session_id)
+    
+    # ===== 消息历史 =====
+    
+    def add_message(
+        self,
+        session_id: str,
+        role: str,
+        content: str,
+        metadata: Optional[Dict] = None,
+    ) -> bool:
+        """添加消息"""
+        return self._backend.add_message(session_id, role, content, metadata)
+    
+    def get_messages(
+        self,
+        session_id: str,
+        limit: Optional[int] = None,
+    ) -> List[Dict]:
+        """获取消息历史"""
+        return self._backend.get_messages(session_id, limit)
+    
+    # ===== 便捷方法 =====
+    
+    def get_state(self, session_id: str) -> Dict:
+        """获取会话状态 (短期记忆)"""
+        return self.get("__state__", session_id) or {}
+    
+    def set_state(self, session_id: str, state: Dict, ttl: Optional[timedelta] = None) -> bool:
+        """设置会话状态"""
+        return self.set("__state__", state, session_id, ttl)
+    
+    def update_state(self, session_id: str, updates: Dict) -> bool:
+        """更新会话状态 (合并)"""
+        current = self.get_state(session_id)
+        current.update(updates)
+        return self.set_state(session_id, current)
+    
+    def close(self) -> None:
+        """关闭连接"""
+        self._backend.close()
+    
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+
+# 全局单例 (懒加载)
+_global_manager: Optional[MemoryManager] = None
+
+
+def get_memory_manager() -> MemoryManager:
+    """获取全局 MemoryManager 单例
+    
+    在 Agent 代码中使用:
+        from ksadk.memory import get_memory_manager
+        
+        memory = get_memory_manager()
+        memory.set("key", "value")
+    """
+    global _global_manager
+    if _global_manager is None:
+        _global_manager = MemoryManager.from_env()
+    return _global_manager