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

ksadk/memory/adk/long_term_memory.py

@@ -0,0 +1,273 @@
+"""LongTermMemory - 长期记忆服务
+
+实现 Google ADK 的 BaseMemoryService 接口，提供跨 session 的记忆存储和检索。
+通过环境变量配置后端类型和参数，由 ADKRunner 自动注入。
+
+参考 VeADK: veadk/memory/long_term_memory.py
+
+环境变量:
+    KSADK_LTM_BACKEND: 后端类型 (local / http)
+    KSADK_LTM_HTTP_URL: HTTP 记忆服务地址
+    KSADK_LTM_HTTP_TOKEN: HTTP 认证 Token
+    KSADK_LTM_TOP_K: 检索数量 (默认 5)
+
+使用示例:
+    # 通过环境变量自动配置 (推荐)
+    # 设置 KSADK_LTM_BACKEND=local 即可
+
+    # 手动创建
+    from ksadk.memory.adk import LongTermMemory
+    ltm = LongTermMemory(backend="local", app_name="my_app")
+"""
+
+import json
+import logging
+import os
+from typing import Any, List, Literal
+
+from google.adk.events.event import Event
+from google.adk.memory.base_memory_service import (
+    BaseMemoryService,
+    SearchMemoryResponse,
+)
+from google.adk.memory.memory_entry import MemoryEntry
+from google.adk.sessions import Session
+from google.genai import types
+from pydantic import BaseModel, ConfigDict, Field
+from typing_extensions import Union, override
+
+from ksadk.memory.adk.backends.base_ltm_backend import BaseLongTermMemoryBackend
+from ksadk.memory.service import LongTermMemoryService
+
+logger = logging.getLogger(__name__)
+
+
+class LongTermMemory(BaseMemoryService, BaseModel):
+    """长期记忆服务 - 跨 session 的记忆管理
+
+    实现 Google ADK 的 BaseMemoryService 接口，支持:
+    - 将 session 中的用户消息持久化到记忆后端
+    - 通过语义/关键词检索历史记忆
+
+    Attributes:
+        backend: 后端类型或实例 ("local" / "http" / BaseLongTermMemoryBackend)
+        backend_config: 后端配置参数
+        top_k: 检索返回的最大记忆条数
+        index: 索引名称 (用于隔离不同应用的数据)
+        app_name: 应用名称 (备选 index)
+
+    Examples:
+        ```python
+        # 本地开发
+        ltm = LongTermMemory(backend="local", app_name="my_app")
+
+        # 远程服务
+        ltm = LongTermMemory(
+            backend="http",
+            backend_config={
+                "base_url": "https://memory.ksyun.com/api/v1",
+                "token": "sk-xxx",
+            },
+            app_name="my_app"
+        )
+        ```
+    """
+
+    backend: Union[
+        Literal["local", "http", "sdk"],
+        BaseLongTermMemoryBackend,
+    ] = "local"
+
+    backend_config: dict = Field(default_factory=dict)
+    top_k: int = 5
+    index: str = ""
+    app_name: str = ""
+
+    _backend: BaseLongTermMemoryBackend = None
+    _service: LongTermMemoryService = None
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    def model_post_init(self, __context: Any) -> None:
+        self._service = LongTermMemoryService(
+            backend=self.backend,
+            backend_config=self.backend_config,
+            top_k=self.top_k,
+            index=self.index,
+            app_name=self.app_name,
+        )
+        self._backend = self._service._backend
+        self.index = self._service.index
+        logger.info(
+            "LongTermMemory initialized: backend=%s, index=%s",
+            self.backend,
+            self.index,
+        )
+
+    def _get_service(self) -> LongTermMemoryService:
+        if self._service is None:
+            self.model_post_init(None)
+        if self._backend is not None and self._service._backend is not self._backend:
+            self._service.backend = self._backend
+            self._service._backend = self._backend
+            self._service.index = getattr(self._backend, "index", self._service.index)
+            self.index = self._service.index
+        return self._service
+
+    def _filter_and_convert_events(self, events: List[Event]) -> List[str]:
+        """过滤并序列化 session 事件
+
+        参考 VeADK 的逻辑:
+        - 只保留用户消息 (author == "user")
+        - 过滤掉空内容、function call、function response
+        - 序列化为 JSON 字符串
+        """
+        final_events = []
+        for event in events:
+            # 过滤: 无内容的事件
+            if not event.content or not event.content.parts:
+                continue
+
+            # 过滤: 只保留用户事件
+            if event.author != "user":
+                continue
+
+            # 过滤: 去除 function call / function response (无 text)
+            if not event.content.parts[0].text:
+                continue
+
+            # 转换: 序列化为 JSON
+            message = event.content.model_dump(exclude_none=True, mode="json")
+            final_events.append(json.dumps(message, ensure_ascii=False))
+
+        return final_events
+
+    @override
+    async def add_session_to_memory(
+        self,
+        session: Session,
+        **kwargs,
+    ):
+        """将 session 中的用户消息保存到长期记忆
+
+        Args:
+            session: Google ADK Session 对象
+        """
+        user_id = session.user_id
+        event_strings = self._filter_and_convert_events(session.events)
+
+        if not event_strings:
+            logger.info(
+                f"No user events to save for session={session.id}"
+            )
+            return
+
+        logger.info(
+            f"Saving {len(event_strings)} events to long term memory: "
+            f"index={self.index}, user_id={user_id}"
+        )
+
+        self._get_service().save_event_strings(
+            user_id=user_id,
+            event_strings=event_strings,
+            metadata=kwargs.get("metadata"),
+        )
+
+        logger.info(
+            f"Saved {len(event_strings)} events to long term memory: "
+            f"index={self.index}, user_id={user_id}"
+        )
+
+    @override
+    async def search_memory(
+        self, *, app_name: str, user_id: str, query: str
+    ) -> SearchMemoryResponse:
+        """检索长期记忆
+
+        Args:
+            app_name: 应用名称
+            user_id: 用户 ID
+            query: 查询文本
+
+        Returns:
+            SearchMemoryResponse: 包含匹配记忆条目的响应
+        """
+        logger.info(f"Searching memory: query='{query[:50]}' user_id={user_id}")
+
+        memory_chunks = []
+        try:
+            memory_chunks = self._get_service().search_entries(
+                query=query, top_k=self.top_k, user_id=user_id
+            )
+        except Exception as e:
+            logger.error(
+                f"Error during memory search: {e}. Returning empty results."
+            )
+
+        # 转换为 MemoryEntry 格式
+        memory_events = []
+        for memory in memory_chunks:
+            try:
+                memory_dict = json.loads(memory)
+                try:
+                    text = memory_dict["parts"][0]["text"]
+                    role = memory_dict.get("role", "user")
+                except (KeyError, IndexError):
+                    logger.warning(
+                        f"Non-standard memory format: {memory[:100]}. Skipping."
+                    )
+                    continue
+            except json.JSONDecodeError:
+                # 非 JSON 格式的记忆字符串，直接作为 text
+                text = memory
+                role = "user"
+
+            memory_events.append(
+                MemoryEntry(
+                    author="user",
+                    content=types.Content(
+                        parts=[types.Part(text=text)], role=role
+                    ),
+                )
+            )
+
+        logger.info(
+            f"Found {len(memory_events)} memory entries for "
+            f"query='{query[:50]}' index={self.index} user_id={user_id}"
+        )
+        return SearchMemoryResponse(memories=memory_events)
+
+    @classmethod
+    def from_env(
+        cls,
+        *,
+        app_name: str = "",
+        backend: str | None = None,
+    ) -> "LongTermMemory":
+        """从环境变量创建 LongTermMemory
+
+        环境变量:
+            KSADK_LTM_BACKEND: local / http / sdk
+            KSADK_LTM_HTTP_URL: HTTP 记忆服务地址
+            KSADK_LTM_HTTP_TOKEN: HTTP 认证 Token
+            KSADK_LTM_ACCESS_KEY: SDK AK (fallback to KSYUN_ACCESS_KEY)
+            KSADK_LTM_SECRET_KEY: SDK SK (fallback to KSYUN_SECRET_KEY)
+            KSADK_LTM_TOP_K: 检索数量 (默认 5)
+            KSADK_LTM_INDEX: 索引名称
+        """
+        service = LongTermMemoryService.from_env(app_name=app_name, backend=backend)
+        backend_name = (
+            backend
+            or os.environ.get("KSADK_LTM_BACKEND", "local")
+        )
+        backend_config = dict(service.backend_config)
+        if "index" in backend_config and backend_name != "local":
+            backend_config.pop("index", None)
+
+        return cls(
+            backend=backend_name,
+            backend_config=backend_config,
+            top_k=service.top_k,
+            index=service.index,
+            app_name=service.app_name,
+        )

ksadk/memory/adk/short_term_memory.py

@@ -0,0 +1,289 @@
+"""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(
+                f"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
+
+            self._session_service = DatabaseSessionService(db_url=normalized_db_url)
+            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()
+        except Exception as e:
+            logger.error(
+                f"Failed to create DatabaseSessionService: {e}. "
+                f"Falling back to InMemorySessionService."
+            )
+            self._session_service = InMemorySessionService()
+
+    @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