AddConfig.py 1.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  1. from typing import Callable, Optional
  2. from embedchain.config.BaseConfig import BaseConfig
  3. from embedchain.helper_classes.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. self.loader = loader
  36. self.chunker = chunker