studiomeyer-aishield 0.1.0__py3-none-any.whl

ai_shield/__init__.py

@@ -0,0 +1,47 @@
+"""ai-shield — LLM input shield for prompt-injection, PII, tool-policy, cost.
+
+Python 1:1 port of ai-shield-core (TypeScript, MIT, 4 audit rounds).
+"""
+
+from __future__ import annotations
+
+from ai_shield.shield import AIShield
+from ai_shield.types import (
+    AuditRecord,
+    BudgetCheckResult,
+    BudgetConfig,
+    BudgetPeriod,
+    CostRecord,
+    Decision,
+    PIIAction,
+    PIIEntity,
+    PIIType,
+    ScannerResult,
+    ScanResult,
+    ToolCall,
+    ToolManifestPin,
+    Violation,
+    ViolationType,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "AIShield",
+    "AuditRecord",
+    "BudgetCheckResult",
+    "BudgetConfig",
+    "BudgetPeriod",
+    "CostRecord",
+    "Decision",
+    "PIIAction",
+    "PIIEntity",
+    "PIIType",
+    "ScanResult",
+    "ScannerResult",
+    "ToolCall",
+    "ToolManifestPin",
+    "Violation",
+    "ViolationType",
+    "__version__",
+]

ai_shield/audit/__init__.py

@@ -0,0 +1,13 @@
+"""Audit subpackage — async batched logger + store interface."""
+
+from __future__ import annotations
+
+from ai_shield.audit.logger import AuditLogger, ConsoleAuditStore, MemoryAuditStore
+from ai_shield.audit.types import AuditStore
+
+__all__ = [
+    "AuditLogger",
+    "AuditStore",
+    "ConsoleAuditStore",
+    "MemoryAuditStore",
+]

ai_shield/audit/logger.py

@@ -0,0 +1,155 @@
+"""Async batched audit logger.
+
+1:1 port of `packages/core/src/audit/logger.ts`.
+
+Inputs are NEVER stored in plain text — only `sha256(input)` and an
+optional `sha256(user_id)[:32]` end up on disk / in the store.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import hashlib
+import sys
+import unicodedata
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any, TextIO
+
+from ai_shield.audit.types import AuditStore
+from ai_shield.types import AuditRecord, Decision, Violation
+
+
+def _hash_input(text: str) -> str:
+    normalized = unicodedata.normalize("NFKD", text).encode("utf-8")
+    return hashlib.sha256(normalized).hexdigest()
+
+
+def _hash_user(user_id: str | None) -> str | None:
+    if user_id is None:
+        return None
+    return hashlib.sha256(user_id.encode("utf-8")).hexdigest()[:32]
+
+
+@dataclass
+class ConsoleAuditStore:
+    """Write JSON lines to a TextIO sink (default stderr)."""
+
+    sink: TextIO = field(default_factory=lambda: sys.stderr)
+
+    async def write(self, record: AuditRecord) -> None:
+        self.sink.write(record.model_dump_json() + "\n")
+        self.sink.flush()
+
+    async def write_batch(self, records: list[AuditRecord]) -> None:
+        for r in records:
+            self.sink.write(r.model_dump_json() + "\n")
+        self.sink.flush()
+
+    async def flush(self) -> None:
+        self.sink.flush()
+
+    async def close(self) -> None:
+        self.sink.flush()
+
+
+@dataclass
+class MemoryAuditStore:
+    """In-process AuditStore — useful for tests and short-lived processes."""
+
+    records: list[AuditRecord] = field(default_factory=list)
+
+    async def write(self, record: AuditRecord) -> None:
+        self.records.append(record)
+
+    async def write_batch(self, records: list[AuditRecord]) -> None:
+        self.records.extend(records)
+
+    async def flush(self) -> None:
+        return None
+
+    async def close(self) -> None:
+        return None
+
+
+class AuditLogger:
+    """Buffer audit records and flush periodically to an AuditStore."""
+
+    def __init__(
+        self,
+        store: AuditStore | None = None,
+        *,
+        flush_interval_seconds: float = 5.0,
+        max_buffer: int = 100,
+    ) -> None:
+        self.store: AuditStore = store if store is not None else ConsoleAuditStore()
+        self._flush_interval = flush_interval_seconds
+        self._max_buffer = max_buffer
+        self._buffer: list[AuditRecord] = []
+        self._lock = asyncio.Lock()
+        self._closed = False
+        self._task: asyncio.Task[None] | None = None
+
+    async def log(
+        self,
+        *,
+        text: str,
+        decision: Decision,
+        violations: list[Violation],
+        score: float,
+        user_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> None:
+        record = AuditRecord(
+            timestamp=datetime.now(timezone.utc).isoformat(),
+            user_id_hash=_hash_user(user_id),
+            input_sha256=_hash_input(text),
+            decision=decision,
+            violations=violations,
+            score=score,
+            metadata=metadata or {},
+        )
+        async with self._lock:
+            self._buffer.append(record)
+            buffer_full = len(self._buffer) >= self._max_buffer
+
+        if buffer_full:
+            await self.flush()
+        else:
+            self._ensure_task()
+
+    def _ensure_task(self) -> None:
+        if self._closed:
+            return
+        if self._task is not None and not self._task.done():
+            return
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            return
+        self._task = loop.create_task(self._auto_flush())
+
+    async def _auto_flush(self) -> None:
+        try:
+            await asyncio.sleep(self._flush_interval)
+            await self.flush()
+        except asyncio.CancelledError:
+            pass
+
+    async def flush(self) -> None:
+        async with self._lock:
+            if not self._buffer:
+                return
+            batch = self._buffer
+            self._buffer = []
+        await self.store.write_batch(batch)
+
+    async def close(self) -> None:
+        self._closed = True
+        if self._task is not None and not self._task.done():
+            self._task.cancel()
+            with contextlib.suppress(asyncio.CancelledError, Exception):
+                await self._task
+        await self.flush()
+        await self.store.close()

ai_shield/audit/types.py

@@ -0,0 +1,19 @@
+"""AuditStore protocol — pluggable backend for AuditLogger.
+
+1:1 port of `packages/core/src/audit/types.ts`.
+"""
+
+from __future__ import annotations
+
+from typing import Protocol
+
+from ai_shield.types import AuditRecord
+
+
+class AuditStore(Protocol):
+    """Pluggable backend for AuditLogger writes."""
+
+    async def write(self, record: AuditRecord) -> None: ...
+    async def write_batch(self, records: list[AuditRecord]) -> None: ...
+    async def flush(self) -> None: ...
+    async def close(self) -> None: ...

ai_shield/cache/__init__.py

@@ -0,0 +1,7 @@
+"""Cache subpackage — TTL + insertion-order LRU."""
+
+from __future__ import annotations
+
+from ai_shield.cache.lru import ScanLRUCache
+
+__all__ = ["ScanLRUCache"]

ai_shield/cache/lru.py

@@ -0,0 +1,89 @@
+"""TTL + insertion-order LRU cache for scan results.
+
+1:1 port of `packages/core/src/cache/lru.ts` — relies on dict-insertion
+order semantics (guaranteed in CPython 3.7+).
+"""
+
+from __future__ import annotations
+
+import time
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+V = TypeVar("V")
+
+
+@dataclass
+class _Entry(Generic[V]):
+    value: V
+    expires_at: float
+
+
+class ScanLRUCache(Generic[V]):
+    """Least-recently-used cache with per-entry TTL."""
+
+    def __init__(self, *, max_size: int = 1000, ttl_ms: int = 300_000) -> None:
+        if max_size < 1:
+            raise ValueError("max_size must be >= 1")
+        if ttl_ms < 0:
+            raise ValueError("ttl_ms must be >= 0")
+        self._max_size = max_size
+        self._ttl_ms = ttl_ms
+        self._data: OrderedDict[str, _Entry[V]] = OrderedDict()
+
+    def get(self, key: str) -> V | None:
+        entry = self._data.get(key)
+        if entry is None:
+            return None
+        if self._expired(entry):
+            del self._data[key]
+            return None
+        # promote to MRU
+        self._data.move_to_end(key)
+        return entry.value
+
+    def set(self, key: str, value: V) -> None:
+        if key in self._data:
+            del self._data[key]
+        elif len(self._data) >= self._max_size:
+            # evict oldest
+            self._data.popitem(last=False)
+        self._data[key] = _Entry(value=value, expires_at=self._now_ms() + self._ttl_ms)
+
+    def has(self, key: str) -> bool:
+        entry = self._data.get(key)
+        if entry is None:
+            return False
+        if self._expired(entry):
+            del self._data[key]
+            return False
+        return True
+
+    def delete(self, key: str) -> bool:
+        if key in self._data:
+            del self._data[key]
+            return True
+        return False
+
+    def clear(self) -> None:
+        self._data.clear()
+
+    def prune(self) -> int:
+        """Remove all expired entries. Returns count removed."""
+        keys = [k for k, e in self._data.items() if self._expired(e)]
+        for k in keys:
+            del self._data[k]
+        return len(keys)
+
+    def __len__(self) -> int:
+        return len(self._data)
+
+    @staticmethod
+    def _now_ms() -> float:
+        return time.monotonic() * 1000.0
+
+    def _expired(self, entry: _Entry[V]) -> bool:
+        if self._ttl_ms == 0:
+            return False
+        return self._now_ms() >= entry.expires_at

ai_shield/cost/__init__.py

@@ -0,0 +1,17 @@
+"""Cost subpackage — tracker + pricing + anomaly."""
+
+from __future__ import annotations
+
+from ai_shield.cost.anomaly import detect_anomaly
+from ai_shield.cost.pricing import MODEL_PRICING, estimate_cost, get_model_pricing
+from ai_shield.cost.tracker import CostTracker, MemoryStore, RedisLike
+
+__all__ = [
+    "MODEL_PRICING",
+    "CostTracker",
+    "MemoryStore",
+    "RedisLike",
+    "detect_anomaly",
+    "estimate_cost",
+    "get_model_pricing",
+]

ai_shield/cost/anomaly.py

@@ -0,0 +1,50 @@
+"""Z-score anomaly detection for cost-spike alerts.
+
+1:1 port of `packages/core/src/cost/anomaly.ts`.
+"""
+
+from __future__ import annotations
+
+import math
+
+from ai_shield.types import AnomalyResult
+
+
+def detect_anomaly(
+    samples: list[float],
+    current: float,
+    *,
+    z_threshold: float = 2.5,
+) -> AnomalyResult:
+    """Return AnomalyResult flagging `current` if z-score exceeds threshold.
+
+    Returns is_anomaly=False when fewer than 3 samples or stdDev=0.
+    """
+    n = len(samples)
+    if n < 3:
+        return AnomalyResult(
+            is_anomaly=False,
+            z_score=0.0,
+            current_value=current,
+            mean=0.0,
+            std_dev=0.0,
+        )
+    mean = sum(samples) / n
+    variance = sum((s - mean) ** 2 for s in samples) / n
+    std = math.sqrt(variance)
+    if std == 0.0:
+        return AnomalyResult(
+            is_anomaly=False,
+            z_score=0.0,
+            current_value=current,
+            mean=mean,
+            std_dev=0.0,
+        )
+    z = (current - mean) / std
+    return AnomalyResult(
+        is_anomaly=abs(z) >= z_threshold,
+        z_score=z,
+        current_value=current,
+        mean=mean,
+        std_dev=std,
+    )

ai_shield/cost/pricing.py

@@ -0,0 +1,62 @@
+"""Model pricing table + estimate_cost helper.
+
+1:1 port of `packages/core/src/cost/pricing.ts`.
+USD per 1M tokens. Update as providers change rates.
+"""
+
+from __future__ import annotations
+
+from ai_shield.types import ModelPricing
+
+MODEL_PRICING: dict[str, ModelPricing] = {
+    # OpenAI
+    "gpt-5.2": ModelPricing(input_per_1m=4.50, output_per_1m=18.00),
+    "gpt-5.0": ModelPricing(input_per_1m=3.00, output_per_1m=12.00),
+    "gpt-4.1": ModelPricing(input_per_1m=2.50, output_per_1m=10.00),
+    "gpt-4o": ModelPricing(input_per_1m=2.50, output_per_1m=10.00),
+    "gpt-4o-mini": ModelPricing(input_per_1m=0.15, output_per_1m=0.60),
+    "o3": ModelPricing(input_per_1m=15.00, output_per_1m=60.00),
+    "o4-mini": ModelPricing(input_per_1m=1.10, output_per_1m=4.40),
+    # Anthropic
+    "claude-opus-4-7": ModelPricing(input_per_1m=15.00, output_per_1m=75.00),
+    "claude-opus-4-6": ModelPricing(input_per_1m=15.00, output_per_1m=75.00),
+    "claude-sonnet-4-7": ModelPricing(input_per_1m=3.00, output_per_1m=15.00),
+    "claude-sonnet-4-5": ModelPricing(input_per_1m=3.00, output_per_1m=15.00),
+    "claude-haiku-4-7": ModelPricing(input_per_1m=0.80, output_per_1m=4.00),
+    # Google
+    "gemini-2.5-pro": ModelPricing(input_per_1m=1.25, output_per_1m=5.00),
+    "gemini-2.5-flash": ModelPricing(input_per_1m=0.075, output_per_1m=0.30),
+    "gemini-2.5-flash-lite": ModelPricing(input_per_1m=0.0375, output_per_1m=0.15),
+    # xAI
+    "grok-4": ModelPricing(input_per_1m=5.00, output_per_1m=15.00),
+    # Mistral
+    "mistral-large": ModelPricing(input_per_1m=2.00, output_per_1m=6.00),
+    "mistral-small": ModelPricing(input_per_1m=0.20, output_per_1m=0.60),
+}
+
+
+def get_model_pricing(model: str) -> ModelPricing:
+    """Return pricing for `model`. Exact match → longest-prefix match → fallback.
+
+    Longest-prefix-first ensures `gpt-4o-mini-2024-07-18` matches `gpt-4o-mini`
+    ($0.15) NOT `gpt-4o` ($2.50) — critical for cost-accuracy with versioned
+    snapshots. Insertion-order would otherwise return the first registered
+    prefix (which is shorter for many providers).
+    """
+    if model in MODEL_PRICING:
+        return MODEL_PRICING[model]
+    # Sort keys by length descending so most-specific prefix wins first.
+    for key in sorted(MODEL_PRICING.keys(), key=len, reverse=True):
+        if model.startswith(key):
+            return MODEL_PRICING[key]
+    return MODEL_PRICING["gpt-4o-mini"]
+
+
+def estimate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
+    """Compute USD estimate for a single LLM call."""
+    if input_tokens < 0 or output_tokens < 0:
+        raise ValueError("Token counts must be non-negative")
+    p = get_model_pricing(model)
+    return (input_tokens / 1_000_000) * p.input_per_1m + (
+        output_tokens / 1_000_000
+    ) * p.output_per_1m

ai_shield/cost/tracker.py

@@ -0,0 +1,175 @@
+"""Async cost tracker — soft/hard budgets, in-memory or Redis backend.
+
+1:1 port of `packages/core/src/cost/tracker.ts`.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+from collections import deque
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Protocol
+
+from ai_shield.cost.pricing import estimate_cost
+from ai_shield.types import (
+    BudgetCheckResult,
+    BudgetConfig,
+    BudgetPeriod,
+    CostRecord,
+)
+
+_DEFAULT_MAX_RECORDS = int(os.environ.get("AI_SHIELD_MAX_RECORDS", "10000"))
+
+
+class RedisLike(Protocol):
+    """Minimal Redis interface — `incrbyfloat` + `expire` + `get`."""
+
+    async def incrbyfloat(self, key: str, value: float) -> float: ...
+    async def expire(self, key: str, seconds: int) -> bool: ...
+    async def get(self, key: str) -> str | None: ...
+
+
+@dataclass
+class MemoryStore:
+    """In-process fallback when Redis is not provided."""
+
+    _data: dict[str, float] = field(default_factory=dict)
+    _expires_at: dict[str, float] = field(default_factory=dict)
+    _lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+
+    async def incrbyfloat(self, key: str, value: float) -> float:
+        async with self._lock:
+            self._sweep_expired()
+            current = self._data.get(key, 0.0)
+            new = current + value
+            self._data[key] = new
+            return new
+
+    async def expire(self, key: str, seconds: int) -> bool:
+        async with self._lock:
+            if key not in self._data:
+                return False
+            loop_time = asyncio.get_event_loop().time()
+            self._expires_at[key] = loop_time + seconds
+            return True
+
+    async def get(self, key: str) -> str | None:
+        async with self._lock:
+            self._sweep_expired()
+            v = self._data.get(key)
+            return None if v is None else str(v)
+
+    def _sweep_expired(self) -> None:
+        now = asyncio.get_event_loop().time()
+        expired = [k for k, t in self._expires_at.items() if t <= now]
+        for k in expired:
+            self._data.pop(k, None)
+            self._expires_at.pop(k, None)
+
+
+def _period_seconds(period: BudgetPeriod) -> int:
+    if period == "hourly":
+        return 3600
+    if period == "daily":
+        return 86400
+    if period == "monthly":
+        return 86400 * 31
+    raise ValueError(f"Unknown period: {period!r}")
+
+
+def _period_key(period: BudgetPeriod, now: datetime | None = None) -> str:
+    moment = now or datetime.now(timezone.utc)
+    if period == "hourly":
+        return moment.strftime("%Y%m%d%H")
+    if period == "daily":
+        return moment.strftime("%Y%m%d")
+    if period == "monthly":
+        return moment.strftime("%Y%m")
+    raise ValueError(f"Unknown period: {period!r}")
+
+
+@dataclass
+class CostTracker:
+    """Track LLM spend with optional Redis backend.
+
+    All writes are atomic at the store level (`INCRBYFLOAT` on Redis,
+    `asyncio.Lock` on the in-memory store).
+    """
+
+    budget: BudgetConfig = field(default_factory=BudgetConfig)
+    store: RedisLike | None = None
+    max_records: int = _DEFAULT_MAX_RECORDS
+
+    _records: deque[CostRecord] = field(init=False)
+    _store: RedisLike = field(init=False)
+
+    def __post_init__(self) -> None:
+        self._records = deque(maxlen=self.max_records)
+        self._store = self.store if self.store is not None else MemoryStore()
+
+    @staticmethod
+    def _key(entity_id: str, period: BudgetPeriod) -> str:
+        return f"ai-shield:cost:{entity_id}:{period}:{_period_key(period)}"
+
+    async def record(
+        self,
+        *,
+        entity_id: str,
+        model: str,
+        input_tokens: int,
+        output_tokens: int,
+        actual_usd: float | None = None,
+    ) -> CostRecord:
+        cost = (
+            actual_usd
+            if actual_usd is not None
+            else estimate_cost(
+                model,
+                input_tokens,
+                output_tokens,
+            )
+        )
+        record = CostRecord(
+            entity_id=entity_id,
+            model=model,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            actual_usd=cost,
+            timestamp=datetime.now(timezone.utc).isoformat(),
+        )
+        self._records.append(record)
+
+        key = self._key(entity_id, self.budget.period)
+        await self._store.incrbyfloat(key, cost)
+        await self._store.expire(key, _period_seconds(self.budget.period))
+        return record
+
+    async def check_budget(self, entity_id: str) -> BudgetCheckResult:
+        key = self._key(entity_id, self.budget.period)
+        raw = await self._store.get(key)
+        spend = float(raw) if raw is not None else 0.0
+
+        soft = self.budget.soft_limit_usd
+        hard = self.budget.hard_limit_usd
+
+        soft_exceeded = soft is not None and spend >= soft
+        hard_exceeded = hard is not None and spend >= hard
+
+        return BudgetCheckResult(
+            allowed=not hard_exceeded,
+            current_spend_usd=spend,
+            limit_usd=hard if hard is not None else soft,
+            period=self.budget.period,
+            soft_exceeded=soft_exceeded,
+            hard_exceeded=hard_exceeded,
+        )
+
+    async def get_current_spend(self, entity_id: str) -> float:
+        key = self._key(entity_id, self.budget.period)
+        raw = await self._store.get(key)
+        return float(raw) if raw is not None else 0.0
+
+    def recent_records(self, n: int = 100) -> list[CostRecord]:
+        return list(self._records)[-n:]