RAG 완벽 가이드: Naive · Advanced · Graph RAG 통합본

"""example_1_naive_rag.py — Naive RAG 5단계 최소 구현"""
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_anthropic import ChatAnthropic
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

# ── 1) 문서 (실무에선 PDF/Notion/DB에서 로드) ─────────────
raw = [
    Document(page_content=(
        "ACME사 연차 정책: 정규직은 입사 1년이 지나면 연 15일의 연차가 부여된다. "
        "근속 3년 이상 16일, 5년 이상 18일, 10년 이상 20일. 미사용 연차는 다음 해 "
        "6월 30일까지 이월 가능, 이후 소멸."),
        metadata={"source": "HR/연차정책_v3.md"}),
    Document(page_content=(
        "특별휴가: 본인 결혼 5일, 자녀 결혼 1일, 배우자 출산 10일, "
        "본인/배우자 부모 사망 5일, 조부모 사망 3일. 부모 환갑·칠순 등 일반 가족행사는 "
        "별도 특별휴가 없이 연차로 처리한다."),
        metadata={"source": "HR/특별휴가.md"}),
    Document(page_content=(
        "재택근무: 정규직 주 2회 가능. 사전 팀장 승인 필요. 화·목 재택 비권장(전사 회의). "
        "신입사원은 입사 후 3개월간 전일 출근."),
        metadata={"source": "HR/재택근무정책.md"}),
]

# ── 2) 청킹 ─────────────────────────────────────────────
splitter = RecursiveCharacterTextSplitter(
    chunk_size=300, chunk_overlap=50,
    separators=["\n\n", "\n", ". ", " ", ""])
chunks = splitter.split_documents(raw)

# ── 3) 임베딩 + 벡터 저장 ────────────────────────────────
emb = HuggingFaceEmbeddings(
    model_name="BAAI/bge-m3",
    encode_kwargs={"normalize_embeddings": True})
vectordb = Chroma.from_documents(chunks, emb, collection_name="acme_hr")

# ── 4) 검색기 ───────────────────────────────────────────
retriever = vectordb.as_retriever(search_kwargs={"k": 3})

# ── 5) 프롬프트 + LLM + 체인 ────────────────────────────
prompt = ChatPromptTemplate.from_messages([
    ("system",
     "ACME사 HR 어시스턴트. 아래 [참고 문서]만 근거로 답하세요. "
     "찾을 수 없으면 '제공된 문서에서 답을 찾을 수 없습니다.'라고 답하세요. "
     "각 주장 끝에 [파일명]으로 출처를 다세요."),
    ("human", "[참고 문서]\n{context}\n\n[질문]\n{question}")])

llm = ChatAnthropic(model="claude-opus-4-7", temperature=0)

def fmt(docs):
    return "\n\n".join(f"[{d.metadata['source']}]\n{d.page_content}" for d in docs)

# LCEL: LangChain Expression Language. `|` 파이프로 컴포넌트를 연결한다.
# 유닉스 셸의 `cat file | grep ... | wc -l` 과 같은 발상.
chain = ({"context": retriever | fmt, "question": RunnablePassthrough()}
         | prompt | llm | StrOutputParser())

# ── 실행 ───────────────────────────────────────────────
for q in [
    "입사 1년차에 부모님 환갑이면 휴가를 며칠 받을 수 있나요?",
    "근속 7년차의 연차는?",
    "신입사원도 재택근무 가능한가요?",
    "회사 점심값 지원되나요?",  # 문서에 없음 → 모른다고 답해야 정상
]:
    print(f"\n━━━ Q: {q}\n> {chain.invoke(q)}")

from langchain_experimental.text_splitter import SemanticChunker
from langchain_community.embeddings import HuggingFaceEmbeddings

splitter = SemanticChunker(
    HuggingFaceEmbeddings(model_name="BAAI/bge-m3"),
    breakpoint_threshold_type="percentile",  # 또는 "standard_deviation", "interquartile"
    breakpoint_threshold_amount=95)
chunks = splitter.create_documents([long_text])

from langchain_core.prompts import ChatPromptTemplate
from langchain_anthropic import ChatAnthropic

llm = ChatAnthropic(model="claude-opus-4-7", temperature=0)
hyde_prompt = ChatPromptTemplate.from_messages([
    ("system", "다음 질문에 대해 사실 여부와 무관하게 그럴듯한 한 문단의 답변을 작성하세요."),
    ("human", "{question}")
])

def hyde_search(question: str, retriever):
    hypothetical = (hyde_prompt | llm).invoke({"question": question}).content
    # 가상 답변을 임베딩해서 검색 (질문 자체로 검색하는 것보다 보통 더 정확)
    return retriever.invoke(hypothetical)

from langchain_community.retrievers import BM25Retriever
from langchain.retrievers import EnsembleRetriever

bm25 = BM25Retriever.from_documents(chunks); bm25.k = 10
dense = vectordb.as_retriever(search_kwargs={"k": 10})
hybrid = EnsembleRetriever(
    retrievers=[bm25, dense],
    weights=[0.4, 0.6])  # 도메인에 따라 튜닝 (코드/식별자 많으면 BM25 비중↑)

from langchain.retrievers.document_compressors import LLMChainExtractor
from langchain.retrievers import ContextualCompressionRetriever

compressor = LLMChainExtractor.from_llm(llm)  # 질문에 답하는 데 필요한 부분만 추출
compressed = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=hybrid)

"""example_2_advanced_rag.py — 실전형 Advanced RAG 통합"""
from __future__ import annotations
from typing import List
from dataclasses import dataclass

from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.retrievers import BM25Retriever
from langchain.retrievers import EnsembleRetriever
from langchain_anthropic import ChatAnthropic
from langchain_core.prompts import ChatPromptTemplate
from sentence_transformers import CrossEncoder

# ════════════════════════════════════════════════════════════════
# 0. 데이터 — 엔지니어링 정책 문서 6종
# ════════════════════════════════════════════════════════════════
KB = [
    {"source": "ENG/Coding_Standards.md",
     "content": "파이썬은 PEP 8과 black 포매터 준수. 라인 100자. 모든 public 함수에 타입힌트. "
                "함수명 snake_case, 클래스 PascalCase, 상수 UPPER_SNAKE_CASE."},
    {"source": "ENG/Code_Review_Policy.md",
     "content": "PR 머지에는 최소 2명 승인 필요. 시니어 1명 필수. 보안 변경은 보안팀 추가 승인. "
                "PR 크기 400줄 권장, 1000줄 초과 시 분할 요청. 리뷰는 2영업일 이내."},
    {"source": "ENG/Deploy_Process.md",
     "content": "프로덕션 배포는 화·수·목 10:00~16:00. 금요일/공휴일 전일 금지. "
                "staging에서 24시간 검증 후 배포. 핫픽스는 CTO 승인 시 시간 제약 우회. "
                "배포 후 30분 모니터링 대기."},
    {"source": "ENG/On_Call_Policy.md",
     "content": "1주 단위 온콜 로테이션. P1 15분, P2 1시간 대응 목표. "
                "야간(22:00~08:00)/주말 온콜은 시간당 별도 보상. 휴가는 사전 교대 필수."},
    {"source": "ENG/Tech_Stack.md",
     "content": "백엔드 Python 3.12 + FastAPI 표준. DB: PostgreSQL 16, 캐시: Redis 7, "
                "MQ: RabbitMQ. 프론트 TypeScript + React 18. AWS(ECS/RDS/S3) + Terraform."},
    {"source": "HR/Remote_Work.md",
     "content": "정규직 주 2회 재택. 팀장 승인 필요. 화·목 재택 비권장. "
                "신입은 입사 3개월 전일 출근. 해외 원격은 별도 승인 + 세무 검토."},
]

# ════════════════════════════════════════════════════════════════
# 1. 인덱싱 — Hybrid 검색을 위해 Dense + BM25 둘 다 구축
# ════════════════════════════════════════════════════════════════
def build_retrievers():
    docs = [Document(page_content=d["content"], metadata={"source": d["source"]}) for d in KB]
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=350, chunk_overlap=70,
        separators=["\n\n", "\n", ". ", " ", ""])
    chunks = splitter.split_documents(docs)

    emb = HuggingFaceEmbeddings(
        model_name="BAAI/bge-m3", encode_kwargs={"normalize_embeddings": True})
    vectordb = Chroma.from_documents(chunks, emb, collection_name="adv_rag")

    dense = vectordb.as_retriever(search_kwargs={"k": 8})
    bm25 = BM25Retriever.from_documents(chunks); bm25.k = 8

    hybrid = EnsembleRetriever(retrievers=[bm25, dense], weights=[0.4, 0.6])
    return hybrid

# ════════════════════════════════════════════════════════════════
# 2. 쿼리 변환 — Multi-Query 생성으로 질문 다양화
# ════════════════════════════════════════════════════════════════
multi_query_prompt = ChatPromptTemplate.from_messages([
    ("system", "사용자 질문을 의미는 유지하되 표현·관점이 다른 3개의 검색 쿼리로 재작성하세요. "
               "각 줄에 하나씩, 번호 없이 출력하세요."),
    ("human", "{question}")
])

llm = ChatAnthropic(model="claude-opus-4-7", temperature=0)

def multi_queries(question: str) -> List[str]:
    raw = (multi_query_prompt | llm).invoke({"question": question}).content
    qs = [q.strip("-•123456789. ").strip() for q in raw.split("\n") if q.strip()]
    return [question] + qs[:3]  # 원본 + 변형 3개

# ════════════════════════════════════════════════════════════════
# 3. 재랭킹 — Cross-encoder로 후보 정렬
# ════════════════════════════════════════════════════════════════
class Reranker:
    def __init__(self, name="BAAI/bge-reranker-v2-m3"):
        self.m = CrossEncoder(name, max_length=512)

    def __call__(self, query: str, docs: List[Document], top_n=4) -> List[Document]:
        if not docs: return []
        scores = self.m.predict([(query, d.page_content) for d in docs])
        ranked = sorted(zip(scores, docs), key=lambda x: x[0], reverse=True)
        # 중복 제거 (page_content 기준)
        seen, out = set(), []
        for s, d in ranked:
            if d.page_content in seen: continue
            seen.add(d.page_content)
            d.metadata["rerank_score"] = float(s)
            out.append(d)
            if len(out) == top_n: break
        return out

# ════════════════════════════════════════════════════════════════
# 4. 답변 생성 — 인용 강제
# ════════════════════════════════════════════════════════════════
answer_prompt = ChatPromptTemplate.from_messages([
    ("system",
     "ACME사 엔지니어링 어시스턴트. 아래 [참고 문서]만 근거로 답하세요.\n"
     "규칙:\n"
     "1) 모든 사실 끝에 [번호] 형식 출처 표시.\n"
     "2) 여러 문서가 뒷받침하면 [1][3]처럼 다수 표기.\n"
     "3) 문서에 없으면 '문서에 명시되어 있지 않습니다'라고 답.\n"
     "4) 간결한 한국어로 핵심만."),
    ("human", "[참고 문서]\n{context}\n\n[질문]\n{question}")
])

@dataclass
class Result:
    answer: str
    sources: List[Document]
    queries_used: List[str]

def make_ctx(docs: List[Document]) -> str:
    return "\n\n".join(
        f"[{i}] (출처: {d.metadata['source']})\n{d.page_content}"
        for i, d in enumerate(docs, 1))

def advanced_rag(question: str, hybrid, reranker) -> Result:
    # ① Multi-Query 변환
    queries = multi_queries(question)

    # ② Hybrid 검색 (각 변형 쿼리마다)
    candidates: List[Document] = []
    seen = set()
    for q in queries:
        for d in hybrid.invoke(q):
            key = d.page_content
            if key not in seen:
                seen.add(key); candidates.append(d)

    # ③ 재랭킹 (원본 질문 기준으로 정밀 정렬)
    top = reranker(question, candidates, top_n=4)

    # ④ 답변 생성 (인용 강제)
    msg = answer_prompt.invoke({"context": make_ctx(top), "question": question})
    ans = llm.invoke(msg).content
    return Result(answer=ans, sources=top, queries_used=queries)

# ════════════════════════════════════════════════════════════════
# 5. Self-Evaluation — Faithfulness 자동 점검
# ════════════════════════════════════════════════════════════════
judge_prompt = ChatPromptTemplate.from_messages([
    ("system", "RAG 답변 충실도(faithfulness) 심사. [참고 문서]에 답변의 모든 사실이 "
               "뒷받침되면 PASS, 하나라도 근거 없으면 FAIL. 첫 줄에 PASS/FAIL, 둘째 줄부터 사유."),
    ("human", "[참고 문서]\n{context}\n\n[답변]\n{answer}\n\n평가:")
])

def judge(res: Result) -> str:
    msg = judge_prompt.invoke({"context": make_ctx(res.sources), "answer": res.answer})
    return llm.invoke(msg).content

# ════════════════════════════════════════════════════════════════
# 6. 실행
# ════════════════════════════════════════════════════════════════
if __name__ == "__main__":
    hybrid = build_retrievers()
    rerank = Reranker()

    for q in [
        "보안 관련 PR 머지하려면 누가 승인해줘야 해?",
        "금요일 오후 핫픽스 배포 가능?",
        "신입사원 재택근무 신청 가능?",
        "DB와 캐시는 뭘 쓰지?",
        "회사 점심 메뉴는?",  # 문서에 없음
    ]:
        print(f"\n{'='*72}\nQ: {q}")
        r = advanced_rag(q, hybrid, rerank)
        print(f"\n변형 쿼리들: {r.queries_used}")
        print(f"\n검색+재랭킹 top-{len(r.sources)}:")
        for i, d in enumerate(r.sources, 1):
            print(f"  [{i}] {d.metadata['source']:28s}  "
                  f"score={d.metadata.get('rerank_score',0):+.2f}")
        print(f"\n> 답변:\n{r.answer}")
        print(f"\nSelf-Eval:\n{judge(r)}")

"""example_3_graph_rag.py — NetworkX 기반 미니 GraphRAG"""
from __future__ import annotations
import json
from typing import List, Tuple, Dict, Set
from dataclasses import dataclass, field

import networkx as nx
from langchain_anthropic import ChatAnthropic
from langchain_core.prompts import ChatPromptTemplate

# ════════════════════════════════════════════════════════════════
# 0. 데이터 — 인물·프로젝트·조직 관계가 풍부한 가상 회사 위키
# ════════════════════════════════════════════════════════════════
DOCS = [
    "김철수는 ACME사의 CTO이며 2019년에 입사했다. 이전에는 BlueTech의 시니어 엔지니어였다. "
    "그는 현재 프로젝트 알파를 총괄하고 있으며, 머신러닝 인프라 팀의 디렉터를 겸임한다.",

    "박영희는 ACME사의 보안팀장으로, 프로젝트 알파의 보안 책임자를 맡고 있다. "
    "이전에 SecureCorp에서 10년간 근무했고, 김철수와는 BlueTech 시절부터 동료였다.",

    "이민수는 ACME사 데이터플랫폼 팀의 시니어 엔지니어이다. 프로젝트 알파의 데이터 파이프라인을 담당하며, "
    "김철수에게 직접 보고한다. 박영희와는 보안 감사 작업으로 협업 중이다.",

    "최지훈은 프로젝트 베타의 PM이다. 베타 프로젝트는 신규 결제 시스템 구축을 목표로 하며, "
    "이민수가 베타에도 일부 참여해 데이터 마이그레이션을 지원하고 있다.",

    "프로젝트 알파는 ACME사의 차세대 추천 시스템이며 2024년 1월에 시작되었다. "
    "프로젝트 베타는 2024년 6월에 시작된 결제 시스템 프로젝트이다. "
    "두 프로젝트는 모두 머신러닝 인프라 팀의 지원을 받는다.",
]

# ════════════════════════════════════════════════════════════════
# 1. LLM으로 엔티티/관계 트리플 추출
# ════════════════════════════════════════════════════════════════
llm = ChatAnthropic(model="claude-opus-4-7", temperature=0)

extract_prompt = ChatPromptTemplate.from_messages([
    ("system",
     "다음 문서에서 엔티티와 관계를 추출해 JSON 트리플 리스트로 출력하세요.\n"
     "각 트리플은 {\"s\": 주체, \"r\": 관계, \"o\": 객체} 형식.\n"
     "엔티티는 사람·조직·프로젝트·역할 등 명확한 개체만 포함.\n"
     "관계는 짧은 동사구로 (예: WORKS_AT, LEADS, IS_CTO_OF, REPORTS_TO, COLLABORATES_WITH).\n"
     "동일 인물의 다른 표현은 같은 이름으로 통일하세요.\n"
     "출력은 순수 JSON 배열, 다른 텍스트 금지."),
    ("human", "{document}")
])

def _safe_json_parse(text: str, default):
    """LLM 응답에서 JSON 부분만 추출 (markdown fence 등 제거)"""
    import re
    m = re.search(r"(\[.*\]|\{.*\})", text, re.DOTALL)
    if not m:
        return default
    try:
        return json.loads(m.group(1))
    except json.JSONDecodeError:
        return default

def extract_triples(doc: str) -> List[Dict]:
    raw = (extract_prompt | llm).invoke({"document": doc}).content
    return _safe_json_parse(raw, default=[])

# ════════════════════════════════════════════════════════════════
# 2. NetworkX 그래프 구축 (+ 원본 문서 인덱스)
# ════════════════════════════════════════════════════════════════
@dataclass
class KnowledgeGraph:
    G: nx.MultiDiGraph = field(default_factory=nx.MultiDiGraph)
    # 엔티티 → 그 엔티티가 등장한 원본 문서 인덱스 집합
    ent2docs: Dict[str, Set[int]] = field(default_factory=dict)
    docs: List[str] = field(default_factory=list)

def build_kg(docs: List[str]) -> KnowledgeGraph:
    kg = KnowledgeGraph(docs=docs)
    for i, d in enumerate(docs):
        triples = extract_triples(d)
        for t in triples:
            s, r, o = t.get("s"), t.get("r"), t.get("o")
            if not (s and r and o): continue
            kg.G.add_edge(s, o, relation=r, doc_idx=i)
            kg.ent2docs.setdefault(s, set()).add(i)
            kg.ent2docs.setdefault(o, set()).add(i)
    return kg

# ════════════════════════════════════════════════════════════════
# 3. 질의에서 엔티티 추출
# ════════════════════════════════════════════════════════════════
query_ent_prompt = ChatPromptTemplate.from_messages([
    ("system", "질문에서 등장한 엔티티(사람·조직·프로젝트·역할 등)만 JSON 배열로 추출하세요. "
               "예: [\"김철수\", \"프로젝트 알파\"]. 다른 텍스트 금지."),
    ("human", "{question}")
])

def extract_query_entities(q: str) -> List[str]:
    raw = (query_ent_prompt | llm).invoke({"question": q}).content
    return _safe_json_parse(raw, default=[])

# ════════════════════════════════════════════════════════════════
# 4. 그래프 탐색 — 질의 엔티티 주변 N-hop 서브그래프
# ════════════════════════════════════════════════════════════════
def find_node(kg: KnowledgeGraph, name: str) -> str | None:
    """정확 매칭 우선, 실패 시 부분 문자열 매칭으로 노드 찾기"""
    if name in kg.G: return name
    for n in kg.G.nodes:
        if name in n or n in name:
            return n
    return None

def subgraph_around(kg: KnowledgeGraph, entities: List[str], hops: int = 2) -> Tuple[nx.MultiDiGraph, Set[int]]:
    """질의 엔티티들을 시드로 N-hop 이웃까지 모은 서브그래프 + 관련 문서 인덱스"""
    seed_nodes = {n for e in entities if (n := find_node(kg, e))}
    if not seed_nodes:
        return nx.MultiDiGraph(), set()

    # 무방향 그래프로 변환해 양방향 BFS
    undirected = kg.G.to_undirected()
    visited = set(seed_nodes)
    frontier = set(seed_nodes)
    for _ in range(hops):
        next_frontier = set()
        for n in frontier:
            if n not in undirected: continue
            next_frontier.update(undirected.neighbors(n))
        frontier = next_frontier - visited
        visited |= frontier

    sub = kg.G.subgraph(visited).copy()
    # 관련 문서 인덱스 모으기
    doc_ids = set()
    for n in visited:
        doc_ids.update(kg.ent2docs.get(n, set()))
    return sub, doc_ids

def serialize_subgraph(sub: nx.MultiDiGraph) -> str:
    """서브그래프를 LLM에 전달할 텍스트로 변환"""
    if sub.number_of_edges() == 0:
        return "(관련 그래프 없음)"
    lines = []
    for u, v, data in sub.edges(data=True):
        lines.append(f"({u}) -[{data['relation']}]-> ({v})")
    return "\n".join(sorted(set(lines)))

# ════════════════════════════════════════════════════════════════
# 5. 답변 생성 — 그래프 + 원본 문서 둘 다 컨텍스트로 제공
# ════════════════════════════════════════════════════════════════
graph_answer_prompt = ChatPromptTemplate.from_messages([
    ("system",
     "당신은 사내 지식 어시스턴트입니다. 아래 [지식 그래프]와 [원본 문서]만 근거로 답하세요.\n"
     "- 그래프의 관계를 따라가며 다중 홉 추론을 수행하세요.\n"
     "- 답변에 사용한 관계는 (A) -[관계]-> (B) 형식으로 인용하세요.\n"
     "- 근거가 부족하면 '제공된 정보로는 답할 수 없습니다.'라고 답하세요."),
    ("human",
     "[지식 그래프]\n{graph}\n\n[원본 문서]\n{docs}\n\n[질문]\n{question}")
])

def graph_rag(question: str, kg: KnowledgeGraph) -> str:
    ents = extract_query_entities(question)
    sub, doc_ids = subgraph_around(kg, ents, hops=2)

    graph_text = serialize_subgraph(sub)
    doc_text = "\n\n".join(f"[doc{i}] {kg.docs[i]}" for i in sorted(doc_ids)) or "(관련 문서 없음)"

    print(f"  · 질의 엔티티: {ents}")
    print(f"  · 서브그래프 노드 {sub.number_of_nodes()}개, 엣지 {sub.number_of_edges()}개")
    print(f"  · 관련 원본 문서: {sorted(doc_ids)}")

    msg = graph_answer_prompt.invoke({
        "graph": graph_text, "docs": doc_text, "question": question})
    return llm.invoke(msg).content

# ════════════════════════════════════════════════════════════════
# 6. 실행
# ════════════════════════════════════════════════════════════════
if __name__ == "__main__":
    print("[1/3] 그래프 인덱싱 중...")
    kg = build_kg(DOCS)
    print(f"     완료. 노드 {kg.G.number_of_nodes()}개, 엣지 {kg.G.number_of_edges()}개\n")

    # 그래프 미리보기
    print("[2/3] 추출된 트리플 (전체):")
    for u, v, data in kg.G.edges(data=True):
        print(f"     ({u}) -[{data['relation']}]-> ({v})  (문서{data['doc_idx']})")

    # 다중 홉이 진짜로 필요한 질문들
    print("\n[3/3] Graph RAG 질의응답:")
    for q in [
        # 1-hop: 단순 사실
        "김철수는 어느 회사의 CTO야?",
        # 2-hop: 다중 홉 — 김철수의 이전 회사 동료가 누구야?
        "박영희와 김철수는 어떻게 알게 됐어?",
        # 관계 교집합: 두 사람의 공통 프로젝트
        "이민수와 박영희가 함께 일하는 영역은?",
        # 다중 홉 + 집계: 한 명이 여러 프로젝트에 연결
        "이민수가 참여한 프로젝트는 어떤 것들이고, 그 프로젝트의 다른 핵심 멤버는 누구야?",
        # 그래프에 없는 정보
        "김철수의 연봉은?",
    ]:
        print(f"\n━━━ Q: {q}")
        print(f"> {graph_rag(q, kg)}")

from langchain_neo4j import Neo4jGraph
from langchain_experimental.graph_transformers import LLMGraphTransformer

graph = Neo4jGraph(url=..., username=..., password=...)
transformer = LLMGraphTransformer(llm=llm)
graph_documents = transformer.convert_to_graph_documents(docs)
graph.add_graph_documents(graph_documents)

from ragas import evaluate
from ragas.metrics import faithfulness, answer_relevancy, context_precision, context_recall
from datasets import Dataset

eval_data = Dataset.from_dict({
    "question": ["입사 1년차 부모님 환갑 휴가는?"],
    "answer": ["일반 가족행사이므로 별도 특별휴가 없이 연차로 처리됩니다 [HR/특별휴가.md]."],
    "contexts": [["부모 환갑·칠순 등 일반 가족행사는 별도 특별휴가 없이 연차로 처리한다."]],
    "ground_truth": ["연차로 처리해야 합니다."],
})

result = evaluate(eval_data, metrics=[
    faithfulness, answer_relevancy, context_precision, context_recall])
print(result)

python3 -m venv .venv && source .venv/bin/activate

pip install \
  langchain>=0.3.0 \
  langchain-community>=0.3.0 \
  langchain-anthropic>=0.3.0 \
  langchain-text-splitters>=0.3.0 \
  langchain-experimental>=0.3.0 \
  chromadb>=0.5.0 \
  sentence-transformers>=3.0.0 \
  rank_bm25>=0.2.2 \
  networkx>=3.2 \
  tiktoken>=0.7.0

export ANTHROPIC_API_KEY="sk-ant-..."

python example_1_naive_rag.py
python example_2_advanced_rag.py
python example_3_graph_rag.py
python example_4_llm_wiki.py

"""example_4_llm_wiki.py — LLM Wiki 패턴 미니 구현
==================================================
단일 파일에서 다음을 시연한다:
  1) 시간순 소스 3개를 차례로 ingest
  2) LLM이 매 ingest마다 entities/projects/concepts 페이지를 작성·갱신
  3) index.md / log.md 자동 유지
  4) 위키 자체를 컨텍스트로 한 시간 진화 질의
"""
from __future__ import annotations
import json, re, datetime, shutil
from pathlib import Path
from typing import Dict
from langchain_anthropic import ChatAnthropic
from langchain_core.prompts import ChatPromptTemplate

# ════════════════════════════════════════════════════════════════
# 0. 샘플 raw 소스 — 가상 회사의 시간순 3문서
# ════════════════════════════════════════════════════════════════
SAMPLE_SOURCES = {
    "2024Q3_strategy.md": """# Q3 2024 전략 회의 요약

CTO 김철수가 발표한 Q3 우선순위:
1. 프로젝트 알파(추천 시스템 개편) 11월 출시 목표
2. 데이터 인프라 팀 인원 50% 증원
3. 보안팀과의 협업 강화

알파 프로젝트는 박영희가 보안 책임자로 합류.
이민수가 데이터 파이프라인을 담당한다.""",

    "2024Q4_alpha_launch.md": """# 프로젝트 알파 출시 회고 (2024-11-30)

11월 15일 출시. 트래픽 30% 증가, 클릭률 12% 상승.

핵심 기여자:
- 김철수 (총괄)
- 박영희 (보안 검수)
- 이민수 (데이터 파이프라인)
- 정현우 (UI/UX, 신규 합류)

실패 요인: 초기 캐시 미스 폭증 → Redis 클러스터 증설로 해결
다음 단계: 베타 프로젝트(결제 시스템) 시작.""",

    "2025Q1_orgchange.md": """# 2025년 1월 조직 개편

- 김철수: CTO 유지. 머신러닝 인프라 팀 디렉터 겸임
- 박영희: 보안팀장으로 승진
- 이민수: 데이터플랫폼 팀 리더로 승진
- 정현우: 알파 팀에서 베타 팀으로 이동
- 최지훈: 베타 팀 PM으로 영입 (외부 입사)

알파 프로젝트는 유지보수 모드로 전환. 베타 프로젝트가 신규 우선순위."""
}

# ════════════════════════════════════════════════════════════════
# 1. 프롬프트 — Ingest용(액션 계획) / Query용
# ════════════════════════════════════════════════════════════════
INGEST_PROMPT = ChatPromptTemplate.from_messages([
    ("system",
     "당신은 위키 편집자입니다. 새 소스 문서를 보고 위키 페이지들을 어떻게 갱신할지 결정합니다.\n"
     "현재 존재하는 위키 페이지 전체와 그 내용이 함께 주어집니다.\n\n"
     "출력은 순수 JSON 배열이어야 합니다 (다른 설명 금지):\n"
     '  [{"op":"create","path":"entities/김철수.md","content":"전체 마크다운"},\n'
     '   {"op":"append","path":"projects/알파.md","content":"추가할 마크다운"}]\n\n'
     "규칙:\n"
     "- 사람: entities/이름.md   프로젝트: projects/이름.md   개념: concepts/이름.md\n"
     "- 새 정보가 있을 때만 액션. 단순 중복이면 빈 배열.\n"
     "- 페이지 본문에 [[다른_페이지명]] 형식 위키링크를 적극 활용.\n"
     "- 페이지 끝에 '> 출처: [소스파일명]' 표기 (append 시에도).\n"
     "- 모순되는 정보가 있으면 'TODO: 모순 검토 - ...' 메모 추가."),
    ("human",
     "[새 소스: {source_name}]\n{source_text}\n\n"
     "[현재 위키]\n{existing_pages}\n\n"
     "위 소스를 반영할 액션 JSON 배열:")
])

QUERY_PROMPT = ChatPromptTemplate.from_messages([
    ("system",
     "당신은 위키 어시스턴트입니다. 아래 [위키 페이지들]만 근거로 질문에 답하세요.\n"
     "각 사실 끝에 [[페이지명]] 형식으로 출처 페이지를 인용하세요.\n"
     "근거가 부족하면 '위키에 정보가 부족합니다'라고 답하세요."),
    ("human", "[위키 페이지들]\n{pages}\n\n[질문]\n{question}")
])


# ════════════════════════════════════════════════════════════════
# 2. WikiAgent — 핵심 로직
# ════════════════════════════════════════════════════════════════
class WikiAgent:
    def __init__(self, root: str):
        self.root = Path(root)
        self.raw_dir = self.root / "raw"
        self.wiki_dir = self.root / "wiki"
        self.llm = ChatAnthropic(model="claude-opus-4-7", temperature=0)

    # ── 셋업: 디렉터리 + 샘플 소스 작성 + 빈 인덱스/로그 ──
    def setup(self):
        if self.root.exists():
            shutil.rmtree(self.root)
        self.raw_dir.mkdir(parents=True)
        self.wiki_dir.mkdir(parents=True)
        for name, content in SAMPLE_SOURCES.items():
            (self.raw_dir / name).write_text(content, encoding="utf-8")
        (self.wiki_dir / "index.md").write_text("# 위키 인덱스\n\n", encoding="utf-8")
        (self.wiki_dir / "log.md").write_text("# 운영 로그\n\n", encoding="utf-8")

    # ── 현재 위키 페이지 모두 (path → content). 인덱스/로그 제외 ──
    def _list_pages(self) -> Dict[str, str]:
        out = {}
        for p in self.wiki_dir.rglob("*.md"):
            rel = p.relative_to(self.wiki_dir).as_posix()
            if rel in ("index.md", "log.md"):
                continue
            out[rel] = p.read_text(encoding="utf-8")
        return out

    @staticmethod
    def _parse_json_array(text: str):
        m = re.search(r"\[.*\]", text, re.DOTALL)
        if not m: return []
        try: return json.loads(m.group(0))
        except json.JSONDecodeError: return []

    # ── Ingest: 한 소스를 위키에 흡수 ──
    def ingest(self, source_name: str):
        print(f"\nIngest: {source_name}")
        source_text = (self.raw_dir / source_name).read_text(encoding="utf-8")

        pages = self._list_pages()
        existing = "(아직 페이지 없음)" if not pages else "\n\n".join(
            f"### {path}\n{content}" for path, content in pages.items())

        msg = INGEST_PROMPT.invoke({
            "source_name": source_name,
            "source_text": source_text,
            "existing_pages": existing,
        })
        actions = self._parse_json_array(self.llm.invoke(msg).content)

        # 액션 실행
        for a in actions:
            target = self.wiki_dir / a["path"]
            target.parent.mkdir(parents=True, exist_ok=True)
            if a["op"] == "create":
                target.write_text(a["content"].rstrip() + "\n", encoding="utf-8")
                print(f"    CREATE  {a['path']}")
            elif a["op"] == "append":
                cur = target.read_text(encoding="utf-8") if target.exists() else ""
                target.write_text(cur.rstrip() + "\n\n" + a["content"].rstrip() + "\n",
                                  encoding="utf-8")
                print(f"    APPEND  {a['path']}")

        self._update_index()
        self._append_log(f"ingest | {source_name} | actions={len(actions)}")

    # ── 인덱스 재생성: 카테고리별 그룹 + 한 줄 요약 ──
    def _update_index(self):
        pages = self._list_pages()
        groups: Dict[str, list] = {}
        for path in sorted(pages):
            cat = path.split("/")[0] if "/" in path else "root"
            groups.setdefault(cat, []).append(path)
        lines = ["# 위키 인덱스",
                 f"\n_갱신: {datetime.date.today()}_  / 페이지 {len(pages)}개\n"]
        for cat, paths in groups.items():
            lines.append(f"\n## {cat}")
            for p in paths:
                first = pages[p].splitlines()[0].lstrip("# ").strip()
                lines.append(f"- [[{p[:-3]}]] — {first}")
        (self.wiki_dir / "index.md").write_text("\n".join(lines) + "\n", encoding="utf-8")

    # ── 로그 append ──
    def _append_log(self, msg: str):
        line = f"## [{datetime.date.today()}] {msg}\n"
        with (self.wiki_dir / "log.md").open("a", encoding="utf-8") as f:
            f.write(line)

    # ── Query: 위키 전체를 컨텍스트로 답변 (소규모 데모) ──
    def query(self, question: str) -> str:
        # 프로덕션에선 인덱스 먼저 보고 페이지 선별 (= LLM이 도구로 page를 읽음)
        # 데모는 소규모이므로 모든 페이지를 한 번에 컨텍스트로
        pages = self._list_pages()
        joined = "\n\n".join(f"### [[{p[:-3]}]]\n{c}" for p, c in pages.items())
        msg = QUERY_PROMPT.invoke({"pages": joined, "question": question})
        return self.llm.invoke(msg).content


# ════════════════════════════════════════════════════════════════
# 3. 메인 실행 — 시간순 ingest 후 진화 질의
# ════════════════════════════════════════════════════════════════
if __name__ == "__main__":
    agent = WikiAgent("./demo_wiki")
    agent.setup()

    # 3개 소스를 시간순으로. 위키가 점점 풍부해지는 걸 관찰.
    for name in ["2024Q3_strategy.md", "2024Q4_alpha_launch.md", "2025Q1_orgchange.md"]:
        agent.ingest(name)

    # 최종 위키 트리
    print("\n최종 위키 구조:")
    for p in sorted(Path("./demo_wiki/wiki").rglob("*.md")):
        rel = p.relative_to("./demo_wiki/wiki")
        print(f"   {rel}  ({p.stat().st_size}B)")

    # 시간 진화에 대한 질의 — RAG로는 매우 어려운 종류
    # (여러 소스에 걸친 인물·역할 변화를 통합해야 함)
    print("\n참고: 위키 쿼리:")
    for q in [
        "프로젝트 알파의 핵심 멤버는 시간이 지나며 어떻게 변했어?",
        "박영희의 역할은 어떻게 진화했지?",
        "이민수가 알파와 베타 양쪽에 모두 관여하나? 어떻게?",
    ]:
        print(f"\n━━━ Q: {q}")
        print(f"> {agent.query(q)}")

# Schema for the LLM Wiki

## Directory layout
- `raw/` — immutable source documents
- `wiki/entities/` — person and organization pages
- `wiki/projects/` — project pages
- `wiki/concepts/` — conceptual / topic pages
- `wiki/index.md` — auto-maintained catalog
- `wiki/log.md` — append-only operation log

## Conventions
- Every fact ends with `> 출처: [filename]`
- Wiki links: `[[Page Name]]` (no .md extension)
- Person pages: H1 = full name, H2 sections: 직책 / 이력 / 프로젝트 / 관계
- Contradictions: insert `TODO: 모순 검토 — ...`

## Ingest workflow
1. Discuss key takeaways with the user first
2. Plan actions (create/append/update)
3. Execute actions
4. Update `index.md`
5. Append a line to `log.md` like `## [YYYY-MM-DD] ingest | <source>`

## Query workflow
1. Read `index.md` first
2. Open only relevant pages (don't load everything)
3. Cite each fact with `[[Page Name]]`
4. Offer to save the answer back as a new wiki page