/cognee/infrastructure/data/chunking/config.py

from functools import lru_cache
from typing import Any

from pydantic_settings import BaseSettings, SettingsConfigDict

from cognee.infrastructure.data.chunking.DefaultChunkEngine import DefaultChunkEngine
from cognee.shared.data_models import ChunkEngine, ChunkStrategy


class ChunkConfig(BaseSettings):
    """
    Manage configuration settings for chunk processing.
    """

    chunk_size: int = 1500
    chunk_overlap: int = 10
    chunk_strategy: ChunkStrategy = ChunkStrategy.PARAGRAPH
    chunk_engine: ChunkEngine = ChunkEngine.DEFAULT_ENGINE

    model_config = SettingsConfigDict(env_file=".env", extra="allow")

    def to_dict(self) -> dict[str, Any]:
        """
        Convert the chunk settings to a dictionary format.

        Returns:
        --------

            - dict: A dictionary representation of the chunk configuration settings.
        """
        return {
            "chunk_size": self.chunk_size,
            "chunk_overlap": self.chunk_overlap,
            "chunk_strategy": self.chunk_strategy,
            "chunk_engine": self.chunk_engine,
        }


@lru_cache
def get_chunk_config() -> ChunkConfig:
    """
    Retrieve the configuration for chunking data, caching the result for efficiency.

    This function creates an instance of the ChunkConfig class, which contains settings such
    as chunk size, overlap, strategy, and engine. The use of lru_cache ensures that
    subsequent calls to this function will return the cached instance, improving performance
    by avoiding re-creation of the object.

    Returns:
    --------

        - ChunkConfig: An instance of the ChunkConfig class containing the chunking
          configuration settings.
    """
    return ChunkConfig()