Repository: tmc/langchaingo
Pull request #127 (2023-06-14 18:31:05)

document loaders: add csv loader
Add initial implementation for a CSV loader.

Python reference: https://github.com/hwchase17/langchain/blob/master/langchain/document_loaders/csv_loader.py

### PR Checklist

- [x] Read the [Contributing documentation](https://github.com/tmc/langchaingo/blob/main/CONTRIBUTING.md).
- [x] Read the [Code of conduct documentation](https://github.com/tmc/langchaingo/blob/main/CODE_OF_CONDUCT.md).
- [x] Name your Pull Request title clearly, concisely, and prefixed with the name of primarily affected package you changed according to [Good commit messages](https://go.dev/doc/contribute#commit_messages) (such as `memory: add interfaces for X, Y` or `util: add whizzbang helpers`).
- [x] Check that there isn't already a PR that solves the problem the same way to avoid creating a duplicate.
- [x] Provide a description in this PR that addresses **what** the PR is solving, or reference the issue that it solves (e.g. `fixes #123`).
- [x] Describes source of new concepts.
- [x] References existing implementations as appropriate.
- [x] Contains test coverage for new functions.
- [x] Passes all [`golangci-lint`](https://golangci-lint.run/) checks.


Patch:
diff --git a/docs/parity_matrix.md b/docs/parity_matrix.md
index b7d55936..e122a5c6 100644
--- a/docs/parity_matrix.md
+++ b/docs/parity_matrix.md
@@ -1,8 +1,8 @@
 # LangChain Go Features
 
-This page provides a list of features and their status indicators for LangChain Go, a natural language processing library built in Golang. Our main focus at the moment is to reach parity with the Python version of LangChain. 
+This page provides a list of features and their status indicators for LangChain Go, a natural language processing library built in Golang. Our main focus at the moment is to reach parity with the Python version of LangChain.
 
-Please note that this page lists the current state of the LangChain Go project, and some features may still be under development or planned for future releases. If you are interested in contributing to a specific integration or feature, feel free to choose from the list of features that are not yet done and let us know so we can help you get started. 
+Please note that this page lists the current state of the LangChain Go project, and some features may still be under development or planned for future releases. If you are interested in contributing to a specific integration or feature, feel free to choose from the list of features that are not yet done and let us know so we can help you get started.
 
 ## Prompt Templates
 | Feature                            | Status |
@@ -107,7 +107,7 @@ Please note that this page lists the current state of the LangChain Go project,
 | College Confidential            | ❌      |
 | Confluence                      | ❌      |
 | CoNLL-U                         | ❌      |
-| CSV Loader                      | ❌      |
+| CSV Loader                      | ✅      |
 | Dataframe                       | ❌      |
 | Diffbot                         | ❌      |
 | Directory                       | ❌      |
diff --git a/documentloaders/csv.go b/documentloaders/csv.go
new file mode 100644
index 00000000..10d96fa9
--- /dev/null
+++ b/documentloaders/csv.go
@@ -0,0 +1,83 @@
+package documentloaders
+
+import (
+	"context"
+	"encoding/csv"
+	"errors"
+	"fmt"
+	"io"
+	"strings"
+
+	"github.com/tmc/langchaingo/schema"
+	"github.com/tmc/langchaingo/textsplitter"
+	"golang.org/x/exp/slices"
+)
+
+// CSV represents a CSV document loader.
+type CSV struct {
+	r       io.Reader
+	columns []string
+}
+
+var _ Loader = CSV{}
+
+// NewCSV creates a new csv loader with an io.Reader and optional column names for filtering.
+func NewCSV(r io.Reader, columns ...string) CSV {
+	return CSV{
+		r:       r,
+		columns: columns,
+	}
+}
+
+// Load reads from the io.Reader and returns a single document with the data.
+func (c CSV) Load(_ context.Context) ([]schema.Document, error) {
+	var header []string
+	var docs []schema.Document
+	var rown int
+
+	rd := csv.NewReader(c.r)
+	for {
+		row, err := rd.Read()
+		if errors.Is(err, io.EOF) {
+			break
+		}
+		if err != nil {
+			return nil, err
+		}
+		if len(header) == 0 {
+			header = append(header, row...)
+			continue
+		}
+
+		var content []string
+		for i, value := range row {
+			if c.columns != nil &&
+				len(c.columns) > 0 &&
+				!slices.Contains[string](c.columns, header[i]) {
+				continue
+			}
+
+			line := fmt.Sprintf("%s: %s", header[i], value)
+			content = append(content, line)
+		}
+
+		rown++
+		docs = append(docs, schema.Document{
+			PageContent: strings.Join(content, "\n"),
+			Metadata:    map[string]any{"row": rown},
+		})
+	}
+
+	return docs, nil
+}
+
+// LoadAndSplit reads text data from the io.Reader and splits it into multiple
+// documents using a text splitter.
+func (c CSV) LoadAndSplit(ctx context.Context, splitter textsplitter.TextSplitter) ([]schema.Document, error) {
+	docs, err := c.Load(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	return textsplitter.SplitDocuments(splitter, docs)
+}


Test patch:
diff --git a/documentloaders/csv_test.go b/documentloaders/csv_test.go
new file mode 100644
index 00000000..ccbaf008
--- /dev/null
+++ b/documentloaders/csv_test.go
@@ -0,0 +1,31 @@
+package documentloaders
+
+import (
+	"context"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCSVLoader(t *testing.T) {
+	t.Parallel()
+	file, err := os.Open("./testdata/test.csv")
+	assert.NoError(t, err)
+
+	loader := NewCSV(file)
+
+	docs, err := loader.Load(context.Background())
+	require.NoError(t, err)
+	require.Len(t, docs, 20)
+
+	expected1 := "name: John Doe\nage: 25\ncity: New York\ncountry: United States"
+	assert.Equal(t, docs[0].PageContent, expected1)
+
+	expected2 := `name: Jane Smith
+age: 32
+city: London
+country: United Kingdom`
+	assert.Equal(t, docs[1].PageContent, expected2)
+}
diff --git a/documentloaders/testdata/test.csv b/documentloaders/testdata/test.csv
new file mode 100644
index 00000000..11a95730
--- /dev/null
+++ b/documentloaders/testdata/test.csv
@@ -0,0 +1,21 @@
+name,age,city,country
+John Doe,25,New York,United States
+Jane Smith,32,London,United Kingdom
+Alex Johnson,42,Sydney,Australia
+Emma Davis,29,Paris,France
+Michael Lee,37,Toronto,Canada
+Sophia Wilson,22,Berlin,Germany
+Matthew Brown,31,Tokyo,Japan
+Olivia Taylor,27,Rome,Italy
+David Anderson,35,Moscow,Russia
+Emily Thomas,30,Madrid,Spain
+Daniel Martinez,26,Mexico City,Mexico
+Isabella Lopez,28,Seoul,South Korea
+Christopher Harris,33,Cairo,Egypt
+Ava Clark,24,Sao Paulo,Brazil
+James Wright,39,Amsterdam,Netherlands
+Mia Hall,23,Stockholm,Sweden
+Andrew Allen,34,Beijing,China
+Abigail Baker,36,New Delhi,India
+Ryan Scott,41,Johannesburg,South Africa
+Grace Green,38,Zurich,Switzerland
diff --git a/textsplitter/recursive_character_test.go b/textsplitter/recursive_character_test.go
index aa855755..242ef0c6 100644
--- a/textsplitter/recursive_character_test.go
+++ b/textsplitter/recursive_character_test.go
@@ -10,14 +10,12 @@ import (
 //nolint:dupword
 func TestRecursiveCharacterSplitter(t *testing.T) {
 	t.Parallel()
-
 	type testCase struct {
 		text         string
 		chunkOverlap int
 		chunkSize    int
 		expectedDocs []schema.Document
 	}
-
 	testCases := []testCase{
 		{
 			text:         "Hi.\nI'm Harrison.\n\nHow?\na\nb",
@@ -37,6 +35,27 @@ func TestRecursiveCharacterSplitter(t *testing.T) {
 				{PageContent: "How?\na\nbHi.\nI'm Harrison.\n\nHow?\na\nb", Metadata: map[string]any{}},
 			},
 		},
+		{
+			text:         "name: Harrison\nage: 30",
+			chunkOverlap: 1,
+			chunkSize:    40,
+			expectedDocs: []schema.Document{
+				{PageContent: "name: Harrison\nage: 30", Metadata: map[string]any{}},
+			},
+		},
+		{
+			text: `name: Harrison
+age: 30
+
+name: Joe
+age: 32`,
+			chunkOverlap: 1,
+			chunkSize:    40,
+			expectedDocs: []schema.Document{
+				{PageContent: "name: Harrison\nage: 30", Metadata: map[string]any{}},
+				{PageContent: "name: Joe\nage: 32", Metadata: map[string]any{}},
+			},
+		},
 		{
 			text: `Hi.
 I'm Harrison.
@@ -71,7 +90,6 @@ Bye!
 			},
 		},
 	}
-
 	splitter := NewRecursiveCharacter()
 	for _, tc := range testCases {
 		splitter.ChunkOverlap = tc.chunkOverlap