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)
	}

}