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
71 lines
1.7 KiB
Go
71 lines
1.7 KiB
Go
package chains
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/tmc/langchaingo/prompts"
|
|
"github.com/tmc/langchaingo/schema"
|
|
)
|
|
|
|
func TestMapRerankInputVariables(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
mapRerankLLMChain := NewLLMChain(
|
|
&testLanguageModel{},
|
|
prompts.NewPromptTemplate("{{.text}} {{.foo}}", []string{"text", "foo"}),
|
|
)
|
|
|
|
c := MapRerankDocuments{
|
|
LLMChain: mapRerankLLMChain,
|
|
DocumentVariableName: "texts",
|
|
LLMChainInputVariableName: "text",
|
|
InputKey: "input",
|
|
}
|
|
|
|
inputKeys := c.GetInputKeys()
|
|
expectedLength := 3
|
|
require.Len(t, inputKeys, expectedLength)
|
|
}
|
|
|
|
func TestMapRerankDocumentsCall(t *testing.T) {
|
|
ctx := context.Background()
|
|
t.Parallel()
|
|
|
|
mapRerankLLMChain := NewLLMChain(
|
|
&testLanguageModel{},
|
|
prompts.NewPromptTemplate("{{.context}}", []string{"context"}),
|
|
)
|
|
|
|
docs := []schema.Document{
|
|
{PageContent: "Test Low\nScore: 20"},
|
|
{PageContent: "Test High\nScore: 100"},
|
|
}
|
|
|
|
mapRerankDocumentsChain := NewMapRerankDocuments(mapRerankLLMChain)
|
|
|
|
// Test that the answer is the highest scoring document.
|
|
answer, err := Run(ctx, mapRerankDocumentsChain, docs)
|
|
|
|
require.NoError(t, err)
|
|
require.Equal(t, "Test High", answer)
|
|
|
|
// Test that the answer cannot be processed if ReturnIntermediateSteps is true.
|
|
mapRerankDocumentsChain.ReturnIntermediateSteps = true
|
|
_, err = Run(ctx, mapRerankDocumentsChain, docs)
|
|
|
|
require.Error(t, err)
|
|
|
|
// Test that an error is returned if the score cannot be processed.
|
|
mapRerankDocumentsChain.ReturnIntermediateSteps = false
|
|
docs = []schema.Document{
|
|
{PageContent: "Test Low\nScore:"},
|
|
{PageContent: "Test High\nScore:"},
|
|
}
|
|
|
|
_, err = Run(ctx, mapRerankDocumentsChain, docs)
|
|
|
|
require.Error(t, err)
|
|
}
|