1"""Select examples based on length."""23import re4from collections.abc import Callable5from typing import Any67from pydantic import BaseModel, Field, model_validator8from typing_extensions import Self910from langchain_core.example_selectors.base import BaseExampleSelector11from langchain_core.prompts.prompt import PromptTemplate121314def _get_length_based(text: str) -> int:15    return len(re.split(r"\n| ", text))161718class LengthBasedExampleSelector(BaseExampleSelector, BaseModel):19    r"""Select examples based on length.2021    Example:22        ```python23        from langchain_core.example_selectors import LengthBasedExampleSelector24        from langchain_core.prompts import PromptTemplate2526        # Define examples27        examples = [28            {"input": "happy", "output": "sad"},29            {"input": "tall", "output": "short"},30            {"input": "fast", "output": "slow"},31        ]3233        # Create prompt template34        example_prompt = PromptTemplate(35            input_variables=["input", "output"],36            template="Input: {input}\nOutput: {output}",37        )3839        # Create selector with max length constraint40        selector = LengthBasedExampleSelector(41            examples=examples,42            example_prompt=example_prompt,43            max_length=50,  # Maximum prompt length44        )4546        # Select examples for a new input47        selected = selector.select_examples({"input": "large", "output": "tiny"})48        # Returns examples that fit within max_length constraint49        ```50    """5152    examples: list[dict[str, Any]]53    """A list of the examples that the prompt template expects."""5455    example_prompt: PromptTemplate56    """Prompt template used to format the examples."""5758    get_text_length: Callable[[str], int] = _get_length_based59    """Function to measure prompt length. Defaults to word count."""6061    max_length: int = 204862    """Max length for the prompt, beyond which examples are cut."""6364    example_text_lengths: list[int] = Field(default_factory=list)65    """Length of each example."""6667    def add_example(self, example: dict[str, str]) -> None:68        """Add new example to list.6970        Args:71            example: A dictionary with keys as input variables72                and values as their values.73        """74        self.examples.append(example)75        string_example = self.example_prompt.format(**example)76        self.example_text_lengths.append(self.get_text_length(string_example))7778    async def aadd_example(self, example: dict[str, str]) -> None:79        """Async add new example to list.8081        Args:82            example: A dictionary with keys as input variables83                and values as their values.84        """85        self.add_example(example)8687    @model_validator(mode="after")88    def post_init(self) -> Self:89        """Validate that the examples are formatted correctly."""90        if self.example_text_lengths:91            return self92        string_examples = [self.example_prompt.format(**eg) for eg in self.examples]93        self.example_text_lengths = [self.get_text_length(eg) for eg in string_examples]94        return self9596    def select_examples(self, input_variables: dict[str, str]) -> list[dict[str, Any]]:97        """Select which examples to use based on the input lengths.9899        Args:100            input_variables: A dictionary with keys as input variables101                and values as their values.102103        Returns:104            A list of examples to include in the prompt.105        """106        inputs = " ".join(input_variables.values())107        remaining_length = self.max_length - self.get_text_length(inputs)108        i = 0109        examples = []110        while remaining_length > 0 and i < len(self.examples):111            new_length = remaining_length - self.example_text_lengths[i]112            if new_length < 0:113                break114            examples.append(self.examples[i])115            remaining_length = new_length116            i += 1117        return examples118119    async def aselect_examples(120        self, input_variables: dict[str, str]121    ) -> list[dict[str, Any]]:122        """Async select which examples to use based on the input lengths.123124        Args:125            input_variables: A dictionary with keys as input variables126                and values as their values.127128        Returns:129            A list of examples to include in the prompt.130        """131        return self.select_examples(input_variables)