480 lines
17 KiB
Python
480 lines
17 KiB
Python
import asyncio
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from apps.web.vue3 import main as api
|
|
from gfa_core.mas.agents.base import AgentMessage, AgentState
|
|
from gfa_core.mas.agents.intent_v2 import IntentAgentV2
|
|
from gfa_core.mas.agents.planning_v2 import PlanningAgentV2
|
|
from gfa_core.mas.context_index import ProjectContextIndex
|
|
from gfa_core.mas.core.fast_intent import fast_intent_recognition
|
|
from gfa_core.mas.stage_catalog import STAGE_DEPENDENCIES, complete_dependencies
|
|
from gfa_core.mas.utils.intent_parser import parse_intent_from_query, resolve_gwas_intent
|
|
|
|
|
|
GWAS_FMT_COMMANDS = [
|
|
"tassel转换",
|
|
"使用 TASSEL 进行 VCF 格式转换",
|
|
"使用 TASSEL 将上传的 VCF 转换为 HapMap",
|
|
"GWAS格式转换",
|
|
"VCF转HapMap",
|
|
"请对外部 VCF 进行 GWAS 格式转换，使用 TASSEL 排序并转换为 HapMap",
|
|
]
|
|
|
|
GENERIC_GWAS_COMMANDS = [
|
|
"gwas",
|
|
"执行gwas分析",
|
|
"gwas流程",
|
|
"运行完整 GWAS",
|
|
]
|
|
|
|
FULL_GWAS_CHAIN = complete_dependencies(["gwas_plot", "gwas_emplot", "gwas_glm"])
|
|
|
|
|
|
class _ShortenedGwasOrchestrator:
|
|
def __init__(self) -> None:
|
|
self.initial_state = None
|
|
|
|
async def plan(self, initial_state: dict) -> dict:
|
|
self.initial_state = initial_state
|
|
return {
|
|
"status": "completed",
|
|
"mode": "orchestrator_v2",
|
|
"final_output": {
|
|
"analysis_plan": {
|
|
"analysis_type": "gene_family",
|
|
"stages": ["gwas_emmax", "gwas_emplot"],
|
|
"stages_detail": [
|
|
api._stage_detail_defaults("gwas_emmax"),
|
|
api._stage_detail_defaults("gwas_emplot"),
|
|
],
|
|
"total_stages": 2,
|
|
"description": "执行gwas_glm分析",
|
|
"params": {"prefix": ""},
|
|
},
|
|
"intent": {
|
|
"stages": ["gwas_emmax", "gwas_emplot"],
|
|
"primary_stage": "gwas_glm",
|
|
"description": "执行gwas_glm分析",
|
|
"params": {},
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize("query", GENERIC_GWAS_COMMANDS)
|
|
def test_agent_plan_keeps_bare_gwas_full_chain_authoritative(monkeypatch, query: str) -> None:
|
|
orchestrator = _ShortenedGwasOrchestrator()
|
|
monkeypatch.setattr(api, "orchestrator", orchestrator)
|
|
monkeypatch.setattr(api, "MAS_AVAILABLE", True)
|
|
|
|
response = asyncio.run(api.agent_plan(api.AgentPlanRequest(message=query)))
|
|
|
|
assert orchestrator.initial_state["results"]["intent"]["stages"] == FULL_GWAS_CHAIN
|
|
assert response["success"] is True
|
|
assert response["intent"]["stages"] == FULL_GWAS_CHAIN
|
|
assert response["plan"]["stages"] == FULL_GWAS_CHAIN
|
|
assert response["plan"]["all_stages"] == FULL_GWAS_CHAIN
|
|
assert [item["id"] for item in response["plan"]["stages_detail"]] == FULL_GWAS_CHAIN
|
|
assert [item["id"] for item in response["plan"]["all_stages_detail"]] == FULL_GWAS_CHAIN
|
|
assert {
|
|
item["id"]: item["dependencies"]
|
|
for item in response["plan"]["stages_detail"]
|
|
} == {
|
|
stage: STAGE_DEPENDENCIES[stage]
|
|
for stage in FULL_GWAS_CHAIN
|
|
}
|
|
assert {
|
|
item["id"]: item["dependencies"]
|
|
for item in response["plan"]["all_stages_detail"]
|
|
} == {
|
|
stage: STAGE_DEPENDENCIES[stage]
|
|
for stage in FULL_GWAS_CHAIN
|
|
}
|
|
assert response["plan"]["analysis_type"] == "gwas"
|
|
assert response["plan"]["prefix"] == "cohort"
|
|
assert response["plan"]["params"]["prefix"] == "cohort"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("query", "target"),
|
|
[
|
|
("使用 TASSEL 将上传的 VCF 转换为 HapMap", "gwas_fmt"),
|
|
("使用 TASSEL 执行 GWAS GLM 分析", "gwas_glm"),
|
|
("执行 EMMAX GWAS 分析", "gwas_emmax"),
|
|
("绘制 EMMAX Manhattan 和 QQ 图", "gwas_emplot"),
|
|
("绘制 GWAS MLM Manhattan 和 QQ 图", "gwas_plot"),
|
|
],
|
|
)
|
|
def test_agent_plan_does_not_expand_explicit_gwas_stage(monkeypatch, query: str, target: str) -> None:
|
|
expected = complete_dependencies([target])
|
|
|
|
class _ExactOrchestrator:
|
|
async def plan(self, initial_state: dict) -> dict:
|
|
intent = initial_state["results"].get("intent") or {
|
|
"stages": expected,
|
|
"primary_stage": target,
|
|
"description": f"执行{target}分析",
|
|
"params": {"prefix": "cohort"},
|
|
"prefix": "cohort",
|
|
}
|
|
stage_details = []
|
|
for stage in expected:
|
|
detail = api._stage_detail_defaults(stage)
|
|
stage_details.append(detail)
|
|
return {
|
|
"status": "completed",
|
|
"mode": "orchestrator_v2",
|
|
"final_output": {
|
|
"analysis_plan": {
|
|
"analysis_type": "gwas",
|
|
"stages": expected,
|
|
"stages_detail": stage_details,
|
|
"total_stages": len(expected),
|
|
"description": f"执行{target}分析",
|
|
"params": {"prefix": "cohort"},
|
|
},
|
|
"intent": intent,
|
|
},
|
|
}
|
|
|
|
monkeypatch.setattr(api, "orchestrator", _ExactOrchestrator())
|
|
monkeypatch.setattr(api, "MAS_AVAILABLE", True)
|
|
|
|
response = asyncio.run(api.agent_plan(api.AgentPlanRequest(message=query)))
|
|
|
|
assert response["plan"]["all_stages"] == expected
|
|
assert response["intent"]["stages"] == expected
|
|
assert response["plan"]["all_stages"] != FULL_GWAS_CHAIN or expected == FULL_GWAS_CHAIN
|
|
assert {
|
|
item["id"]: item["dependencies"]
|
|
for item in response["plan"]["all_stages_detail"]
|
|
} == {
|
|
stage: STAGE_DEPENDENCIES[stage]
|
|
for stage in expected
|
|
}
|
|
|
|
|
|
def _recognize_with_public_agent(query: str) -> dict:
|
|
agent = IntentAgentV2(object())
|
|
state = AgentState(session_id="gwas-intent-test", job_id="job_test")
|
|
message = AgentMessage(
|
|
sender="router",
|
|
receiver="intent",
|
|
message_type="recognize",
|
|
payload={"query": query, "workspace": "/tmp/job_test"},
|
|
)
|
|
|
|
result = asyncio.run(agent.execute(state, message))
|
|
|
|
return result.results["results"]["intent"]
|
|
|
|
|
|
@pytest.mark.parametrize("query", GWAS_FMT_COMMANDS)
|
|
def test_gwas_fmt_commands_use_deterministic_parser(query: str) -> None:
|
|
parsed = parse_intent_from_query(query)
|
|
|
|
assert parsed["stages"] == ["gwas_fmt"]
|
|
assert parsed["primary_stage"] == "gwas_fmt"
|
|
assert parsed["task_kind"] == "pipeline"
|
|
|
|
|
|
@pytest.mark.parametrize("query", GWAS_FMT_COMMANDS)
|
|
def test_gwas_fmt_commands_use_fast_intent_path(query: str) -> None:
|
|
parsed = fast_intent_recognition(query)
|
|
|
|
assert parsed is not None
|
|
assert parsed["stage"] == "gwas_fmt"
|
|
assert parsed["stages"] == ["gwas_fmt"]
|
|
|
|
|
|
@pytest.mark.parametrize("query", GWAS_FMT_COMMANDS)
|
|
def test_gwas_fmt_commands_use_public_intent_agent(query: str) -> None:
|
|
parsed = _recognize_with_public_agent(query)
|
|
|
|
assert parsed["stages"] == ["gwas_fmt"]
|
|
assert parsed["primary_stage"] == "gwas_fmt"
|
|
|
|
|
|
@pytest.mark.parametrize("query", GENERIC_GWAS_COMMANDS)
|
|
def test_generic_gwas_commands_plan_full_chain_in_deterministic_parser(query: str) -> None:
|
|
parsed = parse_intent_from_query(query)
|
|
|
|
assert parsed["stages"] == FULL_GWAS_CHAIN
|
|
assert parsed["primary_stage"] == "gwas_glm"
|
|
assert parsed["need_clarify"] is False
|
|
assert "蛋白" not in (parsed.get("message") or "")
|
|
|
|
|
|
@pytest.mark.parametrize("query", GENERIC_GWAS_COMMANDS)
|
|
def test_generic_gwas_commands_use_full_chain_fast_intent_path(query: str) -> None:
|
|
parsed = fast_intent_recognition(query)
|
|
assert parsed is not None
|
|
assert parsed["stage"] == "gwas_glm"
|
|
assert parsed["stages"] == FULL_GWAS_CHAIN
|
|
|
|
|
|
@pytest.mark.parametrize("query", GENERIC_GWAS_COMMANDS)
|
|
def test_generic_gwas_commands_use_public_agent_pipeline_response(query: str) -> None:
|
|
parsed = _recognize_with_public_agent(query)
|
|
|
|
assert parsed["stages"] == FULL_GWAS_CHAIN
|
|
assert parsed["primary_stage"] == "gwas_glm"
|
|
assert parsed["task_kind"] == "pipeline"
|
|
assert parsed["taskType"] == "pipeline"
|
|
assert "蛋白" not in str(parsed)
|
|
|
|
|
|
def test_gwas_fmt_plan_blocks_when_no_job_vcf_exists(tmp_path: Path) -> None:
|
|
input_dir = tmp_path / "input"
|
|
workspace = tmp_path / "workspace"
|
|
input_dir.mkdir()
|
|
workspace.mkdir()
|
|
agent = PlanningAgentV2.__new__(PlanningAgentV2)
|
|
|
|
issues = agent._check_input_files(
|
|
"gwas_fmt",
|
|
input_dir,
|
|
workspace,
|
|
"cohort",
|
|
context_index=None,
|
|
planned_stages=["gwas_fmt"],
|
|
)
|
|
|
|
assert any(issue["severity"] == "error" for issue in issues)
|
|
assert any(issue.get("code") == "missing_external_vcf" for issue in issues)
|
|
|
|
|
|
def _empty_project_context(tmp_path: Path) -> ProjectContextIndex:
|
|
workspace = tmp_path / "workspace" / "job_test"
|
|
return ProjectContextIndex(
|
|
job_id="job_test",
|
|
workspace=str(workspace),
|
|
input_dir=str(tmp_path / "input" / "job_test"),
|
|
final_dir=str(workspace / "final"),
|
|
logs_dir=str(workspace / "logs"),
|
|
)
|
|
|
|
|
|
def _planning_input_issues(tmp_path: Path, stage: str, planned_stages: list[str]) -> list[dict]:
|
|
context = _empty_project_context(tmp_path)
|
|
agent = PlanningAgentV2.__new__(PlanningAgentV2)
|
|
return agent._check_input_files(
|
|
stage,
|
|
Path(context.input_dir),
|
|
Path(context.workspace),
|
|
"cohort",
|
|
context_index=context,
|
|
planned_stages=planned_stages,
|
|
)
|
|
|
|
|
|
def _root_record(path: Path, role: str) -> dict:
|
|
return {
|
|
"role": role,
|
|
"path": str(path.resolve()),
|
|
"size": path.stat().st_size,
|
|
"sha256": hashlib.sha256(path.read_bytes()).hexdigest(),
|
|
}
|
|
|
|
|
|
def _write_covariate_root_pair(tmp_path: Path, *, audited: bool) -> tuple[Path, Path]:
|
|
input_dir = tmp_path / "input" / "job_test"
|
|
final_dir = tmp_path / "workspace" / "job_test" / "final"
|
|
input_dir.mkdir(parents=True)
|
|
final_dir.mkdir(parents=True)
|
|
q_file = input_dir / "cohort.2.Q"
|
|
fam_file = input_dir / "cohort.fam"
|
|
q_file.write_text("0.4 0.6\n", encoding="utf-8")
|
|
fam_file.write_text("F1 I1 0 0 0 -9\n", encoding="utf-8")
|
|
if audited:
|
|
(final_dir / "cohort.pop_struct.manifest.json").write_text(
|
|
json.dumps({
|
|
"stage": "pop_struct",
|
|
"status": "success",
|
|
"genotype_scope_id": "genotype:" + "a" * 64,
|
|
"artifacts": {
|
|
"admixture_q": _root_record(q_file, "admixture_q"),
|
|
"sample_fam": _root_record(fam_file, "sample_fam"),
|
|
},
|
|
}),
|
|
encoding="utf-8",
|
|
)
|
|
return q_file, fam_file
|
|
|
|
|
|
def test_gwas_cov_planning_resolves_audited_q_fam_pair_without_generic_artifact_errors(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
q_file, fam_file = _write_covariate_root_pair(tmp_path, audited=True)
|
|
# Distinct scan candidates must not override the authoritative upstream Manifest.
|
|
(tmp_path / "input" / "job_test" / "other.Q").write_text("0.2 0.8\n", encoding="utf-8")
|
|
(tmp_path / "input" / "job_test" / "other.fam").write_text(
|
|
"F2 I2 0 0 0 -9\n", encoding="utf-8",
|
|
)
|
|
|
|
issues = _planning_input_issues(tmp_path, "gwas_cov", FULL_GWAS_CHAIN)
|
|
|
|
assert not any(issue["severity"] == "error" for issue in issues)
|
|
resolved = {issue.get("role"): issue for issue in issues if issue.get("category") == "input_resolved"}
|
|
assert set(resolved) == {"admixture_q", "sample_fam"}
|
|
assert resolved["admixture_q"]["suggestion"] == str(q_file.resolve())
|
|
assert resolved["sample_fam"]["suggestion"] == str(fam_file.resolve())
|
|
|
|
|
|
def test_gwas_cov_planning_preserves_unverified_lineage_error_instead_of_missing_roles(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
_write_covariate_root_pair(tmp_path, audited=False)
|
|
|
|
issues = _planning_input_issues(tmp_path, "gwas_cov", FULL_GWAS_CHAIN)
|
|
|
|
assert [issue.get("code") for issue in issues] == ["unverified_covariate_lineage"]
|
|
assert issues[0]["severity"] == "error"
|
|
assert issues[0]["resolver"]["needs_confirmation"] is False
|
|
assert "lineage" in issues[0]["message"]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("case", "expected_code", "needs_confirmation"),
|
|
[
|
|
("invalid_manifest", "invalid_upstream_manifest", False),
|
|
("invalid_binding", "invalid_external_binding", False),
|
|
("ambiguous", "ambiguous_root_input", True),
|
|
],
|
|
)
|
|
def test_gwas_cov_planning_preserves_root_pair_fail_closed_issues(
|
|
tmp_path: Path,
|
|
case: str,
|
|
expected_code: str,
|
|
needs_confirmation: bool,
|
|
) -> None:
|
|
_write_covariate_root_pair(tmp_path, audited=case == "invalid_manifest")
|
|
workspace = tmp_path / "workspace" / "job_test"
|
|
if case == "invalid_manifest":
|
|
manifest = workspace / "final" / "cohort.pop_struct.manifest.json"
|
|
payload = json.loads(manifest.read_text(encoding="utf-8"))
|
|
payload["status"] = "failed"
|
|
manifest.write_text(json.dumps(payload), encoding="utf-8")
|
|
elif case == "invalid_binding":
|
|
(workspace / "external_inputs.manifest.json").write_text(
|
|
json.dumps({"manifest_version": "1.0", "job_id": "job_test", "inputs": "malformed"}),
|
|
encoding="utf-8",
|
|
)
|
|
else:
|
|
(tmp_path / "input" / "job_test" / "other.Q").write_text("0.2 0.8\n", encoding="utf-8")
|
|
|
|
issues = _planning_input_issues(tmp_path, "gwas_cov", FULL_GWAS_CHAIN)
|
|
|
|
matching = [issue for issue in issues if issue.get("code") == expected_code]
|
|
assert matching
|
|
assert all(issue["needs_confirmation"] is needs_confirmation for issue in matching)
|
|
assert all(issue["message"] in [item["message"] for item in issue["resolver"]["issues"]] for issue in matching)
|
|
assert all(issue["severity"] == ("warning" if needs_confirmation else "error") for issue in matching)
|
|
|
|
|
|
def test_full_gwas_dag_treats_derived_roles_as_dependency_provided(tmp_path: Path) -> None:
|
|
expected = {
|
|
("gwas_kinship", "sorted_vcf"): "gwas_fmt",
|
|
("gwas_mlm", "sorted_vcf"): "gwas_fmt",
|
|
("gwas_mlm", "kinship_matrix"): "gwas_kinship",
|
|
("gwas_mlm", "trait_table"): "gwas_trait",
|
|
("gwas_mlm", "covariate_table"): "gwas_cov",
|
|
("gwas_plot", "mlm_pvalues"): "gwas_mlm",
|
|
("gwas_plot", "cmlm_pvalues"): "gwas_mlm",
|
|
("gwas_emmax", "sorted_vcf"): "gwas_fmt",
|
|
("gwas_emmax", "trait_table"): "gwas_trait",
|
|
("gwas_emplot", "emmax_pvalues"): "gwas_emmax",
|
|
("gwas_glm", "sorted_vcf"): "gwas_fmt",
|
|
("gwas_glm", "trait_table"): "gwas_trait",
|
|
("gwas_glm", "covariate_table"): "gwas_cov",
|
|
}
|
|
|
|
for (stage, role), producer in expected.items():
|
|
issues = _planning_input_issues(tmp_path, stage, FULL_GWAS_CHAIN)
|
|
matching = [issue for issue in issues if role in issue.get("message", "")]
|
|
assert matching, (stage, role, issues)
|
|
assert not any(issue["severity"] == "error" for issue in matching)
|
|
assert any(
|
|
issue["severity"] == "warning"
|
|
and issue.get("category") == "dependency"
|
|
and producer in issue.get("message", "")
|
|
for issue in matching
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("stage", "root_roles"),
|
|
[
|
|
("gwas_trait", {"phenotype_table"}),
|
|
("gwas_cov", {"admixture_q", "sample_fam"}),
|
|
],
|
|
)
|
|
def test_full_gwas_dag_keeps_unresolved_root_inputs_blocking(
|
|
tmp_path: Path,
|
|
stage: str,
|
|
root_roles: set[str],
|
|
) -> None:
|
|
issues = _planning_input_issues(tmp_path, stage, FULL_GWAS_CHAIN)
|
|
|
|
if stage == "gwas_cov":
|
|
assert any(
|
|
issue["severity"] == "error" and issue.get("code") == "missing_root_input"
|
|
for issue in issues
|
|
)
|
|
return
|
|
for role in root_roles:
|
|
assert any(
|
|
issue["severity"] == "error" and role in issue.get("message", "")
|
|
for issue in issues
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("query", "stage", "derived_roles"),
|
|
[
|
|
("运行 EMMAX", "gwas_emmax", {"sorted_vcf", "trait_table"}),
|
|
("绘制 EMMAX Manhattan 和 QQ 图", "gwas_emplot", {"emmax_pvalues"}),
|
|
],
|
|
)
|
|
def test_single_gwas_consumer_intent_includes_its_artifact_producers(
|
|
tmp_path: Path,
|
|
query: str,
|
|
stage: str,
|
|
derived_roles: set[str],
|
|
) -> None:
|
|
planned_stages = resolve_gwas_intent(query)["stages"]
|
|
issues = _planning_input_issues(tmp_path, stage, planned_stages)
|
|
|
|
for role in derived_roles:
|
|
assert any(
|
|
issue["severity"] == "warning"
|
|
and issue.get("category") == "dependency"
|
|
and role in issue.get("message", "")
|
|
for issue in issues
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("stage", "derived_roles"),
|
|
[
|
|
("gwas_emmax", {"sorted_vcf", "trait_table"}),
|
|
("gwas_emplot", {"emmax_pvalues"}),
|
|
],
|
|
)
|
|
def test_gwas_consumer_without_planned_producer_still_blocks(
|
|
tmp_path: Path,
|
|
stage: str,
|
|
derived_roles: set[str],
|
|
) -> None:
|
|
issues = _planning_input_issues(tmp_path, stage, [stage])
|
|
|
|
for role in derived_roles:
|
|
assert any(
|
|
issue["severity"] == "error" and role in issue.get("message", "")
|
|
for issue in issues
|
|
)
|