From 0f7248b3e03453defac5b013e07af331cc25c45b Mon Sep 17 00:00:00 2001 From: Lev Neiman Date: Wed, 3 Dec 2025 19:59:15 -0800 Subject: [PATCH 1/2] add Diagrams --- diagrams/01-system-overview.md | 72 +++++++ diagrams/02-deployment-architecture.md | 108 ++++++++++ diagrams/03-entity-relationships.md | 133 ++++++++++++ diagrams/04-control-definition-model.md | 179 ++++++++++++++++ diagrams/05-request-response-models.md | 214 +++++++++++++++++++ diagrams/06-evaluation-flow.md | 221 ++++++++++++++++++++ diagrams/07-agent-lifecycle.md | 193 +++++++++++++++++ diagrams/08-protect-decorator-flow.md | 261 +++++++++++++++++++++++ diagrams/09-plugin-system.md | 258 +++++++++++++++++++++++ diagrams/10-sdk-architecture.md | 266 ++++++++++++++++++++++++ diagrams/README.md | 42 ++++ 11 files changed, 1947 insertions(+) create mode 100644 diagrams/01-system-overview.md create mode 100644 diagrams/02-deployment-architecture.md create mode 100644 diagrams/03-entity-relationships.md create mode 100644 diagrams/04-control-definition-model.md create mode 100644 diagrams/05-request-response-models.md create mode 100644 diagrams/06-evaluation-flow.md create mode 100644 diagrams/07-agent-lifecycle.md create mode 100644 diagrams/08-protect-decorator-flow.md create mode 100644 diagrams/09-plugin-system.md create mode 100644 diagrams/10-sdk-architecture.md create mode 100644 diagrams/README.md diff --git a/diagrams/01-system-overview.md b/diagrams/01-system-overview.md new file mode 100644 index 00000000..c0cf6b2c --- /dev/null +++ b/diagrams/01-system-overview.md @@ -0,0 +1,72 @@ +# 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 + + style client fill:#e1f5fe,stroke:#01579b + style server fill:#f3e5f5,stroke:#4a148c + style external fill:#fff3e0,stroke:#e65100 +``` + +## 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 + + style models fill:#c8e6c9,stroke:#2e7d32 + style engine fill:#bbdefb,stroke:#1565c0 + style sdk fill:#fff9c4,stroke:#f9a825 + style server fill:#f8bbd9,stroke:#c2185b +``` + +## 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..2336bd9b --- /dev/null +++ b/diagrams/02-deployment-architecture.md @@ -0,0 +1,108 @@ +# 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 + + style k8s fill:#e3f2fd,stroke:#1565c0 + style ns fill:#fce4ec,stroke:#c2185b + style deploy fill:#f3e5f5,stroke:#7b1fa2 + style db fill:#e8f5e9,stroke:#388e3c +``` + +## 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 + + style external fill:#fff3e0,stroke:#ef6c00 + style cluster fill:#e8eaf6,stroke:#3949ab +``` + +## 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 + + style config fill:#fff8e1,stroke:#ff8f00 + style server fill:#e1f5fe,stroke:#0277bd +``` + +## 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..f62af075 --- /dev/null +++ b/diagrams/03-entity-relationships.md @@ -0,0 +1,133 @@ +# 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 + + style agent fill:#e3f2fd,stroke:#1565c0 + style policy fill:#f3e5f5,stroke:#7b1fa2 + style controlset fill:#e8f5e9,stroke:#388e3c + style control fill:#fff3e0,stroke:#ef6c00 +``` + +## 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 + + style agents fill:#e3f2fd,stroke:#1565c0 + style policies fill:#f3e5f5,stroke:#7b1fa2 + style controlsets fill:#e8f5e9,stroke:#388e3c + style controls fill:#fff3e0,stroke:#ef6c00 +``` + +## 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..57bfb1f7 --- /dev/null +++ b/diagrams/04-control-definition-model.md @@ -0,0 +1,179 @@ +# 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 + + style action_deny fill:#ffcdd2,stroke:#c62828 + style action_warn fill:#fff9c4,stroke:#f9a825 + style skip fill:#e0e0e0,stroke:#757575 +``` + +## 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..979ffc8e --- /dev/null +++ b/diagrams/05-request-response-models.md @@ -0,0 +1,214 @@ +# 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 + + rect rgb(227, 242, 253) + Note over Client,DB: Agent Registration + Client->>Server: POST /agents/initAgent
InitAgentRequest + Server->>DB: Upsert Agent + Server-->>Client: InitAgentResponse + end + + rect rgb(243, 229, 245) + Note over Client,DB: Policy Assignment + Client->>Server: POST /agents/{id}/policy/{policy_id} + Server->>DB: Update Agent.policy_id + Server-->>Client: SetPolicyResponse + end + + rect rgb(255, 243, 224) + Note over Client,DB: Evaluation + Client->>Server: POST /evaluation
EvaluationRequest + Server->>DB: Fetch Agent's Controls + Server->>Server: Run Control Engine + Server-->>Client: EvaluationResponse + end +``` + +## 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 + + style llm fill:#e8f5e9,stroke:#388e3c + style tool fill:#fff3e0,stroke:#ef6c00 +``` + +## 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..2f5421f5 --- /dev/null +++ b/diagrams/06-evaluation-flow.md @@ -0,0 +1,221 @@ +# 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 + + style request fill:#e3f2fd,stroke:#1565c0 + style engine fill:#f5f5f5,stroke:#616161 + style response fill:#e8f5e9,stroke:#388e3c +``` + +## 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 + + style applicable fill:#c8e6c9,stroke:#2e7d32 + style skipped fill:#ffcdd2,stroke:#c62828 +``` + +## 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 + + style payload fill:#e3f2fd,stroke:#1565c0 + style data fill:#e8f5e9,stroke:#388e3c +``` + +## 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 + + style data fill:#e3f2fd,stroke:#1565c0 + style result fill:#fff3e0,stroke:#ef6c00 +``` + +## 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 + + style safe fill:#c8e6c9,stroke:#2e7d32 + style unsafe fill:#ffcdd2,stroke:#c62828 +``` + +## 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 + + style condition_met fill:#c8e6c9,stroke:#2e7d32 + style condition_not fill:#ffcdd2,stroke:#c62828 +``` diff --git a/diagrams/07-agent-lifecycle.md b/diagrams/07-agent-lifecycle.md new file mode 100644 index 00000000..58a3f5cc --- /dev/null +++ b/diagrams/07-agent-lifecycle.md @@ -0,0 +1,193 @@ +# 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 + + rect rgb(243, 229, 245) + 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 + end + + rect rgb(227, 242, 253) + 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} + end +``` + +## 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 + + style v0 fill:#e0e0e0,stroke:#757575 + style v1 fill:#c8e6c9,stroke:#2e7d32 +``` + +## 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 + + style step6 fill:#c8e6c9,stroke:#2e7d32 +``` diff --git a/diagrams/08-protect-decorator-flow.md b/diagrams/08-protect-decorator-flow.md new file mode 100644 index 00000000..fd50cca1 --- /dev/null +++ b/diagrams/08-protect-decorator-flow.md @@ -0,0 +1,261 @@ +# @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"] + + style code fill:#e3f2fd,stroke:#1565c0 + style decorator fill:#f3e5f5,stroke:#7b1fa2 +``` + +## 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 + + rect rgb(255, 243, 224) + 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 + end + + Dec->>Func: Execute original function + Func-->>Dec: Return value + + rect rgb(232, 245, 233) + 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 + 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 + + style mapping fill:#e8f5e9,stroke:#388e3c +``` + +## 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 + + style result fill:#fff3e0,stroke:#ef6c00 +``` + +## 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 + + style pre fill:#fff3e0,stroke:#ef6c00 + style post fill:#e8f5e9,stroke:#388e3c +``` + +## 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 + + style can fill:#c8e6c9,stroke:#2e7d32 + style cannot fill:#ffcdd2,stroke:#c62828 +``` + +## 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"] + + style block fill:#ffcdd2,stroke:#c62828 +``` + +### 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"] + + style redact fill:#fff3e0,stroke:#ef6c00 +``` + +### 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"] + + style blocked1 fill:#ffcdd2,stroke:#c62828 + style blocked2 fill:#ffcdd2,stroke:#c62828 +``` + +## 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 + + style node1 fill:#e3f2fd,stroke:#1565c0 + style node3 fill:#e3f2fd,stroke:#1565c0 +``` + +## 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 + + style exception fill:#ffcdd2,stroke:#c62828 + style log_warning fill:#fff3e0,stroke:#ef6c00 + style safe fill:#c8e6c9,stroke:#2e7d32 +``` diff --git a/diagrams/09-plugin-system.md b/diagrams/09-plugin-system.md new file mode 100644 index 00000000..f646e5d6 --- /dev/null +++ b/diagrams/09-plugin-system.md @@ -0,0 +1,258 @@ +# 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 + + style engine fill:#e3f2fd,stroke:#1565c0 + style plugins fill:#f3e5f5,stroke:#7b1fa2 + style external fill:#fff3e0,stroke:#ef6c00 +``` + +## 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 + + rect rgb(243, 229, 245) + Note over Plugin,Registry: Registration (at import time) + Plugin->>Registry: @register_plugin decorator + Registry->>Registry: Store plugin class by name + end + + rect rgb(227, 242, 253) + 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 + end +``` + +## 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 + + style config fill:#fff3e0,stroke:#ef6c00 +``` + +## 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 + + style result fill:#c8e6c9,stroke:#2e7d32 + style error fill:#ffcdd2,stroke:#c62828 +``` + +## 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 + + style steps fill:#e3f2fd,stroke:#1565c0 + style structure fill:#f3e5f5,stroke:#7b1fa2 +``` + +## 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 + + style success fill:#c8e6c9,stroke:#2e7d32 + style error_result fill:#fff3e0,stroke:#ef6c00 +``` + +## 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 + + style plugin fill:#e3f2fd,stroke:#1565c0 + style usage fill:#e8f5e9,stroke:#388e3c +``` + +## 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..d9ddd816 --- /dev/null +++ b/diagrams/10-sdk-architecture.md @@ -0,0 +1,266 @@ +# 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 + + style sdk fill:#e3f2fd,stroke:#1565c0 + style operations fill:#f3e5f5,stroke:#7b1fa2 + style plugins_pkg fill:#fff3e0,stroke:#ef6c00 +``` + +## 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 + + style modules fill:#e8f5e9,stroke:#388e3c + style endpoints fill:#e3f2fd,stroke:#1565c0 +``` + +## 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 + + style exports fill:#f5f5f5,stroke:#616161 +``` + +## 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 + + style code fill:#c8e6c9,stroke:#2e7d32 +``` + +### Pattern 2: Direct Import + +```mermaid +flowchart LR + subgraph code["Code"] + import["from agent_control import policies"] + use["policies.create_policy(client, name)"] + end + + style code fill:#e3f2fd,stroke:#1565c0 +``` + +## 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 + + style plugins fill:#fff3e0,stroke:#ef6c00 + style optional_pkg fill:#ffe0b2,stroke:#f57c00 +``` + +## 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 + + style models fill:#c8e6c9,stroke:#2e7d32 + style sdk fill:#fff3e0,stroke:#ef6c00 + style engine fill:#e3f2fd,stroke:#1565c0 + style server fill:#f3e5f5,stroke:#7b1fa2 +``` + +## 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 + + style ready fill:#c8e6c9,stroke:#2e7d32 +``` + +## 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 + + style success fill:#c8e6c9,stroke:#2e7d32 + style client_error fill:#fff3e0,stroke:#ef6c00 + style server_error fill:#ffcdd2,stroke:#c62828 + style network_error fill:#ffcdd2,stroke:#c62828 +``` 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 From 090c83e214ac3c85fb73e40c3aa92b9ae0edb20b Mon Sep 17 00:00:00 2001 From: Lev Neiman Date: Wed, 3 Dec 2025 20:06:44 -0800 Subject: [PATCH 2/2] update style --- diagrams/01-system-overview.md | 9 ------- diagrams/02-deployment-architecture.md | 11 -------- diagrams/03-entity-relationships.md | 10 ------- diagrams/04-control-definition-model.md | 4 --- diagrams/05-request-response-models.md | 15 +++-------- diagrams/06-evaluation-flow.md | 19 -------------- diagrams/07-agent-lifecycle.md | 13 ++------- diagrams/08-protect-decorator-flow.md | 35 ++----------------------- diagrams/09-plugin-system.md | 26 ++---------------- diagrams/10-sdk-architecture.md | 28 -------------------- 10 files changed, 9 insertions(+), 161 deletions(-) diff --git a/diagrams/01-system-overview.md b/diagrams/01-system-overview.md index c0cf6b2c..c4b443b9 100644 --- a/diagrams/01-system-overview.md +++ b/diagrams/01-system-overview.md @@ -30,10 +30,6 @@ flowchart TB engine -.->|plugin calls| luna engine -.->|plugin calls| guardrails engine -.->|plugin calls| custom - - style client fill:#e1f5fe,stroke:#01579b - style server fill:#f3e5f5,stroke:#4a148c - style external fill:#fff3e0,stroke:#e65100 ``` ## Package Structure @@ -55,11 +51,6 @@ flowchart LR models --> server engine --> server sdk -.->|HTTP| server - - style models fill:#c8e6c9,stroke:#2e7d32 - style engine fill:#bbdefb,stroke:#1565c0 - style sdk fill:#fff9c4,stroke:#f9a825 - style server fill:#f8bbd9,stroke:#c2185b ``` ## Key Interactions diff --git a/diagrams/02-deployment-architecture.md b/diagrams/02-deployment-architecture.md index 2336bd9b..2f765486 100644 --- a/diagrams/02-deployment-architecture.md +++ b/diagrams/02-deployment-architecture.md @@ -36,11 +36,6 @@ flowchart TB svc --> deploy deploy --> pg hpa -.->|scales| deploy - - style k8s fill:#e3f2fd,stroke:#1565c0 - style ns fill:#fce4ec,stroke:#c2185b - style deploy fill:#f3e5f5,stroke:#7b1fa2 - style db fill:#e8f5e9,stroke:#388e3c ``` ## Service Communication @@ -67,9 +62,6 @@ flowchart LR ingress -->|HTTP| service service -->|Round Robin| pods pods -->|TCP| postgres - - style external fill:#fff3e0,stroke:#ef6c00 - style cluster fill:#e8eaf6,stroke:#3949ab ``` ## Configuration @@ -91,9 +83,6 @@ flowchart TD env --> app secrets --> app configmap --> app - - style config fill:#fff8e1,stroke:#ff8f00 - style server fill:#e1f5fe,stroke:#0277bd ``` ## Endpoints diff --git a/diagrams/03-entity-relationships.md b/diagrams/03-entity-relationships.md index f62af075..787a8f8e 100644 --- a/diagrams/03-entity-relationships.md +++ b/diagrams/03-entity-relationships.md @@ -14,11 +14,6 @@ flowchart TD agent -->|"has one (optional)"| policy policy -->|"has many"| controlset controlset -->|"has many"| control - - style agent fill:#e3f2fd,stroke:#1565c0 - style policy fill:#f3e5f5,stroke:#7b1fa2 - style controlset fill:#e8f5e9,stroke:#388e3c - style control fill:#fff3e0,stroke:#ef6c00 ``` ## Database Schema (ERD) @@ -116,11 +111,6 @@ flowchart LR cs1 --> c2 cs2 --> c3 cs2 --> c4 - - style agents fill:#e3f2fd,stroke:#1565c0 - style policies fill:#f3e5f5,stroke:#7b1fa2 - style controlsets fill:#e8f5e9,stroke:#388e3c - style controls fill:#fff3e0,stroke:#ef6c00 ``` ## Data Traversal diff --git a/diagrams/04-control-definition-model.md b/diagrams/04-control-definition-model.md index 57bfb1f7..43d1a341 100644 --- a/diagrams/04-control-definition-model.md +++ b/diagrams/04-control-definition-model.md @@ -133,10 +133,6 @@ flowchart TD check_match -->|Yes| action_deny check_match -->|Yes| action_warn check_match -->|Yes| action_log - - style action_deny fill:#ffcdd2,stroke:#c62828 - style action_warn fill:#fff9c4,stroke:#f9a825 - style skip fill:#e0e0e0,stroke:#757575 ``` ## Example Control Definition diff --git a/diagrams/05-request-response-models.md b/diagrams/05-request-response-models.md index 979ffc8e..c97670fe 100644 --- a/diagrams/05-request-response-models.md +++ b/diagrams/05-request-response-models.md @@ -141,27 +141,21 @@ sequenceDiagram participant Server participant DB - rect rgb(227, 242, 253) - Note over Client,DB: Agent Registration + Note over Client,DB: Agent Registration Client->>Server: POST /agents/initAgent
InitAgentRequest Server->>DB: Upsert Agent Server-->>Client: InitAgentResponse - end - rect rgb(243, 229, 245) - Note over Client,DB: Policy Assignment + Note over Client,DB: Policy Assignment Client->>Server: POST /agents/{id}/policy/{policy_id} Server->>DB: Update Agent.policy_id Server-->>Client: SetPolicyResponse - end - rect rgb(255, 243, 224) - Note over Client,DB: Evaluation + Note over Client,DB: Evaluation Client->>Server: POST /evaluation
EvaluationRequest Server->>DB: Fetch Agent's Controls Server->>Server: Run Control Engine Server-->>Client: EvaluationResponse - end ``` ## Payload Discrimination @@ -180,9 +174,6 @@ flowchart LR payload --> check check -->|No| llm check -->|Yes| tool - - style llm fill:#e8f5e9,stroke:#388e3c - style tool fill:#fff3e0,stroke:#ef6c00 ``` ## Control Model (API Response) diff --git a/diagrams/06-evaluation-flow.md b/diagrams/06-evaluation-flow.md index 2f5421f5..d9033306 100644 --- a/diagrams/06-evaluation-flow.md +++ b/diagrams/06-evaluation-flow.md @@ -25,10 +25,6 @@ flowchart TD filter --> loop loop --> process process --> response - - style request fill:#e3f2fd,stroke:#1565c0 - style engine fill:#f5f5f5,stroke:#616161 - style response fill:#e8f5e9,stroke:#388e3c ``` ## Control Filtering Logic @@ -51,9 +47,6 @@ flowchart TD c2 -->|Yes| c3 c3 -->|No| skipped c3 -->|Yes| applicable - - style applicable fill:#c8e6c9,stroke:#2e7d32 - style skipped fill:#ffcdd2,stroke:#c62828 ``` ## Data Selection @@ -73,9 +66,6 @@ flowchart LR payload --> selector selector --> data - - style payload fill:#e3f2fd,stroke:#1565c0 - style data fill:#e8f5e9,stroke:#388e3c ``` ## Evaluator Processing @@ -111,9 +101,6 @@ flowchart TD regex --> result list --> result plugin --> result - - style data fill:#e3f2fd,stroke:#1565c0 - style result fill:#fff3e0,stroke:#ef6c00 ``` ## Result Aggregation @@ -134,9 +121,6 @@ flowchart TD check -->|Yes| unsafe safe --> response unsafe --> response - - style safe fill:#c8e6c9,stroke:#2e7d32 - style unsafe fill:#ffcdd2,stroke:#c62828 ``` ## Complete Sequence @@ -215,7 +199,4 @@ flowchart TD match_on -->|"match"| result_match match_on -->|"no_match"| result_nomatch - - style condition_met fill:#c8e6c9,stroke:#2e7d32 - style condition_not fill:#ffcdd2,stroke:#c62828 ``` diff --git a/diagrams/07-agent-lifecycle.md b/diagrams/07-agent-lifecycle.md index 58a3f5cc..50b4b5ec 100644 --- a/diagrams/07-agent-lifecycle.md +++ b/diagrams/07-agent-lifecycle.md @@ -60,8 +60,7 @@ sequenceDiagram participant Server as Server participant DB as Database - rect rgb(243, 229, 245) - Note over Admin,DB: Setup Phase (One-time) + Note over Admin,DB: Setup Phase (One-time) Admin->>Server: PUT /controls {name} Server->>DB: Create control Admin->>Server: PUT /controls/{id}/data {definition} @@ -76,14 +75,11 @@ sequenceDiagram Server->>DB: Create policy Admin->>Server: POST /policies/{id}/control_sets/{set_id} Server->>DB: Link control set to policy - end - rect rgb(227, 242, 253) - Note over Admin,DB: Assignment Phase + 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} - end ``` ## Agent Data Model @@ -138,9 +134,6 @@ flowchart LR end v0 -->|"Tool schema changed"| v1 - - style v0 fill:#e0e0e0,stroke:#757575 - style v1 fill:#c8e6c9,stroke:#2e7d32 ``` ## Agent States @@ -188,6 +181,4 @@ flowchart TD step3 --> step4 step4 --> step5 step5 --> step6 - - style step6 fill:#c8e6c9,stroke:#2e7d32 ``` diff --git a/diagrams/08-protect-decorator-flow.md b/diagrams/08-protect-decorator-flow.md index fd50cca1..7b72316f 100644 --- a/diagrams/08-protect-decorator-flow.md +++ b/diagrams/08-protect-decorator-flow.md @@ -19,9 +19,6 @@ flowchart LR call["Function Call"] --> decorator decorator --> result["Return Value"] - - style code fill:#e3f2fd,stroke:#1565c0 - style decorator fill:#f3e5f5,stroke:#7b1fa2 ``` ## Execution Sequence @@ -38,21 +35,18 @@ sequenceDiagram Dec->>Dec: Bind arguments to parameters Dec->>Dec: Extract mapped data sources - rect rgb(255, 243, 224) - Note over Dec,Engine: PRE-execution check + 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 - end Dec->>Func: Execute original function Func-->>Dec: Return value - rect rgb(232, 245, 233) - Note over Dec,Engine: POST-execution check + Note over Dec,Engine: POST-execution check Dec->>Dec: Add return value to data Dec->>Engine: Evaluate(data, stage="post") Engine-->>Dec: EvaluationResponse @@ -60,7 +54,6 @@ sequenceDiagram alt is_safe = false Dec-->>App: Raise exception or return error end - end Dec-->>App: Return result ``` @@ -85,8 +78,6 @@ flowchart TD decorator --> mapping function --> mapping - - style mapping fill:#e8f5e9,stroke:#388e3c ``` ## Parameter Binding Process @@ -109,8 +100,6 @@ flowchart LR call --> bind bind --> result - - style result fill:#fff3e0,stroke:#ef6c00 ``` ## Data Available at Each Stage @@ -129,9 +118,6 @@ flowchart TD pre --> exec exec --> post - - style pre fill:#fff3e0,stroke:#ef6c00 - style post fill:#e8f5e9,stroke:#388e3c ``` ## What You CAN and CANNOT Access @@ -151,9 +137,6 @@ flowchart LR n3["Intermediate results"] n4["Unmapped parameters"] end - - style can fill:#c8e6c9,stroke:#2e7d32 - style cannot fill:#ffcdd2,stroke:#c62828 ``` ## Common Patterns @@ -166,8 +149,6 @@ flowchart LR pre -->|Safe| func["Process"] pre -->|Unsafe| block["Block"] func --> output["Output"] - - style block fill:#ffcdd2,stroke:#c62828 ``` ### Pattern 2: Output Filtering @@ -178,8 +159,6 @@ flowchart LR func --> post["POST Check"] post -->|Safe| output["Return Output"] post -->|PII Detected| redact["Redact/Block"] - - style redact fill:#fff3e0,stroke:#ef6c00 ``` ### Pattern 3: Full Pipeline @@ -199,9 +178,6 @@ flowchart LR pre -->|Fail| blocked1["Blocked"] post -->|Fail| blocked2["Blocked"] - - style blocked1 fill:#ffcdd2,stroke:#c62828 - style blocked2 fill:#ffcdd2,stroke:#c62828 ``` ## LangGraph Integration Example @@ -222,9 +198,6 @@ flowchart TD node1 --> node2 node2 --> node3 node3 --> finish - - style node1 fill:#e3f2fd,stroke:#1565c0 - style node3 fill:#e3f2fd,stroke:#1565c0 ``` ## Error Handling @@ -254,8 +227,4 @@ flowchart TD deny --> exception warn --> log_warning log --> log_info - - style exception fill:#ffcdd2,stroke:#c62828 - style log_warning fill:#fff3e0,stroke:#ef6c00 - style safe fill:#c8e6c9,stroke:#2e7d32 ``` diff --git a/diagrams/09-plugin-system.md b/diagrams/09-plugin-system.md index f646e5d6..61607609 100644 --- a/diagrams/09-plugin-system.md +++ b/diagrams/09-plugin-system.md @@ -38,10 +38,6 @@ flowchart TB registry --> registered luna -.-> luna_api guardrails -.-> gr_api - - style engine fill:#e3f2fd,stroke:#1565c0 - style plugins fill:#f3e5f5,stroke:#7b1fa2 - style external fill:#fff3e0,stroke:#ef6c00 ``` ## Plugin Class Hierarchy @@ -93,20 +89,16 @@ sequenceDiagram participant Registry as Plugin Registry participant Engine as Control Engine - rect rgb(243, 229, 245) - Note over Plugin,Registry: Registration (at import time) + Note over Plugin,Registry: Registration (at import time) Plugin->>Registry: @register_plugin decorator Registry->>Registry: Store plugin class by name - end - rect rgb(227, 242, 253) - Note over Engine,Registry: Usage (at evaluation time) + 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 - end ``` ## Plugin Configuration in ControlDefinition @@ -122,8 +114,6 @@ flowchart LR end control --> config - - style config fill:#fff3e0,stroke:#ef6c00 ``` ## Plugin Evaluation Flow @@ -156,9 +146,6 @@ flowchart TD plugin_exec --> external external --> result plugin_exec -.->|exception| error - - style result fill:#c8e6c9,stroke:#2e7d32 - style error fill:#ffcdd2,stroke:#c62828 ``` ## Creating a Custom Plugin @@ -179,9 +166,6 @@ flowchart TD end steps --> structure - - style steps fill:#e3f2fd,stroke:#1565c0 - style structure fill:#f3e5f5,stroke:#7b1fa2 ``` ## Plugin Error Handling @@ -202,9 +186,6 @@ flowchart TD result -->|Yes| success result -->|Exception| error error --> error_result - - style success fill:#c8e6c9,stroke:#2e7d32 - style error_result fill:#fff3e0,stroke:#ef6c00 ``` ## Available Plugins @@ -230,9 +211,6 @@ flowchart LR end plugin -.->|validates| usage - - style plugin fill:#e3f2fd,stroke:#1565c0 - style usage fill:#e8f5e9,stroke:#388e3c ``` ## Timeout Handling diff --git a/diagrams/10-sdk-architecture.md b/diagrams/10-sdk-architecture.md index d9ddd816..889fd047 100644 --- a/diagrams/10-sdk-architecture.md +++ b/diagrams/10-sdk-architecture.md @@ -33,10 +33,6 @@ flowchart TB init --> operations init --> decorator decorator --> evaluation - - style sdk fill:#e3f2fd,stroke:#1565c0 - style operations fill:#f3e5f5,stroke:#7b1fa2 - style plugins_pkg fill:#fff3e0,stroke:#ef6c00 ``` ## Module Responsibilities @@ -72,9 +68,6 @@ flowchart LR m4 --> e3 m5 --> e4 m6 --> e5 - - style modules fill:#e8f5e9,stroke:#388e3c - style endpoints fill:#e3f2fd,stroke:#1565c0 ``` ## Client Usage Pattern @@ -134,8 +127,6 @@ flowchart TD mo5["EvaluationResponse"] end end - - style exports fill:#f5f5f5,stroke:#616161 ``` ## Two Usage Patterns @@ -148,8 +139,6 @@ flowchart LR import["import agent_control"] use["agent_control.policies.create_policy(client, name)"] end - - style code fill:#c8e6c9,stroke:#2e7d32 ``` ### Pattern 2: Direct Import @@ -160,8 +149,6 @@ flowchart LR import["from agent_control import policies"] use["policies.create_policy(client, name)"] end - - style code fill:#e3f2fd,stroke:#1565c0 ``` ## Plugin Module Structure @@ -184,9 +171,6 @@ flowchart TB init --> registry registry --> base optional_pkg -.->|lazy load| registry - - style plugins fill:#fff3e0,stroke:#ef6c00 - style optional_pkg fill:#ffe0b2,stroke:#f57c00 ``` ## Dependency Graph @@ -204,11 +188,6 @@ flowchart BT server --> engine sdk -.->|HTTP| server - - style models fill:#c8e6c9,stroke:#2e7d32 - style sdk fill:#fff3e0,stroke:#ef6c00 - style engine fill:#e3f2fd,stroke:#1565c0 - style server fill:#f3e5f5,stroke:#7b1fa2 ``` ## Initialization Flow @@ -232,8 +211,6 @@ flowchart TD init --> register register --> store store --> ready - - style ready fill:#c8e6c9,stroke:#2e7d32 ``` ## Error Handling @@ -258,9 +235,4 @@ flowchart TD check -->|4xx| client_error check -->|5xx| server_error http -.->|Network| network_error - - style success fill:#c8e6c9,stroke:#2e7d32 - style client_error fill:#fff3e0,stroke:#ef6c00 - style server_error fill:#ffcdd2,stroke:#c62828 - style network_error fill:#ffcdd2,stroke:#c62828 ```