Coverage for mindsdb / integrations / handlers / llama_index_handler / settings.py: 0%

1from typing import List, Optional 

2from pydantic import BaseModel, field_validator, model_validator 

3from pydantic_settings import BaseSettings 

4 

5 

6class LlamaIndexConfig(BaseSettings): 

7 """ 

8 Model for LlamaIndexHandler settings. 

9 

10 Attributes: 

11 default_index_class (str): Default index class. 

12 supported_index_class (List[str]): Supported index classes. 

13 default_reader (str): Default reader. Note this is custom data frame reader. 

14 supported_reader (List[str]): Supported readers. 

15 """ 

16 DEFAULT_INDEX_CLASS: str = "VectorStoreIndex" 

17 SUPPORTED_INDEXES: List[str] = ["VectorStoreIndex"] 

18 DEFAULT_READER: str = "DFReader" 

19 SUPPORTED_READERS: List[str] = ["DFReader", "SimpleWebPageReader"] 

20 

21 

22llama_index_config = LlamaIndexConfig() 

23 

24 

25class LlamaIndexModel(BaseModel): 

26 """ 

27 Model for LlamaIndexHandler. 

28 

29 Attributes: 

30 reader (str): Reader. 

31 index_class (str): Index class. 

32 index (Any): Index. 

33 reader_params (Any): Reader parameters. 

34 index_params (Any): Index parameters. 

35 """ 

36 reader: Optional[str] = None 

37 index_class: Optional[str] = None 

38 input_column: str 

39 openai_api_key: Optional[str] = None 

40 input_column: Optional[str] 

41 mode: Optional[str] = None 

42 user_column: Optional[str] = None 

43 assistant_column: Optional[str] = None 

44 

45 @field_validator('reader') 

46 @classmethod 

47 def validate_reader(cls, value): 

48 if value not in llama_index_config.SUPPORTED_READERS: 

49 raise ValueError(f"Reader {value} is not supported.") 

50 

51 return value 

52 

53 @field_validator('index_class') 

54 @classmethod 

55 def validate_index_class(cls, value): 

56 if value not in llama_index_config.SUPPORTED_INDEXES: 

57 raise ValueError(f"Index class {value} is not supported.") 

58 

59 return value 

60 

61 @model_validator(mode='after') 

62 def validate_mode(self): 

63 if self.mode == "conversational" and not all([self.user_column, self.assistant_column]): 

64 raise ValueError("Conversational mode requires user_column and assistant_column parameter") 

65 

66 return self