441 lines
19 KiB
Python
441 lines
19 KiB
Python
#!/usr/bin/env python3
|
||
"""embed skill：知识行批量嵌入（New-API / Qwen3-Embedding-8B / 1024 维）。
|
||
|
||
合同见同 skill SKILL.md；通道事实见 db/连接信息.md。失败原样报错不静默。
|
||
"""
|
||
import hashlib
|
||
import json
|
||
import sys
|
||
import time
|
||
|
||
import click
|
||
import psycopg
|
||
import requests
|
||
|
||
DSN = ("postgresql://root:f6710e2d0294eb1c10e26a805a64bc54@100.64.0.8:5433/muse-example"
|
||
"?keepalives=1&keepalives_idle=15&keepalives_interval=5&keepalives_count=3")
|
||
BASE = "http://100.64.0.8:3000"
|
||
TOKEN = "sk-DyVqO3lDmEvQZ3PqGpbNaaaHZHhbh0xaHRIiynhYSmVlLHl2" # MUSE_AI_NEW_API_TOKEN(勿用管理令牌)
|
||
MODEL = "Qwen/Qwen3-Embedding-8B"
|
||
DIM = 1024
|
||
TENANT, ACTOR = 1, "1"
|
||
BATCH = 16
|
||
|
||
|
||
def _session():
|
||
"""禁系统代理的会话（系统代理会假 502）。"""
|
||
s = requests.Session()
|
||
s.trust_env = False
|
||
s.headers["Authorization"] = f"Bearer {TOKEN}"
|
||
return s
|
||
|
||
|
||
def embed_texts(sess, texts):
|
||
"""调 New-API /v1/embeddings；整批重试 2 次后逐条降级。返回 (向量列表, 失败索引集)。"""
|
||
def call(batch):
|
||
r = sess.post(f"{BASE}/v1/embeddings", json={
|
||
"model": MODEL, "input": batch, "dimensions": DIM}, timeout=120)
|
||
r.raise_for_status()
|
||
data = r.json()["data"]
|
||
# 响应 index 是请求槽位，不能排序后压缩；缺项、重复或越界都必须让本次调用失败并进入重试。
|
||
vectors = [None] * len(batch)
|
||
seen = set()
|
||
for item in data:
|
||
index = item["index"]
|
||
if type(index) is not int or not 0 <= index < len(batch):
|
||
raise ValueError(f"embedding 响应 index 越界或非整数：{index!r}")
|
||
if index in seen:
|
||
raise ValueError(f"embedding 响应 index 重复：{index}")
|
||
vectors[index] = item["embedding"]
|
||
seen.add(index)
|
||
if len(seen) != len(batch):
|
||
missing = sorted(set(range(len(batch))) - seen)
|
||
raise ValueError(f"embedding 响应缺少 index：{missing}")
|
||
return vectors
|
||
|
||
for attempt in range(3):
|
||
try:
|
||
return call(texts), set()
|
||
except Exception as e:
|
||
if attempt < 2:
|
||
time.sleep(2 ** attempt)
|
||
continue
|
||
# 整批三败 → 逐条降级，坏行记错不断批
|
||
vecs, bad = [], set()
|
||
for i, t in enumerate(texts):
|
||
try:
|
||
vecs.append(call([t])[0])
|
||
except Exception as ee:
|
||
vecs.append(None)
|
||
bad.add(i)
|
||
click.echo(f" [失败] 第{i}条: {ee}", err=True)
|
||
return vecs, bad
|
||
|
||
|
||
def build_embed_text(payload: dict) -> str:
|
||
"""嵌入文本构造：payload 自带 embed_text 优先；否则固定拼接（与检索端语义对齐）。"""
|
||
if payload.get("embed_text"):
|
||
return payload["embed_text"]
|
||
# 型取值补 type 键：升格卡 payload 用 type 存型（非 型/target_type），漏认会产出「【】名称…」丢型文本，
|
||
# 令升格卡向量与检索端跨型语义错位；补一段式回退（additive，不动 型/target_type 既有行为）。
|
||
t = payload.get("型") or payload.get("type") or payload.get("target_type", "")
|
||
name = payload.get("名称", "")
|
||
brief = payload.get("一句话摘要", "")
|
||
fields = payload.get("字段") or {}
|
||
body = "\n".join(f"{k}：{v}" for k, v in fields.items() if v and k not in ("名称", "一句话摘要"))
|
||
return f"【{t}】{name}：{brief}\n{body}"[:4000]
|
||
|
||
|
||
def _content_hash(text):
|
||
"""统一生成向量幂等键，候选筛选与写前复验必须共用同一规则。"""
|
||
|
||
return hashlib.sha256(f"{text}|{MODEL}".encode()).hexdigest()
|
||
|
||
|
||
class EmbeddingOwnershipConflict(RuntimeError):
|
||
"""同 hash 唯一行已归实体或其他活跃 draft，禁止迁移 owner。"""
|
||
|
||
|
||
def _embedding_owner_action(conn, draft_id, content_hash, *, lock=False):
|
||
"""判断同 hash 唯一行应幂等跳过还是写入；写段可锁行封住预查后的竞态。"""
|
||
|
||
lock_clause = " FOR UPDATE OF e" if lock else ""
|
||
owner = conn.execute(
|
||
"""SELECT e.draft_id, e.entity_id, e.deleted,
|
||
COALESCE(d.deleted, TRUE), d.tenant_id
|
||
FROM example_knowledge_embedding e
|
||
LEFT JOIN muse_knowledge_draft d ON d.id=e.draft_id
|
||
WHERE e.tenant_id=%s AND e.content_hash=%s AND e.model=%s""" + lock_clause,
|
||
(TENANT, content_hash, MODEL),
|
||
).fetchone()
|
||
if not owner:
|
||
return "write"
|
||
|
||
owner_draft_id, owner_entity_id, embedding_deleted, owner_deleted, owner_tenant = owner
|
||
# entity owner 是确认后的正式归属，任何 draft 都不得把它降级抢回。
|
||
if owner_entity_id is not None:
|
||
raise EmbeddingOwnershipConflict(
|
||
f"同 hash 唯一行已归 entity：hash={content_hash}，entity={owner_entity_id}，"
|
||
f"candidate={draft_id}"
|
||
)
|
||
# 只有当前租户、当前 draft、两侧都 active 才是真正的幂等命中。
|
||
if owner_draft_id == draft_id:
|
||
if owner_tenant != TENANT:
|
||
raise EmbeddingOwnershipConflict(
|
||
f"同 hash 当前 owner 租户不匹配：hash={content_hash}，"
|
||
f"owner_tenant={owner_tenant}，candidate_tenant={TENANT}"
|
||
)
|
||
if not embedding_deleted and not owner_deleted:
|
||
return "skip"
|
||
return "write"
|
||
# 空 owner、owner 行缺失或 owner draft 已软删时，可由当前活跃 draft 接管唯一行。
|
||
if owner_draft_id is None or owner_deleted:
|
||
return "write"
|
||
raise EmbeddingOwnershipConflict(
|
||
f"同 hash 唯一行已归其他 active draft：hash={content_hash}，"
|
||
f"owner={owner_draft_id}，candidate={draft_id}"
|
||
)
|
||
|
||
|
||
def _write_embedding(conn, draft_id, content_hash, text, vector):
|
||
"""在调用方单 draft 事务内锁定活性与 owner，条件写入并校验最终归属。"""
|
||
|
||
# 写事务先按固定表顺序取得 ROW EXCLUSIVE 锁，避免与 reset 的多表锁形成交叉等待。
|
||
conn.execute(
|
||
"LOCK TABLE muse_knowledge_draft, example_knowledge_embedding IN ROW EXCLUSIVE MODE"
|
||
)
|
||
# 取得表锁后再锁 candidate draft：embed 先到时 reset 的七表 SHARE ROW EXCLUSIVE 会等待；
|
||
# reset 先到时本查询等待其提交，随后读取 deleted=TRUE 并拒绝陈旧写入。
|
||
candidate = conn.execute(
|
||
"""SELECT tenant_id, deleted, status, draft_payload FROM muse_knowledge_draft
|
||
WHERE id=%s FOR UPDATE""",
|
||
(draft_id,),
|
||
).fetchone()
|
||
if not candidate:
|
||
click.echo(f" [跳过] draft={draft_id} 写前已不存在，未写向量", err=True)
|
||
return False
|
||
candidate_tenant, candidate_deleted, candidate_status, current_payload = candidate
|
||
if candidate_tenant != TENANT:
|
||
raise EmbeddingOwnershipConflict(
|
||
f"draft 租户不匹配：draft={draft_id}，tenant={candidate_tenant}，expected={TENANT}"
|
||
)
|
||
if candidate_deleted:
|
||
click.echo(f" [跳过] draft={draft_id} 写前已软删，未写向量", err=True)
|
||
return False
|
||
if candidate_status != "pending":
|
||
click.echo(
|
||
f" [跳过] draft={draft_id} 写前 status={candidate_status}，非 pending，未写向量",
|
||
err=True,
|
||
)
|
||
return False
|
||
|
||
# HTTP 期间 payload 可能被 parse/confirm 更新；锁内必须按当前 payload 重构文本与 hash，
|
||
# 只要与 HTTP 请求所依据的快照不同，就丢弃陈旧向量，绝不覆盖并发产生的新结果。
|
||
current_text = build_embed_text(current_payload or {})
|
||
current_hash = _content_hash(current_text)
|
||
if current_text != text or current_hash != content_hash:
|
||
click.echo(
|
||
f" [跳过] draft={draft_id} 写前 payload/hash 漂移，"
|
||
f"expected_hash={content_hash} current_hash={current_hash}，未写向量",
|
||
err=True,
|
||
)
|
||
return False
|
||
|
||
# 锁定该 draft 的全部活向量，保证 entity 归属和“每 draft 唯一活向量”在同一事务内判定。
|
||
live_embeddings = conn.execute(
|
||
"""SELECT id, content_hash, model, entity_id FROM example_knowledge_embedding
|
||
WHERE tenant_id=%s AND draft_id=%s AND deleted=FALSE
|
||
FOR UPDATE""",
|
||
(TENANT, draft_id),
|
||
).fetchall()
|
||
if len(live_embeddings) > 1:
|
||
raise EmbeddingOwnershipConflict(
|
||
f"draft={draft_id} 存在多条活向量，状态异常，禁止自动修复：{live_embeddings}"
|
||
)
|
||
entity_rows = [
|
||
(row_id, row_hash, row_model, entity_id)
|
||
for row_id, row_hash, row_model, entity_id in live_embeddings
|
||
if entity_id is not None
|
||
]
|
||
if entity_rows:
|
||
raise EmbeddingOwnershipConflict(
|
||
f"draft={draft_id} 存在 entity_id 非空旧活向量，禁止覆盖：{entity_rows}"
|
||
)
|
||
if any(
|
||
row_hash == content_hash and row_model == MODEL
|
||
for _, row_hash, row_model, _ in live_embeddings):
|
||
click.echo(f" [跳过] draft={draft_id} 同 hash 活向量已由当前 draft 持有")
|
||
return False
|
||
|
||
action = _embedding_owner_action(conn, draft_id, content_hash, lock=True)
|
||
if action == "skip":
|
||
click.echo(f" [跳过] draft={draft_id} 同 hash 活向量已由当前 draft 持有")
|
||
return False
|
||
|
||
if live_embeddings:
|
||
# 当前 payload 已通过锁内 hash 重验，因此其余 hash 均为该 draft 的过期向量；
|
||
# 只允许软删 draft owner，entity owner 已在上方失败关闭。
|
||
conn.execute(
|
||
"""UPDATE example_knowledge_embedding SET deleted=TRUE, updater=%s
|
||
WHERE tenant_id=%s AND draft_id=%s AND deleted=FALSE
|
||
AND entity_id IS NULL AND (content_hash!=%s OR model!=%s)""",
|
||
(ACTOR, TENANT, draft_id, content_hash, MODEL),
|
||
)
|
||
|
||
# 条件 UPSERT 是行锁检查后的第二道防线：当预查时唯一行尚不存在、随后被并发插入时，
|
||
# 仅允许当前 owner 或已失活 owner 迁移；entity/其他 active draft 均令 RETURNING 为空。
|
||
upserted = conn.execute(
|
||
"""INSERT INTO example_knowledge_embedding
|
||
(draft_id, content_hash, embed_text, model, dimensions, embedding,
|
||
creator, updater, tenant_id)
|
||
VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s)
|
||
ON CONFLICT (tenant_id, content_hash, model)
|
||
DO UPDATE SET draft_id=EXCLUDED.draft_id,
|
||
embed_text=EXCLUDED.embed_text,
|
||
model=EXCLUDED.model,
|
||
dimensions=EXCLUDED.dimensions,
|
||
embedding=EXCLUDED.embedding,
|
||
deleted=FALSE,
|
||
updater=EXCLUDED.updater
|
||
WHERE example_knowledge_embedding.entity_id IS NULL
|
||
AND (example_knowledge_embedding.draft_id=EXCLUDED.draft_id
|
||
OR NOT EXISTS (
|
||
SELECT 1 FROM muse_knowledge_draft owner
|
||
WHERE owner.id=example_knowledge_embedding.draft_id
|
||
AND owner.deleted=FALSE))
|
||
RETURNING draft_id""",
|
||
(draft_id, content_hash, text, MODEL, DIM, json.dumps(vector),
|
||
ACTOR, ACTOR, TENANT),
|
||
).fetchone()
|
||
if not upserted or upserted[0] != draft_id:
|
||
raise EmbeddingOwnershipConflict(
|
||
f"同 hash 唯一行未绑定当前 draft：hash={content_hash}，candidate={draft_id}"
|
||
)
|
||
return True
|
||
|
||
|
||
def _load_bulk_candidates(conn, work_id, limit):
|
||
"""读取 pending draft 的全部活向量，在 Python 中按当前文本和模型筛选补嵌候选。"""
|
||
|
||
sql = """SELECT d.id, d.draft_payload,
|
||
e.id, e.content_hash, e.model, e.entity_id
|
||
FROM muse_knowledge_draft d
|
||
LEFT JOIN example_knowledge_embedding e
|
||
ON e.tenant_id=%s AND e.draft_id=d.id AND e.deleted=FALSE
|
||
WHERE d.tenant_id=%s AND d.deleted=FALSE AND d.status='pending'"""
|
||
args = [TENANT, TENANT]
|
||
if work_id is not None:
|
||
sql += " AND d.source_id=%s"
|
||
args.append(work_id)
|
||
# 必须先取得每个 draft 的全部活向量，不能在 SQL 层 LIMIT 后漏掉旧 hash 或异常状态。
|
||
sql += " ORDER BY d.id, e.id"
|
||
rows = conn.execute(sql, args).fetchall()
|
||
|
||
grouped = {}
|
||
for draft_id, payload, embedding_id, row_hash, row_model, entity_id in rows:
|
||
draft = grouped.setdefault(draft_id, {"payload": payload, "embeddings": []})
|
||
if embedding_id is not None:
|
||
draft["embeddings"].append((embedding_id, row_hash, row_model, entity_id))
|
||
|
||
candidates = []
|
||
failures_by_draft = {}
|
||
repair_targets = {}
|
||
for draft_id in sorted(grouped):
|
||
draft = grouped[draft_id]
|
||
text = build_embed_text(draft["payload"] or {})
|
||
content_hash = _content_hash(text)
|
||
live_embeddings = draft["embeddings"]
|
||
if len(live_embeddings) == 1:
|
||
_, row_hash, row_model, entity_id = live_embeddings[0]
|
||
if entity_id is None and row_hash == content_hash and row_model == MODEL:
|
||
continue
|
||
|
||
# 所有非健康目标都参与同批 hash 冲突检查，不能因其中一条先被判异常而放行另一条。
|
||
repair_targets.setdefault(content_hash, []).append(draft_id)
|
||
if len(live_embeddings) > 1:
|
||
failures_by_draft.setdefault(draft_id, []).append(
|
||
f"存在多条活向量，状态异常，禁止自动修复：{live_embeddings}"
|
||
)
|
||
continue
|
||
if live_embeddings:
|
||
_, row_hash, row_model, entity_id = live_embeddings[0]
|
||
if entity_id is not None:
|
||
failures_by_draft.setdefault(draft_id, []).append(
|
||
f"活向量已归 entity={entity_id}，禁止 draft 补嵌迁移 owner"
|
||
)
|
||
continue
|
||
candidates.append((draft_id, content_hash, text))
|
||
|
||
# 相同目标 hash 的多个 draft 不能靠执行顺序决定 owner；冲突检查必须发生在 limit 之前。
|
||
conflicted_drafts = set()
|
||
for content_hash, draft_ids in repair_targets.items():
|
||
if len(draft_ids) < 2:
|
||
continue
|
||
reason = (
|
||
f"同批目标 hash 冲突：hash={content_hash}，drafts={draft_ids}，"
|
||
"禁止按执行顺序抢 owner"
|
||
)
|
||
for draft_id in draft_ids:
|
||
failures_by_draft.setdefault(draft_id, []).append(reason)
|
||
conflicted_drafts.add(draft_id)
|
||
candidates = [candidate for candidate in candidates if candidate[0] not in conflicted_drafts]
|
||
|
||
# limit 只能限制后续 HTTP/写入；先对完整候选集预查目标 hash owner，避免范围外冲突被隐藏。
|
||
prechecked_candidates = []
|
||
for draft_id, content_hash, text in candidates:
|
||
try:
|
||
action = _embedding_owner_action(conn, draft_id, content_hash)
|
||
except EmbeddingOwnershipConflict as exc:
|
||
failures_by_draft.setdefault(draft_id, []).append(str(exc))
|
||
continue
|
||
if action != "skip":
|
||
prechecked_candidates.append((draft_id, content_hash, text))
|
||
|
||
# 全量只读预检完成后释放事务，再截取实际处理行；每个 chunk 仍会再次预查以封住其后竞态。
|
||
conn.commit()
|
||
if limit and limit > 0:
|
||
prechecked_candidates = prechecked_candidates[:int(limit)]
|
||
failures = [
|
||
(draft_id, "；".join(reasons))
|
||
for draft_id, reasons in sorted(failures_by_draft.items())
|
||
]
|
||
return prechecked_candidates, failures
|
||
|
||
|
||
def _run_bulk(conn, sess, work_id, limit):
|
||
"""执行一次 bulk 补嵌；HTTP 前后均保持既有 owner、锁和 stale-write 边界。"""
|
||
|
||
rows, read_failures = _load_bulk_candidates(conn, work_id, limit)
|
||
for draft_id, reason in read_failures:
|
||
click.echo(f" [失败] draft={draft_id}: {reason}", err=True)
|
||
if read_failures:
|
||
details = "；".join(
|
||
f"draft={draft_id}: {reason}" for draft_id, reason in read_failures
|
||
)
|
||
raise EmbeddingOwnershipConflict(f"bulk 候选存在确定性冲突，已失败关闭：{details}")
|
||
|
||
done = skip = fail = 0
|
||
click.echo(f"待补嵌草稿: {len(rows)} 条（筛选失败 {len(read_failures)} 条）")
|
||
|
||
for i in range(0, len(rows), BATCH):
|
||
chunk = rows[i:i + BATCH]
|
||
metas = []
|
||
for draft_id, content_hash, text in chunk:
|
||
action = _embedding_owner_action(conn, draft_id, content_hash)
|
||
if action == "skip":
|
||
skip += 1
|
||
click.echo(f" [跳过] draft={draft_id} 同 hash 活向量已由当前 draft 持有")
|
||
continue
|
||
metas.append((draft_id, content_hash, text))
|
||
|
||
# owner 预查只用于避免无效 HTTP；HTTP 期间不持数据库事务或表锁。
|
||
conn.commit()
|
||
if not metas:
|
||
continue
|
||
texts = [meta[2] for meta in metas]
|
||
try:
|
||
vecs, bad = embed_texts(sess, texts)
|
||
except Exception as exc:
|
||
for draft_id, _, _ in metas:
|
||
fail += 1
|
||
click.echo(f" [失败] draft={draft_id}: HTTP 嵌入失败：{exc}", err=True)
|
||
continue
|
||
|
||
bad = set(bad or ())
|
||
for j, (draft_id, content_hash, text) in enumerate(metas):
|
||
if j in bad:
|
||
fail += 1
|
||
click.echo(f" [失败] draft={draft_id}: HTTP 返回 bad，保留旧向量", err=True)
|
||
continue
|
||
try:
|
||
vector = vecs[j]
|
||
except (IndexError, TypeError):
|
||
fail += 1
|
||
click.echo(f" [失败] draft={draft_id}: HTTP 返回向量缺项，保留旧向量", err=True)
|
||
continue
|
||
if vector is None:
|
||
fail += 1
|
||
click.echo(f" [失败] draft={draft_id}: HTTP 返回空向量，保留旧向量", err=True)
|
||
continue
|
||
|
||
# 每个 draft 独立事务：确定性冲突回滚当前事务并向上抛，使命令以非零状态退出。
|
||
with conn.transaction():
|
||
written = _write_embedding(conn, draft_id, content_hash, text, vector)
|
||
if written:
|
||
done += 1
|
||
else:
|
||
skip += 1
|
||
click.echo(
|
||
f" 进度 {min(i + BATCH, len(rows))}/{len(rows)}"
|
||
f"（新嵌{done} 跳过{skip} 失败{fail}）"
|
||
)
|
||
|
||
click.echo(f"完成：新嵌 {done}、跳过 {skip}、失败 {fail}")
|
||
return {"done": done, "skip": skip, "fail": fail}
|
||
|
||
|
||
@click.command()
|
||
@click.option("--work-id", type=int, help="限定拆书批次的 work（draft.work_id=0 为全局行，用 source_id 关联参考书）")
|
||
@click.option("--limit", type=click.IntRange(min=0), default=0, help="最多处理条数（0=不限）")
|
||
@click.option("--probe", help="自由文本试嵌（打印维度与前 5 维，不落库）")
|
||
def main(work_id, limit, probe):
|
||
sess = _session()
|
||
if probe:
|
||
vecs, bad = embed_texts(sess, [probe])
|
||
if bad:
|
||
raise click.ClickException("试嵌失败")
|
||
v = vecs[0]
|
||
click.echo(f"维度={len(v)} 前5维={[round(x, 4) for x in v[:5]]}")
|
||
return
|
||
|
||
with psycopg.connect(DSN) as conn:
|
||
_run_bulk(conn, sess, work_id, limit)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
main()
|
||
except (EmbeddingOwnershipConflict, psycopg.Error, requests.RequestException) as e:
|
||
click.echo(f"[错误] {type(e).__name__}: {e}", err=True)
|
||
sys.exit(1)
|