SmartBizTrip/sql/init_meeting_package_rag.py

# 【数据准备 D：独立脚本】在会务套餐检索前准备 Milvus 集合；不在每轮聊天中执行。
# 运行顺序：D1 读取配置和样例数据 → D2 删除重建集合及索引
# → D3 为套餐生成向量并入库 → D4 用样例查询验证检索。
# 注意 D2 会删除同名集合及其已有数据；本文件不是增量更新脚本。

"""
需求：初始化会务套餐 RAG 系统
思路：
1. 定义会务套餐数据（城市、线路名称、天数、价格、方案亮点等）
2. 使用 Qwen 的 Embedding API 将每个套餐的方案亮点生成向量
3. 在本地 Milvus 中创建 Collection 并插入数据
4. 支持后续通过语义搜索匹配用户意图的会务套餐
"""

import json
import os
import sys
import requests
from pymilvus import connections, Collection, CollectionSchema, FieldSchema, DataType, utility

# 将项目根目录加入 Python 路径，以便导入配置
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

# ==================== Qwen Embedding API ====================
# 使用 DashScope 的 text-embedding-v3 模型生成向量
# 返回 1024 维向量
from config import Config
# 【数据准备 D1】读取向量服务配置；下面 MEETING_PACKAGES 是脚本内置的样例数据。
conf = Config()
DASHSCOPE_API_KEY = conf.embedding_api_key  # 向量模型使用独立的 DashScope API Key
EMBEDDING_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings"


def get_embedding(text: str) -> list:
    """
    调用 Qwen Embedding API 生成文本向量

    参数：
        text (str): 需要生成向量的文本

    返回值：
        list: 1024 维向量
    """
    # 【数据准备 D3/D4 共用】文本 → 向量接口 → 1024 维数值列表。
    # 入库的套餐文本和后续用户查询必须用兼容的模型及维度才能比较。
    # 输入 text 是拼好的套餐文本，或测试用的查询文本；输出一条 1024 维向量。
    # 同样的模型用于入库和查询，使两类文本被映射到可相互比较的向量空间。
    # 调用可能产生外部 API 请求；这个函数没有本地缓存，重复调用也会再次请求。
    headers = {
        "Authorization": f"Bearer {DASHSCOPE_API_KEY}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": "text-embedding-v3",
        "input": [text],
        "dimensions": 1024
    }
    response = requests.post(EMBEDDING_URL, headers=headers, json=payload, timeout=30)
    result = response.json()
    # 响应 data 是按输入文本顺序排列的列表；当前 input 只包含一个 text，因此取索引 0。
    return result["data"][0]["embedding"]


# ==================== 会务套餐数据 ====================
# 这是用于入库的本地套餐样例列表，每个字典是一条套餐；定义列表不会向数据库写入。
# package_id 是业务编号，embedding 尚未在这些样例中生成；后面 insert_meeting_packages 才逐条补上。
MEETING_PACKAGES = [
    {
        "package_id": "CD001",
        "package_name": "成都数字化转型峰会3日会务套餐",
        "city": "成都",
        "days": 3,
        "price": 2380.00,
        "total_seats": 80,
        "remaining_seats": 26,
        "agency": "携程商旅",
        "rating": 4.8,
        "departure_dates": ["2026-05-01", "2026-06-01", "2026-06-15"],
        "highlights": "覆盖机场接送、会场签到、四星商务酒店、数字化转型主论坛、企业案例圆桌、成都高新区企业参访，适合中高层管理者集中参会"
    },
    {
        "package_id": "CD002",
        "package_name": "成都客户拜访2日接待套餐",
        "city": "成都",
        "days": 2,
        "price": 1280.00,
        "total_seats": 40,
        "remaining_seats": 18,
        "agency": "同程商旅",
        "rating": 4.7,
        "departure_dates": ["2026-05-01", "2026-06-01", "2026-06-08"],
        "highlights": "包含机场/高铁站接送、客户会议室预订、商务午餐、企业展厅参观、晚间轻接待路线，适合销售团队与客户成功团队使用"
    },
    {
        "package_id": "CD003",
        "package_name": "成都团队共创工作坊2日套餐",
        "city": "成都",
        "days": 2,
        "price": 980.00,
        "total_seats": 60,
        "remaining_seats": 31,
        "agency": "会小二",
        "rating": 4.6,
        "departure_dates": ["2026-05-02", "2026-06-03", "2026-06-20"],
        "highlights": "提供创意会议空间、分组讨论物料、茶歇、团队共创主持人、晚间团队餐，可按部门目标设计议程"
    },
    {
        "package_id": "BJ001",
        "package_name": "北京政企拜访3日差旅套餐",
        "city": "北京",
        "days": 3,
        "price": 2680.00,
        "total_seats": 50,
        "remaining_seats": 14,
        "agency": "国旅运通商旅",
        "rating": 4.8,
        "departure_dates": ["2026-05-06", "2026-06-02", "2026-06-16"],
        "highlights": "覆盖政企客户拜访、会议室协调、商务酒店、用车排班、材料打印与接待动线，适合法务、销售和公共事务团队"
    },
    {
        "package_id": "BJ002",
        "package_name": "北京研发交流2日套餐",
        "city": "北京",
        "days": 2,
        "price": 1580.00,
        "total_seats": 35,
        "remaining_seats": 12,
        "agency": "飞猪企业商旅",
        "rating": 4.7,
        "departure_dates": ["2026-05-10", "2026-06-05"],
        "highlights": "安排技术交流会议、园区参访、商务茶歇、投影与白板设备、晚间简餐，适合研发团队跨城对齐方案"
    },
    {
        "package_id": "SH001",
        "package_name": "上海总部培训4日会务套餐",
        "city": "上海",
        "days": 4,
        "price": 3280.00,
        "total_seats": 90,
        "remaining_seats": 37,
        "agency": "锦江会务",
        "rating": 4.9,
        "departure_dates": ["2026-05-08", "2026-06-01", "2026-06-18"],
        "highlights": "包含培训教室、签到系统、四星酒店、浦东机场接送、讲师休息室、分组研讨区，适合新人训和区域经理培训"
    },
    {
        "package_id": "SH002",
        "package_name": "上海大型展会参展3日套餐",
        "city": "上海",
        "days": 3,
        "price": 2980.00,
        "total_seats": 70,
        "remaining_seats": 22,
        "agency": "携程商旅",
        "rating": 4.8,
        "departure_dates": ["2026-05-12", "2026-06-10"],
        "highlights": "覆盖展馆通行、展位物料运输、酒店班车、客户邀约晚宴、现场证件管理，适合市场与渠道团队参展"
    },
    {
        "package_id": "HZ001",
        "package_name": "杭州电商标杆研学3日套餐",
        "city": "杭州",
        "days": 3,
        "price": 2180.00,
        "total_seats": 45,
        "remaining_seats": 16,
        "agency": "浙商会务",
        "rating": 4.7,
        "departure_dates": ["2026-05-09", "2026-06-06", "2026-06-21"],
        "highlights": "安排电商园区参访、直播基地交流、业务复盘会议、团队晚餐和西湖周边轻接待，适合运营与增长团队"
    },
    {
        "package_id": "SZ001",
        "package_name": "深圳创新企业参访3日套餐",
        "city": "深圳",
        "days": 3,
        "price": 2580.00,
        "total_seats": 55,
        "remaining_seats": 20,
        "agency": "南方商旅",
        "rating": 4.8,
        "departure_dates": ["2026-05-11", "2026-06-09"],
        "highlights": "覆盖科技园企业参访、创新产品交流、商务用车、酒店会议室、闭门分享会，适合产品、战略和投资团队"
    },
    {
        "package_id": "GZ001",
        "package_name": "广州供应链展会3日套餐",
        "city": "广州",
        "days": 3,
        "price": 2280.00,
        "total_seats": 65,
        "remaining_seats": 28,
        "agency": "广之旅商旅",
        "rating": 4.6,
        "departure_dates": ["2026-05-15", "2026-06-12"],
        "highlights": "适配采购团队参展，包含展馆接送、供应商洽谈室、商务午餐、样品寄送协助和晚间客户接待"
    },
    {
        "package_id": "XA001",
        "package_name": "西安制造业客户答谢2日套餐",
        "city": "西安",
        "days": 2,
        "price": 1380.00,
        "total_seats": 50,
        "remaining_seats": 19,
        "agency": "陕旅会务",
        "rating": 4.6,
        "departure_dates": ["2026-05-18", "2026-06-14"],
        "highlights": "包含客户答谢晚宴、工厂参观车辆、会议设备、伴手礼管理和历史文化轻接待路线，适合制造业销售团队"
    },
    {
        "package_id": "QD001",
        "package_name": "青岛销售年会3日套餐",
        "city": "青岛",
        "days": 3,
        "price": 2480.00,
        "total_seats": 120,
        "remaining_seats": 48,
        "agency": "海滨会务",
        "rating": 4.7,
        "departure_dates": ["2026-05-20", "2026-06-20"],
        "highlights": "支持百人销售年会、主会场搭建、分会场讨论、颁奖晚宴、团队拓展和海边商务晚餐"
    },
    {
        "package_id": "CQ001",
        "package_name": "重庆渠道大会2日套餐",
        "city": "重庆",
        "days": 2,
        "price": 1180.00,
        "total_seats": 100,
        "remaining_seats": 42,
        "agency": "山城会务",
        "rating": 4.6,
        "departure_dates": ["2026-05-16", "2026-06-11"],
        "highlights": "覆盖渠道商签到、主论坛、圆桌沟通、商务晚宴、酒店会议室和机场接送，适合区域渠道大会"
    },
    {
        "package_id": "NJ001",
        "package_name": "南京项目交付驻场5日套餐",
        "city": "南京",
        "days": 5,
        "price": 3680.00,
        "total_seats": 30,
        "remaining_seats": 9,
        "agency": "苏宁商旅",
        "rating": 4.7,
        "departure_dates": ["2026-05-13", "2026-06-07"],
        "highlights": "适合项目交付团队驻场，包含长住商务酒店、每日通勤用车、会议室、打印服务、加班简餐和周报复盘空间"
    },
    {
        "package_id": "XM001",
        "package_name": "厦门高管闭门会3日套餐",
        "city": "厦门",
        "days": 3,
        "price": 3980.00,
        "total_seats": 25,
        "remaining_seats": 7,
        "agency": "建发商旅",
        "rating": 4.9,
        "departure_dates": ["2026-05-22", "2026-06-18"],
        "highlights": "面向高管闭门战略会，包含海景会议室、贵宾接送、保密签到、商务晚宴、会议纪要服务和安静住宿环境"
    },
]


# ==================== Milvus 初始化 ====================
MILVUS_HOST = "localhost"
MILVUS_PORT = 19530
COLLECTION_NAME = "meeting_packages"
DIMENSION = 1024  # text-embedding-v3 输出维度


def create_meeting_package_collection():
    """
    在 Milvus 中创建会务套餐 Collection

    字段设计：
    - id: 主键（自增）
    - package_id: 套餐编号（字符串，如 YN001）
    - package_name: 套餐名称
    - city: 城市
    - days: 天数
    - price: 价格
    - total_seats: 总座位
    - remaining_seats: 剩余座位
    - agency: 差旅服务商
    - rating: 评分
    - departure_dates: 出发日期列表（JSON 字符串）
    - highlights: 方案亮点（用于生成向量）
    - embedding: 向量（1024 维）
    """
    # 连接 Milvus
    # 【数据准备 D2.1】先连接 Milvus，再检查同名集合。
    # 由文件末尾主程序首先调用；输出新建好的 Collection 对象，交给下一步入库函数。
    # 普通 FieldSchema 保存套餐字段，FLOAT_VECTOR 字段承载语义检索向量。
    # 已有集合删除后，新的自增 id 会重新由 Milvus 分配；它与样例中的 package_id 不是一个字段。
    connections.connect(alias="default", host=MILVUS_HOST, port=MILVUS_PORT)
    print(f"已连接 Milvus: {MILVUS_HOST}:{MILVUS_PORT}")

    # 如果已存在则删除
    # 【数据准备 D2.2】已有集合会被删除后重建；此操作不会保留旧集合中的套餐。
    if utility.has_collection(COLLECTION_NAME):
        print(f"集合 {COLLECTION_NAME} 已存在，删除重建")
        utility.drop_collection(COLLECTION_NAME)

    # 定义字段
    # 【数据准备 D2.3】普通字段保存套餐信息，embedding 字段保存用于相似度检索的向量。
    fields = [
        FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
        FieldSchema(name="package_id", dtype=DataType.VARCHAR, max_length=20),
        FieldSchema(name="package_name", dtype=DataType.VARCHAR, max_length=200),
        FieldSchema(name="city", dtype=DataType.VARCHAR, max_length=50),
        FieldSchema(name="days", dtype=DataType.INT64),
        FieldSchema(name="price", dtype=DataType.FLOAT),
        FieldSchema(name="total_seats", dtype=DataType.INT64),
        FieldSchema(name="remaining_seats", dtype=DataType.INT64),
        FieldSchema(name="agency", dtype=DataType.VARCHAR, max_length=100),
        FieldSchema(name="rating", dtype=DataType.FLOAT),
        FieldSchema(name="departure_dates", dtype=DataType.VARCHAR, max_length=200),
        FieldSchema(name="highlights", dtype=DataType.VARCHAR, max_length=1000),
        FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=DIMENSION),
    ]

    # 字段列表组装成集合结构；Collection(name, schema) 才将这个结构用于创建集合。
    schema = CollectionSchema(fields, description="会务套餐信息表（RAG 向量检索）")
    collection = Collection(COLLECTION_NAME, schema)
    print(f"集合 {COLLECTION_NAME} 创建成功")

    # 创建向量索引
    # IVF_FLAT 是选用的向量索引类型；COSINE 是相似度度量；nlist 控制索引分组参数。
    # 这些参数影响检索方式，不是套餐数量、向量维度或模型生成的 temperature。
    index_params = {
        "index_type": "IVF_FLAT",
        "metric_type": "COSINE",
        "params": {"nlist": 64}
    }
    # 【数据准备 D2.4】为 embedding 创建 COSINE 索引，然后返回集合给下一步使用。
    collection.create_index(field_name="embedding", index_params=index_params)
    print("向量索引创建成功 (IVF_FLAT, COSINE)")

    # 返回的是可操作该集合的客户端对象；不是套餐列表，也不是已经生成好的向量。
    return collection


def insert_meeting_packages(collection):
    """
    将会务套餐数据生成向量后插入 Milvus
    """
    # 【数据准备 D3.1】依次处理套餐：拼检索文本 → 请求向量 → 将向量和原字段放入 rows。
    # 输入 collection 是上一步返回的集合；无业务返回值，效果是将 rows 插入并刷入存储。
    # 每个样例先组成 embedding_text，再调用向量接口，最后把向量与业务字段放在同一行中。
    # 全部样例处理完以后才调用一次 collection.insert，前面的 rows.append 仅修改 Python 列表。
    print(f"\n开始为 {len(MEETING_PACKAGES)} 个会务套餐生成向量...")
    rows = []
    # enumerate 给出从 0 开始的索引 i 和当前套餐 group，进度显示用 i+1 转成从 1 开始的计数。
    for i, group in enumerate(MEETING_PACKAGES):
        # 使用"城市 + 套餐名称 + 方案亮点"作为 embedding 输入文本
        # 这是入库向量对应的文本；查询时会把用户需求另行转成向量与它比较。
        # 只把城市、名称和亮点拼成向量文本；价格、余位等普通字段随后原样保存，不通过向量生成。
        embedding_text = f"{group['city']} {group['package_name']} {group['highlights']}"

        if (i + 1) % 5 == 0:
            print(f"  已处理 {i + 1}/{len(MEETING_PACKAGES)} 个套餐...")

        embedding = get_embedding(embedding_text)

        rows.append({
            "package_id": group["package_id"],
            "package_name": group["package_name"],
            "city": group["city"],
            "days": group["days"],
            "price": group["price"],
            "total_seats": group["total_seats"],
            "remaining_seats": group["remaining_seats"],
            "agency": group["agency"],
            "rating": group["rating"],
            # 出发日期原来是列表，集合 schema 对应字段是 VARCHAR，所以先编码成 JSON 字符串保存。
            "departure_dates": json.dumps(group["departure_dates"], ensure_ascii=False),
            "highlights": group["highlights"],
            "embedding": embedding,
        })

    # 批量插入
    # 【数据准备 D3.2】全部向量准备完才批量插入，并 flush；这里没有逐条重试或断点续写。
    collection.insert(rows)
    # flush 要求刷入待持久化数据；它与后续查询前的 load 用途不同，load 是准备读取检索。
    collection.flush()  # 刷盘
    print(f"成功插入 {len(MEETING_PACKAGES)} 个会务套餐数据")


def test_search(collection):
    """
    测试语义搜索：用自然语言查询会务套餐
    """
    # 测试查询
    # 【数据准备 D4.1】每条测试查询也先向量化，然后检索；本函数会实际调用向量接口。
    # 输入是已入库的 collection；输出方式为 print，便于人查看检索候选。
    # 每个 query 单独求向量，检索结果按“查询向量→命中列表”组织，所以有两层 for。
    # limit=2 限制候选条数，不表示验证了两个测试用例，也不是准确率阈值。
    test_queries = [
        "成都办数字化转型峰会，需要接送和酒店",
        "上海大型展会参展，需要展馆通行和客户晚宴",
        "适合销售团队的青岛年会方案",
        "深圳创新企业参访，最好包含闭门分享",
        "南京项目驻场，需要长住酒店和每日通勤",
    ]

    print("\n=== 语义搜索测试 ===")
    # 外层依次执行测试需求；前一条查询完成以后才开始下一条，这里没有并发任务调度。
    for query in test_queries:
        print(f"\n查询: {query}")
        query_embedding = get_embedding(query)

        collection.load()
        # 【数据准备 D4.2】取相似度靠前的两条并打印，供人工核对；不是自动质量评分。
        results = collection.search(
            data=[query_embedding],
            anns_field="embedding",
            param={"metric_type": "COSINE", "params": {"nprobe": 10}},
            limit=2,
            output_fields=["package_id", "package_name", "city", "days", "price", "rating", "agency", "highlights"]
        )

        for hits in results:
            for hit in hits:
                print(f"  [{hit.entity.get('package_id')}] {hit.entity.get('package_name')} "
                      f"- {hit.entity.get('city')} {hit.entity.get('days')}天 "
                      f"¥{hit.entity.get('price')} 评分{hit.entity.get('rating')} "
                      f"距离: {hit.distance:.4f}")


# 【数据准备入口】仅直接运行本脚本时按 D2 → D3 → D4 执行；导入不会进入此区块。
if __name__ == '__main__':
    # 本文件的实际运行入口：先拿到集合对象，再把同一个对象依次传给入库和测试函数。
    # 上面的函数定义在加载文件时被创建，但直到这里调用才真正连接服务、生成向量和写入数据。
    collection = create_meeting_package_collection()
    insert_meeting_packages(collection)
    test_search(collection)