AI-CS/backend/infra/milvus.go

package infra

import (
	"context"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/milvus-io/milvus-sdk-go/v2/client"
)

// MilvusConfig Milvus 连接配置
type MilvusConfig struct {
	Host string
	Port string
}

// GetMilvusConfig 从环境变量读取 Milvus 配置
func GetMilvusConfig() *MilvusConfig {
	host := os.Getenv("MILVUS_HOST")
	if host == "" {
		host = "localhost" // 默认值
	}
	port := os.Getenv("MILVUS_PORT")
	if port == "" {
		port = "19530" // 默认端口
	}
	return &MilvusConfig{
		Host: host,
		Port: port,
	}
}

func envTruthy(key string) bool {
	v := strings.TrimSpace(strings.ToLower(os.Getenv(key)))
	return v == "1" || v == "true" || v == "yes" || v == "on"
}

// IsMilvusDisabled 为 true 时跳过连接 Milvus（不参与 RAG / 向量化，核心 HTTP 仍可启动）。
// 环境变量：MILVUS_DISABLED 或 VECTOR_STORE_DISABLED。
func IsMilvusDisabled() bool {
	return envTruthy("MILVUS_DISABLED") || envTruthy("VECTOR_STORE_DISABLED")
}

// IsMilvusRequired 为 true 时，若 Milvus 不可用则启动失败（便于生产环境强依赖向量库时 fail-fast）。
// 环境变量：MILVUS_REQUIRED。
func IsMilvusRequired() bool {
	return envTruthy("MILVUS_REQUIRED")
}

// NewMilvusClient 创建 Milvus 客户端连接
func NewMilvusClient() (client.Client, error) {
	config := GetMilvusConfig()
	// 构建连接地址
	address := fmt.Sprintf("%s:%s", config.Host, config.Port)
	// 创建客户端
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	milvusClient, err := client.NewClient(
		ctx,
		client.Config{
			Address:  address,
			Username: os.Getenv("MILVUS_USERNAME"), // 可选
			Password: os.Getenv("MILVUS_PASSWORD"), // 可选
		},
	)
	if err != nil {
		return nil, fmt.Errorf("连接 Milvus 失败: %w", err)
	}
	return milvusClient, nil
}

// HealthCheck 检查 Milvus 连接健康状态
func HealthCheck(milvusClient client.Client) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	// 尝试列出集合来验证连接
	_, err := milvusClient.ListCollections(ctx)
	if err != nil {
		return fmt.Errorf("Milvus 健康检查失败: %w", err)
	}
	return nil
}