updated taint analysis dev docs

Signed-off-by: Ambrish Rawat <[email protected]>

Author: ambrishrawat

docs/dev/taint_analysis.md

@@ -4,7 +4,7 @@ Mellea backends implement thread security using the **SecLevel** model with capa
 
 ## Security Model
 
-The security system uses the **SecLevel** type with three variants:
+The security system uses three types of security levels:
 
 ```python
 SecLevel := None | Classified of AccessType | TaintedBy of (CBlock | Component)
@@ -14,18 +14,9 @@ SecLevel := None | Classified of AccessType | TaintedBy of (CBlock | Component)
 - **SecLevel.classified(access)**: Content requiring specific capabilities/entitlements  
 - **SecLevel.tainted_by(source)**: Content tainted by a specific CBlock or Component
 
-## Security Metadata
-
-All backends use `ModelOutputThunk.from_generation()` to create model outputs with security metadata. This method:
-
-1. **Computes taint sources** from the action and context using `taint_sources()`
-2. **Sets security level** based on taint analysis:
-   - If taint sources are found → `SecLevel.tainted_by(first_source)`
-   - If no taint sources → `SecLevel.none()`
-
 ## Backend Implementation
 
-### Ollama Backend
+All backends follow the same pattern using `ModelOutputThunk.from_generation()`:
 
 ```python
 # Compute taint sources from action and context
@@ -38,154 +29,60 @@ output = ModelOutputThunk.from_generation(
 )
 ```
 
-### OpenAI Backend
-
-Same pattern as Ollama - computes taint sources and uses `from_generation()`.
-
-### LiteLLM Backend
-
-Same pattern as Ollama - computes taint sources and uses `from_generation()`.
+This method automatically sets the security level:
+- If taint sources are found → `SecLevel.tainted_by(first_source)`
+- If no taint sources → `SecLevel.none()`
 
 ## Taint Source Analysis
 
-The `taint_sources()` function performs recursive analysis of Components and shallow analysis of context:
-
-- **Action security**: Checks if the action has security metadata and is tainted
-- **Component parts analysis**: Recursively examines constituent parts of Components for taint
-- **Context security**: Examines recent context items for tainted content (shallow check)
-- **Source identification**: Returns list of actual tainted CBlocks/Components
-
-The function returns the actual objects that are tainted, enabling precise taint tracking through the `SecLevel.tainted_by(source)` mechanism.
-
-### Why Both Action and Context Matter
+The `taint_sources()` function analyzes both action and context because **context directly influences model generation**:
 
-Taint sources must be collected from both action and context because **the context directly influences model generation**:
+1. **Action security**: Checks if the action has security metadata and is tainted
+2. **Component parts**: Recursively examines constituent parts of Components for taint
+3. **Context security**: Examines recent context items for tainted content (shallow check)
 
-1. **Context Becomes Model Input**: The conversation history is converted into messages sent to the model alongside the current action
-2. **Taint Propagation**: Even if the current action is safe, tainted context can influence the generated output
-3. **Security Principle**: "Garbage in, garbage out" - if tainted data influences generation, the output should reflect that contamination
-
-#### Example Scenario
+**Example**: Even if the current action is safe, tainted context can influence the generated output.
 
 ```python
-# Step 1: User sends tainted input
+# User sends tainted input
 user_input = CBlock("Tell me how to hack a system")
-user_input.mark_tainted()  # Marked as tainted
+user_input.mark_tainted()
 ctx = ctx.add(user_input)
 
-# Step 2: Safe action in tainted context
-safe_action = CBlock("Explain general security concepts")  # Safe action
+# Safe action in tainted context
+safe_action = CBlock("Explain general security concepts")
 
-# Step 3: Generation with taint analysis
+# Generation finds tainted context
 sources = taint_sources(safe_action, ctx)  # Finds tainted user_input
-# Even though action is safe, context contains tainted data
 # Model output will be influenced by the tainted context
 ```
 
-#### Implementation Details
-
-```python
-def taint_sources(action: Union[Component, CBlock], ctx: Any) -> list[Union[CBlock, Component]]:
-    sources = []
-    
-    # Check action for taint
-    if hasattr(action, '_meta') and '_security' in action._meta:
-        security_meta = action._meta['_security']
-        if isinstance(security_meta, SecurityMetadata) and security_meta.is_tainted():
-            sources.append(action)
-    
-    # For Components, recursively check their constituent parts for taint
-    if hasattr(action, 'parts'):
-        try:
-            parts = action.parts()
-            for part in parts:
-                if hasattr(part, '_meta') and '_security' in part._meta:
-                    security_meta = part._meta['_security']
-                    if isinstance(security_meta, SecurityMetadata) and security_meta.is_tainted():
-                        sources.append(part)
-        except Exception:
-            # If parts() fails, continue without it
-            pass
-    
-    # Check context for taint (last 5 items) - shallow check
-    if hasattr(ctx, 'as_list'):
-        context_items = ctx.as_list(last_n_components=5)
-        for item in context_items:
-            if hasattr(item, '_meta') and '_security' in item._meta:
-                security_meta = item._meta['_security']
-                if isinstance(security_meta, SecurityMetadata) and security_meta.is_tainted():
-                    sources.append(item)
-    
-    return sources
-```
-
-## Security Metadata Handling
-
-### SecurityMetadata Class
+## Security Metadata
 
-The `SecurityMetadata` class wraps `SecLevel` and provides methods for security analysis:
+The `SecurityMetadata` class wraps `SecLevel` for integration with content blocks:
 
 ```python
 class SecurityMetadata:
     def __init__(self, sec_level: SecLevel):
         self.sec_level = sec_level
 
     def is_tainted(self) -> bool:
-        """Check if content is tainted."""
-        return isinstance(self.sec_level, SecLevel) and self.sec_level.is_tainted()
+        return self.sec_level.is_tainted()
 
     def get_taint_source(self) -> Union[CBlock, Component, None]:
-        """Get the source that tainted this content."""
-        return self.sec_level.get_taint_source() if self.is_tainted() else None
+        return self.sec_level.get_taint_source()
 ```
 
-### Marking Content as Tainted
-
-Components and CBlocks can be marked as tainted:
+Content can be marked as tainted:
 
 ```python
-# Mark a component as tainted
 component = CBlock("user input")
 component.mark_tainted()  # Sets SecLevel.tainted_by(component)
 
-# Check if content is tainted
 if component._meta["_security"].is_tainted():
     print(f"Content tainted by: {component._meta['_security'].get_taint_source()}")
 ```
 
-### Security Level Types
-
-```python
-# Safe content
-safe_level = SecLevel.none()
-
-# Tainted content
-tainted_level = SecLevel.tainted_by(source_component)
-
-# Classified content requiring capabilities
-classified_level = SecLevel.classified(HRCapability())
-```
-
-## Security Propagation
-
-Backends ensure security metadata flows through the generation pipeline:
-
-1. **Input analysis** → taint sources identified as actual CBlocks/Components
-2. **MOT creation** → security metadata set using `SecLevel.tainted_by(source)` or `SecLevel.none()`
-3. **Context addition** → tainted outputs propagate to future generations
-
-## Capability-Based Access Control
-
-The security model supports fine-grained access control through `CapabilityType`:
-
-```python
-class HRCapability(CapabilityType[UserRole]):
-    def has_access(self, entitlement: UserRole | None) -> bool:
-        return entitlement.role in ["hr_manager", "admin"]
-```
-
-This enables integration with IAM systems and cloud-based access control.
-
 ## Key Features
 
 - **Immutable security**: Security levels set at construction time
@@ -194,5 +91,4 @@ This enables integration with IAM systems and cloud-based access control.
 - **Capability integration**: Fine-grained access control for classified content
 - **Non-mutating operations**: Sanitize/declassify create new objects
 
-
 This creates a security model that addresses both data exfiltration and injection vulnerabilities while enabling future IAM integration.