SDK Overview

The Saf3AI SDK provides a production-ready interface for integrating AI security, observability, and governance into your applications. This overview covers the core concepts and API surface.

Installation

pip install saf3ai-sdk

Core Concepts

Initialization

The SDK must be initialized once at application startup before any security callbacks are created:

import os
from dotenv import load_dotenv
from saf3ai_sdk import init

load_dotenv()

init(
    service_name=os.getenv("SAF3AI_SERVICE_NAME", "my-agent"),
    framework="adk",  # "adk", "langchain", "rest", etc.
    agent_id="unique-agent-id",
    api_key=os.getenv("SAF3AI_API_KEY"),
    safeai_collector_agent=os.getenv("SAF3AI_COLLECTOR_AGENT"),
)

Security Callbacks

Security callbacks scan prompts and responses in real-time:

from saf3ai_sdk import create_security_callback

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", {})
    return not any(
        result.get("result") == "MATCH_FOUND" 
        for result in detections.values()
    )

security_callback = create_security_callback(
    api_endpoint=os.getenv("SAF3AI_API_ENDPOINT"),
    on_scan_complete=security_policy,
    scan_responses=True,
)

Tracing

The SDK provides comprehensive tracing for AI operations:

from saf3ai_sdk import trace, agent, task, tool

# Trace a function
@trace("my-operation")
def my_function():
    pass

# Mark agent execution
@agent("research-agent")
def research():
    pass

# Mark task execution
@task("summarization")
def summarize():
    pass

# Mark tool execution
@tool("web-search")
def search():
    pass

Direct Scanning

For custom integrations, use the scanner functions directly:

from saf3ai_sdk import scan_prompt, scan_response, scan_prompt_and_response

# Scan a prompt
results = scan_prompt(
    prompt="User input here",
    api_endpoint=os.getenv("SAF3AI_API_ENDPOINT"),
    model_name="gemini-2.5-flash",
    api_key=os.getenv("SAF3AI_API_KEY"),
)

# Scan a response
results = scan_response(
    response="LLM output here",
    api_endpoint=os.getenv("SAF3AI_API_ENDPOINT"),
)

# Scan both in one call
results = scan_prompt_and_response(
    prompt="User input",
    response="LLM output",
    api_endpoint=os.getenv("SAF3AI_API_ENDPOINT"),
)

Framework-Specific Integration

Google ADK

from saf3ai_sdk import init, create_security_callback
from google.adk.agents import LlmAgent

init(service_name="adk-agent", framework="adk", agent_id="my-agent")

security_callback = create_security_callback(
    api_endpoint=os.getenv("SAF3AI_API_ENDPOINT"),
    on_scan_complete=security_policy,
)

agent = LlmAgent(
    name="my_agent",
    model="gemini-2.5-flash",
    before_model_callback=security_callback,
)

LangChain

from saf3ai_sdk import init
from saf3ai_sdk.langchain_callbacks import create_security_callback
from langchain.chat_models import ChatOpenAI

init(service_name="langchain-agent", framework="langchain", agent_id="my-agent")

security_callback = create_security_callback(
    api_endpoint=os.getenv("SAF3AI_API_ENDPOINT"),
    on_scan_complete=security_policy,
)

chat = ChatOpenAI(callbacks=[security_callback])

Custom Frameworks

from saf3ai_sdk import init, create_framework_security_callbacks

init(service_name="custom-agent", framework="rest", agent_id="my-agent")

prompt_cb, response_cb = create_framework_security_callbacks(
    framework='rest',
    api_endpoint=os.getenv("SAF3AI_API_ENDPOINT"),
    agent_identifier='my-agent',
    on_scan_complete=security_policy,
)

Initialization Options

Option	Type	Required	Description
service_name	str	Yes	Name of your service for telemetry
framework	str	Yes	Framework type: “adk”, “langchain”, “rest”, etc.
agent_id	str	Yes	Unique identifier for your agent
api_key	str	Yes	Your Saf3AI API key
api_key_header_name	str	No	Custom header name for API key
safeai_collector_agent	str	No	Collector endpoint for telemetry
debug_mode	bool	No	Enable verbose logging

Security Callback Options

Option	Type	Default	Description
api_endpoint	str	Required	URL of the security scanner API
on_scan_complete	callable	None	Policy function: (text, results, type) -> bool
scan_responses	bool	False	Also scan LLM responses
timeout	int	30	Request timeout in seconds

Scan Results Format

{
    "status": "success",
    "detection_results": {
        "ThreatType": {"result": "MATCH_FOUND" | "NO_MATCH"}
    },
    "custom_rule_matches": [
        {"rule_id": "...", "action": "block", "matched_text": "..."}
    ],
    "OutofScopeAnalysis": {
        "detected_categories": ["Finance", "Legal"]
    },
    "entities": [...],
    "sentiment": {...},
    "scan_metadata": {
        "duration_ms": 245,
        "prompt_length": 150
    }
}

Custom Attributes

Add custom attributes to all spans:

from saf3ai_sdk import set_custom_attributes, get_custom_attributes, clear_custom_attributes

# Set attributes
set_custom_attributes({
    "user_id": "user-123",
    "session_id": "session-456",
    "environment": "production"
})

# Get current attributes
attrs = get_custom_attributes()

# Clear all attributes
clear_custom_attributes()

Session Management

Reset sessions for new conversations:

from saf3ai_sdk import reset_session

# Start a new session
reset_session()

Error Handling

The SDK handles errors gracefully:

try:
    results = scan_prompt(
        prompt="User input",
        api_endpoint="http://scanner.example.com",
    )
    
    if results.get("status") == "error":
        print(f"Scan failed: {results.get('error')}")
    elif results.get("status") == "timeout":
        print("Scan timed out")
    else:
        # Process results
        pass
        
except Exception as e:
    print(f"Unexpected error: {e}")

Disabling the SDK

For testing or local development:

# Via environment variable
# SAF3AI_ENABLED=false

# Or skip initialization
# Don't call init() in development

Supported Frameworks

Framework	Adapter	Status
Google ADK	adk	✅ Production
LangChain	langchain	✅ Production
LlamaIndex	llamaindex	✅ Ready
OpenAI	openai	✅ Ready
Anthropic	anthropic	✅ Ready
Cohere	cohere	✅ Ready
Groq	groq	✅ Ready
Ollama	ollama	✅ Ready
CrewAI	crewai	✅ Ready
AG2 (AutoGen)	ag2	✅ Ready
LiteLLM	litellm	✅ Ready
REST API	rest	✅ Ready

Next Steps

• Quick Start Guide
• LangChain Integration
• Google ADK Integration
• Custom Framework Integration