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

ksadk/mcp_runtime/__init__.py

@@ -0,0 +1,231 @@
+from __future__ import annotations
+
+import json
+import os
+import inspect
+from dataclasses import dataclass
+from typing import Any, Sequence
+from urllib.parse import urlparse
+
+import httpx
+from google.adk.tools.mcp_tool.mcp_session_manager import (
+    CheckableMcpHttpClientFactory,
+    StreamableHTTPConnectionParams,
+)
+from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
+
+MCP_TOOLSET_KEY_ATTR = "_ksadk_mcp_toolset_key"
+
+
+@dataclass(frozen=True)
+class MCPServerConfig:
+    name: str
+    url: str
+    api_key: str | None = None
+    tool_filter: tuple[str, ...] = ()
+    tool_name_prefix: str | None = None
+
+    @property
+    def dedupe_key(self) -> str:
+        return f"{self.url}|{self.tool_name_prefix or ''}"
+
+    @property
+    def headers(self) -> dict[str, str] | None:
+        if not self.api_key:
+            return None
+        return {"Authorization": f"Bearer {self.api_key}"}
+
+
+def mcp_tools_enabled() -> bool:
+    value = os.environ.get("KSADK_ENABLE_MCP_TOOLS", "1").strip().lower()
+    return value not in {"0", "false", "no", "off"}
+
+
+def load_mcp_server_configs(raw_value: str | None = None) -> list[MCPServerConfig]:
+    raw = raw_value if raw_value is not None else os.environ.get("KSADK_MCP_SERVERS", "")
+    raw = raw.strip()
+    if not raw:
+        return []
+
+    try:
+        payload = json.loads(raw)
+    except json.JSONDecodeError as exc:
+        raise ValueError("KSADK_MCP_SERVERS must be a valid JSON array") from exc
+
+    if not isinstance(payload, list):
+        raise ValueError("KSADK_MCP_SERVERS must be a JSON array")
+
+    return [_parse_server_config(item, index) for index, item in enumerate(payload)]
+
+
+def build_connection_params(
+    config: MCPServerConfig,
+    *,
+    httpx_client_factory: CheckableMcpHttpClientFactory | None = None,
+) -> StreamableHTTPConnectionParams:
+    kwargs: dict[str, Any] = {"url": config.url, "headers": config.headers}
+    kwargs["httpx_client_factory"] = httpx_client_factory or _default_httpx_client_factory(config.url)
+    return StreamableHTTPConnectionParams(**kwargs)
+
+
+def build_mcp_toolset(
+    config: MCPServerConfig,
+    *,
+    httpx_client_factory: CheckableMcpHttpClientFactory | None = None,
+) -> McpToolset:
+    toolset = McpToolset(
+        connection_params=build_connection_params(
+            config,
+            httpx_client_factory=httpx_client_factory,
+        ),
+        tool_filter=list(config.tool_filter) or None,
+        tool_name_prefix=config.tool_name_prefix,
+    )
+    setattr(toolset, MCP_TOOLSET_KEY_ATTR, config.dedupe_key)
+    setattr(toolset, "_ksadk_mcp_server_name", config.name)
+    return toolset
+
+
+def load_mcp_toolsets_from_env() -> list[McpToolset]:
+    return [build_mcp_toolset(config) for config in load_mcp_server_configs()]
+
+
+def register_mcp_toolset_descriptors(toolset: Any) -> list[str]:
+    """Best-effort registration of MCP tool descriptors for ksadk tool_search.
+
+    Execution stays with the MCP runtime/toolset. This only makes the tools
+    discoverable in ksadk's unified tool catalog.
+    """
+
+    try:
+        from ksadk.toolsets import register_external_tools
+    except Exception:
+        return []
+
+    tools = _mcp_toolset_tools_sync(toolset)
+    if not tools:
+        return []
+    server_name = str(getattr(toolset, "_ksadk_mcp_server_name", "") or "mcp").strip() or "mcp"
+    return register_external_tools(
+        tools,
+        group=f"mcp:{server_name}",
+        boundary="ksadk_managed_mcp_tool",
+        risk_level="medium",
+        enabled=True,
+    )
+
+
+def _default_httpx_client_factory(url: str) -> CheckableMcpHttpClientFactory:
+    hostname = (urlparse(url).hostname or "").lower()
+    trust_env = hostname not in {"localhost", "127.0.0.1", "::1"}
+
+    def _factory(headers=None, timeout=None, auth=None):
+        return httpx.AsyncClient(
+            headers=headers,
+            timeout=timeout,
+            auth=auth,
+            follow_redirects=True,
+            trust_env=trust_env,
+        )
+
+    return _factory
+
+
+def _mcp_toolset_tools_sync(toolset: Any) -> list[Any]:
+    get_tools = getattr(toolset, "get_tools_with_prefix", None)
+    if not callable(get_tools):
+        return []
+    try:
+        result = get_tools()
+    except Exception:
+        return []
+    if inspect.isawaitable(result):
+        try:
+            import asyncio
+
+            try:
+                asyncio.get_running_loop()
+            except RuntimeError:
+                result = asyncio.run(result)
+            else:
+                close = getattr(result, "close", None)
+                if callable(close):
+                    close()
+                return []
+        except Exception:
+            return []
+    if isinstance(result, Sequence) and not isinstance(result, (str, bytes, bytearray)):
+        return list(result)
+    return []
+
+
+def _parse_server_config(item: Any, index: int) -> MCPServerConfig:
+    if not isinstance(item, dict):
+        raise ValueError(f"KSADK_MCP_SERVERS[{index}] must be an object")
+
+    name = _require_non_empty_string(item.get("name"), f"KSADK_MCP_SERVERS[{index}].name")
+    url = _validate_mcp_url(
+        _require_non_empty_string(item.get("url"), f"KSADK_MCP_SERVERS[{index}].url")
+    )
+    api_key = _optional_string(item.get("api_key"))
+    tool_filter = tuple(
+        _string_list(
+            item.get("tool_filter"),
+            f"KSADK_MCP_SERVERS[{index}].tool_filter",
+        )
+    )
+    tool_name_prefix = _optional_string(item.get("tool_name_prefix"))
+
+    return MCPServerConfig(
+        name=name,
+        url=url,
+        api_key=api_key,
+        tool_filter=tool_filter,
+        tool_name_prefix=tool_name_prefix,
+    )
+
+
+def _require_non_empty_string(value: Any, field_name: str) -> str:
+    if not isinstance(value, str) or not value.strip():
+        raise ValueError(f"{field_name} must be a non-empty string")
+    return value.strip()
+
+
+def _optional_string(value: Any) -> str | None:
+    if value is None:
+        return None
+    if not isinstance(value, str):
+        raise ValueError("Optional MCP config fields must be strings")
+    stripped = value.strip()
+    return stripped or None
+
+
+def _string_list(value: Any, field_name: str) -> Sequence[str]:
+    if value is None:
+        return ()
+    if not isinstance(value, list) or any(
+        not isinstance(item, str) or not item.strip() for item in value
+    ):
+        raise ValueError(f"{field_name} must be a list of non-empty strings")
+    return [item.strip() for item in value]
+
+
+def _validate_mcp_url(url: str) -> str:
+    parsed = urlparse(url)
+    if parsed.scheme not in {"http", "https"} or not parsed.netloc:
+        raise ValueError("MCP server url must be an absolute http(s) URL")
+    if not parsed.path.endswith("/mcp"):
+        raise ValueError("MCP server url must point to a /mcp endpoint")
+    return url
+
+
+__all__ = [
+    "MCPServerConfig",
+    "MCP_TOOLSET_KEY_ATTR",
+    "build_connection_params",
+    "build_mcp_toolset",
+    "load_mcp_server_configs",
+    "load_mcp_toolsets_from_env",
+    "mcp_tools_enabled",
+    "register_mcp_toolset_descriptors",
+]

ksadk/memory/__init__.py

@@ -0,0 +1,57 @@
+"""Memory SDK.
+
+模块组成:
+    - ksadk.memory.manager: KV 存储 + 消息历史
+    - ksadk.memory.adk: ADK 专用 STM/LTM 集成
+    - ksadk.memory.service: 平台级长期记忆 service
+    - ksadk.memory.tool: 跨框架 load/save memory 工具
+
+使用示例:
+    from ksadk.memory import load_memory, save_memory
+    from ksadk.memory import LongTermMemoryService
+"""
+
+from typing import TYPE_CHECKING
+
+from ksadk.memory.manager import MemoryManager, get_memory_manager
+from ksadk.memory.backends.base import BaseMemoryBackend
+from ksadk.memory.backends.memory import InMemoryBackend
+
+if TYPE_CHECKING:
+    from ksadk.memory.service import LongTermMemoryService
+
+__all__ = [
+    "MemoryManager",
+    "get_memory_manager",
+    "BaseMemoryBackend",
+    "InMemoryBackend",
+    "LongTermMemoryService",
+    "load_memory",
+    "save_memory",
+    "create_langchain_tools",
+    "create_adk_tool",
+]
+
+
+def create_langchain_tools():
+    from ksadk.memory.langchain_tool import create_langchain_tools as _create
+
+    return _create()
+
+
+def create_adk_tool():
+    from ksadk.memory.adk_tool import create_adk_tool as _create
+
+    return _create()
+
+
+def __getattr__(name):
+    if name in {"load_memory", "save_memory"}:
+        from ksadk.memory.langchain_tool import load_memory, save_memory
+
+        return {"load_memory": load_memory, "save_memory": save_memory}[name]
+    if name == "LongTermMemoryService":
+        from ksadk.memory.service import LongTermMemoryService
+
+        return LongTermMemoryService
+    raise AttributeError(f"module 'ksadk.memory' has no attribute {name!r}")

ksadk/memory/adk/__init__.py

@@ -0,0 +1,51 @@
+"""ADK 记忆体集成 - 短期记忆 + 长期记忆
+
+提供与 Google ADK 原生接口兼容的记忆体管理能力。
+通过环境变量配置，由 ADKRunner 自动注入。
+
+使用示例:
+    from ksadk.memory.adk import ShortTermMemory, LongTermMemory
+
+    # 短期记忆 (会话管理)
+    stm = ShortTermMemory(backend="local")
+
+    # 长期记忆 (跨 session 检索)
+    ltm = LongTermMemory(backend="local", app_name="my_app")
+
+环境变量:
+    # 短期记忆
+    KSADK_ADK_SESSION_BACKEND=database
+    KSADK_ADK_SESSION_URL=           # ADK 数据库 URL
+    KSADK_SESSION_BACKEND=postgres    # unified fallback
+    KSADK_SESSION_DSN=                # unified fallback
+    KSADK_STM_BACKEND=local          # legacy fallback
+    KSADK_STM_DB_URL=                # legacy fallback
+
+    # 长期记忆
+    KSADK_LTM_BACKEND=http           # local | http
+    KSADK_LTM_HTTP_URL=              # 记忆服务 HTTP 地址
+    KSADK_LTM_HTTP_TOKEN=            # 认证 Token
+    KSADK_LTM_TOP_K=5                # 检索数量
+"""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ksadk.memory.adk.long_term_memory import LongTermMemory
+    from ksadk.memory.adk.short_term_memory import ShortTermMemory
+
+
+# Lazy loading (避免在未安装 google-adk 时导入失败)
+def __getattr__(name):
+    if name == "ShortTermMemory":
+        from ksadk.memory.adk.short_term_memory import ShortTermMemory
+        return ShortTermMemory
+
+    if name == "LongTermMemory":
+        from ksadk.memory.adk.long_term_memory import LongTermMemory
+        return LongTermMemory
+
+    raise AttributeError(f"module 'ksadk.memory.adk' has no attribute '{name}'")
+
+
+__all__ = ["ShortTermMemory", "LongTermMemory"]

ksadk/memory/adk/backends/__init__.py

@@ -0,0 +1,5 @@
+"""ADK 长期记忆后端"""
+
+from ksadk.memory.adk.backends.base_ltm_backend import BaseLongTermMemoryBackend
+
+__all__ = ["BaseLongTermMemoryBackend"]

ksadk/memory/adk/backends/base_ltm_backend.py

@@ -0,0 +1,52 @@
+"""长期记忆后端抽象基类
+
+所有长期记忆后端必须继承此类并实现 save_memory / search_memory 方法。
+
+参考 VeADK: veadk/memory/long_term_memory_backends/base_backend.py
+"""
+
+from abc import ABC, abstractmethod
+from typing import List
+
+from pydantic import BaseModel
+
+
+class BaseLongTermMemoryBackend(ABC, BaseModel):
+    """长期记忆存储后端抽象基类
+
+    Attributes:
+        index: 索引/集合名称，用于隔离不同应用的记忆数据
+    """
+
+    index: str = ""
+
+    @abstractmethod
+    def save_memory(
+        self, user_id: str, event_strings: List[str], **kwargs
+    ) -> bool:
+        """保存记忆
+
+        Args:
+            user_id: 用户 ID
+            event_strings: 序列化的事件字符串列表
+
+        Returns:
+            是否保存成功
+        """
+        pass
+
+    @abstractmethod
+    def search_memory(
+        self, user_id: str, query: str, top_k: int = 5, **kwargs
+    ) -> List[str]:
+        """检索记忆
+
+        Args:
+            user_id: 用户 ID
+            query: 查询文本
+            top_k: 返回最相关的 N 条记忆
+
+        Returns:
+            匹配的记忆字符串列表
+        """
+        pass

ksadk/memory/adk/backends/http_ltm_backend.py

@@ -0,0 +1,183 @@
+"""HTTP 长期记忆后端 - 远程记忆服务
+
+通过 HTTP API 连接金山云记忆服务。
+API 对接细节待提供后实现，当前为框架预留。
+
+环境变量:
+    KSADK_LTM_HTTP_URL: 记忆服务 HTTP 地址
+    KSADK_LTM_HTTP_TOKEN: 认证 Token
+"""
+
+import logging
+from typing import List, Optional
+
+import httpx
+from pydantic import ConfigDict, Field
+
+from ksadk.memory.adk.backends.base_ltm_backend import BaseLongTermMemoryBackend
+
+logger = logging.getLogger(__name__)
+
+
+class HttpLTMBackend(BaseLongTermMemoryBackend):
+    """HTTP 远程记忆服务后端
+
+    连接金山云记忆服务，通过 HTTP API 进行记忆的存储和检索。
+
+    Attributes:
+        base_url: 记忆服务 HTTP 基础地址
+        token: 认证 Token
+        timeout: HTTP 请求超时时间（秒）
+
+    Examples:
+        ```python
+        backend = HttpLTMBackend(
+            index="my_app",
+            base_url="https://memory.ksyun.com/api/v1",
+            token="sk-xxx"
+        )
+        backend.save_memory("user_1", ["用户喜欢Python"])
+        results = backend.search_memory("user_1", "编程语言偏好")
+        ```
+    """
+
+    base_url: str = ""
+    token: str = ""
+    timeout: int = 30
+
+    _client: Optional[httpx.Client] = None
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    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."
+            )
+        logger.info(
+            f"HttpLTMBackend initialized: base_url={self.base_url[:50]}... "
+            f"index={self.index}"
+        )
+
+    @property
+    def client(self) -> httpx.Client:
+        """懒加载 HTTP 客户端"""
+        if self._client is None:
+            headers = {"Content-Type": "application/json"}
+            if self.token:
+                headers["Authorization"] = f"Bearer {self.token}"
+
+            self._client = httpx.Client(
+                base_url=self.base_url,
+                headers=headers,
+                timeout=self.timeout,
+            )
+        return self._client
+
+    def save_memory(
+        self, user_id: str, event_strings: List[str], **kwargs
+    ) -> bool:
+        """保存记忆到远程服务
+
+        TODO: 对接金山云记忆服务 API
+        预期请求格式:
+            POST {base_url}/memories
+            {
+                "index": "...",
+                "user_id": "...",
+                "events": ["...", "..."]
+            }
+        """
+        if not self.base_url:
+            logger.warning("HttpLTMBackend: base_url not configured, skip save.")
+            return False
+
+        try:
+            payload = {
+                "index": self.index,
+                "user_id": user_id,
+                "events": event_strings,
+            }
+            payload.update(kwargs)
+
+            response = self.client.post("/memories", json=payload)
+            response.raise_for_status()
+
+            logger.info(
+                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]}"
+            )
+            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]:
+        """从远程服务检索记忆
+
+        TODO: 对接金山云记忆服务 API
+        预期请求格式:
+            POST {base_url}/memories/search
+            {
+                "index": "...",
+                "user_id": "...",
+                "query": "...",
+                "top_k": 5
+            }
+        预期响应格式:
+            {
+                "memories": ["...", "..."]
+            }
+        """
+        if not self.base_url:
+            logger.warning(
+                "HttpLTMBackend: base_url not configured, return empty results."
+            )
+            return []
+
+        try:
+            payload = {
+                "index": self.index,
+                "user_id": user_id,
+                "query": query,
+                "top_k": top_k,
+            }
+            payload.update(kwargs)
+
+            response = self.client.post("/memories/search", json=payload)
+            response.raise_for_status()
+
+            data = response.json()
+            memories = data.get("memories", [])
+
+            logger.info(
+                f"Retrieved {len(memories)} memories from remote service "
+                f"for user={user_id} query='{query[:50]}'"
+            )
+            return memories
+
+        except httpx.HTTPStatusError as e:
+            logger.error(
+                f"HTTP error searching memory: {e.response.status_code} "
+                f"{e.response.text[:200]}"
+            )
+            return []
+        except Exception as e:
+            logger.error(f"Error searching memory from remote service: {e}")
+            return []
+
+    def close(self) -> None:
+        """关闭 HTTP 客户端"""
+        if self._client:
+            self._client.close()
+            self._client = None

ksadk/memory/adk/backends/inmemory_ltm_backend.py

@@ -0,0 +1,96 @@
+"""InMemory 长期记忆后端 - 开发/测试用
+
+使用简单的内存字典存储和文本匹配检索。
+数据在进程退出后丢失，仅适用于开发和测试场景。
+"""
+
+import logging
+from collections import defaultdict
+from typing import List
+
+from ksadk.memory.adk.backends.base_ltm_backend import BaseLongTermMemoryBackend
+
+logger = logging.getLogger(__name__)
+
+
+class InMemoryLTMBackend(BaseLongTermMemoryBackend):
+    """内存长期记忆后端
+
+    使用 dict 存储记忆，文本关键词匹配检索。
+    适用于开发测试，不提供语义搜索能力。
+
+    Examples:
+        ```python
+        backend = InMemoryLTMBackend(index="my_app")
+        backend.save_memory("user_1", ["我喜欢吃冰淇淋", "今天天气真好"])
+        results = backend.search_memory("user_1", "冰淇淋", top_k=5)
+        ```
+    """
+
+    # Pydantic v2 不允许直接声明 mutable default，使用 model_post_init
+    _storage: dict = None
+
+    def model_post_init(self, __context) -> None:
+        # {user_id: [event_string, ...]}
+        self._storage = defaultdict(list)
+        logger.info(
+            f"InMemoryLTMBackend initialized: index={self.index}"
+        )
+
+    def save_memory(
+        self, user_id: str, event_strings: List[str], **kwargs
+    ) -> bool:
+        """保存记忆到内存"""
+        if not event_strings:
+            return True
+
+        self._storage[user_id].extend(event_strings)
+        logger.debug(
+            f"Saved {len(event_strings)} events for user={user_id}, "
+            f"total={len(self._storage[user_id])}"
+        )
+        return True
+
+    def search_memory(
+        self, user_id: str, query: str, top_k: int = 5, **kwargs
+    ) -> List[str]:
+        """基于关键词匹配检索记忆
+
+        简单实现：对 query 分词后，按匹配关键词数量排序。
+        生产环境应使用向量搜索后端。
+        """
+        user_memories = self._storage.get(user_id, [])
+        if not user_memories:
+            return []
+
+        # 简单的关键词匹配打分
+        query_lower = query.lower()
+        # 按字符分词（支持中英文混合）
+        query_terms = query_lower.split()
+
+        scored = []
+        for memory in user_memories:
+            memory_lower = memory.lower()
+            # 计算匹配分数：完整 query 匹配得高分，部分关键词匹配得低分
+            score = 0
+            if query_lower in memory_lower:
+                score += 10  # 完整匹配
+            for term in query_terms:
+                if term in memory_lower:
+                    score += 1
+            if score > 0:
+                scored.append((score, memory))
+
+        # 如果没有匹配，返回最近的记忆
+        if not scored:
+            return user_memories[-top_k:]
+
+        # 按分数降序排列，取 top_k
+        scored.sort(key=lambda x: x[0], reverse=True)
+        results = [item[1] for item in scored[:top_k]]
+
+        logger.debug(
+            f"Search memory for user={user_id} query='{query[:50]}': "
+            f"found {len(results)} results from {len(user_memories)} total"
+        )
+        return results