@auth0/ai-llamaindex

Auth0 AI for LlamaIndex

@auth0/ai-llamaindex is an SDK for building secure AI-powered applications using Auth0, Okta FGA and LlamaIndex.

Features

• Authorization for RAG: Securely filter documents using Okta FGA as a retriever for RAG applications. This smart retriever performs efficient batch access control checks, ensuring users only see documents they have permission to access.

• Tool Authorization with FGA: Protect AI tool execution with fine-grained authorization policies through Okta FGA integration, controlling which users can invoke specific tools based on custom authorization rules.

• Client Initiated Backchannel Authentication (CIBA): Implement secure, out-of-band user authorization for sensitive AI operations using the CIBA standard, enabling user confirmation without disrupting the main interaction flow.

• Federated API Access: Seamlessly connect to third-party services by leveraging Auth0's Tokens For APIs feature, allowing AI tools to access users' connected services (like Google, Microsoft, etc.) with proper authorization.

• Device Authorization Flow: Support headless and input-constrained environments with the Device Authorization Flow, enabling secure user authentication without direct input capabilities.

Install

[!WARNING] @auth0/ai-llamaindex is currently under heavy development. We strictly follow Semantic Versioning (SemVer), meaning all breaking changes will only occur in major versions. However, please note that during this early phase, major versions may be released frequently as the API evolves. We recommend locking versions when using this in production.

npm install @auth0/ai @auth0/ai-llamaindex

Initialization

Initialize the SDK with your Auth0 credentials:

import { Auth0AI, setAIContext } from "@auth0/ai-llamaindex";

const auth0AI = new Auth0AI({
  // Alternatively, you can use the `AUTH0_DOMAIN`, `AUTH0_CLIENT_ID`, and `AUTH0_CLIENT_SECRET`
  // environment variables.
  auth0: {
    domain: "YOUR_AUTH0_DOMAIN",
    clientId: "YOUR_AUTH0_CLIENT_ID",
    clientSecret: "YOUR_AUTH0_CLIENT_SECRET",
  },
  // store: new MemoryStore(), // Optional: Use a custom store
});

Set the ThreadID:

setAIContext({
  threadID: "...",
});

Calling APIs

The "Tokens for API" feature of Auth0 allows you to exchange refresh tokens for access tokens for third-party APIs. This is useful when you want to use a federated connection (like Google, Facebook, etc.) to authenticate users and then use the access token to call the API on behalf of the user.

First initialize the Federated Connection Authorizer as follows:

import { auth0 } from "./auth0";

export const withGoogleAccess = auth0AI.withTokenForConnection({
  // A function to retrieve the refresh token of the context.
  // In this case we are assuming this agent is running on next.js app
  // using next-auth0. "auth0" here is an instance of next-auth0
  refreshToken: async () => {
    const session = await auth0.getSession();
    const refreshToken = session?.tokenSet.refreshToken!;
    return refreshToken;
  },
  // The connection name.
  connection: 'google-oauth2',
  // The scopes to request.
  scopes: ["https://www.googleapis.com/auth/calendar.freebusy"],
});

Then use the withGoogleAccess to wrap the tool and use getAccessTokenForConnection from the SDK to get the access token.

import { tool } from "llamaindex";
import { getAccessTokenForConnection } from "@auth0/ai-llamaindex";
import { FederatedConnectionError } from "@auth0/ai/interrupts";
import { addHours } from "date-fns";

export const checkUsersCalendar = withGoogleAccess(
  tool(async ({ date }) => {
    const accessToken = getAccessTokenForConnection();
    const url = "https://www.googleapis.com/calendar/v3/freeBusy";
    const body = JSON.stringify({
      timeMin: date,
      timeMax: addHours(date, 1),
      timeZone: "UTC",
      items: [{ id: "primary" }],
    });

    const response = await fetch(url, {
      method: "POST",
      headers: {
        Authorization: `Bearer ${accessToken}`,
        "Content-Type": "application/json",
      },
      body,
    });

    if (!response.ok) {
      if (response.status === 401) {
        throw new FederatedConnectionError(
          `Authorization required to access the Federated Connection`
        );
      }
      throw new Error(
        `Invalid response from Google Calendar API: ${
          response.status
        } - ${await response.text()}`
      );
    }

    const busyResp = await response.json();
    return { available: busyResp.calendars.primary.busy.length === 0 };
  }, {
    description:
      "Check user availability on a given date time on their calendar",
    parameters: z.object({
      date: z.coerce.date(),
    })
  })
);

CIBA: Client-Initiated Backchannel Authentication

CIBA (Client-Initiated Backchannel Authentication) enables secure, user-in-the-loop authentication for sensitive operations. This flow allows you to request user authorization asynchronously and resume execution once authorization is granted.

import { auth0 } from "./auth0";

export const buyStockAuthorizer = auth0AI.withAsyncUserConfirmation({
  // A callback to retrieve the userID from tool context.
  userID: (params: { userID: string }, ctx) => params.userID,

  // The message the user will see on the notification
  bindingMessage: async ({ qty , ticker }) => {
    return `Confirm the purchase of ${qty} ${ticker}`;
  },
  // The scopes and audience to request
  audience: process.env["AUDIENCE"],
  scopes: ["stock:trade"]  
});

Then wrap the tool as follows:

import { tool } from "llamaindex";
import { z } from "zod";
import { getCIBACredentials } from "@auth0/ai-llamaindex";

export const purchaseStock = buyStockAuthorizer(
  tool(async ({ tradeID, userID, ticker, qty }) => {
    const { accessToken } = getCIBACredentials();
    fetch("http://yourapi.com/buy", {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        Authorization: `Bearer ${accessToken}`,
      },
      body: JSON.stringify({ ticker, qty }),
    });
    return `Purchased ${qty} shares of ${ticker}`;
  }, {
    name: "buy_stock",
    description: "Execute a stock purchase given stock ticker and quantity",
    parameters: z.object({
      tradeID: z
        .string()
        .uuid()
        .describe("The unique identifier for the trade provided by the user"),
      userID: z
        .string()
        .describe("The user ID of the user who created the conditional trade"),
      ticker: z.string().describe("The stock ticker to trade"),
      qty: z
        .number()
        .int()
        .positive()
        .describe("The quantity of shares to trade"),
    })
  })
);

Device Flow Authorizer

The Device Flow Authorizer enables secure, user-in-the-loop authentication for devices or tools that cannot directly authenticate users. It uses the OAuth 2.0 Device Authorization Grant to request user authorization and resume execution once authorization is granted.

import { auth0 } from "./auth0";

export const deviceFlowAuthorizer = auth0AI.withDeviceAuthorizationFlow({
  // The scopes and audience to request
  scopes: ["read:data", "write:data"],
  audience: "https://api.example.com",
});

Then wrap the tool as follows:

import { tool } from "llamaindex";
import { z } from "zod";
import { getDeviceAuthorizerCredentials } from "@auth0/ai-llamaindex";

export const fetchData = deviceFlowAuthorizer(
  tool(async ({ resourceID }) => {
    const credentials = getDeviceAuthorizerCredentials();
    const response = await fetch(`https://api.example.com/resource/${resourceID}`, {
      headers: {
        Authorization: `Bearer ${credentials.accessToken}`,
      },
    });

    if (!response.ok) {
      throw new Error(`Failed to fetch resource: ${response.statusText}`);
    }

    return await response.json();
  }, {
    name: "fetch_data",
    description: "Fetch data from a secure API",
    parameters: z.object({
      resourceID: z.string().describe("The ID of the resource to fetch"),
    })
  })
);

This flow is particularly useful for devices or tools that cannot directly authenticate users, such as IoT devices or CLI tools.

FGA

import { Auth0AI } from "@auth0/ai-llamaindex";

const auth0AI = new Auth0AI.FGA({
  apiScheme,
  apiHost,
  storeId,
  credentials: {
    method: CredentialsMethod.ClientCredentials,
    config: {
      apiTokenIssuer,
      clientId,
      clientSecret,
    },
  },
});
// Alternatively you can use env variables: `FGA_API_SCHEME`, `FGA_API_HOST`, `FGA_STORE_ID`, `FGA_API_TOKEN_ISSUER`, `FGA_CLIENT_ID` and `FGA_CLIENT_SECRET`

Then initialize the tool wrapper:

const authorizedTool = fgaAI.withFGA(
  {
    buildQuery: async ({ userID, doc }) => ({
      user: userID,
      object: doc,
      relation: "read",
    }),
  },
  myAITool
);

// Or create a wrapper to apply to tools later
const authorizer = fgaAI.withFGA({
  buildQuery: async ({ userID, doc }) => ({
    user: userID,
    object: doc,
    relation: "read",
  }),
});

const authorizedTool = authorizer(myAITool);

Note: the parameters given to the buildQuery function are the same provided to the tool's execute function.

Feedback

Contributing

We appreciate feedback and contribution to this repo! Before you get started, please see the following:

• Auth0's general contribution guidelines
• Auth0's code of conduct guidelines

Raise an issue

To provide feedback or report a bug, please raise an issue on our issue tracker.

Vulnerability Reporting

Please do not report security vulnerabilities on the public GitHub issue tracker. The Responsible Disclosure Program details the procedure for disclosing security issues.

---

Auth0 is an easy to implement, adaptable authentication and authorization platform. To learn more checkout Why Auth0?

This project is licensed under the Apache 2.0 license. See the LICENSE file for more info.