aiagents4pharma 1.30.0__py3-none-any.whl → 1.30.2__py3-none-any.whl

aiagents4pharma/talk2scholars/tests/test_question_and_answer_tool.py

@@ -3,7 +3,6 @@ Unit tests for question_and_answer tool functionality.
 """
 
 from langchain.docstore.document import Document
-
 from ..tools.pdf import question_and_answer
 from ..tools.pdf.question_and_answer import (
     extract_text_from_pdf_data,
@@ -11,6 +10,8 @@ from ..tools.pdf.question_and_answer import (
     generate_answer,
 )
 
+# pylint: disable=redefined-outer-name,too-few-public-methods
+
 
 def test_extract_text_from_pdf_data():
     """
@@ -46,14 +47,14 @@ DUMMY_PDF_BYTES = (
 )
 
 
-def fake_generate_answer(question, pdf_bytes, _llm_model):
+def fake_generate_answer2(question, pdf_url, _text_embedding_model):
     """
-    Fake generate_answer function to bypass external dependencies.
+    Fake generate_answer2 function to bypass external dependencies.
     """
     return {
         "answer": "Mock answer",
         "question": question,
-        "pdf_bytes_length": len(pdf_bytes),
+        "pdf_url": pdf_url,
     }
 
 
@@ -61,30 +62,31 @@ def test_question_and_answer_tool_success(monkeypatch):
     """
     Test that question_and_answer_tool returns the expected result on success.
     """
-    monkeypatch.setattr(
-        question_and_answer, "generate_answer", fake_generate_answer
-    )
-    # Create a valid state with pdf_data containing both pdf_object and pdf_url,
-    # and include a dummy llm_model.
+    # Patch generate_answer2 because the tool calls that.
+    monkeypatch.setattr(question_and_answer, "generate_answer2", fake_generate_answer2)
+    dummy_text_embedding_model = object()  # Provide a dummy text embedding model.
+    # Create a valid state with pdf_data and include dummy llm_model and text_embedding_model.
     state = {
         "pdf_data": {"pdf_object": DUMMY_PDF_BYTES, "pdf_url": "http://dummy.url"},
         "llm_model": object(),  # Provide a dummy LLM model instance.
+        "text_embedding_model": dummy_text_embedding_model,
     }
     question = "What is in the PDF?"
-    # Call the underlying function directly via .func to bypass the StructuredTool wrapper.
     result = question_and_answer_tool.func(
         question=question, tool_call_id="test_call_id", state=state
     )
     assert result["answer"] == "Mock answer"
     assert result["question"] == question
-    assert result["pdf_bytes_length"] == len(DUMMY_PDF_BYTES)
+    assert result["pdf_url"] == "http://dummy.url"
 
 
 def test_question_and_answer_tool_no_pdf_data():
     """
     Test that an error is returned if the state lacks the 'pdf_data' key.
     """
-    state = {}  # pdf_data key is missing.
+    state = {
+        "text_embedding_model": object(),  # Added to avoid KeyError.
+    }
     question = "Any question?"
     result = question_and_answer_tool.func(
         question=question, tool_call_id="test_call_id", state=state
@@ -97,7 +99,11 @@ def test_question_and_answer_tool_no_pdf_object():
     """
     Test that an error is returned if the pdf_object is missing within pdf_data.
     """
-    state = {"pdf_data": {"pdf_object": None}}
+    state = {
+        "pdf_data": {"pdf_object": None},
+        "text_embedding_model": object(),  # Added to avoid KeyError.
+        "llm_model": object(),  # Dummy LLM model.
+    }
     question = "Any question?"
     result = question_and_answer_tool.func(
         question=question, tool_call_id="test_call_id", state=state
@@ -114,8 +120,9 @@ def test_question_and_answer_tool_no_llm_model():
     Test that an error is returned if the LLM model is missing in the state.
     """
     state = {
-        "pdf_data": {"pdf_object": DUMMY_PDF_BYTES, "pdf_url": "http://dummy.url"}
-        # Note: llm_model is intentionally omitted.
+        "pdf_data": {"pdf_object": DUMMY_PDF_BYTES, "pdf_url": "http://dummy.url"},
+        "text_embedding_model": object(),  # Added to avoid KeyError.
+        # llm_model is intentionally omitted.
     }
     question = "What is in the PDF?"
     result = question_and_answer_tool.func(
@@ -124,6 +131,57 @@ def test_question_and_answer_tool_no_llm_model():
     assert result == {"error": "No LLM model found in state."}
 
 
+def test_generate_answer2_actual(monkeypatch):
+    """
+    Test the actual behavior of generate_answer2 using fake dependencies
+    to exercise its internal logic.
+    """
+
+    # Create a fake PyPDFLoader that does not perform a network call.
+    class FakePyPDFLoader:
+        """class to fake PyPDFLoader"""
+
+        def __init__(self, file_path, headers=None):
+            """Initialize the fake PyPDFLoader."""
+            self.file_path = file_path
+            self.headers = headers
+
+        def lazy_load(self):
+            """Return a list with one fake Document."""
+            # Return a list with one fake Document.
+            return [Document(page_content="Answer for Test question?")]
+
+    monkeypatch.setattr(question_and_answer, "PyPDFLoader", FakePyPDFLoader)
+
+    # Create a fake vector store that returns a controlled result for similarity_search.
+    class FakeVectorStore:
+        """Fake vector store for similarity search."""
+
+        def similarity_search(self, query):
+            """Return a list with one Document containing our expected answer."""
+            # Return a list with one Document containing our expected answer.
+            return [Document(page_content=f"Answer for {query}")]
+
+    monkeypatch.setattr(
+        question_and_answer.InMemoryVectorStore,
+        "from_documents",
+        lambda docs, emb: FakeVectorStore(),
+    )
+
+    # Provide a dummy text embedding model.
+    dummy_text_embedding_model = object()
+    question = "Test question?"
+    pdf_url = "http://dummy.pdf"
+
+    # Call generate_answer2 without triggering an actual network call.
+    result = question_and_answer.generate_answer2(
+        question, pdf_url, dummy_text_embedding_model
+    )
+    # The function should join the page content from the similarity search.
+    expected = "Answer for Test question?"
+    assert result == expected
+
+
 def test_generate_answer(monkeypatch):
     """
     Test generate_answer function with controlled monkeypatched dependencies.
@@ -141,12 +199,15 @@ def test_generate_answer(monkeypatch):
         """
         Fake Annoy.from_documents function that returns a fake vector store.
         """
+
         # pylint: disable=too-few-public-methods, unused-argument
         class FakeVectorStore:
             """Fake vector store for similarity search."""
+
             def similarity_search(self, _question, k):
                 """Return a list with a single dummy Document."""
                 return [Document(page_content="dummy content")]
+
         return FakeVectorStore()
 
     monkeypatch.setattr(
@@ -157,9 +218,11 @@ def test_generate_answer(monkeypatch):
         """
         Fake load_qa_chain function that returns a fake QA chain.
         """
+
         # pylint: disable=too-few-public-methods, unused-argument
         class FakeChain:
             """Fake QA chain for testing generate_answer."""
+
             def invoke(self, **kwargs):
                 """
                 Fake invoke method that returns a mock answer.
@@ -169,6 +232,7 @@ def test_generate_answer(monkeypatch):
                     "answer": "real mock answer",
                     "question": input_data.get("question"),
                 }
+
         return FakeChain()
 
     monkeypatch.setattr(question_and_answer, "load_qa_chain", fake_load_qa_chain)

aiagents4pharma/talk2scholars/tests/test_routing_logic.py

@@ -22,9 +22,9 @@ def mock_router():
     """Creates a mock supervisor router that routes based on keyword matching."""
 
     def mock_supervisor_node(state):
+        """Mock supervisor node that routes based on keyword matching."""
         query = state["messages"][-1].content.lower()
-
-        # Expanded keyword matching for S2 Agent
+        # Define keywords for each sub-agent.
         s2_keywords = [
             "paper",
             "research",
@@ -34,13 +34,19 @@ def mock_router():
             "references",
         ]
         zotero_keywords = ["zotero", "library", "saved papers", "academic library"]
+        pdf_keywords = ["pdf", "document", "read pdf"]
+        paper_download_keywords = ["download", "arxiv", "fetch paper", "paper download"]
 
+        # Priority ordering: Zotero, then paper download, then PDF, then S2.
         if any(keyword in query for keyword in zotero_keywords):
             return Command(goto="zotero_agent")
+        if any(keyword in query for keyword in paper_download_keywords):
+            return Command(goto="paper_download_agent")
+        if any(keyword in query for keyword in pdf_keywords):
+            return Command(goto="pdf_agent")
         if any(keyword in query for keyword in s2_keywords):
             return Command(goto="s2_agent")
-
-        # If no match, default to ending the conversation
+        # Default to end if no keyword matches.
         return Command(goto=END)
 
     return mock_supervisor_node
@@ -55,10 +61,9 @@ def mock_router():
         ("Fetch my academic library.", "zotero_agent"),
         ("Retrieve citations.", "s2_agent"),
         ("Can you get journal articles?", "s2_agent"),
-        (
-            "Completely unrelated query.",
-            "__end__",
-        ),  # NEW: Should trigger the `END` case
+        ("I want to read the PDF document.", "pdf_agent"),
+        ("Download the paper from arxiv.", "paper_download_agent"),
+        ("Completely unrelated query.", "__end__"),
     ],
 )
 def test_routing_logic(mock_state, mock_router, user_query, expected_agent):
@@ -66,6 +71,4 @@ def test_routing_logic(mock_state, mock_router, user_query, expected_agent):
     mock_state["messages"].append(HumanMessage(content=user_query))
     result = mock_router(mock_state)
 
-    print(f"\nDEBUG: Query '{user_query}' routed to: {result.goto}")
-
     assert result.goto == expected_agent, f"Failed for query: {user_query}"

aiagents4pharma/talk2scholars/tests/test_s2_multi.py

@@ -94,7 +94,7 @@ def dummy_requests_post_success(url, headers, params, data, timeout):
             {
                 "paperId": "paperA",
                 "title": "Multi Rec Paper A",
-                "authors": ["Author X"],
+                "authors": [{"name": "Author X", "authorId": "AX"}],
                 "year": 2019,
                 "citationCount": 12,
                 "url": "http://paperA",
@@ -103,7 +103,7 @@ def dummy_requests_post_success(url, headers, params, data, timeout):
             {
                 "paperId": "paperB",
                 "title": "Multi Rec Paper B",
-                "authors": ["Author Y"],
+                "authors": [{"name": "Author Y", "authorId": "AY"}],
                 "year": 2020,
                 "citationCount": 18,
                 "url": "http://paperB",
@@ -112,7 +112,7 @@ def dummy_requests_post_success(url, headers, params, data, timeout):
             {
                 "paperId": "paperC",
                 "title": "Multi Rec Paper C",
-                "authors": None,  # This one should be filtered out.
+                "authors": None,  # This paper should be filtered out.
                 "year": 2021,
                 "citationCount": 25,
                 "url": "http://paperC",
@@ -277,6 +277,29 @@ def test_multi_paper_rec_requests_exception(monkeypatch):
     }
     with pytest.raises(
         RuntimeError,
-        match="Failed to connect to Semantic Scholar API. Please retry the same query.",
+        match="Failed to connect to Semantic Scholar API after 10 attempts."
+        "Please retry the same query.",
+    ):
+        get_multi_paper_recommendations.run(input_data)
+
+
+def test_multi_paper_rec_no_response(monkeypatch):
+    """
+    Test that get_multi_paper_recommendations raises a RuntimeError when no response is obtained.
+    This is simulated by patching 'range' in the underlying function's globals to
+    return an empty iterator,
+    so that the for loop does not iterate and response remains None.
+    """
+    # Patch 'range' in the underlying function's globals (accessed via .func.__globals__)
+    monkeypatch.setitem(
+        get_multi_paper_recommendations.func.__globals__, "range", lambda x: iter([])
+    )
+    tool_call_id = "test_tool_call_id"
+    input_data = {
+        "paper_ids": ["p1", "p2"],
+        "tool_call_id": tool_call_id,
+    }
+    with pytest.raises(
+        RuntimeError, match="Failed to obtain a response from the Semantic Scholar API."
     ):
         get_multi_paper_recommendations.run(input_data)

aiagents4pharma/talk2scholars/tests/test_s2_search.py

@@ -85,7 +85,7 @@ def dummy_requests_get_success(url, params, timeout):
             {
                 "paperId": "1",
                 "title": "Paper 1",
-                "authors": ["Author A"],
+                "authors": [{"name": "Author A", "authorId": "A1"}],
                 "year": 2020,
                 "citationCount": 10,
                 "url": "http://paper1",
@@ -94,7 +94,7 @@ def dummy_requests_get_success(url, params, timeout):
             {
                 "paperId": "2",
                 "title": "Paper 2",
-                "authors": ["Author B"],
+                "authors": [{"name": "Author B", "authorId": "B1"}],
                 "year": 2021,
                 "citationCount": 20,
                 "url": "http://paper2",
@@ -256,7 +256,8 @@ def test_search_tool_requests_exception(monkeypatch):
     tool_call_id = "test_tool_call_id"
     with pytest.raises(
         RuntimeError,
-        match="Failed to connect to Semantic Scholar API. Please retry the same query.",
+        match="Failed to connect to Semantic Scholar API after 10 attempts."
+        "Please retry the same query.",
     ):
         search_tool.run(
             {
@@ -264,3 +265,18 @@ def test_search_tool_requests_exception(monkeypatch):
                 "tool_call_id": tool_call_id,
             }
         )
+
+
+def test_search_tool_no_response(monkeypatch):
+    """
+    Test that search_tool raises a RuntimeError when no response is obtained.
+    This is simulated by patching 'range' in the original function's globals (a dict)
+    so that it returns an empty iterator, leaving response as None.
+    """
+    # Patch 'range' in the original function's globals using setitem.
+    monkeypatch.setitem(search_tool.func.__globals__, "range", lambda x: iter([]))
+    tool_call_id = "test_tool_call_id"
+    with pytest.raises(
+        RuntimeError, match="Failed to obtain a response from the Semantic Scholar API."
+    ):
+        search_tool.run({"query": "test", "tool_call_id": tool_call_id})

aiagents4pharma/talk2scholars/tests/test_s2_single.py

@@ -92,7 +92,7 @@ def dummy_requests_get_success(url, params, timeout):
             {
                 "paperId": "paper1",
                 "title": "Recommended Paper 1",
-                "authors": ["Author A"],
+                "authors": [{"name": "Author A", "authorId": "A1"}],
                 "year": 2020,
                 "citationCount": 15,
                 "url": "http://paper1",
@@ -101,7 +101,7 @@ def dummy_requests_get_success(url, params, timeout):
             {
                 "paperId": "paper2",
                 "title": "Recommended Paper 2",
-                "authors": ["Author B"],
+                "authors": [{"name": "Author B", "authorId": "B1"}],
                 "year": 2021,
                 "citationCount": 25,
                 "url": "http://paper2",
@@ -269,6 +269,30 @@ def test_single_paper_rec_requests_exception(monkeypatch):
     }
     with pytest.raises(
         RuntimeError,
-        match="Failed to connect to Semantic Scholar API. Please retry the same query.",
+        match="Failed to connect to Semantic Scholar API after 10 attempts."
+        "Please retry the same query.",
+    ):
+        get_single_paper_recommendations.run(input_data)
+
+
+def test_single_paper_rec_no_response(monkeypatch):
+    """
+    Test that get_single_paper_recommendations raises a RuntimeError
+    when no response is obtained from the API.
+
+    This is simulated by patching 'range' in the underlying function's globals
+    to return an empty iterator, so the for-loop never iterates and response remains None.
+    """
+    # Patch 'range' in the underlying function's globals (accessed via .func.__globals__)
+    monkeypatch.setitem(
+        get_single_paper_recommendations.func.__globals__, "range", lambda x: iter([])
+    )
+    tool_call_id = "test_tool_call_id"
+    input_data = {
+        "paper_id": "12345",
+        "tool_call_id": tool_call_id,
+    }
+    with pytest.raises(
+        RuntimeError, match="Failed to obtain a response from the Semantic Scholar API."
     ):
         get_single_paper_recommendations.run(input_data)

aiagents4pharma/talk2scholars/tests/test_zotero_agent.py

@@ -10,7 +10,8 @@ from langchain_openai import ChatOpenAI
 from ..agents.zotero_agent import get_app
 from ..state.state_talk2scholars import Talk2Scholars
 
-LLM_MODEL = ChatOpenAI(model='gpt-4o-mini', temperature=0)
+LLM_MODEL = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+
 
 @pytest.fixture(autouse=True)
 def mock_hydra_fixture():
@@ -39,7 +40,7 @@ def mock_tools_fixture():
             "retrieve_semantic_scholar_paper_id"
         ) as mock_s2_retrieve_id,
         mock.patch(
-            "aiagents4pharma.talk2scholars.tools.zotero.zotero_read.zotero_search_tool"
+            "aiagents4pharma.talk2scholars.tools.zotero.zotero_read.zotero_read"
         ) as mock_zotero_query_results,
     ):
         mock_s2_display.return_value = {"result": "Mock Display Result"}