RAG 파이프라인 — 검색 증강 생성

-- ai_parse_document으로 PDF 텍스트 추출
CREATE OR REFRESH MATERIALIZED VIEW parsed_documents AS
SELECT
    file_name,
    ai_parse_document(
        CONCAT('/Volumes/catalog/schema/raw_docs/', file_name),
        'markdown'  -- 구조(제목, 표 등)를 유지하는 markdown 모드
    ) AS content,
    current_timestamp() AS parsed_at
FROM (SELECT file_name FROM LIST('/Volumes/catalog/schema/raw_docs/') WHERE file_name LIKE '%.pdf');

# Python으로 PDF 파싱 (더 세밀한 제어)
import fitz  # PyMuPDF

def parse_pdf(path: str) -> str:
    doc = fitz.open(path)
    text = ""
    for page in doc:
        text += page.get_text("text") + "\n\n"
    return text

from langchain.text_splitter import RecursiveCharacterTextSplitter

# 재귀적 분할 (가장 일반적)
splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,         # 최대 문자 수
    chunk_overlap=200,       # 청크 간 겹침 (맥락 연속성 유지)
    separators=["\n\n", "\n", ". ", " ", ""],  # 분할 우선순위
    length_function=len
)

chunks = splitter.split_text(document_text)

# 결과를 Delta 테이블로 저장
from pyspark.sql import Row

chunk_rows = [
    Row(
        doc_id=f"{doc_id}_chunk_{i}",
        parent_doc_id=doc_id,
        chunk_index=i,
        content=chunk,
        title=doc_title,
        source=doc_source
    )
    for i, chunk in enumerate(chunks)
]

chunk_df = spark.createDataFrame(chunk_rows)
chunk_df.write.format("delta").mode("append").saveAsTable("catalog.schema.document_chunks")

from databricks.vector_search.client import VectorSearchClient

vsc = VectorSearchClient()

# Managed Embeddings 인덱스 생성 (권장)
vsc.create_delta_sync_index(
    endpoint_name="vs-endpoint",
    index_name="catalog.schema.docs_index",
    source_table_name="catalog.schema.document_chunks",
    primary_key="doc_id",
    embedding_source_column="content",
    embedding_model_endpoint_name="databricks-gte-large-en",
    pipeline_type="TRIGGERED",
    columns_to_sync=["doc_id", "parent_doc_id", "content", "title", "source", "chunk_index"]
)

# 인덱스 동기화 실행
vsc.get_index("vs-endpoint", "catalog.schema.docs_index").sync()

import mlflow

@mlflow.trace
def rag_answer(question: str) -> dict:
    """RAG 파이프라인: 질문 → 검색 → 답변"""

    # 1. 관련 문서 검색
    docs = search_documents(question)

    # 2. 프롬프트 구성
    context = format_context(docs)
    prompt = build_prompt(question, context)

    # 3. LLM 호출
    answer = call_llm(prompt)

    # 4. 출처 정보 포함
    sources = [{"title": d["title"], "source": d["source"]} for d in docs]

    return {"answer": answer, "sources": sources}


@mlflow.trace(span_type="RETRIEVER")
def search_documents(query: str, num_results: int = 5) -> list:
    """Vector Search + Reranker로 관련 문서 검색"""
    results = index.similarity_search(
        query_text=query,
        columns=["doc_id", "content", "title", "source"],
        num_results=20,  # 초기 검색은 넓게
        query_options={
            "reranker": {
                "model_name": "databricks-reranker",
                "columns": ["content"],
                "top_k": num_results  # Reranker가 최종 선별
            }
        }
    )
    return [
        {"doc_id": r[0], "content": r[1], "title": r[2], "source": r[3]}
        for r in results["result"]["data_array"]
    ]


def format_context(docs: list) -> str:
    """검색된 문서를 LLM에 전달할 컨텍스트로 포맷팅"""
    return "\n\n---\n\n".join([
        f"[출처: {d['title']}]\n{d['content']}"
        for d in docs
    ])


def build_prompt(question: str, context: str) -> str:
    """시스템 프롬프트 + 컨텍스트 + 질문 조합"""
    return f"""당신은 정확하고 도움이 되는 AI 어시스턴트입니다.
아래 제공된 문서만을 근거로 답변해 주세요.
문서에 없는 내용은 "제공된 문서에서 해당 정보를 찾을 수 없습니다"라고 답변해 주세요.
답변 마지막에 참고한 문서의 출처를 명시해 주세요.

### 참고 문서:
{context}

### 질문:
{question}

### 답변:"""


@mlflow.trace(span_type="LLM")
def call_llm(prompt: str) -> str:
    """Foundation Model API로 LLM 호출"""
    response = client.predict(
        endpoint="databricks-meta-llama-3-3-70b-instruct",
        inputs={
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0.1,
            "max_tokens": 1000
        }
    )
    return response["choices"][0]["message"]["content"]

원본 문서:
"반품 정책: 구매 후 30일 이내 반품 가능합니다.
단, 전자제품은 14일 이내이며, 개봉 시 반품이 불가합니다."

--- 청크 경계에서 잘린 경우 ---
청크 1: "반품 정책: 구매 후 30일 이내 반품 가능합니다."
청크 2: "단, 전자제품은 14일 이내이며, 개봉 시 반품이 불가합니다."

→ 사용자 질문: "전자제품 반품 기한이 어떻게 되나요?"
→ 검색 결과: 청크 2만 검색됨
→ LLM 답변: "14일 이내입니다" (맞지만 "30일"이라는 일반 정책 맥락이 빠짐)

# 해결 방법: 메타데이터 필터링 + 시간 가중치
results = index.similarity_search(
    query_text=query,
    columns=["doc_id", "content", "title", "source", "updated_at"],
    num_results=20,
    filters={
        "updated_at >=": "2024-01-01"  # 최근 문서만 검색
    }
)

import mlflow

# 골든 데이터셋으로 평가
eval_results = mlflow.evaluate(
    model=rag_chain,
    data=golden_dataset,  # 질문 + 기대 답변 + 기대 검색 문서
    model_type="databricks-agent"
)

# 핵심 메트릭 확인
print(f"Retrieval Precision: {eval_results.metrics['retrieval/precision']}")
print(f"Answer Correctness: {eval_results.metrics['answer_correctness/llm_judged/average']}")
print(f"Faithfulness: {eval_results.metrics['faithfulness/llm_judged/average']}")