korasafe-sdk 0.2.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.
- korasafe_sdk-0.2.0/.gitignore +60 -0
- korasafe_sdk-0.2.0/PKG-INFO +131 -0
- korasafe_sdk-0.2.0/README.md +102 -0
- korasafe_sdk-0.2.0/korasafe/__init__.py +37 -0
- korasafe_sdk-0.2.0/korasafe/client.py +192 -0
- korasafe_sdk-0.2.0/korasafe/decorators.py +71 -0
- korasafe_sdk-0.2.0/korasafe/langchain.py +39 -0
- korasafe_sdk-0.2.0/korasafe/llamaindex.py +46 -0
- korasafe_sdk-0.2.0/korasafe/models.py +107 -0
- korasafe_sdk-0.2.0/korasafe/py.typed +1 -0
- korasafe_sdk-0.2.0/korasafe/trace.py +455 -0
- korasafe_sdk-0.2.0/pyproject.toml +69 -0
- korasafe_sdk-0.2.0/tests/test_client.py +246 -0
- korasafe_sdk-0.2.0/tests/test_trace.py +443 -0
|
@@ -0,0 +1,60 @@
|
|
|
1
|
+
.vercel
|
|
2
|
+
.env
|
|
3
|
+
.env.*
|
|
4
|
+
!.env.example
|
|
5
|
+
node_modules
|
|
6
|
+
.DS_Store
|
|
7
|
+
*.log
|
|
8
|
+
*.docx
|
|
9
|
+
reference/plans/
|
|
10
|
+
mockups/
|
|
11
|
+
proto-*.html
|
|
12
|
+
__pycache__/
|
|
13
|
+
*.pyc
|
|
14
|
+
audit-runs/
|
|
15
|
+
.env*.local
|
|
16
|
+
*.vsix
|
|
17
|
+
screenshots-extracted/
|
|
18
|
+
.claude/worktrees/
|
|
19
|
+
.claude/scheduled_tasks.lock
|
|
20
|
+
.codex-memory.md
|
|
21
|
+
.claude2-memory.md
|
|
22
|
+
# local agent tooling (Codex CLI, Claude Code scratch, etc.)
|
|
23
|
+
.agents/
|
|
24
|
+
.~lock.*
|
|
25
|
+
*.tmp
|
|
26
|
+
KoraSafe_*.pdf
|
|
27
|
+
dist/
|
|
28
|
+
.vite/
|
|
29
|
+
public/_app/
|
|
30
|
+
screenshots/
|
|
31
|
+
playwright-report/
|
|
32
|
+
test-results/
|
|
33
|
+
.debug/
|
|
34
|
+
fonts/
|
|
35
|
+
korasafe-fonts.zip
|
|
36
|
+
ziwIH6Qf
|
|
37
|
+
tmp/
|
|
38
|
+
public/images/app/_archive-pre-v2/
|
|
39
|
+
|
|
40
|
+
# SDK build artifacts (Java, .NET, Go, Python)
|
|
41
|
+
target/
|
|
42
|
+
bin/
|
|
43
|
+
obj/
|
|
44
|
+
*.class
|
|
45
|
+
*.dll
|
|
46
|
+
*.exe
|
|
47
|
+
*.so
|
|
48
|
+
*.dylib
|
|
49
|
+
*.jar
|
|
50
|
+
*.nupkg
|
|
51
|
+
*.whl
|
|
52
|
+
go.sum.backup
|
|
53
|
+
.pytest_cache/
|
|
54
|
+
.mypy_cache/
|
|
55
|
+
.coverage
|
|
56
|
+
.claude/launch.json
|
|
57
|
+
|
|
58
|
+
# Build-generated API route registry (#4535). Generated by
|
|
59
|
+
# scripts/generate-hono-routes.mjs via postinstall/prebuild/pretest.
|
|
60
|
+
src/api/_routes.js
|
|
@@ -0,0 +1,131 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: korasafe-sdk
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: KoraSafe Python SDK — inline guardian scans + chain-of-thought trace capture for Python agents
|
|
5
|
+
Project-URL: Homepage, https://korasafe.app
|
|
6
|
+
Project-URL: Repository, https://github.com/korasafe/platform
|
|
7
|
+
Project-URL: Documentation, https://github.com/korasafe/platform/tree/main/docs/sdk/python.md
|
|
8
|
+
Author: KoraSafe
|
|
9
|
+
License: MIT
|
|
10
|
+
Keywords: agents,ai-governance,guardrails,korasafe,langchain,llamaindex
|
|
11
|
+
Classifier: Programming Language :: Python :: 3
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
14
|
+
Classifier: Typing :: Typed
|
|
15
|
+
Requires-Python: >=3.11
|
|
16
|
+
Requires-Dist: httpx<1,>=0.27
|
|
17
|
+
Requires-Dist: pydantic<3,>=2.7
|
|
18
|
+
Provides-Extra: dev
|
|
19
|
+
Requires-Dist: build>=1.2; extra == 'dev'
|
|
20
|
+
Requires-Dist: coverage[toml]>=7.5; extra == 'dev'
|
|
21
|
+
Requires-Dist: mypy>=1.10; extra == 'dev'
|
|
22
|
+
Requires-Dist: pytest>=8.2; extra == 'dev'
|
|
23
|
+
Requires-Dist: ruff>=0.5; extra == 'dev'
|
|
24
|
+
Provides-Extra: langchain
|
|
25
|
+
Requires-Dist: langchain-core>=0.2; extra == 'langchain'
|
|
26
|
+
Provides-Extra: llamaindex
|
|
27
|
+
Requires-Dist: llama-index-core>=0.10; extra == 'llamaindex'
|
|
28
|
+
Description-Content-Type: text/markdown
|
|
29
|
+
|
|
30
|
+
# KoraSafe Python SDK
|
|
31
|
+
|
|
32
|
+
Two surfaces in one package:
|
|
33
|
+
|
|
34
|
+
- `KoraSafeClient` — inline guardian inspection (scan, gate, submit findings). Metadata-only transport.
|
|
35
|
+
- `kora_trace` — capture agent chain-of-thought (plan, LLM calls, tool calls, reasoning, human approvals) and ship to the KoraSafe audit log.
|
|
36
|
+
|
|
37
|
+
```bash
|
|
38
|
+
pip install korasafe-sdk
|
|
39
|
+
export KORASAFE_API_KEY=ks_live_...
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
## kora_trace — chain-of-thought capture
|
|
43
|
+
|
|
44
|
+
10-line FastAPI example:
|
|
45
|
+
|
|
46
|
+
```python
|
|
47
|
+
from fastapi import FastAPI
|
|
48
|
+
from korasafe import init_trace, kora_trace
|
|
49
|
+
|
|
50
|
+
init_trace()
|
|
51
|
+
app = FastAPI()
|
|
52
|
+
|
|
53
|
+
@app.post("/classify")
|
|
54
|
+
def classify(text: str) -> dict[str, str]:
|
|
55
|
+
with kora_trace.run("classify_claim"):
|
|
56
|
+
kora_trace.plan(["look up policy", "score risk", "route"])
|
|
57
|
+
kora_trace.llm_call(provider="openai", model="gpt-4o", input=text, output="tier=gold", input_tokens=120, output_tokens=30)
|
|
58
|
+
kora_trace.tool_call(name="policy_lookup", parameters={"id": "pol-7"}, response={"tier": "gold"}, status="ok")
|
|
59
|
+
kora_trace.human_approval(reviewer_id="user-42", decision="approved")
|
|
60
|
+
return {"tier": "gold"}
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
Events appear in the KoraSafe audit log within 5 seconds.
|
|
64
|
+
|
|
65
|
+
### API
|
|
66
|
+
|
|
67
|
+
| Method | Purpose |
|
|
68
|
+
|---|---|
|
|
69
|
+
| `kora_trace.run(task_name)` | Context manager. Opens a trace; `run_start` + `run_end` events bracket the block. Nested method calls auto-attach via contextvars. |
|
|
70
|
+
| `@kora_trace.trace(task_name)` | Decorator equivalent of `run`. Works on sync and async functions. |
|
|
71
|
+
| `kora_trace.plan(steps, reasoning=None)` | Log initial plan. `steps` can be strings or `{step, description}` dicts. |
|
|
72
|
+
| `kora_trace.llm_call(provider, model, input, output, input_tokens=, output_tokens=, total_tokens=, cost_usd=, duration_ms=, status=, error=)` | Log an LLM invocation. Tokens auto-sum if `total_tokens` omitted. |
|
|
73
|
+
| `kora_trace.tool_call(name, parameters=, response=, status=, duration_ms=, error=)` | Log a tool or external API call. |
|
|
74
|
+
| `kora_trace.human_approval(reviewer_id, decision, notes=, approval_chain=)` | Log a HITL decision. |
|
|
75
|
+
| `kora_trace.flush()` | Force-flush buffered events. |
|
|
76
|
+
| `kora_trace.close()` | Drain buffer + close HTTP client. |
|
|
77
|
+
|
|
78
|
+
`init_trace(**kwargs)` reinitializes the singleton with overrides (`api_key`, `endpoint`, `batch_size=10`, `flush_interval_s=5`, `timeout_s=10`, `max_retries=3`, `disabled=False`, `logger=`, `http_client=`).
|
|
79
|
+
|
|
80
|
+
Calling `plan` / `llm_call` / `tool_call` / `human_approval` outside a `run()` context raises `RuntimeError` — wrap your agent loop first.
|
|
81
|
+
|
|
82
|
+
## KoraSafeClient — guardian inspection
|
|
83
|
+
|
|
84
|
+
```python
|
|
85
|
+
from korasafe import KoraSafeClient, withKoraSafeScan
|
|
86
|
+
|
|
87
|
+
client = KoraSafeClient()
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
@withKoraSafeScan(client=client, context={"system_id": "claims-agent"})
|
|
91
|
+
def answer_claim(prompt: str) -> str:
|
|
92
|
+
return "approved"
|
|
93
|
+
|
|
94
|
+
result = client.scan("Does this contain PII?", {"system_id": "claims-agent"})
|
|
95
|
+
gate = client.gate({"action": "payment_approval", "risk_tier": "high"})
|
|
96
|
+
finding = client.submit_finding({"guardian_id": "pii", "title": "PII found", "severity": "high"})
|
|
97
|
+
```
|
|
98
|
+
|
|
99
|
+
Raw strings passed to `scan()` or the decorator are hashed locally. The SDK sends content hash, byte length, direction, surface, labels, and caller metadata rather than prompt or response bodies.
|
|
100
|
+
|
|
101
|
+
## Frameworks
|
|
102
|
+
|
|
103
|
+
LangChain:
|
|
104
|
+
|
|
105
|
+
```python
|
|
106
|
+
from korasafe import KoraSafeCallback, KoraSafeClient
|
|
107
|
+
|
|
108
|
+
callbacks = [KoraSafeCallback(client=KoraSafeClient(), context={"system_id": "claims-agent"})]
|
|
109
|
+
```
|
|
110
|
+
|
|
111
|
+
LlamaIndex:
|
|
112
|
+
|
|
113
|
+
```python
|
|
114
|
+
from korasafe import KoraSafeClient, KoraSafeLlamaIndexMiddleware
|
|
115
|
+
|
|
116
|
+
query_engine = KoraSafeLlamaIndexMiddleware(KoraSafeClient()).wrap_query_engine(query_engine)
|
|
117
|
+
```
|
|
118
|
+
|
|
119
|
+
## Development
|
|
120
|
+
|
|
121
|
+
```bash
|
|
122
|
+
cd packages/sdk-python
|
|
123
|
+
python -m pip install -e ".[dev]"
|
|
124
|
+
ruff check .
|
|
125
|
+
mypy .
|
|
126
|
+
coverage run -m pytest
|
|
127
|
+
coverage report
|
|
128
|
+
python -m build
|
|
129
|
+
```
|
|
130
|
+
|
|
131
|
+
Publishing uses GitHub OIDC trusted publishing to PyPI from `sdk-python-v*` tags.
|
|
@@ -0,0 +1,102 @@
|
|
|
1
|
+
# KoraSafe Python SDK
|
|
2
|
+
|
|
3
|
+
Two surfaces in one package:
|
|
4
|
+
|
|
5
|
+
- `KoraSafeClient` — inline guardian inspection (scan, gate, submit findings). Metadata-only transport.
|
|
6
|
+
- `kora_trace` — capture agent chain-of-thought (plan, LLM calls, tool calls, reasoning, human approvals) and ship to the KoraSafe audit log.
|
|
7
|
+
|
|
8
|
+
```bash
|
|
9
|
+
pip install korasafe-sdk
|
|
10
|
+
export KORASAFE_API_KEY=ks_live_...
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
## kora_trace — chain-of-thought capture
|
|
14
|
+
|
|
15
|
+
10-line FastAPI example:
|
|
16
|
+
|
|
17
|
+
```python
|
|
18
|
+
from fastapi import FastAPI
|
|
19
|
+
from korasafe import init_trace, kora_trace
|
|
20
|
+
|
|
21
|
+
init_trace()
|
|
22
|
+
app = FastAPI()
|
|
23
|
+
|
|
24
|
+
@app.post("/classify")
|
|
25
|
+
def classify(text: str) -> dict[str, str]:
|
|
26
|
+
with kora_trace.run("classify_claim"):
|
|
27
|
+
kora_trace.plan(["look up policy", "score risk", "route"])
|
|
28
|
+
kora_trace.llm_call(provider="openai", model="gpt-4o", input=text, output="tier=gold", input_tokens=120, output_tokens=30)
|
|
29
|
+
kora_trace.tool_call(name="policy_lookup", parameters={"id": "pol-7"}, response={"tier": "gold"}, status="ok")
|
|
30
|
+
kora_trace.human_approval(reviewer_id="user-42", decision="approved")
|
|
31
|
+
return {"tier": "gold"}
|
|
32
|
+
```
|
|
33
|
+
|
|
34
|
+
Events appear in the KoraSafe audit log within 5 seconds.
|
|
35
|
+
|
|
36
|
+
### API
|
|
37
|
+
|
|
38
|
+
| Method | Purpose |
|
|
39
|
+
|---|---|
|
|
40
|
+
| `kora_trace.run(task_name)` | Context manager. Opens a trace; `run_start` + `run_end` events bracket the block. Nested method calls auto-attach via contextvars. |
|
|
41
|
+
| `@kora_trace.trace(task_name)` | Decorator equivalent of `run`. Works on sync and async functions. |
|
|
42
|
+
| `kora_trace.plan(steps, reasoning=None)` | Log initial plan. `steps` can be strings or `{step, description}` dicts. |
|
|
43
|
+
| `kora_trace.llm_call(provider, model, input, output, input_tokens=, output_tokens=, total_tokens=, cost_usd=, duration_ms=, status=, error=)` | Log an LLM invocation. Tokens auto-sum if `total_tokens` omitted. |
|
|
44
|
+
| `kora_trace.tool_call(name, parameters=, response=, status=, duration_ms=, error=)` | Log a tool or external API call. |
|
|
45
|
+
| `kora_trace.human_approval(reviewer_id, decision, notes=, approval_chain=)` | Log a HITL decision. |
|
|
46
|
+
| `kora_trace.flush()` | Force-flush buffered events. |
|
|
47
|
+
| `kora_trace.close()` | Drain buffer + close HTTP client. |
|
|
48
|
+
|
|
49
|
+
`init_trace(**kwargs)` reinitializes the singleton with overrides (`api_key`, `endpoint`, `batch_size=10`, `flush_interval_s=5`, `timeout_s=10`, `max_retries=3`, `disabled=False`, `logger=`, `http_client=`).
|
|
50
|
+
|
|
51
|
+
Calling `plan` / `llm_call` / `tool_call` / `human_approval` outside a `run()` context raises `RuntimeError` — wrap your agent loop first.
|
|
52
|
+
|
|
53
|
+
## KoraSafeClient — guardian inspection
|
|
54
|
+
|
|
55
|
+
```python
|
|
56
|
+
from korasafe import KoraSafeClient, withKoraSafeScan
|
|
57
|
+
|
|
58
|
+
client = KoraSafeClient()
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
@withKoraSafeScan(client=client, context={"system_id": "claims-agent"})
|
|
62
|
+
def answer_claim(prompt: str) -> str:
|
|
63
|
+
return "approved"
|
|
64
|
+
|
|
65
|
+
result = client.scan("Does this contain PII?", {"system_id": "claims-agent"})
|
|
66
|
+
gate = client.gate({"action": "payment_approval", "risk_tier": "high"})
|
|
67
|
+
finding = client.submit_finding({"guardian_id": "pii", "title": "PII found", "severity": "high"})
|
|
68
|
+
```
|
|
69
|
+
|
|
70
|
+
Raw strings passed to `scan()` or the decorator are hashed locally. The SDK sends content hash, byte length, direction, surface, labels, and caller metadata rather than prompt or response bodies.
|
|
71
|
+
|
|
72
|
+
## Frameworks
|
|
73
|
+
|
|
74
|
+
LangChain:
|
|
75
|
+
|
|
76
|
+
```python
|
|
77
|
+
from korasafe import KoraSafeCallback, KoraSafeClient
|
|
78
|
+
|
|
79
|
+
callbacks = [KoraSafeCallback(client=KoraSafeClient(), context={"system_id": "claims-agent"})]
|
|
80
|
+
```
|
|
81
|
+
|
|
82
|
+
LlamaIndex:
|
|
83
|
+
|
|
84
|
+
```python
|
|
85
|
+
from korasafe import KoraSafeClient, KoraSafeLlamaIndexMiddleware
|
|
86
|
+
|
|
87
|
+
query_engine = KoraSafeLlamaIndexMiddleware(KoraSafeClient()).wrap_query_engine(query_engine)
|
|
88
|
+
```
|
|
89
|
+
|
|
90
|
+
## Development
|
|
91
|
+
|
|
92
|
+
```bash
|
|
93
|
+
cd packages/sdk-python
|
|
94
|
+
python -m pip install -e ".[dev]"
|
|
95
|
+
ruff check .
|
|
96
|
+
mypy .
|
|
97
|
+
coverage run -m pytest
|
|
98
|
+
coverage report
|
|
99
|
+
python -m build
|
|
100
|
+
```
|
|
101
|
+
|
|
102
|
+
Publishing uses GitHub OIDC trusted publishing to PyPI from `sdk-python-v*` tags.
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
from .client import KoraSafeAPIError, KoraSafeClient
|
|
2
|
+
from .decorators import KoraSafeBlockedError, withKoraSafeScan
|
|
3
|
+
from .langchain import KoraSafeCallback
|
|
4
|
+
from .llamaindex import KoraSafeLlamaIndexMiddleware
|
|
5
|
+
from .models import (
|
|
6
|
+
FindingSubmission,
|
|
7
|
+
GateDecision,
|
|
8
|
+
GateResult,
|
|
9
|
+
GuardianFinding,
|
|
10
|
+
ScanContext,
|
|
11
|
+
ScanInput,
|
|
12
|
+
ScanResult,
|
|
13
|
+
SubmittedFinding,
|
|
14
|
+
)
|
|
15
|
+
from .trace import KoraTrace, init_trace, kora_trace
|
|
16
|
+
|
|
17
|
+
__version__ = "0.2.0"
|
|
18
|
+
|
|
19
|
+
__all__ = [
|
|
20
|
+
"FindingSubmission",
|
|
21
|
+
"GateDecision",
|
|
22
|
+
"GateResult",
|
|
23
|
+
"GuardianFinding",
|
|
24
|
+
"KoraSafeAPIError",
|
|
25
|
+
"KoraSafeBlockedError",
|
|
26
|
+
"KoraSafeCallback",
|
|
27
|
+
"KoraSafeClient",
|
|
28
|
+
"KoraSafeLlamaIndexMiddleware",
|
|
29
|
+
"KoraTrace",
|
|
30
|
+
"ScanContext",
|
|
31
|
+
"ScanInput",
|
|
32
|
+
"ScanResult",
|
|
33
|
+
"SubmittedFinding",
|
|
34
|
+
"init_trace",
|
|
35
|
+
"kora_trace",
|
|
36
|
+
"withKoraSafeScan",
|
|
37
|
+
]
|
|
@@ -0,0 +1,192 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
from typing import Any, cast
|
|
5
|
+
|
|
6
|
+
import httpx
|
|
7
|
+
|
|
8
|
+
from .models import (
|
|
9
|
+
FindingSubmission,
|
|
10
|
+
GateDecision,
|
|
11
|
+
GateResult,
|
|
12
|
+
ScanContext,
|
|
13
|
+
ScanInput,
|
|
14
|
+
ScanResult,
|
|
15
|
+
SubmittedFinding,
|
|
16
|
+
)
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class KoraSafeAPIError(RuntimeError):
|
|
20
|
+
def __init__(self, message: str, *, status_code: int = 0, code: str = "UNKNOWN") -> None:
|
|
21
|
+
super().__init__(message)
|
|
22
|
+
self.status_code = status_code
|
|
23
|
+
self.code = code
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class KoraSafeClient:
|
|
27
|
+
def __init__(
|
|
28
|
+
self,
|
|
29
|
+
api_key: str | None = None,
|
|
30
|
+
*,
|
|
31
|
+
base_url: str | None = None,
|
|
32
|
+
timeout: float = 30.0,
|
|
33
|
+
http_client: httpx.Client | None = None,
|
|
34
|
+
async_http_client: httpx.AsyncClient | None = None,
|
|
35
|
+
) -> None:
|
|
36
|
+
resolved_key = api_key or os.getenv("KORASAFE_API_KEY")
|
|
37
|
+
if not resolved_key:
|
|
38
|
+
raise ValueError("KoraSafe API key required; pass api_key or set KORASAFE_API_KEY")
|
|
39
|
+
self.api_key = resolved_key
|
|
40
|
+
self.base_url = (
|
|
41
|
+
base_url or os.getenv("KORASAFE_BASE_URL") or "https://korasafe.app"
|
|
42
|
+
).rstrip("/")
|
|
43
|
+
self.timeout = timeout
|
|
44
|
+
self._client = http_client or httpx.Client(timeout=timeout)
|
|
45
|
+
self._async_client = async_http_client or httpx.AsyncClient(timeout=timeout)
|
|
46
|
+
|
|
47
|
+
def scan(
|
|
48
|
+
self,
|
|
49
|
+
input: str | bytes | dict[str, Any] | ScanInput,
|
|
50
|
+
context: dict[str, Any] | ScanContext | None = None,
|
|
51
|
+
) -> ScanResult:
|
|
52
|
+
payload = self._scan_payload(input, context)
|
|
53
|
+
return ScanResult.model_validate(self._request("POST", "/api/guardian-scan", json=payload))
|
|
54
|
+
|
|
55
|
+
async def async_scan(
|
|
56
|
+
self,
|
|
57
|
+
input: str | bytes | dict[str, Any] | ScanInput,
|
|
58
|
+
context: dict[str, Any] | ScanContext | None = None,
|
|
59
|
+
) -> ScanResult:
|
|
60
|
+
payload = self._scan_payload(input, context)
|
|
61
|
+
return ScanResult.model_validate(
|
|
62
|
+
await self._async_request("POST", "/api/guardian-scan", json=payload)
|
|
63
|
+
)
|
|
64
|
+
|
|
65
|
+
def gate(
|
|
66
|
+
self,
|
|
67
|
+
decision: dict[str, Any] | GateDecision,
|
|
68
|
+
action: str | None = None,
|
|
69
|
+
) -> GateResult:
|
|
70
|
+
payload = self._gate_payload(decision, action)
|
|
71
|
+
return GateResult.model_validate(self._request("POST", "/api/guardian-gate", json=payload))
|
|
72
|
+
|
|
73
|
+
async def async_gate(
|
|
74
|
+
self,
|
|
75
|
+
decision: dict[str, Any] | GateDecision,
|
|
76
|
+
action: str | None = None,
|
|
77
|
+
) -> GateResult:
|
|
78
|
+
payload = self._gate_payload(decision, action)
|
|
79
|
+
return GateResult.model_validate(
|
|
80
|
+
await self._async_request("POST", "/api/guardian-gate", json=payload)
|
|
81
|
+
)
|
|
82
|
+
|
|
83
|
+
def submit_finding(
|
|
84
|
+
self,
|
|
85
|
+
finding: dict[str, Any] | FindingSubmission,
|
|
86
|
+
**kwargs: Any,
|
|
87
|
+
) -> SubmittedFinding:
|
|
88
|
+
payload = self._finding_payload(finding, kwargs)
|
|
89
|
+
return SubmittedFinding.model_validate(self._request("POST", "/api/findings", json=payload))
|
|
90
|
+
|
|
91
|
+
async def async_submit_finding(
|
|
92
|
+
self,
|
|
93
|
+
finding: dict[str, Any] | FindingSubmission,
|
|
94
|
+
**kwargs: Any,
|
|
95
|
+
) -> SubmittedFinding:
|
|
96
|
+
payload = self._finding_payload(finding, kwargs)
|
|
97
|
+
return SubmittedFinding.model_validate(
|
|
98
|
+
await self._async_request("POST", "/api/findings", json=payload)
|
|
99
|
+
)
|
|
100
|
+
|
|
101
|
+
def close(self) -> None:
|
|
102
|
+
self._client.close()
|
|
103
|
+
|
|
104
|
+
async def aclose(self) -> None:
|
|
105
|
+
await self._async_client.aclose()
|
|
106
|
+
|
|
107
|
+
def _scan_payload(
|
|
108
|
+
self,
|
|
109
|
+
input: str | bytes | dict[str, Any] | ScanInput,
|
|
110
|
+
context: dict[str, Any] | ScanContext | None,
|
|
111
|
+
) -> dict[str, Any]:
|
|
112
|
+
scan_input = ScanInput.from_value(input)
|
|
113
|
+
scan_context = context if isinstance(context, ScanContext) else ScanContext.model_validate(
|
|
114
|
+
context or {}
|
|
115
|
+
)
|
|
116
|
+
return {
|
|
117
|
+
"input": scan_input.model_dump(mode="json"),
|
|
118
|
+
"context": scan_context.model_dump(mode="json"),
|
|
119
|
+
}
|
|
120
|
+
|
|
121
|
+
def _gate_payload(
|
|
122
|
+
self,
|
|
123
|
+
decision: dict[str, Any] | GateDecision,
|
|
124
|
+
action: str | None,
|
|
125
|
+
) -> dict[str, Any]:
|
|
126
|
+
if isinstance(decision, GateDecision):
|
|
127
|
+
gate_decision = decision
|
|
128
|
+
else:
|
|
129
|
+
payload = dict(decision)
|
|
130
|
+
if action is not None:
|
|
131
|
+
payload.setdefault("action", action)
|
|
132
|
+
gate_decision = GateDecision.model_validate(payload)
|
|
133
|
+
return gate_decision.model_dump(mode="json")
|
|
134
|
+
|
|
135
|
+
def _finding_payload(
|
|
136
|
+
self,
|
|
137
|
+
finding: dict[str, Any] | FindingSubmission,
|
|
138
|
+
kwargs: dict[str, Any],
|
|
139
|
+
) -> dict[str, Any]:
|
|
140
|
+
payload = (
|
|
141
|
+
finding
|
|
142
|
+
if isinstance(finding, FindingSubmission)
|
|
143
|
+
else FindingSubmission.model_validate({**finding, **kwargs})
|
|
144
|
+
)
|
|
145
|
+
return payload.model_dump(mode="json")
|
|
146
|
+
|
|
147
|
+
def _headers(self) -> dict[str, str]:
|
|
148
|
+
return {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
|
|
149
|
+
|
|
150
|
+
def _request(self, method: str, path: str, *, json: dict[str, Any]) -> dict[str, Any]:
|
|
151
|
+
response = self._client.request(
|
|
152
|
+
method,
|
|
153
|
+
f"{self.base_url}{path}",
|
|
154
|
+
headers=self._headers(),
|
|
155
|
+
json=json,
|
|
156
|
+
)
|
|
157
|
+
return self._decode(response)
|
|
158
|
+
|
|
159
|
+
async def _async_request(
|
|
160
|
+
self,
|
|
161
|
+
method: str,
|
|
162
|
+
path: str,
|
|
163
|
+
*,
|
|
164
|
+
json: dict[str, Any],
|
|
165
|
+
) -> dict[str, Any]:
|
|
166
|
+
response = await self._async_client.request(
|
|
167
|
+
method,
|
|
168
|
+
f"{self.base_url}{path}",
|
|
169
|
+
headers=self._headers(),
|
|
170
|
+
json=json,
|
|
171
|
+
)
|
|
172
|
+
return self._decode(response)
|
|
173
|
+
|
|
174
|
+
def _decode(self, response: httpx.Response) -> dict[str, Any]:
|
|
175
|
+
try:
|
|
176
|
+
data = response.json()
|
|
177
|
+
except ValueError as exc:
|
|
178
|
+
raise KoraSafeAPIError(
|
|
179
|
+
f"Non-JSON response from KoraSafe: HTTP {response.status_code}",
|
|
180
|
+
status_code=response.status_code,
|
|
181
|
+
) from exc
|
|
182
|
+
if response.is_error:
|
|
183
|
+
error = data.get("error", {}) if isinstance(data, dict) else {}
|
|
184
|
+
message = error.get("message") or f"HTTP {response.status_code}"
|
|
185
|
+
code = error.get("code") or "UNKNOWN"
|
|
186
|
+
raise KoraSafeAPIError(message, status_code=response.status_code, code=code)
|
|
187
|
+
if not isinstance(data, dict):
|
|
188
|
+
raise KoraSafeAPIError(
|
|
189
|
+
"KoraSafe response must be a JSON object",
|
|
190
|
+
status_code=response.status_code,
|
|
191
|
+
)
|
|
192
|
+
return cast(dict[str, Any], data)
|
|
@@ -0,0 +1,71 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import functools
|
|
4
|
+
import inspect
|
|
5
|
+
from collections.abc import Callable
|
|
6
|
+
from typing import Any, ParamSpec, TypeVar, cast, overload
|
|
7
|
+
|
|
8
|
+
from .client import KoraSafeClient
|
|
9
|
+
from .models import ScanContext
|
|
10
|
+
|
|
11
|
+
P = ParamSpec("P")
|
|
12
|
+
R = TypeVar("R")
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class KoraSafeBlockedError(RuntimeError):
|
|
16
|
+
pass
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
@overload
|
|
20
|
+
def withKoraSafeScan(handler: Callable[P, R]) -> Callable[P, R]: ...
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@overload
|
|
24
|
+
def withKoraSafeScan(
|
|
25
|
+
handler: None = None,
|
|
26
|
+
*,
|
|
27
|
+
client: KoraSafeClient | None = None,
|
|
28
|
+
context: ScanContext | dict[str, Any] | None = None,
|
|
29
|
+
) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def withKoraSafeScan(
|
|
33
|
+
handler: Callable[P, R] | None = None,
|
|
34
|
+
*,
|
|
35
|
+
client: KoraSafeClient | None = None,
|
|
36
|
+
context: ScanContext | dict[str, Any] | None = None,
|
|
37
|
+
) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
|
|
38
|
+
def decorate(func: Callable[P, R]) -> Callable[P, R]:
|
|
39
|
+
scanner = client or KoraSafeClient()
|
|
40
|
+
if inspect.iscoroutinefunction(func):
|
|
41
|
+
|
|
42
|
+
@functools.wraps(func)
|
|
43
|
+
async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Any:
|
|
44
|
+
result = await scanner.async_scan(_input_summary(args, kwargs), context)
|
|
45
|
+
if not result.allowed:
|
|
46
|
+
raise KoraSafeBlockedError(result.action or "KoraSafe blocked this call")
|
|
47
|
+
return await func(*args, **kwargs)
|
|
48
|
+
|
|
49
|
+
return cast(Callable[P, R], async_wrapper)
|
|
50
|
+
|
|
51
|
+
@functools.wraps(func)
|
|
52
|
+
def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
|
|
53
|
+
result = scanner.scan(_input_summary(args, kwargs), context)
|
|
54
|
+
if not result.allowed:
|
|
55
|
+
raise KoraSafeBlockedError(result.action or "KoraSafe blocked this call")
|
|
56
|
+
return func(*args, **kwargs)
|
|
57
|
+
|
|
58
|
+
return wrapper
|
|
59
|
+
|
|
60
|
+
if handler is not None:
|
|
61
|
+
return decorate(handler)
|
|
62
|
+
return decorate
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
def _input_summary(args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]:
|
|
66
|
+
text = repr({"args": args, "kwargs": kwargs})
|
|
67
|
+
return {
|
|
68
|
+
"content": text,
|
|
69
|
+
"surface": "python-sdk",
|
|
70
|
+
"metadata": {"arg_count": len(args), "kwarg_keys": sorted(kwargs)},
|
|
71
|
+
}
|
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from typing import Any
|
|
4
|
+
|
|
5
|
+
from .client import KoraSafeClient
|
|
6
|
+
|
|
7
|
+
try:
|
|
8
|
+
from langchain_core.callbacks import BaseCallbackHandler # type: ignore[import-not-found]
|
|
9
|
+
except Exception: # pragma: no cover - optional dependency
|
|
10
|
+
BaseCallbackHandler = object
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class KoraSafeCallback(BaseCallbackHandler): # type: ignore[misc]
|
|
14
|
+
def __init__(
|
|
15
|
+
self,
|
|
16
|
+
client: KoraSafeClient | None = None,
|
|
17
|
+
context: dict[str, Any] | None = None,
|
|
18
|
+
) -> None:
|
|
19
|
+
super().__init__()
|
|
20
|
+
self.client = client or KoraSafeClient()
|
|
21
|
+
self.context = context or {}
|
|
22
|
+
|
|
23
|
+
def on_llm_start(self, serialized: dict[str, Any], prompts: list[str], **kwargs: Any) -> None:
|
|
24
|
+
for prompt in prompts:
|
|
25
|
+
self.client.scan(
|
|
26
|
+
{
|
|
27
|
+
"content": prompt,
|
|
28
|
+
"direction": "prompt",
|
|
29
|
+
"surface": "langchain",
|
|
30
|
+
"metadata": {"serialized": serialized},
|
|
31
|
+
},
|
|
32
|
+
{**self.context, "metadata": {"run_id": str(kwargs.get("run_id", ""))}},
|
|
33
|
+
)
|
|
34
|
+
|
|
35
|
+
def on_llm_end(self, response: Any, **kwargs: Any) -> None:
|
|
36
|
+
self.client.scan(
|
|
37
|
+
{"content": repr(response), "direction": "response", "surface": "langchain"},
|
|
38
|
+
{**self.context, "metadata": {"run_id": str(kwargs.get("run_id", ""))}},
|
|
39
|
+
)
|
|
@@ -0,0 +1,46 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from collections.abc import Callable
|
|
4
|
+
from typing import Any, TypeVar
|
|
5
|
+
|
|
6
|
+
from .client import KoraSafeClient
|
|
7
|
+
|
|
8
|
+
R = TypeVar("R")
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class KoraSafeLlamaIndexMiddleware:
|
|
12
|
+
def __init__(
|
|
13
|
+
self,
|
|
14
|
+
client: KoraSafeClient | None = None,
|
|
15
|
+
context: dict[str, Any] | None = None,
|
|
16
|
+
) -> None:
|
|
17
|
+
self.client = client or KoraSafeClient()
|
|
18
|
+
self.context = context or {}
|
|
19
|
+
|
|
20
|
+
def wrap_query_engine(self, query_engine: Any) -> Any:
|
|
21
|
+
original = query_engine.query
|
|
22
|
+
|
|
23
|
+
def query(prompt: str, *args: Any, **kwargs: Any) -> Any:
|
|
24
|
+
self.client.scan(
|
|
25
|
+
{"content": prompt, "surface": "llamaindex", "direction": "prompt"},
|
|
26
|
+
self.context,
|
|
27
|
+
)
|
|
28
|
+
result = original(prompt, *args, **kwargs)
|
|
29
|
+
self.client.scan(
|
|
30
|
+
{"content": repr(result), "surface": "llamaindex", "direction": "response"},
|
|
31
|
+
self.context,
|
|
32
|
+
)
|
|
33
|
+
return result
|
|
34
|
+
|
|
35
|
+
query_engine.query = query
|
|
36
|
+
return query_engine
|
|
37
|
+
|
|
38
|
+
def wrap_callable(self, handler: Callable[..., R]) -> Callable[..., R]:
|
|
39
|
+
def wrapped(*args: Any, **kwargs: Any) -> R:
|
|
40
|
+
self.client.scan(
|
|
41
|
+
{"content": repr({"args": args, "kwargs": kwargs}), "surface": "llamaindex"},
|
|
42
|
+
self.context,
|
|
43
|
+
)
|
|
44
|
+
return handler(*args, **kwargs)
|
|
45
|
+
|
|
46
|
+
return wrapped
|