Authentifizierung für KI-Agents

KI-Agents müssen sich häufig bei anderen Ressourcen authentifizieren, um Aufgaben auszuführen. Beispielsweise muss ein bereitgestellter Agent möglicherweise auf einen Vektorsuchindex zugreifen, um unstrukturierte Daten abzufragen, oder die Aufforderungsregistrierung, um dynamische Eingabeaufforderungen zu laden.

Auf dieser Seite werden die Authentifizierungsmethoden behandelt, die beim Entwickeln und Bereitstellen von Agents mithilfe von Mosaik AI Agent Framework verfügbar sind.

Authentifizierungsmethoden

In der folgenden Tabelle werden die verfügbaren Authentifizierungsmethoden verglichen. Sie können eine der folgenden Ansätze kombinieren und aufeinander abstimmen:

Authentifizieren bei externen Systemen und MCP-Servern

Anleitungen zum Authentifizieren bei externen APIs und MCP-Servern von Ihrem Agent finden Sie unter Verbinden von KI-Agent-Tools mit externen Diensten. Diese Ressourcen können auch im Namen des Agents oder Benutzers abgefragt werden, wie in Authentifizierungsmethoden beschrieben.

Auswählen der richtigen Authentifizierungsmethode für Ihre Ressource

Verwenden Sie dieses Flussdiagramm, um die richtige Authentifizierungsmethode für jede Ressource auszuwählen. Sie können Methoden nach Bedarf kombinieren, und ein Agent kann je nach Anwendungsfall eine andere Methode für jede Ressource verwenden.

1. Ist eine benutzerbasierte Zugriffskontrolle oder eine auf Benutzerattribution basierende Überprüfung erforderlich?

   • Ja → Verwenden der Benutzerauthentifizierung im Auftrag von Benutzern
   • Keine → Weiter zu Schritt 2

2. Unterstützen alle Ressourcen die automatische Authentifizierung?

   • Ja → Automatische Authentifizierungspassthrough verwenden (empfohlen)
   • Keine → manuelle Authentifizierung verwenden

Passthrough für die automatische Authentifizierung

Die automatische Authentifizierungs-Durchleitung ist die einfachste Methode, um auf von Databricks verwaltete Ressourcen zuzugreifen. Deklarieren Sie beim Loggen des Agents Ressourcenabhängigkeiten, und Databricks stellt automatisch kurzlebige Anmeldeinformationen bereit, rotiert sie und verwaltet sie, wenn der Agent bereitgestellt wird.

Dieses Authentifizierungsverhalten ähnelt dem Verhalten "Als Besitzer ausführen" für Databricks-Dashboards. Nachgelagerte Ressourcen wie Unity Catalog-Tabellen werden mithilfe der Anmeldeinformationen eines Dienstprinzipals mit geringsten Berechtigungen nur auf die Ressourcen zugegriffen, die der Agent benötigt.

Funktionsweise des automatischen Authentifizierungs-Passthrough

Wenn ein Agent hinter einem Endpunkt mit automatischer Authentifizierungsweiterleitung eingesetzt wird, führt Databricks die folgenden Schritte aus:

1. Berechtigungsüberprüfung: Databricks überprüft, ob der Ersteller des Endpunkts auf alle Abhängigkeiten zugreifen kann, die während der Agentprotokollierung angegeben wurden.

2. Dienstprinzipalerstellung und -gewährung: Für die Agentmodellversion wird ein Dienstprinzipal erstellt und automatisch Lesezugriff auf Agentressourcen gewährt.

   Hinweis

   Der vom System generierte Dienstprinzipal wird nicht in API- oder UI-Listen angezeigt. Wenn die Version des Agent-Modells aus dem Endpunkt entfernt wird, wird der Dienstprinzipal ebenfalls gelöscht.

3. Bereitstellung und Rotation von Anmeldeinformationen: Kurzlebige Anmeldeinformationen (ein M2M OAuth-Token) für das Diensthauptkonto werden in den Endpunkt eingebunden, sodass Agent-Code auf Databricks-Ressourcen zugreifen kann. Databricks wechselt auch die Anmeldeinformationen, um sicherzustellen, dass Ihr Agent weiterhin sicheren Zugriff auf abhängige Ressourcen hat.

Unterstützte Ressourcen für die automatische Authentifizierungsdurchleitung

In der folgenden Tabelle sind die Databricks-Ressourcen aufgeführt, die die automatische Authentifizierungsdurchführung unterstützen, und die Berechtigungen, die der Ersteller des Endpoints bei der Bereitstellung des Agents besitzen muss.

Implementierung der automatischen Authentifizierungsdurchleitung

Zum Aktivieren der automatischen Authentifizierungsdurchlauf geben Sie abhängige Ressourcen an, wenn Sie den Agent protokollieren. Verwenden Sie den resources Parameter der log_model() API:

import mlflow
from mlflow.models.resources import (
  DatabricksVectorSearchIndex,
  DatabricksServingEndpoint,
  DatabricksSQLWarehouse,
  DatabricksFunction,
  DatabricksGenieSpace,
  DatabricksTable,
  DatabricksUCConnection,
  DatabricksApp,
  DatabricksLakebase
)

with mlflow.start_run():
  logged_agent_info = mlflow.pyfunc.log_model(
    python_model="agent.py",
    artifact_path="agent",
    input_example=input_example,
    example_no_conversion=True,
    # Specify resources for automatic authentication passthrough
    resources=[
      DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
      DatabricksServingEndpoint(endpoint_name="databricks-meta-llama-3-3-70b-instruct"),
      DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"),
      DatabricksSQLWarehouse(warehouse_id="your_warehouse_id"),
      DatabricksFunction(function_name="ml.tools.python_exec"),
      DatabricksGenieSpace(genie_space_id="your_genie_space_id"),
      DatabricksTable(table_name="your_table_name"),
      DatabricksUCConnection(connection_name="your_connection_name"),
      DatabricksApp(app_name="app_name"),
      DatabricksLakebase(database_instance_name="lakebase_instance_name"),
    ]
  )

Authentifizierung im Namen des Benutzers

Mit der Authentifizierung im Auftrag von Benutzern (OBO) kann ein Agent als Databricks-Benutzer fungieren, der die Abfrage ausführt. Dies bietet Folgendes:

• Benutzerspezifischer Zugriff auf vertrauliche Daten
• Feinkörnige Datensteuerelemente, die vom Unity-Katalog erzwungen werden
• Sicherheitstoken sind auf die VON Ihrem Agent deklarierten APIs beschränkt ("downscoped"), wodurch das Missbrauchsrisiko verringert wird.

Anforderungen

• Für die Authentifizierung im Auftrag von Benutzern ist MLflow 2.22.1 und höher erforderlich.
• Die Benutzerauthentifizierung im Auftrag des Benutzers ist standardmäßig deaktiviert und muss von einem Arbeitsbereichsadministrator aktiviert werden. Überprüfen Sie die Sicherheitsüberlegungen , bevor Sie dieses Feature aktivieren.

Unterstützte OBO-Ressourcen

Agents mit OBO-Authentifizierung können auf die folgenden Databricks-Ressourcen zugreifen:

Implementieren der OBO-Authentifizierung

Führen Sie die folgenden Schritte aus, um die Benutzerauthentifizierung im Auftrag von Benutzern zu aktivieren:

1. Aktualisieren Sie SDK-Aufrufe, um anzugeben, dass im Auftrag des Endbenutzers auf Ressourcen zugegriffen wird.
2. Aktualisieren Sie den Agentcode, um den OBO-Zugriff innerhalb der predict Funktion zu initialisieren, nicht in __init__, da die Benutzeridentität nur zur Laufzeit bekannt ist.
3. Deklarieren Sie bei der Protokollierung des Agents für die Bereitstellung die Databricks REST-API-Scopes, die der Agent benötigt.

Die folgenden Codeausschnitte veranschaulichen, wie der Benutzerzugriff im Auftrag des Benutzers auf verschiedene Databricks-Ressourcen konfiguriert wird. Behandeln Sie beim Initialisieren von Tools Berechtigungsfehler ordnungsgemäß, indem Sie die Initialisierung in einen try-except Block umschließen.

Vektorsuch-Retriever-Tool

from databricks.sdk import WorkspaceClient
from databricks_ai_bridge import ModelServingUserCredentials
from databricks_langchain import VectorSearchRetrieverTool

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy = ModelServingUserCredentials())

vector_search_tools = []
# Exclude exception handling if the agent should fail
# when users lack access to all required Databricks resources
try:
  tool = VectorSearchRetrieverTool(
    index_name="<index_name>",
    description="...",
    tool_name="...",
    workspace_client=user_client # Specify the user authorized client
    )
    vector_search_tools.append(tool)
except Exception as e:
    _logger.debug("Skipping adding tool as user does not have permissions")

Vektorsuchclient

from databricks.vector_search.client import VectorSearchClient
from databricks.vector_search.utils import CredentialStrategy

# Configure a VectorSearch Client to use on behalf of end
# user authentication
user_authenticated_vsc = VectorSearchClient(credential_strategy=CredentialStrategy.MODEL_SERVING_USER_CREDENTIALS)
# Exclude exception handling if the agent should fail when
# users lack access to all required Databricks resources
try:
  vs_index = user_authenticated_vsc.get_index(endpoint_name="endpoint_name", index_name="index_name")
  ...
except Exception as e:
  _logger.debug("Skipping Vector Index because user does not have permissions")

MCP

from databricks.sdk import WorkspaceClient
from databricks_ai_bridge import ModelServingUserCredentials
from databricks_mcp import DatabricksMCPClient

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy=ModelServingUserCredentials())

mcp_client = DatabricksMCPClient(
    server_url="<mcp_server_url>",
    workspace_client=user_client, # Specify the user client here
  )

Modellbereitstellungs-Endpunkt

from databricks.sdk import WorkspaceClient
from databricks_ai_bridge import ModelServingUserCredentials

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy=ModelServingUserCredentials())

# Exclude exception handling if the agent should fail
# when users lack access to all required Databricks resources
try:
  user_client.serving_endpoints.query("endpoint_name", input="")
except Exception as e:
  _logger.debug("Skipping Model Serving Endpoint due to no permissions")

UC Connections

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod
from databricks_ai_bridge import ModelServingUserCredentials

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy=ModelServingUserCredentials())

user_client.serving_endpoints.http_request(
  conn="connection_name",
  method=ExternalFunctionRequestHttpMethod.POST,
  path="/api/v1/resource",
  json={"key": "value"},
  headers={"extra_header_key": "extra_header_value"},
)

Genie Spaces (WorkspaceClient)

from databricks_langchain.genie import GenieAgent
from databricks.sdk import WorkspaceClient
from databricks_ai_bridge import ModelServingUserCredentials


# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy=ModelServingUserCredentials())


genie_agent = GenieAgent(
    genie_space_id="space-id",
    genie_agent_name="Genie",
    description="This Genie space has access to sales data in Europe"
    client=user_client
)

# Use the Genie SDK methods available through WorkspaceClient
try:
    response = agent.invoke("Your query here")
except Exception as e:
    _logger.debug("Skipping Genie due to no permissions")

Genie Spaces (LangChain)

from databricks.sdk import WorkspaceClient
from databricks_ai_bridge import ModelServingUserCredentials
from databricks_langchain.genie import GenieAgent

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy=ModelServingUserCredentials())

genie_agent = GenieAgent(
    genie_space_id="<genie_space_id>",
    genie_agent_name="Genie",
    description="Genie_description",
    client=user_client, # Specify the user client here
  )

Initialisieren des Agenten in der Vorhersagefunktion

Da die Identität des Benutzers nur zur Abfragezeit bekannt ist, müssen Sie auf OBO-Ressourcen innerhalb predict oder predict_streamnicht in der Methode des __init__ Agents zugreifen. Dadurch wird sichergestellt, dass Ressourcen zwischen Aufrufen isoliert sind.

from mlflow.pyfunc import ResponsesAgent

class OBOResponsesAgent(ResponsesAgent):
  def initialize_agent():
    user_client = WorkspaceClient(
      credentials_strategy=ModelServingUserCredentials()
    )
    system_authorized_client = WorkspaceClient()
    ### Use the clients above to access resources with either system or user authentication

  def predict(
    self, request
  ) -> ResponsesAgentResponse:
    agent = initialize_agent() # Initialize the Agent in Predict

    agent.predict(request)
    ...

Deklarieren von REST-API-Bereichen beim Protokollieren des Agents

Wenn Sie Ihren OBO-Agent für die Bereitstellung protokollieren, müssen Sie die Databricks-REST-API-Bereiche auflisten, die Ihr Agent im Namen des Benutzers aufruft. Dadurch wird sichergestellt, dass der Agent dem Prinzip der geringsten Berechtigungen folgt: Token sind nur auf die APIs beschränkt, die Ihr Agent benötigt, wodurch die Wahrscheinlichkeit nicht autorisierter Aktionen oder Tokenmissbrauch reduziert wird.

Nachfolgend finden Sie eine Liste der Bereiche, die für den Zugriff auf mehrere gängige Datentypen von Databricks-Ressourcen erforderlich sind:

Um die Authentifizierung im Auftrag von Benutzern zu aktivieren, übergeben Sie einen MLflow AuthPolicy an log_model():

import mlflow
from mlflow.models.auth_policy import AuthPolicy, SystemAuthPolicy, UserAuthPolicy
from mlflow.models.resources import DatabricksServingEndpoint

# System policy: resources accessed with system credentials
system_policy = SystemAuthPolicy(
    resources=[DatabricksServingEndpoint(endpoint_name="my_endpoint")]
)

# User policy: API scopes for OBO access
user_policy = UserAuthPolicy(api_scopes=[
    "serving.serving-endpoints",
    "vectorsearch.vector-search-endpoints",
    "vectorsearch.vector-search-indexes"
])

# Log the agent with both policies
with mlflow.start_run():
    mlflow.pyfunc.log_model(
        name="agent",
        python_model="agent.py",
        auth_policy=AuthPolicy(
            system_auth_policy=system_policy,
            user_auth_policy=user_policy
        )
    )

OBO-Authentifizierung für OpenAI-Clients

Verwenden Sie für Agents, die den OpenAI-Client verwenden, das Databricks SDK, um sich während der Bereitstellung automatisch zu authentifizieren. Databricks SDK verfügt über einen Wrapper zum Erstellen des OpenAI-Clients mit automatisch konfigurierter Authentifizierung: get_open_ai_client()

% pip install databricks-sdk[openai]

from databricks.sdk import WorkspaceClient
def openai_client(self):
  w = WorkspaceClient()
  return w.serving_endpoints.get_open_ai_client()

Geben Sie dann den Model Serving-Endpunkt als Teil von resources an, um sich automatisch bei der Bereitstellung zu authentifizieren.

Überlegungen zur Sicherheit von OBO

Berücksichtigen Sie die folgenden Sicherheitsaspekte, bevor Sie die Authentifizierung im Auftrag des Benutzers mit Agents aktivieren.

Erweiterter Ressourcenzugriff: Agents können im Auftrag von Benutzern auf vertrauliche Ressourcen zugreifen. Während Bereiche APIs einschränken, erlauben Endpunkte möglicherweise mehr Aktionen als der Agent explizit anfordert. Beispielsweise gewährt der serving.serving-endpoints API-Bereich einem Agent die Berechtigung zum Ausführen eines dienstenden Endpunkts im Namen des Benutzers. Der dienstende Endpunkt kann jedoch auf zusätzliche API-Bereiche zugreifen, die der ursprüngliche Agent nicht für die Verwendung autorisiert ist.

OBO-Beispiel-Notebooks

Das folgende Notizbuch zeigt, wie Sie mithilfe der On-Behalf-of-User-Autorisierung einen Agenten mit Vektorsuche erstellen.

Im Auftrag der Benutzerautorisierung mit Vektorsuche

Notebook abrufen

Das folgende Notebook zeigt Ihnen, wie Sie einen Agenten erstellen, der die SQL-Ausführung in einem SQL Warehouse unter Verwendung der On-Behalf-of-User-Autorisierung unterstützt. Auf diese Weise kann der Agent Unity Catalog-Funktionen mithilfe von Benutzeranmeldeinformationen sicher aufrufen. Hinweis: Dies ist derzeit die empfohlene Methode zum Ausführen von UC-Funktionen mit OBO, da Serverless Spark Execution mit OBO noch nicht unterstützt wird.

Im Namen der Benutzerautorisierung mit SQL-Befehlen

Notebook abrufen

Manuelle Authentifizierung

Die manuelle Authentifizierung ermöglicht die explizite Angabe von Anmeldeinformationen während der Agentbereitstellung. Diese Methode hat die größte Flexibilität, erfordert jedoch mehr Einrichtung und fortlaufende Verwaltung von Anmeldeinformationen. Verwenden Sie diese Methode, wenn:

• Die abhängige Ressource unterstützt keine automatische Authentifizierungsweitergabe.
• Der Agent muss andere Anmeldeinformationen als die des Agent-Deployers verwenden.
• Der Agent greift auf externe Ressourcen oder APIs außerhalb von Databricks zu
• Der bereitgestellte Agent greift auf das Prompt-Register zu.

OAuth-Authentifizierung (empfohlen)

OAuth ist der empfohlene Ansatz für die manuelle Authentifizierung, da sie über eine sichere, tokenbasierte Authentifizierung für Dienstprinzipale mit automatischen Tokenaktualisierungsfunktionen verfügt:

1. Erstellen Sie einen Dienstprinzipal und generieren Sie OAuth-Anmeldeinformationen.

2. Gewähren Sie dem Dienstprinzipal Berechtigungen für jede Databricks-Ressource, auf die der Agent privilegierten Zugriff auf Databricks-Ressourcen hat. Um auf das Eingabeaufforderungsprotokoll zuzugreifen, erteilen Sie CREATE FUNCTION, EXECUTE und MANAGE Berechtigungen für das Unity Katalog-Schema zum Speichern von Eingabeaufforderungen.

3. Erstellen Sie Databricks Secrets für die OAuth-Anmeldeinformationen.

4. Konfigurieren Sie die OAuth-Anmeldeinformationen im Agentcode:

   import os
   
   # Configure OAuth authentication for Prompt Registry access
   # Replace with actual secret scope and key names
   secret_scope_name = "your-secret-scope"
   client_id_key = "oauth-client-id"
   client_secret_key = "oauth-client-secret"
   
   os.environ["DATABRICKS_HOST"] = "https://<your-workspace-url>"
   os.environ["DATABRICKS_CLIENT_ID"] = dbutils.secrets.get(scope=secret_scope_name, key=client_id_key)
   os.environ["DATABRICKS_CLIENT_SECRET"] = dbutils.secrets.get(scope=secret_scope_name, key=client_secret_key)
   

5. Verwenden Sie die geheimen Schlüssel, um eine Verbindung mit dem Arbeitsbereich herzustellen:

   w = WorkspaceClient(
     host=os.environ["DATABRICKS_HOST"],
     client_id=os.environ["DATABRICKS_CLIENT_ID"],
     client_secret = os.environ["DATABRICKS_CLIENT_SECRET"]
   )
   

6. Bei der Bereitstellung mit agents.deploy() sollten die OAuth-Anmeldeinformationen als Umgebungsvariablen eingeschlossen werden.

   agents.deploy(
       UC_MODEL_NAME,
       uc_registered_model_info.version,
       environment_vars={
           "DATABRICKS_HOST": "https://<your-workspace-url>",
           "DATABRICKS_CLIENT_ID": f"{{{{secrets/{secret_scope_name}/{client_id_key}}}}}",
           "DATABRICKS_CLIENT_SECRET": f"{{{{secrets/{secret_scope_name}/{client_secret_key}}}}}"
       },
   )
   

PAT-Authentifizierung

Die Authentifizierung für persönliche Zugriffstoken (Personal Access Token, PAT) bietet ein einfacheres Setup für Entwicklungs- und Testumgebungen, erfordert jedoch eine manuelle Verwaltung von Anmeldeinformationen:

1. Abrufen eines PAT mithilfe eines Dienstprinzipals oder persönlichen Kontos:

   Dienstprinzipal (für Sicherheit empfohlen):

   1. Erstellen eines Dienstprinzipals
   2. Gewähren Sie dem Dienstprinzipal Berechtigungen für jede Databricks-Ressource, auf die der Agent privilegierten Zugriff auf Databricks-Ressourcen hat. Um auf das Aufforderungsregister zuzugreifen, erteilen Sie CREATE FUNCTION, EXECUTE und MANAGE Berechtigungen für das Unity-Katalogschema, das zum Speichern von Aufforderungen verwendet wird.
   3. Erstellen Sie einen PAT für den Dienstprinzipal.

   Persönliches Konto:

   1. Erstellen Sie einen PAT für ein persönliches Konto.

2. Speichern Sie den PAT sicher, indem Sie ein Databricks-Geheimnis für das PAT erstellen.

3. Konfigurieren der PAT-Authentifizierung im Agentcode:

   import os
   
   # Configure PAT authentication for Prompt Registry access
   # Replace with your actual secret scope and key names
   secret_scope_name = "your-secret-scope"
   secret_key_name = "your-pat-key"
   
   os.environ["DATABRICKS_HOST"] = "https://<your-workspace-url>"
   os.environ["DATABRICKS_TOKEN"] = dbutils.secrets.get(scope=secret_scope_name, key=secret_key_name)
   
   # Validate configuration
   assert os.environ["DATABRICKS_HOST"], "DATABRICKS_HOST must be set"
   assert os.environ["DATABRICKS_TOKEN"], "DATABRICKS_TOKEN must be set"
   

4. Schließen Sie bei der Bereitstellung des Agents agents.deploy()den PAT als Umgebungsvariable ein:

   agents.deploy(
       UC_MODEL_NAME,
       uc_registered_model_info.version,
       environment_vars={
           "DATABRICKS_HOST": "https://<your-workspace-url>",
           "DATABRICKS_TOKEN": f"{{{{secrets/{secret_scope_name}/{secret_key_name}}}}}"
       },
   )