AI 에이전트 기록 및 등록

Mosaic AI 에이전트 프레임워크를 사용하여 AI 에이전트를 기록합니다. 에이전트 로깅은 개발 프로세스의 기초입니다. 로깅은 에이전트 코드 및 구성의 "지정 시간"을 캡처하므로 구성의 품질을 평가할 수 있습니다.

요구 사항

로깅하기 전에 AI 에이전트를 만듭니다.

코드 기반 로깅

Databricks는 에이전트를 로깅할 때 MLflow의 모델의 코드 기능을 사용하는 것이 좋습니다.

이 방법에서 에이전트의 코드는 Python 파일로 캡처되고 Python 환경은 패키지 목록으로 캡처됩니다. 에이전트가 배포되면 Python 환경이 복원되고 에이전트의 코드가 실행되어 엔드포인트가 호출될 때 호출될 수 있도록 에이전트를 메모리에 로드합니다.

이 접근 방식을 mlflow.models.predict() 같은 배포 전 유효성 검사 API를 사용하여 서비스를 위해 배포할 때 에이전트가 안정적으로 실행되도록 할 수 있습니다.

에이전트 또는 에이전트를 기록하는 코드는 에이전트 코드와 별도의 Notebook에 있어야 합니다. 이 Notebook을 드라이버 Notebook이라고 합니다. Notebook 예제는 Notebook 예제를 참조하세요.

로깅하는 동안 모델 서명 유추

로깅하는 동안 에이전트의 입력 및 출력 스키마를 지정하는 하나의 MLflow 모델 서명을 정의해야 합니다. 서명은 입력 및 출력의 유효성을 검사하여 에이전트가 AI Playground 및 검토 앱과 같은 다운스트림 도구와 올바르게 상호 작용하도록 합니다. 또한 에이전트를 효과적으로 사용하는 방법에 대한 다른 애플리케이션을 안내합니다.

Databricks는 사용자가 제공한 입력 예제에 따라 에이전트의 서명을 자동으로 생성하기 위해 MLflow의 모델 서명 추론 기능을 사용하는 것이 좋습니다. 이 방법은 서명을 수동으로 정의하는 것보다 더 편리합니다.

아래 LangChain 및 PyFunc 예제에서는 모델 서명 추론을 사용합니다.

로깅 시 모델 서명을 직접 명시적으로 정의하려는 경우 MLflow 문서 - 서명사용하여 모델을 기록하는 방법을 참조하세요.

LangChain을 사용하는 코드 기반 로깅

다음 지침 및 코드 샘플에서는 LangChain을 사용하여 에이전트를 기록하는 방법을 보여 줍니다.

1. 코드를 사용하여 Notebook 또는 Python 파일을 만듭니다. 이 예제에서 Notebook 또는 파일의 이름은 agent.py. 노트북이나 파일에는 여기서 lc_agent이라고 불리는 LangChain 에이전트가 포함되어야 합니다.

2. notebook 또는 파일에 mlflow.models.set_model(lc_agent)을 포함하세요.

3. 드라이버 Notebook으로 사용할 새 Notebook을 만듭니다(이 예제에서는 driver.py로 부름).

4. 드라이버 Notebook에서 다음 코드를 사용하여 agent.py 실행하고 결과를 MLflow 모델에 기록합니다.

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

   resources 매개 변수는 벡터 검색 인덱스 또는 기본 모델을 제공하는 서비스 엔드포인트와 같이 에이전트를 제공하는 데 필요한 Databricks 관리형 리소스를 선언합니다. 자세한 내용은 자동 인증 통과대한 리소스 지정을 참조하세요.

5. 모델을 배포합니다. 생성형 AI 응용 프로그램에 대한 에이전트 배포를 참조하세요.

6. 제공하는 환경이 로드되면 agent.py가 실행됩니다.

7. 제공하는 요청이 들어오면 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를 사용하여 에이전트를 기록하는 방법을 보여 줍니다.

1. 코드를 사용하여 Notebook 또는 Python 파일을 만듭니다. 이 예제에서 Notebook 또는 파일의 이름은 agent.py. Notebook 또는 파일에는 PyFuncClass이라는 PyFunc 클래스가 포함되어야 합니다.

2. Notebook 또는 파일에 mlflow.models.set_model(PyFuncClass)을 포함합니다.

3. 드라이버 Notebook으로 사용할 새 Notebook을 만듭니다(이 예제에서는 driver.py로 부름).

4. 드라이버 Notebook에서 다음 코드를 사용하여 agent.py 실행하고 결과를 MLflow 모델에 기록합니다.

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

   resources 매개 변수는 벡터 검색 인덱스 또는 기본 모델을 제공하는 서비스 엔드포인트와 같이 에이전트를 제공하는 데 필요한 Databricks 관리형 리소스를 선언합니다. 자세한 내용은 자동 인증 통과대한 리소스 지정을 참조하세요.

5. 모델을 배포합니다. 생성형 AI 응용 프로그램에 대한 에이전트 배포를 참조하세요.

6. 제공하는 환경이 로드되면 agent.py가 실행됩니다.

7. 제공하는 요청이 들어오면 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)

자동 인증 통과를 위한 리소스 지정

AI 에이전트는 작업을 완료하기 위해 다른 리소스에 인증해야 하는 경우가 많습니다. 예를 들어 에이전트는 구조화되지 않은 데이터를 쿼리하기 위해 Vector Search 인덱스에 액세스해야 할 수 있습니다.

종속 리소스에 대한 인증에 설명된 대로 Model Serving는 에이전트를 배포할 때 Databricks 관리 및 외부 리소스에 대한 인증을 지원합니다.

가장 일반적인 Databricks 리소스 유형의 경우 Databricks는 로깅 중에 에이전트에 대한 리소스 종속성을 미리 선언하는 것을 지원하고 권장합니다. 이렇게 하면 에이전트를 배포할 때 자동 인증 통과 사용할 수 있습니다. Databricks는 에이전트 엔드포인트 내에서 이러한 리소스 종속성에 안전하게 액세스하기 위해 수명이 짧은 자격 증명을 자동으로 프로비전, 회전 및 관리합니다.

자동 인증 통과를 사용하도록 설정하려면 다음 코드와 같이 log_model() API의 resources 매개 변수를 사용하여 종속 리소스를 지정합니다.

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

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

Databricks는 모든 에이전트 유형에 대해 resources을 수동으로 지정할 것을 권장합니다.

다음 표에서는 자동 인증 통과를 지원하는 Databricks 리소스와 리소스를 기록하는 데 필요한 최소 mlflow 버전을 나열합니다.

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)

다음 단계

• AI 에이전트에 추적을 추가합니다.
• AI 에이전트를 배포합니다.