agentkit-sdk-python 0.8.9__tar.gz → 0.9.1__tar.gz

{agentkit_sdk_python-0.8.9 → agentkit_sdk_python-0.9.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentkit-sdk-python
-Version: 0.8.9
+Version: 0.9.1
 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.9.1/agentkit/frameworks/_agentsentry.py

@@ -0,0 +1,192 @@
+"""Invocation-scoped AgentSentry callbacks for migrated LC/LG agents.
+
+No global monkeypatches: only standard LangChain model/tool callbacks are covered.
+Direct Python side effects inside graph nodes are outside this boundary.
+"""
+
+from contextvars import ContextVar
+from functools import wraps
+import json
+import threading
+
+from agentkit.frameworks._common import adk_event
+
+_active_guard = ContextVar("agentkit_agentsentry_guard", default=None)
+
+
+class AgentSentryBlocked(RuntimeError):
+    """Stop framework execution before a protected tool runs."""
+
+
+def _text(content):
+    if isinstance(content, str):
+        return content
+    return "\n".join(
+        item if isinstance(item, str) else item.get("text", "")
+        for item in (content or [])
+        if isinstance(item, str)
+        or isinstance(item, dict)
+        and item.get("type") == "text"
+        and isinstance(item.get("text"), str)
+    )
+
+
+def _key(name, args):
+    return name, json.dumps(
+        args, sort_keys=True, ensure_ascii=False, separators=(",", ":")
+    )
+
+
+def _make_guard(plugin, ctx):
+    # Optional framework dependency is loaded only for the private V3 path.
+    from langchain_core.callbacks import BaseCallbackHandler
+    from agentkit.plugins.agentsentry_v3 import MISSING_CALL, MISSING_PLAN
+
+    class Guard(BaseCallbackHandler):
+        raise_error = True
+
+        def __init__(self):
+            self.plans = {}
+            self.blocked = None
+            self.lock = threading.RLock()
+            self.session = {"SessionID": ctx.session.id, "RunID": ctx.invocation_id}
+
+        def remember(self, message, completed=()):
+            with self.lock:
+                plan = _text(message.content)
+                for call in getattr(message, "tool_calls", ()) or ():
+                    if call.get("id") in completed:
+                        continue
+                    key = _key(call["name"], call["args"])
+                    self.plans.setdefault(key, []).append((plan, call))
+
+        def on_llm_end(self, response, **kwargs):
+            for generations in response.generations:
+                for generation in generations:
+                    message = getattr(generation, "message", None)
+                    if message is not None:
+                        self.remember(message)
+
+        def on_tool_start(self, serialized, input_str, *, inputs=None, **kwargs):
+            with self.lock:
+                if self.blocked:
+                    raise AgentSentryBlocked(self.blocked)
+                try:
+                    args = inputs if isinstance(inputs, dict) else json.loads(input_str)
+                    matches = self.plans.pop(_key(serialized.get("name"), args), [])
+                except (ValueError, TypeError, AttributeError):
+                    matches = []
+                # Ambiguous or missing model/tool correlation must not allow execution.
+                if len(matches) != 1:
+                    self.blocked = MISSING_CALL
+                elif not matches[0][0].strip():
+                    self.blocked = MISSING_PLAN
+                else:
+                    plan, call = matches[0]
+                    self.blocked = plugin._sentry_v3.moderate(
+                        plan,
+                        role="assistant",
+                        hook_name="before_tool_call",
+                        session_info=self.session,
+                        tool_calls=[
+                            {
+                                "ID": call.get("id", ""),
+                                "Type": "function",
+                                "Function": {
+                                    "Name": call["name"],
+                                    "Arguments": json.dumps(args, ensure_ascii=False),
+                                },
+                            }
+                        ],
+                    )
+                if self.blocked:
+                    raise AgentSentryBlocked(self.blocked)
+
+        def on_chat_model_start(self, *args, **kwargs):
+            if self.blocked:
+                raise AgentSentryBlocked(self.blocked)
+
+        def on_llm_start(self, *args, **kwargs):
+            if self.blocked:
+                raise AgentSentryBlocked(self.blocked)
+
+    return Guard()
+
+
+def guarded_config(config=None):
+    guard = _active_guard.get()
+    if guard is None:
+        return config
+    merged = dict(config or {})
+    callbacks = merged.get("callbacks")
+    if callbacks is None or isinstance(callbacks, list):
+        merged["callbacks"] = [*(callbacks or []), guard]
+    else:
+        callbacks = callbacks.copy()
+        callbacks.add_handler(guard, inherit=True)
+        merged["callbacks"] = callbacks
+    return merged
+
+
+def protect_migrated_tools(method):
+    """Use the V3 plugin already attached to ADK App; public paths are unchanged."""
+
+    @wraps(method)
+    async def wrapped(self, ctx):
+        plugins = getattr(getattr(ctx, "plugin_manager", None), "plugins", ())
+        plugin = next((p for p in plugins if getattr(p, "_sentry_v3", None)), None)
+        if plugin is None:
+            async for event in method(self, ctx):
+                yield event
+            return
+        guard = _make_guard(plugin, ctx)
+        events = method(self, ctx)
+        try:
+            while True:
+                token = _active_guard.set(guard)
+                try:
+                    event = await anext(events)
+                except StopAsyncIteration:
+                    break
+                except AgentSentryBlocked:
+                    if not guard.blocked:
+                        raise
+                    break
+                finally:
+                    _active_guard.reset(token)
+                if guard.blocked:
+                    break
+                yield event
+            if guard.blocked:
+                yield adk_event(ctx, self.name, guard.blocked, partial=False)
+        finally:
+            token = _active_guard.set(guard)
+            try:
+                await events.aclose()
+            finally:
+                _active_guard.reset(token)
+
+    return wrapped
+
+
+async def restore_pending_tools(graph, config):
+    """Rebuild pending calls from this graph/thread's checkpoint on resume."""
+    guard = _active_guard.get()
+    if guard is None:
+        return
+    from agentkit.plugins.agentsentry_v3 import MISSING_CALL
+
+    try:
+        snapshot = await graph.aget_state(config)
+        messages = snapshot.values.get("messages", [])
+        completed = set()
+        for message in reversed(messages):
+            if getattr(message, "type", None) == "tool":
+                completed.add(message.tool_call_id)
+            elif getattr(message, "type", None) == "ai":
+                guard.remember(message, completed)
+                return
+    except (AttributeError, TypeError, ValueError):
+        pass
+    guard.blocked = MISSING_CALL
+    raise AgentSentryBlocked(guard.blocked)

{agentkit_sdk_python-0.8.9 → agentkit_sdk_python-0.9.1}/agentkit/frameworks/langchain.py

@@ -27,6 +27,9 @@ except ImportError:  # pragma: no cover - depends on optional packages.
     HumanMessage = None  # type: ignore[assignment]
 
 
+from agentkit.frameworks._agentsentry import guarded_config, protect_migrated_tools
+
+
 class LangChainAgentkitBridge(BaseAgent):
     """Adapt a LangChain Runnable or callable to AgentKit's ADK runtime boundary."""
 
@@ -59,7 +62,7 @@ class LangChainAgentkitBridge(BaseAgent):
             last_error: Exception | None = None
             for candidate in candidates:
                 try:
-                    return await ainvoke(candidate)
+                    return await ainvoke(candidate, **self._security_kwargs())
                 except Exception as exc:
                     if not is_input_shape_error(exc):
                         raise
@@ -72,7 +75,7 @@ class LangChainAgentkitBridge(BaseAgent):
             last_error = None
             for candidate in candidates:
                 try:
-                    return invoke(candidate)
+                    return invoke(candidate, **self._security_kwargs())
                 except Exception as exc:
                     if not is_input_shape_error(exc):
                         raise
@@ -81,6 +84,10 @@ class LangChainAgentkitBridge(BaseAgent):
                 raise last_error
 
         if callable(self._runnable):
+            if self._security_kwargs():
+                raise UnsupportedFrameworkAgentError(
+                    "AgentSentry tool protection requires a Runnable accepting config callbacks"
+                )
             last_error = None
             for candidate in candidates:
                 try:
@@ -104,7 +111,7 @@ class LangChainAgentkitBridge(BaseAgent):
             for candidate in self._input_candidates(payload, text_input):
                 emitted = False
                 try:
-                    async for chunk in astream(candidate):
+                    async for chunk in astream(candidate, **self._security_kwargs()):
                         emitted = True
                         yield chunk
                     return
@@ -124,7 +131,7 @@ class LangChainAgentkitBridge(BaseAgent):
             for candidate in self._input_candidates(payload, text_input):
                 emitted = False
                 try:
-                    for chunk in stream(candidate):
+                    for chunk in stream(candidate, **self._security_kwargs()):
                         emitted = True
                         yield chunk
                     return
@@ -137,6 +144,12 @@ class LangChainAgentkitBridge(BaseAgent):
             if last_error is not None:
                 raise last_error
 
+    @staticmethod
+    def _security_kwargs():
+        config = guarded_config()
+        return {"config": config} if config is not None else {}
+
+    @protect_migrated_tools
     async def _run_async_impl(
         self,
         ctx: InvocationContext,

{agentkit_sdk_python-0.8.9 → agentkit_sdk_python-0.9.1}/agentkit/frameworks/langgraph.py

@@ -37,6 +37,13 @@ from agentkit.frameworks._common import (
 )
 
 
+from agentkit.frameworks._agentsentry import (
+    guarded_config,
+    protect_migrated_tools,
+    restore_pending_tools,
+)
+
+
 LANGGRAPH_PENDING_INTERRUPT_STATE_KEY = "agentkit:langgraph:pending_interrupt"
 _OUTPUT_FIELD_KEYS = ("answer", "output", "final", "response", "text", "messages")
 
@@ -409,14 +416,17 @@ class LangGraphAgentkitBridge(BaseAgent):
         node = metadata.get("langgraph_node")
         return isinstance(node, str) and node in self._stream_nodes
 
+    @protect_migrated_tools
     async def _run_async_impl(
         self,
         ctx: InvocationContext,
     ) -> AsyncGenerator[Event, None]:
-        config = self._graph_config(ctx)
+        config = guarded_config(self._graph_config(ctx))
         resume_pending = self._pending_interrupt(ctx) is not None
         payload = self._run_payload(ctx)
         graph = await self._graph_for_run(config)
+        if resume_pending:
+            await restore_pending_tools(graph, config)
         accumulated_text = ""
         last_update: Any = None
         pending_update_text = ""

agentkit_sdk_python-0.9.1/agentkit/plugins/agentsentry_v3.py

@@ -0,0 +1,175 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""AgentSentry private Moderate V3 transport (Endpoint HMAC credentials).
+
+Only user/before_dispatch and assistant/before_tool_call are supported.
+No retries: a timeout is ambiguous and must not silently execute a tool.
+"""
+
+import json
+import logging
+import os
+import time
+import uuid
+from urllib.parse import urlsplit
+
+import requests
+from volcengine.auth.SignerV4 import SignerV4
+from volcengine.base.Request import Request
+from volcengine.Credentials import Credentials
+
+PATH = "/OpenTOP/V3/Lumen/Moderate"
+BLOCK = "Your request has been blocked by the security policy."
+UNAVAILABLE = "The security service is unavailable. Please try again later."
+MISSING_PLAN = "Tool execution requires a visible model plan; none was provided."
+UNSUPPORTED_INPUT = (
+    "Unsupported input: AgentSentry text moderation requires visible text."
+)
+MISSING_CALL = "Tool execution could not be correlated with a pending model tool call."
+logger = logging.getLogger(__name__)
+
+
+def enabled_from_env():
+    """AGENTSENTRY_ENABLED overrides only private moderation; default is on."""
+    raw = os.getenv("AGENTSENTRY_ENABLED", "true").strip().lower()
+    if raw not in ("1", "true", "yes", "on", "0", "false", "no", "off"):
+        raise ValueError("AGENTSENTRY_ENABLED must be a boolean")
+    return raw in ("1", "true", "yes", "on")
+
+
+class AgentSentryV3:
+    def __init__(
+        self,
+        url,
+        *,
+        account_id,
+        instance_id,
+        access_key,
+        secret_key,
+        region="cn-private",
+        timeout=30,
+    ):
+        parsed = urlsplit(url)
+        if (
+            parsed.scheme not in ("http", "https")
+            or not parsed.hostname
+            or parsed.username
+            or parsed.password
+            or parsed.fragment
+            or parsed.path.rstrip("/") != PATH
+            or parsed.query not in ("", "Version=2025-01-01")
+        ):
+            raise ValueError("Expected AgentSentry V3 Moderate URL")
+        if not all((account_id, instance_id, access_key, secret_key, region)):
+            raise ValueError(
+                "AgentSentry V3 requires account, instance and Endpoint AK/SK"
+            )
+        self.url = url.split("?", 1)[0].rstrip("/") + "?Version=2025-01-01"
+        self.account_id, self.instance_id = account_id, instance_id
+        self.access_key, self.secret_key = access_key, secret_key
+        self.region, self.timeout = region, timeout
+
+    @classmethod
+    def from_env(cls, url, timeout=30):
+        return cls(
+            url,
+            account_id=os.getenv("AGENTSENTRY_ACCOUNT_ID", ""),
+            instance_id=os.getenv("AGENTSENTRY_INSTANCE_ID", ""),
+            access_key=os.getenv("AGENTSENTRY_AK", ""),
+            secret_key=os.getenv("AGENTSENTRY_SK", ""),
+            region=os.getenv("AGENTSENTRY_REGION", "cn-private"),
+            timeout=timeout,
+        )
+
+    def _headers(self, body):
+        parsed = urlsplit(self.url)
+        request = Request()
+        request.set_schema(parsed.scheme)
+        request.set_method("POST")
+        request.set_host(parsed.netloc.lower())
+        request.set_path(PATH)
+        request.set_query({"Version": "2025-01-01"})
+        request.set_headers(
+            {
+                "Content-Type": "application/json",
+                "Host": parsed.netloc.lower(),
+                "X-Top-Account-Id": self.account_id,
+            }
+        )
+        request.set_body(body.decode("utf-8"))
+        SignerV4.sign(
+            request,
+            Credentials(self.access_key, self.secret_key, "clawsentry", self.region),
+        )
+        return dict(request.headers)
+
+    def moderate(self, text, *, role, hook_name, session_info=None, tool_calls=None):
+        if (role, hook_name) not in (
+            ("user", "before_dispatch"),
+            ("assistant", "before_tool_call"),
+        ):
+            raise ValueError("Unsupported AgentSentry V3 role/hook")
+        if not isinstance(text, str) or not text.strip():
+            return MISSING_PLAN if role == "assistant" else UNSUPPORTED_INPUT
+        session = dict(session_info or {})
+        for key in ("SessionID", "RunID", "MessageID"):
+            session[key] = session.get(key) or str(uuid.uuid4())
+        now_ms = time.time_ns() // 1_000_000
+        message = {"Role": role, "MultiPart": [{"Content": text, "ContentType": 1}]}
+        if tool_calls:
+            message["ToolCall"] = tool_calls
+        payload = {
+            "Base": {"InstanceID": self.instance_id},
+            "Message": message,
+            "SessionInfo": session,
+            "SessionID": session["SessionID"],
+            "LocalTimestamp": now_ms // 1000,
+            "LocalTimestampMS": now_ms,
+            "HookName": hook_name,
+            "IsGroup": False,
+        }
+        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
+        try:
+            response = requests.post(
+                self.url,
+                headers=self._headers(body),
+                data=body,
+                timeout=self.timeout,
+                allow_redirects=False,
+            )
+            if response.status_code != 200:
+                return UNAVAILABLE
+            data = response.json()
+            result = data.get("Result") or {}
+            logger.info(
+                "AgentSentry V3 hook=%s request_id=%s decision=%s degraded=%s",
+                hook_name,
+                (data.get("ResponseMetadata") or {}).get("RequestId"),
+                (result.get("Decision") or {}).get("DecisionType"),
+                result.get("Degraded"),
+            )
+            if (data.get("ResponseMetadata") or {}).get("Error"):
+                return UNAVAILABLE
+            if result.get("Degraded") is not False:
+                return UNAVAILABLE
+            decision = (result.get("Decision") or {}).get("DecisionType")
+            if type(decision) is int and decision == 1:
+                return None
+            if type(decision) is int and decision == 2:
+                return BLOCK
+            # REPLACE/unknown decisions are stopped rather than executing unreviewed args.
+            return UNAVAILABLE
+        except (requests.RequestException, ValueError, TypeError, AttributeError):
+            logger.warning("AgentSentry V3 transport or response error")
+            return UNAVAILABLE