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)

rewrite_prompt = ChatPromptTemplate.from_messages([
    ("system", """당신은 검색 쿼리 최적화 전문가입니다.
사용자의 질문을 벡터 데이터베이스 검색에 최적화된 쿼리로 변환하세요.

규칙:
- 모호한 표현을 구체적으로 변환
- 약어를 풀어서 작성
- 핵심 키워드를 포함
- 한 문장으로 작성"""),
    ("human", "원본 질문: {question}\n\n최적화된 검색 쿼리:"),
])

rewrite_chain = rewrite_prompt | llm | StrOutputParser()
optimized = rewrite_chain.invoke({"question": "RAG 좀 알려줘"})
# → "RAG(Retrieval-Augmented Generation)의 정의, 핵심 구성 요소, 동작 원리"

# 기술 문서용 쿼리 재작성
tech_rewrite_prompt = ChatPromptTemplate.from_messages([
    ("system", """당신은 기술 문서 검색 쿼리 최적화 전문가입니다.

규칙:
1. 약어를 풀어서 작성 (예: RAG → Retrieval-Augmented Generation)
2. 기술 용어를 정확히 사용
3. 동의어를 포함하여 검색 범위 확대
4. 구체적인 기술 키워드를 추가

예시:
- "LLM 환각" → "대규모 언어 모델(LLM)의 환각(hallucination) 문제 원인과 해결 방법"
- "벡터DB 비교" → "벡터 데이터베이스(Chroma, Qdrant, Milvus, Weaviate) 성능 비교와 특징"
"""),
    ("human", "원본 질문: {question}\n\n최적화된 검색 쿼리:"),
])

from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma

embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = Chroma(collection_name="docs", embedding_function=embeddings)
retriever = vectorstore.as_retriever(search_kwargs={"k": 4})

# 쿼리 재작성 → 검색 파이프라인
def rewrite_and_retrieve(question: str):
    optimized_query = rewrite_chain.invoke({"question": question})
    results = retriever.invoke(optimized_query)
    return results

results = rewrite_and_retrieve("RAG 좀 알려줘")

from langchain.chat_models import init_chat_model
from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

llm = init_chat_model("gpt-4o-mini", temperature=0)
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = Chroma(collection_name="docs", embedding_function=embeddings)

# 1. 가상 답변 생성
hyde_prompt = ChatPromptTemplate.from_messages([
    ("system", """질문에 대한 답변을 작성하세요.
정확하지 않아도 됩니다. 관련 키워드와 개념을 최대한 포함하세요.
답변은 실제 기술 문서처럼 작성하세요."""),
    ("human", "{question}"),
])

hyde_chain = hyde_prompt | llm | StrOutputParser()

# 2. 가상 답변으로 검색
def hyde_retrieve(question: str, k: int = 4):
    hypothetical_answer = hyde_chain.invoke({"question": question})
    results = vectorstore.similarity_search(hypothetical_answer, k=k)
    return results

results = hyde_retrieve("RAPTOR는 어떻게 동작하나요?")