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

ksadk/memory/adk/backends/http_ltm_backend.py

@@ -12,7 +12,7 @@ import logging
 from typing import List, Optional
 
 import httpx
-from pydantic import ConfigDict, Field
+from pydantic import ConfigDict
 
 from ksadk.memory.adk.backends.base_ltm_backend import BaseLongTermMemoryBackend
 
@@ -52,12 +52,10 @@ class HttpLTMBackend(BaseLongTermMemoryBackend):
     def model_post_init(self, __context) -> None:
         if not self.base_url:
             logger.warning(
-                "HttpLTMBackend: base_url is empty. "
-                "Set KSADK_LTM_HTTP_URL environment variable."
+                "HttpLTMBackend: base_url is empty. " "Set KSADK_LTM_HTTP_URL environment variable."
             )
         logger.info(
-            f"HttpLTMBackend initialized: base_url={self.base_url[:50]}... "
-            f"index={self.index}"
+            f"HttpLTMBackend initialized: base_url={self.base_url[:50]}... " f"index={self.index}"
         )
 
     @property
@@ -75,9 +73,7 @@ class HttpLTMBackend(BaseLongTermMemoryBackend):
             )
         return self._client
 
-    def save_memory(
-        self, user_id: str, event_strings: List[str], **kwargs
-    ) -> bool:
+    def save_memory(self, user_id: str, event_strings: List[str], **kwargs) -> bool:
         """保存记忆到远程服务
 
         TODO: 对接金山云记忆服务 API
@@ -105,24 +101,20 @@ class HttpLTMBackend(BaseLongTermMemoryBackend):
             response.raise_for_status()
 
             logger.info(
-                f"Saved {len(event_strings)} events to remote memory service "
-                f"for user={user_id}"
+                f"Saved {len(event_strings)} events to remote memory service " f"for user={user_id}"
             )
             return True
 
         except httpx.HTTPStatusError as e:
             logger.error(
-                f"HTTP error saving memory: {e.response.status_code} "
-                f"{e.response.text[:200]}"
+                f"HTTP error saving memory: {e.response.status_code} " f"{e.response.text[:200]}"
             )
             return False
         except Exception as e:
             logger.error(f"Error saving memory to remote service: {e}")
             return False
 
-    def search_memory(
-        self, user_id: str, query: str, top_k: int = 5, **kwargs
-    ) -> List[str]:
+    def search_memory(self, user_id: str, query: str, top_k: int = 5, **kwargs) -> List[str]:
         """从远程服务检索记忆
 
         TODO: 对接金山云记忆服务 API
@@ -140,9 +132,7 @@ class HttpLTMBackend(BaseLongTermMemoryBackend):
             }
         """
         if not self.base_url:
-            logger.warning(
-                "HttpLTMBackend: base_url not configured, return empty results."
-            )
+            logger.warning("HttpLTMBackend: base_url not configured, return empty results.")
             return []
 
         try:
@@ -158,7 +148,10 @@ class HttpLTMBackend(BaseLongTermMemoryBackend):
             response.raise_for_status()
 
             data = response.json()
-            memories = data.get("memories", [])
+            raw_memories = data.get("memories", [])
+            memories = (
+                [str(memory) for memory in raw_memories] if isinstance(raw_memories, list) else []
+            )
 
             logger.info(
                 f"Retrieved {len(memories)} memories from remote service "
@@ -168,8 +161,7 @@ class HttpLTMBackend(BaseLongTermMemoryBackend):
 
         except httpx.HTTPStatusError as e:
             logger.error(
-                f"HTTP error searching memory: {e.response.status_code} "
-                f"{e.response.text[:200]}"
+                f"HTTP error searching memory: {e.response.status_code} " f"{e.response.text[:200]}"
             )
             return []
         except Exception as e:

ksadk/memory/adk/backends/inmemory_ltm_backend.py

@@ -8,6 +8,8 @@ import logging
 from collections import defaultdict
 from typing import List
 
+from pydantic import PrivateAttr
+
 from ksadk.memory.adk.backends.base_ltm_backend import BaseLongTermMemoryBackend
 
 logger = logging.getLogger(__name__)
@@ -27,19 +29,13 @@ class InMemoryLTMBackend(BaseLongTermMemoryBackend):
         ```
     """
 
-    # Pydantic v2 不允许直接声明 mutable default，使用 model_post_init
-    _storage: dict = None
+    _storage: defaultdict[str, list[str]] = PrivateAttr(default_factory=lambda: defaultdict(list))
 
     def model_post_init(self, __context) -> None:
         # {user_id: [event_string, ...]}
-        self._storage = defaultdict(list)
-        logger.info(
-            f"InMemoryLTMBackend initialized: index={self.index}"
-        )
+        logger.info(f"InMemoryLTMBackend initialized: index={self.index}")
 
-    def save_memory(
-        self, user_id: str, event_strings: List[str], **kwargs
-    ) -> bool:
+    def save_memory(self, user_id: str, event_strings: List[str], **kwargs) -> bool:
         """保存记忆到内存"""
         if not event_strings:
             return True
@@ -51,9 +47,7 @@ class InMemoryLTMBackend(BaseLongTermMemoryBackend):
         )
         return True
 
-    def search_memory(
-        self, user_id: str, query: str, top_k: int = 5, **kwargs
-    ) -> List[str]:
+    def search_memory(self, user_id: str, query: str, top_k: int = 5, **kwargs) -> List[str]:
         """基于关键词匹配检索记忆
 
         简单实现：对 query 分词后，按匹配关键词数量排序。
@@ -68,7 +62,7 @@ class InMemoryLTMBackend(BaseLongTermMemoryBackend):
         # 按字符分词（支持中英文混合）
         query_terms = query_lower.split()
 
-        scored = []
+        scored: list[tuple[int, str]] = []
         for memory in user_memories:
             memory_lower = memory.lower()
             # 计算匹配分数：完整 query 匹配得高分，部分关键词匹配得低分

ksadk/memory/adk/backends/sdk_ltm_backend.py

@@ -8,8 +8,9 @@
 
 环境变量:
     KSADK_LTM_BACKEND: sdk
-    KSADK_LTM_ACCESS_KEY: AK (可选，默认取 KSYUN_ACCESS_KEY)
-    KSADK_LTM_SECRET_KEY: SK (可选，默认取 KSYUN_SECRET_KEY)
+    KSADK_LTM_ACCESS_KEY: AK (可选，默认取 KSYUN_ACCESS_KEY / KSYUN_ACCESS_KEY_ID)
+    KSADK_LTM_SECRET_KEY: SK (可选，默认取 KSYUN_SECRET_KEY / KSYUN_SECRET_ACCESS_KEY)
+    KSADK_LTM_SESSION_TOKEN: STS 临时会话 token (可选，默认取 KSYUN_SESSION_TOKEN)
     KSADK_LTM_REGION: 区域 (默认 cn-beijing-6)
     KSADK_LTM_ENDPOINT: API 端点 (默认 aicp.api.ksyun.com)
     KSADK_LTM_SCHEME: http/https (默认 https)
@@ -42,6 +43,7 @@ class SdkLTMBackend(BaseLongTermMemoryBackend):
     Attributes:
         access_key: 访问密钥 ID (AK)
         secret_key: 访问密钥 (SK)
+        session_token: STS 临时会话 token
         region: API 区域
         endpoint: API 端点
         scheme: http 或 https
@@ -64,6 +66,7 @@ class SdkLTMBackend(BaseLongTermMemoryBackend):
 
     access_key: str = ""
     secret_key: str = ""
+    session_token: str = ""
     region: str = "cn-beijing-6"
     endpoint: str = "aicp.api.ksyun.com"
     scheme: str = "https"
@@ -105,9 +108,13 @@ class SdkLTMBackend(BaseLongTermMemoryBackend):
             return self._aicp_client
 
         try:
-            from ksyun.common import credential
-            from ksyun.common.profile.client_profile import ClientProfile
-            from ksyun.common.profile.http_profile import HttpProfile
+            from ksyun.common import credential  # type: ignore[import-untyped]
+            from ksyun.common.profile.client_profile import (  # type: ignore[import-untyped]
+                ClientProfile,
+            )
+            from ksyun.common.profile.http_profile import (  # type: ignore[import-untyped]
+                HttpProfile,
+            )
         except ImportError:
             raise ImportError(
                 "kingsoftcloud-sdk-python is required for SDK memory backend. "
@@ -133,7 +140,9 @@ class SdkLTMBackend(BaseLongTermMemoryBackend):
                 "Ensure kingsoftcloud-sdk-python>=1.5.8.94 is installed."
             )
 
-        cred = credential.Credential(self.access_key, self.secret_key)
+        cred = credential.Credential(
+            self.access_key, self.secret_key, self.session_token or None
+        )
 
         http_profile = HttpProfile()
         http_profile.endpoint = self.endpoint
@@ -144,9 +153,7 @@ class SdkLTMBackend(BaseLongTermMemoryBackend):
         client_profile = ClientProfile()
         client_profile.httpProfile = http_profile
 
-        self._aicp_client = aicp_module.AicpClient(
-            cred, self.region, profile=client_profile
-        )
+        self._aicp_client = aicp_module.AicpClient(cred, self.region, profile=client_profile)
 
         # 强制覆写 API 版本为记忆库 API 所需的 2025-11-14
         self._aicp_client._apiVersion = "2025-11-14"
@@ -181,12 +188,14 @@ class SdkLTMBackend(BaseLongTermMemoryBackend):
             if not text:
                 continue
 
-            conversation.append({
-                "Role": role,
-                "CreatedAt": int(time.time() * 1000),
-                "MessageId": str(uuid.uuid4()),
-                "Content": [{"Type": "input_text", "Text": text}],
-            })
+            conversation.append(
+                {
+                    "Role": role,
+                    "CreatedAt": int(time.time() * 1000),
+                    "MessageId": str(uuid.uuid4()),
+                    "Content": [{"Type": "input_text", "Text": text}],
+                }
+            )
         return conversation
 
     def _effective_memory_collection_id(self) -> str:
@@ -195,9 +204,7 @@ class SdkLTMBackend(BaseLongTermMemoryBackend):
     def _effective_scene_id(self) -> str:
         return self.scene_id or DEFAULT_SCENE_ID
 
-    def save_memory(
-        self, user_id: str, event_strings: list[str], **kwargs
-    ) -> bool:
+    def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool:
         """调用 CreateMemorySdk 写入记忆
 
         Args:
@@ -221,7 +228,10 @@ class SdkLTMBackend(BaseLongTermMemoryBackend):
                 logger.info("No valid conversation items to save")
                 return True
 
-            metadata = kwargs.get("metadata") if isinstance(kwargs.get("metadata"), dict) else {}
+            metadata_value = kwargs.get("metadata")
+            metadata: dict[str, Any] = (
+                dict(metadata_value) if isinstance(metadata_value, dict) else {}
+            )
             agent_id = metadata.get("agent_id") or self.agent_id
             session_id = metadata.get("session_id") or kwargs.get("session_id")
             params = {
@@ -235,15 +245,15 @@ class SdkLTMBackend(BaseLongTermMemoryBackend):
                 params["AgentId"] = agent_id
             if session_id:
                 params["SessionId"] = session_id
+            if metadata.get("flush") is True or kwargs.get("flush") is True:
+                params["Flush"] = True
 
             logger.info(
                 f"CreateMemorySdk: memory_collection_id={memory_collection_id}, "
                 f"user_id={user_id}, messages={len(conversation)}"
             )
 
-            response = client.call(
-                "CreateMemorySdk", params, options={"IsPostJson": True}
-            )
+            response = client.call("CreateMemorySdk", params, options={"IsPostJson": True})
             self.last_create_response = self._parse_json_response(response) or {}
             self.last_session_status = {
                 "SessionId": session_id,
@@ -252,8 +262,7 @@ class SdkLTMBackend(BaseLongTermMemoryBackend):
             }
 
             logger.info(
-                f"Saved {len(conversation)} messages to AICP memory service "
-                f"for user={user_id}"
+                f"Saved {len(conversation)} messages to AICP memory service " f"for user={user_id}"
             )
             return True
 
@@ -262,9 +271,7 @@ class SdkLTMBackend(BaseLongTermMemoryBackend):
             logger.error(f"CreateMemorySdk failed: {e}")
             return False
 
-    def search_memory(
-        self, user_id: str, query: str, top_k: int = 5, **kwargs
-    ) -> list[str]:
+    def search_memory(self, user_id: str, query: str, top_k: int = 5, **kwargs) -> list[str]:
         """调用 QueryMemorySdk 检索记忆
 
         Args:
@@ -306,9 +313,7 @@ class SdkLTMBackend(BaseLongTermMemoryBackend):
                 f"user_id={user_id}, query='{query[:50]}'"
             )
 
-            response = client.call(
-                "QueryMemorySdk", params, options={"IsPostJson": True}
-            )
+            response = client.call("QueryMemorySdk", params, options={"IsPostJson": True})
 
             # 解析响应
             memories = self._parse_query_response(response)
@@ -443,12 +448,7 @@ class SdkLTMBackend(BaseLongTermMemoryBackend):
                 parsed.extend(self._parse_memory_item(memory))
             return parsed
 
-        text = (
-            item.get("Content")
-            or item.get("Text")
-            or item.get("Memory")
-            or item.get("Data")
-        )
+        text = item.get("Content") or item.get("Text") or item.get("Memory") or item.get("Data")
         if text is None:
             return []
         if isinstance(text, (dict, list)):

ksadk/memory/adk/long_term_memory.py

@@ -23,19 +23,19 @@
 import json
 import logging
 import os
-from typing import Any, List, Literal
+from typing import Any, List, Literal, Optional, cast
 
-from google.adk.events.event import Event
-from google.adk.memory.base_memory_service import (
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
+from typing_extensions import Union, override
+
+from ksadk.compat.adk_compat import (
     BaseMemoryService,
+    Event,
+    MemoryEntry,
     SearchMemoryResponse,
+    Session,
 )
-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.compat.adk_compat import genai_types as types
 from ksadk.memory.adk.backends.base_ltm_backend import BaseLongTermMemoryBackend
 from ksadk.memory.service import LongTermMemoryService
 
@@ -83,8 +83,8 @@ class LongTermMemory(BaseMemoryService, BaseModel):
     index: str = ""
     app_name: str = ""
 
-    _backend: BaseLongTermMemoryBackend = None
-    _service: LongTermMemoryService = None
+    _backend: Optional[BaseLongTermMemoryBackend] = PrivateAttr(default=None)
+    _service: Optional[LongTermMemoryService] = PrivateAttr(default=None)
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
@@ -105,14 +105,18 @@ class LongTermMemory(BaseMemoryService, BaseModel):
         )
 
     def _get_service(self) -> LongTermMemoryService:
-        if self._service is None:
+        service = self._service
+        if 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
+            service = self._service
+        if service is None:
+            raise RuntimeError("LongTermMemory service initialization failed")
+        if self._backend is not None and service._backend is not self._backend:
+            service.backend = self._backend
+            service._backend = self._backend
+            service.index = getattr(self._backend, "index", service.index)
+            self.index = service.index
+        return service
 
     def _filter_and_convert_events(self, events: List[Event]) -> List[str]:
         """过滤并序列化 session 事件
@@ -157,9 +161,7 @@ class LongTermMemory(BaseMemoryService, BaseModel):
         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}"
-            )
+            logger.info(f"No user events to save for session={session.id}")
             return
 
         logger.info(
@@ -200,9 +202,7 @@ class LongTermMemory(BaseMemoryService, BaseModel):
                 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."
-            )
+            logger.error(f"Error during memory search: {e}. Returning empty results.")
 
         # 转换为 MemoryEntry 格式
         memory_events = []
@@ -213,9 +213,7 @@ class LongTermMemory(BaseMemoryService, BaseModel):
                     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."
-                    )
+                    logger.warning(f"Non-standard memory format: {memory[:100]}. Skipping.")
                     continue
             except json.JSONDecodeError:
                 # 非 JSON 格式的记忆字符串，直接作为 text
@@ -225,9 +223,7 @@ class LongTermMemory(BaseMemoryService, BaseModel):
             memory_events.append(
                 MemoryEntry(
                     author="user",
-                    content=types.Content(
-                        parts=[types.Part(text=text)], role=role
-                    ),
+                    content=types.Content(parts=[types.Part(text=text)], role=role),
                 )
             )
 
@@ -250,22 +246,23 @@ class LongTermMemory(BaseMemoryService, BaseModel):
             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_ACCESS_KEY: SDK AK (fallback to KSYUN_ACCESS_KEY / KSYUN_ACCESS_KEY_ID)
+            KSADK_LTM_SECRET_KEY: SDK SK (fallback to KSYUN_SECRET_KEY / KSYUN_SECRET_ACCESS_KEY)
+            KSADK_LTM_SESSION_TOKEN: SDK STS token (fallback to KSYUN_SESSION_TOKEN)
             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_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)
 
+        if backend_name not in {"local", "http", "sdk"}:
+            raise ValueError(f"Unsupported long-term memory backend: {backend_name}")
+        normalized_backend = cast(Literal["local", "http", "sdk"], backend_name)
         return cls(
-            backend=backend_name,
+            backend=normalized_backend,
             backend_config=backend_config,
             top_k=service.top_k,
             index=service.index,

ksadk/memory/adk/resilient_session_service.py

@@ -4,15 +4,14 @@ import asyncio
 import logging
 from typing import Any, Optional
 
-from google.adk.events.event import Event
-from google.adk.sessions import InMemorySessionService
-from google.adk.sessions.base_session_service import (
+from ksadk.compat.adk_compat import (
     BaseSessionService,
+    Event,
     GetSessionConfig,
+    InMemorySessionService,
     ListSessionsResponse,
+    Session,
 )
-from google.adk.sessions.session import Session
-
 from ksadk.sessions.resilience import is_session_backend_failure
 
 logger = logging.getLogger(__name__)

ksadk/memory/adk/short_term_memory.py

@@ -24,14 +24,15 @@
 
 import logging
 import os
-from typing import Any, Literal, Optional
+from typing import Any, Literal, Optional, cast
 
-from google.adk.sessions import (
+from pydantic import BaseModel, PrivateAttr
+
+from ksadk.compat.adk_compat import (
     BaseSessionService,
     InMemorySessionService,
     Session,
 )
-from pydantic import BaseModel, PrivateAttr
 
 logger = logging.getLogger(__name__)
 
@@ -119,14 +120,9 @@ class ShortTermMemory(BaseModel):
                 logger.info("ShortTermMemory: using InMemorySessionService")
 
             case "sqlite":
-                db_url = _normalize_database_url(
-                    f"sqlite:///{self.local_database_path}"
-                )
+                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}"
-                )
+                logger.info(f"ShortTermMemory: using SQLite at " f"{self.local_database_path}")
 
             case "database":
                 if not self.db_url:
@@ -148,8 +144,7 @@ class ShortTermMemory(BaseModel):
         """初始化数据库 SessionService"""
         normalized_db_url = _normalize_database_url(db_url)
         try:
-            from google.adk.sessions import DatabaseSessionService
-
+            from ksadk.compat.adk_compat import DatabaseSessionService
             from ksadk.memory.adk.resilient_session_service import ResilientADKSessionService
             from ksadk.sessions.resilience import session_backend_timeout_seconds
 
@@ -165,8 +160,7 @@ class ShortTermMemory(BaseModel):
                 DatabaseSessionService(db_url=normalized_db_url, **service_kwargs)
             )
             logger.info(
-                f"ShortTermMemory: using DatabaseSessionService "
-                f"({normalized_db_url[:30]}...)"
+                f"ShortTermMemory: using DatabaseSessionService " f"({normalized_db_url[:30]}...)"
             )
         except ImportError:
             logger.warning(
@@ -218,8 +212,7 @@ class ShortTermMemory(BaseModel):
                 )
                 if session:
                     logger.debug(
-                        f"Session {session_id} already exists "
-                        f"(app={app_name}, user={user_id})"
+                        f"Session {session_id} already exists " f"(app={app_name}, user={user_id})"
                     )
                     return session
             except Exception as e:
@@ -238,10 +231,7 @@ class ShortTermMemory(BaseModel):
                     app_name=app_name,
                     user_id=user_id,
                 )
-            logger.info(
-                f"Created session: id={session.id}, "
-                f"app={app_name}, user={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}")
@@ -297,8 +287,9 @@ class ShortTermMemory(BaseModel):
             else:
                 backend = "local"
 
+        backend_name = cast(Literal["local", "sqlite", "database"], backend)
         return cls(
-            backend=backend,
+            backend=backend_name,
             db_url=db_url,
             local_database_path=db_path,
         )

ksadk/memory/adk_tool.py

@@ -21,11 +21,12 @@ def save_memory(content: str) -> dict:
 
 def create_adk_tool():
     try:
-        from google.adk.tools import FunctionTool
+        from ksadk.compat.adk_compat 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]"
+            "google-adk not installed, returning raw function as tool. "
+            "Install with: pip install ksadk[adk]"
         )
         return save_memory