BidiRag/bdirag/rag_methods/naive_rag.py

# -*- coding: utf-8 -*-
"""Naive RAG - basic semantic retrieval."""
from .base import BaseRAG


class NaiveRAG(BaseRAG):
    def __init__(self, retrieval_prompt_template=None, **kwargs):
        super().__init__(**kwargs)
        self.retrieval_prompt_template = retrieval_prompt_template or (
            "根据以下参考文档，回答问题。\n\n"
            "参考文档：\n{context}\n\n"
            "问题：{query}\n\n"
            "请详细回答，如果参考文档中没有相关信息，请说明无法从文档中找到答案。"
        )

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

    def generate(self, query, context):
        prompt = self.retrieval_prompt_template.format(context=context, query=query)
        return self._call_llm(prompt)