hexgate 0.1.1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- fortify/__init__.py +78 -0
- fortify/adapters/__init__.py +0 -0
- fortify/adapters/google/__init__.py +3 -0
- fortify/adapters/google/runner.py +106 -0
- fortify/adapters/google/tools.py +57 -0
- fortify/adapters/google/wrapper.py +47 -0
- fortify/adapters/langchain/__init__.py +7 -0
- fortify/adapters/langchain/agent.py +133 -0
- fortify/adapters/langchain/tools.py +248 -0
- fortify/adapters/langchain/wrapper.py +70 -0
- fortify/adapters/openai/__init__.py +3 -0
- fortify/adapters/openai/runner.py +125 -0
- fortify/adapters/openai/tools.py +57 -0
- fortify/adapters/openai/wrapper.py +47 -0
- fortify/adapters/pydantic_ai/__init__.py +7 -0
- fortify/adapters/pydantic_ai/agent.py +111 -0
- fortify/adapters/pydantic_ai/tools.py +40 -0
- fortify/adapters/pydantic_ai/wrapper.py +94 -0
- fortify/agents/__init__.py +57 -0
- fortify/agents/builtin/__init__.py +1 -0
- fortify/agents/builtin/researcher/agent.yaml +7 -0
- fortify/agents/builtin/researcher/policy.yaml +10 -0
- fortify/agents/builtin/researcher/system.md +5 -0
- fortify/agents/factory.py +683 -0
- fortify/agents/loader.py +790 -0
- fortify/agents/models.py +17 -0
- fortify/agents/prompts/agent_system.md +14 -0
- fortify/audit.py +292 -0
- fortify/bootstrap.py +32 -0
- fortify/cli/__init__.py +42 -0
- fortify/cli/_common.py +306 -0
- fortify/cli/chat.py +282 -0
- fortify/cli/policy/__init__.py +17 -0
- fortify/cli/policy/main.py +527 -0
- fortify/cli/register/__init__.py +8 -0
- fortify/cli/register/fortify.py +139 -0
- fortify/cli/register/google.py +104 -0
- fortify/cli/register/langchain.py +76 -0
- fortify/cli/register/main.py +113 -0
- fortify/cli/register/manifest.py +65 -0
- fortify/cli/register/models.py +107 -0
- fortify/cli/register/openai.py +78 -0
- fortify/cli/register/pydantic_ai.py +114 -0
- fortify/cli/register/register.py +71 -0
- fortify/cli/serve.py +337 -0
- fortify/cli/state.py +156 -0
- fortify/cloud/__init__.py +32 -0
- fortify/cloud/attenuate.py +105 -0
- fortify/cloud/biscuit.py +172 -0
- fortify/cloud/client.py +327 -0
- fortify/config/__init__.py +1 -0
- fortify/config/settings.py +50 -0
- fortify/runtime/__init__.py +53 -0
- fortify/runtime/command_policy.py +289 -0
- fortify/runtime/context.py +129 -0
- fortify/runtime/sandbox_runtime.py +115 -0
- fortify/runtime/srt.py +88 -0
- fortify/runtime/workspace.py +330 -0
- fortify/security/__init__.py +141 -0
- fortify/security/bundle.py +399 -0
- fortify/security/constraints.py +252 -0
- fortify/security/decision.py +145 -0
- fortify/security/enforcer.py +83 -0
- fortify/security/errors.py +11 -0
- fortify/security/file_scope.py +78 -0
- fortify/security/models.py +59 -0
- fortify/security/policy.py +182 -0
- fortify/security/policy_set.py +266 -0
- fortify/security/rego.py +334 -0
- fortify/security/rego_wasm.py +213 -0
- fortify/security/signing.py +122 -0
- fortify/security/source.py +372 -0
- fortify/security/wasm_engine.py +486 -0
- fortify/streaming/__init__.py +57 -0
- fortify/streaming/events.py +206 -0
- fortify/streaming/normalize.py +429 -0
- fortify/tools/__init__.py +21 -0
- fortify/tools/bash.py +49 -0
- fortify/tools/decorators.py +158 -0
- fortify/tools/fetch.py +72 -0
- fortify/tools/files/__init__.py +9 -0
- fortify/tools/files/_common.py +35 -0
- fortify/tools/files/edit_file.py +57 -0
- fortify/tools/files/glob.py +48 -0
- fortify/tools/files/grep.py +91 -0
- fortify/tools/files/read_file.py +39 -0
- fortify/tools/files/write_file.py +36 -0
- fortify/tools/refund.py +53 -0
- fortify/tools/websearch.py +72 -0
- fortify/tracing/__init__.py +1 -0
- fortify/tracing/langfuse.py +68 -0
- fortify/utils/__init__.py +1 -0
- fortify/utils/retry.py +58 -0
- hexgate-0.1.1.dist-info/METADATA +1440 -0
- hexgate-0.1.1.dist-info/RECORD +98 -0
- hexgate-0.1.1.dist-info/WHEEL +5 -0
- hexgate-0.1.1.dist-info/entry_points.txt +2 -0
- hexgate-0.1.1.dist-info/top_level.txt +1 -0
fortify/__init__.py
ADDED
|
@@ -0,0 +1,78 @@
|
|
|
1
|
+
"""Public package surface for fortify."""
|
|
2
|
+
|
|
3
|
+
from fortify.agents.factory import (
|
|
4
|
+
create_agent,
|
|
5
|
+
enforce_policy,
|
|
6
|
+
invoke_agent,
|
|
7
|
+
stream_agent,
|
|
8
|
+
stream_agent_raw,
|
|
9
|
+
)
|
|
10
|
+
from fortify.agents.loader import (
|
|
11
|
+
clear_registered_agents,
|
|
12
|
+
list_available_agents,
|
|
13
|
+
list_builtin_agents,
|
|
14
|
+
list_local_agents,
|
|
15
|
+
list_registered_agents,
|
|
16
|
+
load_agent,
|
|
17
|
+
load_builtin_agent,
|
|
18
|
+
load_fortify_agent,
|
|
19
|
+
load_local_agent,
|
|
20
|
+
load_registered_agent,
|
|
21
|
+
register_agent,
|
|
22
|
+
unregister_agent,
|
|
23
|
+
)
|
|
24
|
+
from fortify.cli.register import AgentManifest, create_manifest
|
|
25
|
+
from fortify.cloud import FortifyClient, FortifyConfig
|
|
26
|
+
from fortify.runtime import LocalWorkspace, ToolUseContext, User, Workspace
|
|
27
|
+
from fortify.security import AgentPolicy
|
|
28
|
+
from fortify.tools import (
|
|
29
|
+
agent_tool,
|
|
30
|
+
bash,
|
|
31
|
+
edit_file,
|
|
32
|
+
glob,
|
|
33
|
+
grep,
|
|
34
|
+
read_file,
|
|
35
|
+
refund_order,
|
|
36
|
+
write_file,
|
|
37
|
+
)
|
|
38
|
+
from fortify.tools.fetch import fetch
|
|
39
|
+
from fortify.tools.websearch import web_search
|
|
40
|
+
|
|
41
|
+
__all__ = [
|
|
42
|
+
"AgentManifest",
|
|
43
|
+
"AgentPolicy",
|
|
44
|
+
"FortifyClient",
|
|
45
|
+
"FortifyConfig",
|
|
46
|
+
"LocalWorkspace",
|
|
47
|
+
"ToolUseContext",
|
|
48
|
+
"User",
|
|
49
|
+
"Workspace",
|
|
50
|
+
"agent_tool",
|
|
51
|
+
"bash",
|
|
52
|
+
"edit_file",
|
|
53
|
+
"clear_registered_agents",
|
|
54
|
+
"create_agent",
|
|
55
|
+
"create_manifest",
|
|
56
|
+
"enforce_policy",
|
|
57
|
+
"fetch",
|
|
58
|
+
"glob",
|
|
59
|
+
"grep",
|
|
60
|
+
"invoke_agent",
|
|
61
|
+
"list_available_agents",
|
|
62
|
+
"list_builtin_agents",
|
|
63
|
+
"list_local_agents",
|
|
64
|
+
"list_registered_agents",
|
|
65
|
+
"load_agent",
|
|
66
|
+
"load_builtin_agent",
|
|
67
|
+
"load_fortify_agent",
|
|
68
|
+
"load_local_agent",
|
|
69
|
+
"load_registered_agent",
|
|
70
|
+
"register_agent",
|
|
71
|
+
"read_file",
|
|
72
|
+
"refund_order",
|
|
73
|
+
"stream_agent",
|
|
74
|
+
"stream_agent_raw",
|
|
75
|
+
"unregister_agent",
|
|
76
|
+
"web_search",
|
|
77
|
+
"write_file",
|
|
78
|
+
]
|
|
File without changes
|
|
@@ -0,0 +1,106 @@
|
|
|
1
|
+
"""Google ADK ``Runner`` wrapper: opens a :class:`User` scope around each
|
|
2
|
+
``Runner.run*`` call so the wrapped tools' enforcers can resolve the
|
|
3
|
+
active role. Langfuse propagation mirrors User identity into spans.
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
import asyncio
|
|
7
|
+
import os
|
|
8
|
+
from contextlib import contextmanager
|
|
9
|
+
from typing import Any, AsyncGenerator, Generator
|
|
10
|
+
|
|
11
|
+
import nest_asyncio
|
|
12
|
+
from google.adk.agents import BaseAgent
|
|
13
|
+
from google.adk.runners import Runner
|
|
14
|
+
from google.adk.sessions import BaseSessionService
|
|
15
|
+
from google.genai import types
|
|
16
|
+
from langfuse import get_client, propagate_attributes
|
|
17
|
+
from openinference.instrumentation.google_adk import GoogleADKInstrumentor
|
|
18
|
+
|
|
19
|
+
from fortify.adapters.google.wrapper import wrap_google_agent
|
|
20
|
+
from fortify.runtime import User
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class FortifyRunner:
|
|
24
|
+
"""Runner for Google ADK agents with Fortify tool policy and observability."""
|
|
25
|
+
|
|
26
|
+
def __init__(
|
|
27
|
+
self,
|
|
28
|
+
*,
|
|
29
|
+
agent: BaseAgent,
|
|
30
|
+
app_name: str,
|
|
31
|
+
session_service: BaseSessionService,
|
|
32
|
+
api_key: str | None = None,
|
|
33
|
+
**runner_kwargs: Any,
|
|
34
|
+
):
|
|
35
|
+
self.api_key = api_key or os.getenv("FORTIFY_KEY")
|
|
36
|
+
if self.api_key is None:
|
|
37
|
+
raise ValueError(
|
|
38
|
+
"FORTIFY_KEY is not set. Pass api_key= explicitly or set FORTIFY_KEY environment variable."
|
|
39
|
+
)
|
|
40
|
+
# Policy is baked at construction; the Runner is built once and reused
|
|
41
|
+
# since role resolution happens at call time via the User contextvar.
|
|
42
|
+
self._wrapped_agent = wrap_google_agent(agent, api_key=self.api_key)
|
|
43
|
+
self._runner = Runner(
|
|
44
|
+
agent=self._wrapped_agent,
|
|
45
|
+
app_name=app_name,
|
|
46
|
+
session_service=session_service,
|
|
47
|
+
**runner_kwargs,
|
|
48
|
+
)
|
|
49
|
+
self._agent_name = getattr(agent, "name", "default")
|
|
50
|
+
|
|
51
|
+
def _setup_observability(self):
|
|
52
|
+
"""Install Langfuse + GoogleADKInstrumentor (idempotent)."""
|
|
53
|
+
try:
|
|
54
|
+
asyncio.get_running_loop()
|
|
55
|
+
except RuntimeError:
|
|
56
|
+
# No running loop: safe to patch (and only useful for sync entry points).
|
|
57
|
+
# Patching a live loop breaks asyncio.current_task() on Python 3.12+.
|
|
58
|
+
nest_asyncio.apply()
|
|
59
|
+
get_client()
|
|
60
|
+
GoogleADKInstrumentor().instrument()
|
|
61
|
+
|
|
62
|
+
@contextmanager
|
|
63
|
+
def _propagate(self, user: User):
|
|
64
|
+
"""Propagate User identity into Langfuse spans for the block."""
|
|
65
|
+
kwargs: dict[str, Any] = {"tags": [f"google.runner.run.{self._agent_name}"]}
|
|
66
|
+
kwargs["user_id"] = user.user_id
|
|
67
|
+
kwargs["session_id"] = user.session_id
|
|
68
|
+
kwargs["metadata"] = {"user_role": user.role}
|
|
69
|
+
with propagate_attributes(**kwargs):
|
|
70
|
+
yield
|
|
71
|
+
|
|
72
|
+
def run(
|
|
73
|
+
self,
|
|
74
|
+
*,
|
|
75
|
+
new_message: types.Content,
|
|
76
|
+
user: User,
|
|
77
|
+
**kwargs: Any,
|
|
78
|
+
) -> Generator[Any, None, None]:
|
|
79
|
+
"""Run the Google ADK agent synchronously, yielding events."""
|
|
80
|
+
self._setup_observability()
|
|
81
|
+
with user.sync_scope(), self._propagate(user):
|
|
82
|
+
yield from self._runner.run(
|
|
83
|
+
user_id=user.user_id,
|
|
84
|
+
session_id=user.session_id,
|
|
85
|
+
new_message=new_message,
|
|
86
|
+
**kwargs,
|
|
87
|
+
)
|
|
88
|
+
|
|
89
|
+
async def run_async(
|
|
90
|
+
self,
|
|
91
|
+
*,
|
|
92
|
+
new_message: types.Content | None = None,
|
|
93
|
+
user: User,
|
|
94
|
+
**kwargs: Any,
|
|
95
|
+
) -> AsyncGenerator[Any, None]:
|
|
96
|
+
"""Run the Google ADK agent asynchronously, yielding events."""
|
|
97
|
+
self._setup_observability()
|
|
98
|
+
async with user:
|
|
99
|
+
with self._propagate(user):
|
|
100
|
+
async for event in self._runner.run_async(
|
|
101
|
+
user_id=user.user_id,
|
|
102
|
+
session_id=user.session_id,
|
|
103
|
+
new_message=new_message,
|
|
104
|
+
**kwargs,
|
|
105
|
+
):
|
|
106
|
+
yield event
|
|
@@ -0,0 +1,57 @@
|
|
|
1
|
+
"""Google ADK adapter: wrap ``BaseTool`` so ``run_async`` consults a
|
|
2
|
+
:class:`PolicyEnforcer` first. Non-allow outcomes render as markered
|
|
3
|
+
strings the model sees as tool output.
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
from __future__ import annotations
|
|
7
|
+
|
|
8
|
+
import copy
|
|
9
|
+
import functools
|
|
10
|
+
from collections.abc import Callable
|
|
11
|
+
from typing import Any, Union
|
|
12
|
+
|
|
13
|
+
from google.adk.tools.base_tool import BaseTool
|
|
14
|
+
from google.adk.tools.function_tool import FunctionTool
|
|
15
|
+
from google.adk.tools.tool_context import ToolContext
|
|
16
|
+
|
|
17
|
+
from fortify.security.enforcer import PolicyEnforcer
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
ToolEntry = Union[BaseTool, Callable[..., Any]]
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def _normalize(tool: ToolEntry) -> BaseTool:
|
|
24
|
+
"""Coerce a tool entry into a ``BaseTool`` (plain callables → FunctionTool)."""
|
|
25
|
+
if isinstance(tool, BaseTool):
|
|
26
|
+
return tool
|
|
27
|
+
if callable(tool):
|
|
28
|
+
return FunctionTool(func=tool)
|
|
29
|
+
raise TypeError(
|
|
30
|
+
f"Cannot install policy on tool {tool!r}: expected google.adk BaseTool "
|
|
31
|
+
f"or callable, got {type(tool).__name__}."
|
|
32
|
+
)
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def wrap_tool(tool: ToolEntry, enforcer: PolicyEnforcer) -> BaseTool:
|
|
36
|
+
"""Return a copy of ``tool`` with ``run_async`` gated by ``enforcer``."""
|
|
37
|
+
base = _normalize(tool)
|
|
38
|
+
name = base.name
|
|
39
|
+
original_run_async = base.run_async
|
|
40
|
+
|
|
41
|
+
@functools.wraps(original_run_async, updated=())
|
|
42
|
+
async def guarded_run_async(
|
|
43
|
+
*, args: dict[str, Any], tool_context: ToolContext
|
|
44
|
+
) -> Any:
|
|
45
|
+
decision = enforcer.decide(name, args or {})
|
|
46
|
+
if decision.allowed:
|
|
47
|
+
return await original_run_async(args=args, tool_context=tool_context)
|
|
48
|
+
return decision.as_error_message()
|
|
49
|
+
|
|
50
|
+
wrapped = copy.copy(base)
|
|
51
|
+
wrapped.run_async = guarded_run_async
|
|
52
|
+
return wrapped
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def wrap_tools(tools: list[ToolEntry], enforcer: PolicyEnforcer) -> list[BaseTool]:
|
|
56
|
+
"""Return a fresh list of policy-gated copies."""
|
|
57
|
+
return [wrap_tool(t, enforcer) for t in tools]
|
|
@@ -0,0 +1,47 @@
|
|
|
1
|
+
"""Google ADK adapter: build a :class:`PolicySet`, construct one
|
|
2
|
+
:class:`PolicyEnforcer`, and return a clone of the agent whose tools
|
|
3
|
+
are policy-gated. User-agnostic at wrap time — role resolution happens
|
|
4
|
+
inside the enforcer via the :class:`User` contextvar.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
from google.adk.agents import BaseAgent
|
|
10
|
+
|
|
11
|
+
from fortify import audit
|
|
12
|
+
from fortify.adapters.google.tools import wrap_tools
|
|
13
|
+
from fortify.security import AgentPolicy, BaseToolPolicy, PolicySet
|
|
14
|
+
from fortify.security.enforcer import PolicyEnforcer
|
|
15
|
+
from fortify.security.policy_set import DEFAULT_ROLE_NAME
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def build_policy_set(
|
|
19
|
+
api_key: str, # noqa: ARG001 — reserved for the future Fortify-cloud fetch
|
|
20
|
+
agent_name: str, # noqa: ARG001 — same
|
|
21
|
+
tool_names: list[str],
|
|
22
|
+
) -> PolicySet:
|
|
23
|
+
"""Placeholder allow-all one-role bundle. TODO: cloud-fetch via FortifyClient."""
|
|
24
|
+
default_policy = AgentPolicy(
|
|
25
|
+
tools={name: BaseToolPolicy(mode="allow") for name in tool_names}
|
|
26
|
+
)
|
|
27
|
+
return PolicySet({DEFAULT_ROLE_NAME: default_policy})
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def wrap_google_agent(agent: BaseAgent, *, api_key: str) -> BaseAgent:
|
|
31
|
+
"""Return a clone of ``agent`` with policy-gated tools.
|
|
32
|
+
|
|
33
|
+
Caller must open a :class:`User` scope around the run.
|
|
34
|
+
``NEEDS_APPROVAL`` outcomes surface as ``[approval_required]``-prefixed
|
|
35
|
+
strings in tool results; ``[policy_denied]`` for denials.
|
|
36
|
+
"""
|
|
37
|
+
audit_sender = audit.configure(api_key)
|
|
38
|
+
|
|
39
|
+
agent_name = getattr(agent, "name", "default")
|
|
40
|
+
tools = list(getattr(agent, "tools", []) or [])
|
|
41
|
+
tool_names = [getattr(t, "name", getattr(t, "__name__", "tool")) for t in tools]
|
|
42
|
+
policy_set = build_policy_set(api_key, agent_name, tool_names)
|
|
43
|
+
enforcer = PolicyEnforcer(
|
|
44
|
+
policy_set, agent_name=agent_name, audit_sender=audit_sender
|
|
45
|
+
)
|
|
46
|
+
guarded_tools = wrap_tools(tools, enforcer)
|
|
47
|
+
return agent.model_copy(update={"tools": guarded_tools})
|
|
@@ -0,0 +1,133 @@
|
|
|
1
|
+
"""Proxy around a pre-built ``CompiledStateGraph`` for Fortify-aware calls."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from typing import Any, AsyncIterator, Iterator
|
|
6
|
+
|
|
7
|
+
from langchain_core.runnables import RunnableConfig
|
|
8
|
+
from langfuse import get_client, propagate_attributes
|
|
9
|
+
from langfuse.langchain import CallbackHandler
|
|
10
|
+
from langgraph.graph.state import CompiledStateGraph
|
|
11
|
+
|
|
12
|
+
from fortify.runtime import User
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class FortifyLangchainAgent:
|
|
16
|
+
"""Proxy around a ``CompiledStateGraph`` that opens a User scope per call.
|
|
17
|
+
|
|
18
|
+
Tools are already enforcer-installed at construction (by
|
|
19
|
+
:func:`wrap_langchain_agent`). This proxy pushes the active
|
|
20
|
+
:class:`User` onto the contextvar and propagates identity into
|
|
21
|
+
Langfuse spans. ``user`` is per-call, so one proxy serves many
|
|
22
|
+
users concurrently.
|
|
23
|
+
"""
|
|
24
|
+
|
|
25
|
+
def __init__(
|
|
26
|
+
self,
|
|
27
|
+
*,
|
|
28
|
+
agent: CompiledStateGraph,
|
|
29
|
+
api_key: str,
|
|
30
|
+
tool_names: list[str],
|
|
31
|
+
) -> None:
|
|
32
|
+
self._agent = agent
|
|
33
|
+
self._api_key = api_key
|
|
34
|
+
self._tool_names = tool_names
|
|
35
|
+
self._langfuse = get_client()
|
|
36
|
+
self._callback_handler = CallbackHandler()
|
|
37
|
+
|
|
38
|
+
def _propagate_kwargs(self, user: User, method: str) -> dict[str, Any]:
|
|
39
|
+
return {
|
|
40
|
+
"tags": [f"langchain.agent.{method}"],
|
|
41
|
+
"user_id": user.user_id,
|
|
42
|
+
"session_id": user.session_id,
|
|
43
|
+
"metadata": {"user_role": user.role},
|
|
44
|
+
}
|
|
45
|
+
|
|
46
|
+
def _with_callbacks(self, config: RunnableConfig | None) -> RunnableConfig:
|
|
47
|
+
"""Append the Fortify callback handler to ``config['callbacks']``."""
|
|
48
|
+
merged: RunnableConfig = dict(config) if config else {}
|
|
49
|
+
callbacks = list(merged.get("callbacks") or [])
|
|
50
|
+
if self._callback_handler not in callbacks:
|
|
51
|
+
callbacks.append(self._callback_handler)
|
|
52
|
+
merged["callbacks"] = callbacks
|
|
53
|
+
return merged
|
|
54
|
+
|
|
55
|
+
async def ainvoke(
|
|
56
|
+
self,
|
|
57
|
+
input: dict[str, Any],
|
|
58
|
+
*,
|
|
59
|
+
user: User,
|
|
60
|
+
config: RunnableConfig | None = None,
|
|
61
|
+
**kwargs: Any,
|
|
62
|
+
) -> dict[str, Any]:
|
|
63
|
+
"""Invoke the agent asynchronously inside a User scope."""
|
|
64
|
+
async with user:
|
|
65
|
+
with propagate_attributes(**self._propagate_kwargs(user, "ainvoke")):
|
|
66
|
+
return await self._agent.ainvoke(
|
|
67
|
+
input, self._with_callbacks(config), **kwargs
|
|
68
|
+
)
|
|
69
|
+
|
|
70
|
+
def invoke(
|
|
71
|
+
self,
|
|
72
|
+
input: dict[str, Any],
|
|
73
|
+
*,
|
|
74
|
+
user: User,
|
|
75
|
+
config: RunnableConfig | None = None,
|
|
76
|
+
**kwargs: Any,
|
|
77
|
+
) -> dict[str, Any]:
|
|
78
|
+
"""Invoke the agent synchronously inside a User scope."""
|
|
79
|
+
with user.sync_scope():
|
|
80
|
+
with propagate_attributes(**self._propagate_kwargs(user, "invoke")):
|
|
81
|
+
return self._agent.invoke(input, self._with_callbacks(config), **kwargs)
|
|
82
|
+
|
|
83
|
+
async def astream(
|
|
84
|
+
self,
|
|
85
|
+
input: dict[str, Any],
|
|
86
|
+
*,
|
|
87
|
+
user: User,
|
|
88
|
+
config: RunnableConfig | None = None,
|
|
89
|
+
**kwargs: Any,
|
|
90
|
+
) -> AsyncIterator[dict[str, Any]]:
|
|
91
|
+
"""Stream the agent asynchronously inside a User scope."""
|
|
92
|
+
async with user:
|
|
93
|
+
with propagate_attributes(**self._propagate_kwargs(user, "astream")):
|
|
94
|
+
async for chunk in self._agent.astream(
|
|
95
|
+
input, self._with_callbacks(config), **kwargs
|
|
96
|
+
):
|
|
97
|
+
yield chunk
|
|
98
|
+
|
|
99
|
+
def stream(
|
|
100
|
+
self,
|
|
101
|
+
input: dict[str, Any],
|
|
102
|
+
*,
|
|
103
|
+
user: User,
|
|
104
|
+
config: RunnableConfig | None = None,
|
|
105
|
+
**kwargs: Any,
|
|
106
|
+
) -> Iterator[dict[str, Any]]:
|
|
107
|
+
"""Stream the agent synchronously inside a User scope."""
|
|
108
|
+
with user.sync_scope():
|
|
109
|
+
with propagate_attributes(**self._propagate_kwargs(user, "stream")):
|
|
110
|
+
yield from self._agent.stream(
|
|
111
|
+
input, self._with_callbacks(config), **kwargs
|
|
112
|
+
)
|
|
113
|
+
|
|
114
|
+
async def astream_events(
|
|
115
|
+
self,
|
|
116
|
+
input: dict[str, Any],
|
|
117
|
+
version: str,
|
|
118
|
+
*,
|
|
119
|
+
user: User,
|
|
120
|
+
config: RunnableConfig | None = None,
|
|
121
|
+
**kwargs: Any,
|
|
122
|
+
) -> AsyncIterator[dict[str, Any]]:
|
|
123
|
+
"""Stream the agent events asynchronously inside a User scope."""
|
|
124
|
+
async with user:
|
|
125
|
+
with propagate_attributes(**self._propagate_kwargs(user, "astream_events")):
|
|
126
|
+
async for event in self._agent.astream_events(
|
|
127
|
+
input, version, config=self._with_callbacks(config), **kwargs
|
|
128
|
+
):
|
|
129
|
+
yield event
|
|
130
|
+
|
|
131
|
+
def __getattr__(self, name: str) -> Any:
|
|
132
|
+
"""Delegate unknown attributes to the wrapped agent."""
|
|
133
|
+
return getattr(self._agent, name)
|
|
@@ -0,0 +1,248 @@
|
|
|
1
|
+
"""LangChain adapter for :class:`PolicyEnforcer`.
|
|
2
|
+
|
|
3
|
+
:class:`GuardedTool` wraps a ``BaseTool`` (used by
|
|
4
|
+
:meth:`FortifyAgent.enforce_policy`, which rebuilds the graph) and
|
|
5
|
+
carries an optional ``approval_handler`` for inline ``NEEDS_APPROVAL``
|
|
6
|
+
resolution.
|
|
7
|
+
:func:`install_enforcer_on_tool` mutates ``StructuredTool``'s ``func``/
|
|
8
|
+
``coroutine`` in place (used by :func:`wrap_langchain_agent` for
|
|
9
|
+
pre-built ``CompiledStateGraph``s) and always renders non-allow as a
|
|
10
|
+
structured error — approval flows wire in on the host side.
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
from __future__ import annotations
|
|
14
|
+
|
|
15
|
+
import functools
|
|
16
|
+
from collections.abc import Awaitable, Callable
|
|
17
|
+
from inspect import isawaitable
|
|
18
|
+
from typing import Any
|
|
19
|
+
|
|
20
|
+
from langchain_core.tools import BaseTool
|
|
21
|
+
from langchain_core.tools.structured import StructuredTool
|
|
22
|
+
from pydantic import ConfigDict
|
|
23
|
+
|
|
24
|
+
from fortify.agents.factory import ApprovalHandler
|
|
25
|
+
from fortify.security.decision import Decision, DecisionOutcome
|
|
26
|
+
from fortify.security.enforcer import PolicyEnforcer
|
|
27
|
+
from fortify.tools.decorators import TOOL_METADATA_ATTR
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def _copy_tool_metadata(source: Any, target: Any) -> Any:
|
|
31
|
+
"""Copy fortify tool metadata (tracing labels, etc.) onto a wrapper."""
|
|
32
|
+
metadata = getattr(source, TOOL_METADATA_ATTR, None)
|
|
33
|
+
if metadata is not None:
|
|
34
|
+
setattr(target, TOOL_METADATA_ATTR, metadata)
|
|
35
|
+
return target
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def _resolve_approval_sync(handler: ApprovalHandler, decision: Decision) -> bool:
|
|
39
|
+
"""Resolve a NEEDS_APPROVAL decision in a sync caller (rejects coroutines)."""
|
|
40
|
+
if isinstance(handler, bool):
|
|
41
|
+
return handler
|
|
42
|
+
result = handler(decision)
|
|
43
|
+
if isawaitable(result):
|
|
44
|
+
raise RuntimeError(
|
|
45
|
+
"approval_handler returned a coroutine; sync tool invocation cannot "
|
|
46
|
+
"await it — use ainvoke/astream/astream_events"
|
|
47
|
+
)
|
|
48
|
+
return bool(result)
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
async def _resolve_approval_async(handler: ApprovalHandler, decision: Decision) -> bool:
|
|
52
|
+
"""Resolve a NEEDS_APPROVAL decision in an async caller."""
|
|
53
|
+
if isinstance(handler, bool):
|
|
54
|
+
return handler
|
|
55
|
+
result = handler(decision)
|
|
56
|
+
if isawaitable(result):
|
|
57
|
+
result = await result
|
|
58
|
+
return bool(result)
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
class GuardedTool(BaseTool):
|
|
62
|
+
"""LangChain tool wrapper that consults a :class:`PolicyEnforcer`.
|
|
63
|
+
|
|
64
|
+
ALLOW delegates to the wrapped tool; non-ALLOW renders
|
|
65
|
+
``Decision.as_error_payload()`` so the LLM sees governance failures
|
|
66
|
+
as tool output. NEEDS_APPROVAL is treated as denial unless
|
|
67
|
+
``approval_handler`` (callable taking the :class:`Decision`, or a
|
|
68
|
+
``bool`` shorthand) returns truthy.
|
|
69
|
+
"""
|
|
70
|
+
|
|
71
|
+
model_config = ConfigDict(arbitrary_types_allowed=True)
|
|
72
|
+
|
|
73
|
+
wrapped_tool: BaseTool
|
|
74
|
+
enforcer: PolicyEnforcer | None = None
|
|
75
|
+
approval_handler: ApprovalHandler | None = None
|
|
76
|
+
|
|
77
|
+
@classmethod
|
|
78
|
+
def wrap(
|
|
79
|
+
cls,
|
|
80
|
+
tool: BaseTool,
|
|
81
|
+
*,
|
|
82
|
+
enforcer: PolicyEnforcer | None = None,
|
|
83
|
+
approval_handler: ApprovalHandler | None = None,
|
|
84
|
+
) -> "GuardedTool":
|
|
85
|
+
"""Return a GuardedTool delegating to ``tool`` after policy check.
|
|
86
|
+
|
|
87
|
+
Idempotent re-wrap: an existing ``GuardedTool`` is unwrapped once
|
|
88
|
+
so enforcers don't stack; fields fall through unless explicitly
|
|
89
|
+
overridden.
|
|
90
|
+
"""
|
|
91
|
+
if isinstance(tool, cls):
|
|
92
|
+
inner = tool.wrapped_tool
|
|
93
|
+
resolved_enforcer = enforcer if enforcer is not None else tool.enforcer
|
|
94
|
+
resolved_approval = (
|
|
95
|
+
approval_handler
|
|
96
|
+
if approval_handler is not None
|
|
97
|
+
else tool.approval_handler
|
|
98
|
+
)
|
|
99
|
+
else:
|
|
100
|
+
inner = tool
|
|
101
|
+
resolved_enforcer = enforcer
|
|
102
|
+
resolved_approval = approval_handler
|
|
103
|
+
|
|
104
|
+
guarded = cls(
|
|
105
|
+
name=inner.name,
|
|
106
|
+
description=inner.description,
|
|
107
|
+
args_schema=inner.args_schema,
|
|
108
|
+
return_direct=inner.return_direct,
|
|
109
|
+
verbose=inner.verbose,
|
|
110
|
+
callbacks=inner.callbacks,
|
|
111
|
+
tags=inner.tags,
|
|
112
|
+
metadata=inner.metadata,
|
|
113
|
+
handle_tool_error=inner.handle_tool_error,
|
|
114
|
+
handle_validation_error=inner.handle_validation_error,
|
|
115
|
+
response_format=inner.response_format,
|
|
116
|
+
extras=inner.extras,
|
|
117
|
+
wrapped_tool=inner,
|
|
118
|
+
enforcer=resolved_enforcer,
|
|
119
|
+
approval_handler=resolved_approval,
|
|
120
|
+
)
|
|
121
|
+
return _copy_tool_metadata(inner, guarded)
|
|
122
|
+
|
|
123
|
+
async def _arun(self, *args: Any, **kwargs: Any) -> Any:
|
|
124
|
+
if self.enforcer is not None:
|
|
125
|
+
decision = self.enforcer.decide(self.name, kwargs)
|
|
126
|
+
if not decision.allowed:
|
|
127
|
+
if (
|
|
128
|
+
decision.outcome is DecisionOutcome.NEEDS_APPROVAL
|
|
129
|
+
and self.approval_handler is not None
|
|
130
|
+
and await _resolve_approval_async(self.approval_handler, decision)
|
|
131
|
+
):
|
|
132
|
+
pass # approved → fall through and invoke
|
|
133
|
+
else:
|
|
134
|
+
return {"ok": False, "error": decision.as_error_payload()}
|
|
135
|
+
return await self._invoke_wrapped_async(*args, **kwargs)
|
|
136
|
+
|
|
137
|
+
def _run(self, *args: Any, **kwargs: Any) -> Any:
|
|
138
|
+
if self.enforcer is not None:
|
|
139
|
+
decision = self.enforcer.decide(self.name, kwargs)
|
|
140
|
+
if not decision.allowed:
|
|
141
|
+
if (
|
|
142
|
+
decision.outcome is DecisionOutcome.NEEDS_APPROVAL
|
|
143
|
+
and self.approval_handler is not None
|
|
144
|
+
and _resolve_approval_sync(self.approval_handler, decision)
|
|
145
|
+
):
|
|
146
|
+
pass
|
|
147
|
+
else:
|
|
148
|
+
return {"ok": False, "error": decision.as_error_payload()}
|
|
149
|
+
return self._invoke_wrapped_sync(*args, **kwargs)
|
|
150
|
+
|
|
151
|
+
async def _invoke_wrapped_async(self, *args: Any, **kwargs: Any) -> Any:
|
|
152
|
+
"""Call the wrapped tool without re-entering LangChain instrumentation."""
|
|
153
|
+
if isinstance(self.wrapped_tool, StructuredTool):
|
|
154
|
+
if self.wrapped_tool.coroutine is not None:
|
|
155
|
+
return await self.wrapped_tool.coroutine(*args, **kwargs)
|
|
156
|
+
if self.wrapped_tool.func is not None:
|
|
157
|
+
return self.wrapped_tool.func(*args, **kwargs)
|
|
158
|
+
return await self.wrapped_tool._arun(*args, **kwargs)
|
|
159
|
+
|
|
160
|
+
def _invoke_wrapped_sync(self, *args: Any, **kwargs: Any) -> Any:
|
|
161
|
+
if (
|
|
162
|
+
isinstance(self.wrapped_tool, StructuredTool)
|
|
163
|
+
and self.wrapped_tool.func is not None
|
|
164
|
+
):
|
|
165
|
+
return self.wrapped_tool.func(*args, **kwargs)
|
|
166
|
+
return self.wrapped_tool._run(*args, **kwargs)
|
|
167
|
+
|
|
168
|
+
|
|
169
|
+
# ---------------------------------------------------------------------------
|
|
170
|
+
# In-place installer for retrofitting existing CompiledStateGraph tools.
|
|
171
|
+
# ---------------------------------------------------------------------------
|
|
172
|
+
|
|
173
|
+
_ORIGINAL_FUNC_ATTR = "_fortify_original_func"
|
|
174
|
+
_ORIGINAL_COROUTINE_ATTR = "_fortify_original_coroutine"
|
|
175
|
+
_INSTALLED_ATTR = "_fortify_enforcer_installed"
|
|
176
|
+
|
|
177
|
+
|
|
178
|
+
def install_enforcer_on_tool(
|
|
179
|
+
tool: BaseTool,
|
|
180
|
+
*,
|
|
181
|
+
enforcer: PolicyEnforcer,
|
|
182
|
+
) -> BaseTool:
|
|
183
|
+
"""Install :class:`PolicyEnforcer` gating on ``tool`` in place.
|
|
184
|
+
|
|
185
|
+
Same semantics as :class:`GuardedTool` but mutates ``StructuredTool``'s
|
|
186
|
+
``func``/``coroutine`` instead of constructing a wrapper — use when
|
|
187
|
+
the tool is already bound to a ``CompiledStateGraph``. Idempotent:
|
|
188
|
+
re-install restores captured originals first so gates don't stack.
|
|
189
|
+
Non-allow outcomes render as the structured error dict; approval
|
|
190
|
+
flows belong on the host side, not on this in-place installer.
|
|
191
|
+
"""
|
|
192
|
+
name = tool.name
|
|
193
|
+
original_func: Callable[..., Any] | None = getattr(tool, _ORIGINAL_FUNC_ATTR, None)
|
|
194
|
+
if original_func is None:
|
|
195
|
+
original_func = getattr(tool, "func", None)
|
|
196
|
+
original_coroutine: Callable[..., Awaitable[Any]] | None = getattr(
|
|
197
|
+
tool, _ORIGINAL_COROUTINE_ATTR, None
|
|
198
|
+
)
|
|
199
|
+
if original_coroutine is None:
|
|
200
|
+
original_coroutine = getattr(tool, "coroutine", None)
|
|
201
|
+
|
|
202
|
+
if original_func is None and original_coroutine is None:
|
|
203
|
+
raise TypeError(
|
|
204
|
+
f"Cannot install policy on tool {name!r}: it is a "
|
|
205
|
+
f"{type(tool).__name__} without `func`/`coroutine` attributes. "
|
|
206
|
+
"In-place wrapping only supports StructuredTool-style tools."
|
|
207
|
+
)
|
|
208
|
+
|
|
209
|
+
if original_func is not None:
|
|
210
|
+
captured_func = original_func
|
|
211
|
+
|
|
212
|
+
@functools.wraps(captured_func)
|
|
213
|
+
def guarded_func(*args: Any, **kwargs: Any) -> Any:
|
|
214
|
+
decision = enforcer.decide(name, kwargs)
|
|
215
|
+
if decision.allowed:
|
|
216
|
+
return captured_func(*args, **kwargs)
|
|
217
|
+
return {"ok": False, "error": decision.as_error_payload()}
|
|
218
|
+
|
|
219
|
+
setattr(tool, _ORIGINAL_FUNC_ATTR, captured_func)
|
|
220
|
+
tool.func = guarded_func
|
|
221
|
+
|
|
222
|
+
if original_coroutine is not None:
|
|
223
|
+
captured_coroutine = original_coroutine
|
|
224
|
+
|
|
225
|
+
@functools.wraps(captured_coroutine)
|
|
226
|
+
async def guarded_coroutine(*args: Any, **kwargs: Any) -> Any:
|
|
227
|
+
decision = enforcer.decide(name, kwargs)
|
|
228
|
+
if decision.allowed:
|
|
229
|
+
return await captured_coroutine(*args, **kwargs)
|
|
230
|
+
return {"ok": False, "error": decision.as_error_payload()}
|
|
231
|
+
|
|
232
|
+
setattr(tool, _ORIGINAL_COROUTINE_ATTR, captured_coroutine)
|
|
233
|
+
tool.coroutine = guarded_coroutine
|
|
234
|
+
|
|
235
|
+
tool.handle_tool_error = True
|
|
236
|
+
setattr(tool, _INSTALLED_ATTR, True)
|
|
237
|
+
return tool
|
|
238
|
+
|
|
239
|
+
|
|
240
|
+
def install_enforcer_on_tools(
|
|
241
|
+
tools: list[BaseTool],
|
|
242
|
+
*,
|
|
243
|
+
enforcer: PolicyEnforcer,
|
|
244
|
+
) -> list[BaseTool]:
|
|
245
|
+
"""Install enforcement on every StructuredTool-style tool in place."""
|
|
246
|
+
for t in tools:
|
|
247
|
+
install_enforcer_on_tool(t, enforcer=enforcer)
|
|
248
|
+
return tools
|