langchain/docs/examples/prompts/custom_example_selector.ipynb

Custom ExampleSelector

This notebook goes over how to implement a custom ExampleSelector. ExampleSelectors are used to select examples to use in few shot prompts.

An ExampleSelector must implement two methods:

1. An add_example method which takes in an example and adds it into the ExampleSelector
2. A select_examples method which takes in input variables (which are meant to be user input) and returns a list of examples to use in the few shot prompt.

Let's implement a custom ExampleSelector that just selects two examples at random.

In [1]:

from langchain.prompts.example_selector.base import BaseExampleSelector
from typing import Dict, List
import numpy as np

In [2]:

class CustomExampleSelector(BaseExampleSelector):
    
    def __init__(self, examples: List[Dict[str, str]]):
        self.examples = examples
    
    def add_example(self, example: Dict[str, str]) -> None:
        """Add new example to store for a key."""
        self.examples.append(example)

    def select_examples(self, input_variables: Dict[str, str]) -> List[dict]:
        """Select which examples to use based on the inputs."""
        return np.random.choice(self.examples, size=2, replace=False)

In [3]:

examples = [{"foo": "1"}, {"foo": "2"}, {"foo": "3"}]
example_selector = CustomExampleSelector(examples)

Let's now try it out! We can select some examples and try adding examples.

In [4]:

example_selector.select_examples({"foo": "foo"})

Out [4]:

array([{'foo': '2'}, {'foo': '3'}], dtype=object)

In [5]:

example_selector.add_example({"foo": "4"})

In [6]:

example_selector.examples

Out [6]:

[{'foo': '1'}, {'foo': '2'}, {'foo': '3'}, {'foo': '4'}]

In [7]:

example_selector.select_examples({"foo": "foo"})

Out [7]:

array([{'foo': '1'}, {'foo': '4'}], dtype=object)

In [ ]: