aiagents4pharma 1.20.0__py3-none-any.whl → 1.21.0__py3-none-any.whl

aiagents4pharma/talk2scholars/tests/test_s2_tools.py

@@ -4,12 +4,20 @@ Unit tests for S2 tools functionality.
 
 # pylint: disable=redefined-outer-name
 from unittest.mock import patch
-from langchain_core.messages import ToolMessage
+from unittest.mock import MagicMock
 import pytest
-from ..tools.s2.display_results import display_results, NoPapersFoundError
+from langgraph.types import Command
+from ..tools.s2.display_results import (
+    display_results,
+    NoPapersFoundError as raised_error,
+)
 from ..tools.s2.multi_paper_rec import get_multi_paper_recommendations
 from ..tools.s2.search import search_tool
 from ..tools.s2.single_paper_rec import get_single_paper_recommendations
+from ..tools.s2.query_results import query_results, NoPapersFoundError
+from ..tools.s2.retrieve_semantic_scholar_paper_id import (
+    retrieve_semantic_scholar_paper_id,
+)
 
 
 @pytest.fixture
@@ -50,20 +58,29 @@ class TestS2Tools:
     def test_display_results_empty_state(self, initial_state):
         """Verifies display_results tool behavior when state is empty and raises an exception"""
         with pytest.raises(
-            NoPapersFoundError,
-            match="No papers found. A search needs to be performed first.",
+            raised_error,
+            match="No papers found. A search/rec needs to be performed first.",
         ):
-            display_results.invoke({"state": initial_state})
+            display_results.invoke({"state": initial_state, "tool_call_id": "test123"})
 
     def test_display_results_shows_papers(self, initial_state):
         """Verifies display_results tool correctly returns papers from state"""
         state = initial_state.copy()
+        state["last_displayed_papers"] = "papers"
         state["papers"] = MOCK_STATE_PAPER
-        state["multi_papers"] = {}
-        result = display_results.invoke(input={"state": state})
-        assert isinstance(result, dict)
-        assert result["papers"] == MOCK_STATE_PAPER
-        assert result["multi_papers"] == {}
+
+        result = display_results.invoke(
+            input={"state": state, "tool_call_id": "test123"}
+        )
+
+        assert isinstance(result, Command)  # Expect a Command object
+        assert isinstance(result.update, dict)  # Ensure update is a dictionary
+        assert "messages" in result.update
+        assert len(result.update["messages"]) == 1
+        assert (
+            "1 papers found. Papers are attached as an artifact."
+            in result.update["messages"][0].content
+        )
 
     @patch("requests.get")
     def test_search_finds_papers(self, mock_get):
@@ -152,14 +169,12 @@ class TestS2Tools:
         mock_get.return_value.status_code = 200
 
         result = get_single_paper_recommendations.invoke(
-            input={
-                "paper_id": "123",
-                "limit": 1,
-                "tool_call_id": "test123",
-            }
+            input={"paper_id": "123", "limit": 1, "tool_call_id": "test123"}
         )
+
+        assert isinstance(result, Command)
         assert "papers" in result.update
-        assert isinstance(result.update["messages"][0], ToolMessage)
+        assert len(result.update["messages"]) == 1
 
     @patch("requests.get")
     def test_single_paper_rec_with_optional_params(self, mock_get):
@@ -189,14 +204,12 @@ class TestS2Tools:
         mock_post.return_value.status_code = 200
 
         result = get_multi_paper_recommendations.invoke(
-            input={
-                "paper_ids": ["123", "456"],
-                "limit": 1,
-                "tool_call_id": "test123",
-            }
+            input={"paper_ids": ["123", "456"], "limit": 1, "tool_call_id": "test123"}
         )
+
+        assert isinstance(result, Command)
         assert "multi_papers" in result.update
-        assert isinstance(result.update["messages"][0], ToolMessage)
+        assert len(result.update["messages"]) == 1
 
     @patch("requests.post")
     def test_multi_paper_rec_with_optional_params(self, mock_post):
@@ -217,3 +230,126 @@ class TestS2Tools:
         )
         assert "multi_papers" in result.update
         assert len(result.update["messages"]) == 1
+
+    @patch("requests.get")
+    def test_search_tool_finds_papers(self, mock_get):
+        """Verifies search tool finds and formats papers correctly"""
+        mock_get.return_value.json.return_value = MOCK_SEARCH_RESPONSE
+        mock_get.return_value.status_code = 200
+
+        result = search_tool.invoke(
+            input={"query": "machine learning", "limit": 1, "tool_call_id": "test123"}
+        )
+
+        assert isinstance(result, Command)  # Expect a Command object
+        assert "papers" in result.update
+        assert len(result.update["papers"]) > 0
+
+    def test_query_results_empty_state(self, initial_state):
+        """Tests query_results tool behavior when no papers are found."""
+        with pytest.raises(
+            NoPapersFoundError,
+            match="No papers found. A search needs to be performed first.",
+        ):
+            query_results.invoke(
+                {"question": "List all papers", "state": initial_state}
+            )
+
+    @patch(
+        "aiagents4pharma.talk2scholars.tools.s2.query_results.create_pandas_dataframe_agent"
+    )
+    def test_query_results_with_papers(self, mock_create_agent, initial_state):
+        """Tests querying papers when data is available."""
+        state = initial_state.copy()
+        state["last_displayed_papers"] = "papers"
+        state["papers"] = MOCK_STATE_PAPER
+
+        # Mock the dataframe agent instead of the LLM
+        mock_agent = MagicMock()
+        mock_agent.invoke.return_value = {"output": "Mocked response"}
+
+        mock_create_agent.return_value = (
+            mock_agent  # Mock the function returning the agent
+        )
+
+        # Ensure that the output of query_results is correctly structured
+        result = query_results.invoke({"question": "List all papers", "state": state})
+
+        assert isinstance(result, str)  # Ensure output is a string
+        assert result == "Mocked response"  # Validate the expected response
+
+    @patch("requests.get")
+    def test_retrieve_semantic_scholar_paper_id(self, mock_get):
+        """Tests retrieving a paper ID from Semantic Scholar."""
+        mock_get.return_value.json.return_value = MOCK_SEARCH_RESPONSE
+        mock_get.return_value.status_code = 200
+
+        result = retrieve_semantic_scholar_paper_id.invoke(
+            input={"paper_title": "Machine Learning Basics", "tool_call_id": "test123"}
+        )
+
+        assert isinstance(result, Command)
+        assert "messages" in result.update
+        assert (
+            "Paper ID for 'Machine Learning Basics' is: 123"
+            in result.update["messages"][0].content
+        )
+
+    def test_retrieve_semantic_scholar_paper_id_no_results(self):
+        """Test retrieving a paper ID when no results are found."""
+        with pytest.raises(ValueError, match="No papers found for query: UnknownPaper"):
+            retrieve_semantic_scholar_paper_id.invoke(
+                input={"paper_title": "UnknownPaper", "tool_call_id": "test123"}
+            )
+
+    def test_single_paper_rec_invalid_id(self):
+        """Test single paper recommendation with an invalid ID."""
+        with pytest.raises(ValueError, match="Invalid paper ID or API error."):
+            get_single_paper_recommendations.invoke(
+                input={"paper_id": "", "tool_call_id": "test123"}  # Empty ID case
+            )
+
+    @patch("requests.post")
+    def test_multi_paper_rec_no_recommendations(self, mock_post):
+        """Tests behavior when multi-paper recommendation API returns no results."""
+        mock_post.return_value.json.return_value = {"recommendedPapers": []}
+        mock_post.return_value.status_code = 200
+
+        result = get_multi_paper_recommendations.invoke(
+            input={"paper_ids": ["123", "456"], "limit": 1, "tool_call_id": "test123"}
+        )
+
+        assert isinstance(result, Command)
+        assert "messages" in result.update
+        assert (
+            "No recommendations found based on multiple papers."
+            in result.update["messages"][0].content
+        )
+
+    @patch("requests.get")
+    def test_search_no_results(self, mock_get):
+        """Tests behavior when search API returns no results."""
+        mock_get.return_value.json.return_value = {"data": []}
+        mock_get.return_value.status_code = 200
+
+        result = search_tool.invoke(
+            input={"query": "nonexistent topic", "limit": 1, "tool_call_id": "test123"}
+        )
+
+        assert isinstance(result, Command)
+        assert "messages" in result.update
+        assert "No papers found." in result.update["messages"][0].content
+
+    @patch("requests.get")
+    def test_single_paper_rec_no_recommendations(self, mock_get):
+        """Tests behavior when single paper recommendation API returns no results."""
+        mock_get.return_value.json.return_value = {"recommendedPapers": []}
+        mock_get.return_value.status_code = 200
+
+        result = get_single_paper_recommendations.invoke(
+            input={"paper_id": "123", "limit": 1, "tool_call_id": "test123"}
+        )
+
+        assert isinstance(result, Command)
+        assert "messages" in result.update
+        assert "No recommendations found for" in result.update["messages"][0].content

aiagents4pharma/talk2scholars/tools/s2/__init__.py

@@ -1,8 +1,10 @@
-'''
+"""
 This file is used to import all the modules in the package.
-'''
+"""
 
 from . import display_results
 from . import multi_paper_rec
 from . import search
 from . import single_paper_rec
+from . import query_results
+from . import retrieve_semantic_scholar_paper_id

aiagents4pharma/talk2scholars/tools/s2/display_results.py

@@ -1,13 +1,24 @@
 #!/usr/bin/env python3
 
+
 """
-This tool is used to display the table of studies.
+Tool for displaying search or recommendation results.
+
+This module defines a tool that retrieves and displays a table of research papers
+found during searches or recommendations. If no papers are found, an exception is raised
+to signal the need for a new search.
 """
 
+
 import logging
-from typing import Annotated, Dict, Any
+
+from typing import Annotated
+from langchain_core.messages import ToolMessage
 from langchain_core.tools import tool
+from langchain_core.tools.base import InjectedToolCallId
 from langgraph.prebuilt import InjectedState
+from langgraph.types import Command
+
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -15,36 +26,64 @@ logger = logging.getLogger(__name__)
 
 
 class NoPapersFoundError(Exception):
-    """Exception raised when no papers are found in the state."""
+    """
+    Exception raised when no research papers are found in the agent's state.
 
+    This exception helps the language model determine whether a new search
+    or recommendation should be initiated.
 
-@tool("display_results")
-def display_results(state: Annotated[dict, InjectedState]) -> Dict[str, Any]:
+    Example:
+        >>> if not papers:
+        >>>     raise NoPapersFoundError("No papers found. A search is needed.")
     """
-    Display the papers in the state. If no papers are found, raises an exception
-    indicating that a search is needed.
+
+
+@tool("display_results", parse_docstring=True)
+def display_results(
+    tool_call_id: Annotated[str, InjectedToolCallId],
+    state: Annotated[dict, InjectedState],
+) -> Command:
+    """
+    Displays retrieved research papers after a search or recommendation.
+
+    This function retrieves the last displayed research papers from the state and
+    returns them as an artifact for further processing. If no papers are found,
+    it raises a `NoPapersFoundError` to indicate that a new search is needed.
 
     Args:
-        state (dict): The state of the agent containing the papers.
+        tool_call_id (Annotated[str, InjectedToolCallId]): The tool call ID for tracking.
+        state (dict): The agent's state containing retrieved papers.
 
     Returns:
-        dict: A dictionary containing the papers and multi_papers from the state.
+        Command: A command containing a message with the number of displayed papers
+                 and an attached artifact for further reference.
 
     Raises:
-        NoPapersFoundError: If no papers are found in the state.
+        NoPapersFoundError: If no research papers are found in the agent's state.
 
-    Note:
-        The exception allows the LLM to make a more informed decision about initiating a search.
+    Example:
+        >>> state = {"last_displayed_papers": {"paper1": "Title 1", "paper2": "Title 2"}}
+        >>> result = display_results(tool_call_id="123", state=state)
+        >>> print(result.update["messages"][0].content)
+        "2 papers found. Papers are attached as an artifact."
     """
-    logger.info("Displaying papers from the state")
-
-    if not state.get("papers") and not state.get("multi_papers"):
+    logger.info("Displaying papers")
+    context_key = state.get("last_displayed_papers")
+    artifact = state.get(context_key)
+    if not artifact:
         logger.info("No papers found in state, raising NoPapersFoundError")
         raise NoPapersFoundError(
-            "No papers found. A search needs to be performed first."
+            "No papers found. A search/rec needs to be performed first."
         )
-
-    return {
-        "papers": state.get("papers"),
-        "multi_papers": state.get("multi_papers"),
-    }
+    content = f"{len(artifact)} papers found. Papers are attached as an artifact."
+    return Command(
+        update={
+            "messages": [
+                ToolMessage(
+                    content=content,
+                    tool_call_id=tool_call_id,
+                    artifact=artifact,
+                )
+            ],
+        }
+    )

aiagents4pharma/talk2scholars/tools/s2/multi_paper_rec.py

@@ -7,7 +7,7 @@ multi_paper_rec: Tool for getting recommendations
 
 import json
 import logging
-from typing import Annotated, Any, Dict, List, Optional
+from typing import Annotated, Any, List, Optional
 import hydra
 import requests
 from langchain_core.messages import ToolMessage
@@ -52,15 +52,16 @@ with hydra.initialize(version_base=None, config_path="../../configs"):
     cfg = cfg.tools.multi_paper_recommendation
 
 
-@tool(args_schema=MultiPaperRecInput)
+@tool(args_schema=MultiPaperRecInput, parse_docstring=True)
 def get_multi_paper_recommendations(
     paper_ids: List[str],
     tool_call_id: Annotated[str, InjectedToolCallId],
     limit: int = 2,
     year: Optional[str] = None,
-) -> Dict[str, Any]:
+) -> Command[Any]:
     """
-    Get paper recommendations based on multiple papers.
+    Get recommendations for a group of multiple papers using the Semantic Scholar IDs.
+    No other paper IDs are supported.
 
     Args:
         paper_ids (List[str]): The list of paper IDs to base recommendations on.
@@ -72,7 +73,9 @@ def get_multi_paper_recommendations(
     Returns:
         Dict[str, Any]: The recommendations and related information.
     """
-    logging.info("Starting multi-paper recommendations search.")
+    logging.info(
+        "Starting multi-paper recommendations search with paper IDs: %s", paper_ids
+    )
 
     endpoint = cfg.api_endpoint
     headers = cfg.headers
@@ -101,26 +104,50 @@ def get_multi_paper_recommendations(
     data = response.json()
     recommendations = data.get("recommendedPapers", [])
 
+    if not recommendations:
+        return Command(
+            update={  # Place 'messages' inside 'update'
+                "messages": [
+                    ToolMessage(
+                        content="No recommendations found based on multiple papers.",
+                        tool_call_id=tool_call_id,
+                    )
+                ]
+            }
+        )
+
     # Create a dictionary to store the papers
     filtered_papers = {
         paper["paperId"]: {
+            # "semantic_scholar_id": paper["paperId"],  # Store Semantic Scholar ID
             "Title": paper.get("title", "N/A"),
             "Abstract": paper.get("abstract", "N/A"),
             "Year": paper.get("year", "N/A"),
             "Citation Count": paper.get("citationCount", "N/A"),
             "URL": paper.get("url", "N/A"),
+            # "arXiv_ID": paper.get("externalIds", {}).get(
+            #     "ArXiv", "N/A"
+            # ),  # Extract arXiv ID
         }
         for paper in recommendations
-        if paper.get("title") and paper.get("paperId")
+        if paper.get("title") and paper.get("authors")
     }
 
+    content = "Recommendations based on multiple papers was successful."
+    content += " Here is a summary of the recommendations:"
+    content += f"Number of papers found: {len(filtered_papers)}\n"
+    content += f"Query Paper IDs: {', '.join(paper_ids)}\n"
+    content += f"Year: {year}\n" if year else ""
+
     return Command(
         update={
             "multi_papers": filtered_papers,  # Now sending the dictionary directly
+            "last_displayed_papers": "multi_papers",
             "messages": [
                 ToolMessage(
-                    content=f"Search Successful: {filtered_papers}",
-                    tool_call_id=tool_call_id
+                    content=content,
+                    tool_call_id=tool_call_id,
+                    artifact=filtered_papers,
                 )
             ],
         }

aiagents4pharma/talk2scholars/tools/s2/query_results.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+
+"""
+This tool is used to display the table of studies.
+"""
+
+import logging
+from typing import Annotated
+import pandas as pd
+from langchain_experimental.agents import create_pandas_dataframe_agent
+from langchain_core.tools import tool
+from langgraph.prebuilt import InjectedState
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class NoPapersFoundError(Exception):
+    """Exception raised when no papers are found in the state."""
+
+
+@tool("query_results", parse_docstring=True)
+def query_results(question: str, state: Annotated[dict, InjectedState]) -> str:
+    """
+    Query the last displayed papers from the state. If no papers are found,
+    raises an exception.
+
+    Use this also to get the last displayed papers from the state,
+    and then use the papers to get recommendations for a single paper or
+    multiple papers.
+
+    Args:
+        question (str): The question to ask the agent.
+        state (dict): The state of the agent containing the papers.
+
+    Returns:
+        str: A message with the last displayed papers.
+    """
+    logger.info("Querying last displayed papers with question: %s", question)
+    llm_model = state.get("llm_model")
+    if not state.get("last_displayed_papers"):
+        logger.info("No papers displayed so far, raising NoPapersFoundError")
+        raise NoPapersFoundError(
+            "No papers found. A search needs to be performed first."
+        )
+    context_key = state.get("last_displayed_papers")
+    dic_papers = state.get(context_key)
+    df_papers = pd.DataFrame.from_dict(dic_papers, orient="index")
+    df_agent = create_pandas_dataframe_agent(
+        llm_model,
+        allow_dangerous_code=True,
+        agent_type="tool-calling",
+        df=df_papers,
+        max_iterations=5,
+        include_df_in_prompt=True,
+        number_of_head_rows=df_papers.shape[0],
+        verbose=True,
+    )
+    llm_result = df_agent.invoke(question, stream_mode=None)
+    return llm_result["output"]

aiagents4pharma/talk2scholars/tools/s2/retrieve_semantic_scholar_paper_id.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+
+"""
+This tool is used to search for academic papers on Semantic Scholar.
+"""
+
+import logging
+from typing import Annotated, Any
+import hydra
+import requests
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import tool
+from langchain_core.tools.base import InjectedToolCallId
+from langgraph.types import Command
+from pydantic import Field
+
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Load hydra configuration
+with hydra.initialize(version_base=None, config_path="../../configs"):
+    cfg = hydra.compose(
+        config_name="config",
+        overrides=["tools/retrieve_semantic_scholar_paper_id=default"],
+    )
+    cfg = cfg.tools.retrieve_semantic_scholar_paper_id
+
+
+@tool("retrieve_semantic_scholar_paper_id", parse_docstring=True)
+def retrieve_semantic_scholar_paper_id(
+    tool_call_id: Annotated[str, InjectedToolCallId],
+    paper_title: str = Field(
+        description="The title of the paper to search for on Semantic Scholar."
+    ),
+) -> Command[Any]:
+    """
+    This tool can be used to search for a paper on Semantic Scholar
+    and retrieve the paper Semantic Scholar ID.
+
+    This is useful for when an article is retrieved from users Zotero library
+    and the Semantic Scholar ID is needed to retrieve more information about the paper.
+
+    Args:
+        tool_call_id (Annotated[str, InjectedToolCallId]): The tool call ID.
+        paper_title (str): The title of the paper to search for on Semantic Scholar.
+
+    Returns:
+        ToolMessage: A message containing the paper ID.
+    """
+    logger.info("Retrieving ID of paper with title: %s", paper_title)
+    endpoint = cfg.api_endpoint
+    params = {
+        "query": paper_title,
+        "limit": 1,
+        "fields": ",".join(cfg.api_fields),
+    }
+
+    response = requests.get(endpoint, params=params, timeout=10)
+    data = response.json()
+    papers = data.get("data", [])
+    logger.info("Received %d papers", len(papers))
+    if not papers:
+        logger.error("No papers found for query: %s", paper_title)
+        raise ValueError(f"No papers found for query: {paper_title}. Try again.")
+    # Get the paper ID
+    paper_id = papers[0]["paperId"]
+
+    return Command(
+        update={
+            "messages": [
+                ToolMessage(
+                    content=f"Paper ID for '{paper_title}' is: {paper_id}",
+                    tool_call_id=tool_call_id,
+                )
+            ],
+        }
+    )