add: documentloader and textsplitter config

pkg/datastore/defaults/defaults.go

@@ -1,4 +1,11 @@
 package defaults
 
-const EmbeddingDimension int = 1536
-const TopK int = 5
+const (
+	EmbeddingDimension int = 1536
+	TopK               int = 5
+
+	TextSplitterTokenModel    = "gpt-4"
+	TextSplitterChunkSize     = 1024
+	TextSplitterChunkOverlap  = 256
+	TextSplitterTokenEncoding = "cl100k_base"
+)

pkg/datastore/documentloader/documentloaders_test.go

@@ -0,0 +1,54 @@
+package documentloader
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetDocumentLoaderConfig_ValidLoader(t *testing.T) {
+	cfg, err := GetDocumentLoaderConfig("pdf")
+	assert.NoError(t, err)
+	assert.IsTypef(t, PDFOptions{}, cfg, "cfg is not of type PDFOptions")
+}
+
+func TestGetDocumentLoaderConfig_InvalidLoader(t *testing.T) {
+	_, err := GetDocumentLoaderConfig("invalid")
+	assert.Error(t, err)
+}
+
+func TestGetDocumentLoaderFunc_ValidLoaderWithoutConfig(t *testing.T) {
+	_, err := GetDocumentLoaderFunc("plaintext", nil)
+	assert.NoError(t, err)
+}
+
+func TestGetDocumentLoaderFunc_ValidLoaderWithInvalidConfig(t *testing.T) {
+	_, err := GetDocumentLoaderFunc("pdf", "invalid")
+	assert.Error(t, err)
+}
+
+func TestGetDocumentLoaderFunc_ValidLoaderWithValidConfig(t *testing.T) {
+	_, err := GetDocumentLoaderFunc("pdf", PDFOptions{})
+	assert.NoError(t, err)
+}
+
+func TestGetDocumentLoaderFunc_InvalidLoader(t *testing.T) {
+	_, err := GetDocumentLoaderFunc("invalid", nil)
+	assert.Error(t, err)
+}
+
+func TestGetDocumentLoaderFunc_LoadPlainText(t *testing.T) {
+	loaderFunc, _ := GetDocumentLoaderFunc("plaintext", nil)
+	docs, err := loaderFunc(context.Background(), strings.NewReader("test"))
+	assert.NoError(t, err)
+	assert.Len(t, docs, 1)
+	assert.Equal(t, "test", docs[0].Content)
+}
+
+func TestGetDocumentLoaderFunc_LoadPDF(t *testing.T) {
+	loaderFunc, _ := GetDocumentLoaderFunc("pdf", PDFOptions{})
+	_, err := loaderFunc(context.Background(), strings.NewReader("test"))
+	assert.Error(t, err)
+}

pkg/datastore/ingest.go

@@ -26,12 +26,7 @@ import (
 	"strings"
 )
 
-const (
-	defaultTokenModel    = "gpt-4"
-	defaultChunkSize     = 1024
-	defaultChunkOverlap  = 256
-	defaultTokenEncoding = "cl100k_base"
-)
+const ()
 
 var firstclassFileExtensions = map[string]struct{}{
 	".pdf":   {},
@@ -51,7 +46,7 @@ type IngestOpts struct {
 	FileMetadata        *index.FileMetadata
 	IsDuplicateFuncName string
 	IsDuplicateFunc     IsDuplicateFunc
-	TextSplitterOpts    *TextSplitterOpts
+	TextSplitterOpts    *textsplitter.TextSplitterOpts
 }
 
 // Ingest loads a document from a reader and adds it to the dataset.
@@ -270,12 +265,12 @@ func DefaultDocLoaderFunc(filetype string) func(ctx context.Context, reader io.R
 	}
 }
 
-func DefaultTextSplitter(filetype string, textSplitterOpts *TextSplitterOpts) types.TextSplitter {
+func DefaultTextSplitter(filetype string, textSplitterOpts *textsplitter.TextSplitterOpts) types.TextSplitter {
 	if textSplitterOpts == nil {
-		textSplitterOpts = z.Pointer(NewTextSplitterOpts())
+		textSplitterOpts = z.Pointer(textsplitter.NewTextSplitterOpts())
 	}
-	genericTextSplitter := textsplitter.FromLangchain(NewLcgoTextSplitter(*textSplitterOpts))
-	markdownTextSplitter := textsplitter.FromLangchain(NewLcgoMarkdownSplitter(*textSplitterOpts))
+	genericTextSplitter := textsplitter.FromLangchain(textsplitter.NewLcgoTextSplitter(*textSplitterOpts))
+	markdownTextSplitter := textsplitter.FromLangchain(textsplitter.NewLcgoMarkdownSplitter(*textSplitterOpts))
 
 	switch filetype {
 	case ".md", "text/markdown":

pkg/datastore/textsplitter.go

@@ -2,47 +2,9 @@ package datastore
 
 import (
 	vs "github.com/gptscript-ai/knowledge/pkg/vectorstore"
-	lcgosplitter "github.com/tmc/langchaingo/textsplitter"
 	"strings"
 )
 
-type TextSplitterOpts struct {
-	ChunkSize    int    `usage:"Textsplitter Chunk Size" default:"1024" env:"KNOW_TEXTSPLITTER_CHUNK_SIZE" name:"textsplitter-chunk-size"`
-	ChunkOverlap int    `usage:"Textsplitter Chunk Overlap" default:"256" env:"KNOW_TEXTSPLITTER_CHUNK_OVERLAP" name:"textsplitter-chunk-overlap"`
-	ModelName    string `usage:"Textsplitter Model Name" default:"gpt-4" env:"KNOW_TEXTSPLITTER_MODEL_NAME" name:"textsplitter-model-name"`
-	EncodingName string `usage:"Textsplitter Encoding Name" default:"cl100k_base" env:"KNOW_TEXTSPLITTER_ENCODING_NAME" name:"textsplitter-encoding-name"`
-}
-
-// NewTextSplitterOpts returns the default options for a text splitter.
-func NewTextSplitterOpts() TextSplitterOpts {
-	return TextSplitterOpts{
-		ChunkSize:    defaultChunkSize,
-		ChunkOverlap: defaultChunkOverlap,
-		ModelName:    defaultTokenModel,
-		EncodingName: defaultTokenEncoding,
-	}
-}
-
-// NewLcgoTextSplitter returns a new langchain-go text splitter.
-func NewLcgoTextSplitter(opts TextSplitterOpts) lcgosplitter.TokenSplitter {
-	return lcgosplitter.NewTokenSplitter(
-		lcgosplitter.WithChunkSize(opts.ChunkSize),
-		lcgosplitter.WithChunkOverlap(opts.ChunkOverlap),
-		lcgosplitter.WithModelName(opts.ModelName),
-		lcgosplitter.WithEncodingName(opts.EncodingName),
-	)
-}
-
-func NewLcgoMarkdownSplitter(opts TextSplitterOpts) *lcgosplitter.MarkdownTextSplitter {
-	return lcgosplitter.NewMarkdownTextSplitter(
-		lcgosplitter.WithChunkSize(opts.ChunkSize),
-		lcgosplitter.WithChunkOverlap(opts.ChunkOverlap),
-		lcgosplitter.WithModelName(opts.ModelName),
-		lcgosplitter.WithEncodingName(opts.EncodingName),
-		lcgosplitter.WithHeadingHierarchy(true),
-	)
-}
-
 // FilterMarkdownDocsNoContent filters out Markdown documents with no content or only headings
 //
 // TODO: this may be moved into the MarkdownTextSplitter as well

pkg/datastore/textsplitter/textsplitter.go

@@ -0,0 +1,82 @@
+package textsplitter
+
+import (
+	"fmt"
+	"github.com/gptscript-ai/knowledge/pkg/datastore/defaults"
+	vs "github.com/gptscript-ai/knowledge/pkg/vectorstore"
+	lcgosplitter "github.com/tmc/langchaingo/textsplitter"
+)
+
+type SplitterFunc func([]vs.Document) ([]vs.Document, error)
+
+type TextSplitterOpts struct {
+	ChunkSize    int    `usage:"Textsplitter Chunk Size" default:"1024" env:"KNOW_TEXTSPLITTER_CHUNK_SIZE" name:"textsplitter-chunk-size"`
+	ChunkOverlap int    `usage:"Textsplitter Chunk Overlap" default:"256" env:"KNOW_TEXTSPLITTER_CHUNK_OVERLAP" name:"textsplitter-chunk-overlap"`
+	ModelName    string `usage:"Textsplitter Model Name" default:"gpt-4" env:"KNOW_TEXTSPLITTER_MODEL_NAME" name:"textsplitter-model-name"`
+	EncodingName string `usage:"Textsplitter Encoding Name" default:"cl100k_base" env:"KNOW_TEXTSPLITTER_ENCODING_NAME" name:"textsplitter-encoding-name"`
+}
+
+// NewTextSplitterOpts returns the default options for a text splitter.
+func NewTextSplitterOpts() TextSplitterOpts {
+	return TextSplitterOpts{
+		ChunkSize:    defaults.TextSplitterChunkSize,
+		ChunkOverlap: defaults.TextSplitterChunkOverlap,
+		ModelName:    defaults.TextSplitterTokenModel,
+		EncodingName: defaults.TextSplitterTokenEncoding,
+	}
+}
+
+// NewLcgoTextSplitter returns a new langchain-go text splitter.
+func NewLcgoTextSplitter(opts TextSplitterOpts) lcgosplitter.TokenSplitter {
+	return lcgosplitter.NewTokenSplitter(
+		lcgosplitter.WithChunkSize(opts.ChunkSize),
+		lcgosplitter.WithChunkOverlap(opts.ChunkOverlap),
+		lcgosplitter.WithModelName(opts.ModelName),
+		lcgosplitter.WithEncodingName(opts.EncodingName),
+	)
+}
+
+func NewLcgoMarkdownSplitter(opts TextSplitterOpts) *lcgosplitter.MarkdownTextSplitter {
+	return lcgosplitter.NewMarkdownTextSplitter(
+		lcgosplitter.WithChunkSize(opts.ChunkSize),
+		lcgosplitter.WithChunkOverlap(opts.ChunkOverlap),
+		lcgosplitter.WithModelName(opts.ModelName),
+		lcgosplitter.WithEncodingName(opts.EncodingName),
+		lcgosplitter.WithHeadingHierarchy(true),
+	)
+}
+
+func GetTextSplitterConfig(name string) (any, error) {
+	// TODO: expose splitter-specific config, not only our top-level options
+	switch name {
+	case "text", "markdown":
+		return TextSplitterOpts{}, nil
+	default:
+		return nil, fmt.Errorf("unknown text splitter %q", name)
+	}
+}
+
+func GetTextSplitterFunc(name string, config any) (SplitterFunc, error) {
+	switch name {
+	case "text":
+		if config == nil {
+			config = NewTextSplitterOpts()
+		}
+		config, ok := config.(TextSplitterOpts)
+		if !ok {
+			return nil, fmt.Errorf("invalid text splitter configuration")
+		}
+		return FromLangchain(NewLcgoTextSplitter(config)).SplitDocuments, nil
+	case "markdown":
+		if config == nil {
+			config = NewTextSplitterOpts()
+		}
+		config, ok := config.(TextSplitterOpts)
+		if !ok {
+			return nil, fmt.Errorf("invalid markdown text splitter configuration")
+		}
+		return FromLangchain(NewLcgoMarkdownSplitter(config)).SplitDocuments, nil
+	default:
+		return nil, fmt.Errorf("unknown text splitter %q", name)
+	}
+}

pkg/datastore/textsplitter/textsplitter_test.go

@@ -0,0 +1,36 @@
+package textsplitter
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestGetTextSplitterConfigWithValidName(t *testing.T) {
+	_, err := GetTextSplitterConfig("text")
+	assert.NoError(t, err)
+}
+
+func TestGetTextSplitterConfigWithInvalidName(t *testing.T) {
+	_, err := GetTextSplitterConfig("invalid")
+	assert.Error(t, err)
+}
+
+func TestGetTextSplitterFuncWithValidNameAndNilConfig(t *testing.T) {
+	_, err := GetTextSplitterFunc("text", nil)
+	assert.NoError(t, err)
+}
+
+func TestGetTextSplitterFuncWithValidNameAndInvalidConfig(t *testing.T) {
+	_, err := GetTextSplitterFunc("text", "invalid")
+	assert.Error(t, err)
+}
+
+func TestGetTextSplitterFuncWithValidNameAndValidConfig(t *testing.T) {
+	_, err := GetTextSplitterFunc("text", NewTextSplitterOpts())
+	assert.NoError(t, err)
+}
+
+func TestGetTextSplitterFuncWithInvalidName(t *testing.T) {
+	_, err := GetTextSplitterFunc("invalid", nil)
+	assert.Error(t, err)
+}

pkg/datastore/types/types.go

@@ -3,7 +3,6 @@ package types
 import (
 	"context"
 	vs "github.com/gptscript-ai/knowledge/pkg/vectorstore"
-	"io"
 )
 
 type DocumentTransformerFunc func(context.Context, []vs.Document) ([]vs.Document, error)
@@ -12,10 +11,6 @@ type DocumentTransformer interface {
 	Transform(context.Context, []vs.Document) ([]vs.Document, error)
 }
 
-type DocumentLoaderFunc func(context.Context, io.Reader) ([]vs.Document, error)
-
-type TextSplitterFunc func([]vs.Document) ([]vs.Document, error)
-
 type DocumentLoader interface {
 	Load(ctx context.Context) ([]vs.Document, error)
 	LoadAndSplit(ctx context.Context, splitter TextSplitter) ([]vs.Document, error)

pkg/flows/config/config.go

@@ -3,6 +3,7 @@ package config
 import (
 	"encoding/json"
 	"github.com/gptscript-ai/knowledge/pkg/datastore/documentloader"
+	"github.com/gptscript-ai/knowledge/pkg/datastore/textsplitter"
 	"github.com/gptscript-ai/knowledge/pkg/flows"
 	"os"
 	"sigs.k8s.io/yaml"
@@ -79,6 +80,34 @@ func (i *IngestionFlowConfig) AsIngestionFlow() (*flows.IngestionFlow, error) {
 				return nil, err
 			}
 		}
+		loaderFunc, err := documentloader.GetDocumentLoaderFunc(name, cfg)
+		if err != nil {
+			return nil, err
+		}
+		flow.Load = loaderFunc
+	}
+
+	if i.TextSplitter.Name != "" {
+		name := strings.ToLower(strings.Trim(i.TextSplitter.Name, " "))
+		cfg, err := textsplitter.GetTextSplitterConfig(name)
+		if err != nil {
+			return nil, err
+		}
+		if len(i.TextSplitter.Options) > 0 {
+			jsondata, err := json.Marshal(i.TextSplitter.Options)
+			if err != nil {
+				return nil, err
+			}
+			err = json.Unmarshal(jsondata, &cfg)
+			if err != nil {
+				return nil, err
+			}
+		}
+		splitterFunc, err := textsplitter.GetTextSplitterFunc(name, cfg)
+		if err != nil {
+			return nil, err
+		}
+		flow.Split = splitterFunc
 	}
 
 	return flow, nil

pkg/flows/config/config_test.go

@@ -2,13 +2,14 @@ package config
 
 import (
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"testing"
 )
 
 func TestLoadConfigFromValidJSONFile(t *testing.T) {
 	cfg, err := FromFile("testdata/valid.json")
 	assert.NoError(t, err)
-	assert.NotNil(t, cfg)
+	require.NotNil(t, cfg)
 	assert.NotEmpty(t, cfg.Flows)
 	assert.Equal(t, 2, len(cfg.Flows))
 	assert.Equal(t, 1, len(cfg.Flows["flow1"].Ingestion))
@@ -19,7 +20,7 @@ func TestLoadConfigFromValidJSONFile(t *testing.T) {
 func TestLoadConfigFromValidYAMLFile(t *testing.T) {
 	cfg, err := FromFile("testdata/valid.yaml")
 	assert.NoError(t, err)
-	assert.NotNil(t, cfg)
+	require.NotNil(t, cfg)
 	assert.NotEmpty(t, cfg.Flows)
 }
 

pkg/flows/config/testdata/valid.json

@@ -4,8 +4,12 @@
       "ingestion": [
         {
           "filetypes": [".txt", ".md"],
-          "documentLoader": "textLoader",
-          "textSplitter": "simpleSplitter",
+          "documentLoader": {
+            "name": "textLoader"
+          },
+          "textSplitter": {
+            "name": "simpleSplitter"
+          },
           "transformers": ["transformer1", "transformer2"]
         }
       ],
@@ -15,8 +19,12 @@
       "ingestion": [
         {
           "filetypes": [".json"],
-          "documentLoader": "jsonLoader",
-          "textSplitter": "jsonSplitter",
+          "documentLoader": {
+            "name": "jsonLoader"
+          },
+          "textSplitter": {
+            "name": "jsonSplitter"
+          },
           "transformers": ["transformer3"]
         }
       ],