llama-index-trustguard 0.1.0
TrustGuard integration for LlamaIndex - protect your RAG pipeline from malicious content
License: MIT
Requires Python: >=3.9
Keywords: llamaindex, llama-index, rag, ai-security, trustguard, prompt-injection

# LlamaIndex TrustGuard

Security integration for [LlamaIndex](https://www.llamaindex.ai/) - protect your RAG pipeline from prompt injection and malicious content.

## Installation

```bash
pip install llama-index-trustguard
```

## Features

- **TrustGuardReader** - Scan documents before indexing
- **TrustGuardNodePostprocessor** - Scan retrieved nodes before using as context

## Quick Start

### Protected Document Loading

Scan documents for threats before indexing:

```python
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index_trustguard import TrustGuardReader

# Wrap your reader with threat protection
base_reader = SimpleDirectoryReader("./documents")
reader = TrustGuardReader(
    base_reader,
    api_key="ta_xxx...",
    on_threat="filter",  # Skip documents with threats
)

# Only safe documents are loaded
documents = reader.load_data()

# Build index with safe documents only
index = VectorStoreIndex.from_documents(documents)
```

### Protected Query Engine

Scan retrieved nodes before using them as context:

```python
from llama_index.core import VectorStoreIndex
from llama_index_trustguard import TrustGuardNodePostprocessor

# Create your index
index = VectorStoreIndex.from_documents(documents)

# Add TrustGuard postprocessor
postprocessor = TrustGuardNodePostprocessor(
    api_key="ta_xxx...",
    on_threat="filter",  # Filter out threatening nodes
)

# Create protected query engine
query_engine = index.as_query_engine(
    node_postprocessors=[postprocessor]
)

# Poisoned nodes are automatically filtered
response = query_engine.query("What is the company policy?")
```

### Full Protected RAG Pipeline

```python
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index_trustguard import TrustGuardReader, TrustGuardNodePostprocessor

# 1. Protected document loading
reader = TrustGuardReader(
    SimpleDirectoryReader("./knowledge_base"),
    api_key="ta_xxx...",
    on_threat="filter",
)
documents = reader.load_data()

# 2. Build index
index = VectorStoreIndex.from_documents(documents)

# 3. Protected retrieval
postprocessor = TrustGuardNodePostprocessor(
    api_key="ta_xxx...",
    on_threat="filter",
)

query_engine = index.as_query_engine(
    node_postprocessors=[postprocessor]
)

# Your RAG pipeline is now protected at both indexing and retrieval
response = query_engine.query("Tell me about our products")
```

---

## API Reference

### TrustGuardReader

Wraps any LlamaIndex reader with threat scanning.

```python
TrustGuardReader(
    reader: BaseReader,           # The reader to wrap
    api_key: str = None,          # TrustGuard API key
    on_threat: str = "warn",      # "block", "warn", "filter", "tag"
    min_block_level: ThreatLevel = ThreatLevel.HIGH,
    content_type: ContentSource = ContentSource.DOCUMENT,
)
```

**on_threat options:**
- `"block"` - Raise `ThreatInDocumentError` on threat
- `"warn"` - Log warning and continue
- `"filter"` - Skip documents with threats
- `"tag"` - Add threat info to document metadata

**Methods:**
```python
documents = reader.load_data()      # Load and scan documents
stats = reader.get_stats()          # Get scanning statistics
```

### TrustGuardNodePostprocessor

Postprocessor that scans retrieved nodes.

```python
TrustGuardNodePostprocessor(
    api_key: str = None,          # TrustGuard API key
    on_threat: str = "warn",      # Same options as reader
    min_block_level: ThreatLevel = ThreatLevel.HIGH,
)
```

**Usage:**
```python
query_engine = index.as_query_engine(
    node_postprocessors=[postprocessor]
)

# Or with a retriever
retriever = index.as_retriever(
    node_postprocessors=[postprocessor]
)
```

---

## Examples

### Web Content Scanning

```python
from llama_index.readers.web import SimpleWebPageReader
from llama_index_trustguard import TrustGuardReader
from agent_trust import ContentSource

reader = TrustGuardReader(
    SimpleWebPageReader(),
    api_key="ta_xxx...",
    content_type=ContentSource.WEB,  # Optimized for web content
    on_threat="filter",
)

documents = reader.load_data(urls=["https://example.com/docs"])
```

### Tagging Instead of Filtering

```python
# Tag documents with threat info instead of filtering
reader = TrustGuardReader(
    base_reader,
    api_key="ta_xxx...",
    on_threat="tag",
)

documents = reader.load_data()

for doc in documents:
    if doc.metadata.get("trust_guard", {}).get("safe") == False:
        print(f"Document has threats: {doc.metadata['trust_guard']['threats']}")
```

### Strict Mode (Block on Medium Threats)

```python
from agent_trust import ThreatLevel

reader = TrustGuardReader(
    base_reader,
    api_key="ta_xxx...",
    on_threat="block",
    min_block_level=ThreatLevel.MEDIUM,  # Stricter blocking
)
```

---

## Error Handling

```python
from llama_index_trustguard import TrustGuardReader, ThreatInDocumentError

reader = TrustGuardReader(
    base_reader,
    api_key="ta_xxx...",
    on_threat="block",
)

try:
    documents = reader.load_data()
except ThreatInDocumentError as e:
    print(f"Threat in: {e.document_id}")
    print(f"Verdict: {e.guard_result.verdict}")
    print(f"Threats: {[t.pattern_name for t in e.guard_result.threats]}")
```

---

## Statistics

```python
# Reader stats
reader_stats = reader.get_stats()
print(f"Scanned: {reader_stats['scanned_count']}")
print(f"Threats: {reader_stats['threat_count']}")
print(f"Filtered: {reader_stats['filtered_count']}")

# Postprocessor stats
pp_stats = postprocessor.get_stats()
print(f"Nodes scanned: {pp_stats['scanned_count']}")
```

## License

MIT License

## Links

- **TrustAgents:** https://trustagents.dev
- **LlamaIndex:** https://www.llamaindex.ai/
- **GitHub:** https://github.com/jd-delatorre/trustlayer