LangChain RAG 환각 줄이기 - Citations 강제

from langchain_core.documents import Document


def make_docs(chunks, source: str, doc_id: str):
    docs = []
    for i, text in enumerate(chunks):
        docs.append(
            Document(
                page_content=text,
                metadata={
                    "doc_id": doc_id,
                    "chunk_id": i,
                    "source": source,
                },
            )
        )
    return docs

from typing import List, Optional
from pydantic import BaseModel, Field


class Citation(BaseModel):
    doc_id: str
    chunk_id: int
    source: str
    quote: str = Field(..., description="원문에서 그대로 가져온 짧은 인용")


class AnswerSentence(BaseModel):
    text: str
    citations: List[Citation] = Field(default_factory=list)


class RAGAnswer(BaseModel):
    sentences: List[AnswerSentence]
    refused: bool = False
    refusal_reason: Optional[str] = None

from langchain_core.prompts import ChatPromptTemplate

prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            """
너는 RAG 답변 생성기다.
규칙:
- 출력은 지정된 JSON 스키마를 따른다.
- 각 문장에는 최소 1개의 citations가 있어야 한다.
- citations.quote는 제공된 컨텍스트에서 그대로 복사한 짧은 문장 조각이어야 한다.
- 컨텍스트에 근거가 없으면 refused를 true로 하고 refusal_reason에 이유를 적는다.
""",
        ),
        (
            "user",
            """
질문: {question}

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

from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4.1-mini", temperature=0)

structured_llm = llm.with_structured_output(RAGAnswer)

chain = prompt | structured_llm

def format_docs_for_context(docs):
    parts = []
    for d in docs:
        header = f"[doc_id={d.metadata['doc_id']} chunk_id={d.metadata['chunk_id']} source={d.metadata['source']}]"
        parts.append(header + "\n" + d.page_content)
    return "\n\n".join(parts)

import re


def validate_answer(answer: RAGAnswer, docs) -> list[str]:
    errors = []

    # 컨텍스트에서 허용되는 (doc_id, chunk_id) 집합
    allowed = {(d.metadata["doc_id"], d.metadata["chunk_id"]) for d in docs}
    context_text = "\n".join([d.page_content for d in docs])

    for i, s in enumerate(answer.sentences):
        if not s.citations:
            errors.append(f"sentence[{i}] has no citations")
            continue

        for j, c in enumerate(s.citations):
            if (c.doc_id, c.chunk_id) not in allowed:
                errors.append(f"sentence[{i}].citations[{j}] invalid doc_id/chunk_id")
            # quote가 컨텍스트에 실제 포함되는지(너무 엄격하면 false negative 가능)
            q = re.sub(r"\s+", " ", c.quote).strip()
            if q and q not in re.sub(r"\s+", " ", context_text):
                errors.append(f"sentence[{i}].citations[{j}] quote not found in context")

    return errors

from langchain_core.messages import SystemMessage, HumanMessage


def run_with_repair(question: str, retriever, max_retries: int = 2):
    docs = retriever.get_relevant_documents(question)
    context = format_docs_for_context(docs)

    result = chain.invoke({"question": question, "context": context})
    errors = validate_answer(result, docs)

    tries = 0
    while errors and tries < max_retries:
        tries += 1

        repair_prompt = ChatPromptTemplate.from_messages(
            [
                (
                    "system",
                    """
너는 RAG 답변 JSON을 고치는 역할이다.
규칙:
- 기존 답변을 유지하되, 오류가 난 citations만 수정한다.
- quote는 반드시 컨텍스트에서 그대로 복사한다.
- 근거를 찾을 수 없는 문장은 삭제하거나 refused로 전환한다.
""",
                ),
                (
                    "user",
                    """
질문: {question}

컨텍스트:
{context}

기존 JSON:
{bad_json}

검증 오류 목록:
{errors}
""",
                ),
            ]
        )

        repair_chain = repair_prompt | structured_llm
        result = repair_chain.invoke(
            {
                "question": question,
                "context": context,
                "bad_json": result.model_dump_json(),
                "errors": "\n".join(errors),
            }
        )
        errors = validate_answer(result, docs)

    return result, docs, errors

from langchain_core.runnables import RunnablePassthrough

# retriever는 이미 구성되어 있다고 가정

def build_chain(retriever):
    def get_context(q):
        docs = retriever.get_relevant_documents(q)
        return {
            "docs": docs,
            "context": format_docs_for_context(docs),
        }

    rag = (
        RunnablePassthrough()
        .assign(**{"bundle": lambda x: get_context(x["question"])})
        .assign(
            context=lambda x: x["bundle"]["context"],
        )
        | chain
    )
    return rag