MLOps Agent

def check_data_drift(table_name: str = "lgit_pm_inference_results") -> dict:
    """PSI 기반 데이터 드리프트 확인"""
    import numpy as np

    feature_columns = ["air_temperature_k", "process_temperature_k",
                       "rotational_speed_rpm", "torque_nm", "tool_wear_min"]

    train_pdf = spark.table("lgit_pm_training").filter("split='train'").select(*feature_columns).toPandas()

    try:
        infer_pdf = spark.table(table_name).select(*feature_columns).toPandas()
    except Exception:
        return {"drift_detected": False, "message": "추론 테이블이 없습니다.", "psi_values": {}}

    psi_values = {}
    for col in feature_columns:
        breakpoints = np.linspace(
            min(train_pdf[col].min(), infer_pdf[col].min()),
            max(train_pdf[col].max(), infer_pdf[col].max()), 11)
        expected = np.maximum(np.histogram(train_pdf[col], bins=breakpoints)[0] / len(train_pdf), 0.001)
        actual = np.maximum(np.histogram(infer_pdf[col], bins=breakpoints)[0] / len(infer_pdf), 0.001)
        psi_values[col] = float(np.sum((actual - expected) * np.log(actual / expected)))
    drift_detected = any(v > 0.2 for v in psi_values.values())
    return {"drift_detected": drift_detected, "psi_values": psi_values,
            "timestamp": datetime.now().isoformat()}

def trigger_retraining(reason: str = "scheduled") -> dict:
    """모델 재학습 트리거 — 실제 환경에서는 w.jobs.run_now() 호출"""
    print(f"재학습 트리거 — 사유: {reason}")
    return {"status": "triggered", "reason": reason,
            "timestamp": datetime.now().isoformat(),
            "message": f"재학습이 트리거되었습니다. 사유: {reason}"}

from databricks.sdk import WorkspaceClient
w = WorkspaceClient()
run = w.jobs.run_now(
    job_id=178771619413732,  # LGIT_MLOps_PoC_Structured_Pipeline
    notebook_params={"retrain_reason": reason}
)

def run_batch_prediction() -> dict:
    """Champion 모델로 배치 예측 실행"""
    champion_udf = mlflow.pyfunc.spark_udf(spark, model_uri=f"models:/{model_name}@Champion")
    inference_df = spark.table("lgit_pm_training").select("udi", "type", *feature_columns)
    preds = inference_df.withColumn("failure_probability", champion_udf(*feature_columns))
    count = preds.count()
    return {"status": "success", "predictions_count": count,
            "timestamp": datetime.now().isoformat()}

def get_model_status() -> dict:
    """현재 Champion 모델의 상태를 확인합니다."""
    from mlflow import MlflowClient
    client = MlflowClient()
    champion = client.get_model_version_by_alias(model_name, "Champion")
    run = mlflow.get_run(champion.run_id)
    return {
        "model_name": model_name,
        "champion_version": champion.version,
        "val_f1_score": run.data.metrics.get("val_f1_score", "N/A"),
        "val_auc": run.data.metrics.get("val_auc", "N/A"),
        "status": "active"
    }

SYSTEM_PROMPT = """당신은 제조 현장의 MLOps Agent입니다.
다음 도구를 사용하여 예지보전 모델의 운영을 자동화합니다:

1. check_data_drift: 데이터 드리프트 확인
2. trigger_retraining: 모델 재학습 트리거
3. run_batch_prediction: 배치 예측 실행
4. get_model_status: 현재 모델 상태 확인

주요 자동화 규칙:
- 데이터 드리프트가 감지되면 자동으로 재학습을 트리거합니다.
- 스케줄에 따라 배치 예측을 실행합니다.
- 모델 성능이 임계값 이하이면 알림을 생성합니다."""

TOOLS = {
    "check_data_drift": check_data_drift,
    "trigger_retraining": trigger_retraining,
    "run_batch_prediction": run_batch_prediction,
    "get_model_status": get_model_status,
}

# 시나리오 1: 정기 상태 점검
status = get_model_status()
print(f"모델 상태: {json.dumps(status, indent=2, ensure_ascii=False)}")

# 시나리오 2: 드리프트 기반 자동 재학습
drift_result = check_data_drift()
if drift_result["drift_detected"]:
    retrain_result = trigger_retraining(reason="data_drift_detected")

# 시나리오 3: 스케줄 기반 배치 예측
pred_result = run_batch_prediction()

def mlops_agent_workflow(trigger_type: str = "full_cycle"):
    """
    MLOps Agent의 전체 워크플로우를 실행합니다.

    Args:
        trigger_type: 트리거 유형
            - "scheduled": 정기 스케줄 (배치 예측 + 상태 점검)
            - "drift_check": 드리프트 확인 및 필요 시 재학습
            - "full_cycle": 전체 주기 (드리프트 → 재학습 → 예측)
    """
    results = {"trigger_type": trigger_type, "timestamp": datetime.now().isoformat()}

    if trigger_type in ["scheduled", "full_cycle"]:
        results["batch_prediction"] = run_batch_prediction()

    if trigger_type in ["drift_check", "full_cycle"]:
        drift = check_data_drift()
        results["drift_check"] = drift
        if drift["drift_detected"]:
            results["retraining"] = trigger_retraining(reason="auto_drift_detection")

    results["model_status"] = get_model_status()
    return results

# 1. Foundation Model API — 별도 API 키 없이 Databricks 내장 LLM 사용
from langchain_databricks import ChatDatabricks
llm = ChatDatabricks(endpoint="databricks-meta-llama-3-1-70b-instruct")

# 2. UC Function as Tool — SQL 함수를 Agent Tool로 직접 연결
from databricks_langchain import UCFunctionToolkit
tools = UCFunctionToolkit(function_names=["mlops.tools.check_drift", "mlops.tools.trigger_retrain"])

# 3. MLflow Tracing — Agent의 모든 판단을 자동 기록
import mlflow
mlflow.langchain.autolog()  # 이 한 줄로 LangChain/LangGraph 전체 추적!

from langgraph.graph import StateGraph, END
from langchain_databricks import ChatDatabricks
from typing import TypedDict

# 1. Agent 상태 정의
class MLOpsState(TypedDict):
    drift_detected: bool
    max_psi: float
    retrain_needed: bool
    retrain_result: dict
    report: str

# 2. 각 노드(단계) 정의
def observe(state):
    result = check_data_drift()
    return {"drift_detected": result["drift_detected"], "max_psi": result["max_psi"]}

def decide(state):
    return {"retrain_needed": state["drift_detected"] and state["max_psi"] > 0.2}

def act(state):
    if state["retrain_needed"]:
        result = trigger_retraining(reason=f"drift_psi_{state['max_psi']:.2f}")
        return {"retrain_result": result}
    return {"retrain_result": {"status": "skipped"}}

def report(state):
    llm = ChatDatabricks(endpoint="databricks-meta-llama-3-1-70b-instruct")
    summary = llm.invoke(f"MLOps 상태 보고: 드리프트={state['drift_detected']}, "
                         f"PSI={state['max_psi']}, 재학습={state['retrain_result']}")
    return {"report": summary.content}

# 3. 워크플로우 그래프 구성
graph = StateGraph(MLOpsState)
graph.add_node("observe", observe)
graph.add_node("decide", decide)
graph.add_node("act", act)
graph.add_node("report", report)

graph.set_entry_point("observe")
graph.add_edge("observe", "decide")
graph.add_conditional_edges("decide",
    lambda s: "act" if s["retrain_needed"] else "report",
    {"act": "act", "report": "report"})
graph.add_edge("act", "report")
graph.add_edge("report", END)

# 4. Agent 실행
agent = graph.compile()
result = agent.invoke({"drift_detected": False, "max_psi": 0.0,
                       "retrain_needed": False, "retrain_result": {}, "report": ""})