parser

pkg/main.go → examples/normal/main.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"golangchain/pkg/lib"
 	"golangchain/pkg/openai"
+	"golangchain/pkg/parser"
 	"golangchain/pkg/prompt"
 )
 
@@ -13,9 +14,10 @@ func main() {
 		fmt.Println(err)
 	}
 	prompt, err := prompt.NewPromptTemplate("{{.Word}}の意味を教えて。")
+	parser := parser.NewStrOutputParser()
 
 	pipeline := lib.NewPipeline()
-	pipeline.Pipe(prompt).Pipe(llm)
+	pipeline.Pipe(prompt).Pipe(llm).Pipe(parser)
 	m := map[string]string{
 		"Word": "因果応報",
 	}
@@ -25,5 +27,4 @@ func main() {
 		fmt.Println(err)
 	}
 	fmt.Printf("Response: %+v\n", response)
-
 }

pkg/parser/stroutputparser.go

@@ -0,0 +1,23 @@
+package parser
+
+import (
+	"golangchain/pkg/openai"
+)
+
+type StrOutputParser struct {
+}
+
+func NewStrOutputParser() *StrOutputParser {
+	return &StrOutputParser{}
+}
+
+func (p *StrOutputParser) Invoke(input any) (any, error) {
+	var output string
+	res, ok := input.(*openai.Response)
+	if ok {
+		output = res.Choices[0].Message.Content
+	} else {
+		return nil, nil
+	}
+	return output, nil
+}

pkg/parser/stroutputparser_test.go

@@ -0,0 +1,36 @@
+package parser
+
+import (
+	"golangchain/pkg/openai"
+	"testing"
+)
+
+func TestInvoke(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    any
+		expected string
+	}{
+		{
+			name: "template is string",
+			input: &openai.Response{
+				Choices: []openai.Choice{
+					{Message: openai.Message{Role: "assistant", Content: "This is the test response"}},
+				},
+			},
+			expected: "This is the test response",
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			parser := NewStrOutputParser()
+			have, err := parser.Invoke(tc.input)
+			if err != nil {
+				t.Fatalf("Error happens: %v", err)
+			}
+			if have != tc.expected {
+				t.Fatalf("unexpected string: %v != %v", have, tc.expected)
+			}
+		})
+	}
+}