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

ksadk/memory/adk/short_term_memory.py

@@ -0,0 +1,304 @@
+"""ShortTermMemory - 短期记忆 (会话管理)
+
+封装 Google ADK 的 SessionService，提供统一的会话创建和管理接口。
+通过环境变量配置后端类型，由 ADKRunner 自动注入。
+
+参考 VeADK: veadk/memory/short_term_memory.py
+
+环境变量:
+    KSADK_SESSION_BACKEND: 统一 session 后端类型 (memory / local / sqlite / postgres)
+    KSADK_SESSION_DSN: 统一 session 数据库 DSN
+    KSADK_STM_BACKEND: 旧短期记忆后端类型 (local / sqlite / database)
+    KSADK_STM_DB_URL: 旧数据库连接 URL (sqlite/database 时需要)
+
+使用示例:
+    # InMemory (默认，开发测试)
+    stm = ShortTermMemory(backend="local")
+
+    # SQLite 持久化
+    stm = ShortTermMemory(backend="sqlite")
+
+    # 自定义数据库
+    stm = ShortTermMemory(db_url="postgresql+asyncpg://user:pass@host/db")
+"""
+
+import logging
+import os
+from typing import Any, Literal, Optional
+
+from google.adk.sessions import (
+    BaseSessionService,
+    InMemorySessionService,
+    Session,
+)
+from pydantic import BaseModel, PrivateAttr
+
+logger = logging.getLogger(__name__)
+
+
+def _env_first(*names: str) -> str:
+    for name in names:
+        value = str(os.environ.get(name, "")).strip()
+        if value:
+            return value
+    return ""
+
+
+def _normalize_backend_name(backend: str) -> str:
+    normalized = str(backend or "").strip().lower()
+    if normalized == "memory":
+        return "local"
+    if normalized == "postgres":
+        return "database"
+    return normalized
+
+
+def _normalize_database_url(db_url: str) -> str:
+    """Normalize sqlite URLs for google.adk DatabaseSessionService.
+
+    ADK uses SQLAlchemy async engines underneath, so plain `sqlite:///...`
+    URLs fail while `sqlite+aiosqlite:///...` works.
+    """
+    normalized = str(db_url or "").strip()
+    if not normalized:
+        return ""
+    if normalized.startswith("sqlite+aiosqlite:"):
+        return normalized
+    if normalized.startswith("sqlite:"):
+        return "sqlite+aiosqlite:" + normalized[len("sqlite:") :]
+    return normalized
+
+
+def _session_backend_requires_database_url(backend: str) -> bool:
+    return _normalize_backend_name(backend) == "database"
+
+
+class ShortTermMemory(BaseModel):
+    """短期记忆 - 会话管理
+
+    封装 Google ADK 的 SessionService，所有短期记忆内容（系统提示、
+    历史消息、模型回复）都会发送给模型。
+
+    Attributes:
+        backend: 后端类型
+            - "local": InMemorySessionService (默认，开发测试)
+            - "sqlite": DatabaseSessionService (SQLite 本地持久化)
+            - "database": DatabaseSessionService (自定义数据库 URL)
+        db_url: 数据库连接 URL (backend 为 sqlite/database 时使用)
+        local_database_path: SQLite 文件路径 (backend 为 sqlite 时使用)
+
+    Examples:
+        ```python
+        # InMemory (开发)
+        stm = ShortTermMemory(backend="local")
+
+        # SQLite 持久化
+        stm = ShortTermMemory(backend="sqlite", local_database_path="/tmp/ksadk.db")
+
+        # 从环境变量自动创建
+        stm = ShortTermMemory.from_env()
+        ```
+    """
+
+    backend: Literal["local", "sqlite", "database"] = "local"
+    db_url: str = ""
+    local_database_path: str = "/tmp/ksadk_local_database.db"
+
+    _session_service: BaseSessionService = PrivateAttr()
+
+    def model_post_init(self, __context: Any) -> None:
+        # 优先使用 db_url
+        if self.db_url:
+            logger.info("ShortTermMemory: using db_url (ignoring backend option)")
+            self._init_database_service(self.db_url)
+            return
+
+        match self.backend:
+            case "local":
+                self._session_service = InMemorySessionService()
+                logger.info("ShortTermMemory: using InMemorySessionService")
+
+            case "sqlite":
+                db_url = _normalize_database_url(
+                    f"sqlite:///{self.local_database_path}"
+                )
+                self._init_database_service(db_url)
+                logger.info(
+                    f"ShortTermMemory: using SQLite at "
+                    f"{self.local_database_path}"
+                )
+
+            case "database":
+                if not self.db_url:
+                    raise ValueError(
+                        "KSADK_SESSION_DSN is required when ADK session backend "
+                        "resolves to database/postgres"
+                    )
+                else:
+                    self._init_database_service(self.db_url)
+
+            case _:
+                logger.warning(
+                    f"ShortTermMemory: unknown backend '{self.backend}', "
+                    f"falling back to InMemorySessionService"
+                )
+                self._session_service = InMemorySessionService()
+
+    def _init_database_service(self, db_url: str) -> None:
+        """初始化数据库 SessionService"""
+        normalized_db_url = _normalize_database_url(db_url)
+        try:
+            from google.adk.sessions import DatabaseSessionService
+
+            from ksadk.memory.adk.resilient_session_service import ResilientADKSessionService
+            from ksadk.sessions.resilience import session_backend_timeout_seconds
+
+            service_kwargs = {}
+            if normalized_db_url.startswith(("postgresql://", "postgresql+asyncpg://")):
+                timeout = session_backend_timeout_seconds()
+                service_kwargs["connect_args"] = {
+                    "timeout": timeout,
+                    "command_timeout": timeout,
+                }
+
+            self._session_service = ResilientADKSessionService(
+                DatabaseSessionService(db_url=normalized_db_url, **service_kwargs)
+            )
+            logger.info(
+                f"ShortTermMemory: using DatabaseSessionService "
+                f"({normalized_db_url[:30]}...)"
+            )
+        except ImportError:
+            logger.warning(
+                "DatabaseSessionService not available. "
+                "Install google-adk with database support. "
+                "Falling back to InMemorySessionService."
+            )
+            self._session_service = InMemorySessionService()
+            self.backend = "local"
+        except Exception as e:
+            logger.error(
+                f"Failed to create DatabaseSessionService: {e}. "
+                f"Falling back to InMemorySessionService."
+            )
+            self._session_service = InMemorySessionService()
+            self.backend = "local"
+
+    @property
+    def session_service(self) -> BaseSessionService:
+        """获取底层的 SessionService 实例"""
+        return self._session_service
+
+    async def create_session(
+        self,
+        app_name: str,
+        user_id: str,
+        session_id: str = "",
+    ) -> Optional[Session]:
+        """创建或获取 session
+
+        如果 session_id 对应的 session 已存在，直接返回。
+        否则创建新的 session。
+
+        Args:
+            app_name: 应用名称
+            user_id: 用户 ID
+            session_id: 会话 ID (可选，为空时由 ADK 自动生成)
+
+        Returns:
+            Session 对象，创建失败时返回 None
+        """
+        # 尝试获取已存在的 session
+        if session_id:
+            try:
+                session = await self._session_service.get_session(
+                    app_name=app_name,
+                    user_id=user_id,
+                    session_id=session_id,
+                )
+                if session:
+                    logger.debug(
+                        f"Session {session_id} already exists "
+                        f"(app={app_name}, user={user_id})"
+                    )
+                    return session
+            except Exception as e:
+                logger.debug(f"Session {session_id} not found: {e}")
+
+        # 创建新 session
+        try:
+            if session_id:
+                session = await self._session_service.create_session(
+                    app_name=app_name,
+                    user_id=user_id,
+                    session_id=session_id,
+                )
+            else:
+                session = await self._session_service.create_session(
+                    app_name=app_name,
+                    user_id=user_id,
+                )
+            logger.info(
+                f"Created session: id={session.id}, "
+                f"app={app_name}, user={user_id}"
+            )
+            return session
+        except Exception as e:
+            logger.error(f"Failed to create session: {e}")
+            return None
+
+    @classmethod
+    def from_env(cls) -> "ShortTermMemory":
+        """从环境变量创建 ShortTermMemory
+
+        环境变量:
+            KSADK_ADK_SESSION_BACKEND: ADK 专用 backend
+            KSADK_ADK_SESSION_URL: ADK 专用数据库 URL
+            KSADK_ADK_SESSION_PATH: ADK 专用 SQLite 路径
+            KSADK_STM_BACKEND: 平台级 STM backend
+            KSADK_STM_URL / KSADK_STM_DB_URL: 平台级数据库 URL
+            KSADK_STM_PATH / KSADK_STM_DB_PATH: 平台级 SQLite 路径
+            KSADK_SESSION_BACKEND: 统一 session backend fallback
+            KSADK_SESSION_DSN: 统一 session DSN fallback
+        """
+        explicit_backend = _normalize_backend_name(
+            _env_first(
+                "KSADK_ADK_SESSION_BACKEND",
+                "KSADK_STM_BACKEND",
+                "KSADK_SESSION_BACKEND",
+            )
+        )
+        db_url = _normalize_database_url(
+            _env_first(
+                "KSADK_ADK_SESSION_URL",
+                "KSADK_STM_URL",
+                "KSADK_STM_DB_URL",
+                "KSADK_SESSION_DSN",
+            )
+        )
+        configured_db_path = _env_first(
+            "KSADK_ADK_SESSION_PATH",
+            "KSADK_STM_PATH",
+            "KSADK_STM_DB_PATH",
+        )
+        db_path = configured_db_path or "/tmp/ksadk_local_database.db"
+
+        backend = explicit_backend
+        if _session_backend_requires_database_url(backend) and not db_url:
+            raise ValueError(
+                "KSADK_SESSION_DSN is required when ADK session backend "
+                "resolves to database/postgres"
+            )
+        if not backend:
+            if db_url:
+                backend = "database"
+            elif configured_db_path:
+                backend = "sqlite"
+            else:
+                backend = "local"
+
+        return cls(
+            backend=backend,
+            db_url=db_url,
+            local_database_path=db_path,
+        )

ksadk/memory/adk_tool.py

@@ -0,0 +1,31 @@
+"""ADK wrapper for platform save_memory tool."""
+
+from __future__ import annotations
+
+import logging
+
+from ksadk.memory.tool import save_memory as _save_memory
+
+logger = logging.getLogger(__name__)
+
+
+def save_memory(content: str) -> dict:
+    """保存一条长期记忆。"""
+
+    result = _save_memory(content)
+    return {
+        "result": result.get("message", ""),
+        **result,
+    }
+
+
+def create_adk_tool():
+    try:
+        from google.adk.tools import FunctionTool
+
+        return FunctionTool(func=save_memory)
+    except ImportError:
+        logger.warning(
+            "google-adk not installed, returning raw function as tool. Install with: pip install ksadk[adk]"
+        )
+        return save_memory

ksadk/memory/backends/__init__.py

@@ -0,0 +1,6 @@
+"""Memory Backends Package"""
+
+from ksadk.memory.backends.base import BaseMemoryBackend
+from ksadk.memory.backends.memory import InMemoryBackend
+
+__all__ = ["BaseMemoryBackend", "InMemoryBackend"]

ksadk/memory/backends/base.py

@@ -0,0 +1,103 @@
+"""Base Memory Backend - 抽象基类"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Optional, List, Dict
+from datetime import timedelta
+
+
+class BaseMemoryBackend(ABC):
+    """记忆存储后端抽象基类
+    
+    所有后端实现必须继承此类并实现以下方法。
+    """
+    
+    @abstractmethod
+    def get(self, key: str, session_id: Optional[str] = None) -> Optional[Any]:
+        """获取值
+        
+        Args:
+            key: 键名
+            session_id: 可选的会话 ID，用于隔离不同会话的数据
+            
+        Returns:
+            存储的值，不存在返回 None
+        """
+        pass
+    
+    @abstractmethod
+    def set(
+        self,
+        key: str,
+        value: Any,
+        session_id: Optional[str] = None,
+        ttl: Optional[timedelta] = None,
+    ) -> bool:
+        """设置值
+        
+        Args:
+            key: 键名
+            value: 值 (会自动序列化)
+            session_id: 可选的会话 ID
+            ttl: 可选的过期时间
+            
+        Returns:
+            是否成功
+        """
+        pass
+    
+    @abstractmethod
+    def delete(self, key: str, session_id: Optional[str] = None) -> bool:
+        """删除键
+        
+        Args:
+            key: 键名
+            session_id: 可选的会话 ID
+            
+        Returns:
+            是否成功
+        """
+        pass
+    
+    @abstractmethod
+    def exists(self, key: str, session_id: Optional[str] = None) -> bool:
+        """检查键是否存在"""
+        pass
+    
+    @abstractmethod
+    def clear_session(self, session_id: str) -> bool:
+        """清除整个会话的所有数据"""
+        pass
+    
+    # ===== 消息历史相关 (短期记忆) =====
+    
+    @abstractmethod
+    def add_message(
+        self,
+        session_id: str,
+        role: str,
+        content: str,
+        metadata: Optional[Dict] = None,
+    ) -> bool:
+        """添加消息到会话历史"""
+        pass
+    
+    @abstractmethod
+    def get_messages(
+        self,
+        session_id: str,
+        limit: Optional[int] = None,
+    ) -> List[Dict]:
+        """获取会话消息历史
+        
+        Args:
+            session_id: 会话 ID
+            limit: 最近 N 条消息，None 表示全部
+            
+        Returns:
+            消息列表 [{role, content, timestamp, ...}]
+        """
+        pass
+    
+    def close(self) -> None:
+        """关闭连接 (可选实现)"""
+        pass

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