299 lines
13 KiB
Python
299 lines
13 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"]
|
||
return [d["embedding"] for d in sorted(data, key=lambda d: d["index"])]
|
||
|
||
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]
|
||
|
||
|
||
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 = hashlib.sha256(f"{current_text}|{MODEL}".encode()).hexdigest()
|
||
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, entity_id FROM example_knowledge_embedding
|
||
WHERE tenant_id=%s AND draft_id=%s AND deleted=FALSE
|
||
FOR UPDATE""",
|
||
(TENANT, draft_id),
|
||
).fetchall()
|
||
entity_rows = [
|
||
(row_id, row_hash, entity_id)
|
||
for row_id, row_hash, 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 for _, row_hash, _ 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""",
|
||
(ACTOR, TENANT, draft_id, content_hash),
|
||
)
|
||
|
||
# 条件 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
|
||
|
||
|
||
@click.command()
|
||
@click.option("--work-id", type=int, help="限定拆书批次的 work（draft.work_id=0 为全局行，用 source_id 关联参考书）")
|
||
@click.option("--limit", type=int, 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:
|
||
# 待嵌=pending 草稿且无嵌入行
|
||
sql = """SELECT d.id, d.draft_payload FROM muse_knowledge_draft d
|
||
WHERE d.tenant_id=%s AND d.deleted=FALSE AND d.status='pending'
|
||
AND NOT EXISTS (SELECT 1 FROM example_knowledge_embedding e
|
||
WHERE e.tenant_id=%s AND e.draft_id=d.id AND e.deleted=FALSE)"""
|
||
args = [TENANT, TENANT]
|
||
if work_id is not None:
|
||
sql += " AND d.source_id=%s"
|
||
args.append(work_id)
|
||
sql += " ORDER BY d.id"
|
||
if limit:
|
||
sql += f" LIMIT {int(limit)}"
|
||
rows = conn.execute(sql, args).fetchall()
|
||
click.echo(f"待嵌草稿: {len(rows)} 条")
|
||
|
||
done = skip = fail = 0
|
||
for i in range(0, len(rows), BATCH):
|
||
chunk = rows[i:i + BATCH]
|
||
texts, metas = [], []
|
||
for did, payload in chunk:
|
||
text = build_embed_text(payload or {})
|
||
h = hashlib.sha256(f"{text}|{MODEL}".encode()).hexdigest()
|
||
if _embedding_owner_action(conn, did, h) == "skip":
|
||
skip += 1 # 仅当前 active draft 持有 active 唯一行时才幂等跳过。
|
||
continue
|
||
texts.append(text)
|
||
metas.append((did, h, text))
|
||
# owner 预查只用于避免无效 HTTP；明确提交读事务，HTTP 期间不持数据库事务或表锁。
|
||
conn.commit()
|
||
if not texts:
|
||
continue
|
||
vecs, bad = embed_texts(sess, texts)
|
||
for j, (did, h, text) in enumerate(metas):
|
||
if j in bad:
|
||
fail += 1
|
||
continue
|
||
# 每个 draft 独立事务：draft 行锁、owner 行锁、条件 UPSERT 同生共死。
|
||
with conn.transaction():
|
||
written = _write_embedding(conn, did, h, text, vecs[j])
|
||
if written:
|
||
done += 1
|
||
else:
|
||
skip += 1
|
||
click.echo(f" 进度 {min(i + BATCH, len(rows))}/{len(rows)}（新嵌{done} 跳过{skip} 失败{fail}）")
|
||
click.echo(f"完成：新嵌 {done}、跳过 {skip}、失败 {fail}")
|
||
|
||
|
||
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)
|