211 lines
7.0 KiB
Python
211 lines
7.0 KiB
Python
|
"""Runtime provider / data policy enforcement for AI agents.
|
||
|
|
|
||
|
|
Filters messages and context before they reach the LLM based on:
|
||
|
|
- Sensitive fields (``app.core.sensitive_data.SENSITIVE_FIELDS``)
|
||
|
|
- AI use-case metadata (allowed data categories)
|
||
|
|
- Provider compliance (data residency / allowed data classes)
|
||
|
|
|
||
|
|
This is the enforcement layer that guarantees an agent never sends data it
|
||
|
|
is not permitted to process to a provider that is not approved for it.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import logging
|
||
|
|
import uuid
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
|
|
||
|
|
from app.ai.ai_use_case import AIUseCaseMetadata
|
||
|
|
from app.core.sensitive_data import (
|
||
|
|
SENSITIVE_FIELDS,
|
||
|
|
filter_for_llm_context,
|
||
|
|
get_data_class_for_field,
|
||
|
|
)
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
# Data categories that map to entity types for sensitive-field filtering.
|
||
|
|
_CATEGORY_ENTITY_MAP = {
|
||
|
|
"contact_data": "contact",
|
||
|
|
"email_content": "mail_account",
|
||
|
|
"communication": "mail_account",
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
async def enforce_data_policy(
|
||
|
|
db: AsyncSession,
|
||
|
|
tenant_id: uuid.UUID,
|
||
|
|
messages: list[dict[str, Any]],
|
||
|
|
agent_definition: Any,
|
||
|
|
) -> list[dict[str, Any]]:
|
||
|
|
"""Filter messages/context based on the data policy.
|
||
|
|
|
||
|
|
Steps:
|
||
|
|
1. Remove sensitive fields from any dict content in the messages.
|
||
|
|
2. Check AI use-case metadata for allowed data categories.
|
||
|
|
3. Check provider compliance for data residency requirements.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
db: Async DB session (may be ``None`` in tests / mock mode).
|
||
|
|
tenant_id: Tenant ID for provider lookup.
|
||
|
|
messages: The chat messages to filter.
|
||
|
|
agent_definition: AgentDefinition with ``ai_use_case_metadata``.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A new list of messages with disallowed data removed.
|
||
|
|
"""
|
||
|
|
metadata = AIUseCaseMetadata.from_dict(
|
||
|
|
getattr(agent_definition, "ai_use_case_metadata", None)
|
||
|
|
)
|
||
|
|
|
||
|
|
# Provider compliance (data residency / allowed data classes).
|
||
|
|
compliance: dict[str, Any] | None = None
|
||
|
|
if db is not None and tenant_id is not None:
|
||
|
|
try:
|
||
|
|
from app.ai.llm_client import get_provider_compliance
|
||
|
|
|
||
|
|
compliance = await get_provider_compliance(db, tenant_id)
|
||
|
|
except Exception:
|
||
|
|
logger.debug("Failed to load provider compliance — skipping residency check")
|
||
|
|
|
||
|
|
filtered: list[dict[str, Any]] = []
|
||
|
|
for msg in messages:
|
||
|
|
content = msg.get("content", "")
|
||
|
|
if isinstance(content, dict):
|
||
|
|
content = _filter_dict_content(
|
||
|
|
content, metadata, compliance, agent_definition
|
||
|
|
)
|
||
|
|
elif isinstance(content, list):
|
||
|
|
content = [
|
||
|
|
_filter_dict_content(c, metadata, compliance, agent_definition)
|
||
|
|
if isinstance(c, dict)
|
||
|
|
else c
|
||
|
|
for c in content
|
||
|
|
]
|
||
|
|
new_msg = dict(msg)
|
||
|
|
new_msg["content"] = content
|
||
|
|
filtered.append(new_msg)
|
||
|
|
|
||
|
|
return filtered
|
||
|
|
|
||
|
|
|
||
|
|
def _filter_dict_content(
|
||
|
|
data: dict[str, Any],
|
||
|
|
metadata: AIUseCaseMetadata,
|
||
|
|
compliance: dict[str, Any] | None,
|
||
|
|
agent_definition: Any,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
"""Filter a single dict (entity payload) against the data policy."""
|
||
|
|
# 1. Remove sensitive fields (always blocked from LLM context).
|
||
|
|
result = _strip_sensitive_fields(data)
|
||
|
|
|
||
|
|
# 2. Enforce allowed data categories from AI use-case metadata.
|
||
|
|
if metadata.data_categories:
|
||
|
|
result = _filter_by_allowed_categories(result, metadata.data_categories)
|
||
|
|
|
||
|
|
# 3. Provider compliance — block fields whose data class the provider
|
||
|
|
# is not approved to process.
|
||
|
|
if compliance is not None:
|
||
|
|
result = _filter_by_provider_compliance(result, compliance)
|
||
|
|
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def _strip_sensitive_fields(data: dict[str, Any]) -> dict[str, Any]:
|
||
|
|
"""Recursively remove any key that matches a sensitive field name."""
|
||
|
|
sensitive_names = set()
|
||
|
|
for fields in SENSITIVE_FIELDS.values():
|
||
|
|
sensitive_names |= fields
|
||
|
|
|
||
|
|
result: dict[str, Any] = {}
|
||
|
|
for key, value in data.items():
|
||
|
|
if key in sensitive_names:
|
||
|
|
continue
|
||
|
|
if isinstance(value, dict):
|
||
|
|
result[key] = _strip_sensitive_fields(value)
|
||
|
|
elif isinstance(value, list):
|
||
|
|
result[key] = [
|
||
|
|
_strip_sensitive_fields(v) if isinstance(v, dict) else v
|
||
|
|
for v in value
|
||
|
|
]
|
||
|
|
else:
|
||
|
|
result[key] = value
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def _filter_by_allowed_categories(
|
||
|
|
data: dict[str, Any], allowed_categories: list[str]
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
"""Remove entity-type payloads whose category is not allowed.
|
||
|
|
|
||
|
|
Uses the category→entity mapping to decide whether a dict represents a
|
||
|
|
disallowed entity type. Unknown dicts are kept (fail-open for generic
|
||
|
|
context that has no clear entity type).
|
||
|
|
"""
|
||
|
|
# Determine the entity type of this dict by checking for known keys.
|
||
|
|
entity_type = _guess_entity_type(data)
|
||
|
|
if entity_type is None:
|
||
|
|
return data
|
||
|
|
|
||
|
|
category = _entity_to_category(entity_type)
|
||
|
|
if category is not None and category not in allowed_categories:
|
||
|
|
return {}
|
||
|
|
return data
|
||
|
|
|
||
|
|
|
||
|
|
def _filter_by_provider_compliance(
|
||
|
|
data: dict[str, Any], compliance: dict[str, Any]
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
"""Remove fields whose data class the provider may not process."""
|
||
|
|
allowed_classes = compliance.get("allowed_data_classes") or []
|
||
|
|
if not allowed_classes:
|
||
|
|
return data # No restriction configured (fail-open).
|
||
|
|
|
||
|
|
from app.core.sensitive_data import check_provider_compliance
|
||
|
|
|
||
|
|
result: dict[str, Any] = {}
|
||
|
|
for key, value in data.items():
|
||
|
|
if isinstance(value, dict):
|
||
|
|
result[key] = _filter_by_provider_compliance(value, compliance)
|
||
|
|
continue
|
||
|
|
# Determine data class for this field (best-effort).
|
||
|
|
data_class = _guess_data_class(key, value)
|
||
|
|
if check_provider_compliance(allowed_classes, data_class):
|
||
|
|
result[key] = value
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def _guess_entity_type(data: dict[str, Any]) -> str | None:
|
||
|
|
"""Best-effort guess of the entity type from dict keys."""
|
||
|
|
if any(k in data for k in ("email", "smtp_password", "imap_password")):
|
||
|
|
return "mail_account"
|
||
|
|
if any(k in data for k in ("first_name", "last_name", "company_id")):
|
||
|
|
return "contact"
|
||
|
|
if any(k in data for k in ("secret_key", "encryption_key")):
|
||
|
|
return "system_settings"
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _entity_to_category(entity_type: str) -> str | None:
|
||
|
|
"""Map an entity type to a data category."""
|
||
|
|
for category, entity in _CATEGORY_ENTITY_MAP.items():
|
||
|
|
if entity == entity_type:
|
||
|
|
return category
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _guess_data_class(key: str, value: Any) -> str:
|
||
|
|
"""Best-effort data class for a field (defaults to 'internal')."""
|
||
|
|
# Sensitive field names are always critical.
|
||
|
|
for fields in SENSITIVE_FIELDS.values():
|
||
|
|
if key in fields:
|
||
|
|
return "critical"
|
||
|
|
# Heuristic: values that look like credentials/tokens are critical.
|
||
|
|
if isinstance(value, str) and any(
|
||
|
|
marker in key.lower() for marker in ("password", "token", "secret", "key")
|
||
|
|
):
|
||
|
|
return "critical"
|
||
|
|
return "internal"
|