mirror of
https://github.com/vxcontrol/langchaingo.git
synced 2026-07-21 00:45:22 -04:00
Travis Cline
c544fb77bd
* all: add broad httprr coverage, update dependencies, organize go.mod file, bump to 1.23 update go version to 1.23 add lots of test coverage via httprr recordings update dependencies and organize go.mod add testutil/testctr which helps work around a testcontainers-go+colima bug expand the huggingface implementation and tests expand capabilities of the ollama package
60 lines
1.3 KiB
Go
60 lines
1.3 KiB
Go
package chains
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/tmc/langchaingo/prompts"
|
|
"github.com/tmc/langchaingo/schema"
|
|
)
|
|
|
|
func TestMapReduceInputVariables(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
c := MapReduceDocuments{
|
|
LLMChain: NewLLMChain(
|
|
&testLanguageModel{},
|
|
prompts.NewPromptTemplate("{{.text}} {{.foo}}", []string{"text", "foo"}),
|
|
),
|
|
ReduceChain: NewLLMChain(
|
|
&testLanguageModel{},
|
|
prompts.NewPromptTemplate("{{.texts}} {{.baz}}", []string{"texts", "baz"}),
|
|
),
|
|
ReduceDocumentVariableName: "texts",
|
|
LLMChainInputVariableName: "text",
|
|
InputKey: "input",
|
|
}
|
|
|
|
inputKeys := c.GetInputKeys()
|
|
expectedLength := 3
|
|
require.Len(t, inputKeys, expectedLength)
|
|
}
|
|
|
|
func TestMapReduce(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := context.Background()
|
|
|
|
c := NewMapReduceDocuments(
|
|
NewLLMChain(
|
|
&testLanguageModel{},
|
|
prompts.NewPromptTemplate("{{.context}}", []string{"context"}),
|
|
),
|
|
NewStuffDocuments(
|
|
NewLLMChain(
|
|
&testLanguageModel{},
|
|
prompts.NewPromptTemplate("{{.context}}", []string{"context"}),
|
|
),
|
|
),
|
|
)
|
|
|
|
result, err := Run(ctx, c, []schema.Document{
|
|
{PageContent: "foo"},
|
|
{PageContent: "boo"},
|
|
{PageContent: "zoo"},
|
|
{PageContent: "doo"},
|
|
})
|
|
require.NoError(t, err)
|
|
require.Equal(t, "foo\n\nboo\n\nzoo\n\ndoo", result)
|
|
}
|