Ведение журнала и регистрация агентов ИИ

Статья
02/27/2025

Регистрируйте агентов ИИ с помощью платформы Mosaic AI Agent Framework. Логирование агента является основой процесса разработки. Ведение журнала фиксирует "точку во времени" кода и конфигурации агента, чтобы вы могли оценить качество конфигурации.

Требования

Создайте ИИ-агент перед его регистрацией.

Databricks рекомендует установить последнюю версию databricks-sdk.

% pip install databricks-sdk

ведение журнала на основе кода

Databricks рекомендует использовать модели MLflow из функционала кода при ведении журналов агентов.

В этом подходе код агента записывается в виде файла Python, а среда Python записывается в виде списка пакетов. При развертывании агента среда Python восстанавливается, а код агента выполняется для загрузки агента в память, чтобы его можно было вызвать при вызове конечной точки.

Этот подход можно использовать с использованием API для предварительной проверки развертывания, таких как mlflow.models.predict(), чтобы убедиться, что агент работает надежно при развертывании для обслуживания.

Чтобы увидеть пример ведения журнала на основе кода, см. пример записных книжек для разработки ChatAgent.

Определение сигнатуры модели во время логирования

Примечание.

Databricks рекомендует создавать агент с помощью интерфейса ChatAgent. При использовании ChatAgent можно пропустить этот раздел; MLflow автоматически выводит допустимую подпись для агента.

Если вы не используете интерфейс ChatAgent, необходимо применить один из следующих методов, чтобы указать сигнатуру модели MLflow вашего агента при записи в журнал:

Вручную определите сигнатуру
Используйте возможности вывода подписи модели MLflow для автоматического создания подписи агента на основе предоставленного входного примера. Этот подход удобнее, чем вручную определение сигнатуры.

Подпись модели MLflow проверяет входные и выходные данные, чтобы агент правильно взаимодействовал с последующими инструментами, такими как AI Playground и приложение для обзора. Он также поможет другим приложениям эффективно использовать агент.

Примеры LangChain и PyFunc ниже используют вывод на основе сигнатуры модели.

Если вы предпочитаете сами явно определять подпись модели во время записи, обратитесь к документации MLflow, раздел - Как регистрировать модели с сигнатурами.

Ведение журнала на основе кода с помощью LangChain

В следующих инструкциях и примере кода показано, как регистрировать агент с помощью LangChain.

Создайте записную книжку или файл Python с кодом. В этом примере записная книжка или файл именуется agent.py. Записная книжка или файл должны содержать агент LangChain, называемый здесь lc_agent.
Включите mlflow.models.set_model(lc_agent) в записную книжку или файл.
Создайте новую записную книжку, которую можно использовать в качестве управляющей записной книжки (называемой driver.py в этом примере).
В записной книжке драйвера используйте следующий код, чтобы запустить agent.py и записать результаты в модель MLflow:
```
mlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)
```
Параметр resources объявляет управляемые Databricks ресурсы, необходимые для обслуживания агента, например индекс векторного поиска или конечную точку, которая служит базовой моделью. Для получения дополнительной информации см. Указание ресурсов для автоматической передачи данных аутентификации.
Разверните модель. См. раздел "Развертывание агента для создания приложения ИИ".
При загрузке среды выполнения выполняется agent.py.
При поступлении запроса на обслуживание вызывается lc_agent.invoke(...).


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)

Ведение журнала на основе кода с использованием PyFunc ()

В следующих инструкциях и примере кода показано, как регистрировать агент с помощью PyFunc.

Создайте записную книжку или файл Python с кодом. В этом примере записная книжка или файл именуется agent.py. Записная книжка или файл должны содержать класс PyFunc с именем PyFuncClass.
Включите mlflow.models.set_model(PyFuncClass) в записную книжку или файл.
Создайте новую записную книжку, которую можно использовать в качестве управляющей записной книжки (называемой driver.py в этом примере).
В записной книжке драйвера используйте следующий код, чтобы запустить agent.py и записать результаты в модель MLflow:
```
mlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)
```
Параметр resources объявляет управляемые Databricks ресурсы, необходимые для обслуживания агента, например индекс векторного поиска или конечную точку, которая служит базовой моделью. Для получения дополнительной информации см. Указание ресурсов для автоматической передачи данных аутентификации.
Разверните модель. См. раздел "Развертывание агента для создания приложения ИИ".
При загрузке среды выполнения выполняется agent.py.
При поступлении запроса на обслуживание вызывается PyFuncClass.predict(...).

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)

Укажите ресурсы для автоматической сквозной аутентификации

Агенты ИИ часто должны проходить проверку подлинности в других ресурсах для выполнения задач. Например, агенту может потребоваться доступ к индексу векторного поиска для запроса неструктурированных данных.

Как описано в процессе аутентификации для зависимых ресурсов, Model Serving поддерживает аутентификацию как для ресурсов, управляемых Databricks, так и для внешних ресурсов при развертывании агента.

Для наиболее распространенных типов ресурсов Databricks система поддерживает и рекомендует заранее объявлять зависимости ресурсов для агента во время логирования. Это позволяет сквозную автоматическую проверку подлинности при развертывании агента; Databricks автоматически подготавливает, обновляет и управляет кратковременными учетными данными для безопасного доступа к этим зависимостям ресурсов из конечной точки агента.

Чтобы включить автоматический сквозной процесс проверки подлинности, укажите зависимые ресурсы с помощью параметра resources API log_model(), как это показано в следующем коде.

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 рекомендует вручную указать resources для всех вариантов агента.

Примечание.

Если вы не указываете ресурсы при логировании агентов LangChain с помощью mlflow.langchain.log_model(...), MLflow выполняет автоматическое определение ресурсов по мере возможности. Однако это может не записывать все зависимости, что приводит к ошибкам авторизации при обслуживании или запросе агента.

В следующей таблице перечислены ресурсы Databricks, которые поддерживают автоматическую передачу аутентификации, и минимальная версия mlflow, необходимая для доступа к ресурсу.

Тип ресурса	Минимальная версия `mlflow`, необходимая для регистрации ресурса
Индекс векторного поиска	Требуется `mlflow` 2.13.1 или более поздняя версия
Сервисная конечная точка модели	Требуется `mlflow` 2.13.1 или более поздняя версия
Хранилище SQL	Требуется `mlflow` 2.16.1 или более поздняя версия
Функция каталога Unity	Требуется `mlflow` 2.16.1 или более поздняя версия
Пространство Genie	Требуется `mlflow` 2.17.1 или более поздняя версия
Таблица каталога Unity	Требуется `mlflow` 2.18.0 или выше
Подключение каталога Unity	Требуется `mlflow` 2.17.1 или более поздняя версия

Автоматическая проверка подлинности для клиентов OpenAI

Если агент использует клиент OpenAI, используйте пакет SDK Databricks для автоматической проверки подлинности во время развертывания. Пакет SDK Databricks предоставляет оболочку для создания клиента OpenAI с автоматической настройкой авторизации. Выполните следующую команду в записной книжке:

% pip install databricks-sdk[openai]

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

Затем укажите конечную точку обслуживания модели в рамках resources для автоматической проверки подлинности во время развертывания.

регистрация агента в каталоге Unity

Перед развертыванием агента необходимо зарегистрировать агент в каталоге Unity. Регистрация агента оформляет его как модель в каталоге Unity. В результате можно использовать разрешения каталога Unity для авторизации ресурсов в агенте.

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)

Поделиться через