diff --git a/diagrams/01-system-overview.md b/diagrams/01-system-overview.md new file mode 100644 index 00000000..c4b443b9 --- /dev/null +++ b/diagrams/01-system-overview.md @@ -0,0 +1,63 @@ +# System Overview + +High-level view of Agent Protect architecture showing the main components and how they interact. + +## Component Architecture + +```mermaid +flowchart TB + subgraph client["Client Application"] + agent["AI Agent
(LangChain, LangGraph, etc.)"] + sdk["Agent Control SDK"] + end + + subgraph server["Agent Control Server"] + api["REST API
(FastAPI)"] + engine["Control Engine"] + db[(PostgreSQL)] + end + + subgraph external["External Services"] + luna["Luna-2"] + guardrails["Guardrails AI"] + custom["Custom Plugins"] + end + + agent <--> sdk + sdk <-->|HTTP/JSON| api + api <--> engine + api <--> db + engine -.->|plugin calls| luna + engine -.->|plugin calls| guardrails + engine -.->|plugin calls| custom +``` + +## Package Structure + +```mermaid +flowchart LR + subgraph packages["Python Packages"] + models["agent-control-models
━━━━━━━━━━━━━━━
Shared Pydantic models
for API contracts"] + + engine["agent-control-engine
━━━━━━━━━━━━━━━
Control evaluation logic
Evaluators & Selectors"] + + sdk["agent-control
━━━━━━━━━━━━━━━
Python SDK client
@protect decorator"] + + server["agent-control-server
━━━━━━━━━━━━━━━
FastAPI server
Database layer"] + end + + models --> engine + models --> sdk + models --> server + engine --> server + sdk -.->|HTTP| server +``` + +## Key Interactions + +| From | To | Protocol | Purpose | +|------|-----|----------|---------| +| AI Agent | SDK | In-process | Wrap agent calls with protection | +| SDK | Server | HTTP/JSON | Register agents, evaluate requests | +| Server | Database | SQL | Persist agents, policies, controls | +| Engine | Plugins | In-process | Delegate to external evaluators | diff --git a/diagrams/02-deployment-architecture.md b/diagrams/02-deployment-architecture.md new file mode 100644 index 00000000..2f765486 --- /dev/null +++ b/diagrams/02-deployment-architecture.md @@ -0,0 +1,97 @@ +# Deployment Architecture + +How Agent Protect components are deployed in a production environment. + +## Kubernetes Deployment + +```mermaid +flowchart TB + subgraph k8s["Kubernetes Cluster"] + subgraph ns["agent-control namespace"] + subgraph deploy["Deployment: agent-control-server"] + pod1["Pod"] + pod2["Pod"] + pod3["Pod"] + end + + svc["Service
agent-control-server"] + hpa["HPA
Auto-scaling"] + + subgraph db["StatefulSet or Managed"] + pg[(PostgreSQL)] + end + end + + ing["Ingress Controller"] + end + + subgraph apps["Client Applications"] + app1["Agent App 1"] + app2["Agent App 2"] + app3["Agent App N"] + end + + apps --> ing + ing --> svc + svc --> deploy + deploy --> pg + hpa -.->|scales| deploy +``` + +## Service Communication + +```mermaid +flowchart LR + subgraph external["External"] + client["SDK Client"] + end + + subgraph cluster["Kubernetes"] + ingress["Ingress
:443"] + service["Service
:8000"] + + subgraph pods["Pods"] + uvicorn1["Uvicorn
:8000"] + uvicorn2["Uvicorn
:8000"] + end + + postgres[(PostgreSQL
:5432)] + end + + client -->|HTTPS| ingress + ingress -->|HTTP| service + service -->|Round Robin| pods + pods -->|TCP| postgres +``` + +## Configuration + +```mermaid +flowchart TD + subgraph config["Configuration Sources"] + env["Environment Variables
━━━━━━━━━━━━━━━
DATABASE_URL
API_PREFIX
DEBUG"] + + secrets["Kubernetes Secrets
━━━━━━━━━━━━━━━
DB credentials
API keys"] + + configmap["ConfigMap
━━━━━━━━━━━━━━━
Feature flags
Plugin configs"] + end + + subgraph server["Server Pod"] + app["Agent Control Server"] + end + + env --> app + secrets --> app + configmap --> app +``` + +## Endpoints + +| Path | Method | Description | +|------|--------|-------------| +| `/health` | GET | Liveness/readiness probe | +| `/api/v1/agents/*` | * | Agent management | +| `/api/v1/policies/*` | * | Policy management | +| `/api/v1/control-sets/*` | * | Control set management | +| `/api/v1/controls/*` | * | Control management | +| `/api/v1/evaluation` | POST | Control evaluation | diff --git a/diagrams/03-entity-relationships.md b/diagrams/03-entity-relationships.md new file mode 100644 index 00000000..787a8f8e --- /dev/null +++ b/diagrams/03-entity-relationships.md @@ -0,0 +1,123 @@ +# Entity Relationships + +Database schema showing the hierarchical relationship between Agents, Policies, Control Sets, and Controls. + +## Entity Hierarchy + +```mermaid +flowchart TD + agent["🤖 Agent"] + policy["📋 Policy"] + controlset["📦 Control Set"] + control["⚙️ Control"] + + agent -->|"has one (optional)"| policy + policy -->|"has many"| controlset + controlset -->|"has many"| control +``` + +## Database Schema (ERD) + +```mermaid +erDiagram + agents { + UUID agent_uuid PK "Primary key" + VARCHAR name UK "Unique agent name" + JSONB data "Agent metadata & tools" + INT policy_id FK "Optional policy reference" + TIMESTAMP created_at "Creation timestamp" + } + + policies { + INT id PK "Auto-increment" + VARCHAR name UK "Unique policy name" + } + + control_sets { + INT id PK "Auto-increment" + VARCHAR name UK "Unique control set name" + } + + controls { + INT id PK "Auto-increment" + VARCHAR name UK "Unique control name" + JSONB data "ControlDefinition JSON" + } + + policy_control_sets { + INT policy_id PK,FK + INT control_set_id PK,FK + } + + control_set_controls { + INT control_set_id PK,FK + INT control_id PK,FK + } + + agents ||--o| policies : "belongs to" + policies ||--o{ policy_control_sets : "has" + policy_control_sets }o--|| control_sets : "contains" + control_sets ||--o{ control_set_controls : "has" + control_set_controls }o--|| controls : "contains" +``` + +## Relationship Details + +### Agent → Policy (Many-to-One) +- An Agent can optionally have one Policy assigned +- A Policy can be shared across multiple Agents +- When an Agent has no Policy, it has no active controls + +### Policy ↔ Control Set (Many-to-Many) +- A Policy groups multiple Control Sets together +- A Control Set can be reused across multiple Policies +- Association managed via `policy_control_sets` junction table + +### Control Set ↔ Control (Many-to-Many) +- A Control Set groups multiple atomic Controls +- A Control can be reused across multiple Control Sets +- Association managed via `control_set_controls` junction table + +## Example Configuration + +```mermaid +flowchart LR + subgraph agents["Agents"] + a1["customer-service-bot"] + a2["sales-assistant"] + end + + subgraph policies["Policies"] + p1["production-policy"] + end + + subgraph controlsets["Control Sets"] + cs1["pii-protection"] + cs2["content-safety"] + end + + subgraph controls["Controls"] + c1["ssn-detector"] + c2["email-detector"] + c3["profanity-filter"] + c4["toxicity-check"] + end + + a1 --> p1 + a2 --> p1 + p1 --> cs1 + p1 --> cs2 + cs1 --> c1 + cs1 --> c2 + cs2 --> c3 + cs2 --> c4 +``` + +## Data Traversal + +To get all Controls for an Agent: +``` +Agent → Policy → ControlSets (via junction) → Controls (via junction) +``` + +This multi-hop traversal is performed server-side when processing evaluation requests. diff --git a/diagrams/04-control-definition-model.md b/diagrams/04-control-definition-model.md new file mode 100644 index 00000000..43d1a341 --- /dev/null +++ b/diagrams/04-control-definition-model.md @@ -0,0 +1,175 @@ +# Control Definition Model + +Detailed structure of a Control's configuration stored in the `data` JSONB field. + +## ControlDefinition Structure + +```mermaid +classDiagram + class ControlDefinition { + +description: string | null + +enabled: bool + +applies_to: "llm_call" | "tool_call" + +check_stage: "pre" | "post" + +selector: ControlSelector + +evaluator: ControlEvaluator + +action: ControlAction + +tags: list~string~ + } + + class ControlSelector { + +path: string + } + + class ControlAction { + +decision: "allow" | "deny" | "warn" | "log" + } + + ControlDefinition *-- ControlSelector + ControlDefinition *-- ControlAction + ControlDefinition *-- ControlEvaluator + + note for ControlDefinition "Stored in controls.data as JSONB" +``` + +## Evaluator Types + +```mermaid +classDiagram + class ControlEvaluator { + <> + } + + class RegexControlEvaluator { + +type: "regex" + +config: RegexConfig + } + + class ListControlEvaluator { + +type: "list" + +config: ListConfig + } + + class PluginControlEvaluator { + +type: "plugin" + +config: PluginConfig + } + + class CustomControlEvaluator { + +type: "custom" + +config: CustomConfig + } + + ControlEvaluator <|-- RegexControlEvaluator + ControlEvaluator <|-- ListControlEvaluator + ControlEvaluator <|-- PluginControlEvaluator + ControlEvaluator <|-- CustomControlEvaluator +``` + +## Evaluator Configurations + +```mermaid +classDiagram + class RegexConfig { + +pattern: string + +flags: list~string~ + } + + class ListConfig { + +values: list~string~ + +logic: "any" | "all" + +match_on: "match" | "no_match" + +case_sensitive: bool + } + + class PluginConfig { + +plugin_name: string + +plugin_config: dict + } + + class CustomConfig { + +code: string + +language: string + } + + note for RegexConfig "Uses Google RE2 engine" + note for ListConfig "Supports allowlists and blocklists" + note for PluginConfig "Delegates to external systems" +``` + +## Control Flow Decision Tree + +```mermaid +flowchart TD + start["Incoming Request"] + + check_enabled{"enabled?"} + check_stage{"check_stage
matches request?"} + check_type{"applies_to
matches payload?"} + + select["Select data via
selector.path"] + evaluate["Run evaluator"] + + check_match{"Matched?"} + + action_allow["Action: allow"] + action_deny["Action: deny
(is_safe = false)"] + action_warn["Action: warn"] + action_log["Action: log"] + + skip["Skip control"] + + start --> check_enabled + check_enabled -->|No| skip + check_enabled -->|Yes| check_stage + check_stage -->|No| skip + check_stage -->|Yes| check_type + check_type -->|No| skip + check_type -->|Yes| select + select --> evaluate + evaluate --> check_match + check_match -->|No| skip + check_match -->|Yes| action_allow + check_match -->|Yes| action_deny + check_match -->|Yes| action_warn + check_match -->|Yes| action_log +``` + +## Example Control Definition + +```json +{ + "description": "Block outputs containing SSN", + "enabled": true, + "applies_to": "llm_call", + "check_stage": "post", + "selector": { + "path": "output" + }, + "evaluator": { + "type": "regex", + "config": { + "pattern": "\\b\\d{3}-\\d{2}-\\d{4}\\b", + "flags": [] + } + }, + "action": { + "decision": "deny" + }, + "tags": ["pii", "compliance"] +} +``` + +## Check Stage Semantics + +| Stage | When | Use Case | +|-------|------|----------| +| `pre` | Before LLM/tool execution | Validate inputs, block dangerous requests | +| `post` | After LLM/tool execution | Filter outputs, redact PII | + +## Applies To Semantics + +| Type | Payload | Use Case | +|------|---------|----------| +| `llm_call` | LlmCall (input/output text) | Content moderation, PII detection | +| `tool_call` | ToolCall (tool_name, arguments) | Permission enforcement, input validation | diff --git a/diagrams/05-request-response-models.md b/diagrams/05-request-response-models.md new file mode 100644 index 00000000..c97670fe --- /dev/null +++ b/diagrams/05-request-response-models.md @@ -0,0 +1,205 @@ +# Request & Response Models + +API contracts for the Agent Control Server endpoints. + +## Agent Models + +```mermaid +classDiagram + class Agent { + +agent_id: UUID + +agent_name: string + +agent_description: string | null + +agent_created_at: string | null + +agent_updated_at: string | null + +agent_version: string | null + +agent_metadata: dict | null + } + + class AgentTool { + +tool_name: string + +arguments: dict + +output_schema: dict + } + + class InitAgentRequest { + +agent: Agent + +tools: list~AgentTool~ + } + + class InitAgentResponse { + +created: bool + +controls: list~Control~ + } + + class GetAgentResponse { + +agent: Agent + +tools: list~AgentTool~ + } + + InitAgentRequest *-- Agent + InitAgentRequest *-- AgentTool + InitAgentResponse *-- Control + GetAgentResponse *-- Agent + GetAgentResponse *-- AgentTool +``` + +## Evaluation Models + +```mermaid +classDiagram + class EvaluationRequest { + +agent_uuid: UUID + +check_stage: "pre" | "post" + +payload: LlmCall | ToolCall + } + + class LlmCall { + +input: string | null + +output: string | null + } + + class ToolCall { + +tool_name: string + +arguments: dict + +result: any | null + } + + class EvaluationResponse { + +is_safe: bool + +confidence: float + +matches: list~ControlMatch~ | null + } + + class ControlMatch { + +control_id: int + +control_name: string + +action: string + +result: EvaluatorResult + } + + class EvaluatorResult { + +matched: bool + +confidence: float + +message: string | null + +metadata: dict | null + } + + EvaluationRequest *-- LlmCall + EvaluationRequest *-- ToolCall + EvaluationResponse *-- ControlMatch + ControlMatch *-- EvaluatorResult +``` + +## Policy & Control Management Models + +```mermaid +classDiagram + class CreatePolicyRequest { + +name: string + } + class CreatePolicyResponse { + +policy_id: int + } + + class CreateControlSetRequest { + +name: string + } + class CreateControlSetResponse { + +control_set_id: int + } + + class CreateControlRequest { + +name: string + } + class CreateControlResponse { + +control_id: int + } + + class SetControlDataRequest { + +data: ControlDefinition + } + class SetControlDataResponse { + +success: bool + } + + class AssocResponse { + +success: bool + } + + class SetPolicyResponse { + +success: bool + +old_policy_id: int | null + } +``` + +## API Request/Response Flow + +```mermaid +sequenceDiagram + participant Client + participant Server + participant DB + + Note over Client,DB: Agent Registration + Client->>Server: POST /agents/initAgent
InitAgentRequest + Server->>DB: Upsert Agent + Server-->>Client: InitAgentResponse + + Note over Client,DB: Policy Assignment + Client->>Server: POST /agents/{id}/policy/{policy_id} + Server->>DB: Update Agent.policy_id + Server-->>Client: SetPolicyResponse + + Note over Client,DB: Evaluation + Client->>Server: POST /evaluation
EvaluationRequest + Server->>DB: Fetch Agent's Controls + Server->>Server: Run Control Engine + Server-->>Client: EvaluationResponse +``` + +## Payload Discrimination + +The `EvaluationRequest.payload` field is a discriminated union: + +```mermaid +flowchart LR + payload["payload"] + + check{"Has tool_name
field?"} + + llm["LlmCall
━━━━━━━━
input: string
output: string"] + tool["ToolCall
━━━━━━━━
tool_name: string
arguments: dict
result: any"] + + payload --> check + check -->|No| llm + check -->|Yes| tool +``` + +## Control Model (API Response) + +```mermaid +classDiagram + class Control { + +id: int + +name: string + +control: dict + } + + class ControlSet { + +id: int + +name: string + +controls: list~Control~ + } + + class Policy { + +id: int + +name: string + +control_sets: list~ControlSet~ + } + + Policy *-- ControlSet + ControlSet *-- Control + + note for Control "control field contains ControlDefinition" +``` diff --git a/diagrams/06-evaluation-flow.md b/diagrams/06-evaluation-flow.md new file mode 100644 index 00000000..d9033306 --- /dev/null +++ b/diagrams/06-evaluation-flow.md @@ -0,0 +1,202 @@ +# Control Evaluation Flow + +How the Control Engine processes evaluation requests and applies controls. + +## High-Level Flow + +```mermaid +flowchart TD + request["EvaluationRequest
━━━━━━━━━━━━━━━
agent_uuid
check_stage
payload"] + + subgraph engine["Control Engine"] + filter["Filter Applicable
Controls"] + loop["For Each Control"] + + subgraph process["Process Control"] + select["1. Select Data"] + evaluate["2. Evaluate"] + act["3. Record Match"] + end + end + + response["EvaluationResponse
━━━━━━━━━━━━━━━
is_safe
confidence
matches"] + + request --> filter + filter --> loop + loop --> process + process --> response +``` + +## Control Filtering Logic + +```mermaid +flowchart TD + controls["All Agent Controls"] + + c1{"enabled = true?"} + c2{"check_stage
matches request?"} + c3{"applies_to
matches payload?"} + + applicable["Applicable Controls"] + skipped["Skipped"] + + controls --> c1 + c1 -->|No| skipped + c1 -->|Yes| c2 + c2 -->|No| skipped + c2 -->|Yes| c3 + c3 -->|No| skipped + c3 -->|Yes| applicable +``` + +## Data Selection + +The Selector extracts data from the payload using a path expression: + +```mermaid +flowchart LR + subgraph payload["LlmCall Payload"] + input["input: 'Hello world'"] + output["output: 'Response text'"] + end + + selector["Selector
path: 'output'"] + + data["Selected Data:
'Response text'"] + + payload --> selector + selector --> data +``` + +## Evaluator Processing + +```mermaid +flowchart TD + data["Selected Data"] + + type{"Evaluator Type?"} + + subgraph regex["Regex Evaluator"] + re_compile["Compile Pattern
(RE2 Engine)"] + re_match["Search for Match"] + end + + subgraph list["List Evaluator"] + list_norm["Normalize Input"] + list_match["Match Against Values"] + list_logic["Apply Logic
(any/all)"] + end + + subgraph plugin["Plugin Evaluator"] + plugin_load["Load Plugin"] + plugin_call["Call Plugin.evaluate()"] + end + + result["EvaluatorResult
━━━━━━━━━━━━━
matched: bool
confidence: float
message: string
metadata: dict"] + + data --> type + type -->|"regex"| regex + type -->|"list"| list + type -->|"plugin"| plugin + regex --> result + list --> result + plugin --> result +``` + +## Result Aggregation + +```mermaid +flowchart TD + matches["All Control Matches"] + + check{"Any match with
action = 'deny'?"} + + safe["is_safe = true"] + unsafe["is_safe = false"] + + response["EvaluationResponse"] + + matches --> check + check -->|No| safe + check -->|Yes| unsafe + safe --> response + unsafe --> response +``` + +## Complete Sequence + +```mermaid +sequenceDiagram + participant Client as SDK Client + participant API as Server API + participant DB as Database + participant Engine as Control Engine + participant Eval as Evaluator + + Client->>API: POST /evaluation + API->>DB: Get Agent by UUID + API->>DB: Get Controls via Policy→ControlSets + DB-->>API: List of Controls + + API->>Engine: Process(request, controls) + + loop For each applicable control + Engine->>Engine: Select data from payload + Engine->>Eval: Evaluate(data) + Eval-->>Engine: EvaluatorResult + + alt If matched + Engine->>Engine: Record ControlMatch + alt If action = deny + Engine->>Engine: Set is_safe = false + end + end + end + + Engine-->>API: EvaluationResponse + API-->>Client: JSON Response +``` + +## List Evaluator Logic Detail + +```mermaid +flowchart TD + input["Input Values"] + config["Config Values"] + + match_check["Match each input
against config values"] + + matches["Matched Values"] + + logic{"logic setting?"} + + any_check{"Any matched?"} + all_check{"All matched?"} + + condition_met["condition_met = true"] + condition_not["condition_met = false"] + + match_on{"match_on setting?"} + + result_match["matched = condition_met"] + result_nomatch["matched = !condition_met"] + + input --> match_check + config --> match_check + match_check --> matches + matches --> logic + + logic -->|"any"| any_check + logic -->|"all"| all_check + + any_check -->|Yes| condition_met + any_check -->|No| condition_not + all_check -->|Yes| condition_met + all_check -->|No| condition_not + + condition_met --> match_on + condition_not --> match_on + + match_on -->|"match"| result_match + match_on -->|"no_match"| result_nomatch +``` diff --git a/diagrams/07-agent-lifecycle.md b/diagrams/07-agent-lifecycle.md new file mode 100644 index 00000000..50b4b5ec --- /dev/null +++ b/diagrams/07-agent-lifecycle.md @@ -0,0 +1,184 @@ +# Agent Lifecycle + +How agents are registered, configured, and protected throughout their lifecycle. + +## Lifecycle Overview + +```mermaid +stateDiagram-v2 + [*] --> Unregistered + + Unregistered --> Registered: Register Agent + Registered --> Protected: Assign Policy + Protected --> Registered: Remove Policy + Protected --> Protected: Update Policy + Registered --> Registered: Update Tools + + note right of Unregistered: Agent exists in code only + note right of Registered: Agent known to server,\nno active controls + note right of Protected: Agent has policy,\ncontrols are active +``` + +## Registration Flow + +```mermaid +sequenceDiagram + participant App as Application + participant SDK as Agent Control SDK + participant Server as Server + participant DB as Database + + App->>SDK: Initialize agent + + SDK->>Server: POST /agents/initAgent
{agent, tools} + + Server->>DB: Check if agent name exists + + alt Agent is new + DB-->>Server: Not found + Server->>DB: INSERT agent record + Server->>Server: Generate tool schema + Server-->>SDK: {created: true, controls: []} + else Agent exists (same UUID) + DB-->>Server: Found + Server->>DB: UPDATE tools if changed + Server->>DB: Fetch current policy controls + Server-->>SDK: {created: false, controls: [...]} + else Agent exists (different UUID) + DB-->>Server: Found with different UUID + Server-->>SDK: 409 Conflict + end + + SDK-->>App: Registration result +``` + +## Policy Assignment Flow + +```mermaid +sequenceDiagram + participant Admin as Administrator + participant Server as Server + participant DB as Database + + Note over Admin,DB: Setup Phase (One-time) + Admin->>Server: PUT /controls {name} + Server->>DB: Create control + Admin->>Server: PUT /controls/{id}/data {definition} + Server->>DB: Store control config + + Admin->>Server: PUT /control-sets {name} + Server->>DB: Create control set + Admin->>Server: POST /control-sets/{id}/controls/{control_id} + Server->>DB: Link control to set + + Admin->>Server: PUT /policies {name} + Server->>DB: Create policy + Admin->>Server: POST /policies/{id}/control_sets/{set_id} + Server->>DB: Link control set to policy + + Note over Admin,DB: Assignment Phase + Admin->>Server: POST /agents/{agent_id}/policy/{policy_id} + Server->>DB: Update agent.policy_id + Server-->>Admin: {success: true} +``` + +## Agent Data Model + +```mermaid +classDiagram + class AgentRecord { + +agent_uuid: UUID + +name: string + +policy_id: int | null + +data: AgentData + +created_at: datetime + } + + class AgentData { + +agent_metadata: dict + +tools: list~AgentVersionedTool~ + +agent_schema: dict | null + } + + class AgentVersionedTool { + +version: int + +tool: AgentTool + } + + class AgentTool { + +tool_name: string + +arguments: dict + +output_schema: dict + } + + AgentRecord *-- AgentData : "data (JSONB)" + AgentData *-- AgentVersionedTool + AgentVersionedTool *-- AgentTool + + note for AgentRecord "Database table: agents" + note for AgentData "Stored as JSONB in data column" +``` + +## Tool Versioning + +When tools are updated, the system tracks changes: + +```mermaid +flowchart LR + subgraph v0["Version 0"] + t0["search_kb
━━━━━━━
args: {query}"] + end + + subgraph v1["Version 1 (Updated)"] + t1["search_kb
━━━━━━━
args: {query, limit}"] + end + + v0 -->|"Tool schema changed"| v1 +``` + +## Agent States + +| State | policy_id | Controls | Evaluation Behavior | +|-------|-----------|----------|---------------------| +| **Unregistered** | N/A | None | Cannot evaluate (404) | +| **Registered** | `null` | None | Always returns `is_safe: true` | +| **Protected** | `` | Active | Applies all policy controls | + +## Complete Setup Example + +```mermaid +flowchart TD + subgraph step1["1. Create Controls"] + c1["Create 'ssn-detector'"] + c2["Configure regex pattern"] + end + + subgraph step2["2. Create Control Set"] + cs1["Create 'pii-protection'"] + cs2["Add ssn-detector to set"] + end + + subgraph step3["3. Create Policy"] + p1["Create 'production-policy'"] + p2["Add pii-protection to policy"] + end + + subgraph step4["4. Register Agent"] + a1["Register 'my-agent'"] + a2["Provide tool schemas"] + end + + subgraph step5["5. Assign Policy"] + assign["Assign production-policy
to my-agent"] + end + + subgraph step6["6. Protected!"] + eval["Evaluation requests now
apply ssn-detector control"] + end + + step1 --> step2 + step2 --> step3 + step3 --> step4 + step4 --> step5 + step5 --> step6 +``` diff --git a/diagrams/08-protect-decorator-flow.md b/diagrams/08-protect-decorator-flow.md new file mode 100644 index 00000000..7b72316f --- /dev/null +++ b/diagrams/08-protect-decorator-flow.md @@ -0,0 +1,230 @@ +# @protect Decorator Flow + +How the SDK's `@protect` decorator intercepts function calls and enforces controls. + +## Decorator Overview + +```mermaid +flowchart LR + subgraph code["Your Code"] + func["@protect('step-id', input='msg')
async def process(msg):"] + end + + subgraph decorator["Decorator Behavior"] + extract["Extract mapped
parameters"] + pre["PRE check"] + execute["Execute function"] + post["POST check"] + end + + call["Function Call"] --> decorator + decorator --> result["Return Value"] +``` + +## Execution Sequence + +```mermaid +sequenceDiagram + participant App as Application + participant Dec as @protect Decorator + participant Func as Your Function + participant Engine as Control Engine + + App->>Dec: Call decorated function + + Dec->>Dec: Bind arguments to parameters + Dec->>Dec: Extract mapped data sources + + Note over Dec,Engine: PRE-execution check + Dec->>Engine: Evaluate(data, stage="pre") + Engine-->>Dec: EvaluationResponse + + alt is_safe = false + Dec-->>App: Raise exception or return error + end + + Dec->>Func: Execute original function + Func-->>Dec: Return value + + Note over Dec,Engine: POST-execution check + Dec->>Dec: Add return value to data + Dec->>Engine: Evaluate(data, stage="post") + Engine-->>Dec: EvaluationResponse + + alt is_safe = false + Dec-->>App: Raise exception or return error + end + + Dec-->>App: Return result +``` + +## Data Source Mapping + +```mermaid +flowchart TD + subgraph decorator["Decorator Definition"] + dec["@protect('step-id',
input='user_msg',
context='ctx',
output='response')"] + end + + subgraph function["Function Signature"] + func["def process(
user_msg: str,
ctx: dict
) -> str:"] + end + + subgraph mapping["Extracted Data"] + m1["input → user_msg parameter"] + m2["context → ctx parameter"] + m3["output → return value"] + end + + decorator --> mapping + function --> mapping +``` + +## Parameter Binding Process + +```mermaid +flowchart LR + subgraph call["Function Call"] + args["process('Hello', {'user': 'A'})"] + end + + subgraph bind["Signature Binding"] + sig["signature(process)"] + bound["bind(*args, **kwargs)"] + defaults["apply_defaults()"] + end + + subgraph result["Bound Arguments"] + ba["user_msg: 'Hello'
ctx: {'user': 'A'}"] + end + + call --> bind + bind --> result +``` + +## Data Available at Each Stage + +```mermaid +flowchart TD + subgraph pre["PRE Stage"] + pre_data["Available Data:
━━━━━━━━━━━━
✓ All mapped parameters
✗ Return value (not yet)"] + end + + exec["Function Executes"] + + subgraph post["POST Stage"] + post_data["Available Data:
━━━━━━━━━━━━
✓ All mapped parameters
✓ Return value (if mapped)"] + end + + pre --> exec + exec --> post +``` + +## What You CAN and CANNOT Access + +```mermaid +flowchart LR + subgraph can["✓ Accessible"] + p1["Function parameters"] + p2["Return value"] + p3["Default values"] + p4["Nested data in params"] + end + + subgraph cannot["✗ Not Accessible"] + n1["Local variables"] + n2["Outer scope variables"] + n3["Intermediate results"] + n4["Unmapped parameters"] + end +``` + +## Common Patterns + +### Pattern 1: Input Validation + +```mermaid +flowchart LR + input["User Input"] --> pre["PRE Check"] + pre -->|Safe| func["Process"] + pre -->|Unsafe| block["Block"] + func --> output["Output"] +``` + +### Pattern 2: Output Filtering + +```mermaid +flowchart LR + input["Input"] --> func["Process"] + func --> post["POST Check"] + post -->|Safe| output["Return Output"] + post -->|PII Detected| redact["Redact/Block"] +``` + +### Pattern 3: Full Pipeline + +```mermaid +flowchart LR + input["Input"] + pre["PRE Check"] + func["Process"] + post["POST Check"] + output["Output"] + + input --> pre + pre -->|Pass| func + func --> post + post -->|Pass| output + + pre -->|Fail| blocked1["Blocked"] + post -->|Fail| blocked2["Blocked"] +``` + +## LangGraph Integration Example + +```mermaid +flowchart TD + subgraph graph["LangGraph Workflow"] + start["Start"] + + node1["@protect('input-check')
Input Node"] + node2["LLM Node"] + node3["@protect('output-check')
Output Node"] + + finish["End"] + end + + start --> node1 + node1 --> node2 + node2 --> node3 + node3 --> finish +``` + +## Error Handling + +```mermaid +flowchart TD + check["Control Check"] + + result{"Result?"} + + safe["Continue execution"] + + deny["Action: deny"] + warn["Action: warn"] + log["Action: log"] + + exception["Raise RuleViolation"] + log_warning["Log warning,
continue"] + log_info["Log match,
continue"] + + check --> result + result -->|is_safe| safe + result -->|deny match| deny + result -->|warn match| warn + result -->|log match| log + + deny --> exception + warn --> log_warning + log --> log_info +``` diff --git a/diagrams/09-plugin-system.md b/diagrams/09-plugin-system.md new file mode 100644 index 00000000..61607609 --- /dev/null +++ b/diagrams/09-plugin-system.md @@ -0,0 +1,236 @@ +# Plugin System Architecture + +How external evaluators integrate with Agent Control through the plugin system. + +## Plugin Overview + +```mermaid +flowchart TB + subgraph engine["Control Engine"] + eval_factory["Evaluator Factory"] + + subgraph builtin["Built-in Evaluators"] + regex["Regex"] + list["List"] + end + + plugin_wrapper["Plugin Wrapper"] + end + + subgraph plugins["Plugin Registry"] + registry["Plugin Registry"] + + subgraph registered["Registered Plugins"] + luna["galileo-luna2"] + guardrails["guardrails-ai"] + custom["custom-plugin"] + end + end + + subgraph external["External Services"] + luna_api["Luna-2 API"] + gr_api["Guardrails API"] + end + + eval_factory --> builtin + eval_factory --> plugin_wrapper + plugin_wrapper --> registry + registry --> registered + luna -.-> luna_api + guardrails -.-> gr_api +``` + +## Plugin Class Hierarchy + +```mermaid +classDiagram + class PluginMetadata { + +name: string + +version: string + +description: string + +requires_api_key: bool + +timeout_ms: int + +config_schema: dict | null + } + + class PluginEvaluator { + <> + +metadata: PluginMetadata + +evaluate(data, config)*: EvaluatorResult + +get_timeout_seconds(config): float + } + + class Luna2Plugin { + +metadata: PluginMetadata + +evaluate(data, config): EvaluatorResult + } + + class GuardrailsPlugin { + +metadata: PluginMetadata + +evaluate(data, config): EvaluatorResult + } + + class CustomPlugin { + +metadata: PluginMetadata + +evaluate(data, config): EvaluatorResult + } + + PluginEvaluator <|-- Luna2Plugin + PluginEvaluator <|-- GuardrailsPlugin + PluginEvaluator <|-- CustomPlugin + PluginEvaluator *-- PluginMetadata +``` + +## Plugin Registration Flow + +```mermaid +sequenceDiagram + participant Plugin as Plugin Module + participant Registry as Plugin Registry + participant Engine as Control Engine + + Note over Plugin,Registry: Registration (at import time) + Plugin->>Registry: @register_plugin decorator + Registry->>Registry: Store plugin class by name + + Note over Engine,Registry: Usage (at evaluation time) + Engine->>Registry: get_plugin("plugin-name") + Registry-->>Engine: PluginClass + Engine->>Engine: Instantiate plugin + Engine->>Plugin: evaluate(data, config) + Plugin-->>Engine: EvaluatorResult +``` + +## Plugin Configuration in ControlDefinition + +```mermaid +flowchart LR + subgraph control["ControlDefinition"] + evaluator["evaluator:
━━━━━━━━
type: 'plugin'
config: PluginConfig"] + end + + subgraph config["PluginConfig"] + pc["plugin_name: 'galileo-luna2'
plugin_config:
model: 'luna-2-large'
threshold: 0.8
timeout_ms: 5000"] + end + + control --> config +``` + +## Plugin Evaluation Flow + +```mermaid +flowchart TD + request["Evaluation Request"] + + select["Select Data
from Payload"] + + wrapper["Plugin Wrapper"] + + subgraph plugin_exec["Plugin Execution"] + load["Load Plugin
from Registry"] + instantiate["Create Plugin
Instance"] + call["Call evaluate()"] + end + + subgraph external["External Call"] + api["External API
(Luna-2, etc.)"] + end + + result["EvaluatorResult"] + + error["Error Result
(if plugin fails)"] + + request --> select + select --> wrapper + wrapper --> plugin_exec + plugin_exec --> external + external --> result + plugin_exec -.->|exception| error +``` + +## Creating a Custom Plugin + +```mermaid +flowchart TD + subgraph steps["Plugin Development Steps"] + s1["1. Extend PluginEvaluator"] + s2["2. Define metadata"] + s3["3. Implement evaluate()"] + s4["4. Register with decorator"] + end + + subgraph structure["Plugin Structure"] + meta["PluginMetadata
━━━━━━━━━━━━
name
version
description
timeout_ms"] + + eval["evaluate(data, config)
━━━━━━━━━━━━
→ EvaluatorResult"] + end + + steps --> structure +``` + +## Plugin Error Handling + +```mermaid +flowchart TD + call["Plugin.evaluate()"] + + result{"Success?"} + + success["Return EvaluatorResult"] + + error["Exception caught"] + + error_result["EvaluatorResult
━━━━━━━━━━━━
matched: false
confidence: 0.0
message: error details
metadata: {error: ...}"] + + call --> result + result -->|Yes| success + result -->|Exception| error + error --> error_result +``` + +## Available Plugins + +| Plugin | Description | Requires API Key | +|--------|-------------|------------------| +| `galileo-luna2` | Galileo's Luna-2 safety model | Yes | +| `guardrails-ai` | Guardrails AI validators | Depends | +| Custom | User-defined plugins | Varies | + +## Plugin Config Schema + +Plugins can define a JSON Schema for their configuration: + +```mermaid +flowchart LR + subgraph plugin["Plugin Definition"] + meta["metadata.config_schema:
{
'model': {type: 'string'},
'threshold': {type: 'number'}
}"] + end + + subgraph usage["Control Config"] + config["plugin_config:
{
'model': 'luna-2-large',
'threshold': 0.8
}"] + end + + plugin -.->|validates| usage +``` + +## Timeout Handling + +```mermaid +flowchart LR + config["plugin_config"] + meta["metadata.timeout_ms"] + + check{"timeout_ms
in config?"} + + use_config["Use config timeout"] + use_default["Use metadata default"] + + seconds["Convert to seconds
for API calls"] + + config --> check + meta --> check + check -->|Yes| use_config + check -->|No| use_default + use_config --> seconds + use_default --> seconds +``` diff --git a/diagrams/10-sdk-architecture.md b/diagrams/10-sdk-architecture.md new file mode 100644 index 00000000..889fd047 --- /dev/null +++ b/diagrams/10-sdk-architecture.md @@ -0,0 +1,238 @@ +# SDK Architecture + +Structure of the Python SDK (`agent-control` package) and how its modules are organized. + +## Module Overview + +```mermaid +flowchart TB + subgraph sdk["agent_control package"] + init["__init__.py
━━━━━━━━━━━━━
Public API
Convenience functions"] + + client["client.py
━━━━━━━━━━━━━
HTTP client
Connection management"] + + subgraph operations["Operation Modules"] + agents["agents.py"] + policies["policies.py"] + control_sets["control_sets.py"] + controls["controls.py"] + evaluation["evaluation.py"] + end + + decorator["tool_decorator.py
━━━━━━━━━━━━━
@protect decorator"] + + subgraph plugins_pkg["plugins/"] + plugins_init["__init__.py"] + base["base.py"] + registry["registry.py"] + optional["optional/"] + end + end + + init --> client + init --> operations + init --> decorator + decorator --> evaluation +``` + +## Module Responsibilities + +```mermaid +flowchart LR + subgraph modules["SDK Modules"] + direction TB + + m1["client.py
━━━━━━━━━
AgentProtectClient
HTTP session
Health check"] + + m2["agents.py
━━━━━━━━━
register_agent
get_agent"] + + m3["policies.py
━━━━━━━━━
create_policy
add_control_set
remove_control_set
list_control_sets"] + + m4["control_sets.py
━━━━━━━━━
create_control_set
add_control
remove_control
list_controls"] + + m5["controls.py
━━━━━━━━━
create_control
get_control_data
set_control_data"] + + m6["evaluation.py
━━━━━━━━━
evaluate"] + end + + subgraph endpoints["Server Endpoints"] + e1["/agents/*"] + e2["/policies/*"] + e3["/control-sets/*"] + e4["/controls/*"] + e5["/evaluation"] + end + + m2 --> e1 + m3 --> e2 + m4 --> e3 + m5 --> e4 + m6 --> e5 +``` + +## Client Usage Pattern + +```mermaid +sequenceDiagram + participant App as Application + participant SDK as agent_control + participant Client as AgentProtectClient + participant Server as Server + + App->>SDK: import agent_control + + App->>Client: async with AgentProtectClient() as client: + Client->>Client: Create HTTP session + + App->>SDK: agent_control.agents.register_agent(client, ...) + SDK->>Client: POST /agents/initAgent + Client->>Server: HTTP Request + Server-->>Client: Response + Client-->>SDK: Parsed result + SDK-->>App: Return data + + App->>Client: Exit context manager + Client->>Client: Close HTTP session +``` + +## Public API Exports + +```mermaid +flowchart TD + subgraph exports["agent_control exports"] + subgraph functions["Functions"] + f1["init()"] + f2["current_agent()"] + f3["get_agent()"] + f4["protect()"] + end + + subgraph classes["Classes"] + c1["AgentProtectClient"] + end + + subgraph modules_exp["Modules"] + m1["agents"] + m2["policies"] + m3["control_sets"] + m4["controls"] + m5["evaluation"] + end + + subgraph models_exp["Models (re-exported)"] + mo1["Agent"] + mo2["LlmCall"] + mo3["ToolCall"] + mo4["EvaluationRequest"] + mo5["EvaluationResponse"] + end + end +``` + +## Two Usage Patterns + +### Pattern 1: Module-First (Recommended) + +```mermaid +flowchart LR + subgraph code["Code"] + import["import agent_control"] + use["agent_control.policies.create_policy(client, name)"] + end +``` + +### Pattern 2: Direct Import + +```mermaid +flowchart LR + subgraph code["Code"] + import["from agent_control import policies"] + use["policies.create_policy(client, name)"] + end +``` + +## Plugin Module Structure + +```mermaid +flowchart TB + subgraph plugins["agent_control.plugins"] + init["__init__.py
━━━━━━━━━━━━━
get_plugin()
list_plugins()
register_plugin()"] + + base["base.py
━━━━━━━━━━━━━
PluginMetadata
PluginEvaluator"] + + registry["registry.py
━━━━━━━━━━━━━
Plugin storage
Discovery logic"] + + subgraph optional_pkg["optional/"] + luna["luna2.py"] + guardrails["guardrails.py"] + end + end + + init --> registry + registry --> base + optional_pkg -.->|lazy load| registry +``` + +## Dependency Graph + +```mermaid +flowchart BT + models["agent-control-models"] + engine["agent-control-engine"] + sdk["agent-control (SDK)"] + server["agent-control-server"] + + sdk --> models + engine --> models + server --> models + server --> engine + + sdk -.->|HTTP| server +``` + +## Initialization Flow + +```mermaid +flowchart TD + start["Application Start"] + + import_sdk["import agent_control"] + + init["agent_control.init(
agent_name='my-agent',
agent_id=uuid,
tools=[...],
server_url='http://...'
)"] + + register["Register agent with server"] + + store["Store current agent reference"] + + ready["Ready to use @protect"] + + start --> import_sdk + import_sdk --> init + init --> register + register --> store + store --> ready +``` + +## Error Handling + +```mermaid +flowchart TD + call["SDK API Call"] + + http["HTTP Request"] + + check{"Response
Status?"} + + success["Return parsed data"] + + client_error["400-499:
Raise HTTPError"] + server_error["500-599:
Raise HTTPError"] + network_error["Connection Error:
Raise Exception"] + + call --> http + http --> check + check -->|2xx| success + check -->|4xx| client_error + check -->|5xx| server_error + http -.->|Network| network_error +``` diff --git a/diagrams/README.md b/diagrams/README.md new file mode 100644 index 00000000..4fb0b500 --- /dev/null +++ b/diagrams/README.md @@ -0,0 +1,42 @@ +# Agent Protect Architecture Diagrams + +This folder contains architecture diagrams for the Agent Protect system. All diagrams use Mermaid syntax and render natively in GitHub. + +## Diagram Index + +### High-Level Architecture +| Diagram | Description | Audience | +|---------|-------------|----------| +| [01-system-overview.md](./01-system-overview.md) | Bird's eye view of the entire system | Everyone | +| [02-deployment-architecture.md](./02-deployment-architecture.md) | How components are deployed | DevOps, Backend | + +### Data Models +| Diagram | Description | Audience | +|---------|-------------|----------| +| [03-entity-relationships.md](./03-entity-relationships.md) | Database schema and relationships | Backend | +| [04-control-definition-model.md](./04-control-definition-model.md) | Control configuration structure | Backend, SDK Users | +| [05-request-response-models.md](./05-request-response-models.md) | API request/response schemas | SDK Users | + +### Flows & Sequences +| Diagram | Description | Audience | +|---------|-------------|----------| +| [06-evaluation-flow.md](./06-evaluation-flow.md) | How controls are evaluated | Backend | +| [07-agent-lifecycle.md](./07-agent-lifecycle.md) | Agent registration and policy assignment | SDK Users | +| [08-protect-decorator-flow.md](./08-protect-decorator-flow.md) | How @protect decorator works | SDK Users | + +### Component Deep Dives +| Diagram | Description | Audience | +|---------|-------------|----------| +| [09-plugin-system.md](./09-plugin-system.md) | External evaluator integration | Backend, Plugin Authors | +| [10-sdk-architecture.md](./10-sdk-architecture.md) | Python SDK module structure | SDK Users | + +## Viewing Diagrams + +### GitHub +Mermaid diagrams render automatically when viewing markdown files on GitHub. + +### VS Code +Install the "Markdown Preview Mermaid Support" extension. + +### Local +Use the Mermaid CLI or online editor at https://mermaid.live