codekingpro/portable-devtools

1"""Prompt template that contains few shot examples."""2 3from __future__ import annotations4 5from typing import TYPE_CHECKING, Any, Literal6 7from pydantic import (8    BaseModel,9    ConfigDict,10    Field,11    model_validator,12)13from typing_extensions import override14 15from langchain_core._api import deprecated16from langchain_core.example_selectors import BaseExampleSelector17from langchain_core.messages import BaseMessage, get_buffer_string18from langchain_core.prompts.chat import BaseChatPromptTemplate19from langchain_core.prompts.message import BaseMessagePromptTemplate20from langchain_core.prompts.prompt import PromptTemplate21from langchain_core.prompts.string import (22    DEFAULT_FORMATTER_MAPPING,23    StringPromptTemplate,24    check_valid_template,25    get_template_variables,26)27 28if TYPE_CHECKING:29    from pathlib import Path30 31    from typing_extensions import Self32 33 34class _FewShotPromptTemplateMixin(BaseModel):35    """Prompt template that contains few shot examples."""36 37    examples: list[dict] | None = None38    """Examples to format into the prompt.39 40    Either this or `example_selector` should be provided.41    """42 43    example_selector: BaseExampleSelector | None = None44    """`ExampleSelector` to choose the examples to format into the prompt.45 46    Either this or `examples` should be provided.47    """48 49    model_config = ConfigDict(50        arbitrary_types_allowed=True,51        extra="forbid",52    )53 54    @model_validator(mode="before")55    @classmethod56    def check_examples_and_selector(cls, values: dict) -> Any:57        """Check that one and only one of `examples`/`example_selector` are provided.58 59        Args:60            values: The values to check.61 62        Returns:63            The values if they are valid.64 65        Raises:66            ValueError: If neither or both `examples` and `example_selector` are67                provided.68            ValueError: If both `examples` and `example_selector` are provided.69        """70        examples = values.get("examples")71        example_selector = values.get("example_selector")72        if examples and example_selector:73            msg = "Only one of 'examples' and 'example_selector' should be provided"74            raise ValueError(msg)75 76        if examples is None and example_selector is None:77            msg = "One of 'examples' and 'example_selector' should be provided"78            raise ValueError(msg)79 80        return values81 82    def _get_examples(self, **kwargs: Any) -> list[dict]:83        """Get the examples to use for formatting the prompt.84 85        Args:86            **kwargs: Keyword arguments to be passed to the example selector.87 88        Returns:89            List of examples.90 91        Raises:92            ValueError: If neither `examples` nor `example_selector` are provided.93        """94        if self.examples is not None:95            return self.examples96        if self.example_selector is not None:97            return self.example_selector.select_examples(kwargs)98        msg = "One of 'examples' and 'example_selector' should be provided"99        raise ValueError(msg)100 101    async def _aget_examples(self, **kwargs: Any) -> list[dict]:102        """Async get the examples to use for formatting the prompt.103 104        Args:105            **kwargs: Keyword arguments to be passed to the example selector.106 107        Returns:108            List of examples.109 110        Raises:111            ValueError: If neither `examples` nor `example_selector` are provided.112        """113        if self.examples is not None:114            return self.examples115        if self.example_selector is not None:116            return await self.example_selector.aselect_examples(kwargs)117        msg = "One of 'examples' and 'example_selector' should be provided"118        raise ValueError(msg)119 120 121class FewShotPromptTemplate(_FewShotPromptTemplateMixin, StringPromptTemplate):122    """Prompt template that contains few shot examples."""123 124    @classmethod125    def is_lc_serializable(cls) -> bool:126        """Return `False` as this class is not serializable."""127        return False128 129    validate_template: bool = False130    """Whether or not to try validating the template."""131 132    example_prompt: PromptTemplate133    """`PromptTemplate` used to format an individual example."""134 135    suffix: str136    """A prompt template string to put after the examples."""137 138    example_separator: str = "\n\n"139    """String separator used to join the prefix, the examples, and suffix."""140 141    prefix: str = ""142    """A prompt template string to put before the examples."""143 144    template_format: Literal["f-string", "jinja2"] = "f-string"145    """The format of the prompt template.146 147    Options are: `'f-string'`, `'jinja2'`.148    """149 150    def __init__(self, **kwargs: Any) -> None:151        """Initialize the few shot prompt template."""152        if "input_variables" not in kwargs and "example_prompt" in kwargs:153            kwargs["input_variables"] = kwargs["example_prompt"].input_variables154        super().__init__(**kwargs)155 156    @model_validator(mode="after")157    def template_is_valid(self) -> Self:158        """Check that prefix, suffix, and input variables are consistent."""159        if self.validate_template:160            check_valid_template(161                self.prefix + self.suffix,162                self.template_format,163                self.input_variables + list(self.partial_variables),164            )165        elif self.template_format:166            self.input_variables = [167                var168                for var in get_template_variables(169                    self.prefix + self.suffix, self.template_format170                )171                if var not in self.partial_variables172            ]173        return self174 175    model_config = ConfigDict(176        arbitrary_types_allowed=True,177        extra="forbid",178    )179 180    def format(self, **kwargs: Any) -> str:181        """Format the prompt with inputs generating a string.182 183        Use this method to generate a string representation of a prompt.184 185        Args:186            **kwargs: Keyword arguments to use for formatting.187 188        Returns:189            A string representation of the prompt.190        """191        kwargs = self._merge_partial_and_user_variables(**kwargs)192        # Get the examples to use.193        examples = self._get_examples(**kwargs)194        examples = [195            {k: e[k] for k in self.example_prompt.input_variables} for e in examples196        ]197        # Format the examples.198        example_strings = [199            self.example_prompt.format(**example) for example in examples200        ]201        # Create the overall template.202        pieces = [self.prefix, *example_strings, self.suffix]203        template = self.example_separator.join([piece for piece in pieces if piece])204 205        # Format the template with the input variables.206        return DEFAULT_FORMATTER_MAPPING[self.template_format](template, **kwargs)207 208    async def aformat(self, **kwargs: Any) -> str:209        """Async format the prompt with inputs generating a string.210 211        Use this method to generate a string representation of a prompt.212 213        Args:214            **kwargs: Keyword arguments to use for formatting.215 216        Returns:217            A string representation of the prompt.218        """219        kwargs = self._merge_partial_and_user_variables(**kwargs)220        # Get the examples to use.221        examples = await self._aget_examples(**kwargs)222        examples = [223            {k: e[k] for k in self.example_prompt.input_variables} for e in examples224        ]225        # Format the examples.226        example_strings = [227            await self.example_prompt.aformat(**example) for example in examples228        ]229        # Create the overall template.230        pieces = [self.prefix, *example_strings, self.suffix]231        template = self.example_separator.join([piece for piece in pieces if piece])232 233        # Format the template with the input variables.234        return DEFAULT_FORMATTER_MAPPING[self.template_format](template, **kwargs)235 236    @property237    def _prompt_type(self) -> str:238        """Return the prompt type key."""239        return "few_shot"240 241    @deprecated(242        since="1.2.21",243        removal="2.0.0",244        alternative="Use `dumpd`/`dumps` from `langchain_core.load` to serialize "245        "prompts and `load`/`loads` to deserialize them.",246    )247    def save(self, file_path: Path | str) -> None:248        """Save the prompt template to a file.249 250        Args:251            file_path: The path to save the prompt template to.252 253        Raises:254            ValueError: If `example_selector` is provided.255        """256        if self.example_selector:257            msg = "Saving an example selector is not currently supported"258            raise ValueError(msg)259        return super().save(file_path)260 261 262class FewShotChatMessagePromptTemplate(263    BaseChatPromptTemplate, _FewShotPromptTemplateMixin264):265    """Chat prompt template that supports few-shot examples.266 267    The high level structure of produced by this prompt template is a list of messages268    consisting of prefix message(s), example message(s), and suffix message(s).269 270    This structure enables creating a conversation with intermediate examples like:271 272    ```txt273    System: You are a helpful AI Assistant274 275    Human: What is 2+2?276 277    AI: 4278 279    Human: What is 2+3?280 281    AI: 5282 283    Human: What is 4+4?284    ```285 286    This prompt template can be used to generate a fixed list of examples or else to287    dynamically select examples based on the input.288 289    Examples:290        Prompt template with a fixed list of examples (matching the sample291        conversation above):292 293        ```python294        from langchain_core.prompts import (295            FewShotChatMessagePromptTemplate,296            ChatPromptTemplate,297        )298 299        examples = [300            {"input": "2+2", "output": "4"},301            {"input": "2+3", "output": "5"},302        ]303 304        example_prompt = ChatPromptTemplate.from_messages(305            [306                ("human", "What is {input}?"),307                ("ai", "{output}"),308            ]309        )310 311        few_shot_prompt = FewShotChatMessagePromptTemplate(312            examples=examples,313            # This is a prompt template used to format each individual example.314            example_prompt=example_prompt,315        )316 317        final_prompt = ChatPromptTemplate.from_messages(318            [319                ("system", "You are a helpful AI Assistant"),320                few_shot_prompt,321                ("human", "{input}"),322            ]323        )324        final_prompt.format(input="What is 4+4?")325        ```326 327        Prompt template with dynamically selected examples:328 329        ```python330        from langchain_core.prompts import SemanticSimilarityExampleSelector331        from langchain_core.embeddings import OpenAIEmbeddings332        from langchain_core.vectorstores import Chroma333 334        examples = [335            {"input": "2+2", "output": "4"},336            {"input": "2+3", "output": "5"},337            {"input": "2+4", "output": "6"},338            # ...339        ]340 341        to_vectorize = [" ".join(example.values()) for example in examples]342        embeddings = OpenAIEmbeddings()343        vectorstore = Chroma.from_texts(to_vectorize, embeddings, metadatas=examples)344        example_selector = SemanticSimilarityExampleSelector(vectorstore=vectorstore)345 346        from langchain_core import SystemMessage347        from langchain_core.prompts import HumanMessagePromptTemplate348        from langchain_core.prompts.few_shot import FewShotChatMessagePromptTemplate349 350        few_shot_prompt = FewShotChatMessagePromptTemplate(351            # Which variable(s) will be passed to the example selector.352            input_variables=["input"],353            example_selector=example_selector,354            # Define how each example will be formatted.355            # In this case, each example will become 2 messages:356            # 1 human, and 1 AI357            example_prompt=(358                HumanMessagePromptTemplate.from_template("{input}")359                + AIMessagePromptTemplate.from_template("{output}")360            ),361        )362        # Define the overall prompt.363        final_prompt = (364            SystemMessagePromptTemplate.from_template("You are a helpful AI Assistant")365            + few_shot_prompt366            + HumanMessagePromptTemplate.from_template("{input}")367        )368        # Show the prompt369        print(final_prompt.format_messages(input="What's 3+3?"))  # noqa: T201370 371        # Use within an LLM372        from langchain_core.chat_models import ChatAnthropic373 374        chain = final_prompt | ChatAnthropic(model="claude-3-haiku-20240307")375        chain.invoke({"input": "What's 3+3?"})376        ```377    """378 379    input_variables: list[str] = Field(default_factory=list)380    """A list of the names of the variables the prompt template will use to pass to381    the `example_selector`, if provided.382    """383 384    example_prompt: BaseMessagePromptTemplate | BaseChatPromptTemplate385    """The class to format each example."""386 387    @classmethod388    def is_lc_serializable(cls) -> bool:389        """Return `False` as this class is not serializable."""390        return False391 392    model_config = ConfigDict(393        arbitrary_types_allowed=True,394        extra="forbid",395    )396 397    def format_messages(self, **kwargs: Any) -> list[BaseMessage]:398        """Format kwargs into a list of messages.399 400        Args:401            **kwargs: Keyword arguments to use for filling in templates in messages.402 403        Returns:404            A list of formatted messages with all template variables filled in.405        """406        # Get the examples to use.407        examples = self._get_examples(**kwargs)408        examples = [409            {k: e[k] for k in self.example_prompt.input_variables} for e in examples410        ]411        # Format the examples.412        return [413            message414            for example in examples415            for message in self.example_prompt.format_messages(**example)416        ]417 418    async def aformat_messages(self, **kwargs: Any) -> list[BaseMessage]:419        """Async format kwargs into a list of messages.420 421        Args:422            **kwargs: Keyword arguments to use for filling in templates in messages.423 424        Returns:425            A list of formatted messages with all template variables filled in.426        """427        # Get the examples to use.428        examples = await self._aget_examples(**kwargs)429        examples = [430            {k: e[k] for k in self.example_prompt.input_variables} for e in examples431        ]432        # Format the examples.433        return [434            message435            for example in examples436            for message in await self.example_prompt.aformat_messages(**example)437        ]438 439    def format(self, **kwargs: Any) -> str:440        """Format the prompt with inputs generating a string.441 442        Use this method to generate a string representation of a prompt consisting of443        chat messages.444 445        Useful for feeding into a string-based completion language model or debugging.446 447        Args:448            **kwargs: Keyword arguments to use for formatting.449 450        Returns:451            A string representation of the prompt452        """453        messages = self.format_messages(**kwargs)454        return get_buffer_string(messages)455 456    async def aformat(self, **kwargs: Any) -> str:457        """Async format the prompt with inputs generating a string.458 459        Use this method to generate a string representation of a prompt consisting of460        chat messages.461 462        Useful for feeding into a string-based completion language model or debugging.463 464        Args:465            **kwargs: Keyword arguments to use for formatting.466 467        Returns:468            A string representation of the prompt469        """470        messages = await self.aformat_messages(**kwargs)471        return get_buffer_string(messages)472 473    @override474    def pretty_repr(self, html: bool = False) -> str:475        """Return a pretty representation of the prompt template.476 477        Args:478            html: Whether or not to return an HTML formatted string.479 480        Returns:481            A pretty representation of the prompt template.482        """483        raise NotImplementedError484