62 lines
1.5 KiB
Go
62 lines
1.5 KiB
Go
package main
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"github.com/joho/godotenv"
|
|
"github.com/satori/go.uuid"
|
|
"io/ioutil"
|
|
"knowledge/utils"
|
|
"log"
|
|
"os"
|
|
"path/filepath"
|
|
)
|
|
|
|
func main() {
|
|
err := godotenv.Load(".env")
|
|
// 读取环境变量
|
|
OPENAI_API_KEY := os.Getenv("OPENAI_KEY")
|
|
OPENAI_API_BASE := os.Getenv("OPENAI_API_BASE")
|
|
utils.QdrantBase = os.Getenv("QDRANT_BASE")
|
|
utils.QdrantPort = os.Getenv("QDRANT_PORT")
|
|
|
|
gpt := utils.NewChatGptTool(OPENAI_API_BASE, OPENAI_API_KEY)
|
|
|
|
// 定义要遍历的目录路径
|
|
dir := "./"
|
|
|
|
// 调用 Walk 函数遍历目录下的所有文件
|
|
err = filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
|
|
// 判断文件是否为 .txt 文件
|
|
if filepath.Ext(path) == ".txt" {
|
|
// 读取文件内容
|
|
data, err := ioutil.ReadFile(path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
// 打印文件内容
|
|
fileData := string(data)
|
|
response, err := gpt.GetEmbedding(fileData, "text-embedding-ada-002")
|
|
var embeddingResponse utils.EmbeddingResponse
|
|
json.Unmarshal([]byte(response), &embeddingResponse)
|
|
|
|
collectionName := "data_collection"
|
|
id := uuid.NewV4().String()
|
|
points := []map[string]interface{}{
|
|
{
|
|
"id": id,
|
|
"payload": map[string]interface{}{"text": fileData},
|
|
"vector": embeddingResponse.Data[0].Embedding,
|
|
},
|
|
}
|
|
res, err := utils.PutPoints(collectionName, points)
|
|
log.Println(res, err)
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
fmt.Println("Error:", err)
|
|
}
|
|
|
|
}
|