Coverage for src/dataknobs_data/backends/vector_config_mixin.py: 73%

1"""Shared mixin for vector configuration across all backends.""" 

2 

3import logging 

4from typing import Any 

5 

6from ..vector.types import DistanceMetric 

7 

8logger = logging.getLogger(__name__) 

9 

10 

11class VectorConfigMixin: 

12 """Mixin to provide consistent vector configuration across all backends. 

13  

14 This mixin ensures that all backends handle vector-related configuration 

15 parameters in a consistent way, including: 

16 - vector_enabled: Whether vector support is enabled 

17 - vector_metric: The distance metric to use for vector operations 

18  

19 Usage: 

20 Include this mixin in your backend class and call _parse_vector_config() 

21 during initialization to extract vector configuration parameters. 

22 """ 

23 

24 def _parse_vector_config(self, config: dict[str, Any] | None = None) -> None: 

25 """Parse vector-related configuration parameters. 

26  

27 Args: 

28 config: Configuration dictionary (uses self.config if not provided) 

29 """ 

30 # Use provided config or fall back to self.config 

31 config = config if config is not None else getattr(self, 'config', {}) 

32 

33 # Extract vector configuration parameters 

34 self._vector_enabled = config.get("vector_enabled", False) 

35 

36 # Parse distance metric 

37 metric = config.get("vector_metric", "cosine") 

38 if isinstance(metric, str): 

39 try: 

40 self.vector_metric = DistanceMetric(metric.lower()) 

41 except ValueError: 

42 logger.warning(f"Invalid vector metric '{metric}', using cosine") 

43 self.vector_metric = DistanceMetric.COSINE 

44 elif isinstance(metric, DistanceMetric): 

45 self.vector_metric = metric 

46 else: 

47 self.vector_metric = DistanceMetric.COSINE 

48 

49 # Log vector configuration 

50 if self.vector_enabled: 

51 logger.debug( 

52 f"Vector support enabled with metric: {self.vector_metric.value}" 

53 ) 

54 

55 def _init_vector_state(self) -> None: 

56 """Initialize vector-related state variables. 

57  

58 This method should be called after _parse_vector_config() to set up 

59 any additional state needed for vector operations. 

60 """ 

61 # Track vector field dimensions (using underscore for consistency with mixins) 

62 self._vector_dimensions: dict[str, int] = {} 

63 self._vector_fields: dict[str, Any] = {} 

64 

65 # For backends that support native vector operations 

66 self._has_native_vector_support = False 

67 

68 @property 

69 def vector_enabled(self) -> bool: 

70 """Check if vector support is enabled. 

71  

72 Returns: 

73 True if vector operations are enabled for this backend 

74 """ 

75 return getattr(self, '_vector_enabled', False) 

76 

77 @vector_enabled.setter 

78 def vector_enabled(self, value: bool) -> None: 

79 """Set vector support enabled state. 

80  

81 Args: 

82 value: Whether to enable vector support 

83 """ 

84 self._vector_enabled = value