AddConfig.py 1.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. from typing import Callable, Optional
  2. from embedchain.config.BaseConfig import BaseConfig
  3. from embedchain.helper.json_serializable import register_deserializable
  4. @register_deserializable
  5. class ChunkerConfig(BaseConfig):
  6. """
  7. Config for the chunker used in `add` method
  8. """
  9. def __init__(
  10. self,
  11. chunk_size: Optional[int] = None,
  12. chunk_overlap: Optional[int] = None,
  13. length_function: Optional[Callable[[str], int]] = None,
  14. ):
  15. self.chunk_size = chunk_size if chunk_size else 2000
  16. self.chunk_overlap = chunk_overlap if chunk_overlap else 0
  17. self.length_function = length_function if length_function else len
  18. @register_deserializable
  19. class LoaderConfig(BaseConfig):
  20. """
  21. Config for the chunker used in `add` method
  22. """
  23. def __init__(self):
  24. pass
  25. @register_deserializable
  26. class AddConfig(BaseConfig):
  27. """
  28. Config for the `add` method.
  29. """
  30. def __init__(
  31. self,
  32. chunker: Optional[ChunkerConfig] = None,
  33. loader: Optional[LoaderConfig] = None,
  34. ):
  35. """
  36. Initializes a configuration class instance for the `add` method.
  37. :param chunker: Chunker config, defaults to None
  38. :type chunker: Optional[ChunkerConfig], optional
  39. :param loader: Loader config, defaults to None
  40. :type loader: Optional[LoaderConfig], optional
  41. """
  42. self.loader = loader
  43. self.chunker = chunker