LlamaIndex

import os

from atlasent import AtlaSentClient
from atlasent_llamaindex import with_llamaindex_guard
from llama_index.core.tools import FunctionTool

client = AtlaSentClient(api_key=os.environ["ATLASENT_API_KEY"])

def search(query: str) -> str:
    """Search the knowledge base."""
    return f"Results for: {query}"

guarded_search = with_llamaindex_guard(
    search,
    client,
    agent="service:knowledge-bot",
    # agent.tool.invoke (the default action) requires `environment`
    extra_context={"environment": "production"},
    # Action classes require a state snapshot by default (SNAPSHOT_REQUIRED otherwise)
    state_snapshot={"source": "knowledge-bot", "complete": True},
)

tool = FunctionTool.from_defaults(fn=guarded_search)

import { AtlaSentClient } from "@atlasent/sdk";
import { FunctionTool } from "llamaindex";
import { withLlamaIndexGuard } from "@atlasent/llamaindex";

const client = new AtlaSentClient({ apiKey: process.env.ATLASENT_API_KEY! });

const guarded = withLlamaIndexGuard(
  [
    {
      metadata: {
        name: "vector_search",
        description: "Semantic search over the knowledge base",
        parameters: {
          type: "object",
          properties: { query: { type: "string" } },
          required: ["query"],
        },
      },
      execute: async ({ query }) => vectorStore.search(query),
    },
  ],
  client,
  {
    agent: "service:knowledge-bot",
    extraContext: { environment: "production" },
    // Action classes require a state snapshot by default (SNAPSHOT_REQUIRED otherwise)
    stateSnapshot: { source: "knowledge-bot", complete: true },
  },
);

const tools = guarded.map((t) => FunctionTool.from(t.execute, t.metadata));

{ "denied": true, "decision": "deny", "evaluationId": "...", "reason": "..." }