LlamaIndex Integration

pip install raxe[llamaindex]

from llama_index.core import VectorStoreIndex, Settings
from llama_index.core.callbacks import CallbackManager
from raxe import Raxe
from raxe.sdk.integrations import RaxeLlamaIndexCallback

# Initialize RAXE
raxe = Raxe()

# Create callback handler (default: log-only mode)
raxe_callback = RaxeLlamaIndexCallback(raxe_client=raxe)

# Configure LlamaIndex
callback_manager = CallbackManager([raxe_callback])
Settings.callback_manager = callback_manager

# Create and query index
index = VectorStoreIndex.from_documents(documents)
query_engine = index.as_query_engine()

# Queries are automatically scanned
response = query_engine.query("What are the key findings?")

# Blocking mode for queries
callback = RaxeLlamaIndexCallback(
    raxe_client=raxe,
    block_on_query_threats=True,     # Block on input threats
    block_on_response_threats=False,  # Log-only for responses
    scan_agent_actions=True,          # Scan agent tool inputs
)

from raxe.sdk.integrations import RaxeQueryEngineCallback

callback = RaxeQueryEngineCallback(
    raxe_client=raxe,
    block_on_threats=True,  # Block on any threats
)

from raxe.sdk.integrations import RaxeAgentCallback

callback = RaxeAgentCallback(
    raxe_client=raxe,
    block_on_threats=True,
    scan_tool_outputs=True,
)

from llama_index.core.agent import ReActAgent
from llama_index.core.tools import FunctionTool
from llama_index.core.callbacks import CallbackManager
from raxe.sdk.integrations import RaxeAgentCallback

# Create agent callback with blocking
raxe_callback = RaxeAgentCallback(block_on_threats=True)
callback_manager = CallbackManager([raxe_callback])

# Define tools
def calculate(expression: str) -> str:
    return str(eval(expression))

calc_tool = FunctionTool.from_defaults(fn=calculate)

# Create agent
agent = ReActAgent.from_tools(
    tools=[calc_tool],
    callback_manager=callback_manager,
    verbose=True,
)

# Agent interactions are scanned
response = agent.chat("Calculate 2 + 2")

from llama_index.core.instrumentation import get_dispatcher
from raxe.sdk.integrations import RaxeSpanHandler

# Create span handler
span_handler = RaxeSpanHandler(
    block_on_threats=False,  # Log only
    scan_llm_inputs=True,
    scan_llm_outputs=True,
)

# Register with root dispatcher
root_dispatcher = get_dispatcher()
root_dispatcher.add_span_handler(span_handler)

# All operations now traced and scanned

from raxe.sdk.exceptions import SecurityException

try:
    response = query_engine.query("Potentially malicious query")
except SecurityException as e:
    print("Query blocked for security reasons")