Repository: tmc/langchaingo
Pull request #401 (2023-12-05 12:07:58)

documentloaders: add notion directory loader
This PR adds the NotionDirectoryLoader feature to Go, similar to its [Python counterpart](https://github.com/langchain-ai/langchain/blob/master/libs/langchain/langchain/document_loaders/notion.py). The loader manages fetching and arranging local Notion Markdown files from a specified directory.

### 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 the 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 the 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/documentloaders/notion.go b/documentloaders/notion.go
new file mode 100644
index 00000000..0dacbe78
--- /dev/null
+++ b/documentloaders/notion.go
@@ -0,0 +1,57 @@
+package documentloaders
+
+import (
+	"os"
+	"path/filepath"
+
+	"github.com/tmc/langchaingo/schema"
+)
+
+// NotionDirectoryLoader is a document loader that reads content from pages within a Notion Database.
+type NotionDirectoryLoader struct {
+	filePath string
+	encoding string
+}
+
+// NewNotionDirectory creates a new NotionDirectoryLoader with the given file path and encoding.
+func NewNotionDirectory(filePath string, encoding ...string) *NotionDirectoryLoader {
+	defaultEncoding := "utf-8"
+
+	if len(encoding) > 0 {
+		return &NotionDirectoryLoader{
+			filePath: filePath,
+			encoding: encoding[0],
+		}
+	}
+
+	return &NotionDirectoryLoader{
+		filePath: filePath,
+		encoding: defaultEncoding,
+	}
+}
+
+// Load retrieves data from a Notion directory and returns a list of schema.Document objects.
+func (n *NotionDirectoryLoader) Load() ([]schema.Document, error) {
+	files, err := os.ReadDir(n.filePath)
+	if err != nil {
+		return nil, err
+	}
+
+	documents := make([]schema.Document, 0, len(files))
+	for _, file := range files {
+		if file.IsDir() || filepath.Ext(file.Name()) != ".md" {
+			continue
+		}
+
+		filePath := filepath.Join(n.filePath, file.Name())
+		text, err := os.ReadFile(filePath)
+		if err != nil {
+			return nil, err
+		}
+
+		metadata := map[string]interface{}{"source": filePath}
+		documents = append(documents, schema.Document{PageContent: string(text), Metadata: metadata})
+	}
+
+	return documents, nil
+}
diff --git a/llms/ernie/erniellm.go b/llms/ernie/erniellm.go
index f7ce9686..9d8d2ae6 100644
--- a/llms/ernie/erniellm.go
+++ b/llms/ernie/erniellm.go
@@ -160,6 +160,8 @@ func modelToPath(model ModelName) ernieclient.ModelPath {
 		return "completions"
 	case ModelNameERNIEBotTurbo:
 		return "eb-instant"
+	case ModelNameERNIEBot8K:
+		return "ernie_bot_8k"
 	case ModelNameERNIEBotPro:
 		return "completions_pro"
 	case ModelNameBloomz7B:
diff --git a/llms/ernie/erniellm_option.go b/llms/ernie/erniellm_option.go
index 2f923b6e..148e56f3 100644
--- a/llms/ernie/erniellm_option.go
+++ b/llms/ernie/erniellm_option.go
@@ -9,6 +9,7 @@ type ModelName string
 
 const (
 	ModelNameERNIEBot       = "ERNIE-Bot"
+	ModelNameERNIEBot8K     = "ERNIE-Bot-8K"
 	ModelNameERNIEBotTurbo  = "ERNIE-Bot-turbo"
 	ModelNameERNIEBotPro    = "ERNIE-Bot-pro"
 	ModelNameBloomz7B       = "BLOOMZ-7B"
diff --git a/llms/openai/openaillm_chat.go b/llms/openai/openaillm_chat.go
index 128a8120..f4e35819 100644
--- a/llms/openai/openaillm_chat.go
+++ b/llms/openai/openaillm_chat.go
@@ -165,6 +165,12 @@ func messagesToClientMessages(messages []schema.ChatMessage) []*openaiclient.Cha
 			msg.Role = "system"
 		case schema.ChatMessageTypeAI:
 			msg.Role = "assistant"
+			if mm, ok := m.(schema.AIChatMessage); ok && mm.FunctionCall != nil {
+				msg.FunctionCall = &openaiclient.FunctionCall{
+					Name:      mm.FunctionCall.Name,
+					Arguments: mm.FunctionCall.Arguments,
+				}
+			}
 		case schema.ChatMessageTypeHuman:
 			msg.Role = "user"
 		case schema.ChatMessageTypeGeneric:


Test patch:
diff --git a/documentloaders/notion_test.go b/documentloaders/notion_test.go
new file mode 100644
index 00000000..06920e65
--- /dev/null
+++ b/documentloaders/notion_test.go
@@ -0,0 +1,63 @@
+package documentloaders
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/tmc/langchaingo/schema"
+)
+
+func TestNotionDirectoryLoader_Load(t *testing.T) {
+	t.Parallel()
+
+	// Create a temporary test directory
+	tempDir, err := os.MkdirTemp("", "notion_test")
+	require.NoError(t, err)
+	defer os.RemoveAll(tempDir)
+
+	// Create sample Markdown files in the temporary directory
+	testFiles := []struct {
+		name     string
+		content  string
+		expected schema.Document
+	}{
+		{
+			name:    "test1.md",
+			content: "# Test Document 1\nThis is test document 1.",
+			expected: schema.Document{
+				PageContent: "# Test Document 1\nThis is test document 1.",
+				Metadata:    map[string]interface{}{"source": filepath.Join(tempDir, "test1.md")},
+			},
+		},
+		{
+			name:    "test2.md",
+			content: "# Test Document 2\nThis is test document 2.",
+			expected: schema.Document{
+				PageContent: "# Test Document 2\nThis is test document 2.",
+				Metadata:    map[string]interface{}{"source": filepath.Join(tempDir, "test2.md")},
+			},
+		},
+	}
+
+	for _, file := range testFiles {
+		filePath := filepath.Join(tempDir, file.name)
+		err := os.WriteFile(filePath, []byte(file.content), 0o600)
+		require.NoError(t, err)
+	}
+
+	// Create a NotionDirectoryLoader instance
+	loader := NewNotionDirectory(tempDir)
+
+	// Load documents from the test directory
+	docs, err := loader.Load()
+	require.NoError(t, err)
+
+	// Verify the loaded documents match the expected ones
+	require.Len(t, docs, len(testFiles))
+	for i, expected := range testFiles {
+		assert.Equal(t, expected.expected, docs[i])
+	}
+}