diff --git a/application/retriever/retriever_creator.py b/application/retriever/retriever_creator.py index 07be373d..26cb41ca 100644 --- a/application/retriever/retriever_creator.py +++ b/application/retriever/retriever_creator.py @@ -3,18 +3,18 @@ from application.retriever.duckduck_search import DuckDuckSearch from application.retriever.brave_search import BraveRetSearch - class RetrieverCreator: retrievers = { - 'classic': ClassicRAG, - 'duckduck_search': DuckDuckSearch, - 'brave_search': BraveRetSearch, - 'default': ClassicRAG + "classic": ClassicRAG, + "duckduck_search": DuckDuckSearch, + "brave_search": BraveRetSearch, + "default": ClassicRAG, } @classmethod def create_retriever(cls, type, *args, **kwargs): - retiever_class = cls.retrievers.get(type.lower()) + retriever_type = (type or "default").lower() + retiever_class = cls.retrievers.get(retriever_type) if not retiever_class: raise ValueError(f"No retievers class found for type {type}") - return retiever_class(*args, **kwargs) \ No newline at end of file + return retiever_class(*args, **kwargs)