# ============================================================
# [실습] RAG Chain 구축: 환각 방지와 출처
#
# 아래 항목을 직접 구현하세요:
#  1) "아래 문서만 참고하고, 없으면 모르겠다고 답하라"는 환각 방지 프롬프트를 정의하고,
#     RunnablePassthrough.assign 으로 question 을 보존하며 context 를 format_docs(retriever.invoke(...))
#     로 채운 뒤 prompt | llm | StrOutputParser() 로 잇는다.
#  2) 청크마다 [1],[2] 번호를 붙인 컨텍스트와 중복 제거한 출처 리스트를 만들고,
#     retriever 를 한 번만 호출해 question/context/sources 를 동시에 활용하는 체인을 조립한 뒤
#     answer 끝에 "참고 문서" 목록을 덧붙인다.
#  3) 각 문장 끝에 [1],[2] 근거 번호를 달도록 지시하는 프롬프트로 교체하고,
#     청크별 출처를 번호와 1:1 대응시키는 numbered_sources 로 (2)의 체인을 재구성한다.
# ============================================================

from dotenv import load_dotenv

load_dotenv()  # .env 에서 OPENAI_API_KEY 로드

from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough, RunnableLambda

# retriever 구성 (../DATA 로드·청킹 → Chroma → as_retriever(k=3))
# 여기에 코드를 작성하세요 (retriever)

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

# 1) 여기에 코드를 작성하세요
#    - 환각 방지 프롬프트(system: 아래 문서만 참고, 없으면 '모르겠습니다') 정의
#    - format_docs(docs) 정의
#    - chain = RunnablePassthrough.assign(context=...) | prompt | llm | StrOutputParser()
#    - chain.invoke({"question": "NVMe 와 SATA SSD 의 차이?"}) 출력

# 2) 여기에 코드를 작성하세요
#    - format_docs_with_index([1],[2] 번호) / extract_sources(중복 제거)
#    - retrieve_and_split → RunnablePassthrough.assign(answer=...) → append_sources 체인

# 3) 여기에 코드를 작성하세요
#    - 인라인 [n] 인용 프롬프트로 교체 + numbered_sources 로 (2) 체인 재구성
