saf3ai-sdk 0.2.3
Unified observability and security scanning SDK for AI agents across 23+ frameworks
Requires Python: >=3.9

# Saf3AI SDK

[![Python Version](https://img.shields.io/badge/python-3.9%2B-blue.svg)](https://www.python.org/downloads/)
[![License](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE)
[![PyPI](https://img.shields.io/pypi/v/saf3ai-sdk.svg)](https://pypi.org/project/saf3ai-sdk/)

Production-ready observability and security scanning SDK for Google ADK and LangChain agents.

---

## Installation

```bash
pip install saf3ai-sdk
```

---

## Step 1: Configuration

The SDK reads configuration values from environment variables. Set these in your deployment environment using your preferred method (system environment variables, container orchestration, CI/CD pipelines, secret managers, etc.).

### Required Environment Variables

| Variable | Description | Example |
|----------|-------------|---------|
| `SAF3AI_COLLECTOR_AGENT` | Saf3AI Collector endpoint URL | `https://your-collector-endpoint.com` |
| `SAF3AI_SERVICE_NAME` | Name of your service/agent | `my-agent` |
| `SAF3AI_API_KEY` | Your organization API key | `your-api-key-here` |
| `SAF3AI_API_KEY_HEADER` | HTTP header name for API key (optional, defaults to `X-API-Key`) | `X-API-Key` |
| `SAF3AI_API_ENDPOINT` | Saf3AI Scanner endpoint URL | `https://your-scanner-endpoint.com` |

### Setting Environment Variables

**Linux/macOS:**
```bash
export SAF3AI_COLLECTOR_AGENT=https://your-collector-endpoint.com
export SAF3AI_SERVICE_NAME=my-agent
export SAF3AI_API_KEY=your-api-key-here
export SAF3AI_API_ENDPOINT=https://your-scanner-endpoint.com
```

**Windows (PowerShell):**
```powershell
$env:SAF3AI_COLLECTOR_AGENT="https://your-collector-endpoint.com"
$env:SAF3AI_SERVICE_NAME="my-agent"
$env:SAF3AI_API_KEY="your-api-key-here"
$env:SAF3AI_API_ENDPOINT="https://your-scanner-endpoint.com"
```

**Docker:**
```dockerfile
ENV SAF3AI_COLLECTOR_AGENT=https://your-collector-endpoint.com
ENV SAF3AI_SERVICE_NAME=my-agent
ENV SAF3AI_API_KEY=your-api-key-here
ENV SAF3AI_API_ENDPOINT=https://your-scanner-endpoint.com
```

**Kubernetes (ConfigMap/Secret):**
```yaml
apiVersion: v1
kind: Secret
metadata:
  name: saf3ai-config
type: Opaque
stringData:
  SAF3AI_COLLECTOR_AGENT: https://your-collector-endpoint.com
  SAF3AI_SERVICE_NAME: my-agent
  SAF3AI_API_KEY: your-api-key-here
  SAF3AI_API_ENDPOINT: https://your-scanner-endpoint.com
```

---

## Step 2: Initialize SDK

```python
import os
from saf3ai_sdk import init

# Initialize SDK
init(
    service_name=os.getenv("SAF3AI_SERVICE_NAME", "my-agent"),
    framework="adk",  # Use "adk" for Google ADK or "langchain" for LangChain
    agent_id="unique-agent-id",
    api_key=os.getenv("SAF3AI_API_KEY"),
    api_key_header_name=os.getenv("SAF3AI_API_KEY_HEADER", "X-API-Key"),
    safeai_collector_agent=os.getenv("SAF3AI_COLLECTOR_AGENT"),
)
```

**Important:** Call `init()` once at the start of your application, before creating any agents or LLMs.

---

## Step 3: Define Security Policy

```python
def security_policy(text: str, scan_results: dict, text_type: str) -> bool:
    """
    Return True to allow, False to block.
    """
    detections = scan_results.get("detection_results", {})
    
    # Block if any threat is found
    for threat_type, result in detections.items():
        if result.get("result") == "MATCH_FOUND":
            return False
    
    return True
```

---

## Google ADK Integration

### Step 4: Create Security Callback

```python
from saf3ai_sdk import create_security_callback
from google.adk.agents import LlmAgent

# Create security callback
security_callback = create_security_callback(
    api_endpoint=os.getenv("SAF3AI_API_ENDPOINT"),
    on_scan_complete=security_policy,
    scan_responses=True,  # Optional: also scan AI responses
)

# Create ADK agent with callback
agent = LlmAgent(
    name="my_agent",
    model="gemini-2.5-flash",
    before_model_callback=security_callback,
)

# Use agent
response = agent.run("Hello, how are you?")
```

### Complete ADK Example

```python
import os
from saf3ai_sdk import init, create_security_callback
from google.adk.agents import LlmAgent

# Step 1: Initialize SDK
init(
    service_name=os.getenv("SAF3AI_SERVICE_NAME", "adk-agent"),
    framework="adk",
    agent_id="my-adk-agent",
    api_key=os.getenv("SAF3AI_API_KEY"),
    api_key_header_name=os.getenv("SAF3AI_API_KEY_HEADER", "X-API-Key"),
    safeai_collector_agent=os.getenv("SAF3AI_COLLECTOR_AGENT"),
)

# Step 3: Define security policy
def security_policy(text: str, scan_results: dict, text_type: str) -> bool:
    detections = scan_results.get("detection_results", {})
    return not any(
        result.get("result") == "MATCH_FOUND" 
        for result in detections.values()
    )

# Step 4: Create security callback
security_callback = create_security_callback(
    api_endpoint=os.getenv("SAF3AI_API_ENDPOINT"),
    on_scan_complete=security_policy,
    scan_responses=True,
)

# Step 5: Create ADK agent with callback
agent = LlmAgent(
    name="my_agent",
    model="gemini-2.5-flash",
    before_model_callback=security_callback,
)

# Step 6: Use agent
response = agent.run("Hello, how are you?")
```

---

## LangChain Integration

### Step 4: Create Security Callback

```python
from saf3ai_sdk.langchain_callbacks import create_security_callback
from langchain.chat_models import ChatOpenAI
from langchain.chains import ConversationChain

# Create security callback
security_callback = create_security_callback(
    api_endpoint=os.getenv("SAF3AI_API_ENDPOINT"),
    on_scan_complete=security_policy,
    scan_responses=True,  # Optional: also scan AI responses
    conversation_id="conv-123",  # Optional: for conversation stitching
)

# Create LangChain chain with callback
chat = ChatOpenAI(
    openai_api_key=os.getenv("OPENAI_API_KEY"),
    callbacks=[security_callback],
)

chain = ConversationChain(llm=chat)

# Use chain with error handling
try:
    response = chain.run("Hello, how are you?")
except ValueError as e:
    if "cannot assist" in str(e).lower():
        print("Request blocked by security policy")
    else:
        raise
```

### Complete LangChain Example

```python
import os
import uuid
from saf3ai_sdk import init
from saf3ai_sdk.langchain_callbacks import create_security_callback
from langchain.chat_models import ChatOpenAI
from langchain.chains import ConversationChain

# Step 1: Initialize SDK
init(
    service_name=os.getenv("SAF3AI_SERVICE_NAME", "langchain-agent"),
    framework="langchain",
    agent_id="my-langchain-agent",
    api_key=os.getenv("SAF3AI_API_KEY"),
    api_key_header_name=os.getenv("SAF3AI_API_KEY_HEADER", "X-API-Key"),
    safeai_collector_agent=os.getenv("SAF3AI_COLLECTOR_AGENT"),
)

# Step 3: Define security policy
def security_policy(text: str, scan_results: dict, text_type: str) -> bool:
    detections = scan_results.get("detection_results", {})
    return not any(
        result.get("result") == "MATCH_FOUND" 
        for result in detections.values()
    )

# Step 4: Create security callback
conversation_id = str(uuid.uuid4())  # Generate or get from session
security_callback = create_security_callback(
    api_endpoint=os.getenv("SAF3AI_API_ENDPOINT"),
    on_scan_complete=security_policy,
    scan_responses=True,
    conversation_id=conversation_id,
)

# Step 5: Create LangChain chain with callback
chat = ChatOpenAI(
    openai_api_key=os.getenv("OPENAI_API_KEY"),
    callbacks=[security_callback],
)

chain = ConversationChain(llm=chat)

# Step 6: Use chain with error handling
try:
    response = chain.run("Hello, how are you?")
    print(response)
except ValueError as e:
    if "cannot assist" in str(e).lower():
        print("Request blocked by security policy")
    else:
        raise
```

### LangChain with Agent Executor

```python
from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain.tools import Tool

# Create agent with security callback
agent = create_openai_tools_agent(
    llm=chat,
    tools=[your_tools],
)

executor = AgentExecutor(
    agent=agent,
    tools=[your_tools],
    callbacks=[security_callback],
    verbose=True,
)

# Run agent
try:
    response = executor.invoke({"input": "User query here"})
except ValueError as e:
    if "cannot assist" in str(e).lower():
        print("Request blocked by security policy")
    else:
        raise
```

---

## Security Policy Examples

### Basic Policy (Block All Threats)

```python
def basic_policy(text: str, scan_results: dict, text_type: str) -> bool:
    """Block any detected threats."""
    detections = scan_results.get("detection_results", {})
    return not any(
        result.get("result") == "MATCH_FOUND" 
        for result in detections.values()
    )
```

### Selective Policy (Block Specific Threats)

```python
def selective_policy(text: str, scan_results: dict, text_type: str) -> bool:
    """Block only specific threat types."""
    detections = scan_results.get("detection_results", {})
    blocked_types = {"CSAM", "Dangerous", "HateSpeech"}
    
    for threat_type, result in detections.items():
        if threat_type in blocked_types and result.get("result") == "MATCH_FOUND":
            return False
    return True
```

---

## Troubleshooting

**SDK not initializing:**
- Verify all required environment variables are set in your deployment environment
- Check that environment variables are accessible to your application process
- Ensure environment variables are set before the application starts (not loaded at runtime)
- Verify variable names match exactly (case-sensitive)
- For containerized deployments, ensure environment variables are passed to containers
- For Kubernetes, verify ConfigMaps/Secrets are mounted and accessible
- Use `os.getenv("VARIABLE_NAME")` to debug and verify values are being read correctly

**Callbacks not working:**
- Verify SDK is initialized before creating callbacks
- Check that `framework` parameter matches your framework ("adk" or "langchain")
- Ensure framework is installed: `pip install langchain` or `pip install google-adk`
- Verify callbacks are added to agent/chain before invocation

**Environment variables not being read:**
- Ensure environment variables are exported/set in the same shell/process that runs your application
- For systemd services, set variables in the service file or use `EnvironmentFile`
- For Docker Compose, verify variables are defined in the `environment:` section of your service definition or passed via `env_file:` (prefer explicit `environment:` for production)
- For Kubernetes, ensure ConfigMaps/Secrets are properly mounted and referenced in pod specifications
- For CI/CD pipelines, ensure variables are set in pipeline configuration (GitHub Actions secrets, GitLab CI variables, etc.)
- For cloud platforms (AWS, GCP, Azure), use their respective secret management services (Secrets Manager, Secret Manager, Key Vault)
- Check for typos in variable names (they are case-sensitive)
- Verify the application has permissions to read environment variables
- Use `print(os.getenv("VARIABLE_NAME"))` or logging to verify values are accessible at runtime

**Conversation ID not stitching (LangChain):**
- Pass `conversation_id` to `create_security_callback()`
- Use the same `conversation_id` across multiple calls

---

## License

MIT License - see [LICENSE](LICENSE) file for details.

---

## Links

- **PyPI**: [https://pypi.org/project/saf3ai-sdk/](https://pypi.org/project/saf3ai-sdk/)