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.
@@ -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