Databricks Advanced RAG 구현 로드맵

문서에 전문 용어가 많은가?
  ├─ Yes → Hybrid Retriever (BM25 + Dense)
  │         └─ 한국어인가? → Yes → Kiwi 토크나이저 적용
  └─ No → Dense Retriever
              └─ 정밀도가 중요한가?
                   ├─ Yes → + Reranking 추가
                   └─ No → Dense만으로 충분

# Phase 1 구현 예시
from langchain_databricks import DatabricksVectorSearch, ChatDatabricks
from langchain.chains import RetrievalQA

# 기본 Vector Search Retriever
retriever = DatabricksVectorSearch(
    endpoint="vs-endpoint",
    index_name="catalog.schema.docs_index",
    columns=["content", "source"],
).as_retriever(search_kwargs={"k": 5, "query_type": "hybrid"})

# 기본 RAG 체인
llm = ChatDatabricks(endpoint="databricks-claude-sonnet-4")
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=retriever,
    return_source_documents=True,
)

result = qa_chain.invoke({"query": "Unity Catalog 권한 설정 방법"})

# Phase 2 구현 예시
from kiwipiepy import Kiwi
from langchain.retrievers import EnsembleRetriever, ContextualCompressionRetriever
from langchain_community.retrievers import BM25Retriever
from langchain_community.cross_encoders import HuggingFaceCrossEncoder
from langchain.retrievers.document_compressors import CrossEncoderReranker

kiwi = Kiwi()

# 1. Kiwi 기반 BM25 Retriever
bm25 = BM25Retriever.from_documents(
    documents,
    preprocess_func=lambda t: [tok.form for tok in kiwi.tokenize(t)
                                if tok.tag in {"NNG","NNP","VV","VA","SL"}],
    k=20
)

# 2. Vector Search Retriever
vs = DatabricksVectorSearch(
    endpoint="vs-endpoint",
    index_name="catalog.schema.docs_index",
    columns=["content", "source"],
).as_retriever(search_kwargs={"k": 20})

# 3. 앙상블 (RRF - Reciprocal Rank Fusion)
# RRF는 여러 검색기의 결과를 순위 기반으로 합치는 알고리즘입니다.
# 각 문서의 순위(rank)를 1/(k+rank) 형태로 점수화한 뒤 합산합니다.
# k는 보통 60으로 설정하며, 상위 순위에 높은 가중치를 부여합니다.
# weights=[0.4, 0.6]은 BM25에 40%, Dense에 60% 가중치를 적용합니다.
ensemble = EnsembleRetriever(retrievers=[bm25, vs], weights=[0.4, 0.6])

# 4. Cross-Encoder Reranking
cross_encoder = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-v2-m3")
reranker = CrossEncoderReranker(model=cross_encoder, top_n=5)

# 5. 최종 Retriever
final_retriever = ContextualCompressionRetriever(
    base_compressor=reranker,
    base_retriever=ensemble,
)

# Phase 3: Agentic RAG with LangGraph
from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated

class RAGState(TypedDict):
    question: str
    rewritten_query: str
    route: str
    documents: list
    grade: str
    answer: str
    iteration: int

def route_question(state: RAGState) -> RAGState:
    """질문 의도 분석 → 라우팅"""
    route = route_query(state["question"])
    return {** state, "route": route["source"]}

def retrieve(state: RAGState) -> RAGState:
    """라우팅된 소스에서 검색"""
    if state["route"] == "VECTOR_DB":
        docs = final_retriever.invoke(state["rewritten_query"])
    elif state["route"] == "SQL_DB":
        docs = sql_retriever.invoke(state["rewritten_query"])
    else:
        docs = web_retriever.invoke(state["rewritten_query"])
    return {** state, "documents": docs}

def grade_documents(state: RAGState) -> RAGState:
    """검색 결과 품질 평가 (CRAG)"""
    result = corrective_rag(state["question"], state["documents"], llm)
    return {** state, "grade": result["strategy"], "documents": result["docs"]}

def generate(state: RAGState) -> RAGState:
    """답변 생성 + 충실도 검증 (Self-RAG)"""
    answer = self_rag(state["question"], final_retriever, llm)
    return {** state, "answer": answer}

# LangGraph 워크플로 구성
workflow = StateGraph(RAGState)
workflow.add_node("rewrite", lambda s: {** s, "rewritten_query": rewrite_query(s["question"])})
workflow.add_node("route", route_question)
workflow.add_node("retrieve", retrieve)
workflow.add_node("grade", grade_documents)
workflow.add_node("generate", generate)

workflow.set_entry_point("rewrite")
workflow.add_edge("rewrite", "route")
workflow.add_edge("route", "retrieve")
workflow.add_edge("retrieve", "grade")
workflow.add_edge("grade", "generate")
workflow.add_edge("generate", END)

app = workflow.compile()