Spaces:
Build error
Build error
File size: 3,033 Bytes
7a529e5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """
Core data models for AEGIS-Env.
Contains Pydantic schemas and Enums utilized across the environment.
"""
from typing import Optional, List
from enum import Enum
from pydantic import BaseModel, ConfigDict
class WorkerRole(str, Enum):
"""Worker role types in the enterprise fleet."""
SEC_OPS = "sec-ops" # Security operations - monitors access, permissions
DATA_ENG = "data-eng" # Data engineering - SQL, ETL, pipelines
API_DEV = "api-dev" # API development - REST/GraphQL endpoints
CODE_REVIEWER = "code-reviewer" # Code review - security, quality
DESIGNER = "designer" # UI/UX design - design tokens, accessibility
PLANNER = "planner" # Project planning - task decomposition
CUSTOMER_SUPPORT = "customer-support" # Customer facing - PII exposure risk
ADMIN = "admin" # Admin operations - system commands
GENERAL_DEV = "general-dev" # General development - default role
class Decision(str, Enum):
"""Actions the oversight agent can take."""
ALLOW = "ALLOW"
BLOCK = "BLOCK"
ESCALATE = "ESCALATE"
class ViolationType(str, Enum):
"""Categorized policy violations."""
NONE = "none"
PII_LEAK = "pii_leak"
PROMPT_INJECTION = "prompt_injection"
HALLUCINATED_FEATURE = "hallucinated_feature"
UNSAFE_CODE = "unsafe_code"
DATA_EXFILTRATION = "data_exfiltration_intent"
AUTHORITY_ESCALATION = "authority_escalation"
COMPOUND_VIOLATION = "compound_violation"
OVERSEER_MANIPULATION = "overseer_manipulation"
class AEGISAction(BaseModel):
"""The action schema output by the oversight agent."""
decision: Decision
confidence: float
violation_type: ViolationType
policy_rule_cited: Optional[str] = None
evidence_quote: Optional[str] = None
explanation: str
remediation: Optional[str] = None
model_config = ConfigDict(extra="forbid")
class AEGISObservation(BaseModel):
"""The observation schema provided to the oversight agent."""
# BUG-1 / Layer-1 FIX: All fields have defaults to prevent 422 crashes
# on partial observations sent by the training loop.
worker_id: str = "WORKER-1"
worker_role: WorkerRole = WorkerRole.GENERAL_DEV
turn_number: int = 1
worker_cot_trace: str = ""
worker_output: str = ""
policy_ruleset: List[dict] = []
state_buffer: List[str] = []
scenario_type: str = "SINGLE_TURN"
turns_remaining: int = 15
api_call_log: List[dict] = []
db_query_trace: List[str] = []
memory_context: str = ""
success: bool = False
model_config = ConfigDict(extra="forbid")
class AEGISState(BaseModel):
"""The internal state representation of the environment episode."""
episode_id: str
scenario_id: str
curriculum_level: int
step_count: int = 0
ground_truth: dict
world_db_state: dict
memory_ledger_size: int = 0
total_reward: float = 0.0
earliest_detectable_turn: Optional[int] = None
detection_turn: Optional[int] = None
model_config = ConfigDict(extra="forbid")
|