agentkit-sdk-python 0.8.6__tar.gz → 0.8.8__tar.gz

{agentkit_sdk_python-0.8.6 → agentkit_sdk_python-0.8.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentkit-sdk-python
-Version: 0.8.6
+Version: 0.8.8
 Summary: Python SDK for transforming any AI agent into a production-ready application. Framework-agnostic primitives for runtime, memory, authentication, and tools with volcengine-managed infrastructure.
 Author-email: Xiangrui Cheng <innsdcc@gmail.com>, Yumeng Bao <baoyumeng.123@gmail.com>, Yaozheng Fang <fangyozheng@gmail.com>, Guodong Li <cu.eric.lee@gmail.com>
 License:                                  Apache License

{agentkit_sdk_python-0.8.6 → agentkit_sdk_python-0.8.8}/agentkit/apps/agent_server_app/agent_server_app.py

@@ -126,13 +126,14 @@ def _resolve_memory_service(root_agent: BaseAgent) -> Any:
 
 def _create_a2a_runner(
     *,
-    root_agent: BaseAgent,
+    entry: BaseAgent | App,
     short_term_memory: BaseSessionService | Any | None,
     session_service: BaseSessionService,
     memory_service: Any,
     artifact_service: InMemoryArtifactService,
     credential_service: InMemoryCredentialService,
 ) -> Any:
+    root_agent = entry.root_agent if isinstance(entry, App) else entry
     if VeadkRunner is not None and (
         _is_veadk_agent(root_agent) or _is_veadk_short_term_memory(short_term_memory)
     ):
@@ -145,6 +146,15 @@ def _create_a2a_runner(
             session_service=session_service,
             memory_service=memory_service,
             credential_service=credential_service,
+            **({"plugins": entry.plugins} if isinstance(entry, App) else {}),
+        )
+    if isinstance(entry, App):
+        return Runner(
+            app=entry,
+            session_service=session_service,
+            memory_service=memory_service,
+            artifact_service=artifact_service,
+            credential_service=credential_service,
         )
     return Runner(
         agent=root_agent,
@@ -291,7 +301,7 @@ class AgentkitAgentServerApp(BaseAgentkitApp):
         )
 
         runner = _create_a2a_runner(
-            root_agent=root_agent,
+            entry=entry,
             short_term_memory=short_term_memory,
             session_service=session_service,
             memory_service=memory_service,

{agentkit_sdk_python-0.8.6 → agentkit_sdk_python-0.8.8}/agentkit/apps/langgraph_server_app/langgraph_server_app.py

@@ -8,7 +8,7 @@ import logging
 import os
 import sys
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 from uuid import NAMESPACE_URL, uuid4, uuid5
 
 import uvicorn
@@ -21,6 +21,9 @@ from typing_extensions import override
 
 from agentkit.apps.base_app import BaseAgentkitApp
 
+if TYPE_CHECKING:
+    from agentkit.plugins.llm_shield import LLMShieldPlugin
+
 logger = logging.getLogger(__name__)
 
 _OUTPUT_FIELD_KEYS = ("answer", "output", "final", "response", "text", "messages")
@@ -483,6 +486,7 @@ class AgentkitLangGraphServerApp(BaseAgentkitApp):
         graph_id: str | None = None,
         input_key: str | None = None,
         allow_blocking: bool | None = None,
+        llm_shield: LLMShieldPlugin | None = None,
     ) -> None:
         super().__init__()
         self.config_path = Path(config_path).resolve()
@@ -490,6 +494,7 @@ class AgentkitLangGraphServerApp(BaseAgentkitApp):
         self.config = _normalize_config_paths(self.config_path.parent, raw_config)
         self.graph_id = _resolve_graph_id(self.config, graph_id)
         self.input_key = input_key
+        self.llm_shield = llm_shield
         _extend_python_path(self.config_path.parent, self.config)
         _prepare_langgraph_environment(self.config_path, self.config, allow_blocking=allow_blocking)
         _materialize_lazy_graph_exports(self.config.get("graphs"))
@@ -546,34 +551,79 @@ class AgentkitLangGraphServerApp(BaseAgentkitApp):
         )
 
     async def _stream_events(self, req: AgentkitRunRequest):
+        input_text = _content_text(req.new_message)
+        if self.llm_shield is not None:
+            replacement = await self.llm_shield.moderate_text(input_text, role="user")
+            if replacement is not None:
+                yield _agentkit_text_event(self.graph_id, replacement, partial=False)
+                return
         await self._ensure_thread(req.app_name or self.graph_id, req.user_id, req.session_id)
         accumulated_text = ""
         final_text = ""
-        async for chunk in self._client.runs.stream(
-            **self._run_kwargs(req),
-            stream_mode=["messages", "updates", "values"],
-            version="v2",
-        ):
-            mode, data = _stream_chunk(chunk)
-            if mode == "messages" or mode.startswith("messages/"):
-                text = _stream_message_text(data)
-                delta = _chunk_delta(accumulated_text, text)
-                if delta:
-                    accumulated_text += delta
-                    yield _agentkit_text_event(self.graph_id, delta, partial=True)
-                if text:
-                    final_text = text
-                continue
-            if mode in {"updates", "values"}:
-                interrupt = _interrupt_payload(data)
-                if interrupt is not None:
-                    yield _interrupt_event(self.graph_id, interrupt)
-                    return
-                text = _update_output_text(data)
-                if text:
-                    final_text = text
+        try:
+            async for chunk in self._client.runs.stream(
+                **self._run_kwargs(req),
+                stream_mode=["messages", "updates", "values"],
+                version="v2",
+            ):
+                mode, data = _stream_chunk(chunk)
+                if mode == "messages" or mode.startswith("messages/"):
+                    text = _stream_message_text(data)
+                    delta = _chunk_delta(accumulated_text, text)
+                    if delta:
+                        buffered_text = accumulated_text + delta
+                        if (
+                            self.llm_shield is not None
+                            and len(buffered_text.encode("utf-8")) > self.llm_shield.max_output_bytes
+                        ):
+                            yield _agentkit_text_event(
+                                self.graph_id,
+                                self.llm_shield.unavailable_message,
+                                partial=False,
+                            )
+                            return
+                        accumulated_text = buffered_text
+                        if self.llm_shield is None:
+                            yield _agentkit_text_event(self.graph_id, delta, partial=True)
+                    if text:
+                        final_text = text
+                    continue
+                if mode in {"updates", "values"}:
+                    interrupt = _interrupt_payload(data)
+                    if interrupt is not None:
+                        yield _interrupt_event(self.graph_id, interrupt)
+                        return
+                    text = _update_output_text(data)
+                    if text:
+                        if (
+                            self.llm_shield is not None
+                            and len(text.encode("utf-8")) > self.llm_shield.max_output_bytes
+                        ):
+                            yield _agentkit_text_event(
+                                self.graph_id,
+                                self.llm_shield.unavailable_message,
+                                partial=False,
+                            )
+                            return
+                        final_text = text
+        except Exception as exc:
+            if self.llm_shield is None:
+                raise
+            logger.warning(
+                "LangGraph Server stream failed while LLM Shield was enabled (%s)",
+                type(exc).__name__,
+            )
+            yield _agentkit_text_event(
+                self.graph_id,
+                self.llm_shield.unavailable_message,
+                partial=False,
+            )
+            return
         final_text = final_text or accumulated_text
         if final_text:
+            if self.llm_shield is not None:
+                replacement = await self.llm_shield.moderate_text(final_text, role="assistant")
+                final_text = replacement or final_text
             yield _agentkit_text_event(self.graph_id, final_text, partial=False)
 
     def _attach_agentkit_routes(self) -> None:

{agentkit_sdk_python-0.8.6 → agentkit_sdk_python-0.8.8}/agentkit/auth/admin.py

@@ -47,9 +47,15 @@ POLICY_NAME = "agentkit_cli_sandbox_access"
 SANDBOX_ACCESS_POLICY = "AgentKitSandboxAccess"
 WELL_KNOWN_KEY = ".well-known/agentkit-cli"
 ROLE_ACTIONS = (
-    "agentkit:CreateSession", "agentkit:GetSession", "agentkit:DeleteSession",
-    "agentkit:GetSessionLogs", "agentkit:SetSessionTtl",
-    "agentkit:ListTools", "agentkit:GetTool",
+    "agentkit:CreateSession",
+    "agentkit:GetSession",
+    "agentkit:DeleteSession",
+    "agentkit:PauseSession",
+    "agentkit:ResumeSession",
+    "agentkit:GetSessionLogs",
+    "agentkit:SetSessionTtl",
+    "agentkit:ListTools",
+    "agentkit:GetTool",
 )
 
 

agentkit_sdk_python-0.8.8/agentkit/plugins/__init__.py

@@ -0,0 +1,5 @@
+"""Concrete AgentKit runtime plugins."""
+
+from agentkit.plugins.llm_shield import LLMShieldPlugin
+
+__all__ = ["LLMShieldPlugin"]

agentkit_sdk_python-0.8.8/agentkit/plugins/llm_shield.py

@@ -0,0 +1,392 @@
+"""LLM Shield integration for AgentKit migration runtimes."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Literal
+
+import requests
+from google.adk.events import Event
+from google.adk.plugins import BasePlugin
+from google.genai import types
+from volcengine.Credentials import Credentials as SigningCredentials
+from volcengine.auth.SignerV4 import SignerV4
+from volcengine.base.Request import Request as SigningRequest
+
+from agentkit.platform.configuration import VolcConfiguration
+from agentkit.utils.http_defaults import http_timeout
+
+logger = logging.getLogger(__name__)
+
+LLM_SHIELD_BLOCK_MESSAGE = (
+    "Your request has been blocked by the security policy. "
+    "Please modify it and try again."
+)
+LLM_SHIELD_UNAVAILABLE_MESSAGE = (
+    "The security service is unavailable. Please try again later."
+)
+
+_ACTION = "Moderate"
+_API_VERSION = "2025-08-31"
+_DEFAULT_REGION = "cn-beijing"
+_SERVICE = "llmshield"
+_PATH = "/v2/moderate"
+_INVOCATION_STATE_KEY = "__agentkit_llm_shield__"
+_TRUE_VALUES = frozenset({"1", "true", "yes", "on"})
+_FALSE_VALUES = frozenset({"", "0", "false", "no", "off"})
+
+
+class _Decision(Enum):
+    ALLOW = "allow"
+    BLOCK = "block"
+    ERROR = "error"
+
+
+@dataclass(frozen=True)
+class _ModerationResult:
+    decision: _Decision
+
+
+class LLMShieldPlugin(BasePlugin):
+    """Moderate AgentKit text boundaries with Volcengine LLM Shield.
+
+    Applications opt in by attaching the plugin to the ADK ``App.plugins``
+    boundary; :meth:`from_env` reads the migration runtime environment, where a
+    configured ``TOOL_LLM_SHIELD_APP_ID`` enables moderation and
+    ``ENABLE_LLM_SHIELD`` only overrides that default. Existing AgentKit
+    applications are unchanged until they construct this plugin.
+    """
+
+    unavailable_message = LLM_SHIELD_UNAVAILABLE_MESSAGE
+
+    def __init__(
+        self,
+        *,
+        app_id: str,
+        region: str = _DEFAULT_REGION,
+        timeout: float | None = None,
+        base_url: str | None = None,
+        max_output_bytes: int = 256 * 1024,
+    ) -> None:
+        super().__init__(name="LLMShieldPlugin")
+        if not app_id.strip():
+            raise ValueError("LLM Shield app_id must not be empty")
+        if not region.strip():
+            raise ValueError("LLM Shield region must not be empty")
+        if max_output_bytes <= 0:
+            raise ValueError("max_output_bytes must be greater than zero")
+
+        self.app_id = app_id.strip()
+        self.region = region.strip()
+        self.timeout = timeout if timeout is not None else http_timeout()
+        self.base_url = (
+            base_url or f"https://{self.region}.sdk.access.llm-shield.omini-shield.com"
+        ).rstrip("/")
+        self.max_output_bytes = max_output_bytes
+        self._configuration = VolcConfiguration(region=self.region)
+
+    @classmethod
+    def from_env(cls) -> LLMShieldPlugin | None:
+        """Build the plugin from the migration runtime environment.
+
+        A configured ``TOOL_LLM_SHIELD_APP_ID`` is the enablement signal, so
+        migrations do not need a separate switch. ``ENABLE_LLM_SHIELD`` is an
+        optional override: ``0/false/no/off`` forces the plugin off even when an
+        App ID is configured, while ``1/true/yes/on`` requires an App ID and
+        fails fast without one. Most configurations require only the App ID;
+        request authentication uses AgentKit's AK/SK, STS, or Runtime IAM
+        credential chain.
+        """
+
+        raw_override = os.getenv("ENABLE_LLM_SHIELD", "").strip().lower()
+        app_id = os.getenv("TOOL_LLM_SHIELD_APP_ID", "").strip()
+        if os.getenv("TOOL_LLM_SHIELD_API_KEY", "").strip():
+            logger.warning(
+                "TOOL_LLM_SHIELD_API_KEY is deprecated and ignored; "
+                "LLM Shield authenticates with AgentKit credentials"
+            )
+        forced_off = bool(raw_override) and raw_override in _FALSE_VALUES
+        required_on = raw_override in _TRUE_VALUES
+        if raw_override and not (forced_off or required_on):
+            raise ValueError(
+                "ENABLE_LLM_SHIELD must be one of 1/true/yes/on or 0/false/no/off"
+            )
+        if forced_off:
+            return None
+
+        if not app_id:
+            if required_on:
+                raise ValueError(
+                    "TOOL_LLM_SHIELD_APP_ID is required when ENABLE_LLM_SHIELD=true"
+                )
+            return None
+
+        region = os.getenv("TOOL_LLM_SHIELD_REGION", "").strip()
+        region = region or _DEFAULT_REGION
+
+        return cls(
+            app_id=app_id,
+            region=region,
+        )
+
+    async def moderate_text(
+        self, text: str, *, role: Literal["user", "assistant"]
+    ) -> str | None:
+        """Return replacement text when content must not pass unchanged."""
+
+        if not text:
+            return None
+        if role not in {"user", "assistant"}:
+            raise ValueError("LLM Shield role must be 'user' or 'assistant'")
+        if role == "assistant" and len(text.encode("utf-8")) > self.max_output_bytes:
+            logger.warning("LLM Shield output exceeded the moderation size limit")
+            return LLM_SHIELD_UNAVAILABLE_MESSAGE
+        result = await asyncio.to_thread(self._moderate_sync, text, role)
+        if result.decision is _Decision.ALLOW:
+            return None
+        if result.decision is _Decision.BLOCK:
+            return LLM_SHIELD_BLOCK_MESSAGE
+        return LLM_SHIELD_UNAVAILABLE_MESSAGE
+
+    def _moderate_sync(
+        self, text: str, role: Literal["user", "assistant"]
+    ) -> _ModerationResult:
+        body = json.dumps(
+            {
+                "Message": {
+                    "Role": role,
+                    "Content": text,
+                    "ContentType": 1,
+                },
+                "Scene": self.app_id,
+            },
+            ensure_ascii=False,
+        )
+        try:
+            headers = self._request_headers(body)
+            response = requests.post(
+                f"{self.base_url}{_PATH}",
+                headers=headers,
+                params={"Action": _ACTION, "Version": _API_VERSION},
+                data=body,
+                timeout=self.timeout,
+            )
+        except requests.Timeout:
+            logger.warning("LLM Shield moderation request timed out")
+            return _ModerationResult(_Decision.ERROR)
+        except requests.RequestException:
+            logger.warning("LLM Shield moderation network request failed")
+            return _ModerationResult(_Decision.ERROR)
+        except Exception:
+            logger.warning("LLM Shield moderation authentication failed")
+            return _ModerationResult(_Decision.ERROR)
+
+        if response.status_code != 200:
+            logger.warning(
+                "LLM Shield moderation returned HTTP %s", response.status_code
+            )
+            return _ModerationResult(_Decision.ERROR)
+        try:
+            payload = response.json()
+        except ValueError:
+            logger.warning("LLM Shield moderation returned invalid JSON")
+            return _ModerationResult(_Decision.ERROR)
+
+        if not isinstance(payload, dict):
+            logger.warning("LLM Shield moderation returned an invalid response")
+            return _ModerationResult(_Decision.ERROR)
+        result = payload.get("Result")
+        decision = result.get("Decision") if isinstance(result, dict) else None
+        decision_type = (
+            decision.get("DecisionType") if isinstance(decision, dict) else None
+        )
+        try:
+            normalized_decision = int(decision_type)
+        except (TypeError, ValueError):
+            normalized_decision = None
+        if normalized_decision == 1:
+            return _ModerationResult(_Decision.ALLOW)
+        if normalized_decision == 2:
+            return _ModerationResult(_Decision.BLOCK)
+        logger.warning("LLM Shield moderation returned an unknown decision")
+        return _ModerationResult(_Decision.ERROR)
+
+    def _request_headers(self, body: str) -> dict[str, str]:
+        credentials = self._configuration.get_service_credentials("llm_shield")
+        host = self.base_url.split("://", 1)[-1].split("/", 1)[0]
+        request = SigningRequest()
+        request.set_schema("https")
+        request.set_method("POST")
+        request.set_host(host)
+        request.set_path(_PATH)
+        request.set_headers({"Content-Type": "application/json", "Host": host})
+        request.set_query({"Action": _ACTION, "Version": _API_VERSION})
+        request.set_body(body)
+        SignerV4.sign(
+            request,
+            SigningCredentials(
+                credentials.access_key,
+                credentials.secret_key,
+                _SERVICE,
+                self.region,
+                credentials.session_token or "",
+            ),
+        )
+        request.headers.update({"X-Top-Service": _SERVICE, "X-Top-Region": self.region})
+        return dict(request.headers)
+
+    async def before_run_callback(self, *, invocation_context) -> types.Content | None:
+        text = _content_text(getattr(invocation_context, "user_content", None))
+        replacement = await self.moderate_text(text, role="user")
+        if replacement is None:
+            return None
+        _invocation_state(invocation_context)["terminal"] = True
+        return _model_content(replacement)
+
+    async def before_tool_callback(
+        self,
+        *,
+        tool: Any,
+        tool_args: dict[str, Any],
+        tool_context: Any,
+    ) -> dict[str, Any] | None:
+        del tool, tool_context
+        text = _tool_payload_text(tool_args)
+        if text is None:
+            return {"result": LLM_SHIELD_UNAVAILABLE_MESSAGE}
+        replacement = await self.moderate_text(text, role="user")
+        return {"result": replacement} if replacement is not None else None
+
+    async def after_tool_callback(
+        self,
+        *,
+        tool: Any,
+        tool_args: dict[str, Any],
+        tool_context: Any,
+        result: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        del tool, tool_args, tool_context
+        text = _tool_payload_text(result)
+        if text is None:
+            return {"result": LLM_SHIELD_UNAVAILABLE_MESSAGE}
+        replacement = await self.moderate_text(text, role="assistant")
+        return {"result": replacement} if replacement is not None else None
+
+    async def on_event_callback(
+        self, *, invocation_context, event: Event
+    ) -> Event | None:
+        state = _invocation_state(invocation_context)
+        if state.get("terminal"):
+            return event
+        if getattr(event, "author", None) == "user":
+            return event
+        content = getattr(event, "content", None)
+        if getattr(content, "role", None) == "user":
+            return event
+        text = _content_text(content)
+        if getattr(event, "partial", None) is True:
+            if text:
+                buffered = _merge_stream_text(str(state.get("output", "")), text)
+                if len(buffered.encode("utf-8")) > self.max_output_bytes:
+                    state.pop("output", None)
+                    state["overflowed"] = True
+                else:
+                    state["output"] = buffered
+            return _replace_event_text(event, "") if text else event
+
+        buffered = str(state.pop("output", ""))
+        if state.pop("overflowed", False):
+            return _replace_event_text(event, LLM_SHIELD_UNAVAILABLE_MESSAGE)
+        final_text = _final_stream_text(buffered, text)
+        if not final_text:
+            return event
+        replacement = await self.moderate_text(final_text, role="assistant")
+        return _replace_event_text(event, replacement or final_text)
+
+    async def after_run_callback(self, *, invocation_context) -> None:
+        agent_states = getattr(invocation_context, "agent_states", None)
+        if isinstance(agent_states, dict):
+            agent_states.pop(_INVOCATION_STATE_KEY, None)
+
+
+def _invocation_state(invocation_context: Any) -> dict[str, Any]:
+    """Return state owned by this plugin for one ADK invocation."""
+
+    agent_states = getattr(invocation_context, "agent_states", None)
+    if not isinstance(agent_states, dict):
+        agent_states = {}
+        setattr(invocation_context, "agent_states", agent_states)
+    state = agent_states.get(_INVOCATION_STATE_KEY)
+    if not isinstance(state, dict):
+        state = {}
+        agent_states[_INVOCATION_STATE_KEY] = state
+    return state
+
+
+def _content_text(content: Any) -> str:
+    if content is None:
+        return ""
+    parts = getattr(content, "parts", None) or []
+    texts = [str(part.text) for part in parts if getattr(part, "text", None)]
+    return "\n".join(texts)
+
+
+def _model_content(text: str) -> types.Content:
+    return types.Content(role="model", parts=[types.Part(text=text)])
+
+
+def _tool_payload_text(value: Any) -> str | None:
+    try:
+        return json.dumps(
+            value,
+            ensure_ascii=False,
+            separators=(",", ":"),
+            default=str,
+        )
+    except (TypeError, ValueError, RecursionError):
+        logger.warning("LLM Shield could not serialize a tool payload")
+        return None
+
+
+def _merge_stream_text(buffered: str, text: str) -> str:
+    if not buffered:
+        return text
+    if text.startswith(buffered):
+        return text
+    if buffered.endswith(text):
+        return buffered
+    return buffered + text
+
+
+def _final_stream_text(buffered: str, final_text: str) -> str:
+    if not final_text:
+        return buffered
+    if final_text.startswith(buffered) or buffered.endswith(final_text):
+        return final_text if len(final_text) >= len(buffered) else buffered
+    return final_text
+
+
+def _replace_event_text(event: Event, text: str) -> Event:
+    updated = event.model_copy(deep=True)
+    content = getattr(updated, "content", None)
+    if content is None:
+        updated.content = _model_content(text)
+        return updated
+
+    replaced = False
+    parts = []
+    for part in content.parts or []:
+        if getattr(part, "text", None) is not None:
+            parts.append(part.model_copy(update={"text": text if not replaced else ""}))
+            replaced = True
+        else:
+            parts.append(part)
+    if not replaced:
+        parts.insert(0, types.Part(text=text))
+    updated.content = content.model_copy(update={"parts": parts})
+    return updated