기본 개념과 자동 트레이싱

import mlflow
import openai

# 한 줄로 자동 트레이싱 활성화
mlflow.openai.autolog()

# 이후 모든 OpenAI 호출이 자동으로 추적됩니다
client = openai.OpenAI()
response = client.chat.completions.create(
    model="gpt-4",
    messages=[{"role": "user", "content": "데이터 레이크하우스란?"}]
)
# → MLflow UI에서 입력, 출력, 토큰 사용량, 지연시간을 확인할 수 있습니다

# LangChain 자동 트레이싱
mlflow.langchain.autolog()

from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate

chain = ChatPromptTemplate.from_template("Explain {topic} in Korean") | ChatOpenAI()
result = chain.invoke({"topic": "data lakehouse"})
# → Chain의 각 단계(프롬프트 → LLM 호출)가 자동으로 추적됩니다

import mlflow

@mlflow.trace
def rag_pipeline(question: str) -> str:
    """전체 RAG 파이프라인"""
    docs = retrieve_documents(question)
    answer = generate_answer(question, docs)
    return answer

@mlflow.trace(span_type="RETRIEVER", name="document_search")
def retrieve_documents(query: str) -> list:
    """Vector Search에서 관련 문서 검색"""
    results = vector_index.similarity_search(
        query_text=query,
        num_results=5
    )
    return results

@mlflow.trace(span_type="LLM", name="llm_call")
def generate_answer(question: str, context: list) -> str:
    """LLM에 질문과 컨텍스트를 전달하여 답변 생성"""
    prompt = f"다음 문서를 참고하여 질문에 답해 주세요.\n\n문서:\n{context}\n\n질문: {question}"
    response = client.chat.completions.create(
        model="databricks-meta-llama-3-3-70b-instruct",
        messages=[{"role": "user", "content": prompt}]
    )
    return response.choices[0].message.content

# 호출하면 전체 Trace가 자동 생성됩니다
answer = rag_pipeline("반품 정책이 어떻게 되나요?")

import mlflow

with mlflow.start_span(name="custom_pipeline") as span:
    span.set_inputs({"question": question})

    with mlflow.start_span(name="preprocessing", span_type="CHAIN") as prep_span:
        prep_span.set_inputs({"raw_question": question})
        cleaned = preprocess(question)
        prep_span.set_outputs({"cleaned_question": cleaned})

    with mlflow.start_span(name="search", span_type="RETRIEVER") as search_span:
        search_span.set_inputs({"query": cleaned})
        docs = search(cleaned)
        search_span.set_outputs({"num_results": len(docs)})
        search_span.set_attributes({"search_score": docs[0].score})

    with mlflow.start_span(name="generation", span_type="LLM") as llm_span:
        llm_span.set_inputs({"prompt": prompt, "context_length": len(docs)})
        answer = generate(prompt, docs)
        llm_span.set_outputs({"answer": answer})
        llm_span.set_attributes({
            "token_usage.input": 850,
            "token_usage.output": 120,
            "model": "llama-3.3-70b"
        })

    span.set_outputs({"answer": answer})