Merge branch 'fix/averiewang/uploadfile-by-io' into 'main' (merge request !76)

averiewang · averiewang · commit d8618745115f · 2025-02-26T14:16:17.000Z
fix: upload file by io.Reader
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,8 @@
 # Changelog
 
+## v1.6.2
+* fix: upload file by io.Reader
+
 ## v1.6.1
 * feat: support configuring autoId for the primary key in collection, which can only be set to uuid
 * feat: support a new field indexType json, allowing for one layer of JSON Fields. JSON keys mush be string, while values ​​can be string, uint64 or array
diff --git a/example/collection_upload_file_demo/main.go b/example/collection_upload_file_demo/main.go
@@ -82,8 +82,17 @@ func (d *Demo) CreateDBAndCollection(ctx context.Context, database, collection s
 func (d *Demo) UploadFile(ctx context.Context, database, collection, localFilePath string) error {
 	appendKeywordsToChunk := false
 	appendTitleToChunk := false
+
+	// filename := filepath.Base(localFilePath)
+	// fd, err := os.Open(localFilePath)
+	// if err != nil {
+	// 	return err
+	// }
+
 	param := tcvectordb.UploadFileParams{
 		LocalFilePath: localFilePath,
+		//FileName: filename,
+		//Reader:   fd,
 		SplitterPreprocess: ai_document_set.DocumentSplitterPreprocess{
 			AppendKeywordsToChunk: &appendKeywordsToChunk,
 			AppendTitleToChunk:    &appendTitleToChunk,
@@ -160,7 +169,6 @@ func main() {
 
 	ctx := context.Background()
 	testVdb, err := NewDemo("vdb http url or ip and port", "vdb username", "key get from web console")
-	printErr(err)
 	err = testVdb.CreateDBAndCollection(ctx, database, collectionName)
 	printErr(err)
 	err = testVdb.UploadFile(ctx, database, collectionName, localFilePath)
diff --git a/tcvectordb/ai_document_sets.go b/tcvectordb/ai_document_sets.go
@@ -19,6 +19,7 @@
 package tcvectordb
 
 import (
+	"bytes"
 	"context"
 	"encoding/base64"
 	"encoding/json"
@@ -607,10 +608,11 @@ func (i *implementerAIDocumentSets) LoadAndSplitText(ctx context.Context, param
 		return nil, BaseDbTypeError
 	}
 
-	size, err := i.loadAndSplitTextCheckParams(&param)
+	size, reader, err := i.loadAndSplitTextCheckParams(&param)
 	if err != nil {
 		return nil, err
 	}
+	defer reader.Close()
 
 	res, err := i.GetCosTmpSecret(ctx, GetCosTmpSecretParams{
 		DocumentSetName: param.DocumentSetName,
@@ -666,33 +668,48 @@ func (i *implementerAIDocumentSets) LoadAndSplitText(ctx context.Context, param
 		return nil, fmt.Errorf("cos header for param MetaData is too large, it can not be more than 2k")
 	}
 
-	opt := &cos.MultiUploadOptions{
-		OptIni: &cos.InitiateMultipartUploadOptions{
-			nil,
-			&cos.ObjectPutHeaderOptions{
-				XCosMetaXXX: &header,
-				//Listener:    &cos.DefaultProgressListener{},
+	if param.LocalFilePath != "" {
+		// upload file by reading local file path, which supports multi parts uploading
+		opt := &cos.MultiUploadOptions{
+			OptIni: &cos.InitiateMultipartUploadOptions{
+				nil,
+				&cos.ObjectPutHeaderOptions{
+					XCosMetaXXX: &header,
+					//Listener:    &cos.DefaultProgressListener{},
+				},
 			},
-		},
-		// Whether to enable resume from breakpoint, default is false
-		CheckPoint: true,
-		PartSize:   5,
-	}
+			// Whether to enable resume from breakpoint, default is false
+			CheckPoint: true,
+			PartSize:   5,
+		}
 
-	_, _, err = c.Object.Upload(ctx, res.UploadPath, param.LocalFilePath, opt)
-	if err != nil {
-		return nil, err
+		_, _, err = c.Object.Upload(ctx, res.UploadPath, param.LocalFilePath, opt)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		// upload file by io.reader
+		opt := &cos.ObjectPutOptions{
+			ObjectPutHeaderOptions: &cos.ObjectPutHeaderOptions{
+				ContentLength: size,
+				XCosMetaXXX:   &header,
+			},
+		}
+		_, err = c.Object.Put(ctx, res.UploadPath, reader, opt)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	result = new(LoadAndSplitTextResult)
 	result.GetCosTmpSecretResult = *res
 	return result, nil
 }
 
-func (i *implementerAIDocumentSets) loadAndSplitTextCheckParams(param *LoadAndSplitTextParams) (size int64, err error) {
+func (i *implementerAIDocumentSets) loadAndSplitTextCheckParams(param *LoadAndSplitTextParams) (size int64, reader io.ReadCloser, err error) {
 	if param.DocumentSetName == "" {
 		if param.LocalFilePath == "" {
-			return 0, errors.New("need param: DocumentSetName")
+			return 0, nil, errors.New("need param: DocumentSetName")
 		}
 		param.DocumentSetName = filepath.Base(param.LocalFilePath)
 	}
@@ -706,17 +723,33 @@ func (i *implementerAIDocumentSets) loadAndSplitTextCheckParams(param *LoadAndSp
 			"because only markdown filetype supports defining ChunkSplitter")
 	}
 
-	fileInfo, err := os.Stat(param.LocalFilePath)
-	if err != nil {
-		return 0, errors.Errorf("get file size failed. err: %v", err.Error())
+	if param.LocalFilePath != "" {
+		fd, err := os.Open(param.LocalFilePath)
+		if err != nil {
+			return 0, nil, err
+		}
+		reader = fd
+		fstat, err := fd.Stat()
+		if err != nil {
+			return 0, nil, err
+		}
+		size = fstat.Size()
+	} else {
+		bytesBuf := bytes.NewBuffer(nil)
+		written, err := io.Copy(bytesBuf, param.Reader)
+		if err != nil {
+			return 0, nil, err
+		}
+
+		size = written
+		reader = io.NopCloser(bytesBuf)
 	}
-	size = fileInfo.Size()
 
 	if size == 0 {
-		return 0, errors.New("file size cannot be 0")
+		return 0, nil, errors.New("file size cannot be 0")
 	}
 
-	return size, nil
+	return size, reader, nil
 }
 
 func (i *implementerAIDocumentSets) toDocumentSet(item ai_document_set.QueryDocumentSet) *AIDocumentSet {
diff --git a/tcvectordb/base_flat.go b/tcvectordb/base_flat.go
@@ -19,10 +19,12 @@
 package tcvectordb
 
 import (
+	"bytes"
 	"context"
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
+	"io"
 	"log"
 	"net/http"
 	"net/url"
@@ -91,7 +93,6 @@ type FlatInterface interface {
 	// [ChangePassword] changes the password for the specific user.
 	ChangePassword(ctx context.Context, param ChangePasswordParams) error
 
-	// []
 	UploadFile(ctx context.Context, databaseName, collectionName string, param UploadFileParams) (result *UploadFileResult, err error)
 
 	GetImageUrl(ctx context.Context, databaseName, collectionName string,
@@ -334,6 +335,7 @@ func (i *implementerFlatDocument) ChangePassword(ctx context.Context, param Chan
 type UploadFileParams struct {
 	FileName           string
 	LocalFilePath      string
+	Reader             io.Reader
 	SplitterPreprocess ai_document_set.DocumentSplitterPreprocess
 	EmbeddingModel     string
 	ParsingProcess     *api.ParsingProcess
@@ -355,10 +357,10 @@ func (i *implementerFlatDocument) UploadFile(ctx context.Context, databaseName,
 	return uploadFile(ctx, i, databaseName, collectionName, param)
 }
 
-func checkUploadFileParam(ctx context.Context, param *UploadFileParams) (size int64, err error) {
+func checkUploadFileParam(ctx context.Context, param *UploadFileParams) (size int64, reader io.ReadCloser, err error) {
 	if param.FileName == "" {
 		if param.LocalFilePath == "" {
-			return 0, errors.New("need param: FileName or LocalFilePath")
+			return 0, nil, errors.New("need param: FileName or LocalFilePath")
 		}
 		param.FileName = filepath.Base(param.LocalFilePath)
 	}
@@ -371,24 +373,41 @@ func checkUploadFileParam(ctx context.Context, param *UploadFileParams) (size in
 		log.Printf("[Warning] %s", "param SplitterPreprocess.ChunkSplitter will be ommitted, "+
 			"because only markdown filetype supports defining ChunkSplitter")
 	}
+	if param.LocalFilePath != "" {
+		fd, err := os.Open(param.LocalFilePath)
+		if err != nil {
+			return 0, nil, err
+		}
+		reader = fd
+		fstat, err := fd.Stat()
+		if err != nil {
+			return 0, nil, err
+		}
+		size = fstat.Size()
+	} else {
+		bytesBuf := bytes.NewBuffer(nil)
+		written, err := io.Copy(bytesBuf, param.Reader)
+		if err != nil {
+			return 0, nil, err
+		}
 
-	fileInfo, err := os.Stat(param.LocalFilePath)
-	if err != nil {
-		return 0, errors.Errorf("get file size failed. err: %v", err.Error())
+		size = written
+		reader = io.NopCloser(bytesBuf)
 	}
-	size = fileInfo.Size()
 
 	if size == 0 {
-		return 0, errors.New("file size cannot be 0")
+		return 0, nil, errors.New("file size cannot be 0")
 	}
-	return size, nil
+
+	return size, reader, nil
 }
 func uploadFile(ctx context.Context, cli SdkClient, databaseName, collectionName string,
 	param UploadFileParams) (result *UploadFileResult, err error) {
-	size, err := checkUploadFileParam(ctx, &param)
+	size, reader, err := checkUploadFileParam(ctx, &param)
 	if err != nil {
 		return nil, err
 	}
+	defer reader.Close()
 
 	req := new(document.UploadUrlReq)
 	req.Database = databaseName
@@ -457,22 +476,36 @@ func uploadFile(ctx context.Context, cli SdkClient, databaseName, collectionName
 		return nil, fmt.Errorf("cos header for param MetaData is too large, it can not be more than 2k")
 	}
 
-	opt := &cos.MultiUploadOptions{
-		OptIni: &cos.InitiateMultipartUploadOptions{
-			nil,
-			&cos.ObjectPutHeaderOptions{
-				XCosMetaXXX: &header,
-				//Listener:    &cos.DefaultProgressListener{},
+	if param.LocalFilePath != "" {
+		opt := &cos.MultiUploadOptions{
+			OptIni: &cos.InitiateMultipartUploadOptions{
+				nil,
+				&cos.ObjectPutHeaderOptions{
+					XCosMetaXXX: &header,
+					//Listener:    &cos.DefaultProgressListener{},
+				},
 			},
-		},
-		// Whether to enable resume from breakpoint, default is false
-		CheckPoint: true,
-		PartSize:   5,
-	}
+			// Whether to enable resume from breakpoint, default is false
+			CheckPoint: true,
+			PartSize:   5,
+		}
 
-	_, _, err = c.Object.Upload(ctx, res.UploadPath, param.LocalFilePath, opt)
-	if err != nil {
-		return nil, err
+		_, _, err = c.Object.Upload(ctx, res.UploadPath, param.LocalFilePath, opt)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+
+		opt := &cos.ObjectPutOptions{
+			ObjectPutHeaderOptions: &cos.ObjectPutHeaderOptions{
+				ContentLength: size,
+				XCosMetaXXX:   &header,
+			},
+		}
+		_, err = c.Object.Put(ctx, res.UploadPath, reader, opt)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	return result, nil
diff --git a/tcvectordb/version.go b/tcvectordb/version.go
@@ -18,4 +18,4 @@
 
 package tcvectordb
 
-const SDKVersion = "v1.6.1"
+const SDKVersion = "v1.6.2"
diff --git a/test/aidb_test.go b/test/aidb_test.go
@@ -22,6 +22,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"log"
+	"os"
 	"strings"
 	"testing"
 	"time"
@@ -171,6 +172,41 @@ func TestLoadAndSplitText(t *testing.T) {
 	t.Logf("%+v", result)
 }
 
+func TestLoadAndSplitTextByReader(t *testing.T) {
+	defer cli.Close()
+
+	col := cli.AIDatabase(aiDatabase).CollectionView(collectionViewName)
+
+	metaData := map[string]interface{}{
+		// 元数据只支持string、uint64类型的值
+		"author_name": "sam",
+		"fileKey":     1024}
+
+	appendTitleToChunk := false
+	appendKeywordsToChunk := true
+	chunkSplitter := "\n\n"
+
+	fd, err := os.Open("../example/demo_files/tcvdb.md")
+	printErr(err)
+
+	result, err := col.LoadAndSplitText(ctx, tcvectordb.LoadAndSplitTextParams{
+		DocumentSetName: "tcvdb.md",
+		Reader:          fd,
+		//LocalFilePath: "../example/demo_files/tcvdb.md",
+		MetaData: metaData,
+		SplitterPreprocess: ai_document_set.DocumentSplitterPreprocess{
+			ChunkSplitter:         &chunkSplitter,
+			AppendTitleToChunk:    &appendTitleToChunk,
+			AppendKeywordsToChunk: &appendKeywordsToChunk,
+		},
+		ParsingProcess: &api.ParsingProcess{
+			ParsingType: string(tcvectordb.VisionModelParsing),
+		},
+	})
+	printErr(err)
+	t.Logf("%+v", result)
+}
+
 func TestAIGetDocumentSet(t *testing.T) {
 	time.Sleep(10 * time.Second)
 	col := cli.AIDatabase(aiDatabase).CollectionView(collectionViewName)

Original file line number	Diff line number	Diff line change
`@@ -18,4 +18,4 @@`
`18`	`18`
`19`	`19`	`package tcvectordb`
`20`	`20`
`21`		`-const SDKVersion = "v1.6.1"`
	`21`	`+const SDKVersion = "v1.6.2"`