group_fqcd_jr/tools/setup_milvus_profile_collection.py

"""幂等创建长期记忆画像向量集合 `user_long_term_memory_v1`。

写入侧是 `app/infrastructure/milvus_profile_projection.py`，字段必须与之一致。

用法：
    .\\.venv\\Scripts\\python.exe tools\\setup_milvus_profile_collection.py

安全口径（与 `setup_milvus_knowledge_collections.py` 相同：这是**共享** Milvus 实例，
里面还有别的项目在用的集合）：
- 只碰 `user_long_term_memory_v1` 这一个集合，绝不 list 后批量删除；
- 集合已存在时**只做结构比对并报告**，不覆盖、不删重建；
- 结构不一致时明确报错退出，由人决定怎么处理，避免静默丢数据。
"""

import asyncio
import sys
from typing import Any

COLLECTION = "user_long_term_memory_v1"
VECTOR_FIELD = "embedding"
VECTOR_DIM = 1024
PRIMARY_FIELD = "memory_uuid"

#: (字段名, 最大长度)。`memory_uuid` 是主键（UUID 字符串）。
VARCHAR_FIELDS: tuple[tuple[str, int], ...] = (
    ("memory_uuid", 64),
    ("content", 2048),
    ("memory_type", 32),
    ("memory_key", 64),
    ("status", 16),
)
#: INT64 字段；可空，因为 `valid_until_ts` 对永久记忆为空。
INT64_FIELDS: tuple[str, ...] = (
    "customer_id",
    "version",
    "valid_until_ts",
    "updated_at_ts",
)
FLOAT_FIELDS: tuple[str, ...] = ("confidence",)


def _build_schema() -> Any:
    from pymilvus import DataType, MilvusClient  # type: ignore[import-untyped]

    schema = MilvusClient.create_schema(auto_id=False, enable_dynamic_field=False)
    for name, max_length in VARCHAR_FIELDS:
        schema.add_field(
            field_name=name,
            datatype=DataType.VARCHAR,
            max_length=max_length,
            is_primary=(name == PRIMARY_FIELD),
            nullable=False,
        )
    for name in INT64_FIELDS:
        # 可空：`valid_until_ts` 对永久记忆必须能不写。
        schema.add_field(
            field_name=name, datatype=DataType.INT64, nullable=name != "customer_id"
        )
    for name in FLOAT_FIELDS:
        schema.add_field(field_name=name, datatype=DataType.DOUBLE, nullable=False)
    schema.add_field(
        field_name=VECTOR_FIELD, datatype=DataType.FLOAT_VECTOR, dim=VECTOR_DIM
    )
    return schema


def _build_index_params() -> Any:
    from pymilvus import MilvusClient  # type: ignore[import-untyped]

    index_params = MilvusClient.prepare_index_params()
    index_params.add_index(
        field_name=VECTOR_FIELD,
        index_name="profile_vector_index",
        index_type="AUTOINDEX",
        metric_type="COSINE",
    )
    return index_params


def describe_mismatch(described: dict[str, Any]) -> list[str]:
    """比对已存在集合与期望结构，返回差异列表（一致时为空）。"""
    problems: list[str] = []
    actual = {field["name"]: field for field in described.get("fields", [])}
    expected_names = (
        [name for name, _ in VARCHAR_FIELDS] + list(INT64_FIELDS)
        + list(FLOAT_FIELDS) + [VECTOR_FIELD]
    )
    for name in expected_names:
        if name not in actual:
            problems.append(f"缺少字段 {name}")
    for name, _ in VARCHAR_FIELDS:
        if name in actual and not actual[name].get("is_primary") and name == PRIMARY_FIELD:
            problems.append(f"{name} 不是主键")
    vector = actual.get(VECTOR_FIELD)
    if vector is not None:
        params = vector.get("params") or {}
        dim = params.get("dim")
        if dim is not None and int(dim) != VECTOR_DIM:
            problems.append(f"{VECTOR_FIELD} 维度是 {dim}，期望 {VECTOR_DIM}")
    return problems


async def ensure_collection(uri: str, token: str = "") -> str:
    """返回 `created` / `exists` / `conflict:<原因>`。"""
    from pymilvus import AsyncMilvusClient  # type: ignore[import-untyped]

    client = AsyncMilvusClient(uri=uri, token=token or None)
    try:
        if await client.has_collection(COLLECTION):
            described = await client.describe_collection(COLLECTION)
            problems = describe_mismatch(described)
            if problems:
                return "conflict:" + "; ".join(problems)
            return "exists"
        await client.create_collection(
            collection_name=COLLECTION,
            schema=_build_schema(),
            index_params=_build_index_params(),
        )
        return "created"
    finally:
        close = getattr(client, "close", None)
        if close is not None:
            await close()


async def main() -> int:
    from app.core.config import get_settings

    settings = get_settings()
    uri = settings.milvus_uri
    if not uri:
        print("未配置 milvus_uri，无法创建集合")
        return 1
    outcome = await ensure_collection(uri, settings.milvus_token or "")
    print(f"{COLLECTION}: {outcome}")
    if outcome.startswith("conflict:"):
        print("结构不一致，未做任何修改。请人工确认后再处理。")
        return 2
    return 0


if __name__ == "__main__":
    sys.exit(asyncio.run(main()))