axiorank 0.1.0__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.
@@ -0,0 +1,37 @@
1
+ # dependencies
2
+ node_modules/
3
+ .pnpm-store/
4
+
5
+ # builds
6
+ .next/
7
+ out/
8
+ dist/
9
+ build/
10
+
11
+ # env / secrets
12
+ .env
13
+ .env.local
14
+ .env*.local
15
+
16
+ # logs
17
+ *.log
18
+ npm-debug.log*
19
+ pnpm-debug.log*
20
+
21
+ # misc
22
+ .DS_Store
23
+ .vercel
24
+ .turbo
25
+ coverage/
26
+ *.tsbuildinfo
27
+ next-env.d.ts
28
+
29
+ # python (packages/sdk-python)
30
+ __pycache__/
31
+ *.py[cod]
32
+ .venv/
33
+ venv/
34
+ .pytest_cache/
35
+ .mypy_cache/
36
+ .ruff_cache/
37
+ *.egg-info
@@ -0,0 +1,17 @@
1
+ # Changelog
2
+
3
+ All notable changes to the `axiorank` Python SDK are documented here.
4
+
5
+ ## 0.1.0
6
+
7
+ Initial release. Parity with `@axiorank/sdk` (TypeScript).
8
+
9
+ - `AxioRank` (sync) and `AsyncAxioRank` (async) clients over `httpx`.
10
+ - `tool_call` / `enforce` — route a tool call through the gateway; `enforce`
11
+ raises `AxioRankDeniedError` on a `deny` verdict. Transparently waits out a
12
+ `require_approval` hold and resolves to the final `allow` / `deny`.
13
+ - `verify_card` / `enforce_card` — preflight an external MCP server or A2A
14
+ agent before trusting it.
15
+ - Typed result dataclasses and the full error taxonomy.
16
+ - LangChain integration (`axiorank[langchain]`): `AxioRankCallbackHandler`,
17
+ `AxioRankAsyncCallbackHandler`, and `guard_tool`.
axiorank-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 AxioRank
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,146 @@
1
+ Metadata-Version: 2.4
2
+ Name: axiorank
3
+ Version: 0.1.0
4
+ Summary: Official Python SDK for AxioRank — route AI agent tool calls through your agent firewall.
5
+ Project-URL: Homepage, https://axiorank.com
6
+ Project-URL: Documentation, https://app.axiorank.com/docs
7
+ Project-URL: Source, https://github.com/frostyhand/AxioRank
8
+ Project-URL: Issues, https://github.com/frostyhand/AxioRank/issues
9
+ Author: AxioRank
10
+ License: MIT
11
+ License-File: LICENSE
12
+ Keywords: agent,ai,axiorank,firewall,gateway,llm,security
13
+ Classifier: Development Status :: 4 - Beta
14
+ Classifier: Intended Audience :: Developers
15
+ Classifier: License :: OSI Approved :: MIT License
16
+ Classifier: Operating System :: OS Independent
17
+ Classifier: Programming Language :: Python :: 3
18
+ Classifier: Programming Language :: Python :: 3.9
19
+ Classifier: Programming Language :: Python :: 3.10
20
+ Classifier: Programming Language :: Python :: 3.11
21
+ Classifier: Programming Language :: Python :: 3.12
22
+ Classifier: Programming Language :: Python :: 3.13
23
+ Classifier: Topic :: Security
24
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
25
+ Classifier: Typing :: Typed
26
+ Requires-Python: >=3.9
27
+ Requires-Dist: httpx<1,>=0.27
28
+ Provides-Extra: dev
29
+ Requires-Dist: mypy>=1.8; extra == 'dev'
30
+ Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
31
+ Requires-Dist: pytest>=8; extra == 'dev'
32
+ Requires-Dist: ruff>=0.4; extra == 'dev'
33
+ Provides-Extra: langchain
34
+ Requires-Dist: langchain-core>=0.3; extra == 'langchain'
35
+ Description-Content-Type: text/markdown
36
+
37
+ # axiorank
38
+
39
+ Official Python SDK for [AxioRank](https://axiorank.com) — route your AI
40
+ agent's tool calls through your agent firewall so policies are enforced, risk is
41
+ scored, and every call is audited.
42
+
43
+ Parity with the TypeScript [`@axiorank/sdk`](https://www.npmjs.com/package/@axiorank/sdk).
44
+
45
+ ## Install
46
+
47
+ ```bash
48
+ pip install axiorank
49
+ ```
50
+
51
+ ## Quickstart (sync)
52
+
53
+ ```python
54
+ import os
55
+ from axiorank import AxioRank
56
+
57
+ axio = AxioRank(api_key=os.environ["AXIORANK_API_KEY"]) # looks like axr_live_...
58
+
59
+ # Get the decision and act on it yourself:
60
+ result = axio.tool_call("github.push", {"repo": "myrepo"})
61
+ if result.decision == "deny":
62
+ print(f"Blocked: {result.reason} (risk {result.risk})")
63
+ else:
64
+ ... # proceed with the real tool call
65
+ ```
66
+
67
+ ### `enforce()` — guard in one line
68
+
69
+ ```python
70
+ from axiorank import AxioRank, AxioRankDeniedError
71
+
72
+ axio = AxioRank(api_key=os.environ["AXIORANK_API_KEY"])
73
+
74
+ try:
75
+ axio.enforce("aws.delete_bucket", {"name": "prod-data"})
76
+ delete_bucket("prod-data") # only runs if AxioRank allowed it
77
+ except AxioRankDeniedError as e:
78
+ log.warning("AxioRank blocked the call: %s", e.result.reason)
79
+ ```
80
+
81
+ A `require_approval` policy holds the call for a human; `tool_call`/`enforce`
82
+ transparently wait out the hold and resolve to the final `allow`/`deny`.
83
+
84
+ ## Quickstart (async)
85
+
86
+ ```python
87
+ import os
88
+ from axiorank import AsyncAxioRank
89
+
90
+ async def main():
91
+ async with AsyncAxioRank(api_key=os.environ["AXIORANK_API_KEY"]) as axio:
92
+ result = await axio.tool_call("stripe.refund", {"amount": 5000})
93
+ print(result.decision, result.risk)
94
+ ```
95
+
96
+ ## Preflight an external server before trusting it
97
+
98
+ ```python
99
+ result = axio.verify_card(url="https://mcp.acme.com")
100
+ print(result.decision, result.identity.signature_valid, result.protocol)
101
+
102
+ # Or guard in one line — raises AxioRankCardDeniedError on `deny`:
103
+ axio.enforce_card(url="https://mcp.acme.com")
104
+ ```
105
+
106
+ ## LangChain
107
+
108
+ ```bash
109
+ pip install "axiorank[langchain]"
110
+ ```
111
+
112
+ ```python
113
+ from axiorank import AxioRank
114
+ from axiorank.integrations.langchain import AxioRankCallbackHandler
115
+
116
+ axio = AxioRank(api_key=os.environ["AXIORANK_API_KEY"])
117
+
118
+ # Every tool the agent runs is checked against your policies first; a blocked
119
+ # tool raises and the step fails.
120
+ agent_executor.invoke(
121
+ {"input": "..."},
122
+ config={"callbacks": [AxioRankCallbackHandler(axio)]},
123
+ )
124
+ ```
125
+
126
+ Prefer a model-readable refusal over a raised exception? Wrap individual tools:
127
+
128
+ ```python
129
+ from axiorank.integrations.langchain import guard_tool
130
+
131
+ safe_tool = guard_tool(my_tool, axio, on_deny="return")
132
+ ```
133
+
134
+ ## Configuration
135
+
136
+ | Argument | Default | Notes |
137
+ | ------------------ | ----------------------------- | ------------------------------------------------ |
138
+ | `api_key` | — (required) | Your agent's key (`axr_live_...`). |
139
+ | `base_url` | `https://app.axiorank.com` | Point at your own deployment. |
140
+ | `timeout` | `10.0` | Per-request timeout, in seconds. |
141
+ | `approval_timeout` | `300.0` | Max wait for a human to resolve a hold, seconds. |
142
+ | `client` | a fresh `httpx.Client` | Inject your own `httpx` client (tests, proxies). |
143
+
144
+ ## License
145
+
146
+ MIT
@@ -0,0 +1,110 @@
1
+ # axiorank
2
+
3
+ Official Python SDK for [AxioRank](https://axiorank.com) — route your AI
4
+ agent's tool calls through your agent firewall so policies are enforced, risk is
5
+ scored, and every call is audited.
6
+
7
+ Parity with the TypeScript [`@axiorank/sdk`](https://www.npmjs.com/package/@axiorank/sdk).
8
+
9
+ ## Install
10
+
11
+ ```bash
12
+ pip install axiorank
13
+ ```
14
+
15
+ ## Quickstart (sync)
16
+
17
+ ```python
18
+ import os
19
+ from axiorank import AxioRank
20
+
21
+ axio = AxioRank(api_key=os.environ["AXIORANK_API_KEY"]) # looks like axr_live_...
22
+
23
+ # Get the decision and act on it yourself:
24
+ result = axio.tool_call("github.push", {"repo": "myrepo"})
25
+ if result.decision == "deny":
26
+ print(f"Blocked: {result.reason} (risk {result.risk})")
27
+ else:
28
+ ... # proceed with the real tool call
29
+ ```
30
+
31
+ ### `enforce()` — guard in one line
32
+
33
+ ```python
34
+ from axiorank import AxioRank, AxioRankDeniedError
35
+
36
+ axio = AxioRank(api_key=os.environ["AXIORANK_API_KEY"])
37
+
38
+ try:
39
+ axio.enforce("aws.delete_bucket", {"name": "prod-data"})
40
+ delete_bucket("prod-data") # only runs if AxioRank allowed it
41
+ except AxioRankDeniedError as e:
42
+ log.warning("AxioRank blocked the call: %s", e.result.reason)
43
+ ```
44
+
45
+ A `require_approval` policy holds the call for a human; `tool_call`/`enforce`
46
+ transparently wait out the hold and resolve to the final `allow`/`deny`.
47
+
48
+ ## Quickstart (async)
49
+
50
+ ```python
51
+ import os
52
+ from axiorank import AsyncAxioRank
53
+
54
+ async def main():
55
+ async with AsyncAxioRank(api_key=os.environ["AXIORANK_API_KEY"]) as axio:
56
+ result = await axio.tool_call("stripe.refund", {"amount": 5000})
57
+ print(result.decision, result.risk)
58
+ ```
59
+
60
+ ## Preflight an external server before trusting it
61
+
62
+ ```python
63
+ result = axio.verify_card(url="https://mcp.acme.com")
64
+ print(result.decision, result.identity.signature_valid, result.protocol)
65
+
66
+ # Or guard in one line — raises AxioRankCardDeniedError on `deny`:
67
+ axio.enforce_card(url="https://mcp.acme.com")
68
+ ```
69
+
70
+ ## LangChain
71
+
72
+ ```bash
73
+ pip install "axiorank[langchain]"
74
+ ```
75
+
76
+ ```python
77
+ from axiorank import AxioRank
78
+ from axiorank.integrations.langchain import AxioRankCallbackHandler
79
+
80
+ axio = AxioRank(api_key=os.environ["AXIORANK_API_KEY"])
81
+
82
+ # Every tool the agent runs is checked against your policies first; a blocked
83
+ # tool raises and the step fails.
84
+ agent_executor.invoke(
85
+ {"input": "..."},
86
+ config={"callbacks": [AxioRankCallbackHandler(axio)]},
87
+ )
88
+ ```
89
+
90
+ Prefer a model-readable refusal over a raised exception? Wrap individual tools:
91
+
92
+ ```python
93
+ from axiorank.integrations.langchain import guard_tool
94
+
95
+ safe_tool = guard_tool(my_tool, axio, on_deny="return")
96
+ ```
97
+
98
+ ## Configuration
99
+
100
+ | Argument | Default | Notes |
101
+ | ------------------ | ----------------------------- | ------------------------------------------------ |
102
+ | `api_key` | — (required) | Your agent's key (`axr_live_...`). |
103
+ | `base_url` | `https://app.axiorank.com` | Point at your own deployment. |
104
+ | `timeout` | `10.0` | Per-request timeout, in seconds. |
105
+ | `approval_timeout` | `300.0` | Max wait for a human to resolve a hold, seconds. |
106
+ | `client` | a fresh `httpx.Client` | Inject your own `httpx` client (tests, proxies). |
107
+
108
+ ## License
109
+
110
+ MIT
@@ -0,0 +1,74 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "axiorank"
7
+ dynamic = ["version"]
8
+ description = "Official Python SDK for AxioRank — route AI agent tool calls through your agent firewall."
9
+ readme = "README.md"
10
+ requires-python = ">=3.9"
11
+ license = { text = "MIT" }
12
+ authors = [{ name = "AxioRank" }]
13
+ keywords = ["axiorank", "ai", "agent", "firewall", "security", "gateway", "llm"]
14
+ classifiers = [
15
+ "Development Status :: 4 - Beta",
16
+ "Intended Audience :: Developers",
17
+ "License :: OSI Approved :: MIT License",
18
+ "Operating System :: OS Independent",
19
+ "Programming Language :: Python :: 3",
20
+ "Programming Language :: Python :: 3.9",
21
+ "Programming Language :: Python :: 3.10",
22
+ "Programming Language :: Python :: 3.11",
23
+ "Programming Language :: Python :: 3.12",
24
+ "Programming Language :: Python :: 3.13",
25
+ "Topic :: Security",
26
+ "Topic :: Software Development :: Libraries :: Python Modules",
27
+ "Typing :: Typed",
28
+ ]
29
+ dependencies = ["httpx>=0.27,<1"]
30
+
31
+ [project.optional-dependencies]
32
+ # Framework integrations (lazy-imported; install what you use).
33
+ langchain = ["langchain-core>=0.3"]
34
+ # Tooling for local development and CI.
35
+ dev = [
36
+ "pytest>=8",
37
+ "pytest-asyncio>=0.23",
38
+ "mypy>=1.8",
39
+ "ruff>=0.4",
40
+ ]
41
+
42
+ [project.urls]
43
+ Homepage = "https://axiorank.com"
44
+ Documentation = "https://app.axiorank.com/docs"
45
+ Source = "https://github.com/frostyhand/AxioRank"
46
+ Issues = "https://github.com/frostyhand/AxioRank/issues"
47
+
48
+ [tool.hatch.version]
49
+ path = "src/axiorank/__init__.py"
50
+
51
+ [tool.hatch.build.targets.wheel]
52
+ packages = ["src/axiorank"]
53
+
54
+ [tool.hatch.build.targets.sdist]
55
+ include = ["src/axiorank", "README.md", "CHANGELOG.md", "LICENSE"]
56
+
57
+ [tool.pytest.ini_options]
58
+ asyncio_mode = "auto"
59
+ testpaths = ["tests"]
60
+
61
+ [tool.ruff]
62
+ line-length = 100
63
+ target-version = "py39"
64
+ src = ["src", "tests"]
65
+
66
+ [tool.ruff.lint]
67
+ select = ["E", "F", "I", "UP", "B"]
68
+
69
+ # Type-checked at 3.10 (the floor this mypy supports); runtime stays 3.9-safe,
70
+ # enforced by ruff's py39 target + `from __future__ import annotations`.
71
+ [tool.mypy]
72
+ python_version = "3.10"
73
+ strict = true
74
+ files = ["src/axiorank"]
@@ -0,0 +1,64 @@
1
+ """AxioRank Python SDK — route AI agent tool calls through your agent firewall.
2
+
3
+ Parity with ``@axiorank/sdk`` (TypeScript). Sync and async clients::
4
+
5
+ from axiorank import AxioRank, AsyncAxioRank
6
+
7
+ axio = AxioRank(api_key="axr_live_...")
8
+ result = axio.tool_call("github.push", {"repo": "myrepo"})
9
+
10
+ See https://app.axiorank.com/docs for the full guide.
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ from ._async import AsyncAxioRank
16
+ from ._errors import (
17
+ AxioRankAuthError,
18
+ AxioRankCardDeniedError,
19
+ AxioRankDeniedError,
20
+ AxioRankError,
21
+ AxioRankRequestError,
22
+ )
23
+ from ._sync import AxioRank
24
+ from ._types import (
25
+ CardAuth,
26
+ CardCapabilities,
27
+ CardCapabilitySample,
28
+ CardDecision,
29
+ CardIdentity,
30
+ CardVerifyResult,
31
+ Decision,
32
+ Severity,
33
+ SignalCategory,
34
+ ToolCallResult,
35
+ ToolCallSignal,
36
+ )
37
+
38
+ __version__ = "0.1.0"
39
+
40
+ __all__ = [
41
+ "__version__",
42
+ # clients
43
+ "AxioRank",
44
+ "AsyncAxioRank",
45
+ # errors
46
+ "AxioRankError",
47
+ "AxioRankAuthError",
48
+ "AxioRankRequestError",
49
+ "AxioRankDeniedError",
50
+ "AxioRankCardDeniedError",
51
+ # result types
52
+ "ToolCallResult",
53
+ "ToolCallSignal",
54
+ "CardVerifyResult",
55
+ "CardIdentity",
56
+ "CardCapabilities",
57
+ "CardCapabilitySample",
58
+ "CardAuth",
59
+ # literal aliases
60
+ "Decision",
61
+ "CardDecision",
62
+ "SignalCategory",
63
+ "Severity",
64
+ ]
@@ -0,0 +1,187 @@
1
+ """Asynchronous client over ``httpx.AsyncClient``."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import asyncio
6
+ import time
7
+ from collections.abc import Mapping
8
+ from types import TracebackType
9
+ from typing import Any
10
+
11
+ import httpx
12
+
13
+ from ._base import (
14
+ build_headers,
15
+ drop_none,
16
+ normalize_base_url,
17
+ process_response,
18
+ resolved_hold,
19
+ )
20
+ from ._constants import (
21
+ APPROVAL_POLL_BACKOFF,
22
+ APPROVAL_POLL_TIMEOUT,
23
+ APPROVALS_PATH,
24
+ DEFAULT_APPROVAL_TIMEOUT,
25
+ DEFAULT_TIMEOUT,
26
+ TOOL_CALL_PATH,
27
+ VERIFY_CARD_PATH,
28
+ )
29
+ from ._errors import (
30
+ AxioRankAuthError,
31
+ AxioRankCardDeniedError,
32
+ AxioRankDeniedError,
33
+ AxioRankError,
34
+ AxioRankRequestError,
35
+ )
36
+ from ._types import CardVerifyResult, ToolCallResult
37
+
38
+
39
+ class AsyncAxioRank:
40
+ """Route AI agent tool calls through your AxioRank gateway (asyncio).
41
+
42
+ Example::
43
+
44
+ from axiorank import AsyncAxioRank
45
+
46
+ async with AsyncAxioRank(api_key=os.environ["AXIORANK_API_KEY"]) as axio:
47
+ await axio.enforce("aws.delete_bucket", {"name": "prod"})
48
+ """
49
+
50
+ def __init__(
51
+ self,
52
+ api_key: str,
53
+ *,
54
+ base_url: str | None = None,
55
+ timeout: float | None = None,
56
+ approval_timeout: float | None = None,
57
+ client: httpx.AsyncClient | None = None,
58
+ ) -> None:
59
+ if not api_key:
60
+ raise AxioRankError("AxioRank: `api_key` is required")
61
+ self._api_key = api_key
62
+ self._base_url = normalize_base_url(base_url)
63
+ self._timeout = DEFAULT_TIMEOUT if timeout is None else timeout
64
+ self._approval_timeout = (
65
+ DEFAULT_APPROVAL_TIMEOUT if approval_timeout is None else approval_timeout
66
+ )
67
+ self._headers = build_headers(api_key)
68
+ self._client = client if client is not None else httpx.AsyncClient()
69
+ self._owns_client = client is None
70
+
71
+ # ── lifecycle ────────────────────────────────────────────────
72
+ async def aclose(self) -> None:
73
+ """Close the underlying httpx client (only if the SDK created it)."""
74
+ if self._owns_client:
75
+ await self._client.aclose()
76
+
77
+ async def __aenter__(self) -> AsyncAxioRank:
78
+ return self
79
+
80
+ async def __aexit__(
81
+ self,
82
+ exc_type: type[BaseException] | None,
83
+ exc: BaseException | None,
84
+ tb: TracebackType | None,
85
+ ) -> None:
86
+ await self.aclose()
87
+
88
+ # ── outbound tool calls ──────────────────────────────────────
89
+ async def tool_call(
90
+ self, tool: str, arguments: Mapping[str, Any] | None = None
91
+ ) -> ToolCallResult:
92
+ """Send a tool call to the gateway and return the policy decision.
93
+
94
+ Resolves normally for an explicit ``deny``. A ``require_approval`` hold
95
+ is waited out transparently, so callers only ever see ``allow``/``deny``.
96
+ """
97
+ if not tool:
98
+ raise AxioRankError("AxioRank: `tool` is required")
99
+ body = await self._post(TOOL_CALL_PATH, {"tool": tool, "arguments": arguments or {}})
100
+ if body.get("decision") == "hold" and body.get("approvalId"):
101
+ return await self._wait_for_approval(str(body["approvalId"]), body)
102
+ return ToolCallResult.from_dict(body)
103
+
104
+ async def enforce(
105
+ self, tool: str, arguments: Mapping[str, Any] | None = None
106
+ ) -> ToolCallResult:
107
+ """Like :meth:`tool_call`, but raise :class:`AxioRankDeniedError` on deny."""
108
+ result = await self.tool_call(tool, arguments)
109
+ if result.decision == "deny":
110
+ raise AxioRankDeniedError(result)
111
+ return result
112
+
113
+ # ── preflight (card verification) ────────────────────────────
114
+ async def verify_card(
115
+ self,
116
+ *,
117
+ url: str | None = None,
118
+ document: Any | None = None,
119
+ protocol: str | None = None,
120
+ ) -> CardVerifyResult:
121
+ """Preflight an external MCP server / A2A agent before trusting it."""
122
+ if url is None and document is None:
123
+ raise AxioRankError("AxioRank: `url` or `document` is required")
124
+ body = await self._post(
125
+ VERIFY_CARD_PATH,
126
+ drop_none({"url": url, "document": document, "protocol": protocol}),
127
+ )
128
+ return CardVerifyResult.from_dict(body)
129
+
130
+ async def enforce_card(
131
+ self,
132
+ *,
133
+ url: str | None = None,
134
+ document: Any | None = None,
135
+ protocol: str | None = None,
136
+ ) -> CardVerifyResult:
137
+ """Like :meth:`verify_card`, but raise on a ``deny`` verdict."""
138
+ result = await self.verify_card(url=url, document=document, protocol=protocol)
139
+ if result.decision == "deny":
140
+ raise AxioRankCardDeniedError(result)
141
+ return result
142
+
143
+ # ── internals ────────────────────────────────────────────────
144
+ async def _post(self, path: str, payload: Mapping[str, Any]) -> dict[str, Any]:
145
+ try:
146
+ response = await self._client.post(
147
+ f"{self._base_url}{path}",
148
+ headers=self._headers,
149
+ json=payload,
150
+ timeout=self._timeout,
151
+ )
152
+ except httpx.TimeoutException as err:
153
+ raise AxioRankRequestError(
154
+ f"AxioRank: request timed out after {self._timeout}s"
155
+ ) from err
156
+ except httpx.RequestError as err:
157
+ raise AxioRankRequestError(f"AxioRank: network error — {err}") from err
158
+ return process_response(response)
159
+
160
+ async def _wait_for_approval(self, approval_id: str, held: Mapping[str, Any]) -> ToolCallResult:
161
+ """Poll the approvals endpoint until a held call resolves (the gateway
162
+ long-polls, so this is cheap). After ``approval_timeout`` give up and
163
+ return ``deny`` — the gateway also auto-denies after its own TTL."""
164
+ url = f"{self._base_url}{APPROVALS_PATH}/{approval_id}"
165
+ deadline = time.monotonic() + self._approval_timeout
166
+
167
+ while time.monotonic() < deadline:
168
+ try:
169
+ response = await self._client.get(
170
+ url, headers=self._headers, timeout=APPROVAL_POLL_TIMEOUT
171
+ )
172
+ if response.status_code == 401:
173
+ raise AxioRankAuthError()
174
+ status: Any = response.json()
175
+ except AxioRankAuthError:
176
+ raise
177
+ except Exception:
178
+ # Transient network/timeout while polling — back off and retry.
179
+ await asyncio.sleep(APPROVAL_POLL_BACKOFF)
180
+ continue
181
+
182
+ decision = status.get("decision") if isinstance(status, dict) else None
183
+ if decision and decision != "hold":
184
+ reason = status.get("reason") if isinstance(status, dict) else None
185
+ return resolved_hold(held, decision, reason)
186
+
187
+ return resolved_hold(held, "deny", "AxioRank: approval timed out")
@@ -0,0 +1,74 @@
1
+ """Transport-agnostic helpers shared by the sync and async clients.
2
+
3
+ The only thing that differs between :class:`AxioRank` and
4
+ :class:`AsyncAxioRank` is *how* a request is awaited — building the request and
5
+ mapping the response are identical (httpx reads the body during the request, so
6
+ ``response.json()`` is synchronous on both clients). That shared logic lives
7
+ here.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ from collections.abc import Mapping
13
+ from typing import Any
14
+
15
+ import httpx
16
+
17
+ from ._constants import DEFAULT_BASE_URL
18
+ from ._errors import AxioRankAuthError, AxioRankRequestError
19
+ from ._types import ToolCallResult, _parse_signals
20
+
21
+
22
+ def normalize_base_url(base_url: str | None) -> str:
23
+ """Strip trailing slashes; fall back to the public default."""
24
+ return (base_url or DEFAULT_BASE_URL).rstrip("/")
25
+
26
+
27
+ def build_headers(api_key: str) -> dict[str, str]:
28
+ return {
29
+ "content-type": "application/json",
30
+ "authorization": f"Bearer {api_key}",
31
+ }
32
+
33
+
34
+ def drop_none(payload: dict[str, Any]) -> dict[str, Any]:
35
+ """Drop ``None`` values so the wire matches JS's ``JSON.stringify`` (which
36
+ omits ``undefined``) — the gateway's optional fields reject ``null``."""
37
+ return {k: v for k, v in payload.items() if v is not None}
38
+
39
+
40
+ def process_response(response: httpx.Response) -> dict[str, Any]:
41
+ """Map an httpx response to a parsed body dict, raising the SDK's errors.
42
+
43
+ Mirrors the TS client: 401 → :class:`AxioRankAuthError`; any other non-2xx
44
+ → :class:`AxioRankRequestError` (carrying the server's ``error`` message and
45
+ the status code); a 2xx with an unparseable body resolves to ``{}``.
46
+ """
47
+ if response.status_code == 401:
48
+ raise AxioRankAuthError()
49
+
50
+ body: Any
51
+ try:
52
+ body = response.json()
53
+ except Exception:
54
+ body = None
55
+
56
+ if not response.is_success:
57
+ message = body.get("error") if isinstance(body, dict) else None
58
+ if not message:
59
+ message = f"request failed with status {response.status_code}"
60
+ raise AxioRankRequestError(f"AxioRank: {message}", response.status_code)
61
+
62
+ return body if isinstance(body, dict) else {}
63
+
64
+
65
+ def resolved_hold(held: Mapping[str, Any], decision: str, reason: str | None) -> ToolCallResult:
66
+ """Build the final result for a resolved approval hold, carrying forward the
67
+ risk / audit-log id / signals from the original held response."""
68
+ return ToolCallResult(
69
+ decision=decision, # type: ignore[arg-type]
70
+ reason=reason or held.get("reason", ""),
71
+ risk=int(held.get("risk") or 0),
72
+ audit_log_id=held.get("auditLogId", ""),
73
+ signals=_parse_signals(held.get("signals")),
74
+ )
@@ -0,0 +1,26 @@
1
+ """Defaults and gateway paths, kept in lockstep with the TypeScript SDK.
2
+
3
+ Timeouts are expressed in **seconds** (httpx's unit), whereas the TS SDK uses
4
+ milliseconds — the values are otherwise identical.
5
+ """
6
+
7
+ DEFAULT_BASE_URL = "https://app.axiorank.com"
8
+
9
+ # Per-request timeout. TS: DEFAULT_TIMEOUT_MS = 10_000.
10
+ DEFAULT_TIMEOUT = 10.0
11
+
12
+ # Max time to wait for a human to resolve a `require_approval` hold before the
13
+ # call resolves to `deny`. TS: DEFAULT_APPROVAL_TIMEOUT_MS = 300_000.
14
+ DEFAULT_APPROVAL_TIMEOUT = 300.0
15
+
16
+ # Per-poll timeout while waiting on an approval. MUST stay larger than the
17
+ # gateway's server-side long-poll window (~8s) — and is deliberately NOT the
18
+ # 10s request timeout, or a held call would abort mid-long-poll. TS: 20_000.
19
+ APPROVAL_POLL_TIMEOUT = 20.0
20
+
21
+ # Back-off after a transient error while polling an approval. TS: delay(1_000).
22
+ APPROVAL_POLL_BACKOFF = 1.0
23
+
24
+ TOOL_CALL_PATH = "/api/gateway/tool-call"
25
+ VERIFY_CARD_PATH = "/api/gateway/verify-card"
26
+ APPROVALS_PATH = "/api/gateway/approvals"
@@ -0,0 +1,43 @@
1
+ """Error taxonomy — a 1:1 port of `@axiorank/sdk`'s `errors.ts`."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import TYPE_CHECKING
6
+
7
+ if TYPE_CHECKING:
8
+ from ._types import CardVerifyResult, ToolCallResult
9
+
10
+
11
+ class AxioRankError(Exception):
12
+ """Base class for all errors raised by the SDK."""
13
+
14
+
15
+ class AxioRankAuthError(AxioRankError):
16
+ """Raised when the API key is missing, invalid, or revoked (HTTP 401)."""
17
+
18
+ def __init__(self, message: str = "Invalid or missing AxioRank API key") -> None:
19
+ super().__init__(message)
20
+
21
+
22
+ class AxioRankRequestError(AxioRankError):
23
+ """Raised for malformed requests, timeouts, and other non-2xx responses."""
24
+
25
+ def __init__(self, message: str, status: int | None = None) -> None:
26
+ super().__init__(message)
27
+ self.status: int | None = status
28
+
29
+
30
+ class AxioRankDeniedError(AxioRankError):
31
+ """Raised by :meth:`AxioRank.enforce` when the gateway denies the call."""
32
+
33
+ def __init__(self, result: ToolCallResult) -> None:
34
+ super().__init__(f"AxioRank denied tool call: {result.reason}")
35
+ self.result: ToolCallResult = result
36
+
37
+
38
+ class AxioRankCardDeniedError(AxioRankError):
39
+ """Raised by :meth:`AxioRank.enforce_card` when the verdict is ``deny``."""
40
+
41
+ def __init__(self, result: CardVerifyResult) -> None:
42
+ super().__init__(f"AxioRank denied card preflight: {result.reason}")
43
+ self.result: CardVerifyResult = result
@@ -0,0 +1,184 @@
1
+ """Synchronous client over ``httpx.Client``."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import time
6
+ from collections.abc import Mapping
7
+ from types import TracebackType
8
+ from typing import Any
9
+
10
+ import httpx
11
+
12
+ from ._base import (
13
+ build_headers,
14
+ drop_none,
15
+ normalize_base_url,
16
+ process_response,
17
+ resolved_hold,
18
+ )
19
+ from ._constants import (
20
+ APPROVAL_POLL_BACKOFF,
21
+ APPROVAL_POLL_TIMEOUT,
22
+ APPROVALS_PATH,
23
+ DEFAULT_APPROVAL_TIMEOUT,
24
+ DEFAULT_TIMEOUT,
25
+ TOOL_CALL_PATH,
26
+ VERIFY_CARD_PATH,
27
+ )
28
+ from ._errors import (
29
+ AxioRankAuthError,
30
+ AxioRankCardDeniedError,
31
+ AxioRankDeniedError,
32
+ AxioRankError,
33
+ AxioRankRequestError,
34
+ )
35
+ from ._types import CardVerifyResult, ToolCallResult
36
+
37
+
38
+ class AxioRank:
39
+ """Route AI agent tool calls through your AxioRank gateway (synchronous).
40
+
41
+ Example::
42
+
43
+ from axiorank import AxioRank
44
+
45
+ axio = AxioRank(api_key=os.environ["AXIORANK_API_KEY"])
46
+ result = axio.tool_call("github.push", {"repo": "myrepo"})
47
+ if result.decision == "deny":
48
+ ...
49
+ """
50
+
51
+ def __init__(
52
+ self,
53
+ api_key: str,
54
+ *,
55
+ base_url: str | None = None,
56
+ timeout: float | None = None,
57
+ approval_timeout: float | None = None,
58
+ client: httpx.Client | None = None,
59
+ ) -> None:
60
+ if not api_key:
61
+ raise AxioRankError("AxioRank: `api_key` is required")
62
+ self._api_key = api_key
63
+ self._base_url = normalize_base_url(base_url)
64
+ self._timeout = DEFAULT_TIMEOUT if timeout is None else timeout
65
+ self._approval_timeout = (
66
+ DEFAULT_APPROVAL_TIMEOUT if approval_timeout is None else approval_timeout
67
+ )
68
+ self._headers = build_headers(api_key)
69
+ self._client = client if client is not None else httpx.Client()
70
+ self._owns_client = client is None
71
+
72
+ # ── lifecycle ────────────────────────────────────────────────
73
+ def close(self) -> None:
74
+ """Close the underlying httpx client (only if the SDK created it)."""
75
+ if self._owns_client:
76
+ self._client.close()
77
+
78
+ def __enter__(self) -> AxioRank:
79
+ return self
80
+
81
+ def __exit__(
82
+ self,
83
+ exc_type: type[BaseException] | None,
84
+ exc: BaseException | None,
85
+ tb: TracebackType | None,
86
+ ) -> None:
87
+ self.close()
88
+
89
+ # ── outbound tool calls ──────────────────────────────────────
90
+ def tool_call(self, tool: str, arguments: Mapping[str, Any] | None = None) -> ToolCallResult:
91
+ """Send a tool call to the gateway and return the policy decision.
92
+
93
+ Resolves normally for an explicit ``deny``. A ``require_approval`` hold
94
+ is waited out transparently, so callers only ever see ``allow``/``deny``.
95
+ """
96
+ if not tool:
97
+ raise AxioRankError("AxioRank: `tool` is required")
98
+ body = self._post(TOOL_CALL_PATH, {"tool": tool, "arguments": arguments or {}})
99
+ if body.get("decision") == "hold" and body.get("approvalId"):
100
+ return self._wait_for_approval(str(body["approvalId"]), body)
101
+ return ToolCallResult.from_dict(body)
102
+
103
+ def enforce(self, tool: str, arguments: Mapping[str, Any] | None = None) -> ToolCallResult:
104
+ """Like :meth:`tool_call`, but raise :class:`AxioRankDeniedError` on deny."""
105
+ result = self.tool_call(tool, arguments)
106
+ if result.decision == "deny":
107
+ raise AxioRankDeniedError(result)
108
+ return result
109
+
110
+ # ── preflight (card verification) ────────────────────────────
111
+ def verify_card(
112
+ self,
113
+ *,
114
+ url: str | None = None,
115
+ document: Any | None = None,
116
+ protocol: str | None = None,
117
+ ) -> CardVerifyResult:
118
+ """Preflight an external MCP server / A2A agent before trusting it."""
119
+ if url is None and document is None:
120
+ raise AxioRankError("AxioRank: `url` or `document` is required")
121
+ body = self._post(
122
+ VERIFY_CARD_PATH,
123
+ drop_none({"url": url, "document": document, "protocol": protocol}),
124
+ )
125
+ return CardVerifyResult.from_dict(body)
126
+
127
+ def enforce_card(
128
+ self,
129
+ *,
130
+ url: str | None = None,
131
+ document: Any | None = None,
132
+ protocol: str | None = None,
133
+ ) -> CardVerifyResult:
134
+ """Like :meth:`verify_card`, but raise on a ``deny`` verdict."""
135
+ result = self.verify_card(url=url, document=document, protocol=protocol)
136
+ if result.decision == "deny":
137
+ raise AxioRankCardDeniedError(result)
138
+ return result
139
+
140
+ # ── internals ────────────────────────────────────────────────
141
+ def _post(self, path: str, payload: Mapping[str, Any]) -> dict[str, Any]:
142
+ try:
143
+ response = self._client.post(
144
+ f"{self._base_url}{path}",
145
+ headers=self._headers,
146
+ json=payload,
147
+ timeout=self._timeout,
148
+ )
149
+ except httpx.TimeoutException as err:
150
+ raise AxioRankRequestError(
151
+ f"AxioRank: request timed out after {self._timeout}s"
152
+ ) from err
153
+ except httpx.RequestError as err:
154
+ raise AxioRankRequestError(f"AxioRank: network error — {err}") from err
155
+ return process_response(response)
156
+
157
+ def _wait_for_approval(self, approval_id: str, held: Mapping[str, Any]) -> ToolCallResult:
158
+ """Poll the approvals endpoint until a held call resolves (the gateway
159
+ long-polls, so this is cheap). After ``approval_timeout`` give up and
160
+ return ``deny`` — the gateway also auto-denies after its own TTL."""
161
+ url = f"{self._base_url}{APPROVALS_PATH}/{approval_id}"
162
+ deadline = time.monotonic() + self._approval_timeout
163
+
164
+ while time.monotonic() < deadline:
165
+ try:
166
+ response = self._client.get(
167
+ url, headers=self._headers, timeout=APPROVAL_POLL_TIMEOUT
168
+ )
169
+ if response.status_code == 401:
170
+ raise AxioRankAuthError()
171
+ status: Any = response.json()
172
+ except AxioRankAuthError:
173
+ raise
174
+ except Exception:
175
+ # Transient network/timeout while polling — back off and retry.
176
+ time.sleep(APPROVAL_POLL_BACKOFF)
177
+ continue
178
+
179
+ decision = status.get("decision") if isinstance(status, dict) else None
180
+ if decision and decision != "hold":
181
+ reason = status.get("reason") if isinstance(status, dict) else None
182
+ return resolved_hold(held, decision, reason)
183
+
184
+ return resolved_hold(held, "deny", "AxioRank: approval timed out")
@@ -0,0 +1,175 @@
1
+ """Result types — frozen dataclasses mirroring `@axiorank/sdk`'s `types.ts`.
2
+
3
+ Wire payloads use camelCase keys (``auditLogId``); the dataclasses expose
4
+ snake_case attributes (``audit_log_id``). ``from_dict`` bridges the two and is
5
+ deliberately tolerant of missing / unknown fields so a newer gateway never
6
+ breaks an older SDK.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ from collections.abc import Mapping
12
+ from dataclasses import dataclass
13
+ from typing import Any, Literal
14
+
15
+ # Policy decision the gateway surfaces to a caller. (`hold` is internal — the
16
+ # client waits it out and only ever returns `allow`/`deny`.)
17
+ Decision = Literal["allow", "deny"]
18
+ # Three-way verdict for a card preflight. Precedence: deny > review > allow.
19
+ CardDecision = Literal["allow", "review", "deny"]
20
+ Severity = Literal["low", "medium", "high", "critical"]
21
+ # Category of a content-inspection finding. The last two are inbound-only.
22
+ SignalCategory = Literal[
23
+ "secret",
24
+ "pii",
25
+ "destructive",
26
+ "injection",
27
+ "egress",
28
+ "bot_spoof",
29
+ "rate_abuse",
30
+ ]
31
+
32
+
33
+ @dataclass(frozen=True)
34
+ class ToolCallSignal:
35
+ """A redacted content-inspection finding on a tool-call payload."""
36
+
37
+ detector: str
38
+ category: str
39
+ severity: str
40
+ label: str
41
+ location: str
42
+ evidence: str
43
+
44
+ @classmethod
45
+ def from_dict(cls, d: Mapping[str, Any]) -> ToolCallSignal:
46
+ return cls(
47
+ detector=d.get("detector", ""),
48
+ category=d.get("category", ""),
49
+ severity=d.get("severity", ""),
50
+ label=d.get("label", ""),
51
+ location=d.get("location", ""),
52
+ evidence=d.get("evidence", ""),
53
+ )
54
+
55
+
56
+ def _parse_signals(raw: Any) -> list[ToolCallSignal] | None:
57
+ if not raw:
58
+ return None
59
+ return [ToolCallSignal.from_dict(s) for s in raw]
60
+
61
+
62
+ @dataclass(frozen=True)
63
+ class ToolCallResult:
64
+ """Outcome of routing a tool call through the gateway."""
65
+
66
+ decision: Decision
67
+ reason: str
68
+ risk: int
69
+ audit_log_id: str
70
+ signals: list[ToolCallSignal] | None = None
71
+
72
+ @classmethod
73
+ def from_dict(cls, d: Mapping[str, Any]) -> ToolCallResult:
74
+ return cls(
75
+ decision=d.get("decision", "deny"),
76
+ reason=d.get("reason", ""),
77
+ risk=int(d.get("risk") or 0),
78
+ audit_log_id=d.get("auditLogId", ""),
79
+ signals=_parse_signals(d.get("signals")),
80
+ )
81
+
82
+
83
+ @dataclass(frozen=True)
84
+ class CardIdentity:
85
+ name: str | None
86
+ provider: str | None
87
+ signed: bool
88
+ # True = verified · False = forged · None = unsigned / key unresolvable.
89
+ signature_valid: bool | None
90
+ key_source: str # "embedded" | "jwks" | "none"
91
+ key_domain_bound: bool
92
+ key_id: str | None
93
+ source_url: str
94
+
95
+ @classmethod
96
+ def from_dict(cls, d: Mapping[str, Any]) -> CardIdentity:
97
+ return cls(
98
+ name=d.get("name"),
99
+ provider=d.get("provider"),
100
+ signed=bool(d.get("signed", False)),
101
+ signature_valid=d.get("signatureValid"),
102
+ key_source=d.get("keySource", "none"),
103
+ key_domain_bound=bool(d.get("keyDomainBound", False)),
104
+ key_id=d.get("keyId"),
105
+ source_url=d.get("sourceUrl", ""),
106
+ )
107
+
108
+
109
+ @dataclass(frozen=True)
110
+ class CardCapabilitySample:
111
+ kind: str
112
+ name: str
113
+
114
+ @classmethod
115
+ def from_dict(cls, d: Mapping[str, Any]) -> CardCapabilitySample:
116
+ return cls(kind=d.get("kind", ""), name=d.get("name", ""))
117
+
118
+
119
+ @dataclass(frozen=True)
120
+ class CardCapabilities:
121
+ count: int
122
+ sample: list[CardCapabilitySample]
123
+
124
+ @classmethod
125
+ def from_dict(cls, d: Mapping[str, Any]) -> CardCapabilities:
126
+ return cls(
127
+ count=int(d.get("count") or 0),
128
+ sample=[CardCapabilitySample.from_dict(s) for s in (d.get("sample") or [])],
129
+ )
130
+
131
+
132
+ @dataclass(frozen=True)
133
+ class CardAuth:
134
+ schemes: list[str]
135
+ protected_resource: bool
136
+
137
+ @classmethod
138
+ def from_dict(cls, d: Mapping[str, Any]) -> CardAuth:
139
+ return cls(
140
+ schemes=list(d.get("schemes") or []),
141
+ protected_resource=bool(d.get("protectedResource", False)),
142
+ )
143
+
144
+
145
+ @dataclass(frozen=True)
146
+ class CardVerifyResult:
147
+ """Outcome of preflighting an external MCP server / A2A agent."""
148
+
149
+ decision: CardDecision
150
+ reason: str
151
+ risk: int
152
+ # The wire carries ~20 protocol ids (a2a, mcp, oauth, did, x402, …); typed
153
+ # as a plain ``str`` so the SDK never lags the gateway's protocol list.
154
+ protocol: str
155
+ identity: CardIdentity
156
+ capabilities: CardCapabilities
157
+ auth: CardAuth
158
+ warnings: list[str]
159
+ card_id: str
160
+ signals: list[ToolCallSignal] | None = None
161
+
162
+ @classmethod
163
+ def from_dict(cls, d: Mapping[str, Any]) -> CardVerifyResult:
164
+ return cls(
165
+ decision=d.get("decision", "deny"),
166
+ reason=d.get("reason", ""),
167
+ risk=int(d.get("risk") or 0),
168
+ protocol=d.get("protocol", ""),
169
+ identity=CardIdentity.from_dict(d.get("identity") or {}),
170
+ capabilities=CardCapabilities.from_dict(d.get("capabilities") or {}),
171
+ auth=CardAuth.from_dict(d.get("auth") or {}),
172
+ warnings=list(d.get("warnings") or []),
173
+ card_id=d.get("cardId", ""),
174
+ signals=_parse_signals(d.get("signals")),
175
+ )
@@ -0,0 +1,6 @@
1
+ """Framework integrations for the AxioRank SDK.
2
+
3
+ Each integration lives in its own module and lazily imports the framework it
4
+ wraps, so installing ``axiorank`` never pulls in LangChain, OpenAI, etc. Install
5
+ only what you use, e.g. ``pip install axiorank[langchain]``.
6
+ """
@@ -0,0 +1,147 @@
1
+ """LangChain integration for AxioRank.
2
+
3
+ Two ways to put your agent's tools behind the gateway:
4
+
5
+ * :class:`AxioRankCallbackHandler` / :class:`AxioRankAsyncCallbackHandler` —
6
+ zero-touch. Attach the handler and *every* tool the agent runs is checked
7
+ first; a denied call raises and the step fails.
8
+ * :func:`guard_tool` — wrap a single tool. A denied call can either raise or
9
+ return a model-readable refusal (``on_deny="return"``) so the agent recovers.
10
+
11
+ Requires ``langchain-core`` (``pip install 'axiorank[langchain]'``).
12
+ """
13
+
14
+ from __future__ import annotations
15
+
16
+ import json
17
+ from typing import Any, Callable
18
+
19
+ try:
20
+ from langchain_core.callbacks import AsyncCallbackHandler, BaseCallbackHandler
21
+ from langchain_core.tools import BaseTool, StructuredTool
22
+ except ImportError as exc: # pragma: no cover - exercised only without the extra
23
+ raise ImportError(
24
+ "The AxioRank LangChain integration requires `langchain-core`. "
25
+ "Install it with: pip install 'axiorank[langchain]'"
26
+ ) from exc
27
+
28
+ from .._async import AsyncAxioRank
29
+ from .._errors import AxioRankDeniedError
30
+ from .._sync import AxioRank
31
+ from .._types import ToolCallResult
32
+
33
+
34
+ def _extract_call(serialized: Any, input_str: str, inputs: Any) -> tuple[str, dict[str, Any]]:
35
+ """Pull the tool name and an arguments dict out of a callback's payload."""
36
+ name = ""
37
+ if isinstance(serialized, dict):
38
+ raw_name = serialized.get("name")
39
+ if not raw_name:
40
+ ident = serialized.get("id")
41
+ if isinstance(ident, list) and ident:
42
+ raw_name = ident[-1]
43
+ name = str(raw_name or "")
44
+
45
+ if isinstance(inputs, dict):
46
+ return name, dict(inputs)
47
+
48
+ # Older LangChain passes only the stringified input — best-effort decode.
49
+ try:
50
+ parsed = json.loads(input_str)
51
+ except Exception:
52
+ return name, {"input": input_str}
53
+ return name, parsed if isinstance(parsed, dict) else {"input": parsed}
54
+
55
+
56
+ def _denied_message(name: str, result: ToolCallResult) -> str:
57
+ return f"AxioRank blocked the tool call `{name}`: {result.reason} (risk {result.risk})."
58
+
59
+
60
+ class AxioRankCallbackHandler(BaseCallbackHandler):
61
+ """Check every tool call against AxioRank before it runs (synchronous).
62
+
63
+ Attach to an agent run; a ``deny`` verdict raises
64
+ :class:`~axiorank.AxioRankDeniedError`, aborting the tool step.
65
+ """
66
+
67
+ # Propagate our exception instead of letting LangChain swallow it.
68
+ raise_error: bool = True
69
+
70
+ def __init__(self, client: AxioRank) -> None:
71
+ self._client = client
72
+
73
+ def on_tool_start(self, serialized: dict[str, Any], input_str: str, **kwargs: Any) -> None:
74
+ name, args = _extract_call(serialized, input_str, kwargs.get("inputs"))
75
+ result = self._client.tool_call(name, args)
76
+ if result.decision == "deny":
77
+ raise AxioRankDeniedError(result)
78
+
79
+
80
+ class AxioRankAsyncCallbackHandler(AsyncCallbackHandler):
81
+ """Async counterpart of :class:`AxioRankCallbackHandler`."""
82
+
83
+ raise_error: bool = True
84
+
85
+ def __init__(self, client: AsyncAxioRank) -> None:
86
+ self._client = client
87
+
88
+ async def on_tool_start(
89
+ self, serialized: dict[str, Any], input_str: str, **kwargs: Any
90
+ ) -> None:
91
+ name, args = _extract_call(serialized, input_str, kwargs.get("inputs"))
92
+ result = await self._client.tool_call(name, args)
93
+ if result.decision == "deny":
94
+ raise AxioRankDeniedError(result)
95
+
96
+
97
+ def guard_tool(
98
+ tool: BaseTool,
99
+ client: AxioRank | None = None,
100
+ *,
101
+ async_client: AsyncAxioRank | None = None,
102
+ on_deny: str = "raise",
103
+ ) -> StructuredTool:
104
+ """Wrap ``tool`` so its arguments are checked by AxioRank before it runs.
105
+
106
+ Provide ``client`` (sync), ``async_client`` (async), or both. On a ``deny``:
107
+ ``on_deny="raise"`` raises :class:`~axiorank.AxioRankDeniedError`;
108
+ ``on_deny="return"`` returns a short, model-readable refusal string instead.
109
+ """
110
+ if client is None and async_client is None:
111
+ raise ValueError("guard_tool: provide `client` and/or `async_client`")
112
+ if on_deny not in ("raise", "return"):
113
+ raise ValueError("guard_tool: `on_deny` must be 'raise' or 'return'")
114
+
115
+ name = tool.name
116
+ func: Callable[..., Any] | None = None
117
+ coroutine: Callable[..., Any] | None = None
118
+
119
+ if client is not None:
120
+ sync_client = client
121
+
122
+ def func(**kwargs: Any) -> Any:
123
+ result = sync_client.tool_call(name, kwargs)
124
+ if result.decision == "deny":
125
+ if on_deny == "return":
126
+ return _denied_message(name, result)
127
+ raise AxioRankDeniedError(result)
128
+ return tool.run(kwargs)
129
+
130
+ if async_client is not None:
131
+ a_client = async_client
132
+
133
+ async def coroutine(**kwargs: Any) -> Any:
134
+ result = await a_client.tool_call(name, kwargs)
135
+ if result.decision == "deny":
136
+ if on_deny == "return":
137
+ return _denied_message(name, result)
138
+ raise AxioRankDeniedError(result)
139
+ return await tool.arun(kwargs)
140
+
141
+ return StructuredTool.from_function(
142
+ func=func,
143
+ coroutine=coroutine,
144
+ name=name,
145
+ description=tool.description,
146
+ args_schema=tool.args_schema,
147
+ )
File without changes