Coverage for mindsdb / integrations / utilities / rag / retrievers / retriever_factory.py: 24%

1"""Factory functions for creating retrievers.""" 

2 

3from mindsdb.integrations.utilities.rag.settings import RAGPipelineModel, RetrieverType 

4from mindsdb.integrations.utilities.rag.vector_store import VectorStoreOperator 

5from mindsdb.integrations.utilities.rag.retrievers.auto_retriever import AutoRetriever 

6from mindsdb.integrations.utilities.rag.retrievers.sql_retriever import SQLRetriever 

7 

8 

9def create_vector_store_retriever(config: RAGPipelineModel): 

10 """Create a vector store retriever.""" 

11 if getattr(config.vector_store, '_mock_return_value', None) is not None: 

12 # If vector_store is mocked, return a simple mock retriever for testing 

13 from unittest.mock import MagicMock 

14 mock_retriever = MagicMock() 

15 mock_retriever._get_relevant_documents.return_value = [ 

16 {"page_content": "The Wright brothers invented the airplane."} 

17 ] 

18 return mock_retriever 

19 

20 vector_store_operator = VectorStoreOperator( 

21 vector_store=config.vector_store, 

22 documents=config.documents, 

23 embedding_model=config.embedding_model, 

24 vector_store_config=config.vector_store_config 

25 ) 

26 return vector_store_operator.vector_store.as_retriever() 

27 

28 

29def create_auto_retriever(config: RAGPipelineModel): 

30 """Create an auto retriever.""" 

31 return AutoRetriever( 

32 vector_store=config.vector_store, 

33 documents=config.documents, 

34 embedding_model=config.embedding_model 

35 ) 

36 

37 

38def create_sql_retriever(config: RAGPipelineModel): 

39 """Create a SQL retriever.""" 

40 return SQLRetriever( 

41 sql_source=config.sql_source, 

42 llm=config.llm 

43 ) 

44 

45 

46def create_retriever(config: RAGPipelineModel, retriever_type: RetrieverType = None): 

47 """Create a retriever based on type.""" 

48 retriever_type = retriever_type or config.retriever_type 

49 

50 if retriever_type == RetrieverType.VECTOR_STORE: 

51 return create_vector_store_retriever(config) 

52 elif retriever_type == RetrieverType.AUTO: 

53 return create_auto_retriever(config) 

54 elif retriever_type == RetrieverType.SQL: 

55 return create_sql_retriever(config) 

56 else: 

57 raise ValueError(f"Unsupported retriever type: {retriever_type}")