Skip to content

Commit

Permalink
Merge pull request #549 from tmc/llms-add-helpers
Browse files Browse the repository at this point in the history
llms: Add convenience helper for text parts
  • Loading branch information
tmc authored Jan 26, 2024
2 parents 59a4f99 + b91b7f6 commit 43478a1
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 0 deletions.
15 changes: 15 additions & 0 deletions llms/generatecontent.go
Original file line number Diff line number Diff line change
Expand Up @@ -82,3 +82,18 @@ type ContentChoice struct {
// FuncCall is non-nil when the model asks to invoke a function/tool.
FuncCall *schema.FunctionCall
}

// TextParts is a helper function to create a MessageContent with a role and a
// list of text parts.
func TextParts(role schema.ChatMessageType, parts ...string) MessageContent {
result := MessageContent{
Role: role,
Parts: []ContentPart{},
}
for _, part := range parts {
result.Parts = append(result.Parts, TextContent{
Text: part,
})
}
return result
}
39 changes: 39 additions & 0 deletions llms/generatecontent_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
package llms

import (
"reflect"
"testing"

"github.com/tmc/langchaingo/schema"
)

func TestTextParts(t *testing.T) {
t.Parallel()
type args struct {
role schema.ChatMessageType
parts []string
}
tests := []struct {
name string
args args
want MessageContent
}{
{"basics", args{schema.ChatMessageTypeHuman, []string{"a", "b", "c"}}, MessageContent{
Role: schema.ChatMessageTypeHuman,
Parts: []ContentPart{
TextContent{Text: "a"},
TextContent{Text: "b"},
TextContent{Text: "c"},
},
}},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
if got := TextParts(tt.args.role, tt.args.parts...); !reflect.DeepEqual(got, tt.want) {
t.Errorf("TextParts() = %v, want %v", got, tt.want)
}
})
}
}

0 comments on commit 43478a1

Please sign in to comment.