ai-cs/backend/service/rag/vector_store.go

package rag

import (
	"context"
	"errors"
	"fmt"
	"strconv"

	"github.com/2930134478/AI-CS/backend/infra"
)

// ErrVectorStoreUnavailable 向量库未启用或未连接（写入/索引前会返回该错误）。
var ErrVectorStoreUnavailable = errors.New("向量数据库未启用或未连接")

// VectorStoreService 向量存储服务（业务层）
type VectorStoreService struct {
	vectorStore *infra.VectorStore
}

// NewVectorStoreService 创建向量存储服务实例（vectorStore 可为 nil，表示无向量库降级模式）。
func NewVectorStoreService(vectorStore *infra.VectorStore) *VectorStoreService {
	return &VectorStoreService{
		vectorStore: vectorStore,
	}
}

// IsAvailable 当前是否已连接可用的 Milvus 向量存储。
func (s *VectorStoreService) IsAvailable() bool {
	return s != nil && s.vectorStore != nil
}

// UpsertVector 插入或更新单个向量
func (s *VectorStoreService) UpsertVector(ctx context.Context, documentID string, knowledgeBaseID string, content string, chunkDBID string, vector []float32) error {
	if s.vectorStore == nil {
		return ErrVectorStoreUnavailable
	}
	return s.vectorStore.UpsertVector(ctx, documentID, knowledgeBaseID, content, chunkDBID, vector)
}

// UpsertVectors 批量插入或更新向量
func (s *VectorStoreService) UpsertVectors(ctx context.Context, documentIDs []string, knowledgeBaseIDs []string, contents []string, vectors [][]float32, chunkDBIDs []string) error {
	if s.vectorStore == nil {
		return ErrVectorStoreUnavailable
	}
	return s.vectorStore.UpsertVectors(ctx, documentIDs, knowledgeBaseIDs, contents, vectors, chunkDBIDs)
}

// SearchVectors 搜索相似向量
func (s *VectorStoreService) SearchVectors(ctx context.Context, queryVector []float32, topK int, knowledgeBaseID *string) ([]SearchResult, error) {
	if s.vectorStore == nil {
		return []SearchResult{}, nil
	}
	results, err := s.vectorStore.SearchVectors(ctx, queryVector, topK, knowledgeBaseID)
	if err != nil {
		return nil, fmt.Errorf("向量检索失败: %w", err)
	}

	searchResults := make([]SearchResult, len(results))
	for i, r := range results {
		searchResults[i] = SearchResult{
			DocumentID:      r.DocumentID,
			KnowledgeBaseID: r.KnowledgeBaseID,
			Content:         r.Content,
			Score:           r.Score,
		}
	}

	return searchResults, nil
}

// DeleteVector 删除向量
func (s *VectorStoreService) DeleteVector(ctx context.Context, documentID string) error {
	if s.vectorStore == nil {
		return nil
	}
	return s.vectorStore.DeleteVector(ctx, documentID)
}

// DeleteVectors 批量删除向量
func (s *VectorStoreService) DeleteVectors(ctx context.Context, documentIDs []string) error {
	if s.vectorStore == nil {
		return nil
	}
	return s.vectorStore.DeleteVectors(ctx, documentIDs)
}

// DeleteVectorByChunkID 按 chunk_db_id 删除单条向量
func (s *VectorStoreService) DeleteVectorByChunkID(ctx context.Context, chunkDBID string) error {
	if s.vectorStore == nil {
		return nil
	}
	return s.vectorStore.DeleteVectorByChunkID(ctx, chunkDBID)
}

// ConvertDocumentID 将 uint 转换为 string
func ConvertDocumentID(id uint) string {
	return strconv.FormatUint(uint64(id), 10)
}

// ConvertKnowledgeBaseID 将 uint 转换为 string
func ConvertKnowledgeBaseID(id uint) string {
	return strconv.FormatUint(uint64(id), 10)
}