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

aiagents4pharma/talk2scholars/tests/test_routing_logic.py

@@ -23,8 +23,7 @@ def mock_router():
 
     def mock_supervisor_node(state):
         query = state["messages"][-1].content.lower()
-
-        # Expanded keyword matching for S2 Agent
+        # Define keywords for each sub-agent.
         s2_keywords = [
             "paper",
             "research",
@@ -34,13 +33,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 +60,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):

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_read.py

@@ -26,7 +26,7 @@ dummy_cfg = SimpleNamespace(tools=SimpleNamespace(zotero_read=dummy_zotero_read_
 
 
 class TestZoteroSearchTool(unittest.TestCase):
-    """test for Zotero search tool"""
+    """Tests for Zotero search tool."""
 
     @patch(
         "aiagents4pharma.talk2scholars.tools.zotero.zotero_read.get_item_collections"
@@ -41,7 +41,7 @@ class TestZoteroSearchTool(unittest.TestCase):
         mock_zotero_class,
         mock_get_item_collections,
     ):
-        """test valid query"""
+        """Test valid query returns correct Command output."""
         # Setup Hydra mocks
         mock_hydra_compose.return_value = dummy_cfg
         mock_hydra_init.return_value.__enter__.return_value = None
@@ -116,7 +116,7 @@ class TestZoteroSearchTool(unittest.TestCase):
         mock_zotero_class,
         mock_get_item_collections,
     ):
-        """test empty query fetches all items"""
+        """Test empty query fetches all items."""
         mock_hydra_compose.return_value = dummy_cfg
         mock_hydra_init.return_value.__enter__.return_value = None
 
@@ -166,7 +166,7 @@ class TestZoteroSearchTool(unittest.TestCase):
         mock_zotero_class,
         mock_get_item_collections,
     ):
-        """test no items returned from Zotero"""
+        """Test no items returned from Zotero."""
         mock_hydra_compose.return_value = dummy_cfg
         mock_hydra_init.return_value.__enter__.return_value = None
 
@@ -199,7 +199,10 @@ class TestZoteroSearchTool(unittest.TestCase):
         mock_zotero_class,
         mock_get_item_collections,
     ):
-        """test no matching papers returned from Zotero"""
+        """
+        Test that when non-research items (e.g. attachments, notes) are returned,
+        they are still included since filtering is disabled.
+        """
         mock_hydra_compose.return_value = dummy_cfg
         mock_hydra_init.return_value.__enter__.return_value = None
 
@@ -240,9 +243,13 @@ class TestZoteroSearchTool(unittest.TestCase):
             "tool_call_id": tool_call_id,
             "limit": 2,
         }
-        with self.assertRaises(RuntimeError) as context:
-            zotero_search_tool.run(tool_input)
-        self.assertIn("No matching papers returned from Zotero", str(context.exception))
+        # Instead of expecting a RuntimeError, we now expect both items to be returned.
+        result = zotero_search_tool.run(tool_input)
+        update = result.update
+        filtered_papers = update["zotero_read"]
+        self.assertIn("paper1", filtered_papers)
+        self.assertIn("paper2", filtered_papers)
+        self.assertEqual(len(filtered_papers), 2)
 
     @patch(
         "aiagents4pharma.talk2scholars.tools.zotero.zotero_read.get_item_collections"
@@ -257,7 +264,7 @@ class TestZoteroSearchTool(unittest.TestCase):
         mock_zotero_class,
         mock_get_item_collections,
     ):
-        """test items API exception"""
+        """Test items API exception is properly raised."""
         mock_hydra_compose.return_value = dummy_cfg
         mock_hydra_init.return_value.__enter__.return_value = None
         mock_get_item_collections.return_value = {}
@@ -306,7 +313,7 @@ class TestZoteroSearchTool(unittest.TestCase):
                     "url": "http://example.com",
                     "itemType": "journalArticle",
                 }
-            },  # missing key triggers line 136
+            },  # Missing 'key' field
             {
                 "data": {
                     "key": "paper_valid",

aiagents4pharma/talk2scholars/tools/paper_download/abstract_downloader.py

@@ -9,6 +9,8 @@ inherit from this class and implement its methods.
 
 from abc import ABC, abstractmethod
 from typing import Any, Dict
+
+
 class AbstractPaperDownloader(ABC):
     """
     Abstract base class for scholarly paper downloaders.

aiagents4pharma/talk2scholars/tools/paper_download/arxiv_downloader.py

@@ -8,6 +8,7 @@ downloads the corresponding PDF.
 By using an abstract base class, this implementation is extendable to other
 APIs like PubMed, IEEE Xplore, etc.
 """
+
 import xml.etree.ElementTree as ET
 from typing import Any, Dict
 import logging
@@ -19,6 +20,7 @@ from .abstract_downloader import AbstractPaperDownloader
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+
 class ArxivPaperDownloader(AbstractPaperDownloader):
     """
     Downloader class for arXiv.
@@ -35,13 +37,13 @@ class ArxivPaperDownloader(AbstractPaperDownloader):
         """
         with hydra.initialize(version_base=None, config_path="../../configs"):
             cfg = hydra.compose(
-                config_name="config",
-                overrides=["tools/download_arxiv_paper=default"]
+                config_name="config", overrides=["tools/download_arxiv_paper=default"]
             )
             self.api_url = cfg.tools.download_arxiv_paper.api_url
             self.request_timeout = cfg.tools.download_arxiv_paper.request_timeout
             self.chunk_size = cfg.tools.download_arxiv_paper.chunk_size
             self.pdf_base_url = cfg.tools.download_arxiv_paper.pdf_base_url
+
     def fetch_metadata(self, paper_id: str) -> Dict[str, Any]:
         """
         Fetch metadata from arXiv for a given paper ID.
@@ -95,11 +97,16 @@ class ArxivPaperDownloader(AbstractPaperDownloader):
         logger.info("Downloading PDF from: %s", pdf_url)
         pdf_response = requests.get(pdf_url, stream=True, timeout=self.request_timeout)
         pdf_response.raise_for_status()
+        # print (pdf_response)
 
         # Combine the PDF data from chunks.
         pdf_object = b"".join(
-            chunk for chunk in pdf_response.iter_content(chunk_size=self.chunk_size) if chunk
-            )
+            chunk
+            for chunk in pdf_response.iter_content(chunk_size=self.chunk_size)
+            if chunk
+        )
+        # print (pdf_object)
+        print("PDF_URL", pdf_url)
 
         return {
             "pdf_object": pdf_object,

aiagents4pharma/talk2scholars/tools/paper_download/download_arxiv_input.py

@@ -14,16 +14,19 @@ from langgraph.types import Command
 # Local import from the same package:
 from .arxiv_downloader import ArxivPaperDownloader
 
+
 class DownloadArxivPaperInput(BaseModel):
     """
     Input schema for the arXiv paper download tool.
     (Optional: if you decide to keep Pydantic validation in the future)
     """
+
     arxiv_id: str = Field(
         description="The arXiv paper ID used to retrieve the paper details and PDF."
-        )
+    )
     tool_call_id: Annotated[str, InjectedToolCallId]
 
+
 @tool(args_schema=DownloadArxivPaperInput, parse_docstring=True)
 def download_arxiv_paper(
     arxiv_id: str,
@@ -49,6 +52,7 @@ def download_arxiv_paper(
 
     # If the downloader fails or the arxiv_id is invalid, this might raise an error
     pdf_data = downloader.download_pdf(arxiv_id)
+    # print (pdf_data)
 
     content = f"Successfully downloaded PDF for arXiv ID {arxiv_id}"
 

aiagents4pharma/talk2scholars/tools/pdf/question_and_answer.py

@@ -2,8 +2,8 @@
 """
 question_and_answer: Tool for performing Q&A on PDF documents using retrieval augmented generation.
 
-This module provides functionality to extract text from PDF binary data, split it into 
-chunks, retrieve relevant segments via a vector store, and generate an answer to a 
+This module provides functionality to extract text from PDF binary data, split it into
+chunks, retrieve relevant segments via a vector store, and generate an answer to a
 user-provided question using a language model chain.
 """
 
@@ -18,13 +18,15 @@ import hydra
 from langchain.chains.question_answering import load_qa_chain
 from langchain.docstore.document import Document
 from langchain.text_splitter import CharacterTextSplitter
-from langchain_community.vectorstores import Annoy
-from langchain_openai import OpenAIEmbeddings
 from langchain_core.language_models.chat_models import BaseChatModel
-
+from langchain_core.vectorstores import InMemoryVectorStore
 from langchain_core.messages import ToolMessage
 from langchain_core.tools import tool
 from langchain_core.tools.base import InjectedToolCallId
+from langchain_core.embeddings import Embeddings
+from langchain_community.vectorstores import Annoy
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_openai import OpenAIEmbeddings
 from langgraph.types import Command
 from langgraph.prebuilt import InjectedState
 
@@ -35,10 +37,13 @@ logger.setLevel(logging.INFO)
 
 # Load configuration using Hydra.
 with hydra.initialize(version_base=None, config_path="../../configs"):
-    cfg = hydra.compose(config_name="config", overrides=["tools/question_and_answer=default"])
+    cfg = hydra.compose(
+        config_name="config", overrides=["tools/question_and_answer=default"]
+    )
     cfg = cfg.tools.question_and_answer
     logger.info("Loaded Question and Answer tool configuration.")
 
+
 class QuestionAndAnswerInput(BaseModel):
     """
     Input schema for the PDF Question and Answer tool.
@@ -47,12 +52,12 @@ class QuestionAndAnswerInput(BaseModel):
         question (str): The question to ask regarding the PDF content.
         tool_call_id (str): Unique identifier for the tool call, injected automatically.
     """
-    question: str = Field(
-        description="The question to ask regarding the PDF content."
-    )
+
+    question: str = Field(description="The question to ask regarding the PDF content.")
     tool_call_id: Annotated[str, InjectedToolCallId]
     state: Annotated[dict, InjectedState]
 
+
 def extract_text_from_pdf_data(pdf_bytes: bytes) -> str:
     """
     Extract text content from PDF binary data.
@@ -73,7 +78,10 @@ def extract_text_from_pdf_data(pdf_bytes: bytes) -> str:
         text += page_text
     return text
 
-def generate_answer(question: str, pdf_bytes: bytes, llm_model: BaseChatModel) -> Dict[str, Any]:
+
+def generate_answer(
+    question: str, pdf_bytes: bytes, llm_model: BaseChatModel
+) -> Dict[str, Any]:
     """
     Generate an answer for a question using retrieval augmented generation on PDF content.
 
@@ -92,9 +100,7 @@ def generate_answer(question: str, pdf_bytes: bytes, llm_model: BaseChatModel) -
     text = extract_text_from_pdf_data(pdf_bytes)
     logger.info("Extracted text from PDF.")
     text_splitter = CharacterTextSplitter(
-        separator="\n",
-        chunk_size=cfg.chunk_size,
-        chunk_overlap=cfg.chunk_overlap
+        separator="\n", chunk_size=cfg.chunk_size, chunk_overlap=cfg.chunk_overlap
     )
     chunks = text_splitter.split_text(text)
     documents: List[Document] = [Document(page_content=chunk) for chunk in chunks]
@@ -102,10 +108,7 @@ def generate_answer(question: str, pdf_bytes: bytes, llm_model: BaseChatModel) -
 
     embeddings = OpenAIEmbeddings(openai_api_key=cfg.openai_api_key)
     vector_store = Annoy.from_documents(documents, embeddings)
-    search_results = vector_store.similarity_search(
-        question,
-        k=cfg.num_retrievals
-    )
+    search_results = vector_store.similarity_search(question, k=cfg.num_retrievals)
     logger.info("Retrieved %d relevant document chunks.", len(search_results))
     # Use the provided llm_model to build the QA chain.
     qa_chain = load_qa_chain(llm_model, chain_type=cfg.qa_chain_type)
@@ -114,6 +117,49 @@ def generate_answer(question: str, pdf_bytes: bytes, llm_model: BaseChatModel) -
     )
     return answer
 
+
+def generate_answer2(
+    question: str, pdf_url: str, text_embedding_model: Embeddings
+) -> Dict[str, Any]:
+    """
+    Generate an answer for a question using retrieval augmented generation on PDF content.
+
+    This function extracts text from the PDF data, splits the text into manageable chunks,
+    performs a similarity search to retrieve the most relevant segments, and then uses a
+    question-answering chain (built using the provided llm_model) to generate an answer.
+
+    Args:
+        question (str): The question to be answered.
+        pdf_bytes (bytes): The binary content of the PDF document.
+        llm_model (BaseChatModel): The language model instance to use for answering.
+
+    Returns:
+        Dict[str, Any]: A dictionary containing the answer generated by the language model.
+    """
+    # text = extract_text_from_pdf_data(pdf_bytes)
+    # logger.info("Extracted text from PDF.")
+    logger.log(logging.INFO, "searching the article with the question: %s", question)
+    # Load the article
+    # loader = PyPDFLoader(state['pdf_file_name'])
+    # loader = PyPDFLoader("https://arxiv.org/pdf/2310.08365")
+    loader = PyPDFLoader(pdf_url)
+    # Load the pages of the article
+    pages = []
+    for page in loader.lazy_load():
+        pages.append(page)
+    # Set up text embedding model
+    # text_embedding_model = state['text_embedding_model']
+    # text_embedding_model = OpenAIEmbeddings(openai_api_key=cfg.openai_api_key)
+    logging.info("Loaded text embedding model %s", text_embedding_model)
+    # Create a vector store from the pages
+    vector_store = InMemoryVectorStore.from_documents(pages, text_embedding_model)
+    # Search the article with the question
+    docs = vector_store.similarity_search(question)
+    # Return the content of the pages
+    return "\n".join([doc.page_content for doc in docs])
+    # return answer
+
+
 @tool(args_schema=QuestionAndAnswerInput)
 def question_and_answer_tool(
     question: str,
@@ -124,7 +170,7 @@ def question_and_answer_tool(
     Answer a question using PDF content stored in the state via retrieval augmented generation.
 
     This tool retrieves the PDF binary data from the state (under the key "pdf_data"), extracts its
-    textual content, and generates an answer to the specified question. It also extracts the 
+    textual content, and generates an answer to the specified question. It also extracts the
     llm_model (of type BaseChatModel) from the state to use for answering.
 
     Args:
@@ -138,15 +184,15 @@ def question_and_answer_tool(
         Dict[str, Any]: A dictionary containing the generated answer or an error message.
     """
     logger.info("Starting PDF Question and Answer tool using PDF data from state.")
+    # print (state['text_embedding_model'])
+    text_embedding_model = state["text_embedding_model"]
     pdf_state = state.get("pdf_data")
     if not pdf_state:
         error_msg = "No pdf_data found in state."
         logger.error(error_msg)
         return Command(
             update={
-                "messages": [
-                    ToolMessage(content=error_msg, tool_call_id=tool_call_id)
-                ]
+                "messages": [ToolMessage(content=error_msg, tool_call_id=tool_call_id)]
             }
         )
     pdf_bytes = pdf_state.get("pdf_object")
@@ -155,16 +201,17 @@ def question_and_answer_tool(
         logger.error(error_msg)
         return Command(
             update={
-                "messages": [
-                    ToolMessage(content=error_msg, tool_call_id=tool_call_id)
-                ]
+                "messages": [ToolMessage(content=error_msg, tool_call_id=tool_call_id)]
             }
         )
+    pdf_url = pdf_state.get("pdf_url")
     # Retrieve llm_model from state; use a default if not provided.
     llm_model = state.get("llm_model")
     if not llm_model:
         logger.error("Missing LLM model instance in state.")
         return {"error": "No LLM model found in state."}
-    answer = generate_answer(question, pdf_bytes, llm_model)
-    logger.info("Generated answer: %s", answer)
+    # answer = generate_answer(question, pdf_bytes, llm_model)
+    print(pdf_url)
+    answer = generate_answer2(question, pdf_url, text_embedding_model)
+    # logger.info("Generated answer: %s", answer)
     return answer