RAG 체인 구축

# Agent Framework 통합의 핵심: ChatAgent 인터페이스
# 이 인터페이스를 구현하면 Databricks 생태계 전체와 호환됩니다
from mlflow.pyfunc import ChatAgent

class MyRAGAgent(ChatAgent):
    def predict(self, messages, context=None, custom_inputs=None):
        # 여기에 RAG 로직 구현
        pass

%pip install langchain langchain-community databricks-vectorsearch mlflow
dbutils.library.restartPython()

from langchain_community.chat_models import ChatDatabricks

llm = ChatDatabricks(
    endpoint="databricks-meta-llama-3-3-70b-instruct",
    temperature=0.1, max_tokens=1024
)
response = llm.invoke("Databricks란 무엇인가요?")

from databricks.vector_search.client import VectorSearchClient
from langchain_community.vectorstores import DatabricksVectorSearch

client = VectorSearchClient()
index = client.get_index(endpoint_name="rag-vs-endpoint",
                          index_name="catalog.schema.document_chunks_index")

retriever = DatabricksVectorSearch(
    index=index, text_column="content", columns=["chunk_id", "content", "source"]
).as_retriever(search_kwargs={"k": 5})

docs = retriever.invoke("Vector Search 설정 방법")

# 나쁜 예: 지시가 모호하고 환각 방지 장치가 없음
bad_prompt = "다음 컨텍스트를 참고해서 답변해줘: {context}\n질문: {question}"

# 좋은 예: 명확한 역할, 제약 조건, 출력 형식 지정
good_prompt = """당신은 Databricks 기술 문서를 기반으로 답변하는 도우미입니다.

규칙:
- 제공된 컨텍스트 정보만을 사용하여 답변하세요.
- 컨텍스트에 없는 내용은 "해당 정보를 찾을 수 없습니다"라고 답변하세요.
- 답변에 출처 문서를 명시하세요.
- 한국어로 답변하세요.

컨텍스트:
{context}"""

from langchain_core.prompts import ChatPromptTemplate

prompt_template = ChatPromptTemplate.from_messages([
    ("system", """당신은 Databricks 기술 문서를 기반으로 답변하는 도우미입니다.

규칙:
- 제공된 컨텍스트 정보만을 사용하여 답변하세요.
- 컨텍스트에 없는 내용은 "해당 정보를 찾을 수 없습니다"라고 답변하세요.
- 답변에 출처 문서를 명시하세요.
- 한국어로 답변하세요.

컨텍스트:
{context}"""),
    ("human", "{question}")
])

from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

def format_docs(docs):
    """검색된 문서를 프롬프트에 주입할 형태로 포매팅합니다.

    각 문서 사이에 '---' 구분자를 넣는 이유: LLM이 서로 다른 출처의 정보를
    명확히 구분할 수 있도록 돕습니다. 구분자가 없으면 LLM이 여러 문서의
    내용을 혼합하여 잘못된 추론을 할 수 있습니다.
    출처(source)를 함께 포함하면 LLM이 답변에 출처를 인용할 수 있습니다.
    """
    return "\n\n---\n\n".join([
        f"[출처: {doc.metadata.get('source', 'N/A')}]\n{doc.page_content}"
        for doc in docs
    ])

rag_chain = (
    {
        "context": retriever | format_docs,
        "question": RunnablePassthrough()
    }
    | prompt_template
    | llm
    | StrOutputParser()
)

# 체인 실행
answer = rag_chain.invoke("Databricks에서 Vector Search Index를 생성하는 방법은?")
print(answer)

from mlflow.pyfunc import ChatAgent
from mlflow.types.agent import ChatAgentMessage, ChatAgentResponse

class DatabricksRAGAgent(ChatAgent):
    def __init__(self):
        self.llm = ChatDatabricks(endpoint="databricks-meta-llama-3-3-70b-instruct", temperature=0.1)
        client = VectorSearchClient()
        index = client.get_index(endpoint_name="rag-vs-endpoint",
                                  index_name="catalog.schema.document_chunks_index")
        self.retriever = DatabricksVectorSearch(
            index=index, text_column="content", columns=["chunk_id", "content", "source"]
        ).as_retriever(search_kwargs={"k": 5})

    def predict(self, messages, context=None, custom_inputs=None):
        question = messages[-1]["content"]
        docs = self.retriever.invoke(question)
        ctx = "\n\n".join([d.page_content for d in docs])
        prompt = prompt_template.invoke({"context": ctx, "question": question})
        response = self.llm.invoke(prompt)
        return ChatAgentResponse(
            messages=[ChatAgentMessage(role="assistant", content=response.content)]
        )

import mlflow
mlflow.set_registry_uri("databricks-uc")

with mlflow.start_run(run_name="rag-chain-v1"):
    model_info = mlflow.pyfunc.log_model(
        artifact_path="rag_chain",
        python_model=DatabricksRAGAgent(),
        input_example={"messages": [{"role": "user", "content": "Vector Search란?"}]},
        registered_model_name="catalog.schema.rag_agent"
    )
print(f"모델 URI: {model_info.model_uri}")

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

# 체인 실행 시 각 스팬이 자동 기록됨
# - retriever: 검색된 문서 목록, 소요 시간
# - prompt: 최종 프롬프트 내용
# - llm: LLM 입력/출력, 토큰 수, 소요 시간
answer = rag_chain.invoke("Vector Search 인덱스 유형을 설명해주세요")