Agent authentication¶

AI agents often need to authenticate to external resources to complete tasks. For example, a deployed agent might need to access external APIs, databases, or cloud services to retrieve data or perform operations.

This documentation provides comprehensive guidance for implementing authentication in DataRobot Agent Templates, covering API tokens, authorization context, OAuth 2.0, and security best practices.

Authentication methods¶

This section provides an overview of the different authentication methods available in the framework:

API token authentication¶

API token authentication is the most common method for authenticating with DataRobot services and external APIs. It uses bearer tokens passed in headers or environment variables.

DataRobot API authentication¶

Configure API tokens using environment variables or programmatically. The DATAROBOT_API_TOKEN is available in the API keys and tools section of your account settings.

DATAROBOT_API_TOKEN=<your_api_key>
DATAROBOT_ENDPOINT=https://app.datarobot.com

from agent import MyAgent

agent = MyAgent(
    api_key="your_api_key",
    api_base="https://app.datarobot.com"
)

External API authentication¶

When creating custom tools that need to authenticate with external APIs, retrieve API keys from environment variables within your tool's run() method. This keeps credentials out of your source code and follows the same security pattern used throughout the agent templates. The example below shows the standard pattern: retrieve the API key using os.getenv(), then include it in the request headers to the external API.

import os
import requests
from crewai.tools import BaseTool

class ExternalAPITool(BaseTool):
    def run(self, query: str) -> str:
        api_key = os.getenv("EXTERNAL_API_KEY")
        headers = {"Authorization": f"Bearer {api_key}"}
        response = requests.get(
            "https://api.external-service.com/data",
            headers=headers,
            params={"query": query}
        )
        return response.json()

Authorization context¶

The framework provides an authorization context system for propagating authentication information to downstream tools and services. This allows tokens and credentials to be automatically passed between agent tools without manual configuration.

Initialize authorization context in your agent's chat() function using resolve_authorization_context() from the datarobot-genai package. The function returns the authorization context dictionary, which should be assigned to completion_create_params["authorization_context"]. Tools can then retrieve the context using get_authorization_context() from the datarobot SDK:

from datarobot_genai.core.chat import resolve_authorization_context
from datarobot.models.genai.agent.auth import get_authorization_context

def chat(completion_create_params, load_model_result, **kwargs):
    # Initialize the authorization context for downstream agents and tools
    completion_create_params["authorization_context"] = resolve_authorization_context(
        completion_create_params, **kwargs
    )
    # ... rest of chat function

# In your tools
auth_context = get_authorization_context()
access_token = auth_context.get("access_token")

OAuth 2.0 authentication¶

For external services that support OAuth 2.0, implement OAuth flows in your tools. Create an OAuth application in the service's developer console and set environment variables:

OAUTH_CLIENT_ID=<your_client_id>
OAUTH_CLIENT_SECRET=<your_client_secret>
OAUTH_REDIRECT_URI=<your_redirect_uri>
OAUTH_SCOPE=<required_scopes>

Then, use those environment variables in the OAuth implementation. This example demonstrates a complete OAuth 2.0 authorization code flow for a CrewAI tool. The tool inherits from crewai.tools.BaseTool (the standard base class for tools in CrewAI agent templates), manages its own token lifecycle with in-memory caching, and follows the repository's pattern of retrieving credentials from environment variables.

import os
import time
import requests
from urllib.parse import urlencode
from crewai.tools import BaseTool

class ExampleToolWithOAuth(BaseTool):
    def __init__(self):
        super().__init__()
        self.client_id = os.getenv("OAUTH_CLIENT_ID")
        self.client_secret = os.getenv("OAUTH_CLIENT_SECRET")
        self.redirect_uri = os.getenv("OAUTH_REDIRECT_URI")
        self._access_token = None
        self._token_expires_at = None

    def get_authorization_url(self) -> str:
        params = {
            "client_id": self.client_id,
            "redirect_uri": self.redirect_uri,
            "response_type": "code",
            "scope": "read:data"
        }
        return f"https://oauth.provider.com/authorize?{urlencode(params)}"

    def exchange_code_for_token(self, code: str) -> dict:
        response = requests.post(
            "https://oauth.provider.com/token",
            data={
                "grant_type": "authorization_code",
                "client_id": self.client_id,
                "client_secret": self.client_secret,
                "code": code,
                "redirect_uri": self.redirect_uri
            }
        )
        token_data = response.json()
        self._access_token = token_data.get("access_token")
        expires_in = token_data.get("expires_in", 3600)
        self._token_expires_at = time.time() + expires_in
        return token_data

    def get_cached_access_token(self) -> str:
        if self._access_token and self._token_expires_at and time.time() < self._token_expires_at:
            return self._access_token
        # Token expired or not set - refresh or re-authenticate
        # In production, implement token refresh logic here
        raise ValueError("Access token expired. Re-authenticate to get a new token.")

    def run(self, query: str) -> str:
        access_token = self.get_cached_access_token()
        response = requests.get(
            "https://api.provider.com/data",
            headers={"Authorization": f"Bearer {access_token}"},
            params={"query": query}
        )
        return response.json()

Security best practices¶

Following security best practices is essential when handling authentication in production environments. Adhere to the following guidelines:

• API token authentication
  • Store tokens in environment variables; never hard code secrets in source code.
  • Use secure token storage solutions in production environments.
  • Implement token rotation and enforce token expiration policies.
  • Validate authorization context before using it in tools.
  • Follow least-privilege access principles.
  • Log authentication events for audit purposes.
• OAuth 2.0
  • Use HTTPS for all OAuth communications.
  • Validate the state parameter to prevent Cross-Site Request Forgery (CSRF) attacks.
  • Store refresh tokens securely.
  • Handle token expiration and refresh logic reliably.
  • Validate authorization context before use.
• General Security
  • Use separate environments for development and production.
  • Implement robust secret management practices.
  • Follow container security best practices.
  • Conduct regular security audits and apply updates.

Troubleshooting authentication issues¶

This section helps you diagnose and resolve common authentication problems when developing or deploying agents.

Common issues¶

The following sections describe common authentication errors and how to resolve them:

Missing API token¶

This error occurs when the DataRobot API token is not configured.

Issue: Error: Missing DataRobot API token. Set the DATAROBOT_API_TOKEN environment variable

Solution: Set the DATAROBOT_API_TOKEN environment variable to use your DataRobot API key.

Invalid endpoint¶

This error occurs when the DataRobot endpoint is missing or incorrectly configured.

Issue: Error: Missing DataRobot endpoint. Set the DATAROBOT_ENDPOINT environment variable

Solution: Set the correct DATAROBOT_ENDPOINT environment variable.

Authorization context not set¶

This error occurs when tools try to access the authorization context before it has been initialized.

Issue: Error: Authorization context not available for tool

Solution: Ensure resolve_authorization_context() is called in your agent's chat() function and the result is assigned to completion_create_params["authorization_context"].

OAuth token expired¶

This error occurs when an OAuth access token has expired and needs to be refreshed.

Issue: Error: 401 Unauthorized

Solution: Implement token refresh logic or re-authenticate.

Debugging tips¶

Useful techniques for debugging authentication issues:

Enable verbose logging¶

Enable verbose logging to get detailed information about authentication operations. To do so, search for where MyAgent is instantiated in custom.py and set verbose=True.

from agent import MyAgent

# ...

agent = MyAgent(verbose=True, **completion_create_params)

Check environment variables¶

Verify that required environment variables are set correctly.

import os
print(f"API Token: {os.getenv('DATAROBOT_API_TOKEN')}")
print(f"Endpoint: {os.getenv('DATAROBOT_ENDPOINT')}")

Test authentication¶

Test authentication using the AgentEnvironment class from datarobot_genai.core.cli:

from datarobot_genai.core.cli import AgentEnvironment

try:
    env = AgentEnvironment()
    print("Authentication successful")
except ValueError as e:
    print(f"Authentication failed: {e}")

Validate authorization context¶

Check authorization context using get_authorization_context() from datarobot.models.genai.agent.auth:

from datarobot.models.genai.agent.auth import get_authorization_context

auth_context = get_authorization_context()
print(f"Authorization context: {auth_context}")