Coverage for mindsdb / integrations / handlers / chromadb_handler / settings.py: 76%

1import difflib 

2from typing import Any 

3 

4from pydantic import BaseModel, model_validator 

5 

6 

7class ChromaHandlerConfig(BaseModel): 

8 """ 

9 Configuration for VectorStoreHandler. 

10 """ 

11 

12 vector_store: str 

13 persist_directory: str = None 

14 host: str = None 

15 port: str = None 

16 password: str = None 

17 distance: str = 'cosine' 

18 

19 class Config: 

20 extra = "forbid" 

21 

22 @model_validator(mode="before") 

23 @classmethod 

24 def check_param_typos(cls, values: Any) -> Any: 

25 """Check if there are any typos in the parameters.""" 

26 

27 expected_params = cls.model_fields.keys() 

28 for key in values.keys(): 

29 if key not in expected_params: 29 ↛ 30line 29 didn't jump to line 30 because the condition on line 29 was never true

30 close_matches = difflib.get_close_matches( 

31 key, expected_params, cutoff=0.4 

32 ) 

33 if close_matches: 

34 raise ValueError( 

35 f"Unexpected parameter '{key}'. Did you mean '{close_matches[0]}'?" 

36 ) 

37 else: 

38 raise ValueError(f"Unexpected parameter '{key}'.") 

39 return values 

40 

41 @model_validator(mode="before") 

42 @classmethod 

43 def check_config(cls, values: Any) -> Any: 

44 """Check if config is valid.""" 

45 

46 vector_store = values.get("vector_store") 

47 host = values.get("host") 

48 port = values.get("port") 

49 persist_directory = values.get("persist_directory") 

50 

51 if bool(port) != bool(host) or (host and persist_directory): 51 ↛ 52line 51 didn't jump to line 52 because the condition on line 51 was never true

52 raise ValueError( 

53 f"For {vector_store} handler - host and port must be provided together. " 

54 f"Additionally, if host and port are provided, persist_directory should not be provided." 

55 ) 

56 

57 if persist_directory and (host or port): 57 ↛ 58line 57 didn't jump to line 58 because the condition on line 57 was never true

58 raise ValueError( 

59 f"For {vector_store} handler - if persistence_folder is provided, " 

60 f"host, port should not be provided." 

61 ) 

62 

63 return values