agentengine-sdk-python 0.6.4__py3-none-any.whl → 0.7.0__py3-none-any.whl

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

ksadk/memory/service.py

@@ -0,0 +1,172 @@
+"""Platform-level long-term memory service helpers."""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+from typing import Any
+
+from ksadk.common.aicp_env import resolve_aicp_connection
+from ksadk.memory.adk.backends.base_ltm_backend import BaseLongTermMemoryBackend
+from ksadk.memory.ltm_backend_factory import get_long_term_memory_backend_cls
+
+logger = logging.getLogger(__name__)
+
+
+def format_memory_entries(entries: list[str]) -> str:
+    if not entries:
+        return "未找到相关长期记忆。"
+
+    formatted_entries: list[str] = []
+    for index, entry in enumerate(entries, 1):
+        text = entry
+        try:
+            payload = json.loads(entry)
+        except (json.JSONDecodeError, TypeError):
+            payload = None
+
+        if isinstance(payload, dict):
+            parts = payload.get("parts") or []
+            if parts and isinstance(parts[0], dict):
+                text = str(parts[0].get("text") or text)
+        formatted_entries.append(f"[{index}] {text}")
+
+    return "\n\n".join(formatted_entries)
+
+
+class LongTermMemoryService:
+    def __init__(
+        self,
+        *,
+        backend: str | BaseLongTermMemoryBackend = "local",
+        backend_config: dict[str, Any] | None = None,
+        top_k: int = 5,
+        index: str = "",
+        app_name: str = "",
+    ):
+        self.backend = backend
+        self.backend_config = dict(backend_config or {})
+        self.top_k = top_k
+        self.app_name = app_name
+        backend_index = getattr(backend, "index", "") if isinstance(backend, BaseLongTermMemoryBackend) else ""
+        self.index = index or backend_index or app_name or "default_app"
+        self._backend = self._resolve_backend()
+
+    @classmethod
+    def from_env(
+        cls,
+        *,
+        app_name: str = "",
+        backend: str | None = None,
+    ) -> "LongTermMemoryService":
+        resolved_backend = backend or os.environ.get("KSADK_LTM_BACKEND", "local")
+        top_k = int(os.environ.get("KSADK_LTM_TOP_K", "5"))
+        index = os.environ.get("KSADK_LTM_INDEX", "")
+        app_name = os.environ.get("KSADK_LTM_APP_NAME", "") or app_name
+
+        backend_config: dict[str, Any] = {}
+        if resolved_backend == "http":
+            backend_config = {
+                "base_url": os.environ.get("KSADK_LTM_HTTP_URL", ""),
+                "token": os.environ.get("KSADK_LTM_HTTP_TOKEN", ""),
+            }
+        elif resolved_backend == "sdk":
+            connection = resolve_aicp_connection("KSADK_LTM")
+            backend_config = {
+                "access_key": (
+                    os.environ.get("KSADK_LTM_ACCESS_KEY")
+                    or os.environ.get("KSYUN_ACCESS_KEY", "")
+                ),
+                "secret_key": (
+                    os.environ.get("KSADK_LTM_SECRET_KEY")
+                    or os.environ.get("KSYUN_SECRET_KEY", "")
+                ),
+                "region": connection["region"],
+                "endpoint": connection["endpoint"],
+                "scheme": connection["scheme"],
+                "namespace": os.environ.get("KSADK_LTM_NAMESPACE", ""),
+                "agent_id": os.environ.get("KSADK_LTM_AGENT_ID", ""),
+                "scene_id": os.environ.get("KSADK_LTM_SCENE_ID", "_sys_general"),
+            }
+
+        return cls(
+            backend=resolved_backend,
+            backend_config=backend_config,
+            top_k=top_k,
+            index=index,
+            app_name=app_name,
+        )
+
+    @staticmethod
+    def is_configured() -> bool:
+        return bool(str(os.environ.get("KSADK_LTM_BACKEND", "")).strip())
+
+    def _resolve_backend(self) -> BaseLongTermMemoryBackend:
+        if isinstance(self.backend, BaseLongTermMemoryBackend):
+            return self.backend
+
+        backend_cls = get_long_term_memory_backend_cls(str(self.backend))
+        config = dict(self.backend_config)
+        config.setdefault("index", self.index)
+        return backend_cls(**config)
+
+    def search_entries(self, *, user_id: str, query: str, top_k: int | None = None) -> list[str]:
+        return self._backend.search_memory(
+            user_id=user_id,
+            query=query,
+            top_k=top_k if top_k is not None else self.top_k,
+        )
+
+    def search_text(self, *, user_id: str, query: str, top_k: int | None = None) -> str:
+        try:
+            return format_memory_entries(
+                self.search_entries(user_id=user_id, query=query, top_k=top_k)
+            )
+        except Exception as exc:
+            logger.error("load_memory failed: %s", exc)
+            return f"长期记忆检索失败: {exc}"
+
+    def save_event_strings(
+        self,
+        *,
+        user_id: str,
+        event_strings: list[str],
+        metadata: dict[str, Any] | None = None,
+    ) -> bool:
+        return bool(
+            self._backend.save_memory(
+                user_id=user_id,
+                event_strings=event_strings,
+                metadata=metadata or {},
+            )
+        )
+
+    def save_text(self, *, user_id: str, content: str, metadata: dict[str, Any] | None = None) -> bool:
+        payload = {
+            "role": "user",
+            "parts": [{"text": content}],
+            "metadata": metadata or {},
+        }
+        return self.save_event_strings(
+            user_id=user_id,
+            event_strings=[json.dumps(payload, ensure_ascii=False)],
+            metadata=metadata,
+        )
+
+    def build_context(
+        self,
+        *,
+        user_id: str,
+        query: str,
+        top_k: int | None = None,
+    ) -> dict[str, str] | None:
+        normalized = str(query or "").strip()
+        if not normalized:
+            return None
+        if not self.is_configured():
+            return None
+        return {
+            "query": normalized,
+            "formatted_text": self.search_text(user_id=user_id, query=normalized, top_k=top_k),
+        }

ksadk/memory/tool.py

@@ -0,0 +1,96 @@
+"""Cross-framework long-term memory tools."""
+
+from __future__ import annotations
+
+import logging
+
+from ksadk.memory.service import LongTermMemoryService
+from ksadk.runtime_context import get_current_invocation_context
+
+logger = logging.getLogger(__name__)
+
+_service: LongTermMemoryService | None = None
+
+
+def _get_or_create_service() -> LongTermMemoryService:
+    global _service
+    if _service is None:
+        _service = LongTermMemoryService.from_env()
+    return _service
+
+
+def load_memory(query: str) -> str:
+    context = get_current_invocation_context()
+    if context is None:
+        return "长期记忆检索失败: 缺少运行时上下文。"
+
+    try:
+        service = _get_or_create_service()
+        return service.search_text(user_id=context.user_id, query=query)
+    except Exception as exc:
+        logger.error("load_memory failed: %s", exc)
+        return f"长期记忆检索失败: {exc}"
+
+
+def _is_sdk_backend(service: LongTermMemoryService) -> bool:
+    backend = getattr(service, "_backend", None)
+    return backend.__class__.__name__ == "SdkLTMBackend"
+
+
+def _sdk_pending_result(service: LongTermMemoryService, context) -> dict:
+    backend = getattr(service, "_backend", None)
+    session_status = None
+    get_session_status = getattr(backend, "get_session_status", None)
+    if callable(get_session_status):
+        session_status = get_session_status(
+            user_id=context.user_id,
+            session_id=context.session_id,
+        )
+
+    result = {
+        "ok": False,
+        "status": "accepted_not_extracted",
+        "message": "记忆保存请求已被后端受理，但尚未抽取成可检索记忆。",
+    }
+    if isinstance(session_status, dict):
+        result["session_state"] = session_status.get("State")
+        result["session_id"] = session_status.get("SessionId") or context.session_id
+    elif context.session_id:
+        result["session_id"] = context.session_id
+    return result
+
+
+def save_memory(content: str) -> dict:
+    context = get_current_invocation_context()
+    if context is None:
+        return {
+            "ok": False,
+            "status": "failed",
+            "message": "记忆保存失败: 缺少运行时上下文。",
+        }
+
+    try:
+        service = _get_or_create_service()
+        ok = service.save_text(
+            user_id=context.user_id,
+            content=content,
+            metadata={
+                "agent_id": context.agent_id,
+                "session_id": context.session_id,
+                "runner_type": context.runner_type,
+            },
+        )
+        if ok:
+            if _is_sdk_backend(service):
+                entries = service.search_entries(user_id=context.user_id, query=content, top_k=1)
+                if not entries:
+                    return _sdk_pending_result(service, context)
+            return {"ok": True, "status": "persisted", "message": "记忆已保存。"}
+
+        backend = getattr(service, "_backend", None)
+        last_error = str(getattr(backend, "last_error", "") or "").strip()
+        message = f"记忆保存失败: {last_error}" if last_error else "记忆保存失败。"
+        return {"ok": False, "status": "failed", "message": message}
+    except Exception as exc:
+        logger.error("save_memory failed: %s", exc)
+        return {"ok": False, "status": "failed", "message": f"记忆保存失败: {exc}"}