Guardrails와 트러블슈팅

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import (
    AiGatewayConfig,
    AiGatewayGuardrails,
    AiGatewayGuardrailParameters,
    AiGatewayRateLimit,
    AiGatewayUsageTrackingConfig
)

w = WorkspaceClient()

# Guardrails가 포함된 엔드포인트 설정
w.serving_endpoints.put_ai_gateway(
    name="customer-support-agent",
    guardrails=AiGatewayGuardrails(
        input=AiGatewayGuardrailParameters(
            safety=True,       # 유해 콘텐츠 필터링
            pii={"behavior": "BLOCK"},  # PII 정보 차단
            valid_topics=["customer support", "order inquiry", "product info"],
            invalid_topics=["politics", "religion", "competitor products"]
        ),
        output=AiGatewayGuardrailParameters(
            safety=True,
            pii={"behavior": "BLOCK"}
        )
    ),
    rate_limits=[
        AiGatewayRateLimit(
            calls=100,
            renewal_period="minute",
            key="user"  # 사용자별 분당 100회 제한
        )
    ],
    usage_tracking_config=AiGatewayUsageTrackingConfig(enabled=True)
)

from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

# 엔드포인트 상태 조회
endpoint = w.serving_endpoints.get("customer-support-agent")
print(f"State: {endpoint.state.ready}")
print(f"Config update: {endpoint.state.config_update}")

# 엔드포인트 로그 조회
logs = w.serving_endpoints.logs(
    name="customer-support-agent",
    served_model_name="customer_support_agent-3"
)
print(logs.logs)

import mlflow

# 로깅된 모델 로드 및 테스트
model = mlflow.pyfunc.load_model(model_info.model_uri)

# 테스트 입력
test_input = {
    "messages": [
        {"role": "user", "content": "주문번호 12345의 배송 상태를 알려주세요."}
    ]
}

# 예측 실행
response = model.predict(test_input)
print(response)

# ============================================
# 1. 에이전트 코드 작성 (agent.py)
# ============================================
# (위의 agent.py 참고)

# ============================================
# 2. MLflow로 모델 로깅 + UC 등록
# ============================================
import mlflow
from mlflow.models.resources import (
    DatabricksServingEndpoint,
    DatabricksVectorSearchIndex
)

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", "databricks-sdk>=0.30.0"],
        resources=[
            DatabricksServingEndpoint(endpoint_name="databricks-meta-llama-3-3-70b-instruct"),
            DatabricksVectorSearchIndex(index_name="catalog.schema.docs_index"),
        ],
        input_example={"messages": [{"role": "user", "content": "안녕하세요"}]}
    )

# ============================================
# 3. 로컬 테스트
# ============================================
model = mlflow.pyfunc.load_model(model_info.model_uri)
result = model.predict({"messages": [{"role": "user", "content": "테스트"}]})
print(f"로컬 테스트 결과: {result}")

# ============================================
# 4. 배포 (엔드포인트 + Review App)
# ============================================
from databricks import agents

deployment = agents.deploy(
    model_name="catalog.schema.customer_support_agent",
    model_version=model_info.registered_model_version
)
print(f"배포 완료!")
print(f"   Endpoint: {deployment.endpoint_name}")
print(f"   Review App: {deployment.review_app_url}")

# ============================================
# 5. Review App 설정
# ============================================
agents.set_permissions(
    model_name="catalog.schema.customer_support_agent",
    users=["tester@company.com"],
    permission="CAN_QUERY"
)

# ============================================
# 6. 프로덕션 엔드포인트 호출 테스트
# ============================================
from databricks.sdk import WorkspaceClient

w = WorkspaceClient()
response = w.serving_endpoints.query(
    name=deployment.endpoint_name,
    input={
        "messages": [
            {"role": "user", "content": "최근 주문 상태를 확인하고 싶습니다."}
        ]
    }
)
print(f"에이전트 응답: {response.choices[0].message.content}")