Uwierzytelnianie agentów sztucznej inteligencji

Agenci sztucznej inteligencji często muszą uwierzytelniać się w innych zasobach, aby wykonywać zadania. Na przykład wdrożony agent może wymagać dostępu do indeksu wyszukiwania wektorowego w celu wykonywania zapytań dotyczących danych bez struktury lub do Rejestru monitu w celu wczytania dynamicznych monitów.

Na tej stronie omówiono metody uwierzytelniania dostępne podczas opracowywania i wdrażania agentów przy użyciu platformy agentów Mozaiki AI.

Metody uwierzytelniania

W poniższej tabeli porównaliśmy dostępne metody uwierzytelniania. Można mieszać i dopasowywać dowolne z tych podejść:

Uwierzytelnianie w systemach zewnętrznych i serwerach MCP

Aby uzyskać wskazówki dotyczące uwierzytelniania w zewnętrznych interfejsach API i serwerach MCP z agenta, zobacz Connect AI agent tools to external services (Łączenie narzędzi agenta sztucznej inteligencji z usługami zewnętrznymi). Te zasoby mogą być również odpytywane na rzecz agenta lub użytkownika, jak opisano w Metodach uwierzytelniania.

Wybieranie właściwej metody uwierzytelniania dla zasobu

Użyj tego schematu blokowego, aby wybrać odpowiednią metodę uwierzytelniania dla każdego zasobu. W zależności od potrzeb można łączyć metody, a agent może użyć innej metody dla każdego zasobu w zależności od jego przypadku użycia.

1. Czy wymagana jest kontrola dostępu dla każdego użytkownika lub audyt z przypisaniem do użytkownika?

   • Tak → użyj uwierzytelniania w imieniu użytkownika
   • Brak → przejdź do kroku 2

2. Czy wszystkie zasoby obsługują automatyczne uwierzytelnianie?

   • Tak → Użyj Przekazywanie Automatycznego Uwierzytelniania (zalecane)
   • Brak → korzystanie z uwierzytelniania ręcznego

Automatyczne przekazywanie uwierzytelniania

Automatyczne przekazywanie uwierzytelnienia jest najprostszą metodą uzyskiwania dostępu do zasobów zarządzanych przez Databricks. Zadeklaruj zależności zasobów podczas logowania agenta, a usługa Databricks automatycznie zapewnia, rotuje i zarządza tymczasowymi poświadczeniami podczas wdrażania agenta.

To zachowanie uwierzytelniania jest podobne do zachowania "Uruchom jako właściciel" dla pulpitów nawigacyjnych Databricks. Zasoby podrzędne, takie jak tabele Unity Catalog, są dostępne przy użyciu poświadczeń głównego serwisu o najniższych uprawnieniach dostępowych tylko do zasobów, których potrzebuje agent.

Jak działa automatyczne przekazywanie uwierzytelniania

Gdy agent jest obsługiwany przy użyciu automatycznego przekazywania uwierzytelnienia na poziomie punktu końcowego, usługa Databricks wykonuje następujące kroki:

1. Weryfikacja uprawnień: usługa Databricks sprawdza, czy twórca punktu końcowego może uzyskać dostęp do wszystkich zależności określonych podczas rejestrowania agenta.

2. Tworzenie i przydzielanie obiektu zabezpieczeń: Obiekt zabezpieczeń jest tworzony dla wersji modelu agenta i jest automatycznie przyznawany dostęp do odczytu do zasobów agenta.

   Uwaga / Notatka

   Systemowo wygenerowana główna usługa nie pojawia się na listach API ani interfejsu użytkownika. Jeśli wersja modelu agenta zostanie usunięta z punktu końcowego, główny składnik usługi zostanie również usunięty.

3. Aprowizowanie i rotacja poświadczeń: poświadczenia krótkotrwałe (token OAuth M2M) dla jednostki usługi są wstrzykiwane do punktu końcowego, co umożliwia kodowi agenta uzyskiwanie dostępu do zasobów Databricks. Usługa Databricks zmienia również poświadczenia, aby zapewnić agentowi ciągły i bezpieczny dostęp do zasobów zależnych.

Obsługiwane zasoby na potrzeby przekazywania uwierzytelniania automatycznego

W poniższej tabeli wymieniono zasoby Databricks, obsługujące automatyczne przekazywanie uwierzytelniania, oraz uprawnienia, które twórca punktu końcowego musi posiadać podczas wdrażania agenta.

Implementowanie automatycznego przekazywania uwierzytelniania

Aby włączyć automatyczne przekazywanie uwierzytelnienia, określ zasoby zależne podczas rejestrowania agenta. Użyj parametru resources API interfejsu log_model().

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"),
    ]
  )

Uwierzytelnianie w imieniu użytkownika

Uwierzytelnianie w imieniu użytkownika (OBO) umożliwia agentowi działanie jako użytkownik platformy Databricks, który uruchamia zapytanie. Zapewnia to:

• Dostęp poszczególnych użytkowników do poufnych danych
• Szczegółowe kontrolki danych wymuszane przez wykaz aparatu Unity
• Tokeny zabezpieczające są ograniczone (zredukowane zakresowo) tylko do interfejsów API zadeklarowanych przez agenta, co zmniejsza ryzyko ich niewłaściwego użycia.

Requirements

• Uwierzytelnianie w imieniu użytkownika wymaga środowiska MLflow 2.22.1 lub nowszego.
• Uwierzytelnianie w imieniu użytkownika jest domyślnie wyłączone i musi być włączone przez administratora obszaru roboczego. Przed włączeniem tej funkcji zapoznaj się z zagadnieniami dotyczącymi zabezpieczeń .

Zasoby obsługiwane przez OBO

Agenci z uwierzytelnianiem OBO mogą uzyskać dostęp do następujących zasobów usługi Databricks:

Implementowanie uwierzytelniania OBO

Aby włączyć uwierzytelnianie w imieniu użytkownika, wykonaj następujące kroki:

1. Zaktualizuj wywołania zestawu SDK, aby określić, że zasoby są dostępne w imieniu użytkownika końcowego.
2. Zaktualizuj kod agenta, aby zainicjować dostęp OBO wewnątrz predict funkcji, a nie w __init__systemie , ponieważ tożsamość użytkownika jest znana tylko w czasie wykonywania.
3. Podczas rejestrowania agenta do wdrożenia zadeklaruj zakresy interfejsu API REST usługi Databricks wymagane przez agenta.

Poniższe fragmenty kodu pokazują, jak skonfigurować dostęp w imieniu użytkownika do różnych zasobów usługi Databricks. Podczas inicjowania narzędzi zadbaj o łagodne obsługiwanie błędów uprawnień, zawijając inicjowanie w bloku try-except.

Narzędzie do wyszukiwania wektorowego

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")

Klient wyszukiwania wektorowego

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
  )

Punkt końcowy obsługi modelu

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")

Połączenia UC

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
  )

Inicjowanie agenta w funkcji predict

Ponieważ tożsamość użytkownika jest znana tylko w czasie zapytania, musisz uzyskać dostęp do zasobów OBO wewnątrz predict lub predict_stream, a nie w metodzie agenta __init__ . Dzięki temu zasoby są izolowane między wywołaniami.

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)
    ...

Deklarowanie zakresów interfejsu API REST podczas rejestrowania agenta

Podczas rejestrowania agenta OBO do wdrożenia należy wyświetlić listę zakresów interfejsu API REST usługi Databricks, które agent wywołuje w imieniu użytkownika. Dzięki temu agent jest zgodny z zasadą najniższych uprawnień: tokeny są ograniczone tylko do interfejsów API, których wymaga agent, co zmniejsza prawdopodobieństwo nieautoryzowanego działania lub nieprawidłowego użycia tokenu.

Poniżej znajduje się lista zakresów wymaganych do uzyskania dostępu do kilku typowych typów zasobów usługi Databricks:

Aby włączyć uwierzytelnianie w imieniu użytkownika, przekaż MLflow AuthPolicy do 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
        )
    )

Uwierzytelnianie OBO dla klientów OpenAI

W przypadku agentów korzystających z klienta OpenAI użyj zestawu SDK usługi Databricks do automatycznego uwierzytelniania podczas wdrażania. Zestaw SDK Databricks ma opakowanie, które automatycznie konfiguruje uwierzytelnianie przy konstruowaniu klienta OpenAI: 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()

Następnie określ punkt końcowy obsługujący model w ramach resources do automatycznego uwierzytelniania w czasie wdrażania.

Zagadnienia dotyczące zabezpieczeń OBO

Przed włączeniem uwierzytelniania w imieniu użytkownika przy użyciu agentów należy wziąć pod uwagę następujące zagadnienia dotyczące zabezpieczeń.

Rozszerzony dostęp do zasobów: Agenci mogą uzyskiwać dostęp do poufnych zasobów w imieniu użytkowników. Chociaż zakresy ograniczają interfejsy API, punkty końcowe mogą zezwalać na więcej akcji niż jawne żądania agenta. Na przykład zakres interfejsu serving.serving-endpoints API przyznaje agentowi uprawnienia do uruchamiania punktu końcowego obsługującego w imieniu użytkownika. Jednak punkt końcowy obsługujący może uzyskać dostęp do dodatkowych zakresów interfejsu API, których oryginalny agent nie ma autoryzacji do użycia.

Przykładowe notesy OBO

W poniższym notatniku pokazano, jak utworzyć agenta z funkcją wyszukiwania wektorowego przy użyciu autoryzacji w imieniu użytkownika.

W imieniu użytkownika autoryzacji za pomocą wyszukiwania wektorowego

Pobierz laptopa

W poniższym notesie pokazano, jak utworzyć agenta obsługującego wykonywanie SQL w usłudze SQL Warehouse przy użyciu autoryzacji w imieniu użytkownika. Dzięki temu agent może bezpiecznie wywoływać funkcje Unity Catalog przy użyciu danych uwierzytelniających użytkownika. Uwaga: jest to obecnie zalecany sposób wykonywania funkcji UC z OBO, ponieważ bezserwerowa realizacja platformy Spark z OBO nie jest jeszcze obsługiwana.

Autoryzacja w imieniu użytkownika przy użyciu wykonania poleceń SQL

Pobierz laptopa

Uwierzytelnianie ręczne

Uwierzytelnianie ręczne pozwala na wyraźne określenie poświadczeń podczas wdrażania agenta. Ta metoda ma największą elastyczność, ale wymaga większej konfiguracji i ciągłego zarządzania poświadczeniami. Użyj tej metody, gdy:

• Zasób zależny nie obsługuje automatycznego przekazywania uwierzytelnienia
• Agent musi używać poświadczeń innych niż poświadczenia osoby wdrażającej agenta
• Agent uzyskuje dostęp do zasobów zewnętrznych lub interfejsów API poza usługą Databricks
• Wdrożony agent uzyskuje dostęp do rejestru monitów

Uwierzytelnianie OAuth (zalecane)

Uwierzytelnianie OAuth jest zalecanym podejściem do uwierzytelniania ręcznego, ponieważ ma bezpieczne uwierzytelnianie oparte na tokenach dla jednostek usługi z funkcjami automatycznego odświeżania tokenów:

1. Utwórz jednostkę usługi i wygeneruj poświadczenia OAuth.

2. Przyznaj jednostce usługi uprawnienia do dowolnego zasobu usługi Databricks, do którego agent ma dostęp do uprawnień dostępu do zasobów usługi Databricks. Aby uzyskać dostęp do rejestru monitów, udziel CREATE FUNCTION, EXECUTE, i MANAGE uprawnień w schemacie Unity Catalog do przechowywania monitów.

3. Utwórz tajne dane Databricks dla poświadczeń OAuth.

4. Skonfiguruj poświadczenia OAuth w kodzie agenta:

   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. Użyj tajnych danych, aby połączyć się ze środowiskiem pracy.

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

6. Podczas wdrażania z użyciem agents.deploy(), dołącz poświadczenia OAuth jako zmienne środowiskowe:

   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}}}}}"
       },
   )
   

Uwierzytelnianie za pomocą tokenu pat

Uwierzytelnianie osobistego tokenu dostępu (PAT) zapewnia prostszą konfigurację środowisk programistycznych i testowych, chociaż wymaga bardziej ręcznego zarządzania poświadczeniami:

1. Zdobyć token dostępu osobistego (PAT) przy użyciu jednostki usługi lub konta osobistego:

   Service principal (zalecane dla bezpieczeństwa):

   1. Utwórz jednostkę usługi.
   2. Przyznaj jednostce usługi uprawnienia do dowolnego zasobu usługi Databricks, do którego agent ma dostęp do uprawnień dostępu do zasobów usługi Databricks. Aby uzyskać dostęp do rejestru monitów, przyznaj CREATE FUNCTION, EXECUTE i MANAGE uprawnienia do schematu Unity Catalogu używanego do przechowywania monitów.
   3. Utwórz identyfikator PAT dla jednostki usługi.

   Konto osobiste:

   1. Utwórz osobisty token dostępu dla konta.

2. Bezpieczne przechowywanie kodu PAT przez utworzenie wpisu tajnego usługi Databricks dla PAT.

3. Skonfiguruj uwierzytelnianie pat w kodzie agenta:

   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. Podczas wdrażania agenta przy użyciu polecenia agents.deploy()dołącz pat jako zmienną środowiskową:

   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}}}}}"
       },
   )