모델 로깅과 등록

import mlflow
from databricks.sdk import WorkspaceClient
from langchain_community.chat_models import ChatDatabricks
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.prompts import ChatPromptTemplate

# 에이전트 정의
class CustomerSupportAgent(mlflow.pyfunc.PythonModel):
    def __init__(self):
        self.llm = ChatDatabricks(
            endpoint="databricks-meta-llama-3-3-70b-instruct",
            temperature=0.1
        )

    def predict(self, context, model_input, params=None):
        messages = model_input.get("messages", [])
        response = self.llm.invoke(messages)
        return {"content": response.content}

# 진입점 지정 — 이 파일이 모델로 사용될 때 이 클래스가 로드됩니다
mlflow.models.set_model(CustomerSupportAgent())

import mlflow

# MLflow 실험 설정
mlflow.set_experiment("/Users/user@example.com/customer-support-agent")

# 에이전트 로깅
with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        artifact_path="agent",
        python_model="agent.py",  # 에이전트 코드 파일 경로
        registered_model_name="catalog.schema.customer_support_agent",
        pip_requirements=[
            "langchain>=0.3.0",
            "langchain-community>=0.3.0",
            "databricks-sdk>=0.30.0",
        ],
        input_example={
            "messages": [
                {"role": "user", "content": "주문 상태를 확인하고 싶습니다."}
            ]
        }
    )
    print(f"Model URI: {model_info.model_uri}")

# 리소스 의존성 명시 (배포 시 자동으로 권한 설정)
from mlflow.models.resources import (
    DatabricksServingEndpoint,
    DatabricksVectorSearchIndex
)

with mlflow.start_run():
    mlflow.pyfunc.log_model(
        artifact_path="agent",
        python_model="agent.py",
        registered_model_name="catalog.schema.my_agent",
        resources=[
            DatabricksServingEndpoint(endpoint_name="databricks-meta-llama-3-3-70b-instruct"),
            DatabricksVectorSearchIndex(index_name="catalog.schema.docs_index"),
        ]
    )

from mlflow import MlflowClient

client = MlflowClient()

# 특정 버전에 앨리어스 지정
client.set_registered_model_alias(
    name="catalog.schema.customer_support_agent",
    alias="champion",
    version=3
)

# 앨리어스로 모델 조회
champion_version = client.get_model_version_by_alias(
    name="catalog.schema.customer_support_agent",
    alias="champion"
)
print(f"Champion version: {champion_version.version}")