""" Milvus 客户端封装 """ import logging from typing import List, Dict, Optional, Any import numpy as np from pymilvus import MilvusClient, FieldSchema, CollectionSchema, DataType, connections, Collection from app.core.config import MILVUS_URL, MILVUS_TOKEN, MILVUS_ALIAS from app.service.recommendation_system.config import MILVUS_COLLECTION_SKETCH_VECTORS, RECOMMENDATION_CONFIG logger = logging.getLogger(__name__) # Milvus 客户端（单例） _milvus_client = None def get_milvus_client() -> MilvusClient: """获取 Milvus 客户端（单例模式）""" global _milvus_client if _milvus_client is None: try: _milvus_client = MilvusClient( uri=MILVUS_URL, token=MILVUS_TOKEN, db_name=MILVUS_ALIAS ) logger.info("Milvus 客户端连接成功") except Exception as e: logger.error(f"Milvus 客户端连接失败: {e}") raise return _milvus_client def create_collection(): """ 创建 Milvus 集合 sketch_vectors 集合结构： - path (PK, varchar(512)) - 主键，MinIO 逻辑 URL - sys_file_id (int64, 可为NULL) - 系统文件ID - style (varchar(50), 可为NULL) - 风格样式 - category (varchar(100), 可为NULL) - 类别 - is_system_sketch (int8, 默认 1) - 标记字段：1-系统图，0-用户图 - deprecated (int8, 默认 0) - 是否废弃 - feature_vector (FloatVector(2048)) - 2048维特征向量 """ client = get_milvus_client() # 检查集合是否已存在 collections = client.list_collections() if MILVUS_COLLECTION_SKETCH_VECTORS in collections: logger.info(f"集合 {MILVUS_COLLECTION_SKETCH_VECTORS} 已存在") return try: # 解析 Milvus URL # 处理 http://host.docker.internal:19530 格式 url_clean = MILVUS_URL.replace("http://", "").replace("https://", "") if ":" in url_clean: host, port_str = url_clean.split(":", 1) port = int(port_str) else: host = url_clean port = 19530 # 使用传统 API 创建集合（更可靠） # 连接到 Milvus（如果未连接） try: connections.connect( alias=MILVUS_ALIAS, host=host, port=port, token=MILVUS_TOKEN if MILVUS_TOKEN else None ) logger.info(f"已连接到 Milvus: {host}:{port}") except Exception as conn_e: # 如果连接已存在，忽略错误 if "already exists" in str(conn_e).lower() or "Connection already exists" in str(conn_e): logger.info("Milvus 连接已存在") else: logger.warning(f"连接 Milvus 时出现警告: {conn_e}") # 定义字段 fields = [ FieldSchema(name="path", dtype=DataType.VARCHAR, is_primary=True, max_length=512), FieldSchema(name="sys_file_id", dtype=DataType.INT64), FieldSchema(name="style", dtype=DataType.VARCHAR, max_length=50), FieldSchema(name="category", dtype=DataType.VARCHAR, max_length=50), FieldSchema(name="is_system_sketch", dtype=DataType.INT8), FieldSchema(name="deprecated", dtype=DataType.INT8), FieldSchema( name="feature_vector", dtype=DataType.FLOAT_VECTOR, dim=RECOMMENDATION_CONFIG["vector_dim"] ) ] # 创建 schema schema = CollectionSchema( fields=fields, description="Sketch vectors collection for recommendation system" ) # 创建集合 collection = Collection( name=MILVUS_COLLECTION_SKETCH_VECTORS, schema=schema, using=MILVUS_ALIAS ) # 创建索引 # 注意：使用 IP（内积）作为度量类型，与搜索时保持一致 # 如果向量已归一化，IP 等价于 COSINE index_params = { "metric_type": "IP", # 内积（Inner Product） "index_type": "IVF_FLAT", "params": {"nlist": 1024} } collection.create_index( field_name="feature_vector", index_params=index_params ) logger.info(f"集合 {MILVUS_COLLECTION_SKETCH_VECTORS} 创建成功") except Exception as e: logger.error(f"创建集合失败: {e}", exc_info=True) raise def insert_vectors(data: List[Dict[str, Any]]): """ 批量插入向量到 Milvus Args: data: 数据列表，每个元素包含： - path: str - sys_file_id: int (可选) - style: str (可选) - category: str (可选) - is_system_sketch: int (默认 1) - deprecated: int (默认 0) - feature_vector: List[float] (2048维) """ if not data: return client = get_milvus_client() try: client.insert( collection_name=MILVUS_COLLECTION_SKETCH_VECTORS, data=data ) logger.info(f"成功插入 {len(data)} 条向量数据") except Exception as e: logger.error(f"插入向量失败: {e}", exc_info=True) raise def query_vectors_by_paths(paths: List[str]) -> Dict[str, Dict]: """ 根据 path 列表批量查询向量 Args: paths: path 列表 Returns: {path: {feature_vector: [...], ...}} 字典 """ if not paths: return {} client = get_milvus_client() try: # 构建查询表达式 # 使用 filter 参数而不是 expr（根据 pymilvus MilvusClient API） # 对于字符串列表，使用单引号包裹每个值 path_list = ", ".join([f"'{p}'" for p in paths]) filter_expr = f"path in [{path_list}]" results = client.query( collection_name=MILVUS_COLLECTION_SKETCH_VECTORS, filter=filter_expr, output_fields=["path", "feature_vector", "style", "category", "sys_file_id", "is_system_sketch", "deprecated"] ) # 转换为字典 result_dict = {} for r in results: result_dict[r["path"]] = r return result_dict except Exception as e: logger.error(f"查询向量失败: {e}", exc_info=True) return {} def search_similar_vectors( query_vector: np.ndarray, category: str, topk: int = 500, style: Optional[str] = None ) -> List[Dict]: """ 向量相似度检索 Args: query_vector: 查询向量（2048维） category: 类别过滤 topk: 返回数量 style: 风格过滤（可选） Returns: 检索结果列表，每个元素包含 path, score, style, category 等字段 """ client = get_milvus_client() try: # 构建过滤表达式 # 使用 filter 参数而不是 expr（根据 pymilvus MilvusClient API） filter_expr = f"category == '{category}' && deprecated == 0" if style: filter_expr += f" && style == '{style}'" # 搜索 results = client.search( collection_name=MILVUS_COLLECTION_SKETCH_VECTORS, data=[query_vector.tolist()], anns_field="feature_vector", search_params={"metric_type": "IP", "params": {"nprobe": 10}}, limit=topk, filter=filter_expr, output_fields=["path", "style", "category", "sys_file_id"] ) # 格式化结果 formatted_results = [] if results and len(results) > 0: for hit in results[0]: formatted_results.append({ "path": hit.get("entity", {}).get("path", ""), "score": hit.get("distance", 0.0), "style": hit.get("entity", {}).get("style", ""), "category": hit.get("entity", {}).get("category", ""), "sys_file_id": hit.get("entity", {}).get("sys_file_id") }) return formatted_results except Exception as e: logger.error(f"向量检索失败: {e}", exc_info=True) return [] def query_random_candidates(category: str, style: Optional[str] = None, limit: int = 10) -> List[Dict]: """ 随机查询候选（用于探索分支） Args: category: 类别 style: 风格（可选） limit: 返回数量 Returns: 候选列表 """ client = get_milvus_client() try: # 构建过滤表达式 filter_expr = f"category == '{category}' && deprecated == 0" if style: filter_expr += f" && style == '{style}'" # 查询所有符合条件的记录 results = client.query( collection_name=MILVUS_COLLECTION_SKETCH_VECTORS, filter=filter_expr, output_fields=["path", "style", "category"], limit=10000 # 先查询大量数据，然后随机选择 ) # 随机选择 if len(results) > limit: import random results = random.sample(results, limit) return results except Exception as e: logger.error(f"随机查询候选失败: {e}", exc_info=True) return []