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 )