mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 01:11:59 +00:00
ThearchyHelios
8d9a712eae
- 新增标准业务实体（CanonicalEntity）模型、确定性跨源编码翻译、后台候选实体扫描与文档驱动实体智能助手 - 新增指标本体图谱（factor / reconciles_with / upstream_driver）、DAG 防环校验以及 ontology_neighbors / source_health 内置工具 - 新增跨源对账（analysis_reconcile）与滞后共变（analysis_cochange）分析算子，打通 [[h:H1]] 数据源健康证据链、报告导出与本体消融评测 - 数据目录新增「标准业务实体」视图与候选审阅工作台，口径页集成指标本体拓扑面板与端到端浏览器检查
824 lines
32 KiB
Python
824 lines
32 KiB
Python
"""标准业务实体 Agent 化治理单元测试（E1 + E2）：附件/知识库联动、增量补丁护栏、交互式实体助手 SSE 与后台扫描发现任务。"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import sqlite3
|
||
import uuid
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import pytest
|
||
from httpx import ASGITransport, AsyncClient
|
||
from langchain_core.messages import AIMessage, AIMessageChunk
|
||
from sqlalchemy import delete, select
|
||
|
||
from app.data import catalog_assistant, entities, entity_assistant, entity_discovery, entity_docs, entity_patch
|
||
from app.data.engine import CATALOG_PROFILE_OPTION, MASK_COLUMNS_OPTION
|
||
from app.db.base import SessionLocal
|
||
from app.db.models import Chunk, Document, Workflow
|
||
from app.db.models_ontology import CanonicalEntity, EntityDiscoveryJob
|
||
from app.main import app
|
||
from tests.fixtures.sources import forget_source, sqlite_source
|
||
|
||
|
||
@dataclass
|
||
class _Turn:
|
||
text: list[str] = field(default_factory=list)
|
||
calls: list[tuple[str, dict[str, Any]]] = field(default_factory=list)
|
||
usage: tuple[int, int] = (80, 30)
|
||
|
||
|
||
class _ScriptedModel:
|
||
def __init__(self, turns: list[_Turn] | None = None, *, invoke_reply: str = "") -> None:
|
||
self.turns = list(turns or [])
|
||
self.invoke_reply = invoke_reply
|
||
self.seen: list[list[Any]] = []
|
||
self.tools: list[Any] | None = None
|
||
|
||
def bind_tools(self, tools: list[Any], **_: Any) -> "_ScriptedModel":
|
||
self.tools = tools
|
||
return self
|
||
|
||
async def astream(self, messages: list[Any]):
|
||
self.seen.append(list(messages))
|
||
turn = self.turns.pop(0)
|
||
for piece in turn.text:
|
||
yield AIMessageChunk(content=piece)
|
||
for i, (name, args) in enumerate(turn.calls):
|
||
yield AIMessageChunk(
|
||
content="",
|
||
tool_call_chunks=[
|
||
{
|
||
"name": name,
|
||
"args": json.dumps(args, ensure_ascii=False),
|
||
"id": f"tc_{len(self.seen)}_{i}",
|
||
"index": i,
|
||
}
|
||
],
|
||
)
|
||
yield AIMessageChunk(
|
||
content="",
|
||
usage_metadata={
|
||
"input_tokens": turn.usage[0],
|
||
"output_tokens": turn.usage[1],
|
||
"total_tokens": sum(turn.usage),
|
||
},
|
||
response_metadata={"finish_reason": "stop"},
|
||
)
|
||
|
||
async def ainvoke(self, messages: list[Any]) -> AIMessage:
|
||
self.seen.append(list(messages))
|
||
return AIMessage(content=self.invoke_reply)
|
||
|
||
|
||
async def _make_sqlite_source(
|
||
tmp_path: Path,
|
||
*,
|
||
name: str,
|
||
setup_sql: str,
|
||
profiling: bool = True,
|
||
mask_columns: str = "",
|
||
) -> str:
|
||
db_file = tmp_path / f"{name}_{uuid.uuid4().hex[:6]}.db"
|
||
conn = sqlite3.connect(db_file)
|
||
try:
|
||
conn.executescript(setup_sql)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
options: dict[str, Any] = {CATALOG_PROFILE_OPTION: {"enabled": profiling}}
|
||
if mask_columns:
|
||
options[MASK_COLUMNS_OPTION] = mask_columns
|
||
|
||
return await sqlite_source(name, db_file, options=options, schema=True)
|
||
|
||
|
||
@pytest.fixture
|
||
async def multi_source_env(tmp_path: Path):
|
||
mes_name = f"mes_ag_{uuid.uuid4().hex[:6]}"
|
||
erp_name = f"erp_ag_{uuid.uuid4().hex[:6]}"
|
||
wf_id = f"wf_ag_{uuid.uuid4().hex[:6]}"
|
||
|
||
mes_id = await _make_sqlite_source(
|
||
tmp_path,
|
||
name=mes_name,
|
||
setup_sql="""
|
||
CREATE TABLE dwd_workorder (
|
||
line_code TEXT,
|
||
line_name TEXT,
|
||
operator_phone TEXT,
|
||
good_qty INTEGER
|
||
);
|
||
INSERT INTO dwd_workorder VALUES
|
||
('L1', '一号产线', '13900000011', 100),
|
||
('L2', '二号产线', '13900000022', 95);
|
||
""",
|
||
profiling=True,
|
||
mask_columns="operator_phone",
|
||
)
|
||
erp_id = await _make_sqlite_source(
|
||
tmp_path,
|
||
name=erp_name,
|
||
setup_sql="""
|
||
CREATE TABLE dwd_settlement (
|
||
workshop_line TEXT,
|
||
settled_qty INTEGER
|
||
);
|
||
INSERT INTO dwd_settlement VALUES
|
||
('LINE_01', 98),
|
||
('LINE_02', 95);
|
||
""",
|
||
profiling=True,
|
||
)
|
||
|
||
async with SessionLocal() as session:
|
||
wf = Workflow(
|
||
id=wf_id,
|
||
name=f"良率核算_{wf_id}",
|
||
version=1,
|
||
graph={
|
||
"nodes": [
|
||
{
|
||
"id": "in_1",
|
||
"type": "input",
|
||
"data": {"config": {"fields": [{"name": "p_line", "label": "产线编码", "type": "string"}]}},
|
||
},
|
||
{
|
||
"id": "card_yield",
|
||
"type": "metrics",
|
||
"data": {"config": {"metrics": [{"id": "fpy", "name": "一次良率"}]}},
|
||
},
|
||
],
|
||
"edges": [],
|
||
},
|
||
)
|
||
session.add(wf)
|
||
await session.commit()
|
||
|
||
created_doc_ids: list[str] = []
|
||
try:
|
||
yield {
|
||
"mes_id": mes_id,
|
||
"mes_name": mes_name,
|
||
"erp_id": erp_id,
|
||
"erp_name": erp_name,
|
||
"wf_id": wf_id,
|
||
"doc_ids": created_doc_ids,
|
||
}
|
||
finally:
|
||
await forget_source(mes_id)
|
||
await forget_source(erp_id)
|
||
async with SessionLocal() as session:
|
||
await session.execute(delete(CanonicalEntity))
|
||
await session.execute(delete(EntityDiscoveryJob))
|
||
await session.execute(delete(Workflow).where(Workflow.id == wf_id))
|
||
if created_doc_ids:
|
||
await session.execute(delete(Chunk).where(Chunk.document_id.in_(created_doc_ids)))
|
||
await session.execute(delete(Document).where(Document.id.in_(created_doc_ids)))
|
||
await session.commit()
|
||
|
||
|
||
async def test_attachment_upload_and_optional_save_to_kb(multi_source_env):
|
||
csv_bytes = "标准编码,产线名称,MES编码,ERP编码\nL1,一号产线,L1,LINE_01\nL2,二号产线,L2,LINE_02\n".encode("utf-8")
|
||
transport = ASGITransport(app=app)
|
||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||
# 1) 不勾选存入知识库：提取表格内容但不创建 Document
|
||
r1 = await client.post(
|
||
"/api/ontology/entities/attachments",
|
||
files={"file": ("line_mapping.csv", csv_bytes, "text/csv")},
|
||
data={"save_to_kb": "false"},
|
||
)
|
||
assert r1.status_code == 201, r1.text
|
||
d1 = r1.json()
|
||
assert d1["saved_to_kb"] is False
|
||
assert d1["document_id"] is None
|
||
assert d1["chunk_count"] >= 1
|
||
|
||
# 2) 勾选存入知识库：同步创建 Document 并返回 document_id
|
||
r2 = await client.post(
|
||
"/api/ontology/entities/attachments",
|
||
files={"file": ("line_mapping_kb.csv", csv_bytes, "text/csv")},
|
||
data={"save_to_kb": "true", "collection": "default"},
|
||
)
|
||
assert r2.status_code == 201, r2.text
|
||
d2 = r2.json()
|
||
assert d2["saved_to_kb"] is True
|
||
assert d2["document_id"]
|
||
multi_source_env["doc_ids"].append(d2["document_id"])
|
||
|
||
async with SessionLocal() as session:
|
||
doc_row = await session.get(Document, d2["document_id"])
|
||
assert doc_row is not None
|
||
assert "LINE_01" in doc_row.content
|
||
|
||
# 同时传 attachment_id 与对应 document_id 时自动去重为 1 份
|
||
bundles = await entity_docs.load_doc_bundles(
|
||
session,
|
||
attachment_ids=[d2["attachment_id"]],
|
||
document_ids=[d2["document_id"]],
|
||
)
|
||
assert len(bundles) == 1
|
||
hits = entity_docs.search_doc_bundles(bundles, query="LINE_01")
|
||
assert len(hits) == 1
|
||
assert "文档《line_mapping_kb.csv》" in hits[0]["citation"]
|
||
|
||
|
||
async def test_entity_patch_guardrails_masked_columns_and_cas(multi_source_env):
|
||
mes = multi_source_env["mes_name"]
|
||
erp = multi_source_env["erp_name"]
|
||
wf_id = multi_source_env["wf_id"]
|
||
|
||
async with SessionLocal() as session:
|
||
# 1) 校验不存在的列、不存在的参数与遮罩列枚举阻断
|
||
bad_plan = await entity_patch.plan_entity_patch(
|
||
session,
|
||
[
|
||
{"op": "upsert_entity", "entity_key": "phone_ent", "label": "作业员号码"},
|
||
{
|
||
"op": "upsert_binding",
|
||
"entity_key": "phone_ent",
|
||
"source_name": mes,
|
||
"table_name": "dwd_workorder",
|
||
"column_name": "operator_phone",
|
||
},
|
||
{
|
||
"op": "upsert_member",
|
||
"entity_key": "phone_ent",
|
||
"value": "P1",
|
||
"label": "张三",
|
||
"source_codes": {mes: "13900000011"},
|
||
},
|
||
{
|
||
"op": "upsert_binding",
|
||
"entity_key": "phone_ent",
|
||
"source_name": mes,
|
||
"table_name": "dwd_workorder",
|
||
"column_name": "no_such_col",
|
||
},
|
||
{
|
||
"op": "upsert_caliber_param",
|
||
"entity_key": "phone_ent",
|
||
"workflow_id": wf_id,
|
||
"param_name": "wrong_param",
|
||
},
|
||
],
|
||
mode="confirmed",
|
||
)
|
||
assert len(bad_plan.problems) == 3
|
||
assert any("遮罩列" in p for p in bad_plan.problems)
|
||
assert any("no_such_col" in p for p in bad_plan.problems)
|
||
assert any("wrong_param" in p for p in bad_plan.problems)
|
||
|
||
# 2) mode="confirmed" 正常落库，随后测试 expected_versions CAS 乐观锁与 mode="proposed" 隔离
|
||
res = await entity_patch.apply_entity_patch(
|
||
session,
|
||
[
|
||
{
|
||
"op": "upsert_entity",
|
||
"entity_key": "prod_line_t",
|
||
"label": "产线",
|
||
"status": "proposed", # 模型自报状态将被忽略，confirmed 模式下按 confirmed 落库
|
||
},
|
||
{
|
||
"op": "upsert_binding",
|
||
"entity_key": "prod_line_t",
|
||
"source_name": mes,
|
||
"table_name": "dwd_workorder",
|
||
"column_name": "line_code",
|
||
"reason": "使用者说：工单产线列是 line_code",
|
||
},
|
||
],
|
||
mode="confirmed",
|
||
actor="builder_a",
|
||
)
|
||
await session.commit()
|
||
assert res["created_entities"] == 1
|
||
ent = res["entities"][0]
|
||
assert ent["bindings"][0]["status"] == "confirmed"
|
||
|
||
# 将一条 ERP 绑定预先设为 rejected，验证 mode="proposed" 绝不覆盖 confirmed 或复活 rejected
|
||
ent_row = await entities.get_entity_by_key(session, "prod_line_t")
|
||
assert ent_row is not None
|
||
await entities.update_entity(
|
||
session,
|
||
ent_row.id,
|
||
{
|
||
"version": ent_row.version,
|
||
"bindings": [
|
||
*ent_row.bindings,
|
||
{
|
||
"source_name": erp,
|
||
"table_name": "dwd_settlement",
|
||
"column_name": "settled_qty",
|
||
"status": "rejected",
|
||
"source": "human",
|
||
},
|
||
],
|
||
},
|
||
)
|
||
await session.commit()
|
||
|
||
# 3) mode="proposed"：尝试重提 rejected 列 + 新增 workshop_line + 自封 status="confirmed"
|
||
prop_res = await entity_patch.apply_entity_patch(
|
||
session,
|
||
[
|
||
{
|
||
"op": "upsert_binding",
|
||
"entity_key": "prod_line_t",
|
||
"source_name": erp,
|
||
"table_name": "dwd_settlement",
|
||
"column_name": "settled_qty",
|
||
"status": "confirmed",
|
||
},
|
||
{
|
||
"op": "upsert_binding",
|
||
"entity_key": "prod_line_t",
|
||
"source_name": erp,
|
||
"table_name": "dwd_settlement",
|
||
"column_name": "workshop_line",
|
||
"status": "confirmed",
|
||
"reason": "跨源目录对齐：结算产线",
|
||
},
|
||
{
|
||
"op": "upsert_caliber_param",
|
||
"entity_key": "prod_line_t",
|
||
"workflow_id": wf_id,
|
||
"card_id": "card_yield",
|
||
"param_name": "p_line",
|
||
"status": "confirmed",
|
||
},
|
||
],
|
||
mode="proposed",
|
||
)
|
||
await session.commit()
|
||
assert prop_res["bindings_changed"] == 1
|
||
assert prop_res["caliber_params_changed"] == 1
|
||
updated_ent = prop_res["entities"][0]
|
||
by_col = {b["column_name"]: b["status"] for b in updated_ent["bindings"]}
|
||
assert by_col["line_code"] == "confirmed"
|
||
assert by_col["settled_qty"] == "rejected"
|
||
assert by_col["workshop_line"] == "proposed"
|
||
assert updated_ent["caliber_params"][0]["status"] == "proposed"
|
||
|
||
# 4) CAS 版本冲突应抛 EntityConflict
|
||
with pytest.raises(entities.EntityConflict):
|
||
await entity_patch.apply_entity_patch(
|
||
session,
|
||
[{"op": "upsert_entity", "entity_key": "prod_line_t", "label": "新产线名"}],
|
||
mode="confirmed",
|
||
expected_versions={"prod_line_t": 1},
|
||
)
|
||
|
||
|
||
async def test_entity_assistant_sse_and_apply_patch(multi_source_env, monkeypatch):
|
||
mes = multi_source_env["mes_name"]
|
||
erp = multi_source_env["erp_name"]
|
||
wf_id = multi_source_env["wf_id"]
|
||
|
||
scripted = _ScriptedModel(
|
||
turns=[
|
||
_Turn(
|
||
calls=[
|
||
("inspect_sources", {"keyword": "line"}),
|
||
("read_document", {"query": "LINE_01"}),
|
||
("run_query", {"source_name": mes, "sql": "SELECT line_code, operator_phone FROM dwd_workorder"}),
|
||
]
|
||
),
|
||
_Turn(
|
||
calls=[
|
||
(
|
||
"propose_entity_changes",
|
||
{
|
||
"operations": [
|
||
{
|
||
"op": "upsert_entity",
|
||
"entity_key": "workshop_line_ent",
|
||
"label": "车间产线",
|
||
"aliases": ["产线"],
|
||
"reason": "使用者说：建立车间产线标准实体",
|
||
},
|
||
{
|
||
"op": "upsert_binding",
|
||
"entity_key": "workshop_line_ent",
|
||
"source_name": mes,
|
||
"table_name": "dwd_workorder",
|
||
"column_name": "line_code",
|
||
"reason": "数据探查显示：工单表产线列为 line_code",
|
||
},
|
||
{
|
||
"op": "upsert_binding",
|
||
"entity_key": "workshop_line_ent",
|
||
"source_name": erp,
|
||
"table_name": "dwd_settlement",
|
||
"column_name": "workshop_line",
|
||
"reason": "文档《对照表.md》写着：ERP 结算表对应 workshop_line",
|
||
},
|
||
{
|
||
"op": "upsert_member",
|
||
"entity_key": "workshop_line_ent",
|
||
"value": "L1",
|
||
"label": "一号产线",
|
||
"source_codes": {mes: "L1", erp: "LINE_01"},
|
||
"reason": "文档《对照表.md》写着：L1 对应 LINE_01",
|
||
},
|
||
{
|
||
"op": "upsert_caliber_param",
|
||
"entity_key": "workshop_line_ent",
|
||
"workflow_id": wf_id,
|
||
"card_id": "card_yield",
|
||
"param_name": "p_line",
|
||
"reason": "口径入口对齐",
|
||
},
|
||
]
|
||
},
|
||
)
|
||
]
|
||
),
|
||
_Turn(text=["已核对实表与对照文档，生成车间产线标准实体及跨源编码提案，请确认保存。"]),
|
||
]
|
||
)
|
||
|
||
async def _fake_resolve(_session):
|
||
return scripted, "mock-draft-model"
|
||
|
||
monkeypatch.setattr(entity_assistant, "resolve_model", _fake_resolve)
|
||
|
||
transport = ASGITransport(app=app)
|
||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||
att_res = await client.post(
|
||
"/api/ontology/entities/attachments",
|
||
files={"file": ("对照表.md", "一号产线：MES 为 L1，ERP dwd_settlement.workshop_line 为 LINE_01".encode("utf-8"), "text/markdown")},
|
||
data={"save_to_kb": "false"},
|
||
)
|
||
assert att_res.status_code == 201
|
||
att_id = att_res.json()["attachment_id"]
|
||
|
||
sse_res = await client.post(
|
||
"/api/ontology/entities/assistant",
|
||
json={
|
||
"messages": [{"role": "user", "content": "请根据对照表帮我关联车间产线实体"}],
|
||
"attachment_ids": [att_id],
|
||
},
|
||
)
|
||
assert sse_res.status_code == 200
|
||
events = [
|
||
json.loads(line[5:].strip())
|
||
for line in sse_res.text.splitlines()
|
||
if line.startswith("data:")
|
||
]
|
||
types = [ev["type"] for ev in events]
|
||
assert "query" in types
|
||
assert "proposal" in types
|
||
assert types[-1] == "done"
|
||
|
||
# 确认只读查询中的敏感列已被遮罩为「已遮罩」
|
||
q_ev = next(ev for ev in events if ev["type"] == "query")
|
||
assert q_ev["source"] == mes
|
||
assert all(row[1] == "已遮罩" for row in q_ev["rows"])
|
||
|
||
# 提案尚未落库
|
||
list_before = await client.get("/api/ontology/entities")
|
||
assert all(e["key"] != "workshop_line_ent" for e in list_before.json()["items"])
|
||
|
||
# 人工确认并应用提案
|
||
prop_ev = next(ev for ev in events if ev["type"] == "proposal")
|
||
ops = [c["op_payload"] for c in prop_ev["changes"]]
|
||
apply_res = await client.post(
|
||
"/api/ontology/entities/apply-patch",
|
||
json={"operations": ops, "expected_versions": prop_ev["expected_versions"]},
|
||
headers={"X-Actor": "builder_one"},
|
||
)
|
||
assert apply_res.status_code == 200, apply_res.text
|
||
applied = apply_res.json()
|
||
assert applied["created_entities"] == 1
|
||
saved = applied["entities"][0]
|
||
assert saved["key"] == "workshop_line_ent"
|
||
assert all(b["status"] == "confirmed" for b in saved["bindings"])
|
||
assert saved["members"][0]["source_codes"] == {mes: "L1", erp: "LINE_01"}
|
||
|
||
|
||
async def test_entity_discovery_job_lifecycle_and_proposed_isolation(multi_source_env, monkeypatch):
|
||
mes = multi_source_env["mes_name"]
|
||
erp = multi_source_env["erp_name"]
|
||
|
||
discovery_json = json.dumps(
|
||
{
|
||
"operations": [
|
||
{
|
||
"op": "upsert_entity",
|
||
"entity_key": "discovered_line",
|
||
"label": "自动发现产线",
|
||
"description": "跨 MES 与 ERP 的产线对象",
|
||
"aliases": ["产线"],
|
||
"reason": "跨源目录对齐：工单与结算产线",
|
||
},
|
||
{
|
||
"op": "upsert_binding",
|
||
"entity_key": "discovered_line",
|
||
"entity_label": "自动发现产线",
|
||
"source_name": mes,
|
||
"table_name": "dwd_workorder",
|
||
"column_name": "line_code",
|
||
"status": "confirmed", # 后台发现必须强制降为 proposed
|
||
"reason": "跨源目录对齐：MES 工单产线",
|
||
},
|
||
{
|
||
"op": "upsert_binding",
|
||
"entity_key": "discovered_line",
|
||
"entity_label": "自动发现产线",
|
||
"source_name": erp,
|
||
"table_name": "dwd_settlement",
|
||
"column_name": "workshop_line",
|
||
"reason": "跨源目录对齐：ERP 结算产线",
|
||
},
|
||
{
|
||
"op": "upsert_member",
|
||
"entity_key": "discovered_line",
|
||
"entity_label": "自动发现产线",
|
||
"value": "L1",
|
||
"label": "一号产线",
|
||
"source_codes": {mes: "L1", erp: "LINE_01"},
|
||
"reason": "文档对照",
|
||
},
|
||
]
|
||
},
|
||
ensure_ascii=False,
|
||
)
|
||
fake_model = _ScriptedModel(invoke_reply=f"```json\n{discovery_json}\n```")
|
||
|
||
async def _fake_resolve(_session):
|
||
return fake_model, "mock-discovery-model"
|
||
|
||
monkeypatch.setattr(catalog_assistant, "resolve_model", _fake_resolve)
|
||
|
||
transport = ASGITransport(app=app)
|
||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||
r = await client.post(
|
||
"/api/ontology/discovery-jobs",
|
||
json={
|
||
"source_ids": [multi_source_env["mes_id"], multi_source_env["erp_id"]],
|
||
"note": "重点扫描产线",
|
||
},
|
||
headers={"X-Actor": "builder_scan"},
|
||
)
|
||
assert r.status_code == 202, r.text
|
||
job_id = r.json()["id"]
|
||
|
||
# 轮询等待完成
|
||
final_job = None
|
||
for _ in range(40):
|
||
jr = await client.get(f"/api/ontology/discovery-jobs/{job_id}")
|
||
assert jr.status_code == 200
|
||
body = jr.json()
|
||
if body["status"] in ("succeeded", "failed", "cancelled"):
|
||
final_job = body
|
||
break
|
||
await asyncio.sleep(0.05)
|
||
|
||
assert final_job is not None
|
||
assert final_job["status"] == "succeeded", final_job
|
||
assert final_job["summary"]["created_entities"] == 1
|
||
assert final_job["summary"]["proposed_bindings"] == 2
|
||
assert final_job["summary"]["added_members"] == 1
|
||
|
||
ents = (await client.get("/api/ontology/entities")).json()["items"]
|
||
disc = next(e for e in ents if e["key"] == "discovered_line")
|
||
assert all(b["status"] == "proposed" and b["source"] == "ai" for b in disc["bindings"])
|
||
|
||
# 运行期 confirmed_bindings 必须为 0（状态门禁隔离）
|
||
assert entities.confirmed_bindings(disc) == []
|
||
|
||
|
||
async def test_entity_discovery_job_busy_and_cancel(multi_source_env, monkeypatch):
|
||
entered = asyncio.Event()
|
||
|
||
class _SlowModel:
|
||
async def ainvoke(self, _messages):
|
||
entered.set()
|
||
await asyncio.sleep(30)
|
||
return AIMessage(content='{"operations": []}')
|
||
|
||
async def _fake_resolve(_session):
|
||
return _SlowModel(), "mock-slow-model"
|
||
|
||
monkeypatch.setattr(catalog_assistant, "resolve_model", _fake_resolve)
|
||
|
||
transport = ASGITransport(app=app)
|
||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||
r1 = await client.post(
|
||
"/api/ontology/discovery-jobs",
|
||
json={"source_ids": [multi_source_env["mes_id"]]},
|
||
)
|
||
assert r1.status_code == 202
|
||
job_id = r1.json()["id"]
|
||
|
||
await asyncio.wait_for(entered.wait(), timeout=3.0)
|
||
|
||
# 并发再次发起应返回 409 discovery_job_busy
|
||
r2 = await client.post(
|
||
"/api/ontology/discovery-jobs",
|
||
json={"source_ids": [multi_source_env["mes_id"]]},
|
||
)
|
||
assert r2.status_code == 409
|
||
assert r2.json().get("code") == "discovery_job_busy"
|
||
|
||
# 主动停止扫描
|
||
rc = await client.post(f"/api/ontology/discovery-jobs/{job_id}/cancel")
|
||
assert rc.status_code == 200
|
||
assert rc.json()["status"] == "cancelled"
|
||
|
||
|
||
async def test_entity_assistant_source_scope_isolates_unrelated_domains(multi_source_env):
|
||
"""进入单个数据源目录页时，传入 source_ids 仅暴露所选数据源的表、查询权限与关联工作流，隔离其他业务域。"""
|
||
mes_id = multi_source_env["mes_id"]
|
||
mes_name = multi_source_env["mes_name"]
|
||
erp_id = multi_source_env["erp_id"]
|
||
erp_name = multi_source_env["erp_name"]
|
||
wf_mes_id = f"wf_mes_only_{uuid.uuid4().hex[:6]}"
|
||
wf_erp_id = f"wf_erp_only_{uuid.uuid4().hex[:6]}"
|
||
|
||
async with SessionLocal() as session:
|
||
# 构造一个仅归属 MES 域的实体与两个分别绑定不同数据源工具的工作流
|
||
await entities.create_entity(
|
||
session,
|
||
{
|
||
"key": "mes_only_line",
|
||
"label": "车间专属产线",
|
||
"bindings": [
|
||
{
|
||
"source_id": mes_id,
|
||
"source_name": mes_name,
|
||
"table_name": "dwd_workorder",
|
||
"column_name": "line_code",
|
||
"status": "confirmed",
|
||
"source": "human",
|
||
}
|
||
],
|
||
},
|
||
actor="tester",
|
||
)
|
||
session.add_all(
|
||
[
|
||
Workflow(
|
||
id=wf_mes_id,
|
||
name="车间工单良率",
|
||
version=1,
|
||
graph={
|
||
"nodes": [
|
||
{
|
||
"id": "in_m",
|
||
"type": "input",
|
||
"data": {"config": {"fields": [{"name": "scx_code", "label": "产线", "type": "string"}]}},
|
||
},
|
||
{
|
||
"id": "q_m",
|
||
"type": "tool",
|
||
"data": {"config": {"tool": f"db_query__{mes_name}", "args": {"sql": "SELECT 1"}}},
|
||
},
|
||
{
|
||
"id": "c_m",
|
||
"type": "metrics",
|
||
"data": {"config": {"metrics": [{"id": "m1", "name": "良率"}]}},
|
||
},
|
||
],
|
||
"edges": [],
|
||
},
|
||
),
|
||
Workflow(
|
||
id=wf_erp_id,
|
||
name="结算金额口径",
|
||
version=1,
|
||
graph={
|
||
"nodes": [
|
||
{
|
||
"id": "in_e",
|
||
"type": "input",
|
||
"data": {"config": {"fields": [{"name": "store_id", "label": "门店", "type": "string"}]}},
|
||
},
|
||
{
|
||
"id": "q_e",
|
||
"type": "tool",
|
||
"data": {"config": {"tool": f"db_query__{erp_name}", "args": {"sql": "SELECT 1"}}},
|
||
},
|
||
{
|
||
"id": "c_e",
|
||
"type": "metrics",
|
||
"data": {"config": {"metrics": [{"id": "e1", "name": "结算量"}]}},
|
||
},
|
||
],
|
||
"edges": [],
|
||
},
|
||
),
|
||
]
|
||
)
|
||
await session.commit()
|
||
|
||
try:
|
||
async with SessionLocal() as session:
|
||
ctx = await entity_assistant.prepare(session, source_ids=[erp_id])
|
||
|
||
# 1. 只加载 erp_name，不包含 mes_name
|
||
assert list(ctx.sources.keys()) == [erp_name.lower()]
|
||
assert ctx.queryable_sources == [erp_name]
|
||
|
||
# 2. 系统提示词不泄露 mes_name 的物理表名 dwd_workorder
|
||
assert "dwd_workorder" not in ctx.system
|
||
assert erp_name in ctx.system
|
||
|
||
# 3. inspect_sources 只返回 erp 的表与 erp 关联的口径工作流，不包含车间工单良率工作流
|
||
tools = entity_assistant._EntityTools(ctx)
|
||
inspected = tools.inspect("", "")
|
||
assert "dwd_settlement" in inspected
|
||
assert "dwd_workorder" not in inspected
|
||
assert wf_erp_id in inspected
|
||
assert wf_mes_id not in inspected
|
||
|
||
# 4. 对未纳入范围的 mes_name 发起只读查询会被立即拒绝，不尝试连接该库
|
||
ev, msg = await tools.query(mes_name, "SELECT * FROM dwd_workorder LIMIT 5")
|
||
assert ev is None
|
||
assert erp_name in msg
|
||
finally:
|
||
async with SessionLocal() as session:
|
||
await session.execute(delete(Workflow).where(Workflow.id.in_([wf_mes_id, wf_erp_id])))
|
||
await session.commit()
|
||
|
||
|
||
async def test_entity_discovery_streaming_progress_and_summary(multi_source_env, monkeypatch):
|
||
"""扫描发现任务支持 astream 流式输出并在 summary 中记录量化扫描统计与新增实体清单。"""
|
||
mes_id = multi_source_env["mes_id"]
|
||
mes_name = multi_source_env["mes_name"]
|
||
stream_payload = json.dumps(
|
||
{
|
||
"operations": [
|
||
{
|
||
"op": "upsert_entity",
|
||
"entity_key": "stream_line",
|
||
"label": "流式识别产线",
|
||
"reason": "目录语义对齐：产线编码",
|
||
},
|
||
{
|
||
"op": "upsert_binding",
|
||
"entity_key": "stream_line",
|
||
"entity_label": "流式识别产线",
|
||
"source_name": mes_name,
|
||
"table_name": "dwd_workorder",
|
||
"column_name": "line_code",
|
||
"reason": "目录语义对齐：产线字段",
|
||
},
|
||
]
|
||
},
|
||
ensure_ascii=False,
|
||
)
|
||
|
||
class _StreamingDiscoveryModel:
|
||
async def astream(self, _messages):
|
||
mid = len(stream_payload) // 2
|
||
yield AIMessageChunk(content=stream_payload[:mid])
|
||
await asyncio.sleep(0.02)
|
||
yield AIMessageChunk(content=stream_payload[mid:])
|
||
|
||
async def _fake_resolve(_session):
|
||
return _StreamingDiscoveryModel(), "mock-stream-model"
|
||
|
||
monkeypatch.setattr(catalog_assistant, "resolve_model", _fake_resolve)
|
||
|
||
transport = ASGITransport(app=app)
|
||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||
r = await client.post(
|
||
"/api/ontology/discovery-jobs",
|
||
json={"source_ids": [mes_id]},
|
||
)
|
||
assert r.status_code == 202
|
||
job_id = r.json()["id"]
|
||
|
||
final_job = None
|
||
for _ in range(40):
|
||
jr = await client.get(f"/api/ontology/discovery-jobs/{job_id}")
|
||
body = jr.json()
|
||
if body["status"] in ("succeeded", "failed", "cancelled"):
|
||
final_job = body
|
||
break
|
||
await asyncio.sleep(0.05)
|
||
|
||
assert final_job is not None
|
||
assert final_job["status"] == "succeeded", final_job
|
||
assert isinstance(final_job["started_at"], str) and (
|
||
final_job["started_at"].endswith("+00:00") or final_job["started_at"].endswith("Z")
|
||
)
|
||
assert isinstance(final_job["finished_at"], str) and (
|
||
final_job["finished_at"].endswith("+00:00") or final_job["finished_at"].endswith("Z")
|
||
)
|
||
summary = final_job["summary"]
|
||
assert summary["progress_pct"] == 100
|
||
assert summary["sources_scanned"] == 1
|
||
assert summary["source_names"] == [mes_name]
|
||
assert summary["tables_scanned"] >= 1
|
||
assert summary["ops_detected"] == 2
|
||
assert any(t["key"] == "stream_line" for t in summary["touched_entities"])
|
||
|
||
|
||
|