from langchain_core.prompts import ChatPromptTemplate

prompt = ChatPromptTemplate.from_messages([
    ("system", """당신은 주어진 컨텍스트를 기반으로 질문에 답변하는 AI 어시스턴트입니다.

규칙:
- 컨텍스트에 포함된 정보만 사용하여 답변하세요
- 컨텍스트에 답변이 없으면 "제공된 문서에서 관련 정보를 찾을 수 없습니다"라고 답변하세요
- 답변에 출처를 명시하세요

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

prompt_with_citation = ChatPromptTemplate.from_messages([
    ("system", """주어진 컨텍스트를 참고하여 질문에 답변하세요.

규칙:
- 각 주장 뒤에 [출처: 문서명] 형식으로 출처를 표기하세요
- 여러 문서의 정보를 종합할 경우 모든 출처를 명시하세요
- 컨텍스트에 없는 정보는 사용하지 마세요

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

from langchain_core.output_parsers import PydanticOutputParser
from pydantic import BaseModel, Field
from typing import List

class RAGAnswer(BaseModel):
    answer: str = Field(description="질문에 대한 답변")
    sources: List[str] = Field(description="답변에 사용된 출처 목록")
    confidence: str = Field(description="답변 신뢰도: high, medium, low")

parser = PydanticOutputParser(pydantic_object=RAGAnswer)

prompt_structured = ChatPromptTemplate.from_messages([
    ("system", """컨텍스트를 기반으로 질문에 답변하세요.

{format_instructions}

컨텍스트:
{context}"""),
    ("human", "{question}"),
]).partial(format_instructions=parser.get_format_instructions())

import tiktoken

def count_tokens(text: str, model: str = "gpt-4o-mini") -> int:
    """텍스트의 토큰 수를 계산합니다."""
    encoding = tiktoken.encoding_for_model(model)
    return len(encoding.encode(text))

def build_context(documents, max_tokens: int = 3000, model: str = "gpt-4o-mini"):
    """토큰 제한 내에서 컨텍스트를 구성합니다."""
    context_parts = []
    total_tokens = 0

    for i, doc in enumerate(documents):
        doc_text = f"[문서 {i+1}] {doc.page_content}"
        doc_tokens = count_tokens(doc_text, model)

        if total_tokens + doc_tokens > max_tokens:
            break

        context_parts.append(doc_text)
        total_tokens += doc_tokens

    return "\n\n".join(context_parts)

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

llm = init_chat_model("gpt-4o-mini", temperature=0)
compressor = LLMChainExtractor.from_llm(llm)

compression_retriever = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=retriever,
)

# 압축된 결과 검색
compressed_docs = compression_retriever.invoke("RAG의 생성 단계란?")

def format_context_with_sources(documents):
    """출처 정보가 포함된 컨텍스트를 구성합니다."""
    formatted = []
    for i, doc in enumerate(documents, 1):
        source = doc.metadata.get("source", "unknown")
        page = doc.metadata.get("page", "")
        source_info = f"{source}" + (f", p.{page}" if page else "")
        formatted.append(f"[{i}] (출처: {source_info})\n{doc.page_content}")
    return "\n\n---\n\n".join(formatted)

def generate_with_citations(question, documents, llm):
    """출처가 포함된 답변을 생성합니다."""
    context = format_context_with_sources(documents)

    prompt = ChatPromptTemplate.from_messages([
        ("system", """컨텍스트의 각 문서에는 번호가 부여되어 있습니다.
답변 시 해당 정보의 출처를 [1], [2] 형식으로 인라인 인용하세요.

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

    chain = prompt | llm | StrOutputParser()
    return chain.invoke({"context": context, "question": question})

from langchain.chat_models import init_chat_model
from langchain_core.output_parsers import StrOutputParser

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

prompt = ChatPromptTemplate.from_messages([
    ("system", "컨텍스트를 참고하여 답변하세요.\n\n{context}"),
    ("human", "{question}"),
])

chain = prompt | llm | StrOutputParser()

# 스트리밍 출력
for chunk in chain.stream({"context": context, "question": question}):
    print(chunk, end="", flush=True)

from langchain.chat_models import init_chat_model
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

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

def generate(state: GraphState) -> GraphState:
    """검색된 문서를 기반으로 답변을 생성합니다."""
    question = state["question"]
    documents = state["documents"]

    # 컨텍스트 구성
    context = "\n\n".join(
        f"[{i+1}] {doc.page_content}"
        for i, doc in enumerate(documents)
    )

    prompt = ChatPromptTemplate.from_messages([
        ("system", """다음 컨텍스트를 참고하여 질문에 답변하세요.
컨텍스트에 없는 내용은 답변하지 마세요.

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

    chain = prompt | llm | StrOutputParser()
    generation = chain.invoke({"context": context, "question": question})
    return {"generation": generation}