aw-sdk 0.1.0__py3-none-any.whl

agentwatch/__init__.py

@@ -0,0 +1,4 @@
+from .wrapper import WatchedOpenAI, wrap, analyze_text, INGEST_URL, AgentBudgetExceeded, AgentBudgetCheckUnavailable
+
+__version__ = "0.1.0"
+__all__ = ["WatchedOpenAI", "wrap", "analyze_text", "INGEST_URL", "AgentBudgetExceeded", "AgentBudgetCheckUnavailable"]

agentwatch/wrapper.py

@@ -0,0 +1,499 @@
+from __future__ import annotations
+
+import base64
+import math
+import queue
+import re
+import logging
+import threading
+import time
+import uuid
+from typing import Any, Dict, List, Optional, Set
+
+_log = logging.getLogger("agentwatch")
+
+import httpx
+from openai import OpenAI as _OpenAI
+
+# ---------------------------------------------------------------------------
+# Risk detection — aligned with the TypeScript classifier (classifier.ts)
+# ---------------------------------------------------------------------------
+
+EMAIL_RE = re.compile(r"\b[A-Z0-9._%+-]{1,64}@[A-Z0-9.-]{1,253}\.[A-Z]{2,63}\b", re.IGNORECASE)
+SSN_RE = re.compile(r"\b(?!000|666|9\d{2})\d{3}[- ]?(?!00)\d{2}[- ]?(?!0000)\d{4}\b")
+AWS_ACCESS_KEY_RE = re.compile(r"\b(?:AKIA|ASIA|AIDA|AROA|AGPA|ANPA)[A-Z0-9]{16}\b")
+STRIPE_SECRET_KEY_RE = re.compile(r"\bsk_(?:live|test)_[A-Za-z0-9]{16,}\b")
+GITHUB_TOKEN_RE = re.compile(r"\b(?:gh[pousr]_[A-Za-z0-9_]{36,255}|github_pat_[A-Za-z0-9_]{22,255})\b")
+CREDIT_CARD_CANDIDATE_RE = re.compile(r"(?=(?:^|[^\d])((?:\d[ -]?){13,19})(?!\d))")
+JWT_CANDIDATE_RE = re.compile(r"\b[A-Za-z0-9_-]{10,512}\.[A-Za-z0-9_-]{10,4096}\.[A-Za-z0-9_-]{16,1024}\b")
+
+MAX_SCAN_CHARS = 256 * 1024
+EDGE_SCAN_CHARS = MAX_SCAN_CHARS // 2
+MAX_STRUCTURED_CANDIDATES = 256
+HIGH_ENTROPY_MIN_SCORE = 3.25
+JWT_SIGNATURE_MIN_ENTROPY = 3.5
+
+TAG_ORDER = (
+    "PII_EMAIL",
+    "PII_SSN",
+    "FINANCIAL_CREDIT_CARD",
+    "SECRET_AWS_ACCESS_KEY",
+    "SECRET_STRIPE",
+    "SECRET_GITHUB",
+    "SECRET_JWT",
+)
+
+
+def analyze_text(text: str) -> List[str]:
+    """Scan text for PII and secret risks. Returns canonical tag names."""
+    if not text:
+        return []
+
+    scan = _bounded_scan_text(text)
+    risks: Set[str] = set()
+
+    if EMAIL_RE.search(scan):
+        risks.add("PII_EMAIL")
+    if SSN_RE.search(scan):
+        risks.add("PII_SSN")
+    if _has_credit_card(scan):
+        risks.add("FINANCIAL_CREDIT_CARD")
+    if _has_aws_access_key(scan):
+        risks.add("SECRET_AWS_ACCESS_KEY")
+    if STRIPE_SECRET_KEY_RE.search(scan):
+        risks.add("SECRET_STRIPE")
+    if GITHUB_TOKEN_RE.search(scan):
+        risks.add("SECRET_GITHUB")
+    if _has_high_entropy_jwt(scan):
+        risks.add("SECRET_JWT")
+
+    return [tag for tag in TAG_ORDER if tag in risks]
+
+
+# ---------------------------------------------------------------------------
+# Classifier helpers
+# ---------------------------------------------------------------------------
+
+def _bounded_scan_text(text: str) -> str:
+    if len(text) <= MAX_SCAN_CHARS:
+        return text
+    return f"{text[:EDGE_SCAN_CHARS]}\n{text[-EDGE_SCAN_CHARS:]}"
+
+
+def _has_credit_card(text: str) -> bool:
+    checked = 0
+    for match in CREDIT_CARD_CANDIDATE_RE.finditer(text):
+        checked += 1
+        if checked > MAX_STRUCTURED_CANDIDATES:
+            return False
+        digits = _digits_only(match.group(1))
+        if _is_likely_credit_card(digits):
+            return True
+    return False
+
+
+def _is_likely_credit_card(digits: str) -> bool:
+    return 13 <= len(digits) <= 19 and _has_known_card_prefix(digits) and _passes_luhn(digits)
+
+
+def _has_known_card_prefix(digits: str) -> bool:
+    prefix2 = _safe_int(digits[:2])
+    prefix3 = _safe_int(digits[:3])
+    prefix4 = _safe_int(digits[:4])
+    prefix6 = _safe_int(digits[:6])
+    length = len(digits)
+
+    if digits.startswith("4") and length in (13, 16, 19):
+        return True
+    if length == 15 and prefix2 in (34, 37):
+        return True
+    if length == 16 and (51 <= prefix2 <= 55 or 2221 <= prefix4 <= 2720):
+        return True
+    if length in (16, 19) and (
+        digits.startswith("6011")
+        or digits.startswith("65")
+        or 644 <= prefix3 <= 649
+        or 622126 <= prefix6 <= 622925
+    ):
+        return True
+    return length in (16, 17, 18, 19) and 3528 <= prefix4 <= 3589
+
+
+def _passes_luhn(digits: str) -> bool:
+    total = 0
+    should_double = False
+    for char in reversed(digits):
+        digit = ord(char) - 48
+        if digit < 0 or digit > 9:
+            return False
+        if should_double:
+            digit *= 2
+            if digit > 9:
+                digit -= 9
+        total += digit
+        should_double = not should_double
+    return total % 10 == 0
+
+
+def _has_aws_access_key(text: str) -> bool:
+    for match in AWS_ACCESS_KEY_RE.finditer(text):
+        if _shannon_entropy(match.group(0)) >= HIGH_ENTROPY_MIN_SCORE:
+            return True
+    return False
+
+
+def _has_high_entropy_jwt(text: str) -> bool:
+    checked = 0
+    for match in JWT_CANDIDATE_RE.finditer(text):
+        checked += 1
+        if checked > MAX_STRUCTURED_CANDIDATES:
+            return False
+        if _is_high_entropy_jwt(match.group(0)):
+            return True
+    return False
+
+
+def _is_high_entropy_jwt(token: str) -> bool:
+    parts = token.split(".")
+    if len(parts) != 3 or not _looks_like_jwt_header(parts[0]):
+        return False
+    return _shannon_entropy(parts[2]) >= JWT_SIGNATURE_MIN_ENTROPY
+
+
+def _looks_like_jwt_header(segment: str) -> bool:
+    decoded = _base64url_decode(segment)
+    return '"alg"' in decoded and '"typ"' in decoded
+
+
+def _base64url_decode(value: str) -> str:
+    normalized = value.replace("-", "+").replace("_", "/")
+    padding = (4 - len(normalized) % 4) % 4
+    try:
+        return base64.b64decode(normalized + "=" * padding).decode("utf-8", errors="replace")
+    except Exception:
+        return ""
+
+
+def _shannon_entropy(value: str) -> float:
+    if not value:
+        return 0.0
+    counts: Dict[str, int] = {}
+    for char in value:
+        counts[char] = counts.get(char, 0) + 1
+    entropy = 0.0
+    length = len(value)
+    for count in counts.values():
+        probability = count / length
+        entropy -= probability * math.log2(probability)
+    return entropy
+
+
+def _digits_only(value: str) -> str:
+    return "".join(c for c in value if "0" <= c <= "9")
+
+
+def _safe_int(value: str) -> int:
+    return int(value) if value.isdigit() else -1
+
+
+def _extract_messages_text(messages: Any) -> str:
+    """Recursively extract all text content from an OpenAI messages list."""
+    if not messages or not isinstance(messages, list):
+        return ""
+    chunks: List[str] = []
+    for msg in messages:
+        if not isinstance(msg, dict):
+            continue
+        role = msg.get("role")
+        if isinstance(role, str):
+            chunks.append(role)
+        content = msg.get("content")
+        if isinstance(content, str):
+            chunks.append(content)
+        elif isinstance(content, list):
+            for part in content:
+                if isinstance(part, dict):
+                    for key in ("text", "input", "content"):
+                        val = part.get(key)
+                        if isinstance(val, str):
+                            chunks.append(val)
+    return "\n".join(chunks)
+
+
+def _extract_completion_text(response: Any) -> str:
+    """Extract text from an OpenAI ChatCompletion response."""
+    choices = getattr(response, "choices", None)
+    if not choices:
+        return ""
+    chunks: List[str] = []
+    for choice in choices:
+        message = getattr(choice, "message", None)
+        if message is not None:
+            content = getattr(message, "content", None)
+            if isinstance(content, str):
+                chunks.append(content)
+    return "\n".join(chunks)
+
+
+# ---------------------------------------------------------------------------
+# Async logger — single worker thread + bounded queue (never blocks caller)
+# ---------------------------------------------------------------------------
+
+INGEST_URL = "https://agentwatch-edge-proxy.agentwatch-proxy.workers.dev/v1/ingest"
+MAX_PENDING_LOGS = 1024
+
+
+class AgentBudgetExceeded(Exception):
+    """Raised when an agent session exceeds its configured token budget limit."""
+    def __init__(self, session_id: str, spent: float, limit: float) -> None:
+        super().__init__(f"Agent budget exceeded for session '{session_id}': spent ${spent:.4f}, limit ${limit:.4f}")
+        self.session_id = session_id
+        self.spent = spent
+        self.limit = limit
+
+
+class AgentBudgetCheckUnavailable(Exception):
+    """Raised when enforcement_fail_open=False and the budget check endpoint is unreachable."""
+    def __init__(self, session_id: str, reason: str) -> None:
+        super().__init__(f"AgentWatch budget check unavailable for session '{session_id}': {reason}")
+        self.session_id = session_id
+        self.reason = reason
+
+
+class _AsyncLogger:
+    """Background log dispatcher. Single daemon thread, bounded queue."""
+
+    def __init__(self, api_key: str, ingest_url: str = INGEST_URL, timeout: float = 2.0, max_pending: int = MAX_PENDING_LOGS) -> None:
+        self._api_key = api_key
+        self._ingest_url = ingest_url
+        self._timeout = timeout
+        self._queue: queue.Queue[Dict[str, Any]] = queue.Queue(maxsize=max_pending)
+        self._started = False
+        self._lock = threading.Lock()
+
+    def submit(self, payload: Dict[str, Any]) -> None:
+        try:
+            self._ensure_started()
+            self._queue.put_nowait(payload)
+        except Exception:
+            _log.debug("agentwatch: failed to enqueue log", exc_info=True)
+
+    def _ensure_started(self) -> None:
+        if self._started:
+            return
+        with self._lock:
+            if self._started:
+                return
+            thread = threading.Thread(target=self._run, name="agentwatch-logger", daemon=True)
+            thread.start()
+            self._started = True
+
+    def _run(self) -> None:
+        while True:
+            try:
+                payload = self._queue.get()
+                self._send(payload)
+            except Exception:
+                _log.debug("agentwatch: log processing error", exc_info=True)
+            finally:
+                try:
+                    self._queue.task_done()
+                except Exception:
+                    pass
+
+    def _send(self, payload: Dict[str, Any]) -> None:
+        try:
+            httpx.post(
+                self._ingest_url,
+                json=payload,
+                headers={"Authorization": f"Bearer {self._api_key}"},
+                timeout=self._timeout,
+            )
+        except Exception:
+            _log.debug("agentwatch: failed to send telemetry", exc_info=True)
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+class WatchedOpenAI(_OpenAI):
+    """Drop-in replacement for openai.OpenAI with built-in telemetry."""
+
+    def __init__(
+        self,
+        *args: Any,
+        agentwatch_api_key: str,
+        agentwatch_project: Optional[str] = None,
+        agentwatch_team: Optional[str] = None,
+        agentwatch_session_id: Optional[str] = None,
+        agentwatch_session_budget_usd: Optional[float] = None,
+        agentwatch_monthly_budget_usd: Optional[float] = None,
+        agentwatch_enforcement_mode: bool = False,
+        agentwatch_enforcement_fail_open: bool = True,
+        ingest_url: str = INGEST_URL,
+        timeout_seconds: float = 2.0,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self._logger = _AsyncLogger(api_key=agentwatch_api_key, ingest_url=ingest_url, timeout=timeout_seconds)
+        self._project = agentwatch_project
+        self._team = agentwatch_team
+        self._session_id = agentwatch_session_id or str(uuid.uuid4())
+        self._session_budget_usd = agentwatch_session_budget_usd
+        self._monthly_budget_usd = agentwatch_monthly_budget_usd
+        self._enforcement_mode = agentwatch_enforcement_mode
+        self._enforcement_fail_open = agentwatch_enforcement_fail_open
+        self._iteration_index = 0
+        self._orig_create = self.chat.completions.create
+        self.chat.completions.create = self._watched_create
+
+    def _watched_create(self, *args: Any, **kwargs: Any) -> Any:
+        if self._enforcement_mode and self._session_id and self._session_budget_usd is not None:
+            base_url = self._logger._ingest_url.rsplit("/", 1)[0]
+            budget_url = f"{base_url}/budget-check?session_id={self._session_id}&limit_usd={self._session_budget_usd}"
+            try:
+                resp = httpx.get(
+                    budget_url,
+                    headers={"Authorization": f"Bearer {self._logger._api_key}"},
+                    timeout=self._logger._timeout,
+                )
+                resp.raise_for_status()
+                status = resp.json()
+                if status.get("exceeded"):
+                    raise AgentBudgetExceeded(
+                        session_id=self._session_id,
+                        spent=status.get("spent_usd", 0.0),
+                        limit=status.get("limit_usd", self._session_budget_usd),
+                    )
+            except AgentBudgetExceeded:
+                raise
+            except Exception as e:
+                if not self._enforcement_fail_open:
+                    raise AgentBudgetCheckUnavailable(session_id=self._session_id, reason=str(e)) from e
+                else:
+                    _log.debug("agentwatch: budget check failed, failing open", exc_info=True)
+
+        prompt_text = _extract_messages_text(kwargs.get("messages"))
+        prompt_risks = set(analyze_text(prompt_text))
+
+        t0 = time.time()
+        response = self._orig_create(*args, **kwargs)
+        latency = int((time.time() - t0) * 1000)
+
+        # Scan completion text for leaked PII/secrets
+        completion_text = _extract_completion_text(response)
+        completion_risks = set(analyze_text(completion_text))
+        combined_risks = [tag for tag in TAG_ORDER if tag in (prompt_risks | completion_risks)]
+
+        usage = getattr(response, "usage", None)
+        prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0
+        completion_tokens = getattr(usage, "completion_tokens", 0) or 0
+
+        payload: Dict[str, Any] = {
+            "model": kwargs.get("model", "unknown"),
+            "latency_ms": latency,
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "identified_risks": combined_risks,
+        }
+        if self._project:
+            payload["project"] = self._project
+        if self._team:
+            payload["team"] = self._team
+        if self._session_id:
+            payload["session_id"] = self._session_id
+            self._iteration_index += 1
+            payload["iteration_index"] = self._iteration_index
+
+        self._logger.submit(payload)
+
+        return response
+
+
+def wrap(
+    openai_client: Any,
+    *,
+    agentwatch_api_key: str,
+    agentwatch_project: Optional[str] = None,
+    agentwatch_team: Optional[str] = None,
+    agentwatch_session_id: Optional[str] = None,
+    agentwatch_session_budget_usd: Optional[float] = None,
+    agentwatch_monthly_budget_usd: Optional[float] = None,
+    agentwatch_enforcement_mode: bool = False,
+    agentwatch_enforcement_fail_open: bool = True,
+    ingest_url: str = INGEST_URL,
+    timeout_seconds: float = 2.0,
+) -> Any:
+    """Wrap an existing OpenAI client with AgentWatch telemetry via composition."""
+    logger = _AsyncLogger(api_key=agentwatch_api_key, ingest_url=ingest_url, timeout=timeout_seconds)
+    _project = agentwatch_project
+    _team = agentwatch_team
+    _session_id = agentwatch_session_id or str(uuid.uuid4())
+    # Use a mutable list so the inner function can modify the iteration count
+    _state = {"iteration_index": 0}
+    orig_create = openai_client.chat.completions.create
+
+    def watched_create(*args: Any, **kwargs: Any) -> Any:
+        if agentwatch_enforcement_mode and _session_id and agentwatch_session_budget_usd is not None:
+            base_url = logger._ingest_url.rsplit("/", 1)[0]
+            budget_url = f"{base_url}/budget-check?session_id={_session_id}&limit_usd={agentwatch_session_budget_usd}"
+            try:
+                resp = httpx.get(
+                    budget_url,
+                    headers={"Authorization": f"Bearer {logger._api_key}"},
+                    timeout=logger._timeout,
+                )
+                resp.raise_for_status()
+                status = resp.json()
+                if status.get("exceeded"):
+                    raise AgentBudgetExceeded(
+                        session_id=_session_id,
+                        spent=status.get("spent_usd", 0.0),
+                        limit=status.get("limit_usd", agentwatch_session_budget_usd),
+                    )
+            except AgentBudgetExceeded:
+                raise
+            except Exception as e:
+                if not agentwatch_enforcement_fail_open:
+                    raise AgentBudgetCheckUnavailable(session_id=_session_id, reason=str(e)) from e
+                else:
+                    _log.debug("agentwatch: budget check failed, failing open", exc_info=True)
+
+        prompt_text = _extract_messages_text(kwargs.get("messages"))
+        prompt_risks = set(analyze_text(prompt_text))
+
+        t0 = time.time()
+        response = orig_create(*args, **kwargs)
+        latency = int((time.time() - t0) * 1000)
+
+        completion_text = _extract_completion_text(response)
+        completion_risks = set(analyze_text(completion_text))
+        combined_risks = [tag for tag in TAG_ORDER if tag in (prompt_risks | completion_risks)]
+
+        usage = getattr(response, "usage", None)
+        prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0
+        completion_tokens_val = getattr(usage, "completion_tokens", 0) or 0
+
+        payload: Dict[str, Any] = {
+            "model": kwargs.get("model", "unknown"),
+            "latency_ms": latency,
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens_val,
+            "identified_risks": combined_risks,
+        }
+        if _project:
+            payload["project"] = _project
+        if _team:
+            payload["team"] = _team
+        if _session_id:
+            payload["session_id"] = _session_id
+            _state["iteration_index"] += 1
+            payload["iteration_index"] = _state["iteration_index"]
+
+        logger.submit(payload)
+
+        return response
+
+    openai_client.chat.completions.create = watched_create
+    return openai_client

aw_sdk-0.1.0.dist-info/METADATA

@@ -0,0 +1,9 @@
+Metadata-Version: 2.4
+Name: aw-sdk
+Version: 0.1.0
+Summary: Lightweight drop-in wrapper for the OpenAI Python client that logs telemetry and detects PII risks.
+License: MIT
+Project-URL: Homepage, https://github.com/agentwatch/agentwatch
+Requires-Python: >=3.9
+Requires-Dist: openai
+Requires-Dist: httpx

aw_sdk-0.1.0.dist-info/RECORD

@@ -0,0 +1,6 @@
+agentwatch/__init__.py,sha256=R67VdcBDPU3C9yOsUr1eqCOl9w3WYEDobvRsj4BW18k,260
+agentwatch/wrapper.py,sha256=sdy1sxXBiO2DsD9hQtBr36cOS7RK-BreoyKDil6idt0,18222
+aw_sdk-0.1.0.dist-info/METADATA,sha256=gtfRaAH-Xvs2nZ6DnLQsZKrxWMe-G-SFIRc0h9KUmfM,302
+aw_sdk-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+aw_sdk-0.1.0.dist-info/top_level.txt,sha256=IML7iKg-Q0BkJznZO5NsCBtfy5WUnk3Z4q7HzdK2keI,11
+aw_sdk-0.1.0.dist-info/RECORD,,

aw_sdk-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

aw_sdk-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+agentwatch