Spaces:
Running
Running
File size: 4,304 Bytes
c690006 ff16d8e c690006 fa696e8 c690006 fa696e8 c690006 fa696e8 c690006 fa696e8 c690006 fa696e8 c690006 ebc6429 c690006 ebc6429 fa696e8 c690006 ebc6429 fa696e8 ebc6429 fa696e8 c690006 ebc6429 c690006 ebc6429 c690006 ebc6429 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
"""Unit tests for HypothesisAgent."""
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
# Skip all tests if agent_framework not installed (optional dep)
pytest.importorskip("agent_framework")
from agent_framework import AgentRunResponse # noqa: E402
from src.agents.hypothesis_agent import HypothesisAgent # noqa: E402
from src.utils.models import ( # noqa: E402
Citation,
Evidence,
HypothesisAssessment,
MechanismHypothesis,
)
@pytest.fixture
def sample_evidence():
return [
Evidence(
content="Testosterone activates androgen receptors...",
citation=Citation(
source="pubmed",
title="Testosterone and Libido",
url="https://pubmed.ncbi.nlm.nih.gov/12345/",
date="2023",
),
)
]
@pytest.fixture
def mock_assessment():
return HypothesisAssessment(
hypotheses=[
MechanismHypothesis(
drug="Testosterone",
target="Androgen Receptor",
pathway="Dopamine modulation",
effect="Enhanced sexual desire in HSDD",
confidence=0.75,
search_suggestions=["testosterone libido mechanism", "HSDD treatment"],
)
],
primary_hypothesis=None,
knowledge_gaps=["Clinical trial data needed"],
recommended_searches=["testosterone HSDD clinical trial"],
)
@pytest.mark.asyncio
async def test_hypothesis_agent_generates_hypotheses(sample_evidence, mock_assessment):
"""HypothesisAgent should generate mechanistic hypotheses."""
store = {"current": sample_evidence, "hypotheses": []}
with patch("src.agents.hypothesis_agent.get_model") as mock_get_model:
with patch("src.agents.hypothesis_agent.Agent") as mock_agent_class:
mock_get_model.return_value = MagicMock() # Mock model
mock_result = MagicMock()
mock_result.output = mock_assessment
# pydantic-ai Agent returns an object with .output for structured output
mock_agent_class.return_value.run = AsyncMock(return_value=mock_result)
agent = HypothesisAgent(store)
response = await agent.run("testosterone libido")
assert isinstance(response, AgentRunResponse)
assert "Androgen" in response.messages[0].text
assert len(store["hypotheses"]) == 1
assert store["hypotheses"][0].drug == "Testosterone"
@pytest.mark.asyncio
async def test_hypothesis_agent_no_evidence():
"""HypothesisAgent should handle empty evidence gracefully."""
store = {"current": [], "hypotheses": []}
# No need to mock Agent/get_model - empty evidence returns early
agent = HypothesisAgent(store)
response = await agent.run("test query")
assert "No evidence" in response.messages[0].text
assert len(store["hypotheses"]) == 0
@pytest.mark.asyncio
async def test_hypothesis_agent_uses_embeddings(sample_evidence, mock_assessment):
"""HypothesisAgent should pass embeddings to prompt formatter."""
store = {"current": sample_evidence, "hypotheses": []}
mock_embeddings = MagicMock()
with patch("src.agents.hypothesis_agent.get_model") as mock_get_model:
with patch("src.agents.hypothesis_agent.Agent") as mock_agent_class:
# Mock format_hypothesis_prompt to check if embeddings were passed
with patch("src.agents.hypothesis_agent.format_hypothesis_prompt") as mock_format:
mock_get_model.return_value = MagicMock() # Mock model
mock_format.return_value = "Prompt"
mock_result = MagicMock()
mock_result.output = mock_assessment
mock_agent_class.return_value.run = AsyncMock(return_value=mock_result)
agent = HypothesisAgent(store, embedding_service=mock_embeddings)
await agent.run("query")
mock_format.assert_called_once()
_args, kwargs = mock_format.call_args
assert kwargs["embeddings"] == mock_embeddings
assert _args[0] == "query" # First positional arg is query
assert _args[1] == sample_evidence # Second positional arg is evidence
|