# -*- coding: utf-8 -*-
"""HyDE RAG - generate hypothetical document to improve embedding."""
from .base import BaseRAG


class HyDERAG(BaseRAG):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def _generate_hypothetical_doc(self, query):
        prompt = (
            "请根据以下问题，生成一个假设的、包含相关信息的答案文档。\n\n"
            "问题：{}\n\n"
            "假设文档：".format(query)
        )
        return self._call_llm(prompt)

    def retrieve(self, query, k=10):
        hypothetical_doc = self._generate_hypothetical_doc(query)
        hypo_embedding = self.embedding_model.embed_query(hypothetical_doc)
        return self._deduplicate_results(self.vector_store.similarity_search(hypo_embedding, k), k)

    def generate(self, query, context):
        prompt = (
            "根据以下参考文档（基于假设文档检索），回答问题。\n\n"
            "参考文档：\n{}\n\n"
            "问题：{}\n\n"
            "请详细回答。".format(context, query)
        )
        return self._call_llm(prompt)