from .retriever import CustomRetriever
from .generator import CustomGenerator


class MyRagModel:
    def __init__(self, retriever: CustomRetriever, generator: CustomGenerator):
        self.retriever = retriever
        self.generator = generator

    def get_answer(self, query):
        context = self.retriever.retrieve(query)
        response = self.generator.generate(query, context)
        return response