import mlflow
from langchain.chains import RetrievalQA
from langchain_community.vectorstores import Chroma

# 한 줄로 전체 애플리케이션 계측
mlflow.langchain.autolog()

@mlflow.trace(name="customer_support")
def answer_question(question, customer_tier="standard"):
    vectorstore = Chroma.from_documents(documents, embeddings)
    qa_chain = RetrievalQA.from_chain_type(
        llm=ChatOpenAI(temperature=0),
        retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
        return_source_documents=True
    )
    # Tracing이 전체 실행 트리를 자동으로 캡처
    result = qa_chain({"query": question})
    # 트레이스에 비즈니스 컨텍스트 추가
    mlflow.update_current_trace(
        tags={
            "customer_tier": customer_tier,
            "question_category": classify_question(question)
        }
    )
    return result["result"]

import mlflow.pyfunc

class CustomerServiceBot(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        # GenAI 애플리케이션의 모든 컴포넌트 로드
        self.llm = load_model_from_artifacts(context.artifacts["llm_config"])
        self.vector_store = initialize_vector_store(context.artifacts["knowledge_base"])
        self.prompt_template = load_prompt_template(context.artifacts["prompt_template"])

    def predict(self, context, model_input):
        # 애플리케이션 로직
        query = model_input["query"][0]
        relevant_docs = self.vector_store.similarity_search(query, k=3)
        formatted_prompt = self.prompt_template.format(
            query=query,
            context="\n".join([doc.page_content for doc in relevant_docs])
        )
        response = self.llm.predict(formatted_prompt)
        return {
            "response": response,
            "sources": [doc.metadata for doc in relevant_docs]
        }

# 완전한 애플리케이션을 패키징하고 버전 관리
with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        artifact_path="customer_service_bot",
        python_model=CustomerServiceBot(),
        artifacts={
            "llm_config": "configs/llm_config.yaml",
            "knowledge_base": "data/knowledge_embeddings.pkl",
            "prompt_template": "prompts/customer_service_v2.txt"
        },
        pip_requirements=["openai", "langchain", "chromadb"],
        signature=mlflow.models.infer_signature(example_input, example_output)
    )
    # 배포를 위해 Model Registry에 등록
    registered_model = mlflow.register_model(
        model_uri=model_info.model_uri,
        name="customer_service_bot",
        tags={"version": "v2.1", "eval_score": "0.87"}
    )

pip install -U mlflow

import mlflow

# 첫 번째 GenAI 실험 생성
mlflow.set_experiment("my_genai_prototype")

# OpenAI에 대한 자동 Tracing 활성화
# (지원되는 20개 이상의 Tracing 통합 중 하나를 선택하여 자동 Tracing 활성화)
mlflow.openai.autolog()

# 기존 GenAI 코드가 이제 자동으로 트레이스를 생성합니다.
# 지원되는 라이브러리에는 추가 계측이 필요 없습니다.