file.go

// file.go
package knowledgesdk

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"gorm.io/gorm"
	"io"
	"mime/multipart"
	"net/http"
	"strings"
)

// TikaConfig represents configuration for Apache Tika service
type TikaConfig struct {
	URL string // Tika server URL, e.g., "http://localhost:9998"
}

// DefaultTikaConfig returns default Tika configuration
func DefaultTikaConfig() TikaConfig {
	return TikaConfig{
		URL: "http://localhost:9998", // Default Tika server URL
	}
}

// FileContent represents the content extracted from a file
type FileContent struct {
	Content     string      `json:"content"`
	Metadata    MetadataMap `json:"metadata"`
	ContentType string      `json:"contentType"`
}

// AddFile 添加文件到知识库
func (k *KnowledgeSDK) AddFile(
	ctx context.Context,
	kbID string,
	fileName string,
	fileData []byte,
	tikaConfig TikaConfig,
	chunkConfig ChunkConfig,
	creatorID string,
) (*Document, error) {
	// 创建带有初始状态的文档
	doc := NewDocument(kbID, fileName, "", creatorID)
	doc.Status = String(DocStatusUploadSuccess) // 上传成功，等待抽取
	doc.Metadata = make(map[string]string)

	// 创建文档记录
	if err := k.db.WithContext(ctx).Create(doc).Error; err != nil {
		return nil, fmt.Errorf("创建文档记录失败: %w", err)
	}

	// 尝试提取文件内容
	fileContent, err := k.ExtractFileContent(ctx, fileName, fileData, tikaConfig)
	if err != nil {
		// 更新状态为抽取失败
		k.UpdateDocumentStatus(ctx, StringValue(doc.ID), DocStatusExtractFailed)
		return doc, fmt.Errorf("提取文件内容失败: %w", err)
	}

	// 提取成功，更新文档内容和状态
	err = k.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		// 更新文档内容和元数据
		doc.OriginalContent = String(fileContent.Content)
		doc.ContentType = String(fileContent.ContentType)
		doc.Metadata = fileContent.Metadata
		doc.Status = String(DocStatusExtractSuccess) // 更新为抽取成功状态

		if err := tx.Save(doc).Error; err != nil {
			return fmt.Errorf("更新文档内容失败: %w", err)
		}

		// 进行文档分块
		if err := k.splitDocumentIntoChunks(ctx, tx, StringValue(doc.ID), fileContent.Content, chunkConfig); err != nil {
			// 如果分块失败，更新状态
			doc.Status = String(DocStatusSplitFailed)
			tx.Save(doc)
			return fmt.Errorf("文档分块失败: %w", err)
		}

		// 更新文档状态为等待索引
		doc.Status = String(DocStatusSplitSuccess)
		if err := tx.Save(doc).Error; err != nil {
			return fmt.Errorf("更新文档状态失败: %w", err)
		}

		return nil
	})

	if err != nil {
		return doc, err
	}

	return doc, nil
}

// AddFileFromReader 从读取器添加文件到知识库
func (k *KnowledgeSDK) AddFileFromReader(
	ctx context.Context,
	kbID string,
	fileName string,
	reader io.Reader,
	tikaConfig TikaConfig,
	chunkConfig ChunkConfig,
	creatorID string,
) (*Document, error) {
	// 读取文件数据
	fileData, err := io.ReadAll(reader)
	if err != nil {
		return nil, fmt.Errorf("读取文件数据失败: %w", err)
	}

	return k.AddFile(ctx, kbID, fileName, fileData, tikaConfig, chunkConfig, creatorID)
}

// AddFileFromMultipart 从HTTP multipart上传添加文件到知识库
func (k *KnowledgeSDK) AddFileFromMultipart(
	ctx context.Context,
	kbID string,
	file *multipart.FileHeader,
	tikaConfig TikaConfig,
	chunkConfig ChunkConfig,
	creatorID string,
) (*Document, error) {
	// 打开上传的文件
	src, err := file.Open()
	if err != nil {
		return nil, fmt.Errorf("打开上传文件失败: %w", err)
	}
	defer src.Close()

	return k.AddFileFromReader(ctx, kbID, file.Filename, src, tikaConfig, chunkConfig, creatorID)
}

// ExtractFileContent 使用Apache Tika提取文件内容和元数据
func (k *KnowledgeSDK) ExtractFileContent(
	ctx context.Context,
	fileName string,
	fileData []byte,
	tikaConfig TikaConfig,
) (*FileContent, error) {
	// 准备请求Tika
	url := fmt.Sprintf("%s/tika", tikaConfig.URL)
	req, err := http.NewRequestWithContext(ctx, http.MethodPut, url, bytes.NewReader(fileData))
	if err != nil {
		return nil, fmt.Errorf("创建Tika请求失败: %w", err)
	}

	// 设置请求头，要求JSON格式以同时获取内容和元数据
	req.Header.Set("Accept", "application/json; charset=utf-8")

	// 发送请求到Tika
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("调用Tika服务失败: %w", err)
	}
	defer resp.Body.Close()

	// 检查响应
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("tika服务返回错误: %s, %s", resp.Status, string(body))
	}

	// 解析JSON响应
	var jsonResponse map[string]interface{}
	if err := json.NewDecoder(resp.Body).Decode(&jsonResponse); err != nil {
		return nil, fmt.Errorf("解析Tika响应失败: %w", err)
	}

	// 从JSON中提取正文内容
	content, ok := jsonResponse["X-TIKA:content"].(string)
	if !ok {
		return nil, fmt.Errorf("tika响应中未找到内容")
	}
	//fmt.Printf("提取的内容:%+v", jsonResponse)
	// 转换元数据为 MetadataMap
	metadata := make(MetadataMap)
	for k, v := range jsonResponse {
		if strings.HasPrefix(k, "X-TIKA:") {
			continue
		}
		if str, ok := v.(string); ok {
			metadata[k] = str
		} else {
			// 将非字符串值转换为JSON字符串
			jsonBytes, err := json.Marshal(v)
			if err == nil {
				metadata[k] = string(jsonBytes)
			}
		}
	}

	return &FileContent{
		Content:     content,
		Metadata:    metadata,
		ContentType: metadata["Content-Type"],
	}, nil
}

// ExtractFileContentFromReader 从io.Reader中提取文件内容
func (k *KnowledgeSDK) ExtractFileContentFromReader(
	ctx context.Context,
	fileName string,
	reader io.Reader,
	tikaConfig TikaConfig,
) (*FileContent, error) {
	// 从读取器中读取所有数据
	fileData, err := io.ReadAll(reader)
	if err != nil {
		return nil, fmt.Errorf("读取文件数据失败: %w", err)
	}

	// 使用现有方法提取内容
	return k.ExtractFileContent(ctx, fileName, fileData, tikaConfig)
}

// ExtractFileContentFromMultipart 从multipart.FileHeader中提取文件内容
func (k *KnowledgeSDK) ExtractFileContentFromMultipart(
	ctx context.Context,
	file *multipart.FileHeader,
	tikaConfig TikaConfig,
) (*FileContent, error) {
	// 打开上传的文件
	src, err := file.Open()
	if err != nil {
		return nil, fmt.Errorf("打开上传文件失败: %w", err)
	}
	defer src.Close()

	// 使用从reader提取的方法
	return k.ExtractFileContentFromReader(ctx, file.Filename, src, tikaConfig)
}

// ExtractFileContentFromURL 从URL提取文件内容
func (k *KnowledgeSDK) ExtractFileContentFromURL(
	ctx context.Context,
	fileURL string,
	tikaConfig TikaConfig,
) (*FileContent, error) {
	// 创建HTTP请求
	req, err := http.NewRequestWithContext(ctx, "GET", fileURL, nil)
	if err != nil {
		return nil, fmt.Errorf("创建HTTP请求失败: %w", err)
	}

	// 发送请求
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("下载文件失败: %w", err)
	}
	defer resp.Body.Close()

	// 检查响应
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("下载文件失败，状态码: %d", resp.StatusCode)
	}

	// 获取文件名
	fileName := ""
	// 尝试从Content-Disposition头获取文件名
	contentDisposition := resp.Header.Get("Content-Disposition")
	if contentDisposition != "" {
		if strings.Contains(contentDisposition, "filename=") {
			parts := strings.Split(contentDisposition, "filename=")
			if len(parts) > 1 {
				fileName = strings.Trim(parts[1], "\"' ")
			}
		}
	}

	// 如果无法从头信息获取文件名，则从URL中提取
	if fileName == "" {
		parts := strings.Split(fileURL, "/")
		if len(parts) > 0 {
			fileName = parts[len(parts)-1]
			// 处理URL参数
			fileName = strings.Split(fileName, "?")[0]
		}
	}

	// 如果仍然无法获取文件名，使用一个默认名称
	if fileName == "" {
		fileName = "downloaded_file"
	}

	// 使用从reader提取的方法
	return k.ExtractFileContentFromReader(ctx, fileName, resp.Body, tikaConfig)
}

// GetDocumentMetadata 获取文档元数据
func (k *KnowledgeSDK) GetDocumentMetadata(ctx context.Context, docID string) (MetadataMap, error) {
	// 首先检查文档是否存在
	var doc Document
	if err := k.db.WithContext(ctx).Where("document_id = ?", docID).First(&doc).Error; err != nil {
		return nil, fmt.Errorf("获取文档失败: %w", err)
	}

	return doc.Metadata, nil
}

// GetFileMetadata 提取文件的元数据，不保存文件
func (k *KnowledgeSDK) GetFileMetadata(
	ctx context.Context,
	fileName string,
	fileData []byte,
	tikaConfig TikaConfig,
) (MetadataMap, error) {
	// 准备请求Tika
	url := fmt.Sprintf("%s/tika/text", tikaConfig.URL)
	req, err := http.NewRequestWithContext(ctx, "PUT", url, bytes.NewReader(fileData))
	if err != nil {
		return nil, fmt.Errorf("创建Tika请求失败: %w", err)
	}

	// 设置请求头
	req.Header.Set("Accept", "application/json; charset=utf-8")

	// 发送请求到Tika
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("调用Tika服务失败: %w", err)
	}
	defer resp.Body.Close()

	// 检查响应
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("tika服务返回错误: %s, %s", resp.Status, string(body))
	}

	// 解析JSON响应
	var jsonResponse map[string]interface{}
	if err := json.NewDecoder(resp.Body).Decode(&jsonResponse); err != nil {
		return nil, fmt.Errorf("解析Tika响应失败: %w", err)
	}

	// 转换元数据为 MetadataMap
	metadata := make(MetadataMap)
	for k, v := range jsonResponse {
		if strings.HasPrefix(k, "X-TIKA:") {
			continue
		}
		if str, ok := v.(string); ok {
			metadata[k] = str
		} else {
			// 将非字符串值转换为JSON字符串
			jsonBytes, err := json.Marshal(v)
			if err == nil {
				metadata[k] = string(jsonBytes)
			}
		}
	}

	return metadata, nil
}

// GetFileMetadataFromReader 从io.Reader中提取文件的元数据
func (k *KnowledgeSDK) GetFileMetadataFromReader(
	ctx context.Context,
	fileName string,
	reader io.Reader,
	tikaConfig TikaConfig,
) (MetadataMap, error) {
	// 读取文件数据
	fileData, err := io.ReadAll(reader)
	if err != nil {
		return nil, fmt.Errorf("读取文件数据失败: %w", err)
	}

	return k.GetFileMetadata(ctx, fileName, fileData, tikaConfig)
}

// GetFileMetadataFromMultipart 从multipart.FileHeader中提取文件的元数据
func (k *KnowledgeSDK) GetFileMetadataFromMultipart(
	ctx context.Context,
	file *multipart.FileHeader,
	tikaConfig TikaConfig,
) (MetadataMap, error) {
	// 打开上传的文件
	src, err := file.Open()
	if err != nil {
		return nil, fmt.Errorf("打开上传文件失败: %w", err)
	}
	defer src.Close()

	return k.GetFileMetadataFromReader(ctx, file.Filename, src, tikaConfig)
}

// AddFileWithoutCreator 不带creatorID的兼容方法
func (k *KnowledgeSDK) AddFileWithoutCreator(
	ctx context.Context,
	kbID string,
	fileName string,
	fileData []byte,
	tikaConfig TikaConfig,
	chunkConfig ChunkConfig,
) (*Document, error) {
	return k.AddFile(ctx, kbID, fileName, fileData, tikaConfig, chunkConfig, "")
}

// AddFileFromReaderWithoutCreator 不带creatorID的兼容方法
func (k *KnowledgeSDK) AddFileFromReaderWithoutCreator(
	ctx context.Context,
	kbID string,
	fileName string,
	reader io.Reader,
	tikaConfig TikaConfig,
	chunkConfig ChunkConfig,
) (*Document, error) {
	return k.AddFileFromReader(ctx, kbID, fileName, reader, tikaConfig, chunkConfig, "")
}

// AddFileFromMultipartWithoutCreator 不带creatorID的兼容方法
func (k *KnowledgeSDK) AddFileFromMultipartWithoutCreator(
	ctx context.Context,
	kbID string,
	file *multipart.FileHeader,
	tikaConfig TikaConfig,
	chunkConfig ChunkConfig,
) (*Document, error) {
	return k.AddFileFromMultipart(ctx, kbID, file, tikaConfig, chunkConfig, "")
}