from pydantic import BaseModel, Field

from .llm import GenerationConfig


class KGEnrichmentSettings(BaseModel):
    """Settings for knowledge graph enrichment."""

    max_knowledge_triples: int = Field(
        default=100,
        description="The maximum number of knowledge triples to extract from each chunk.",
    )

    generation_config: GenerationConfig = Field(
        default_factory=GenerationConfig,
        description="Configuration for text generation during graph enrichment.",
    )
    leiden_params: dict = Field(
        default_factory=dict,
        description="Parameters for the Leiden algorithm.",
    )
