124 lines
5.6 KiB
Python
124 lines
5.6 KiB
Python
#!/usr/bin/env python3
|
||
"""search skill：知识向量检索（余弦召回 + 授权过滤 + aiContext 字段裁剪）。
|
||
|
||
合同见同 skill SKILL.md。查询嵌入与知识行同模型同维（复用 embed skill 的实现）。
|
||
"""
|
||
import json
|
||
import pathlib
|
||
import sys
|
||
|
||
import click
|
||
import psycopg
|
||
|
||
# 复用 embed skill 的通道实现（同模型同维，语义对齐）
|
||
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2] / "embed" / "scripts"))
|
||
from embed_drafts import _session, embed_texts # noqa: E402
|
||
|
||
DSN = "postgresql://root:f6710e2d0294eb1c10e26a805a64bc54@100.64.0.8:5433/muse-example"
|
||
TENANT = 1
|
||
|
||
|
||
def load_ai_context(conn):
|
||
"""读库内 23 型的字段级 aiContext 细则：{target_type: {field: true/false/[用途]}}。"""
|
||
rows = conn.execute(
|
||
"""SELECT s.target_type, v.policy_snapshot->'fieldAiContext'
|
||
FROM muse_meta_schema s
|
||
JOIN muse_meta_schema_version sv ON sv.id = s.active_version_id
|
||
JOIN muse_meta_visibility_policy v ON v.schema_version_id = sv.id
|
||
WHERE s.tenant_id=%s AND s.deleted=FALSE""", (TENANT,)).fetchall()
|
||
return {t: (m or {}) for t, m in rows}
|
||
|
||
|
||
def visible(ai_rule, purpose):
|
||
"""aiContext 判定：true 全用途可见；false 不可见；[用途] 仅列出的可见；无规则默认可见。"""
|
||
if ai_rule is None:
|
||
return True
|
||
if isinstance(ai_rule, bool):
|
||
return ai_rule
|
||
return purpose in ai_rule
|
||
|
||
|
||
@click.command()
|
||
@click.argument("intent")
|
||
@click.option("--scope", type=click.Choice(["admin", "work"]), default="admin", show_default=True,
|
||
help="admin=治理面(含草稿); work=作品面(仅已确认+已绑定)")
|
||
@click.option("--work-id", type=int, help="scope=work 时必填")
|
||
@click.option("--type", "ttype", help="限定型（如 craft/combat/emotion/scene_pattern/trope）")
|
||
@click.option("--purpose", default="generation", show_default=True,
|
||
type=click.Choice(["generation", "planning", "detection", "extraction"]))
|
||
@click.option("--top", default=5, show_default=True)
|
||
@click.option("--json", "as_json", is_flag=True)
|
||
def main(intent, scope, work_id, ttype, purpose, top, as_json):
|
||
if scope == "work" and not work_id:
|
||
raise click.ClickException("--scope work 必须带 --work-id（授权过滤依赖绑定关系）")
|
||
|
||
vecs, bad = embed_texts(_session(), [intent])
|
||
if bad:
|
||
raise click.ClickException("查询嵌入失败")
|
||
qvec = json.dumps(vecs[0])
|
||
|
||
with psycopg.connect(DSN) as conn:
|
||
ai_rules = load_ai_context(conn)
|
||
if scope == "admin":
|
||
# 治理面：draft(pending/confirmed)+entity 全量
|
||
sql = """SELECT 'draft' AS src, d.id, d.draft_payload AS payload, d.status,
|
||
1 - (e.embedding <=> %s::vector) AS score
|
||
FROM example_knowledge_embedding e
|
||
JOIN muse_knowledge_draft d ON d.id = e.draft_id
|
||
WHERE e.tenant_id=%s AND e.deleted=FALSE AND d.deleted=FALSE"""
|
||
args = [qvec, TENANT]
|
||
else:
|
||
# 作品面：仅 active entity 且其 kb 已绑定到该作品（授权在查询层强制）
|
||
sql = """SELECT 'entity' AS src, en.id,
|
||
jsonb_build_object('型', en.entity_type, '名称', en.normalized_name,
|
||
'一句话摘要', en.description, '字段', en.attributes) AS payload,
|
||
en.status, 1 - (e.embedding <=> %s::vector) AS score
|
||
FROM example_knowledge_embedding e
|
||
JOIN muse_knowledge_entity en ON en.id = e.entity_id
|
||
JOIN muse_knowledge_binding b ON b.kb_id = en.kb_id AND b.work_id = %s
|
||
AND b.binding_status='active' AND b.deleted=FALSE AND b.tenant_id=%s
|
||
WHERE e.tenant_id=%s AND e.deleted=FALSE AND en.deleted=FALSE AND en.status='active'"""
|
||
args = [qvec, work_id, TENANT, TENANT]
|
||
if ttype:
|
||
sql += (" AND d.draft_payload->>'型' = %s" if scope == "admin"
|
||
else " AND en.entity_type = %s")
|
||
args.append(ttype)
|
||
sql += " ORDER BY score DESC LIMIT %s"
|
||
args.append(top)
|
||
rows = conn.execute(sql, args).fetchall()
|
||
|
||
results = []
|
||
for src, rid, payload, status, score in rows:
|
||
p = payload or {}
|
||
t = p.get("型", "?")
|
||
rules = ai_rules.get(t, {})
|
||
fields = p.get("字段") or {}
|
||
vis = {k: v for k, v in fields.items() if visible(rules.get(k), purpose)}
|
||
cut = sorted(set(fields) - set(vis))
|
||
results.append({"来源": f"{src}#{rid}", "型": t, "名称": p.get("名称"), "状态": status,
|
||
"相似度": round(float(score), 4), "一句话摘要": p.get("一句话摘要"),
|
||
"可见字段": vis, "出处": p.get("出处"), "裁剪回显": cut})
|
||
|
||
if as_json:
|
||
click.echo(json.dumps(results, ensure_ascii=False, indent=1))
|
||
return
|
||
for i, r in enumerate(results, 1):
|
||
click.echo(f"── {i}. [{r['相似度']}] {r['型']} · {r['名称']}（{r['状态']}，{r['来源']}）")
|
||
click.echo(f" 摘要: {r['一句话摘要']}")
|
||
for k, v in (r["可见字段"] or {}).items():
|
||
click.echo(f" {k}: {str(v)[:120]}")
|
||
if r["出处"]:
|
||
click.echo(f" 出处: {r['出处']}")
|
||
if r["裁剪回显"]:
|
||
click.echo(f" [裁剪回显·{purpose} 不可见] {','.join(r['裁剪回显'])}")
|
||
if not results:
|
||
click.echo("（无召回）")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
main()
|
||
except psycopg.Error as e:
|
||
click.echo(f"[db错误] {type(e).__name__}: {e}", err=True)
|
||
sys.exit(1)
|