AI Agent Policy Engine

A production-grade embedded policy engine that sits between an LLM agent's proposed actions and actual execution. Every action passes through the policy layer, which inspects and enforces data classification rules before anything happens.

The Problem

In large enterprise AI systems, autonomous agents have access to tools that can read, write, and transmit data. Without a guard layer, agents can inadvertently:

• Exfiltrate PII (emails, SSNs, phone numbers) to external APIs
• Leak credentials (AWS keys, GitHub tokens, JWTs, database passwords) in log output or LLM context
• Send sensitive internal data to unauthorized endpoints
• Violate compliance requirements (GDPR, HIPAA, SOC2) by exposing regulated data

Traditional approaches — training prompts, output filters, or post-hoc auditing — are brittle, easy to bypass, and provide no enforcement guarantees.

What This Solves

This policy engine provides a deterministic, low-latency enforcement layer that classifies every piece of data flowing through an AI agent before it reaches any tool or external system. It answers questions like:

• "Does this email body contain an SSN?" → Deny the send action
• "Does this log message contain an AWS secret key?" → Redact before writing
• "Does this API payload contain a credit card number?" → Flag for human review
• "Is this content clean?" → Allow with no latency penalty on repeat content

Architecture

Agent proposes action
       │
       ▼
┌─────────────────────┐
│  PolicyEngine       │  ← Embedded in agent process
│  ┌───────────────┐  │
│  │  Content Hash │  │  ← Fast path: cache hit → 50µs decision
│  │  Cache (LRU)  │  │
│  └───────┬───────┘  │
│          ▼ (miss)   │
│  ┌───────────────┐  │
│  │  Detectors    │  │  ← Parallel regex + entropy detection
│  │  (async)      │  │
│  └───────┬───────┘  │
│          ▼          │
│  ┌───────────────┐  │
│  │  Rules Engine  │  │  ← YAML-defined rules, priority-ordered
│  └───────┬───────┘  │
│          ▼          │
│  ┌───────────────┐  │
│  │  Decision     │  │  ← allow / deny / flag / redact
│  └───────────────┘  │
└─────────────────────┘
       │
       ▼
  Action executes (or blocked)

Dependencies & Integrations

Core Libraries

Library	Purpose
Pydantic v2	Configuration management — all config models are Pydantic models with validation. EngineConfig.from_yaml() loads and validates policy definitions at startup.
PyYAML	Policy definition format — rules and detector configs are authored in YAML files, making them accessible to non-developer security/compliance teams.

Observability Stack

Library	Purpose
Prometheus Client	Metrics — exposes policy_engine_decisions_total{effect} and policy_engine_cache_{hits,misses}_total counters for dashboards and alerting.
OpenTelemetry API + SDK	Distributed tracing — emits one root span per evaluate() call with child spans per policy and per detector. Enables latency breakdowns in Jaeger, Datadog, etc.
Structured Audit Logging	JSON-per-line audit log records every evaluation with action, decision, findings, latency, and environment context. Ship via filebeat/fluentd to SIEM.

Detection Engine

Module	What It Detects
PII Detector	Emails, SSNs, credit card numbers, US phone numbers, IP addresses
Secret Detector	AWS access keys, GitHub tokens, JWTs, Slack tokens, Stripe keys, private key headers, generic API keys
Entropy Detector	High-Shannon-entropy strings — catches unknown or custom secret formats by measuring randomness

All detectors use pre-compiled regex (compiled once at import time) and run in parallel via asyncio.gather. No ML models, no external API calls, no cold starts.

Caching Layer

Component	Purpose
SHA-256 Content Hash	Deterministic canonicalization + hashing of action payloads. Same content → same hash → cached decision.
Async LRU Cache	Thread-safe LRU cache with configurable TTL. Target: >80% hit rate for templated/repeated content.

Project Structure

policy_engine/
├── __init__.py              # Public API: PolicyEngine, EngineConfig, Action, Context, Decision
├── engine.py                # Core evaluation loop, cache orchestration
├── config.py                # Pydantic v2 config models, YAML loader
├── exceptions.py            # Custom error hierarchy
├── models/
│   ├── action.py            # Action schema (type, resource, payload)
│   ├── context.py           # Evaluation context (subject, environment)
│   └── decision.py          # Decision model + Finding model
├── policies/
│   ├── base.py              # PolicyType ABC
│   ├── registry.py          # Action-pattern → policy chain mapping
│   └── data_classification/
│       ├── policy.py        # DataClassificationPolicy implementation
│       ├── rules.py         # Priority-ordered rule engine with redaction
│       └── detectors/
│           ├── base.py      # Detector ABC
│           ├── pii.py       # PII regex patterns
│           ├── secrets.py   # Secret/key regex patterns
│           ├── entropy.py   # Shannon entropy-based detection
│           └── composite.py # Parallel dispatcher + deduplication
├── cache/
│   ├── content_hash.py      # SHA-256 content hashing
│   └── lru.py               # Async-safe LRU cache with TTL
├── observability/
│   ├── metrics.py           # Prometheus counters
│   ├── logging.py           # Structured JSON audit log
│   └── tracing.py           # OpenTelemetry spans
└── tests/
    ├── test_detectors.py    # Unit tests for each detector
    ├── test_rules.py        # Rule engine evaluation tests
    ├── test_cache.py        # LRU cache + content hash tests
    ├── test_config.py       # Config loading and validation
    ├── test_integration.py  # Full pipeline integration tests
    └── test_benchmarks.py   # Performance benchmarks

Getting Your Agent Ready

1. Install

pip install policy-engine

For development:

pip install -e ".[dev]"

2. Create a policy file

policies.yaml:

data_classification:
  enabled: true
  detectors:
    - pii
    - secrets
    - entropy

  rules:
    - name: block_critical
      detectors: [pii, secrets]
      conditions:
        severity: CRITICAL
      effect: deny
      reason: "Critical data detected — blocked"
      priority: 100

    - name: flag_high
      detectors: [pii, secrets]
      conditions:
        severity: HIGH
      effect: flag
      reason: "Sensitive data detected — review required"
      priority: 80

    - name: allow_default
      effect: allow
      reason: "Passed all policy checks"
      priority: 0

  cache:
    ttl_seconds: 60
    max_entries: 10000

observability:
  metrics_enabled: true
  audit_log_enabled: true
  tracing_enabled: false
  audit_log_path: "/var/log/policy-engine/audit.log"

3. Integrate into your agent

import asyncio
from policy_engine import PolicyEngine, EngineConfig, Action, Context

# Load config once at startup
config = EngineConfig.from_yaml("policies.yaml")
engine = PolicyEngine(config)

# In your agent loop, wrap every tool call:
async def agent_execute(tool_name: str, payload: dict, user_id: str) -> str:
    decision = await engine.evaluate(
        Action(type="tool_call", resource=tool_name, payload=payload),
        Context(subject={"id": user_id}, environment="production"),
    )

    if decision.effect == "deny":
        return f"Blocked: {decision.reason}"
    if decision.effect == "flag":
        logger.warning(f"Flagged action: {decision}")
    if decision.effect == "redact":
        payload = decision.redacted_payload

    return await actual_tool_executor(tool_name, payload)

4. (Optional) Set up monitoring

from prometheus_client import start_http_server
start_http_server(8000)  # Exposes metrics at /metrics

Add alerts on:

• rate(policy_engine_decisions_total{effect="deny"}[5m]) > 0 — actions being blocked
• rate(policy_engine_cache_hits_total[5m]) / rate(policy_engine_cache_misses_total[5m]) < 3 — cache health

5. Run tests

pytest tests/ -v

Performance Targets

Scenario	Latency Target	Mechanism
Cache hit	<50µs p50, <200µs p99	LRU dict lookup
Cache miss (clean)	<2ms p50, <10ms p99	Parallel regex on KB-scale text
Throughput (single process)	10K-20K eval/s	Async event loop + caching

Extending

Add a custom detector:

from policy_engine.policies.data_classification.detectors.base import Detector
from policy_engine.models.decision import Finding

class InternalIdDetector(Detector):
    name = "internal_ids"

    async def detect(self, text: str) -> list[Finding]:
        # Custom detection logic
        ...

Register it by adding "internal_ids" to detectors in your policy YAML.

Add a new policy type (permissions, budget, approval):

from policy_engine.policies.base import PolicyType

class BudgetPolicy(PolicyType):
    name = "budget"

    async def evaluate(self, action, context) -> Decision:
        ...

The PolicyRegistry evaluates policies in priority order — any DENY short-circuits the chain.