73 lines
2.6 KiB
Python
73 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
from gfa_core.mas.agents.base import AgentMessage, AgentState
|
|
from gfa_core.mas.agents.intent_v2 import IntentAgentV2
|
|
from gfa_core.mas.core.fast_intent import fast_intent_recognition
|
|
from gfa_core.mas.stage_catalog import complete_dependencies
|
|
from gfa_core.mas.utils.intent_parser import parse_intent_from_query
|
|
|
|
|
|
FULL_TARGETS = ["gwas_plot", "gwas_emplot", "gwas_glm"]
|
|
FULL_CHAIN = complete_dependencies(FULL_TARGETS)
|
|
|
|
CASES = [
|
|
("使用 TASSEL 计算亲缘关系矩阵", "gwas_kinship"),
|
|
("准备 GWAS 表型数据", "gwas_trait"),
|
|
("根据 Q 矩阵和 FAM 准备 GWAS 协变量", "gwas_cov"),
|
|
("使用 TASSEL 执行 GWAS MLM 分析", "gwas_mlm"),
|
|
("绘制 GWAS MLM Manhattan 和 QQ 图", "gwas_plot"),
|
|
("执行 EMMAX GWAS 分析", "gwas_emmax"),
|
|
("绘制 EMMAX Manhattan 和 QQ 图", "gwas_emplot"),
|
|
("使用 TASSEL 执行 GWAS GLM 分析", "gwas_glm"),
|
|
]
|
|
|
|
|
|
def _public(query: str) -> dict:
|
|
agent = IntentAgentV2(object())
|
|
state = AgentState(session_id="gwas-chain-intent", 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", "执行gwas分析", "运行完整 GWAS", "gwas流程"])
|
|
def test_full_gwas_intent_is_the_complete_dependency_dag(query: str) -> None:
|
|
parsed = parse_intent_from_query(query)
|
|
fast = fast_intent_recognition(query)
|
|
public = _public(query)
|
|
|
|
assert parsed["stages"] == FULL_CHAIN
|
|
assert parsed["primary_stage"] == "gwas_glm"
|
|
assert parsed["need_clarify"] is False
|
|
assert fast is not None
|
|
assert fast["stages"] == FULL_CHAIN
|
|
assert fast["stage"] == "gwas_glm"
|
|
assert public["stages"] == FULL_CHAIN
|
|
assert public["primary_stage"] == "gwas_glm"
|
|
assert "蛋白" not in str((parsed, fast, public))
|
|
|
|
|
|
@pytest.mark.parametrize(("query", "target"), CASES)
|
|
def test_specific_gwas_intents_share_one_target_and_dependency_closure(query: str, target: str) -> None:
|
|
expected = complete_dependencies([target])
|
|
parsed = parse_intent_from_query(query)
|
|
fast = fast_intent_recognition(query)
|
|
public = _public(query)
|
|
|
|
assert parsed["stages"] == expected
|
|
assert parsed["primary_stage"] == target
|
|
assert fast is not None
|
|
assert fast["stages"] == expected
|
|
assert fast["stage"] == target
|
|
assert public["stages"] == expected
|
|
assert public["primary_stage"] == target
|