ory-llamaindex 0.13.9__tar.gz
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.
- ory_llamaindex-0.13.9/.gitignore +36 -0
- ory_llamaindex-0.13.9/PKG-INFO +46 -0
- ory_llamaindex-0.13.9/README.md +31 -0
- ory_llamaindex-0.13.9/pyproject.toml +23 -0
- ory_llamaindex-0.13.9/src/ory_llamaindex/__init__.py +14 -0
- ory_llamaindex-0.13.9/src/ory_llamaindex/tools.py +209 -0
- ory_llamaindex-0.13.9/tests/test_realsdk.py +73 -0
- ory_llamaindex-0.13.9/tests/test_tools.py +167 -0
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
node_modules/
|
|
2
|
+
dist/
|
|
3
|
+
*.tsbuildinfo
|
|
4
|
+
.env
|
|
5
|
+
.env.local
|
|
6
|
+
|
|
7
|
+
# Python (uv workspace under python/)
|
|
8
|
+
.venv/
|
|
9
|
+
__pycache__/
|
|
10
|
+
*.egg-info/
|
|
11
|
+
.pytest_cache/
|
|
12
|
+
.ruff_cache/
|
|
13
|
+
build/
|
|
14
|
+
*.pyc
|
|
15
|
+
|
|
16
|
+
# Debug logs
|
|
17
|
+
*.log
|
|
18
|
+
|
|
19
|
+
# Harness sandbox dirs
|
|
20
|
+
.sandbox/
|
|
21
|
+
|
|
22
|
+
# Staged install-surface repo contents (see scripts/sync-install-surfaces.mjs)
|
|
23
|
+
.install-surfaces/
|
|
24
|
+
|
|
25
|
+
# Local dev environment (local Ory stack + Verdaccio registry)
|
|
26
|
+
.ory-dev/
|
|
27
|
+
|
|
28
|
+
# Worktrees for changes
|
|
29
|
+
.worktrees/
|
|
30
|
+
.claude/worktrees/
|
|
31
|
+
|
|
32
|
+
# Gemini CLI extension assets — materialized at install time from
|
|
33
|
+
# @ory/argus templates. The repo source is the canonical templates;
|
|
34
|
+
# these subdirs are generated wherever the install runs.
|
|
35
|
+
packages/gemini-cli/gemini-extension/skills/
|
|
36
|
+
packages/gemini-cli/gemini-extension/commands/
|
|
@@ -0,0 +1,46 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: ory-llamaindex
|
|
3
|
+
Version: 0.13.9
|
|
4
|
+
Summary: Ory Agent Security for LlamaIndex — per-tool authorization, tracing, and identity propagation by wrapping FunctionTools. Built on ory-argus.
|
|
5
|
+
Author: Ory
|
|
6
|
+
License-Expression: Apache-2.0
|
|
7
|
+
Keywords: agent,ai,authorization,llamaindex,ory,permissions
|
|
8
|
+
Requires-Python: >=3.10
|
|
9
|
+
Requires-Dist: llama-index-core>=0.11
|
|
10
|
+
Requires-Dist: ory-argus<1,>=0.8
|
|
11
|
+
Provides-Extra: dev
|
|
12
|
+
Requires-Dist: pytest>=8; extra == 'dev'
|
|
13
|
+
Requires-Dist: ruff>=0.6; extra == 'dev'
|
|
14
|
+
Description-Content-Type: text/markdown
|
|
15
|
+
|
|
16
|
+
# ory-llamaindex
|
|
17
|
+
|
|
18
|
+
Ory Agent Security for [LlamaIndex](https://docs.llamaindex.ai).
|
|
19
|
+
|
|
20
|
+
Wraps your LlamaIndex tools so every call is authorized against Ory Permissions, traced, and
|
|
21
|
+
tied to the user → agent identity — built on [`ory-argus`](https://pypi.org/project/ory-argus/).
|
|
22
|
+
|
|
23
|
+
```bash
|
|
24
|
+
pip install ory-llamaindex
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
```python
|
|
28
|
+
from llama_index.core.agent.workflow import FunctionAgent
|
|
29
|
+
from ory_llamaindex import guard_tools
|
|
30
|
+
|
|
31
|
+
agent = FunctionAgent(tools=guard_tools([search_tool, email_tool]), llm=llm)
|
|
32
|
+
```
|
|
33
|
+
|
|
34
|
+
In **enforce** mode (`ORY_PERMISSION_MODE=enforce`) a denied tool raises an error surfaced to
|
|
35
|
+
the agent; the tool never runs. In **observe** mode (default) it runs and a
|
|
36
|
+
`permission.observe_deny` span is recorded.
|
|
37
|
+
|
|
38
|
+
For tracing across paths that bypass the wrapped tools, register the instrumentation handler:
|
|
39
|
+
|
|
40
|
+
```python
|
|
41
|
+
from llama_index.core.instrumentation import get_dispatcher
|
|
42
|
+
from ory_llamaindex import ory_event_handler
|
|
43
|
+
get_dispatcher().add_event_handler(ory_event_handler())
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
Credentials come from the shared `~/.config/ory-agent-plugins/config.json`.
|
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
# ory-llamaindex
|
|
2
|
+
|
|
3
|
+
Ory Agent Security for [LlamaIndex](https://docs.llamaindex.ai).
|
|
4
|
+
|
|
5
|
+
Wraps your LlamaIndex tools so every call is authorized against Ory Permissions, traced, and
|
|
6
|
+
tied to the user → agent identity — built on [`ory-argus`](https://pypi.org/project/ory-argus/).
|
|
7
|
+
|
|
8
|
+
```bash
|
|
9
|
+
pip install ory-llamaindex
|
|
10
|
+
```
|
|
11
|
+
|
|
12
|
+
```python
|
|
13
|
+
from llama_index.core.agent.workflow import FunctionAgent
|
|
14
|
+
from ory_llamaindex import guard_tools
|
|
15
|
+
|
|
16
|
+
agent = FunctionAgent(tools=guard_tools([search_tool, email_tool]), llm=llm)
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
In **enforce** mode (`ORY_PERMISSION_MODE=enforce`) a denied tool raises an error surfaced to
|
|
20
|
+
the agent; the tool never runs. In **observe** mode (default) it runs and a
|
|
21
|
+
`permission.observe_deny` span is recorded.
|
|
22
|
+
|
|
23
|
+
For tracing across paths that bypass the wrapped tools, register the instrumentation handler:
|
|
24
|
+
|
|
25
|
+
```python
|
|
26
|
+
from llama_index.core.instrumentation import get_dispatcher
|
|
27
|
+
from ory_llamaindex import ory_event_handler
|
|
28
|
+
get_dispatcher().add_event_handler(ory_event_handler())
|
|
29
|
+
```
|
|
30
|
+
|
|
31
|
+
Credentials come from the shared `~/.config/ory-agent-plugins/config.json`.
|
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "ory-llamaindex"
|
|
7
|
+
version = "0.13.9"
|
|
8
|
+
description = "Ory Agent Security for LlamaIndex — per-tool authorization, tracing, and identity propagation by wrapping FunctionTools. Built on ory-argus."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.10"
|
|
11
|
+
license = "Apache-2.0"
|
|
12
|
+
authors = [{ name = "Ory" }]
|
|
13
|
+
keywords = ["ory", "llamaindex", "authorization", "agent", "ai", "permissions"]
|
|
14
|
+
dependencies = [
|
|
15
|
+
"ory-argus>=0.8,<1",
|
|
16
|
+
"llama-index-core>=0.11",
|
|
17
|
+
]
|
|
18
|
+
|
|
19
|
+
[project.optional-dependencies]
|
|
20
|
+
dev = ["pytest>=8", "ruff>=0.6"]
|
|
21
|
+
|
|
22
|
+
[tool.hatch.build.targets.wheel]
|
|
23
|
+
packages = ["src/ory_llamaindex"]
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
"""Ory Agent Security for LlamaIndex.
|
|
2
|
+
|
|
3
|
+
from ory_llamaindex import guard_tools
|
|
4
|
+
|
|
5
|
+
agent = FunctionAgent(tools=guard_tools([search_tool, email_tool]), llm=llm)
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
from .tools import guard_tool_async_fn, guard_tool_fn, guard_tools, ory_event_handler
|
|
11
|
+
|
|
12
|
+
__version__ = "0.13.9"
|
|
13
|
+
|
|
14
|
+
__all__ = ["guard_tools", "guard_tool_fn", "guard_tool_async_fn", "ory_event_handler", "__version__"]
|
|
@@ -0,0 +1,209 @@
|
|
|
1
|
+
"""Ory Agent Security for LlamaIndex.
|
|
2
|
+
|
|
3
|
+
LlamaIndex's per-tool gate point is the ``FunctionTool``'s underlying callables.
|
|
4
|
+
``guard_tools(tools)`` wraps each tool's sync function (``_fn`` / ``fn``, the ``tool.call()``
|
|
5
|
+
path) *and* its async function (``_async_fn`` / ``async_fn``, the ``tool.acall()`` path used
|
|
6
|
+
by FunctionAgent workflows) so every call is authorized against Ory Permissions, traced, and
|
|
7
|
+
(on first call) preceded by the Ory session gates:
|
|
8
|
+
|
|
9
|
+
- allow / observe / fail-open / interactive → run the tool and record ``tool.complete``.
|
|
10
|
+
- deny (enforce) → raise ``OryDenialError`` so the agent surfaces the denial and the tool
|
|
11
|
+
never runs.
|
|
12
|
+
|
|
13
|
+
Pair with :func:`ory_event_handler` (LlamaIndex instrumentation dispatcher) for trace-only
|
|
14
|
+
coverage of paths that bypass the wrapped tools. Gate logic lives in ``ory_argus``;
|
|
15
|
+
``llama_index`` is imported lazily.
|
|
16
|
+
"""
|
|
17
|
+
|
|
18
|
+
from __future__ import annotations
|
|
19
|
+
|
|
20
|
+
import inspect
|
|
21
|
+
import sys
|
|
22
|
+
from collections.abc import Callable
|
|
23
|
+
from typing import Any
|
|
24
|
+
|
|
25
|
+
from ory_argus import (
|
|
26
|
+
DenialContext,
|
|
27
|
+
GateResult,
|
|
28
|
+
OryAgentClient,
|
|
29
|
+
OryDenialError,
|
|
30
|
+
complete,
|
|
31
|
+
gate,
|
|
32
|
+
session_start,
|
|
33
|
+
)
|
|
34
|
+
|
|
35
|
+
HARNESS = "llamaindex"
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def _gate_or_deny(
|
|
39
|
+
client: OryAgentClient,
|
|
40
|
+
*,
|
|
41
|
+
tool_name: str,
|
|
42
|
+
args: tuple,
|
|
43
|
+
kwargs: dict,
|
|
44
|
+
can_block: bool,
|
|
45
|
+
project_url: str | None,
|
|
46
|
+
state: dict,
|
|
47
|
+
) -> None:
|
|
48
|
+
"""Run session-start (once) + the permission gate; raise OryDenialError on enforce-deny."""
|
|
49
|
+
if not state.get("started"):
|
|
50
|
+
state["started"] = True
|
|
51
|
+
try:
|
|
52
|
+
session_start(client, harness=HARNESS, project_url=project_url)
|
|
53
|
+
except Exception as err: # noqa: BLE001
|
|
54
|
+
client.logger.warn("session_start.failed", {"message": str(err)})
|
|
55
|
+
result: GateResult = gate(
|
|
56
|
+
client, harness=HARNESS, tool_name=tool_name,
|
|
57
|
+
tool_args={"args": args, "kwargs": kwargs}, can_block=can_block,
|
|
58
|
+
)
|
|
59
|
+
if result.blocked:
|
|
60
|
+
raise OryDenialError(
|
|
61
|
+
DenialContext(tool=tool_name, subject_id=result.subject, namespace=result.namespace)
|
|
62
|
+
)
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
def guard_tool_fn(
|
|
66
|
+
client: OryAgentClient,
|
|
67
|
+
*,
|
|
68
|
+
tool_name: str,
|
|
69
|
+
fn: Callable[..., Any],
|
|
70
|
+
can_block: bool = True,
|
|
71
|
+
project_url: str | None = None,
|
|
72
|
+
_session_state: dict | None = None,
|
|
73
|
+
) -> Callable[..., Any]:
|
|
74
|
+
"""Return a gated wrapper around a tool callable. SDK-free — testable directly."""
|
|
75
|
+
state = _session_state if _session_state is not None else {}
|
|
76
|
+
|
|
77
|
+
def _wrapped(*args: Any, **kwargs: Any) -> Any:
|
|
78
|
+
_gate_or_deny(
|
|
79
|
+
client, tool_name=tool_name, args=args, kwargs=kwargs,
|
|
80
|
+
can_block=can_block, project_url=project_url, state=state,
|
|
81
|
+
)
|
|
82
|
+
output = fn(*args, **kwargs)
|
|
83
|
+
complete(client, tool_name=tool_name, output=output)
|
|
84
|
+
return output
|
|
85
|
+
|
|
86
|
+
return _wrapped
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
def guard_tool_async_fn(
|
|
90
|
+
client: OryAgentClient,
|
|
91
|
+
*,
|
|
92
|
+
tool_name: str,
|
|
93
|
+
fn: Callable[..., Any],
|
|
94
|
+
can_block: bool = True,
|
|
95
|
+
project_url: str | None = None,
|
|
96
|
+
_session_state: dict | None = None,
|
|
97
|
+
) -> Callable[..., Any]:
|
|
98
|
+
"""Return a gated *async* wrapper around a tool's async callable (``tool.acall()`` path).
|
|
99
|
+
|
|
100
|
+
Same semantics as :func:`guard_tool_fn` — the gate runs (and an enforce-deny raises
|
|
101
|
+
``OryDenialError``) BEFORE the tool body executes; observe/fail-open proceed with spans.
|
|
102
|
+
"""
|
|
103
|
+
state = _session_state if _session_state is not None else {}
|
|
104
|
+
|
|
105
|
+
async def _awrapped(*args: Any, **kwargs: Any) -> Any:
|
|
106
|
+
_gate_or_deny(
|
|
107
|
+
client, tool_name=tool_name, args=args, kwargs=kwargs,
|
|
108
|
+
can_block=can_block, project_url=project_url, state=state,
|
|
109
|
+
)
|
|
110
|
+
output = fn(*args, **kwargs)
|
|
111
|
+
if inspect.isawaitable(output):
|
|
112
|
+
output = await output
|
|
113
|
+
complete(client, tool_name=tool_name, output=output)
|
|
114
|
+
return output
|
|
115
|
+
|
|
116
|
+
return _awrapped
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
def _warn_wrap_failed(client: OryAgentClient, event: str, tool_name: str, err: Exception) -> None:
|
|
120
|
+
"""Wrapping failed: the tool will run UNGATED. Log + always warn on stderr (never stdout)."""
|
|
121
|
+
client.logger.warn(event, {"tool": tool_name, "message": str(err)})
|
|
122
|
+
print(
|
|
123
|
+
f"[ory] WARNING: failed to wrap tool '{tool_name}' ({err}); "
|
|
124
|
+
"it will run UNGATED by Ory permission checks.",
|
|
125
|
+
file=sys.stderr,
|
|
126
|
+
)
|
|
127
|
+
|
|
128
|
+
|
|
129
|
+
def _tool_name(tool: Any) -> str:
|
|
130
|
+
meta = getattr(tool, "metadata", None)
|
|
131
|
+
return (
|
|
132
|
+
getattr(meta, "name", None)
|
|
133
|
+
or getattr(tool, "name", None)
|
|
134
|
+
or type(tool).__name__
|
|
135
|
+
)
|
|
136
|
+
|
|
137
|
+
|
|
138
|
+
def guard_tools(
|
|
139
|
+
tools: list,
|
|
140
|
+
*,
|
|
141
|
+
client: OryAgentClient | None = None,
|
|
142
|
+
can_block: bool = True,
|
|
143
|
+
project_url: str | None = None,
|
|
144
|
+
) -> list:
|
|
145
|
+
"""Wrap each LlamaIndex tool's sync AND async callables with the Ory gate. Returns the same list."""
|
|
146
|
+
c = client or OryAgentClient.from_env(HARNESS)
|
|
147
|
+
state: dict = {}
|
|
148
|
+
for tool in tools:
|
|
149
|
+
name = _tool_name(tool)
|
|
150
|
+
|
|
151
|
+
# Sync path: tool.call() → _fn / fn.
|
|
152
|
+
fn = getattr(tool, "_fn", None) or getattr(tool, "fn", None)
|
|
153
|
+
if callable(fn):
|
|
154
|
+
wrapped = guard_tool_fn(
|
|
155
|
+
c, tool_name=name, fn=fn, can_block=can_block,
|
|
156
|
+
project_url=project_url, _session_state=state,
|
|
157
|
+
)
|
|
158
|
+
try:
|
|
159
|
+
if getattr(tool, "_fn", None) is not None:
|
|
160
|
+
tool._fn = wrapped # type: ignore[attr-defined]
|
|
161
|
+
else:
|
|
162
|
+
tool.fn = wrapped # type: ignore[attr-defined]
|
|
163
|
+
except Exception as err: # noqa: BLE001
|
|
164
|
+
_warn_wrap_failed(c, "llamaindex.wrap_tool.failed", name, err)
|
|
165
|
+
|
|
166
|
+
# Async path: tool.acall() → _async_fn / async_fn. FunctionTool.from_defaults
|
|
167
|
+
# materializes _async_fn from the ORIGINAL fn, so wrapping _fn alone leaves
|
|
168
|
+
# acall() (the FunctionAgent workflow path) ungated.
|
|
169
|
+
afn = getattr(tool, "_async_fn", None) or getattr(tool, "async_fn", None)
|
|
170
|
+
if callable(afn):
|
|
171
|
+
awrapped = guard_tool_async_fn(
|
|
172
|
+
c, tool_name=name, fn=afn, can_block=can_block,
|
|
173
|
+
project_url=project_url, _session_state=state,
|
|
174
|
+
)
|
|
175
|
+
try:
|
|
176
|
+
if getattr(tool, "_async_fn", None) is not None:
|
|
177
|
+
tool._async_fn = awrapped # type: ignore[attr-defined]
|
|
178
|
+
else:
|
|
179
|
+
tool.async_fn = awrapped # type: ignore[attr-defined]
|
|
180
|
+
except Exception as err: # noqa: BLE001
|
|
181
|
+
_warn_wrap_failed(c, "llamaindex.wrap_async_tool.failed", name, err)
|
|
182
|
+
return tools
|
|
183
|
+
|
|
184
|
+
|
|
185
|
+
def ory_event_handler(*, client: OryAgentClient | None = None):
|
|
186
|
+
"""Build a LlamaIndex instrumentation event handler that traces tool calls (trace-only)."""
|
|
187
|
+
from llama_index.core.instrumentation.event_handlers import BaseEventHandler # type: ignore
|
|
188
|
+
|
|
189
|
+
c = client or OryAgentClient.from_env(HARNESS)
|
|
190
|
+
|
|
191
|
+
class _OryEventHandler(BaseEventHandler): # type: ignore[misc]
|
|
192
|
+
@classmethod
|
|
193
|
+
def class_name(cls) -> str:
|
|
194
|
+
return "OryEventHandler"
|
|
195
|
+
|
|
196
|
+
def handle(self, event, **kwargs): # noqa: ANN001
|
|
197
|
+
name = type(event).__name__
|
|
198
|
+
if "ToolCall" in name and "Result" not in name:
|
|
199
|
+
c.tracer.record(
|
|
200
|
+
"tool.invoke", "ok",
|
|
201
|
+
attributes={"toolName": getattr(event, "tool_name", "unknown"), "source": "event"},
|
|
202
|
+
)
|
|
203
|
+
elif "ToolCall" in name and "Result" in name:
|
|
204
|
+
c.tracer.record("tool.complete", "ok", attributes={"source": "event"})
|
|
205
|
+
|
|
206
|
+
return _OryEventHandler()
|
|
207
|
+
|
|
208
|
+
|
|
209
|
+
__all__ = ["HARNESS", "guard_tool_async_fn", "guard_tool_fn", "guard_tools", "ory_event_handler"]
|
|
@@ -0,0 +1,73 @@
|
|
|
1
|
+
"""Real-SDK test: wrap an actual LlamaIndex ``FunctionTool`` and enforce.
|
|
2
|
+
|
|
3
|
+
Runs only when ``llama_index.core`` is importable; the default workspace venv skips it.
|
|
4
|
+
Complements ``test_tools.py`` (fake tool): validates that wrapping a *real* ``FunctionTool``'s
|
|
5
|
+
callable actually intercepts it. The permission decision is stubbed (no backend).
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
import asyncio
|
|
11
|
+
|
|
12
|
+
import pytest
|
|
13
|
+
|
|
14
|
+
pytest.importorskip("llama_index.core")
|
|
15
|
+
|
|
16
|
+
from llama_index.core.tools import FunctionTool # noqa: E402
|
|
17
|
+
|
|
18
|
+
from ory_argus.adapters import GateResult # noqa: E402
|
|
19
|
+
from ory_argus.denial import OryDenialError # noqa: E402
|
|
20
|
+
from ory_argus.permissions import PermissionDecision # noqa: E402
|
|
21
|
+
from ory_argus.testing import create_mock_client # noqa: E402
|
|
22
|
+
from ory_llamaindex import guard_tools # noqa: E402
|
|
23
|
+
from ory_llamaindex import tools as tools_mod # noqa: E402
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def _stub_gate(monkeypatch, *, blocked):
|
|
27
|
+
monkeypatch.setattr(
|
|
28
|
+
tools_mod, "gate",
|
|
29
|
+
lambda client, **kw: GateResult(
|
|
30
|
+
proceed=not blocked, blocked=blocked,
|
|
31
|
+
decision=PermissionDecision(kind="deny" if blocked else "allow"),
|
|
32
|
+
subject="User:t", namespace="AgentTools",
|
|
33
|
+
denial_message="Ory: permission denied" if blocked else None,
|
|
34
|
+
),
|
|
35
|
+
)
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def _real_tool():
|
|
39
|
+
return FunctionTool.from_defaults(fn=lambda: "RAN", name="MyRealTool")
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
def test_real_functiontool_blocked_on_deny(monkeypatch):
|
|
43
|
+
_stub_gate(monkeypatch, blocked=True)
|
|
44
|
+
tool = _real_tool()
|
|
45
|
+
guard_tools([tool], client=create_mock_client(harness="llamaindex"))
|
|
46
|
+
wrapped = getattr(tool, "_fn", None) or getattr(tool, "fn", None)
|
|
47
|
+
with pytest.raises(OryDenialError):
|
|
48
|
+
wrapped()
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
def test_real_functiontool_runs_on_allow(monkeypatch):
|
|
52
|
+
_stub_gate(monkeypatch, blocked=False)
|
|
53
|
+
tool = _real_tool()
|
|
54
|
+
guard_tools([tool], client=create_mock_client(harness="llamaindex"))
|
|
55
|
+
wrapped = getattr(tool, "_fn", None) or getattr(tool, "fn", None)
|
|
56
|
+
assert wrapped() == "RAN"
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
def test_real_functiontool_acall_blocked_on_deny(monkeypatch):
|
|
60
|
+
# acall() invokes _async_fn (the FunctionAgent workflow path) — it must be gated too.
|
|
61
|
+
_stub_gate(monkeypatch, blocked=True)
|
|
62
|
+
tool = _real_tool()
|
|
63
|
+
guard_tools([tool], client=create_mock_client(harness="llamaindex"))
|
|
64
|
+
with pytest.raises(OryDenialError):
|
|
65
|
+
asyncio.run(tool.acall())
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
def test_real_functiontool_acall_runs_on_allow(monkeypatch):
|
|
69
|
+
_stub_gate(monkeypatch, blocked=False)
|
|
70
|
+
tool = _real_tool()
|
|
71
|
+
guard_tools([tool], client=create_mock_client(harness="llamaindex"))
|
|
72
|
+
result = asyncio.run(tool.acall())
|
|
73
|
+
assert "RAN" in str(result)
|
|
@@ -0,0 +1,167 @@
|
|
|
1
|
+
"""Hermetic tests for the LlamaIndex tool-wrap gate (SDK-free).
|
|
2
|
+
|
|
3
|
+
Covers the native veto shape (``OryDenialError`` raised from the wrapped ``_fn``) and
|
|
4
|
+
request-shape extraction (``metadata.name`` on a FunctionTool-shaped object). The full
|
|
5
|
+
permission decision matrix is core-owned (``ory-argus/tests/test_adapters.py``).
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
import asyncio
|
|
11
|
+
from types import SimpleNamespace
|
|
12
|
+
|
|
13
|
+
import pytest
|
|
14
|
+
|
|
15
|
+
from ory_argus.client import PrincipalIdentity
|
|
16
|
+
from ory_argus.denial import OryDenialError
|
|
17
|
+
from ory_argus.testing import (
|
|
18
|
+
create_mock_client,
|
|
19
|
+
get_trace_spans,
|
|
20
|
+
stub_permission_allowed,
|
|
21
|
+
stub_permission_denied,
|
|
22
|
+
)
|
|
23
|
+
from ory_llamaindex import guard_tool_fn, guard_tools # type: ignore
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
@pytest.fixture(autouse=True)
|
|
27
|
+
def _isolate(monkeypatch, tmp_path):
|
|
28
|
+
monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path))
|
|
29
|
+
monkeypatch.delenv("ORY_PERMISSION_MODE", raising=False)
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def client_with_user():
|
|
33
|
+
c = create_mock_client(harness="llamaindex")
|
|
34
|
+
c.set_user_principal(PrincipalIdentity(subject="user:alice"))
|
|
35
|
+
return c
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def fake_tool(name, fn):
|
|
39
|
+
# Mirrors FunctionTool: metadata.name + _fn.
|
|
40
|
+
return SimpleNamespace(metadata=SimpleNamespace(name=name), _fn=fn)
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
def fake_async_tool(name, fn, afn):
|
|
44
|
+
# Mirrors FunctionTool.from_defaults: metadata.name + _fn + _async_fn (acall path).
|
|
45
|
+
return SimpleNamespace(metadata=SimpleNamespace(name=name), _fn=fn, _async_fn=afn)
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
def test_guard_tool_fn_allow():
|
|
49
|
+
c = client_with_user()
|
|
50
|
+
stub_permission_allowed(c)
|
|
51
|
+
wrapped = guard_tool_fn(c, tool_name="search", fn=lambda q: f"hit:{q}")
|
|
52
|
+
assert wrapped("x") == "hit:x"
|
|
53
|
+
assert get_trace_spans(c, "tool.invoke")[0].attributes["allowed"] is True
|
|
54
|
+
assert get_trace_spans(c, "tool.complete")
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def test_guard_tools_wraps_fn_and_enforces(monkeypatch):
|
|
58
|
+
monkeypatch.setenv("ORY_PERMISSION_MODE", "enforce")
|
|
59
|
+
c = client_with_user()
|
|
60
|
+
stub_permission_denied(c)
|
|
61
|
+
ran = {}
|
|
62
|
+
|
|
63
|
+
def rm():
|
|
64
|
+
ran["yes"] = True
|
|
65
|
+
return "deleted"
|
|
66
|
+
|
|
67
|
+
tool = fake_tool("rm", rm)
|
|
68
|
+
guard_tools([tool], client=c)
|
|
69
|
+
# Native veto shape: the replaced _fn raises OryDenialError at the tool boundary.
|
|
70
|
+
with pytest.raises(OryDenialError):
|
|
71
|
+
tool._fn()
|
|
72
|
+
assert "yes" not in ran # underlying fn never invoked
|
|
73
|
+
assert get_trace_spans(c, "tool.block")[0].attributes["blocked"] is True
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def test_guard_tools_observe_runs():
|
|
77
|
+
c = client_with_user()
|
|
78
|
+
stub_permission_denied(c)
|
|
79
|
+
tool = fake_tool("rm", lambda: "ran")
|
|
80
|
+
guard_tools([tool], client=c)
|
|
81
|
+
assert tool._fn() == "ran"
|
|
82
|
+
assert get_trace_spans(c, "permission.observe_deny")
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
def test_guard_tools_async_path_enforce_blocks_before_run(monkeypatch):
|
|
86
|
+
monkeypatch.setenv("ORY_PERMISSION_MODE", "enforce")
|
|
87
|
+
c = client_with_user()
|
|
88
|
+
stub_permission_denied(c)
|
|
89
|
+
ran = {}
|
|
90
|
+
|
|
91
|
+
async def arm():
|
|
92
|
+
ran["yes"] = True
|
|
93
|
+
return "deleted"
|
|
94
|
+
|
|
95
|
+
tool = fake_async_tool("rm", lambda: "deleted", arm)
|
|
96
|
+
guard_tools([tool], client=c)
|
|
97
|
+
with pytest.raises(OryDenialError):
|
|
98
|
+
asyncio.run(tool._async_fn())
|
|
99
|
+
assert "yes" not in ran
|
|
100
|
+
assert get_trace_spans(c, "tool.block")[0].attributes["blocked"] is True
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
def test_guard_tools_async_path_allow_runs():
|
|
104
|
+
c = client_with_user()
|
|
105
|
+
stub_permission_allowed(c)
|
|
106
|
+
|
|
107
|
+
async def search(q):
|
|
108
|
+
return f"hit:{q}"
|
|
109
|
+
|
|
110
|
+
tool = fake_async_tool("search", lambda q: f"hit:{q}", search)
|
|
111
|
+
guard_tools([tool], client=c)
|
|
112
|
+
assert asyncio.run(tool._async_fn("x")) == "hit:x"
|
|
113
|
+
assert get_trace_spans(c, "tool.invoke")[0].attributes["allowed"] is True
|
|
114
|
+
assert get_trace_spans(c, "tool.complete")
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
def test_guard_tools_async_path_observe_runs():
|
|
118
|
+
c = client_with_user()
|
|
119
|
+
stub_permission_denied(c)
|
|
120
|
+
|
|
121
|
+
async def rm():
|
|
122
|
+
return "ran"
|
|
123
|
+
|
|
124
|
+
tool = fake_async_tool("rm", lambda: "ran", rm)
|
|
125
|
+
guard_tools([tool], client=c)
|
|
126
|
+
assert asyncio.run(tool._async_fn()) == "ran"
|
|
127
|
+
assert get_trace_spans(c, "permission.observe_deny")
|
|
128
|
+
|
|
129
|
+
|
|
130
|
+
def test_guard_tools_async_only_tool_is_gated(monkeypatch):
|
|
131
|
+
monkeypatch.setenv("ORY_PERMISSION_MODE", "enforce")
|
|
132
|
+
c = client_with_user()
|
|
133
|
+
stub_permission_denied(c)
|
|
134
|
+
|
|
135
|
+
async def rm():
|
|
136
|
+
return "ran"
|
|
137
|
+
|
|
138
|
+
tool = SimpleNamespace(metadata=SimpleNamespace(name="rm"), _async_fn=rm)
|
|
139
|
+
guard_tools([tool], client=c)
|
|
140
|
+
with pytest.raises(OryDenialError):
|
|
141
|
+
asyncio.run(tool._async_fn())
|
|
142
|
+
|
|
143
|
+
|
|
144
|
+
def test_guard_tools_wrap_failure_warns_on_stderr(capsys):
|
|
145
|
+
c = client_with_user()
|
|
146
|
+
stub_permission_allowed(c)
|
|
147
|
+
|
|
148
|
+
class FrozenTool:
|
|
149
|
+
metadata = SimpleNamespace(name="frozen")
|
|
150
|
+
|
|
151
|
+
@property
|
|
152
|
+
def _fn(self):
|
|
153
|
+
return lambda: "x"
|
|
154
|
+
|
|
155
|
+
guard_tools([FrozenTool()], client=c)
|
|
156
|
+
err = capsys.readouterr().err
|
|
157
|
+
assert "UNGATED" in err
|
|
158
|
+
assert "frozen" in err
|
|
159
|
+
|
|
160
|
+
|
|
161
|
+
def test_guard_tools_uses_metadata_name():
|
|
162
|
+
c = client_with_user()
|
|
163
|
+
stub_permission_allowed(c)
|
|
164
|
+
tool = fake_tool("custom_name", lambda: "ok")
|
|
165
|
+
guard_tools([tool], client=c)
|
|
166
|
+
tool._fn()
|
|
167
|
+
assert get_trace_spans(c, "tool.invoke")[0].attributes["toolName"] == "custom_name"
|