Repository: tmc/langchaingo
Pull request #293 (2023-09-18 06:17:06)

add context prompt in vertexai chat
vertexai chat-bison parameters has an optional context prompt. added system message to be the context in the request

in langchain > vertexai.py SystemMessage type is being set as context prompt:
![image](https://github.com/tmc/langchaingo/assets/82501805/aeb67d41-17e6-4fa0-b5f6-0058633543e1)

### 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/chains/options.go b/chains/options.go
index 11e9e6ec..28756ddb 100644
--- a/chains/options.go
+++ b/chains/options.go
@@ -132,6 +132,7 @@ func getLLMCallOptions(options ...ChainCallOption) []llms.CallOption {
 		llms.WithStopWords(opts.StopWords),
 		llms.WithStreamingFunc(opts.StreamingFunc),
 		llms.WithTopK(opts.TopK),
+		llms.WithTopP(opts.TopP),
 		llms.WithSeed(opts.Seed),
 		llms.WithMinLength(opts.MinLength),
 		llms.WithMaxLength(opts.MaxLength),
diff --git a/chains/stuff_documents.go b/chains/stuff_documents.go
index 0996d377..9df7abda 100644
--- a/chains/stuff_documents.go
+++ b/chains/stuff_documents.go
@@ -51,28 +51,25 @@ func NewStuffDocuments(llmChain *LLMChain) StuffDocuments {
 }
 
 // Call handles the inner logic of the StuffDocuments chain.
-func (c StuffDocuments) Call(ctx context.Context, values map[string]any, options ...ChainCallOption) (map[string]any, error) { //nolint: lll
+func (c StuffDocuments) Call(
+	ctx context.Context, values map[string]any, options ...ChainCallOption,
+) (map[string]any, error) {
 	docs, ok := values[c.InputKey].([]schema.Document)
 	if !ok {
 		return nil, fmt.Errorf("%w: %w", ErrInvalidInputValues, ErrInputValuesWrongType)
 	}
 
-	var text string
-	for _, doc := range docs {
-		text += doc.PageContent + c.Separator
-	}
-
 	inputValues := make(map[string]any)
 	for key, value := range values {
 		inputValues[key] = value
 	}
 
-	inputValues[c.DocumentVariableName] = text
+	inputValues[c.DocumentVariableName] = c.joinDocuments(docs)
 	return Call(ctx, c.LLMChain, inputValues, options...)
 }
 
 // GetMemory returns a simple memory.
-func (c StuffDocuments) GetMemory() schema.Memory { //nolint:ireturn
+func (c StuffDocuments) GetMemory() schema.Memory {
 	return memory.NewSimple()
 }
 
@@ -85,3 +82,16 @@ func (c StuffDocuments) GetInputKeys() []string {
 func (c StuffDocuments) GetOutputKeys() []string {
 	return append([]string{}, c.LLMChain.GetOutputKeys()...)
 }
+
+// joinDocuments joins the documents with the separator.
+func (c StuffDocuments) joinDocuments(docs []schema.Document) string {
+	var text string
+	docLen := len(docs)
+	for k, doc := range docs {
+		text += doc.PageContent
+		if k != docLen-1 {
+			text += c.Separator
+		}
+	}
+	return text
+}
diff --git a/llms/vertexai/vertexai_palm_llm_chat.go b/llms/vertexai/vertexai_palm_llm_chat.go
index 9504e575..1741a115 100644
--- a/llms/vertexai/vertexai_palm_llm_chat.go
+++ b/llms/vertexai/vertexai_palm_llm_chat.go
@@ -48,10 +48,11 @@ func (o *Chat) Generate(ctx context.Context, messageSets [][]schema.ChatMessage,
 
 	generations := make([]*llms.Generation, 0, len(messageSets))
 	for _, messages := range messageSets {
-		msgs := toClientChatMessage(messages)
+		msgs, chatContext := toClientChatMessage(messages)
 		result, err := o.client.CreateChat(ctx, &vertexaiclient.ChatRequest{
 			Temperature: opts.Temperature,
 			Messages:    msgs,
+			Context:     chatContext,
 		})
 		if err != nil {
 			return nil, err
@@ -78,13 +79,20 @@ func (o *Chat) GetNumTokens(text string) int {
 	return llms.CountTokens(vertexaiclient.TextModelName, text)
 }
 
-func toClientChatMessage(messages []schema.ChatMessage) []*vertexaiclient.ChatMessage {
+func toClientChatMessage(messages []schema.ChatMessage) ([]*vertexaiclient.ChatMessage, string) {
 	msgs := make([]*vertexaiclient.ChatMessage, len(messages))
+	chatContext := ""
 	for i, m := range messages {
 		msg := &vertexaiclient.ChatMessage{
 			Content: m.GetContent(),
 		}
 		typ := m.GetType()
+		// system message aka context/prompt will not be included in the message list
+		if i == 0 && typ == schema.ChatMessageTypeSystem {
+			chatContext = msg.GetContent()
+			continue
+		}
+
 		switch typ {
 		case schema.ChatMessageTypeSystem:
 			msg.Author = botAuthor
@@ -102,7 +110,7 @@ func toClientChatMessage(messages []schema.ChatMessage) []*vertexaiclient.ChatMe
 		}
 		msgs[i] = msg
 	}
-	return msgs
+	return msgs, chatContext
 }
 
 // NewChat returns a new VertexAI PaLM Chat LLM.
diff --git a/llms/vertexai/vertexai_palm_llm_option.go b/llms/vertexai/vertexai_palm_llm_option.go
index a142800d..b655e4f1 100644
--- a/llms/vertexai/vertexai_palm_llm_option.go
+++ b/llms/vertexai/vertexai_palm_llm_option.go
@@ -1,10 +1,12 @@
 package vertexai
 
 import (
+	"net/http"
 	"os"
 	"sync"
 
 	"google.golang.org/api/option"
+	"google.golang.org/grpc"
 )
 
 const (
@@ -62,6 +64,18 @@ func WithCredentialsJSON(json []byte) Option {
 	return convertByteArrayOption(option.WithCredentialsJSON)(json)
 }
 
+func WithGRPCDialOption(opt grpc.DialOption) Option {
+	return func(opts *options) {
+		opts.clientOptions = append(opts.clientOptions, option.WithGRPCDialOption(opt))
+	}
+}
+
+func WithHTTPClient(client *http.Client) Option {
+	return func(opts *options) {
+		opts.clientOptions = append(opts.clientOptions, option.WithHTTPClient(client))
+	}
+}
+
 func convertStringOption(fopt func(string) option.ClientOption) func(string) Option {
 	return func(param string) Option {
 		return func(opts *options) {
diff --git a/prompts/templates.go b/prompts/templates.go
index 11ea24b3..523154ff 100644
--- a/prompts/templates.go
+++ b/prompts/templates.go
@@ -35,7 +35,7 @@ var defaultformatterMapping = map[TemplateFormat]interpolator{ //nolint:gocheckn
 func interpolateGoTemplate(tmpl string, values map[string]any) (string, error) {
 	parsedTmpl, err := template.New("template").
 		Option("missingkey=error").
-		Funcs(sprig.FuncMap()).
+		Funcs(sprig.TxtFuncMap()).
 		Parse(tmpl)
 	if err != nil {
 		return "", err


Test patch:
diff --git a/chains/stuff_documents_test.go b/chains/stuff_documents_test.go
index effb6ec6..25e1a75e 100644
--- a/chains/stuff_documents_test.go
+++ b/chains/stuff_documents_test.go
@@ -44,3 +44,53 @@ func TestStuffDocuments(t *testing.T) {
 		require.True(t, ok)
 	}
 }
+
+func TestStuffDocuments_joinDocs(t *testing.T) {
+	t.Parallel()
+
+	testcases := []struct {
+		name string
+		docs []schema.Document
+		want string
+	}{
+		{
+			name: "empty",
+			docs: []schema.Document{},
+			want: "",
+		},
+		{
+			name: "single",
+			docs: []schema.Document{
+				{PageContent: "foo"},
+			},
+			want: "foo",
+		},
+		{
+			name: "multiple",
+			docs: []schema.Document{
+				{PageContent: "foo"},
+				{PageContent: "bar"},
+			},
+			want: "foo\n\nbar",
+		},
+		{
+			name: "multiple with separator",
+			docs: []schema.Document{
+				{PageContent: "foo"},
+				{PageContent: "bar\n\n"},
+			},
+			want: "foo\n\nbar\n\n",
+		},
+	}
+
+	chain := NewStuffDocuments(&LLMChain{})
+
+	for _, tc := range testcases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			got := chain.joinDocuments(tc.docs)
+			require.Equal(t, tc.want, got)
+		})
+	}
+}