aiagents4pharma 1.28.0__py3-none-any.whl → 1.30.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- aiagents4pharma/talk2scholars/agents/__init__.py +1 -0
- aiagents4pharma/talk2scholars/agents/main_agent.py +35 -209
- aiagents4pharma/talk2scholars/agents/paper_download_agent.py +86 -0
- aiagents4pharma/talk2scholars/agents/s2_agent.py +10 -6
- aiagents4pharma/talk2scholars/agents/zotero_agent.py +12 -6
- aiagents4pharma/talk2scholars/configs/agents/talk2scholars/main_agent/default.yaml +2 -48
- aiagents4pharma/talk2scholars/configs/agents/talk2scholars/paper_download_agent/__init__.py +3 -0
- aiagents4pharma/talk2scholars/configs/agents/talk2scholars/s2_agent/default.yaml +5 -28
- aiagents4pharma/talk2scholars/configs/agents/talk2scholars/zotero_agent/default.yaml +5 -21
- aiagents4pharma/talk2scholars/configs/config.yaml +3 -0
- aiagents4pharma/talk2scholars/configs/tools/__init__.py +1 -0
- aiagents4pharma/talk2scholars/configs/tools/download_arxiv_paper/__init__.py +3 -0
- aiagents4pharma/talk2scholars/configs/tools/multi_paper_recommendation/default.yaml +1 -1
- aiagents4pharma/talk2scholars/configs/tools/search/default.yaml +1 -1
- aiagents4pharma/talk2scholars/configs/tools/single_paper_recommendation/default.yaml +1 -1
- aiagents4pharma/talk2scholars/configs/tools/zotero_read/default.yaml +42 -1
- aiagents4pharma/talk2scholars/configs/tools/zotero_write/__inti__.py +3 -0
- aiagents4pharma/talk2scholars/state/state_talk2scholars.py +1 -0
- aiagents4pharma/talk2scholars/tests/test_main_agent.py +186 -111
- aiagents4pharma/talk2scholars/tests/test_paper_download_agent.py +142 -0
- aiagents4pharma/talk2scholars/tests/test_paper_download_tools.py +154 -0
- aiagents4pharma/talk2scholars/tests/test_s2_display.py +74 -0
- aiagents4pharma/talk2scholars/tests/test_s2_multi.py +282 -0
- aiagents4pharma/talk2scholars/tests/test_s2_query.py +78 -0
- aiagents4pharma/talk2scholars/tests/test_s2_retrieve.py +65 -0
- aiagents4pharma/talk2scholars/tests/test_s2_search.py +266 -0
- aiagents4pharma/talk2scholars/tests/test_s2_single.py +274 -0
- aiagents4pharma/talk2scholars/tests/test_zotero_path.py +57 -0
- aiagents4pharma/talk2scholars/tests/test_zotero_read.py +412 -0
- aiagents4pharma/talk2scholars/tests/test_zotero_write.py +626 -0
- aiagents4pharma/talk2scholars/tools/paper_download/__init__.py +17 -0
- aiagents4pharma/talk2scholars/tools/paper_download/abstract_downloader.py +43 -0
- aiagents4pharma/talk2scholars/tools/paper_download/arxiv_downloader.py +108 -0
- aiagents4pharma/talk2scholars/tools/paper_download/download_arxiv_input.py +60 -0
- aiagents4pharma/talk2scholars/tools/s2/multi_paper_rec.py +50 -34
- aiagents4pharma/talk2scholars/tools/s2/retrieve_semantic_scholar_paper_id.py +8 -8
- aiagents4pharma/talk2scholars/tools/s2/search.py +36 -23
- aiagents4pharma/talk2scholars/tools/s2/single_paper_rec.py +44 -38
- aiagents4pharma/talk2scholars/tools/zotero/__init__.py +2 -0
- aiagents4pharma/talk2scholars/tools/zotero/utils/__init__.py +5 -0
- aiagents4pharma/talk2scholars/tools/zotero/utils/zotero_path.py +63 -0
- aiagents4pharma/talk2scholars/tools/zotero/zotero_read.py +64 -19
- aiagents4pharma/talk2scholars/tools/zotero/zotero_write.py +247 -0
- {aiagents4pharma-1.28.0.dist-info → aiagents4pharma-1.30.0.dist-info}/METADATA +6 -5
- {aiagents4pharma-1.28.0.dist-info → aiagents4pharma-1.30.0.dist-info}/RECORD +48 -30
- aiagents4pharma/talk2scholars/tests/test_call_s2.py +0 -100
- aiagents4pharma/talk2scholars/tests/test_call_zotero.py +0 -94
- aiagents4pharma/talk2scholars/tests/test_s2_tools.py +0 -355
- aiagents4pharma/talk2scholars/tests/test_zotero_tool.py +0 -171
- {aiagents4pharma-1.28.0.dist-info → aiagents4pharma-1.30.0.dist-info}/LICENSE +0 -0
- {aiagents4pharma-1.28.0.dist-info → aiagents4pharma-1.30.0.dist-info}/WHEEL +0 -0
- {aiagents4pharma-1.28.0.dist-info → aiagents4pharma-1.30.0.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,74 @@
|
|
1
|
+
"""
|
2
|
+
Unit tests for S2 tools functionality.
|
3
|
+
"""
|
4
|
+
|
5
|
+
# pylint: disable=redefined-outer-name
|
6
|
+
import pytest
|
7
|
+
from langgraph.types import Command
|
8
|
+
from ..tools.s2.display_results import (
|
9
|
+
display_results,
|
10
|
+
NoPapersFoundError as raised_error,
|
11
|
+
)
|
12
|
+
|
13
|
+
|
14
|
+
@pytest.fixture
|
15
|
+
def initial_state():
|
16
|
+
"""Provides an empty initial state for tests."""
|
17
|
+
return {"papers": {}, "multi_papers": {}}
|
18
|
+
|
19
|
+
|
20
|
+
# Fixed test data for deterministic results
|
21
|
+
MOCK_SEARCH_RESPONSE = {
|
22
|
+
"data": [
|
23
|
+
{
|
24
|
+
"paperId": "123",
|
25
|
+
"title": "Machine Learning Basics",
|
26
|
+
"abstract": "An introduction to ML",
|
27
|
+
"year": 2023,
|
28
|
+
"citationCount": 100,
|
29
|
+
"url": "https://example.com/paper1",
|
30
|
+
"authors": [{"name": "Test Author"}],
|
31
|
+
}
|
32
|
+
]
|
33
|
+
}
|
34
|
+
|
35
|
+
MOCK_STATE_PAPER = {
|
36
|
+
"123": {
|
37
|
+
"Title": "Machine Learning Basics",
|
38
|
+
"Abstract": "An introduction to ML",
|
39
|
+
"Year": 2023,
|
40
|
+
"Citation Count": 100,
|
41
|
+
"URL": "https://example.com/paper1",
|
42
|
+
}
|
43
|
+
}
|
44
|
+
|
45
|
+
|
46
|
+
class TestS2Tools:
|
47
|
+
"""Unit tests for individual S2 tools"""
|
48
|
+
|
49
|
+
def test_display_results_empty_state(self, initial_state):
|
50
|
+
"""Verifies display_results tool behavior when state is empty and raises an exception"""
|
51
|
+
with pytest.raises(
|
52
|
+
raised_error,
|
53
|
+
match="No papers found. A search/rec needs to be performed first.",
|
54
|
+
):
|
55
|
+
display_results.invoke({"state": initial_state, "tool_call_id": "test123"})
|
56
|
+
|
57
|
+
def test_display_results_shows_papers(self, initial_state):
|
58
|
+
"""Verifies display_results tool correctly returns papers from state"""
|
59
|
+
state = initial_state.copy()
|
60
|
+
state["last_displayed_papers"] = "papers"
|
61
|
+
state["papers"] = MOCK_STATE_PAPER
|
62
|
+
|
63
|
+
result = display_results.invoke(
|
64
|
+
input={"state": state, "tool_call_id": "test123"}
|
65
|
+
)
|
66
|
+
|
67
|
+
assert isinstance(result, Command) # Expect a Command object
|
68
|
+
assert isinstance(result.update, dict) # Ensure update is a dictionary
|
69
|
+
assert "messages" in result.update
|
70
|
+
assert len(result.update["messages"]) == 1
|
71
|
+
assert (
|
72
|
+
"1 papers found. Papers are attached as an artifact."
|
73
|
+
in result.update["messages"][0].content
|
74
|
+
)
|
@@ -0,0 +1,282 @@
|
|
1
|
+
"""
|
2
|
+
Unit tests for S2 tools functionality.
|
3
|
+
"""
|
4
|
+
|
5
|
+
import json
|
6
|
+
from types import SimpleNamespace
|
7
|
+
import pytest
|
8
|
+
import requests
|
9
|
+
from langgraph.types import Command
|
10
|
+
from langchain_core.messages import ToolMessage
|
11
|
+
import hydra
|
12
|
+
from aiagents4pharma.talk2scholars.tools.s2.multi_paper_rec import (
|
13
|
+
get_multi_paper_recommendations,
|
14
|
+
)
|
15
|
+
|
16
|
+
# --- Dummy Hydra Config Setup ---
|
17
|
+
|
18
|
+
|
19
|
+
class DummyHydraContext:
|
20
|
+
"""dummy context manager for mocking Hydra's initialize and compose functions."""
|
21
|
+
|
22
|
+
def __enter__(self):
|
23
|
+
"""enter function that returns None."""
|
24
|
+
return None
|
25
|
+
|
26
|
+
def __exit__(self, exc_type, exc_val, traceback):
|
27
|
+
"""exit function that does nothing."""
|
28
|
+
return None
|
29
|
+
|
30
|
+
|
31
|
+
# Create a dummy configuration that mimics the expected hydra config.
|
32
|
+
dummy_config = SimpleNamespace(
|
33
|
+
tools=SimpleNamespace(
|
34
|
+
multi_paper_recommendation=SimpleNamespace(
|
35
|
+
api_endpoint="http://dummy.endpoint/multi",
|
36
|
+
headers={"Content-Type": "application/json"},
|
37
|
+
api_fields=["paperId", "title", "authors"],
|
38
|
+
request_timeout=10,
|
39
|
+
)
|
40
|
+
)
|
41
|
+
)
|
42
|
+
|
43
|
+
# --- Dummy Response Classes and Functions for requests.post ---
|
44
|
+
|
45
|
+
|
46
|
+
class DummyResponse:
|
47
|
+
"""A dummy response class for mocking HTTP responses."""
|
48
|
+
|
49
|
+
def __init__(self, json_data, status_code=200):
|
50
|
+
"""Initialize a DummyResponse with the given JSON data and status code."""
|
51
|
+
self._json_data = json_data
|
52
|
+
self.status_code = status_code
|
53
|
+
|
54
|
+
def json(self):
|
55
|
+
"""Return the JSON data from the response."""
|
56
|
+
return self._json_data
|
57
|
+
|
58
|
+
def raise_for_status(self):
|
59
|
+
"""raise an HTTP error for status codes >= 400."""
|
60
|
+
if self.status_code >= 400:
|
61
|
+
raise requests.HTTPError("HTTP Error")
|
62
|
+
|
63
|
+
|
64
|
+
def test_dummy_response_no_error():
|
65
|
+
"""Test that raise_for_status does not raise an exception for a successful response."""
|
66
|
+
# Create a DummyResponse with a successful status code.
|
67
|
+
response = DummyResponse({"data": "success"}, status_code=200)
|
68
|
+
# Calling raise_for_status should not raise an exception and should return None.
|
69
|
+
assert response.raise_for_status() is None
|
70
|
+
|
71
|
+
|
72
|
+
def test_dummy_response_raise_error():
|
73
|
+
"""Test that raise_for_status raises an exception for a failing response."""
|
74
|
+
# Create a DummyResponse with a failing status code.
|
75
|
+
response = DummyResponse({"error": "fail"}, status_code=400)
|
76
|
+
# Calling raise_for_status should raise an HTTPError.
|
77
|
+
with pytest.raises(requests.HTTPError):
|
78
|
+
response.raise_for_status()
|
79
|
+
|
80
|
+
|
81
|
+
def dummy_requests_post_success(url, headers, params, data, timeout):
|
82
|
+
"""dummy_requests_post_success"""
|
83
|
+
# Record call parameters for assertions.
|
84
|
+
dummy_requests_post_success.called_url = url
|
85
|
+
dummy_requests_post_success.called_headers = headers
|
86
|
+
dummy_requests_post_success.called_params = params
|
87
|
+
dummy_requests_post_success.called_data = data
|
88
|
+
dummy_requests_post_success.called_timeout = timeout
|
89
|
+
|
90
|
+
# Simulate a valid API response with three recommended papers;
|
91
|
+
# one paper missing authors should be filtered out.
|
92
|
+
dummy_data = {
|
93
|
+
"recommendedPapers": [
|
94
|
+
{
|
95
|
+
"paperId": "paperA",
|
96
|
+
"title": "Multi Rec Paper A",
|
97
|
+
"authors": ["Author X"],
|
98
|
+
"year": 2019,
|
99
|
+
"citationCount": 12,
|
100
|
+
"url": "http://paperA",
|
101
|
+
"externalIds": {"ArXiv": "arxivA"},
|
102
|
+
},
|
103
|
+
{
|
104
|
+
"paperId": "paperB",
|
105
|
+
"title": "Multi Rec Paper B",
|
106
|
+
"authors": ["Author Y"],
|
107
|
+
"year": 2020,
|
108
|
+
"citationCount": 18,
|
109
|
+
"url": "http://paperB",
|
110
|
+
"externalIds": {},
|
111
|
+
},
|
112
|
+
{
|
113
|
+
"paperId": "paperC",
|
114
|
+
"title": "Multi Rec Paper C",
|
115
|
+
"authors": None, # This one should be filtered out.
|
116
|
+
"year": 2021,
|
117
|
+
"citationCount": 25,
|
118
|
+
"url": "http://paperC",
|
119
|
+
"externalIds": {"ArXiv": "arxivC"},
|
120
|
+
},
|
121
|
+
]
|
122
|
+
}
|
123
|
+
return DummyResponse(dummy_data)
|
124
|
+
|
125
|
+
|
126
|
+
def dummy_requests_post_unexpected(url, headers, params, data, timeout):
|
127
|
+
"""dummy_requests_post_unexpected"""
|
128
|
+
dummy_requests_post_unexpected.called_url = url
|
129
|
+
dummy_requests_post_unexpected.called_headers = headers
|
130
|
+
dummy_requests_post_unexpected.called_params = params
|
131
|
+
dummy_requests_post_unexpected.called_data = data
|
132
|
+
dummy_requests_post_unexpected.called_timeout = timeout
|
133
|
+
# Simulate a response missing the 'recommendedPapers' key.
|
134
|
+
return DummyResponse({"error": "Invalid format"})
|
135
|
+
|
136
|
+
|
137
|
+
def dummy_requests_post_no_recs(url, headers, params, data, timeout):
|
138
|
+
"""dummy_requests_post_no_recs"""
|
139
|
+
dummy_requests_post_no_recs.called_url = url
|
140
|
+
dummy_requests_post_no_recs.called_headers = headers
|
141
|
+
dummy_requests_post_no_recs.called_params = params
|
142
|
+
dummy_requests_post_no_recs.called_data = data
|
143
|
+
dummy_requests_post_no_recs.called_timeout = timeout
|
144
|
+
# Simulate a response with an empty recommendations list.
|
145
|
+
return DummyResponse({"recommendedPapers": []})
|
146
|
+
|
147
|
+
|
148
|
+
def dummy_requests_post_exception(url, headers, params, data, timeout):
|
149
|
+
"""dummy_requests_post_exception"""
|
150
|
+
dummy_requests_post_exception.called_url = url
|
151
|
+
dummy_requests_post_exception.called_headers = headers
|
152
|
+
dummy_requests_post_exception.called_params = params
|
153
|
+
dummy_requests_post_exception.called_data = data
|
154
|
+
dummy_requests_post_exception.called_timeout = timeout
|
155
|
+
# Simulate a network exception.
|
156
|
+
raise requests.exceptions.RequestException("Connection error")
|
157
|
+
|
158
|
+
|
159
|
+
# --- Pytest Fixture to Patch Hydra ---
|
160
|
+
@pytest.fixture(autouse=True)
|
161
|
+
def patch_hydra(monkeypatch):
|
162
|
+
"""Patch Hydra's initialize and compose functions to return dummy objects."""
|
163
|
+
# Patch hydra.initialize to return our dummy context manager.
|
164
|
+
monkeypatch.setattr(
|
165
|
+
hydra, "initialize", lambda version_base, config_path: DummyHydraContext()
|
166
|
+
)
|
167
|
+
# Patch hydra.compose to return our dummy config.
|
168
|
+
monkeypatch.setattr(hydra, "compose", lambda config_name, overrides: dummy_config)
|
169
|
+
|
170
|
+
|
171
|
+
# --- Test Cases ---
|
172
|
+
|
173
|
+
|
174
|
+
def test_multi_paper_rec_success(monkeypatch):
|
175
|
+
"""
|
176
|
+
Test that get_multi_paper_recommendations returns a valid Command object
|
177
|
+
when the API response is successful. Also, ensure that recommendations missing
|
178
|
+
required fields (like authors) are filtered out.
|
179
|
+
"""
|
180
|
+
monkeypatch.setattr(requests, "post", dummy_requests_post_success)
|
181
|
+
|
182
|
+
tool_call_id = "test_tool_call_id"
|
183
|
+
input_data = {
|
184
|
+
"paper_ids": ["p1", "p2"],
|
185
|
+
"tool_call_id": tool_call_id,
|
186
|
+
"limit": 2,
|
187
|
+
"year": "2020",
|
188
|
+
}
|
189
|
+
# Call the tool using .run() with a dictionary input.
|
190
|
+
result = get_multi_paper_recommendations.run(input_data)
|
191
|
+
|
192
|
+
# Validate that the result is a Command with the expected update structure.
|
193
|
+
assert isinstance(result, Command)
|
194
|
+
update = result.update
|
195
|
+
assert "multi_papers" in update
|
196
|
+
|
197
|
+
papers = update["multi_papers"]
|
198
|
+
# Papers with valid 'title' and 'authors' should be included.
|
199
|
+
assert "paperA" in papers
|
200
|
+
assert "paperB" in papers
|
201
|
+
# Paper "paperC" is missing authors and should be filtered out.
|
202
|
+
assert "paperC" not in papers
|
203
|
+
|
204
|
+
# Check that a ToolMessage is included in the messages.
|
205
|
+
messages = update.get("messages", [])
|
206
|
+
assert len(messages) == 1
|
207
|
+
msg = messages[0]
|
208
|
+
assert isinstance(msg, ToolMessage)
|
209
|
+
assert "Recommendations based on multiple papers were successful" in msg.content
|
210
|
+
|
211
|
+
# Verify that the correct parameters were sent to requests.post.
|
212
|
+
called_params = dummy_requests_post_success.called_params
|
213
|
+
assert called_params["limit"] == 2 # Should be min(limit, 500)
|
214
|
+
assert called_params["fields"] == "paperId,title,authors"
|
215
|
+
# The year parameter should be present.
|
216
|
+
assert called_params["year"] == "2020"
|
217
|
+
|
218
|
+
# Also check the payload sent in the data.
|
219
|
+
sent_payload = json.loads(dummy_requests_post_success.called_data)
|
220
|
+
assert sent_payload["positivePaperIds"] == ["p1", "p2"]
|
221
|
+
assert sent_payload["negativePaperIds"] == []
|
222
|
+
|
223
|
+
|
224
|
+
def test_multi_paper_rec_unexpected_format(monkeypatch):
|
225
|
+
"""
|
226
|
+
Test that get_multi_paper_recommendations raises a RuntimeError when the API
|
227
|
+
response does not include the expected 'recommendedPapers' key.
|
228
|
+
"""
|
229
|
+
monkeypatch.setattr(requests, "post", dummy_requests_post_unexpected)
|
230
|
+
tool_call_id = "test_tool_call_id"
|
231
|
+
input_data = {
|
232
|
+
"paper_ids": ["p1", "p2"],
|
233
|
+
"tool_call_id": tool_call_id,
|
234
|
+
}
|
235
|
+
with pytest.raises(
|
236
|
+
RuntimeError,
|
237
|
+
match=(
|
238
|
+
"Unexpected response from Semantic Scholar API. The results could not be "
|
239
|
+
"retrieved due to an unexpected format. "
|
240
|
+
"Please modify your search query and try again."
|
241
|
+
),
|
242
|
+
):
|
243
|
+
get_multi_paper_recommendations.run(input_data)
|
244
|
+
|
245
|
+
|
246
|
+
def test_multi_paper_rec_no_recommendations(monkeypatch):
|
247
|
+
"""
|
248
|
+
Test that get_multi_paper_recommendations raises a RuntimeError when the API
|
249
|
+
returns no recommendations.
|
250
|
+
"""
|
251
|
+
monkeypatch.setattr(requests, "post", dummy_requests_post_no_recs)
|
252
|
+
tool_call_id = "test_tool_call_id"
|
253
|
+
input_data = {
|
254
|
+
"paper_ids": ["p1", "p2"],
|
255
|
+
"tool_call_id": tool_call_id,
|
256
|
+
}
|
257
|
+
with pytest.raises(
|
258
|
+
RuntimeError,
|
259
|
+
match=(
|
260
|
+
"No recommendations were found for your query. Consider refining your search "
|
261
|
+
"by using more specific keywords or different terms."
|
262
|
+
),
|
263
|
+
):
|
264
|
+
get_multi_paper_recommendations.run(input_data)
|
265
|
+
|
266
|
+
|
267
|
+
def test_multi_paper_rec_requests_exception(monkeypatch):
|
268
|
+
"""
|
269
|
+
Test that get_multi_paper_recommendations raises a RuntimeError when requests.post
|
270
|
+
throws an exception.
|
271
|
+
"""
|
272
|
+
monkeypatch.setattr(requests, "post", dummy_requests_post_exception)
|
273
|
+
tool_call_id = "test_tool_call_id"
|
274
|
+
input_data = {
|
275
|
+
"paper_ids": ["p1", "p2"],
|
276
|
+
"tool_call_id": tool_call_id,
|
277
|
+
}
|
278
|
+
with pytest.raises(
|
279
|
+
RuntimeError,
|
280
|
+
match="Failed to connect to Semantic Scholar API. Please retry the same query.",
|
281
|
+
):
|
282
|
+
get_multi_paper_recommendations.run(input_data)
|
@@ -0,0 +1,78 @@
|
|
1
|
+
"""
|
2
|
+
Unit tests for S2 tools functionality.
|
3
|
+
"""
|
4
|
+
|
5
|
+
# pylint: disable=redefined-outer-name
|
6
|
+
from unittest.mock import patch
|
7
|
+
from unittest.mock import MagicMock
|
8
|
+
import pytest
|
9
|
+
from ..tools.s2.query_results import query_results, NoPapersFoundError
|
10
|
+
|
11
|
+
|
12
|
+
@pytest.fixture
|
13
|
+
def initial_state():
|
14
|
+
"""Provides an empty initial state for tests."""
|
15
|
+
return {"papers": {}, "multi_papers": {}}
|
16
|
+
|
17
|
+
|
18
|
+
# Fixed test data for deterministic results
|
19
|
+
MOCK_SEARCH_RESPONSE = {
|
20
|
+
"data": [
|
21
|
+
{
|
22
|
+
"paperId": "123",
|
23
|
+
"title": "Machine Learning Basics",
|
24
|
+
"abstract": "An introduction to ML",
|
25
|
+
"year": 2023,
|
26
|
+
"citationCount": 100,
|
27
|
+
"url": "https://example.com/paper1",
|
28
|
+
"authors": [{"name": "Test Author"}],
|
29
|
+
}
|
30
|
+
]
|
31
|
+
}
|
32
|
+
|
33
|
+
MOCK_STATE_PAPER = {
|
34
|
+
"123": {
|
35
|
+
"Title": "Machine Learning Basics",
|
36
|
+
"Abstract": "An introduction to ML",
|
37
|
+
"Year": 2023,
|
38
|
+
"Citation Count": 100,
|
39
|
+
"URL": "https://example.com/paper1",
|
40
|
+
}
|
41
|
+
}
|
42
|
+
|
43
|
+
|
44
|
+
class TestS2Tools:
|
45
|
+
"""Unit tests for individual S2 tools"""
|
46
|
+
|
47
|
+
def test_query_results_empty_state(self, initial_state):
|
48
|
+
"""Tests query_results tool behavior when no papers are found."""
|
49
|
+
with pytest.raises(
|
50
|
+
NoPapersFoundError,
|
51
|
+
match="No papers found. A search needs to be performed first.",
|
52
|
+
):
|
53
|
+
query_results.invoke(
|
54
|
+
{"question": "List all papers", "state": initial_state}
|
55
|
+
)
|
56
|
+
|
57
|
+
@patch(
|
58
|
+
"aiagents4pharma.talk2scholars.tools.s2.query_results.create_pandas_dataframe_agent"
|
59
|
+
)
|
60
|
+
def test_query_results_with_papers(self, mock_create_agent, initial_state):
|
61
|
+
"""Tests querying papers when data is available."""
|
62
|
+
state = initial_state.copy()
|
63
|
+
state["last_displayed_papers"] = "papers"
|
64
|
+
state["papers"] = MOCK_STATE_PAPER
|
65
|
+
|
66
|
+
# Mock the dataframe agent instead of the LLM
|
67
|
+
mock_agent = MagicMock()
|
68
|
+
mock_agent.invoke.return_value = {"output": "Mocked response"}
|
69
|
+
|
70
|
+
mock_create_agent.return_value = (
|
71
|
+
mock_agent # Mock the function returning the agent
|
72
|
+
)
|
73
|
+
|
74
|
+
# Ensure that the output of query_results is correctly structured
|
75
|
+
result = query_results.invoke({"question": "List all papers", "state": state})
|
76
|
+
|
77
|
+
assert isinstance(result, str) # Ensure output is a string
|
78
|
+
assert result == "Mocked response" # Validate the expected response
|
@@ -0,0 +1,65 @@
|
|
1
|
+
"""
|
2
|
+
Unit tests for S2 tools functionality.
|
3
|
+
"""
|
4
|
+
|
5
|
+
# pylint: disable=redefined-outer-name
|
6
|
+
from unittest.mock import patch
|
7
|
+
import pytest
|
8
|
+
from langgraph.types import Command
|
9
|
+
from ..tools.s2.retrieve_semantic_scholar_paper_id import (
|
10
|
+
retrieve_semantic_scholar_paper_id,
|
11
|
+
)
|
12
|
+
|
13
|
+
|
14
|
+
# Fixed test data for deterministic results
|
15
|
+
MOCK_SEARCH_RESPONSE = {
|
16
|
+
"data": [
|
17
|
+
{
|
18
|
+
"paperId": "123",
|
19
|
+
"title": "Machine Learning Basics",
|
20
|
+
"abstract": "An introduction to ML",
|
21
|
+
"year": 2023,
|
22
|
+
"citationCount": 100,
|
23
|
+
"url": "https://example.com/paper1",
|
24
|
+
"authors": [{"name": "Test Author"}],
|
25
|
+
}
|
26
|
+
]
|
27
|
+
}
|
28
|
+
|
29
|
+
MOCK_STATE_PAPER = {
|
30
|
+
"123": {
|
31
|
+
"Title": "Machine Learning Basics",
|
32
|
+
"Abstract": "An introduction to ML",
|
33
|
+
"Year": 2023,
|
34
|
+
"Citation Count": 100,
|
35
|
+
"URL": "https://example.com/paper1",
|
36
|
+
}
|
37
|
+
}
|
38
|
+
|
39
|
+
|
40
|
+
class TestS2Tools:
|
41
|
+
"""Unit tests for individual S2 tools"""
|
42
|
+
|
43
|
+
@patch("requests.get")
|
44
|
+
def test_retrieve_semantic_scholar_paper_id(self, mock_get):
|
45
|
+
"""Tests retrieving a paper ID from Semantic Scholar."""
|
46
|
+
mock_get.return_value.json.return_value = MOCK_SEARCH_RESPONSE
|
47
|
+
mock_get.return_value.status_code = 200
|
48
|
+
|
49
|
+
result = retrieve_semantic_scholar_paper_id.invoke(
|
50
|
+
input={"paper_title": "Machine Learning Basics", "tool_call_id": "test123"}
|
51
|
+
)
|
52
|
+
|
53
|
+
assert isinstance(result, Command)
|
54
|
+
assert "messages" in result.update
|
55
|
+
assert (
|
56
|
+
"Paper ID for 'Machine Learning Basics' is: 123"
|
57
|
+
in result.update["messages"][0].content
|
58
|
+
)
|
59
|
+
|
60
|
+
def test_retrieve_semantic_scholar_paper_id_no_results(self):
|
61
|
+
"""Test retrieving a paper ID when no results are found."""
|
62
|
+
with pytest.raises(ValueError, match="No papers found for query: UnknownPaper"):
|
63
|
+
retrieve_semantic_scholar_paper_id.invoke(
|
64
|
+
input={"paper_title": "UnknownPaper", "tool_call_id": "test123"}
|
65
|
+
)
|