Loading…

Towards Security-Auditable LLM Agents: Unified Graph Representations | AI | graphwiz.ai

Towards Security-Auditable LLM Agents: Unified Graph Representations

How knowledge graphs can make AI agents transparent, traceable, and truly auditable

The Black Box Problem in Agentic AI

Enterprise adoption of agentic AI is accelerating, but security teams are hitting a wall. Traditional AI systems have a black box problem — inputs go in, outputs come out, and what happens in between is opaque. With agentic systems, this problem is amplified:

Multi-step reasoning — Agents chain tools, make decisions, and transform data across many steps
Dynamic behavior — The same prompt can produce different action sequences depending on state
External integrations — Agents read from databases, write to APIs, send emails, modify files
Compound decisions — Later steps depend on the results of earlier ones in non-linear ways

When something goes wrong — a data leak, a security violation, a harmful action — forensic analysis is nearly impossible. What tools did the agent call? What data did it access? What decisions did it make, and why?

The Need for Auditability

For regulated industries (finance, healthcare, government), auditability is non-negotiable. Organizations need to:

Trace — Every action an agent takes must be traceable back to its cause
Explain — Every decision must be explainable in human-understandable terms
Reproduce — Given the same inputs and state, the same outputs should be achieved
Verify — Security policies must be verifiable before, during, and after execution
Rollback — Harmful actions must be reversible with minimal collateral damage

Traditional logging approaches fail because they're linear (log lines in sequence) while agent execution is hierarchical and graph-structured. We need a representation that captures the true complexity of agent behavior.

Unified Graph Representation: The Solution

A unified graph representation models agent execution as a directed acyclic graph (DAG) where:

Nodes represent: user inputs, LLM generations, tool calls, decisions, data flows
Edges represent: causal relationships, data dependencies, control flow
Metadata on nodes/edges: timestamps, actors, permissions, data classifications, policy violations
Properties on nodes: content, type, status (pending/executed/failed), security labels

The Agent Execution Graph

┌─────────────────────────────────────────────────────────────────┐
│                        AGENT EXECUTION GRAPH                    │
└───────────────────────┬─────────────────────────────────────────┘
                        │
    ┌───────────────┐     │     ┌───────────────┐
    │               │     │     │               │
    │  USER INPUT   │─────┼─────│   LLM THOUGHT │
    │  (Root node)  │     │     │   (Decision)  │
    │               │     │     │               │
    └───────┬───────┘     │     └───────┬───────┘
            │             │             │
            │   data      │   causes    │
            ▼             │             ▼
    ┌───────────────┐     │     ┌───────────────┐
    │               │     │     │               │
    │ TOOL CALL     │─────┼─────│   CONDITION   │
    │ (DB Query)    │     │     │   (If-Then)   │
    │               │     │     │               │
    └───────┬───────┘     │     └───────┬───────┘
            │             │             │
            │   returns   │   satisfies ▼
            ▼             │     ┌───────────────┐
    ┌───────────────┐     │     │               │
    │               │     │     │    ACTION     │
    │ DB RESULT     │─────┼─────│ (Send Email)  │
    │  (Data)       │     │     │               │
    │               │     │     └───────┬───────┘
    └───────────────┘     │             │
                        │   produces   ▼
                        │     ┌───────────────┐
                        │     │               │
                        └─────│   USER OUTPUT  │
                              │ (Final result) │
                              │               │
                              └───────────────┘

What Makes This Graph-Powerful

Causality — Every node has a clear causal chain back to the root
Data flow — Data items (inputs, intermediate results, outputs) are first-class nodes
Control flow — Decision points and branches are explicitly represented
Temporal order — Nodes have timestamps; edges can represent timing constraints
Security context — Every node carries security labels (permission level, data classification)
Audit trail — Query the graph to answer: "How did we get this result?"

Implementation: Building the Security Graph

Step 1: Instrument the Agent Framework

The first requirement is instrumenting your agent framework to emit structured execution data. This typically involves:

import uuid
from datetime import datetime
from typing import Any, Optional
import json

class AuditNode:
    """Represents a node in the execution graph."""
    
    def __init__(
        self,
        node_type: str,
        content: Any,
        parent_id: Optional[str] = None,
        node_id: Optional[str] = None,
    ):
        self.node_id = node_id or str(uuid.uuid4())
        self.node_type = node_type  # input, llm, tool, decision, data, output
        self.content = content
        self.timestamp = datetime.utcnow().isoformat() + "Z"
        self.parent_id = parent_id
        self.children_ids = []
        self.metadata = {
            "security_level": "unknown",
            "data_classification": None,
            "policy_violations": [],
        }
    
    def add_child(self, child: "AuditNode"):
        child.parent_id = self.node_id
        self.children_ids.append(child.node_id)
        return child
    
    def to_dict(self):
        return {
            "node_id": self.node_id,
            "node_type": self.node_type,
            "content_hash": hash_content(self.content) if self.content else None,
            "timestamp": self.timestamp,
            "parent_id": self.parent_id,
            "children_ids": self.children_ids,
            "metadata": self.metadata,
        }


class SecurityGraph:
    """Accumulates and queries the execution graph."""
    
    def __init__(self):
        self.nodes = {}
        self.root = None
    
    def add_node(self, node: AuditNode):
        self.nodes[node.node_id] = node
        if self.root is None:
            self.root = node
        if node.parent_id:
            self.nodes[node.parent_id].children_ids.append(node.node_id)
        return node
    
    def to_neo4j_cypher(self):
        """Export graph to Neo4j Cypher for storage."""
        cypher = []
        # Create nodes
        for node_id, node in self.nodes.items():
            props = {
                "node_id": node.node_id,
                "node_type": node.node_type,
                "timestamp": node.timestamp,
                "security_level": node.metadata.get("security_level", "unknown"),
            }
            cypher.append(
                f"CREATE (n:{node.node_type} {{ {json.dumps(props)} }})"
            )
        # Create relationships
        for node_id, node in self.nodes.items():
            for child_id in node.children_ids:
                cypher.append(
                    f"MATCH (a), (b) WHERE a.node_id = '{node.node_id}' AND b.node_id = '{child_id}' "
                    f"CREATE (a)-[:CAUSES]->(b)"
                )
        return "\n".join(cypher)


# Instrumented LangChain/CrewAI/AutoGen agent wrapper
audit_graph = SecurityGraph()

class AuditableAgent:
    def __init__(self, llm, tools=[]):
        self.llm = llm
        self.tools = tools
        
    def run(self, prompt: str):
        # Root node: user input
        input_node = audit_graph.add_node(
            AuditNode("input", prompt, node_type="user_input")
        )
        current_node = input_node
        
        # LLM reasoning
        llm_node = current_node.add_child(
            AuditNode("llm", {"prompt": prompt, "response": None})
        )
        audit_graph.add_node(llm_node)
        
        response = self.llm.predict(prompt)
        llm_node.content["response"] = response
        current_node = llm_node
        
        # Tool calls
        for tool_name, args in self._extract_tool_calls(response):
            tool_node = current_node.add_child(
                AuditNode("tool_call", {"tool": tool_name, "args": args})
            )
            audit_graph.add_node(tool_node)
            
            result = self._call_tool(tool_name, args)
            result_node = tool_node.add_child(
                AuditNode("tool_result", result)
            )
            audit_graph.add_node(result_node)
            current_node = result_node
        
        return audit_graph

Step 2: Capture the Full Context

For true auditability, don't just capture execution flow — capture all contextual information:

Input context: Prompt, parameters, session metadata
LLM context: Model version, temperature, system prompt
Tool context: Which tools were available, with what permissions
Data context: Data sources accessed, data classifications
Environment context: Running in sandbox? Production? Which user?
Security context: API keys used, rate limits, authentication status

{
  "node_id": "llm_001",
  "node_type": "llm",
  "content": {
    "prompt": "What's the latest database backup?",
    "response": "The latest backup was on 2026-08-25..."
  },
  "timestamp": "2026-08-27T10:30:00Z",
  "context": {
    "model": "gpt-4o-2026-08",
    "temperature": 0,
    "system_prompt_fingerprint": "sha256:abcd1234",
    "available_tools": ["db_query", "file_read", "email_send"],
    "user_id": "user_42",
    "session_id": "sess_99",
    "permissions": ["db:read", "files:read"]
  }
}

Step 3: Security Labels and Policies

Every node should carry security labels that enable policy enforcement at query time:

# Augment nodes with security labels
class SecurityLabeler:
    def label_node(self, node: AuditNode):
        if node.node_type == "user_input":
            node.metadata["security_level"] = "user_controlled"
            node.metadata["data_classification"] = "untrusted"
        elif node.node_type == "llm":
            node.metadata["security_level"] = "processed"
        elif node.node_type == "tool_call":
            if "db" in node.content["tool"]:
                node.metadata["security_level"] = "high"
                node.metadata["data_classification"] = "confidential"
        elif node.node_type == "tool_result":
            # Inherit from parent tool call
            parent = self._find_parent(node)
            node.metadata["security_level"] = parent.metadata.get("security_level")

# Policy checker
class PolicyChecker:
    def check_graph(self, graph: SecurityGraph) -> List[PolicyViolation]:
        violations = []
        for node in graph.nodes.values():
            v = self._check_node(node)
            if v:
                violations.append(v)
        return violations
    
    def _check_node(self, node: AuditNode) -> Optional[PolicyViolation]:
        # Policy: No untrusted data in LLM prompts without sanitization
        if node.node_type == "llm" and node.metadata.get("data_classification") == "untrusted":
            if not node.metadata.get("sanitized", False):
                return PolicyViolation(
                    node=node,
                    policy="untrusted_prompt_policy",
                    message="Untrusted data used in LLM prompt without sanitization"
                )
        # Policy: High security tools require user approval
        if (node.node_type == "tool_call" and 
            node.metadata.get("security_level") == "high"):
            if not node.metadata.get("user_approved", False):
                return PolicyViolation(
                    node=node,
                    policy="high_security_approval", 
                    message="High security tool called without user approval"
                )
        return None

Step 4: Store in Neo4j for Query Power

Neo4j is the ideal graph database for security audit data because it:

Natively supports the query patterns we need (path finding, reachability, pattern matching)
Has built-in security features (RBAC, -based access control)
Supports temporal data (APOC library for time-series queries)
Scales well for large audit graphs
Provides Cypher — a declarative query language perfect for security audits

// Create schema for agent execution graph
CREATE CONSTRAINT FOR (n:input) REQUIRE n.node_id IS UNIQUE;
CREATE CONSTRAINT FOR (n:llm) REQUIRE n.node_id IS UNIQUE;
CREATE CONSTRAINT FOR (n:tool_call) REQUIRE n.node_id IS UNIQUE;
CREATE CONSTRAINT FOR (n:tool_result) REQUIRE n.node_id IS UNIQUE;
CREATE CONSTRAINT FOR (n:output) REQUIRE n.node_id IS UNIQUE;

// Create a node with security labels
CREATE (n:llm {
  node_id: 'llm_001',
  timestamp: '2026-08-27T10:30:00Z',
  security_level: 'processed',
  data_classification: 'internal'
})

// Create relationship with security context
MATCH (a:input {node_id: 'input_001'}), (b:llm {node_id: 'llm_001'})
CREATE (a)-[:CAUSES {
  security_level: 'inherited',
  data_flow: 'unmodified'
}]->(b)

Security Audit Queries

With the execution graph stored in Neo4j, auditors can answer practically any question:

1. What did the agent access?

// Find all data accessed during an execution
MATCH path = (input:input)-[:CAUSES*]->(n)
WHERE input.node_id = 'user_query_123'
AND (n:tool_result OR n:input)
RETURN n.timestamp, n.node_type, n.content as data, n.security_level
ORDER BY n.timestamp

2. Has any PII been leaked?

// Find all ouput nodes that inherited from PII inputs
MATCH (input:input {data_classification: 'PII'})-[:CAUSES*]->(output:output)
RETURN input.node_id, output.node_id, output.timestamp, output.content

3. Was a high-security tool called without approval?

// Find tool calls with high security but no approval
MATCH (tc:tool_call)
WHERE tc.security_level = 'high' AND NOT(tc.user_approved = true)
RETURN tc.node_id, tc.content.tool, tc.timestamp

4. What is the full chain of custody for a specific result?

// Trace backwards from a result to its origins
MATCH path = (result:tool_result {node_id: 'result_456'})<-[:CAUSES*]-()<-[:CAUSES]-(origin)
RETURN nodes(path), edges(path)
ORDER BY length(path) DESC
LIMIT 1

5. Find all interactions with a specific data source

// All nodes related to the 'payments' database
MATCH path = (n)-[:CAUSES*0..5]-()
WHERE n.content.tool = 'db_query' AND toLower(n.content.args.query) CONTAINS 'payments'
RETURN nodes(path)

6. Identify potential prompt injection

// Find LLM calls that received content from untrusted sources
MATCH (untrusted)-[:CAUSES*]->(llm:llm)
WHERE untrusted.data_classification = 'untrusted' 
  AND untrusted.node_type IN ['input', 'tool_result']
  AND llm.metadata.sanitized = false
RETURN untrusted.node_id, llm.node_id, llm.timestamp

Real-World Use Cases

1. Financial Services — Trade Surveillance

A banking application uses auditable agents for trade recommendations. The audit graph captures:

User inputs: Investment objectives, risk tolerance
Market data: Real-time prices, news feeds (with source attribution)
LLM reasoning: Analysis and recommendation logic
Tool calls: Actual trade executions
Output: Recommendations or actions

Auditors query: "Show me all trades for user_x on date_y, with explanation of reasoning."

2. Healthcare — Patient Data Access

A medical diagnosis assistant uses the security graph to:

Track PHI access: Every piece of patient health information is labeled as PHI
Enforce need-to-know: Only authorized personnel can access specific patient records
Audit data flow: Patient data can only flow to approved nodes (medical LLM, specific tools)
Prevent leakage: Queries ensure PHI never appears in chat logs or unsecured outputs

Compliance query: "Prove that patient_123's diagnosis never left the secure zone."

3. Government — Classified Information

For classified deployments, the audit graph:

Classifies all nodes: Each node carries top-secret/secret/confidential/unclassified
Clearance checking: Only nodes at clearance level <= user clearance are retrievable
Need-to-know enforcement: Additional access controls beyond classification
Change tracking: Every modification to classified data is recorded

Security query: "Show me all actions user_456 took with top-secret data in the past 30 days."

Implementing in Production

Architecture

┌──────────────────┐    ┌──────────────────┐    ┌──────────────────┐
│                  │    │                  │    │                  │
│  Agent Framework │    │  Audit Collector │    │   Neo4j Graph    │
│  (CrewAI/LangGraph)   │       (Sidecar)   │    │    Database      │
│                  │    │                  │    │                  │
└────────┬─────────┘    └────────┬─────────┘    └────────┬─────────┘
         │                       │                        │
         │      Audit Events      │                        │
         │─────────────────────────>│                        │
         │                       │                        │
         │                       │      Graph Operations    │
         │                       │─────────────────────────>│
         │                       │                        │
└────────────────────────────────────────────────────────────────┘

┌──────────────────┐    ┌──────────────────┐
│                  │    │                  │
│  Audit Dashboard │    │  Alert System    │
│  (Grafana)       │    │  (PrometheusRule)│
│                  │    │                  │
└────────┬─────────┘    └────────┬─────────┘
         │                       │
         │      Queries          │
         └───────────────────────┘

Performance Considerations

Audit graph storage can grow large quickly. Mitigations:

Retention policies — Archive old graphs after N days, keep only metadata
Compaction — Collapse simple execution paths into single nodes
Sampling — For non-critical agents, sample audit data (1 in 10 executions)
Sharding — Store graphs by agent type, user, or time period

Integration with Existing Security Stack

SIEM systems — Forward audit events to Splunk, QRadar, etc.
SOAR platforms — Trigger remediation workflows based on policy violations
IAM systems — Use existing identity providers for user context
DLP systems — Integrate data classification with existing DLP policies

Challenges

1. Privacy Preservation

Audit graphs contain sensitive information by design. Solutions:

Redaction — Strip PII before storage
Encryption — Encrypt node content at rest
Access controls — Role-based access to different graph segments

2. Performance Overhead

Capturing audit data adds latency. Solutions:

Async capture — Log audit events asynchronously
Batch writes — Collect events and batch-write to Neo4j
Selective auditing — audit only high-security agents or randomly sample

3. Complexity

Building a comprehensive security graph is complex. Solutions:

Phased rollout — Start with high-risk agents first
Open standards — Use shared schema and query patterns
Commercial tools — Several vendors offer agent audit solutions

Conclusion

Unified graph representations transform agentic AI from opaque black boxes into transparent, auditable systems. By capturing the complete causal chain of agent execution — inputs, reasoning, decisions, and actions — in a queryable graph, organizations can finally achieve:

✅ Traceability — Every action has a clear lineage
✅ Accountability — Every decision can be explained
✅ Compliance — Regulatory requirements can be verified
✅ Security — Policy violations can be detected and prevented
✅ Forensics — Incident investigations are possible

The future of secure agentic AI isn't more guardrails — it's better visibility. And the best visibility comes from modeling agent behavior as a graph.

Research sources: Towards Security-Auditable LLM Agents: A Unified Graph Representation (arXiv:2605.06812), AgentRiskBOM: A Risk-Scoping Security Bill of Materials for Agentic AI Systems (arXiv:2606.21877).

Building your own auditable agents? Our Knowledge Graph Toolkit for Python Developers provides the foundation you need.