张胜宇
9675df8453
- 分支内容对齐 qyqy_develop b6ec3aa，树完全一致（同步后 git diff 为空） - 覆盖本轮全部交付：客服 Agent 重构（安全路由 / 五出口 / 记忆与画像 / RAG 全链路） + 开发文档 62 份编号体系（D1.1 v1.17 索引） + 新增 D2.10-客服Agent端到端答辩文档-2026-09-21.html - 基线：e239eb7（2026-09-17 品牌口径统一快照），本提交为其直接后继
151 lines
5.3 KiB
Python
151 lines
5.3 KiB
Python
"""幂等创建长期记忆画像向量集合 `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 pathlib import Path
|
||
from typing import Any
|
||
|
||
# 与 setup_milvus_knowledge_collections.py 同一条引导：脚本必须能直接在仓库根目录用
|
||
# 项目解释器执行（见模块 docstring 的用法）。缺这一行会在根目录直接报
|
||
# ModuleNotFoundError: No module named app（2026-09-19 S-7 实测）。
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||
|
||
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()))
|