diff --git a/README.md b/README.md index e90e307..f2efb47 100644 --- a/README.md +++ b/README.md @@ -66,3 +66,13 @@ Commit with conventional messages (`fix: ...` is a patch, `feat: ...` a minor ve pip install -e ".[dev]" pytest tests ``` + +## Privacy + +With `capture=True`, personal data is scrubbed on your machine before anything is sent: +emails, payment cards, phone numbers and account numbers become keyed tokens +(``), the same value always giving the same token, so AgentCompile can still +match values across a conversation without seeing them. The key stays with you: +`AGENTCOMPILE_SCRUB_KEY` (set the same one on all your servers), else one created once in +`~/.agentcompile/scrub.key`. `scrub=False` turns it off. Live decisions (`/v1/decide`) need +real values to act on a customer's request; they are used in memory and stored scrubbed. diff --git a/src/agentcompile/__init__.py b/src/agentcompile/__init__.py index 35658e8..79d8c0e 100644 --- a/src/agentcompile/__init__.py +++ b/src/agentcompile/__init__.py @@ -19,6 +19,7 @@ from ._conversation import conversation from ._core import LIVE, SHADOW, Proxy, Router, async_create, sync_create from ._decide import AsyncDecider, Decider, Settings +from ._scrub import load_key from ._trail import OnEvent, Trail __all__ = ["LIVE", "SHADOW", "Settings", "conversation", "flush", "wrap"] @@ -42,6 +43,7 @@ def wrap( company: str | None = None, http_client: Any = None, capture: bool | None = None, + scrub: bool = True, ) -> Any: """Wrap an OpenAI- or Anthropic-style client; returns an object you use exactly like it. @@ -51,6 +53,9 @@ def wrap( trail: a path, True for ~/.agentcompile/trail.jsonl, or False. on_event: called per call. capture: send each call's request and answer to AgentCompile in the background, so it can find the jobs your agent repeats (opt-in; or AGENTCOMPILE_CAPTURE=1). Never slows a call. + scrub: with capture, turn emails, cards, phones and account numbers into keyed tokens on + this machine before anything is sent (default True; the key is AGENTCOMPILE_SCRUB_KEY or + one created once in ~/.agentcompile/). """ settings = Settings( base_url=base_url or os.environ.get("AGENTCOMPILE_URL", DEFAULT_URL), @@ -63,7 +68,11 @@ def wrap( capture = os.environ.get("AGENTCOMPILE_CAPTURE", "") in ("1", "true", "yes") # One sender per wrapped client; it always sends from its own thread (sync client). capturer = ( - Capturer(settings, http_client if _is_sync_http(http_client) else None) + Capturer( + settings, + http_client if _is_sync_http(http_client) else None, + load_key() if scrub else None, + ) if capture else None ) diff --git a/src/agentcompile/_capture.py b/src/agentcompile/_capture.py index 3030b89..182624c 100644 --- a/src/agentcompile/_capture.py +++ b/src/agentcompile/_capture.py @@ -20,6 +20,7 @@ from ._decide import CONVERSATION_HEADER, Settings, _request from ._payload import jsonable, payload +from ._scrub import scrub_call MAX_QUEUE = 2000 BATCH = 50 @@ -35,8 +36,15 @@ def flush_all(timeout: float = 5.0) -> None: class Capturer: - def __init__(self, settings: Settings, http: httpx.Client | None = None) -> None: + def __init__( + self, + settings: Settings, + http: httpx.Client | None = None, + scrub_key: bytes | None = None, + ) -> None: self.settings = settings + # None: send as is (scrub=False). Otherwise personal data is tokenized before queueing. + self.scrub_key = scrub_key self._http = http self._queue: collections.deque[dict[str, Any]] = collections.deque(maxlen=MAX_QUEUE) self._lock = threading.Lock() @@ -59,13 +67,19 @@ def add( """Queue one call. Never raises. A streamed answer is not captured yet: the record keeps the request and says so.""" try: + request, answer = payload(kwargs), None if stream else jsonable(response) + if self.scrub_key is not None: + request = scrub_call(request, self.scrub_key) + answer = scrub_call(answer, self.scrub_key) record = { "provider": provider, "conversation_id": conversation_id, "timestamp": datetime.now(timezone.utc).isoformat(), - "request": payload(kwargs), - "response": None if stream else jsonable(response), + "request": request, + "response": answer, } + if self.scrub_key is not None: + record["scrubbed"] = True if stream: record["stream"] = True with self._lock: diff --git a/src/agentcompile/_scrub.py b/src/agentcompile/_scrub.py new file mode 100644 index 0000000..81120f0 --- /dev/null +++ b/src/agentcompile/_scrub.py @@ -0,0 +1,164 @@ +"""Scrubbing on your machine: personal data in captured calls becomes keyed tokens before +anything is sent (on by default with capture; `wrap(..., scrub=False)` turns it off). + +The same rules and token format as AgentCompile's server (``: the first 8 hex +characters of HMAC-SHA256(key, normalized value)), so the same value always becomes the same +token and AgentCompile can still match values across a conversation without ever seeing them. +Emails always; payment cards (13-19 digits passing Luhn); phone-shaped numbers; and values +under keys that name a card, phone or account outright. Ids are left alone. + +The key never leaves your machine: AGENTCOMPILE_SCRUB_KEY (set the same one on every server, +so tokens match across your fleet), else a random key created once in ~/.agentcompile/. +""" + +from __future__ import annotations + +import hashlib +import hmac +import json +import os +import re +import secrets +from pathlib import Path +from typing import Any + +_EMAIL_RE = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}") +_CARD_RE = re.compile( + r"(? bytes: + """AGENTCOMPILE_SCRUB_KEY, else ~/.agentcompile/scrub.key (created once, owner-only).""" + env = os.environ.get("AGENTCOMPILE_SCRUB_KEY") + if env: + return env.encode() + path = ( + Path(os.environ.get("AGENTCOMPILE_HOME", Path.home() / ".agentcompile")) / "scrub.key" + ) + try: + return path.read_bytes().strip() + except FileNotFoundError: + path.parent.mkdir(parents=True, exist_ok=True) + key = secrets.token_hex(32) + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + with os.fdopen(fd, "w") as fh: + fh.write(key) + return key.encode() + + +def _token(kind: str, normalized: str, key: bytes) -> str: + digest = hmac.new(key, normalized.encode(), hashlib.sha256).hexdigest()[:8] + return f"<{kind}:{digest}>" + + +def _luhn_ok(digits: str) -> bool: + total = 0 + for i, ch in enumerate(reversed(digits)): + digit = int(ch) + if i % 2 == 1: + digit *= 2 + if digit > 9: + digit -= 9 + total += digit + return total % 10 == 0 + + +def scrub_text(text: str, key: bytes, *, numbers: bool = True) -> str: + text = _EMAIL_RE.sub(lambda m: _token("email", m.group(0).strip().lower(), key), text) + if not numbers: + return text + + def card(match: re.Match[str]) -> str: + digits = re.sub(r"\D", "", match.group(0)) + if 13 <= len(digits) <= 19 and _luhn_ok(digits): + return _token("card", digits, key) + return match.group(0) + + def phone(match: re.Match[str]) -> str: + digits = re.sub(r"\D", "", match.group(0)) + if 10 <= len(digits) <= 15: + return _token("phone", digits, key) + return match.group(0) + + return _PHONE_RE.sub(phone, _CARD_RE.sub(card, text)) + + +def _phone_field(value: str, key: bytes) -> str: + digits = re.sub(r"\D", "", value) + if 7 <= len(digits) <= 15 and not _EMAIL_RE.search(value): + return _token("phone", digits, key) + return scrub_text(value, key) + + +def _is_card_field(field: str) -> bool: + words = {w.lower() for w in re.findall(r"[A-Za-z][a-z]*", field)} + return bool(field) and bool(words & {"card", "cc", "pan"}) + + +def _sensitive_field(value: Any, field: str, key: bytes) -> str | None: + raw = ( + str(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else value + ) + if not isinstance(raw, str) or not field: + return None + digits = re.sub(r"\D", "", raw) + if _ACCOUNT_KEY.search(field) and len(digits) >= 4: + return _token("account", digits, key) + if _is_card_field(field) and 13 <= len(digits) <= 19 and _luhn_ok(digits): + return _token("card", digits, key) + if _PHONE_KEY.search(field) and 7 <= len(digits) <= 15 and not _EMAIL_RE.search(raw): + return _token("phone", digits, key) + return None + + +def scrub_value(value: Any, key: bytes, field: str = "") -> Any: + sensitive = _sensitive_field(value, field, key) + if sensitive is not None: + return sensitive + if isinstance(value, str): + if field and _PHONE_KEY.search(field): + return _phone_field(value, key) + id_like = bool(field and _ID_KEY.search(field) and not _is_card_field(field)) + return scrub_text(value, key, numbers=not id_like) + if isinstance(value, dict): + return {k: scrub_value(v, key, str(k)) for k, v in value.items()} + if isinstance(value, list): + return [scrub_value(v, key, field) for v in value] + return value + + +def scrub_tool_output(output: str, key: bytes) -> str: + """A tool result or arguments string: JSON is scrubbed field by field, text as text.""" + try: + parsed = json.loads(output) + except ValueError: + return scrub_text(output, key) + separators = (", ", ": ") if (", " in output or ": " in output) else (",", ":") + return json.dumps(scrub_value(parsed, key), separators=separators) + + +def scrub_call(value: Any, key: bytes, field: str = "") -> Any: + """A captured request or answer: every string scrubbed; a string holding JSON (tool + arguments, tool results) scrubbed as JSON, field names included.""" + if isinstance(value, str) and value[:1] in ("{", "["): + return scrub_tool_output(value, key) + if isinstance(value, dict): + return {k: scrub_call(v, key, str(k)) for k, v in value.items()} + if isinstance(value, list): + return [scrub_call(v, key, field) for v in value] + return scrub_value(value, key, field) diff --git a/tests/data/scrub-vectors.json b/tests/data/scrub-vectors.json new file mode 100644 index 0000000..190a0ca --- /dev/null +++ b/tests/data/scrub-vectors.json @@ -0,0 +1,141 @@ +{ + "key": "test-key-0123456789abcdef", + "text": [ + { + "input": "my email is Mia.Garcia2723@Example.com, thanks", + "output": "my email is , thanks" + }, + { + "input": "card 4111 1111 1111 1111 and 4111-1111-1111-1111 and 4111111111111111", + "output": "card and and " + }, + { + "input": "not a card: 1234567812345678", + "output": "not a card: 1234567812345678" + }, + { + "input": "call +44 20 7946 0958 or (415) 555-0199 or 415.555.0199 or 415-555-0199 or +14155550199", + "output": "call or or or or " + }, + { + "input": "order #W1234567 placed 2024-05-27, ids 1684786391 5052031638", + "output": "order #W1234567 placed 2024-05-27, ids 1684786391 5052031638" + }, + { + "input": "price 1234567890.12 and plain 4155550199", + "output": "price 1234567890.12 and plain 4155550199" + }, + { + "input": "no personal data here", + "output": "no personal data here" + }, + { + "input": "two emails a@b.com and A@B.COM", + "output": "two emails and " + } + ], + "value": [ + { + "input": { + "email": "x@y.io", + "order_id": "4111111111111111", + "phone": "4155550199", + "card_number": "4111111111111111" + }, + "output": { + "email": "", + "order_id": "4111111111111111", + "phone": "", + "card_number": "" + } + }, + { + "input": { + "account_number": "000123456789", + "iban": "GB82WEST12345698765432", + "note": "reach me at 415-555-0199" + }, + "output": { + "account_number": "", + "iban": "", + "note": "reach me at " + } + }, + { + "input": { + "items": [ + { + "sku": "4111111111111111", + "contact": "+1 415 555 0199" + } + ], + "mobile": 4155550199, + "cc": 4111111111111111 + }, + "output": { + "items": [ + { + "sku": "4111111111111111", + "contact": "" + } + ], + "mobile": "", + "cc": "" + } + }, + { + "input": { + "user": { + "emails": [ + "p@q.co", + "not-an-email" + ], + "ts": "1684786391" + } + }, + "output": { + "user": { + "emails": [ + "", + "not-an-email" + ], + "ts": "1684786391" + } + } + }, + { + "input": [ + 1, + 2.5, + true, + null, + "plain" + ], + "output": [ + 1, + 2.5, + true, + null, + "plain" + ] + } + ], + "tool_output": [ + { + "input": "{\"user_id\": \"u1\", \"email\": \"mia@example.com\", \"phone\": \"415-555-0199\"}", + "output": "{\"user_id\": \"u1\", \"email\": \"\", \"phone\": \"\"}" + }, + { + "input": "{\"a\":1,\"card\":\"4111111111111111\"}", + "output": "{\"a\":1,\"card\":\"\"}" + }, + { + "input": "not json: mail me@x.org", + "output": "not json: mail " + }, + { + "input": "[{\"order_id\":\"#W1\",\"payment\":{\"card_number\":\"4111 1111 1111 1111\"}}]", + "output": "[{\"order_id\":\"#W1\",\"payment\":{\"card_number\":\"\"}}]" + } + ] +} \ No newline at end of file diff --git a/tests/test_scrub.py b/tests/test_scrub.py new file mode 100644 index 0000000..ea854b4 --- /dev/null +++ b/tests/test_scrub.py @@ -0,0 +1,140 @@ +"""Scrubbing on the customer's machine: identical tokens to AgentCompile's server (shared +vectors, generated by the server's own scrubber), applied to captured calls before sending.""" + +from __future__ import annotations + +import json +import os +import stat +from pathlib import Path +from typing import Any + +import httpx +import openai +import pytest + +import agentcompile +from agentcompile import _capture, _scrub + +VECTORS = json.loads((Path(__file__).parent / "data" / "scrub-vectors.json").read_text()) +KEY = VECTORS["key"].encode() + + +@pytest.mark.parametrize("case", VECTORS["text"], ids=lambda c: c["input"][:30]) +def test_text_matches_the_server(case: dict[str, Any]) -> None: + assert _scrub.scrub_text(case["input"], KEY) == case["output"] + + +@pytest.mark.parametrize("case", VECTORS["value"], ids=lambda c: str(c["input"])[:30]) +def test_values_match_the_server(case: dict[str, Any]) -> None: + assert _scrub.scrub_value(case["input"], KEY) == case["output"] + + +@pytest.mark.parametrize("case", VECTORS["tool_output"], ids=lambda c: c["input"][:30]) +def test_tool_outputs_match_the_server(case: dict[str, Any]) -> None: + assert _scrub.scrub_tool_output(case["input"], KEY) == case["output"] + + +def test_the_key_comes_from_the_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AGENTCOMPILE_SCRUB_KEY", "fleet-key") + assert _scrub.load_key() == b"fleet-key" + + +def test_otherwise_a_key_is_created_once_owner_only( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.delenv("AGENTCOMPILE_SCRUB_KEY", raising=False) + monkeypatch.setenv("AGENTCOMPILE_HOME", str(tmp_path)) + first = _scrub.load_key() + assert len(first) == 64 and _scrub.load_key() == first + mode = os.stat(tmp_path / "scrub.key").st_mode + assert stat.S_IMODE(mode) == 0o600 + + +MESSAGES = [ + {"role": "user", "content": "I'm mia@example.com, card 4111 1111 1111 1111"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "t1", + "type": "function", + "function": {"name": "find_user", "arguments": '{"email": "mia@example.com"}'}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "t1", + "content": '{"user_id": "u1", "phone": "415-555-0199"}', + }, +] + + +class Service: + def __init__(self) -> None: + self.sent: list[dict[str, Any]] = [] + + def __call__(self, request: httpx.Request) -> httpx.Response: + if request.url.path == "/v1/capture": + self.sent += json.loads(request.content)["exchanges"] + return httpx.Response(200, json={}) + return httpx.Response(200, json={"action": "forward", "reason": "x"}) + + +def _send(scrub: bool, monkeypatch: pytest.MonkeyPatch) -> list[dict[str, Any]]: + monkeypatch.setenv("AGENTCOMPILE_SCRUB_KEY", VECTORS["key"]) + service = Service() + completion = { + "id": "c1", + "object": "chat.completion", + "created": 0, + "model": "m", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Done, mia@example.com"}, + "finish_reason": "stop", + } + ], + } + provider = openai.OpenAI( + api_key="sk", + base_url="http://provider.test/v1", + http_client=httpx.Client( + transport=httpx.MockTransport(lambda r: httpx.Response(200, json=completion)) + ), + ) + client = agentcompile.wrap( + provider, + base_url="http://ac.test", + trail=False, + capture=True, + scrub=scrub, + http_client=httpx.Client(transport=httpx.MockTransport(service)), + ) + client.chat.completions.create(model="m", messages=MESSAGES, conversation_id="c1") + agentcompile.flush(5) + _capture._ALL.clear() + return service.sent + + +def test_captured_calls_leave_the_machine_scrubbed(monkeypatch: pytest.MonkeyPatch) -> None: + [record] = _send(True, monkeypatch) + sent = json.dumps(record) + assert "mia@example.com" not in sent + assert "4111 1111 1111 1111" not in sent and "415-555-0199" not in sent + email = _scrub.scrub_text("mia@example.com", KEY) + messages = record["request"]["messages"] + assert email in messages[0]["content"] + args = json.loads(messages[1]["tool_calls"][0]["function"]["arguments"]) + assert args == {"email": email} # the same token in what was said and what was passed + assert json.loads(messages[2]["content"])["user_id"] == "u1" # ids untouched + assert email in record["response"]["choices"][0]["message"]["content"] + assert record["scrubbed"] is True + + +def test_scrub_false_sends_calls_as_they_are(monkeypatch: pytest.MonkeyPatch) -> None: + [record] = _send(False, monkeypatch) + assert "mia@example.com" in json.dumps(record) and "scrubbed" not in record