Repository: tmc/langchaingo (go)

Issue:
feat: FewShotPromptTemplate
It's a clean way to have `FewShotPromptTemplate`, to implement the ConstitutionalChain, and also enable people to build prompts in that way.
[source](https://github.com/hwchase17/langchain/blob/master/langchain/prompts/few_shot.py)

I'd like to work on it.

Pull request description:
prompts: FewShotPrompt implementation
### Description
Added FewShotPrompt. Fixes #196.
[Python source](https://github.com/hwchase17/langchain/blob/master/langchain/prompts/few_shot.py)

### PR Checklist

- [x] Read the [Contributing documentation](https://github.com/tmc/langchaingo/blob/main/CONTRIBUTING.md).
- [x] Read the [Code of conduct documentation](https://github.com/tmc/langchaingo/blob/main/CODE_OF_CONDUCT.md).
- [x] Name your Pull Request title clearly, concisely, and prefixed with the name of the primarily affected package you changed according to [Good commit messages](https://go.dev/doc/contribute#commit_messages) (such as `memory: add interfaces for X, Y` or `util: add whizzbang helpers`).
- [x] Check that there isn't already a PR that solves the problem the same way to avoid creating a duplicate.
- [x] Provide a description in this PR that addresses **what** the PR is solving, or reference the issue that it solves (e.g. `Fixes #123`).
- [x] Describes the source of new concepts.
- [x] References existing implementations as appropriate.
- [x] Contains test coverage for new functions.
- [x] Passes all [`golangci-lint`](https://golangci-lint.run/) checks.


Patch:
diff --git a/prompts/few_shot.go b/prompts/few_shot.go
new file mode 100644
index 00000000..571129c1
--- /dev/null
+++ b/prompts/few_shot.go
@@ -0,0 +1,165 @@
+package prompts
+
+import (
+	"errors"
+	"fmt"
+	"strings"
+)
+
+var (
+	// ErrNoExample is returned when none of the Examples and ExampleSelector are provided.
+	ErrNoExample = errors.New("no example is provided")
+	// ErrExamplesAndExampleSelectorProvided is returned when there are no Examples and ExampleSelector.
+	ErrExamplesAndExampleSelectorProvided = errors.New("only one of 'Examples' and 'example_selector' should be" +
+		" provided")
+)
+
+// FewShotPrompt contains fields for a few-shot prompt.
+type FewShotPrompt struct {
+	// Examples to format into the prompt. Either this or ExamplePrompt should be provided.
+	Examples []map[string]string
+	// ExampleSelector to choose the examples to format into the prompt. Either this or Examples should be provided.
+	ExampleSelector ExampleSelector
+	// ExamplePrompt is used to format an individual example.
+	ExamplePrompt PromptTemplate
+	// A prompt template string to put before the examples.
+	Prefix string
+	// A prompt template string to put after the examples.
+	Suffix string
+	// A list of the names of the variables the prompt template expects.
+	InputVariables map[string]any
+	// Represents a map of variable names to values or functions that return values. If the value is a function, it will
+	// be called when the prompt template is rendered.
+	PartialVariables map[string]any
+	// String separator used to join the prefix, the examples, and suffix.
+	ExampleSeparator string
+	// The format of the prompt template. Options are: 'f-string', 'jinja2'.
+	TemplateFormat TemplateFormat
+	// Whether to try validating the template.
+	ValidateTemplate bool
+}
+
+// NewFewShotPrompt creates a new few-shot prompt with the given input. It returns error if there is no example, both
+// examples and exampleSelector are provided, or CheckValidTemplate returns err when ValidateTemplate is true.
+func NewFewShotPrompt(examplePrompt PromptTemplate, examples []map[string]string, exampleSelector ExampleSelector,
+	prefix string, suffix string, input map[string]interface{}, partialInput map[string]interface{},
+	exampleSeparator string, templateFormat TemplateFormat, validateTemplate bool,
+) (*FewShotPrompt, error) {
+	err := validateExamples(examples, exampleSelector)
+	if err != nil {
+		return nil, err
+	}
+	prompt := &FewShotPrompt{
+		ExamplePrompt:    examplePrompt,
+		Prefix:           prefix,
+		Suffix:           suffix,
+		InputVariables:   input,
+		PartialVariables: partialInput,
+		Examples:         examples,
+		ExampleSelector:  exampleSelector,
+		ExampleSeparator: "\n\n",
+		TemplateFormat:   templateFormat,
+		ValidateTemplate: validateTemplate,
+	}
+	if exampleSeparator != "" {
+		prompt.ExampleSeparator = exampleSeparator
+	}
+
+	if prompt.ValidateTemplate {
+		err := CheckValidTemplate(prompt.Prefix+prompt.Suffix, prompt.TemplateFormat, append(getMapKeys(input),
+			getMapKeys(partialInput)...))
+		if err != nil {
+			return nil, err
+		}
+	}
+	return prompt, nil
+}
+
+// validateExamples validates the provided example and exampleSelector. One of them must be provided only.
+func validateExamples(examples []map[string]string, exampleSelector ExampleSelector) error {
+	if examples != nil && exampleSelector != nil {
+		return ErrExamplesAndExampleSelectorProvided
+	} else if examples == nil && exampleSelector == nil {
+		return ErrNoExample
+	}
+	return nil
+}
+
+// getExamples returns the provided examples or returns error when there is no example.
+func (p *FewShotPrompt) getExamples(input map[string]string) ([]map[string]string, error) {
+	switch {
+	case p.Examples != nil:
+		return p.Examples, nil
+	case p.ExampleSelector != nil:
+		return p.ExampleSelector.SelectExamples(input), nil
+	default:
+		return nil, ErrNoExample
+	}
+}
+
+// Format assembles and formats the pieces of the prompt with the given input values and partial values.
+func (p *FewShotPrompt) Format(values map[string]interface{}) (string, error) {
+	resolvedValues, err := resolvePartialValues(p.PartialVariables, values)
+	if err != nil {
+		return "", err
+	}
+	stringResolvedValues := map[string]string{}
+	for k, v := range resolvedValues {
+		strVal, ok := v.(string)
+		if !ok {
+			return "", fmt.Errorf("%w: %T", ErrInvalidPartialVariableType, v)
+		}
+		stringResolvedValues[k] = strVal
+	}
+	examples, err := p.getExamples(stringResolvedValues)
+	if err != nil {
+		return "", err
+	}
+	exampleStrings := make([]string, len(examples))
+
+	for i, example := range examples {
+		exampleMap := make(map[string]interface{})
+		for k, v := range example {
+			exampleMap[k] = v
+		}
+
+		res, err := p.ExamplePrompt.Format(exampleMap)
+		if err != nil {
+			return "", err
+		}
+		exampleStrings[i] = res
+	}
+
+	template := assemblePieces(p.Prefix, p.Suffix, exampleStrings, p.ExampleSeparator)
+	return defaultformatterMapping[p.TemplateFormat](template, resolvedValues)
+}
+
+// assemblePieces assembles the pieces of the few-shot prompt.
+func assemblePieces(prefix, suffix string, exampleStrings []string, separator string) string {
+	const additionalCapacity = 2
+	pieces := make([]string, 0, len(exampleStrings)+additionalCapacity)
+	if prefix != "" {
+		pieces = append(pieces, prefix)
+	}
+
+	for _, elem := range exampleStrings {
+		if elem != "" {
+			pieces = append(pieces, elem)
+		}
+	}
+
+	if suffix != "" {
+		pieces = append(pieces, suffix)
+	}
+
+	return strings.Join(pieces, separator)
+}
+
+// getMapKeys returns the keys of the provided map.
+func getMapKeys(inputMap map[string]any) []string {
+	keys := make([]string, 0, len(inputMap))
+	for k := range inputMap {
+		keys = append(keys, k)
+	}
+	return keys
+}


Test patch:
diff --git a/prompts/few_shot_test.go b/prompts/few_shot_test.go
new file mode 100644
index 00000000..8f2e5a1a
--- /dev/null
+++ b/prompts/few_shot_test.go
@@ -0,0 +1,156 @@
+package prompts
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+)
+
+// nolint: funlen
+func TestFewShotPrompt_Format(t *testing.T) {
+	examplePrompt := NewPromptTemplate("{{.question}}: {{.answer}}", []string{"question", "answer"})
+	t.Parallel()
+	testCases := []struct {
+		name             string
+		examplePrompt    PromptTemplate
+		examples         []map[string]string
+		prefix           string
+		suffix           string
+		input            map[string]interface{}
+		partialInput     map[string]interface{}
+		exampleSeparator string
+		templateFormat   TemplateFormat
+		validateTemplate bool
+		wantErr          bool
+		expected         string
+	}{
+		{
+			"prefix only", examplePrompt,
+			[]map[string]string{},
+			"This is a {{.foo}} test.", "",
+			map[string]interface{}{"foo": "bar"},
+			nil,
+			"",
+			TemplateFormatGoTemplate,
+			true,
+			false,
+			"This is a bar test.",
+		},
+		{
+			"suffix only", examplePrompt,
+			[]map[string]string{},
+			"", "This is a {{.foo}} test.",
+			map[string]interface{}{"foo": "bar"},
+			nil,
+			"",
+			TemplateFormatGoTemplate,
+			true,
+			false,
+			"This is a bar test.",
+		},
+		{
+			"insufficient InputVariables w err",
+			examplePrompt,
+			[]map[string]string{},
+			"",
+			"This is a {{.foo}} test.",
+			map[string]interface{}{"bar": "bar"},
+			nil,
+			"",
+			TemplateFormatGoTemplate,
+			true,
+			true,
+			`template: template:1:12: executing "template" at <.foo>: map has no entry for key "foo"`,
+		},
+		{
+			"inputVariables neither Examples nor ExampleSelector w err",
+			examplePrompt,
+			nil,
+			"",
+			"",
+			map[string]interface{}{"bar": "bar"},
+			nil,
+			"",
+			TemplateFormatGoTemplate,
+			true,
+			true,
+			ErrNoExample.Error(),
+		},
+		{
+			"functionality test",
+			examplePrompt,
+			[]map[string]string{{"question": "foo", "answer": "bar"}, {"question": "baz", "answer": "foo"}},
+			"This is a test about {{.content}}.",
+			"Now you try to talk about {{.new_content}}.",
+			map[string]interface{}{"content": "animals", "new_content": "party"},
+			nil,
+			"\n",
+			TemplateFormatGoTemplate,
+			true,
+			false,
+			"This is a test about animals.\nfoo: bar\nbaz: foo\nNow you try to talk about party.",
+		},
+		{
+			"functionality test with partial input",
+			examplePrompt,
+			[]map[string]string{{"question": "foo", "answer": "bar"}, {"question": "baz", "answer": "foo"}},
+			"This is a test about {{.content}}.",
+			"Now you try to talk about {{.new_content}}.",
+			map[string]interface{}{"content": "animals"},
+			map[string]interface{}{"new_content": func() string { return "party" }},
+			"\n",
+			TemplateFormatGoTemplate,
+			true,
+			false,
+			"This is a test about animals.\nfoo: bar\nbaz: foo\nNow you try to talk about party.",
+		},
+		{
+			"invalid template w err",
+			examplePrompt,
+			[]map[string]string{{"question": "foo", "answer": "bar"}, {"question": "baz", "answer": "foo"}},
+			"This is a test about {{.wrong_content}}.",
+			"Now you try to talk about {{.new_content}}.",
+			map[string]interface{}{"content": "animals"},
+			map[string]interface{}{"new_content": func() string { return "party" }},
+			"\n",
+			TemplateFormatGoTemplate,
+			true,
+			true,
+			"template: template:1:23: executing \"template\" at <.wrong_content>: map has no entry for key " +
+				"\"wrong_content\"",
+		},
+	}
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			t.Helper()
+			p, err := NewFewShotPrompt(tc.examplePrompt, tc.examples, nil, tc.prefix, tc.suffix, tc.input,
+				tc.partialInput, tc.exampleSeparator, tc.templateFormat, tc.validateTemplate)
+			if tc.wantErr {
+				checkError(t, err, tc.expected)
+				return
+			}
+			fp, err := p.Format(tc.input)
+			if checkError(t, err, "") {
+				return
+			}
+			got := fmt.Sprint(fp)
+			if diff := cmp.Diff(tc.expected, got); diff != "" {
+				t.Errorf("unexpected prompt output (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
+
+func checkError(t *testing.T, err error, expected string) bool {
+	t.Helper()
+	if err != nil {
+		if expected != "" && err.Error() != expected {
+			t.Errorf("unexpected error: got %q, want %q", err.Error(), expected)
+		}
+		return true
+	}
+	return false
+}