BidiRag/bdirag/rag_methods/step_back_rag.py

# -*- coding: utf-8 -*-
"""Step-Back RAG - ask a more general question for context, then specific."""
from .base import BaseRAG


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

    def _generate_step_back_query(self, query):
        prompt = (
            "针对以下具体问题，提出一个更宏观、更通用的背景问题。\n\n"
            "具体问题：{}\n\n"
            "宏观问题：".format(query)
        )
        return self._call_llm(prompt)

    def retrieve(self, query, k=10):
        step_back_query = self._generate_step_back_query(query)
        
        q1_embedding = self.embedding_model.embed_query(query)
        specific_results = self.vector_store.similarity_search(q1_embedding, k // 2)
        
        q2_embedding = self.embedding_model.embed_query(step_back_query)
        general_results = self.vector_store.similarity_search(q2_embedding, k // 2)
        
        combined = specific_results + general_results
        return self._deduplicate_results(combined, k)

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