documentloaders: ms office docs #1068
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,228 @@ | ||
| package documentloaders | ||
|
|
||
| import ( | ||
| "archive/zip" | ||
| "bytes" | ||
| "context" | ||
| "fmt" | ||
| "io" | ||
| "path/filepath" | ||
| "strings" | ||
|
|
||
| "github.com/richardlehane/mscfb" | ||
| "github.com/tealeg/xlsx" | ||
| "github.com/tmc/langchaingo/schema" | ||
| "github.com/tmc/langchaingo/textsplitter" | ||
| ) | ||
|
|
||
| var _ Loader = Office{} | ||
|
|
||
| // Office loads text data from an io.Reader. | ||
| type Office struct { | ||
| reader io.ReaderAt | ||
| size int64 | ||
| fileType string | ||
| } | ||
|
|
||
| // NewOffice creates a new text loader with an io.Reader, filename and file size. | ||
| func NewOffice(reader io.ReaderAt, filename string, size int64) Office { | ||
| return Office{ | ||
| reader: reader, | ||
| size: size, | ||
| fileType: strings.ToLower(filepath.Ext(filename)), | ||
| } | ||
| } | ||
|
|
||
| // Load reads from the io.Reader for the MS Office data and returns the raw document data. | ||
| // nolint | ||
| func (loader Office) Load(ctx context.Context) ([]schema.Document, error) { | ||
| switch loader.fileType { | ||
| case ".doc": | ||
| return loader.loadDoc() | ||
| case ".docx": | ||
| return loader.loadDocx() | ||
| case ".xls", ".xlsx": | ||
| return loader.loadExcel() | ||
| case ".ppt": | ||
| // parsing for old PPTs is same as for old DOCs | ||
| return loader.loadDoc() | ||
| case ".pptx": | ||
| return loader.loadPPTX() | ||
| default: | ||
| return nil, fmt.Errorf("unsupported file type: %s", loader.fileType) | ||
| } | ||
| } | ||
|
|
||
| // LoadAndSplit reads from the io.Reader for the MS Office data and returns the raw document data | ||
| // and splits it into multiple documents using a text splitter. | ||
| func (loader Office) LoadAndSplit(ctx context.Context, splitter textsplitter.TextSplitter) ([]schema.Document, error) { | ||
| docs, err := loader.Load(ctx) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| return textsplitter.SplitDocuments(splitter, docs) | ||
| } | ||
|
|
||
| func (loader Office) loadDoc() ([]schema.Document, error) { | ||
| doc, err := mscfb.New(io.NewSectionReader(loader.reader, 0, loader.size)) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to read DOC file: %w", err) | ||
| } | ||
|
|
||
| var text strings.Builder | ||
| for entry, err := doc.Next(); err == nil; entry, err = doc.Next() { | ||
| // nolint | ||
| if entry.Name == "WordDocument" { | ||
| buf := make([]byte, entry.Size) | ||
| i, err := doc.Read(buf) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("error reading WordDocument stream: %w", err) | ||
| } | ||
| if i > 0 { | ||
| // Process the binary content | ||
| for j := 0; j < i; j++ { | ||
| // Extract readable ASCII text | ||
| if buf[j] >= 32 && buf[j] <= 126 { | ||
|
There was a problem hiding this comment. I think these numbers be unexported constants |
||
| text.WriteByte(buf[j]) | ||
| } else if buf[j] == 13 || buf[j] == 10 { | ||
| text.WriteByte('\n') | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| documents := []schema.Document{ | ||
| { | ||
| PageContent: text.String(), | ||
| Metadata: map[string]interface{}{ | ||
| "fileType": loader.fileType, | ||
| }, | ||
| }, | ||
| } | ||
|
|
||
| return documents, nil | ||
| } | ||
|
|
||
| func (loader Office) loadExcel() ([]schema.Document, error) { | ||
| buf := bytes.NewBuffer(make([]byte, 0, loader.size)) | ||
| if _, err := io.Copy(buf, io.NewSectionReader(loader.reader, 0, loader.size)); err != nil { | ||
| return nil, fmt.Errorf("failed to copy Excel content: %w", err) | ||
| } | ||
|
|
||
| xlFile, err := xlsx.OpenBinary(buf.Bytes()) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to read Excel file: %w", err) | ||
| } | ||
|
|
||
| docs := make([]schema.Document, 0, len(xlFile.Sheets)) | ||
| for i, sheet := range xlFile.Sheets { | ||
| var text strings.Builder | ||
| for _, row := range sheet.Rows { | ||
| for _, cell := range row.Cells { | ||
| text.WriteString(cell.String() + "\t") | ||
| } | ||
| text.WriteString("\n") | ||
| } | ||
|
|
||
| docs = append(docs, schema.Document{ | ||
| PageContent: text.String(), | ||
| Metadata: map[string]interface{}{ | ||
| "fileType": loader.fileType, | ||
| "sheetName": sheet.Name, | ||
| "sheetIndex": i, | ||
| }, | ||
| }) | ||
| } | ||
|
|
||
| return docs, nil | ||
| } | ||
|
|
||
| func (loader Office) loadPPTX() ([]schema.Document, error) { | ||
| buf := bytes.NewBuffer(make([]byte, 0, loader.size)) | ||
| if _, err := io.Copy(buf, io.NewSectionReader(loader.reader, 0, loader.size)); err != nil { | ||
| return nil, fmt.Errorf("failed to copy content: %w", err) | ||
| } | ||
|
|
||
| zipReader, err := zip.NewReader(bytes.NewReader(buf.Bytes()), loader.size) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to read PPTX file as ZIP: %w", err) | ||
| } | ||
|
|
||
| var text strings.Builder | ||
| for _, file := range zipReader.File { | ||
| // PPTX stores slide content in ppt/slides/slide*.xml files | ||
| if strings.HasPrefix(file.Name, "ppt/slides/slide") && strings.HasSuffix(file.Name, ".xml") { | ||
| rc, err := file.Open() | ||
| if err != nil { | ||
| return nil, fmt.Errorf("error opening slide XML: %w", err) | ||
| } | ||
| defer rc.Close() | ||
|
|
||
| content, err := io.ReadAll(rc) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("error reading content: %w", err) | ||
| } | ||
|
|
||
| content = bytes.ReplaceAll(content, []byte("<"), []byte(" <")) | ||
| content = bytes.ReplaceAll(content, []byte(">"), []byte("> ")) | ||
| text.Write(content) | ||
| text.WriteString("\n--- Next Slide ---\n") | ||
| } | ||
| } | ||
|
|
||
| documents := []schema.Document{ | ||
| { | ||
| PageContent: text.String(), | ||
| Metadata: map[string]interface{}{ | ||
| "fileType": loader.fileType, | ||
| }, | ||
| }, | ||
| } | ||
|
|
||
| return documents, nil | ||
| } | ||
|
|
||
| func (loader Office) loadDocx() ([]schema.Document, error) { | ||
| buf := bytes.NewBuffer(make([]byte, 0, loader.size)) | ||
| if _, err := io.Copy(buf, io.NewSectionReader(loader.reader, 0, loader.size)); err != nil { | ||
| return nil, fmt.Errorf("failed to copy content: %w", err) | ||
| } | ||
|
|
||
| zipReader, err := zip.NewReader(bytes.NewReader(buf.Bytes()), loader.size) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to read DOCX file as ZIP: %w", err) | ||
| } | ||
|
|
||
| var text strings.Builder | ||
| for _, file := range zipReader.File { | ||
| if file.Name == "word/document.xml" { | ||
| rc, err := file.Open() | ||
| if err != nil { | ||
| return nil, fmt.Errorf("error opening document.xml: %w", err) | ||
| } | ||
| defer rc.Close() | ||
|
|
||
| content, err := io.ReadAll(rc) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("error reading content: %w", err) | ||
| } | ||
|
|
||
| content = bytes.ReplaceAll(content, []byte("<"), []byte(" <")) | ||
| content = bytes.ReplaceAll(content, []byte(">"), []byte("> ")) | ||
| text.Write(content) | ||
| } | ||
| } | ||
|
|
||
| documents := []schema.Document{ | ||
| { | ||
| PageContent: text.String(), | ||
| Metadata: map[string]interface{}{ | ||
| "fileType": loader.fileType, | ||
| }, | ||
| }, | ||
| } | ||
|
|
||
| return documents, nil | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,92 @@ | ||
| package documentloaders | ||
|
|
||
| import ( | ||
| "context" | ||
| "os" | ||
| "strings" | ||
| "testing" | ||
|
|
||
| "github.com/stretchr/testify/assert" | ||
| "github.com/stretchr/testify/require" | ||
| ) | ||
|
|
||
| func TestMSOfficeLoader(test *testing.T) { | ||
| test.Parallel() | ||
|
|
||
| docExpectedContent := "This is a .doc test file." | ||
| docxExpectedContent := "This is a .docx test file." | ||
| xlsxExpectedContent := "This is an .xlsx test file" | ||
| pptxExpectedContent := "This is a .pptx test file" | ||
|
|
||
| test.Run("Load .doc", func(t *testing.T) { | ||
| t.Parallel() | ||
|
|
||
| file, err := os.Open("./testdata/test.doc") | ||
| require.NoError(t, err) | ||
| defer file.Close() | ||
|
|
||
| fileInfo, err := file.Stat() | ||
| require.NoError(t, err) | ||
|
|
||
| loader := NewOffice(file, fileInfo.Name(), fileInfo.Size()) | ||
| docs, err := loader.Load(context.Background()) | ||
| require.NoError(t, err) | ||
|
|
||
| assert.Len(t, docs, 1) | ||
| assert.True(t, strings.Contains(docs[0].PageContent, docExpectedContent)) | ||
| }) | ||
|
|
||
| test.Run("Load .docx", func(t *testing.T) { | ||
| t.Parallel() | ||
|
|
||
| file, err := os.Open("./testdata/test.docx") | ||
| require.NoError(t, err) | ||
| defer file.Close() | ||
|
|
||
| fileInfo, err := file.Stat() | ||
| require.NoError(t, err) | ||
|
|
||
| loader := NewOffice(file, fileInfo.Name(), fileInfo.Size()) | ||
| docs, err := loader.Load(context.Background()) | ||
| require.NoError(t, err) | ||
|
|
||
| assert.Len(t, docs, 1) | ||
| assert.True(t, strings.Contains(docs[0].PageContent, docxExpectedContent)) | ||
| }) | ||
|
|
||
| test.Run("Load .xlsx", func(t *testing.T) { | ||
| t.Parallel() | ||
|
|
||
| file, err := os.Open("./testdata/test.xlsx") | ||
| require.NoError(t, err) | ||
| defer file.Close() | ||
|
|
||
| fileInfo, err := file.Stat() | ||
| require.NoError(t, err) | ||
|
|
||
| loader := NewOffice(file, fileInfo.Name(), fileInfo.Size()) | ||
| docs, err := loader.Load(context.Background()) | ||
| require.NoError(t, err) | ||
|
|
||
| assert.Len(t, docs, 1) | ||
| assert.True(t, strings.Contains(docs[0].PageContent, xlsxExpectedContent)) | ||
| }) | ||
|
|
||
| test.Run("Load .pptx", func(t *testing.T) { | ||
| t.Parallel() | ||
|
|
||
| file, err := os.Open("./testdata/test.pptx") | ||
| require.NoError(t, err) | ||
| defer file.Close() | ||
|
|
||
| fileInfo, err := file.Stat() | ||
| require.NoError(t, err) | ||
|
|
||
| loader := NewOffice(file, fileInfo.Name(), fileInfo.Size()) | ||
| docs, err := loader.Load(context.Background()) | ||
| require.NoError(t, err) | ||
|
|
||
| assert.Len(t, docs, 1) | ||
| assert.True(t, strings.Contains(docs[0].PageContent, pptxExpectedContent)) | ||
| }) | ||
| } |
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can an error be returned by doc.Next() because of an other reason then there being no more entries?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Umm that's a good question. I'll have to check and get back to you with that. Will submit fix next week ;)