Logowanie i rejestrowanie agentów sztucznej inteligencji

Rejestrowanie agentów sztucznej inteligencji przy użyciu Mosaic AI Agent Framework. Rejestrowanie agenta jest podstawą procesu programowania. Rejestrowanie zapisuje "punkt w czasie" kodu i konfiguracji agenta, co pozwala ocenić jakość konfiguracji.

Wymagania

Przed zalogowaniem utwórz agenta sztucznej inteligencji.

Databricks zaleca zainstalowanie najnowszej wersji databricks-sdk.

% pip install databricks-sdk

rejestrowanie oparte na kodzie

Databricks zaleca korzystanie z funkcjonalności 'Modele z kodu' w MLflow podczas rejestrowania agentów.

W tym podejściu kod agenta jest przechwytywany jako plik w języku Python, a środowisko języka Python jest przechwytywane jako lista pakietów. Po wdrożeniu agenta środowisko języka Python zostanie przywrócone, a kod agenta zostanie wykonany w celu załadowania agenta do pamięci, aby można go było wywołać po wywołaniu punktu końcowego.

Takie podejście można połączyć z użyciem interfejsów API do weryfikacji przed wdrożeniem, takich jak mlflow.models.predict(), aby upewnić się, że agent działa niezawodnie, gdy jest wdrażany do obsługi.

Aby zobaczyć przykład rejestrowania opartego na kodzie, zobacz ChatAgent authoring example notebooks.

Określanie sygnatury modelu podczas zapisywania

Uwaga

Databricks zaleca tworzenie agenta za pomocą interfejsu ChatAgent. Jeśli korzystasz z aplikacji ChatAgent, możesz pominąć tę sekcję; Narzędzie MLflow automatycznie wywnioskuje prawidłowy podpis agenta.

Jeśli nie używasz interfejsu ChatAgent, musisz skorzystać z jednej z poniższych metod, aby określić Sygnaturę Modelu MLflow swojego agenta w momencie rejestrowania.

1. Ręczne definiowanie podpisu
2. Użyj funkcji wnioskowania sygnatury modelu MLflow, aby automatycznie wygenerować podpis agenta na podstawie podanego przykładu danych wejściowych. Takie podejście jest wygodniejsze niż ręczne definiowanie podpisu.

Podpis modelu MLflow weryfikuje dane wejściowe i wyjściowe, aby upewnić się, że agent prawidłowo współdziała z narzędziami podrzędnymi, takimi jak AI Playground i aplikacja do przeglądu. Wskazuje również innym aplikacjom, jak efektywnie korzystać z agenta.

Poniższe przykłady LangChain i PyFunc wykorzystują inferencję sygnatury modelu.

Jeśli wolisz jawnie zdefiniować podpis modelu w czasie rejestrowania, zobacz dokumentacji platformy MLflow — jak rejestrować modele z podpisami.

Rejestrowanie oparte na kodzie za pomocą usługi LangChain

Poniższe instrukcje i przykładowy kod pokazują, jak zarejestrować agenta za pomocą pakietu LangChain.

1. Utwórz notes lub plik w języku Python przy użyciu kodu. W tym przykładzie notes lub plik ma nazwę agent.py. Notatnik lub plik musi zawierać agenta LangChain, określanego tutaj jako lc_agent.

2. Uwzględnij mlflow.models.set_model (lc_agent) w notesie lub pliku.

3. Utwórz nowy notes, który będzie służył jako notes sterownika (nazywany driver.py w tym przykładzie).

4. W notesie sterownika użyj następującego kodu, aby uruchomić agent.py i zarejestrować wyniki w modelu MLflow:

   mlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)
   

   Parametr resources deklaruje zasoby zarządzane przez usługę Databricks potrzebne do obsługi agenta, takie jak indeks wyszukiwania wektorowego lub obsługujący punkt końcowy obsługujący model podstawowy. Aby uzyskać dalsze informacje, zobacz Określanie zasobów na potrzeby przekazywania automatycznego uwierzytelniania.

5. Wdróż model. Zobacz Wdrażanie agenta na potrzeby generowania aplikacji sztucznej inteligencji.

6. Po załadowaniu środowiska serwującego, wykonywane jest agent.py.

7. Gdy pojawia się żądanie usługi, lc_agent.invoke(...) jest wywoływane.

import mlflow

code_path = "/Workspace/Users/first.last/agent.py"
config_path = "/Workspace/Users/first.last/config.yml"

# Input example used by MLflow to infer Model Signature
input_example = {
  "messages": [
    {
      "role": "user",
      "content": "What is Retrieval-augmented Generation?",
    }
  ]
}

# example using langchain
with mlflow.start_run():
  logged_agent_info = mlflow.langchain.log_model(
    lc_model=code_path,
    model_config=config_path, # If you specify this parameter, this configuration is used by agent code. The development_config is overwritten.
    artifact_path="agent", # This string is used as the path inside the MLflow model where artifacts are stored
    input_example=input_example, # Must be a valid input to the agent
    example_no_conversion=True, # Required
  )

print(f"MLflow Run: {logged_agent_info.run_id}")
print(f"Model URI: {logged_agent_info.model_uri}")

# To verify that the model has been logged correctly, load the agent and call `invoke`:
model = mlflow.langchain.load_model(logged_agent_info.model_uri)
model.invoke(example)

rejestrowanie oparte na kodzie za pomocą funkcji PyFunc

Poniższe instrukcje i przykładowy kod pokazują, jak zarejestrować agenta za pomocą narzędzia PyFunc.

1. Utwórz notes lub plik w języku Python przy użyciu kodu. W tym przykładzie notes lub plik ma nazwę agent.py. Notatnik lub plik musi zawierać klasę PyFunc o nazwie PyFuncClass.

2. Dołącz mlflow.models.set_model(PyFuncClass) do notesu lub pliku.

3. Utwórz nowy notes, który będzie służył jako notes sterownika (nazywany driver.py w tym przykładzie).

4. W notesie sterownika użyj następującego kodu, aby uruchomić agent.py i zarejestrować wyniki w modelu MLflow:

   mlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)
   

   Parametr resources deklaruje zasoby zarządzane przez usługę Databricks potrzebne do obsługi agenta, takie jak indeks wyszukiwania wektorowego lub obsługujący punkt końcowy obsługujący model podstawowy. Więcej informacji można znaleźć w Określaniu zasobów dla automatycznego przekazywania uwierzytelnienia.

5. Wdróż model. Zobacz Wdrażanie agenta na potrzeby generowania aplikacji sztucznej inteligencji.

6. Po załadowaniu środowiska serwującego, wykonywane jest agent.py.

7. Gdy pojawia się żądanie obsługi, PyFuncClass.predict(...) jest wywoływane.

import mlflow
from mlflow.models.resources import (
    DatabricksServingEndpoint,
    DatabricksVectorSearchIndex,
)

code_path = "/Workspace/Users/first.last/agent.py"
config_path = "/Workspace/Users/first.last/config.yml"

# Input example used by MLflow to infer Model Signature
input_example = {
  "messages": [
    {
      "role": "user",
      "content": "What is Retrieval-augmented Generation?",
    }
  ]
}

with mlflow.start_run():
  logged_agent_info = mlflow.pyfunc.log_model(
    python_model=agent_notebook_path,
    artifact_path="agent",
    input_example=input_example,
    resources=resources_path,
    example_no_conversion=True,
    resources=[
      DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
      DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
    ]
  )

print(f"MLflow Run: {logged_agent_info.run_id}")
print(f"Model URI: {logged_agent_info.model_uri}")

# To verify that the model has been logged correctly, load the agent and call `invoke`:
model = mlflow.pyfunc.load_model(logged_agent_info.model_uri)
model.invoke(example)

Określ zasoby do automatycznego przekazywania danych uwierzytelniających

Agenci sztucznej inteligencji często muszą uwierzytelniać się w innych zasobach, aby wykonywać zadania. Na przykład agent może potrzebować dostępu do indeksu wyszukiwania wektorowego w celu wykonywania zapytań dotyczących danych bez struktury.

Zgodnie z opisem w Authentication for dependent resources, funkcja Model Serving obsługuje uwierzytelnianie zarówno w zasobach zarządzanych przez Databricks, jak i w zewnętrznych zasobach podczas wdrażania agenta.

W przypadku najczęściej spotykanych typów zasobów Databricks obsługuje i zaleca deklarowanie zależności zasobów dla agenta już na etapie rejestrowania. To umożliwia automatyczne przekazywanie uwierzytelniania podczas wdrażania agenta, a usługa Databricks automatycznie przydziela, rotuje i zarządza krótkotrwałymi poświadczeniami w celu bezpiecznego uzyskania dostępu do tych zależności zasobów z poziomu punktu końcowego agenta.

Aby włączyć automatyczne uwierzytelnianie, określ zasoby zależne, używając parametru resources w API log_model(), jak pokazano w poniższym kodzie.

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

with mlflow.start_run():
  logged_agent_info = mlflow.pyfunc.log_model(
    python_model=agent_notebook_path,
    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-mixtral-8x7b-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"),
    ]
  )

Databricks zaleca ręczne określenie resources dla wszystkich wersji agentów.

Uwaga

Jeśli nie określisz zasobów podczas rejestrowania agentów LangChain przy użyciu mlflow.langchain.log_model(...), narzędzie MLflow podejmuje maksymalne starania w celu automatycznego wnioskowania zasobów. Może to jednak nie przechwytywać wszystkich zależności, co powoduje błędy autoryzacji podczas obsługi lub wykonywania zapytań dotyczących agenta.

W poniższej tabeli wymieniono zasoby usługi Databricks, które obsługują automatyczne przekazywanie uwierzytelniania, oraz minimalną wersję mlflow wymaganą do logowania zasobu.

Typ zasobu	Minimalna wersja mlflow wymagana do zarejestrowania zasobu
Indeks wyszukiwania wektorowego	Wymaga mlflow 2.13.1 lub nowszej
Punkt końcowy obsługujący model	Wymaga mlflow 2.13.1 lub nowszej
SQL Warehouse	Wymaga mlflow 2.16.1 lub nowszego
Unity Catalog, funkcja	Wymaga mlflow 2.16.1 lub nowszego
Przestrzeń Genie	Wymaga mlflow 2.17.1 lub nowszej
Tabela katalogu Unity	Wymaga mlflow 2.18.0 lub nowszego
Połączenie z katalogiem Unity	Wymaga mlflow 2.17.1 lub nowszej

Automatyczne uwierzytelnianie dla klientów openAI

Jeśli agent używa klienta OpenAI, użyj zestawu SDK usługi Databricks do automatycznego uwierzytelniania podczas wdrażania. Zestaw SDK usługi Databricks udostępnia warstwę do tworzenia klienta OpenAI z automatycznie skonfigurowaną autoryzacją. Uruchom poniższe w notebooku:

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

Zarejestrować agenta w Unity Catalog

Przed wdrożeniem agenta należy go zarejestrować w Unity Catalog. Rejestrowanie agenta tworzy z niego model w katalogu Unity. Dzięki temu można użyć uprawnień katalogu Unity do autoryzacji zasobów w agencie.

import mlflow

mlflow.set_registry_uri("databricks-uc")

catalog_name = "test_catalog"
schema_name = "schema"
model_name = "agent_name"

model_name = catalog_name + "." + schema_name + "." + model_name
uc_model_info = mlflow.register_model(model_uri=logged_agent_info.model_uri, name=model_name)

Następne kroki

• Dodaj ślady agentowi AI.
• Wdrażanie agenta sztucznej inteligencji.