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

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),
                 )
             )
 
@@ -256,16 +252,16 @@ class LongTermMemory(BaseMemoryService, BaseModel):
             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

ksadk/memory/backends/base.py

@@ -1,29 +1,29 @@
 """Base Memory Backend - 抽象基类"""
 
 from abc import ABC, abstractmethod
-from typing import Any, Optional, List, Dict
 from datetime import timedelta
+from typing import Any, Dict, List, Optional
 
 
 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,
@@ -33,43 +33,43 @@ class BaseMemoryBackend(ABC):
         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,
@@ -80,7 +80,7 @@ class BaseMemoryBackend(ABC):
     ) -> bool:
         """添加消息到会话历史"""
         pass
-    
+
     @abstractmethod
     def get_messages(
         self,
@@ -88,16 +88,16 @@ class BaseMemoryBackend(ABC):
         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

@@ -1,49 +1,48 @@
 """In-Memory Backend - 内存存储 (开发/测试用)"""
 
-import json
-from datetime import datetime, timedelta
-from typing import Any, Optional, List, Dict
 from collections import defaultdict
+from datetime import datetime, timedelta
+from typing import Any, Dict, List, Optional
 
 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,
@@ -52,32 +51,32 @@ class InMemoryBackend(BaseMemoryBackend):
         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,
@@ -85,28 +84,28 @@ class InMemoryBackend(BaseMemoryBackend):
         content: str,
         metadata: Optional[Dict] = None,
     ) -> bool:
-        msg = {
+        msg: Dict[str, Any] = {
             "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

@@ -2,20 +2,20 @@
 
 import json
 from datetime import datetime, timedelta
-from typing import Any, Optional, List, Dict
+from typing import Any, Dict, List, Optional
 
 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",
@@ -26,39 +26,40 @@ class RedisBackend(BaseMemoryBackend):
         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,
@@ -67,40 +68,40 @@ class RedisBackend(BaseMemoryBackend):
         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
-    
+        return bool(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
-    
+        return bool(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,
@@ -108,38 +109,38 @@ class RedisBackend(BaseMemoryBackend):
         content: str,
         metadata: Optional[Dict] = None,
     ) -> bool:
-        msg = {
+        msg: dict[str, Any] = {
             "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()

ksadk/memory/langchain_tool.py

@@ -31,7 +31,8 @@ def create_langchain_tools():
         return load_memory_tool, save_memory_tool
     except ImportError:
         logger.warning(
-            "langchain-core not installed, returning raw functions. Install with: pip install langchain-core"
+            "langchain-core not installed, returning raw functions. "
+            "Install with: pip install langchain-core"
         )
         return _load_memory, _save_memory