aiagents4pharma 1.19.1__py3-none-any.whl → 1.20.1__py3-none-any.whl

aiagents4pharma/talk2scholars/tests/{test_langgraph.py → test_s2_tools.py}

@@ -1,59 +1,21 @@
 """
-Unit and integration tests for Talk2Scholars system.
-Each test focuses on a single, specific functionality.
-Tests are deterministic and independent of each other.
+Unit tests for S2 tools functionality.
 """
 
-from unittest.mock import Mock, patch
+# pylint: disable=redefined-outer-name
+from unittest.mock import patch
+from langchain_core.messages import ToolMessage
 import pytest
-from langchain_core.messages import AIMessage, HumanMessage
-import hydra
-from hydra.core.global_hydra import GlobalHydra
-from omegaconf import DictConfig, OmegaConf
-
-from ..agents.main_agent import get_app, make_supervisor_node
-from ..state.state_talk2scholars import replace_dict, Talk2Scholars
-from ..tools.s2.display_results import display_results
+from ..tools.s2.display_results import display_results, NoPapersFoundError
 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
 
-# pylint: disable=redefined-outer-name
-
-
-@pytest.fixture(autouse=True)
-def hydra_setup():
-    """Setup and cleanup Hydra for tests."""
-    GlobalHydra.instance().clear()
-    with hydra.initialize(version_base=None, config_path="../configs"):
-        yield
-
 
 @pytest.fixture
-def mock_cfg() -> DictConfig:
-    """Create a mock configuration for testing."""
-    config = {
-        "agents": {
-            "talk2scholars": {
-                "main_agent": {
-                    "state_modifier": "Test prompt for main agent",
-                    "temperature": 0,
-                },
-                "s2_agent": {
-                    "temperature": 0,
-                    "s2_agent": "Test prompt for s2 agent",
-                },
-            }
-        },
-        "tools": {
-            "search": {
-                "api_endpoint": "https://api.semanticscholar.org/graph/v1/paper/search",
-                "default_limit": 2,
-                "api_fields": ["paperId", "title", "abstract", "year", "authors"],
-            }
-        },
-    }
-    return OmegaConf.create(config)
+def initial_state():
+    """Provides an empty initial state for tests."""
+    return {"papers": {}, "multi_papers": {}}
 
 
 # Fixed test data for deterministic results
@@ -82,70 +44,26 @@ MOCK_STATE_PAPER = {
 }
 
 
-@pytest.fixture
-def initial_state() -> Talk2Scholars:
-    """Create a base state for tests"""
-    return Talk2Scholars(
-        messages=[],
-        papers={},
-        is_last_step=False,
-        current_agent=None,
-        llm_model="gpt-4o-mini",
-        next="",
-    )
-
-
-class TestMainAgent:
-    """Unit tests for main agent functionality"""
-
-    def test_supervisor_routes_search_to_s2(
-        self, initial_state: Talk2Scholars, mock_cfg
-    ):
-        """Verifies that search-related queries are routed to S2 agent"""
-        llm_mock = Mock()
-        llm_mock.invoke.return_value = AIMessage(content="Search initiated")
-
-        # Extract the main_agent config
-        supervisor = make_supervisor_node(
-            llm_mock, mock_cfg.agents.talk2scholars.main_agent
-        )
-        state = initial_state
-        state["messages"] = [HumanMessage(content="search for papers")]
-
-        result = supervisor(state)
-        assert result.goto == "s2_agent"
-        assert not result.update["is_last_step"]
-        assert result.update["current_agent"] == "s2_agent"
-
-    def test_supervisor_routes_general_to_end(
-        self, initial_state: Talk2Scholars, mock_cfg
-    ):
-        """Verifies that non-search queries end the conversation"""
-        llm_mock = Mock()
-        llm_mock.invoke.return_value = AIMessage(content="General response")
-
-        # Extract the main_agent config
-        supervisor = make_supervisor_node(
-            llm_mock, mock_cfg.agents.talk2scholars.main_agent
-        )
-        state = initial_state
-        state["messages"] = [HumanMessage(content="What is ML?")]
-
-        result = supervisor(state)
-        assert result.goto == "__end__"
-        assert result.update["is_last_step"]
-
-
 class TestS2Tools:
     """Unit tests for individual S2 tools"""
 
-    def test_display_results_shows_papers(self, initial_state: Talk2Scholars):
+    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.",
+        ):
+            display_results.invoke({"state": initial_state})
+
+    def test_display_results_shows_papers(self, initial_state):
         """Verifies display_results tool correctly returns papers from state"""
-        state = initial_state
+        state = initial_state.copy()
         state["papers"] = MOCK_STATE_PAPER
+        state["multi_papers"] = {}
         result = display_results.invoke(input={"state": state})
-        assert result == MOCK_STATE_PAPER
         assert isinstance(result, dict)
+        assert result["papers"] == MOCK_STATE_PAPER
+        assert result["multi_papers"] == {}
 
     @patch("requests.get")
     def test_search_finds_papers(self, mock_get):
@@ -171,6 +89,60 @@ class TestS2Tools:
         assert paper["Title"] == "Machine Learning Basics"
         assert paper["Year"] == 2023
 
+    @patch("requests.get")
+    def test_search_finds_papers_with_year(self, mock_get):
+        """Verifies search tool works with year parameter"""
+        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,
+                "year": "2023-",
+                "tool_call_id": "test123",
+                "id": "test123",
+            }
+        )
+
+        assert "papers" in result.update
+        assert "messages" in result.update
+        papers = result.update["papers"]
+        assert isinstance(papers, dict)
+        assert len(papers) > 0
+
+    @patch("requests.get")
+    def test_search_filters_invalid_papers(self, mock_get):
+        """Verifies search tool properly filters papers without title or authors"""
+        mock_response = {
+            "data": [
+                {
+                    "paperId": "123",
+                    "abstract": "An introduction to ML",
+                    "year": 2023,
+                    "citationCount": 100,
+                    "url": "https://example.com/paper1",
+                    # Missing title and authors
+                },
+                MOCK_SEARCH_RESPONSE["data"][0],  # This one is valid
+            ]
+        }
+        mock_get.return_value.json.return_value = mock_response
+        mock_get.return_value.status_code = 200
+
+        result = search_tool.invoke(
+            input={
+                "query": "machine learning",
+                "limit": 2,
+                "tool_call_id": "test123",
+                "id": "test123",
+            }
+        )
+
+        assert "papers" in result.update
+        papers = result.update["papers"]
+        assert len(papers) == 1  # Only the valid paper should be included
+
     @patch("requests.get")
     def test_single_paper_rec_basic(self, mock_get):
         """Tests basic single paper recommendation functionality"""
@@ -184,11 +156,10 @@ class TestS2Tools:
                 "paper_id": "123",
                 "limit": 1,
                 "tool_call_id": "test123",
-                "id": "test123",
             }
         )
         assert "papers" in result.update
-        assert len(result.update["messages"]) == 1
+        assert isinstance(result.update["messages"][0], ToolMessage)
 
     @patch("requests.get")
     def test_single_paper_rec_with_optional_params(self, mock_get):
@@ -222,11 +193,10 @@ class TestS2Tools:
                 "paper_ids": ["123", "456"],
                 "limit": 1,
                 "tool_call_id": "test123",
-                "id": "test123",
             }
         )
-        assert "papers" in result.update
-        assert len(result.update["messages"]) == 1
+        assert "multi_papers" in result.update
+        assert isinstance(result.update["messages"][0], ToolMessage)
 
     @patch("requests.post")
     def test_multi_paper_rec_with_optional_params(self, mock_post):
@@ -245,47 +215,5 @@ class TestS2Tools:
                 "id": "test123",
             }
         )
-        assert "papers" in result.update
+        assert "multi_papers" in result.update
         assert len(result.update["messages"]) == 1
-
-
-def test_state_replace_dict():
-    """Verifies state dictionary replacement works correctly"""
-    existing = {"key1": "value1", "key2": "value2"}
-    new = {"key3": "value3"}
-    result = replace_dict(existing, new)
-    assert result == new
-    assert isinstance(result, dict)
-
-
-@pytest.mark.integration
-def test_end_to_end_search_workflow(initial_state: Talk2Scholars, mock_cfg):
-    """Integration test: Complete search workflow"""
-    with (
-        patch("requests.get") as mock_get,
-        patch("langchain_openai.ChatOpenAI") as mock_llm,
-        patch("hydra.compose", return_value=mock_cfg),
-        patch("hydra.initialize"),
-    ):
-        mock_get.return_value.json.return_value = MOCK_SEARCH_RESPONSE
-        mock_get.return_value.status_code = 200
-
-        llm_instance = Mock()
-        llm_instance.invoke.return_value = AIMessage(content="Search completed")
-        mock_llm.return_value = llm_instance
-
-        app = get_app("test_integration")
-        test_state = initial_state
-        test_state["messages"] = [HumanMessage(content="search for ML papers")]
-
-        config = {
-            "configurable": {
-                "thread_id": "test_integration",
-                "checkpoint_ns": "test",
-                "checkpoint_id": "test123",
-            }
-        }
-
-        response = app.invoke(test_state, config)
-        assert "papers" in response
-        assert len(response["messages"]) > 0

aiagents4pharma/talk2scholars/tests/test_state.py

@@ -0,0 +1,14 @@
+"""
+Tests for state management functionality.
+"""
+
+from ..state.state_talk2scholars import replace_dict
+
+
+def test_state_replace_dict():
+    """Verifies state dictionary replacement works correctly"""
+    existing = {"key1": "value1", "key2": "value2"}
+    new = {"key3": "value3"}
+    result = replace_dict(existing, new)
+    assert result == new
+    assert isinstance(result, dict)

aiagents4pharma/talk2scholars/tools/s2/display_results.py

@@ -1,11 +1,11 @@
 #!/usr/bin/env python3
 
-'''
+"""
 This tool is used to display the table of studies.
-'''
+"""
 
 import logging
-from typing import Annotated
+from typing import Annotated, Dict, Any
 from langchain_core.tools import tool
 from langgraph.prebuilt import InjectedState
 
@@ -13,13 +13,38 @@ from langgraph.prebuilt import InjectedState
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-@tool('display_results')
-def display_results(state: Annotated[dict, InjectedState]):
+
+class NoPapersFoundError(Exception):
+    """Exception raised when no papers are found in the state."""
+
+
+@tool("display_results")
+def display_results(state: Annotated[dict, InjectedState]) -> Dict[str, Any]:
     """
-    Display the papers in the state.
+    Display the papers in the state. If no papers are found, raises an exception
+    indicating that a search is needed.
 
     Args:
-        state (dict): The state of the agent.
+        state (dict): The state of the agent containing the papers.
+
+    Returns:
+        dict: A dictionary containing the papers and multi_papers from the state.
+
+    Raises:
+        NoPapersFoundError: If no papers are found in the state.
+
+    Note:
+        The exception allows the LLM to make a more informed decision about initiating a search.
     """
     logger.info("Displaying papers from the state")
-    return state["papers"]
+
+    if not state.get("papers") and not state.get("multi_papers"):
+        logger.info("No papers found in state, raising NoPapersFoundError")
+        raise NoPapersFoundError(
+            "No papers found. A search needs to be performed first."
+        )
+
+    return {
+        "papers": state.get("papers"),
+        "multi_papers": state.get("multi_papers"),
+    }

aiagents4pharma/talk2scholars/tools/s2/multi_paper_rec.py

@@ -9,7 +9,6 @@ import json
 import logging
 from typing import Annotated, Any, Dict, List, Optional
 import hydra
-import pandas as pd
 import requests
 from langchain_core.messages import ToolMessage
 from langchain_core.tools import tool
@@ -18,6 +17,11 @@ from langgraph.types import Command
 from pydantic import BaseModel, Field
 
 
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
 class MultiPaperRecInput(BaseModel):
     """Input schema for multiple paper recommendations tool."""
 
@@ -110,31 +114,14 @@ def get_multi_paper_recommendations(
         if paper.get("title") and paper.get("paperId")
     }
 
-    # Create a DataFrame from the dictionary
-    df = pd.DataFrame.from_dict(filtered_papers, orient="index")
-    # print("Created DataFrame with results:")
-    logging.info("Created DataFrame with results: %s", df)
-
-    # Format papers for state update
-    papers = [
-        f"Paper ID: {paper_id}\n"
-        f"Title: {paper_data['Title']}\n"
-        f"Abstract: {paper_data['Abstract']}\n"
-        f"Year: {paper_data['Year']}\n"
-        f"Citations: {paper_data['Citation Count']}\n"
-        f"URL: {paper_data['URL']}\n"
-        for paper_id, paper_data in filtered_papers.items()
-    ]
-
-    # Convert DataFrame to markdown table
-    markdown_table = df.to_markdown(tablefmt="grid")
-    logging.info("Search results: %s", papers)
-
     return Command(
         update={
-            "papers": filtered_papers,  # Now sending the dictionary directly
+            "multi_papers": filtered_papers,  # Now sending the dictionary directly
             "messages": [
-                ToolMessage(content=markdown_table, tool_call_id=tool_call_id)
+                ToolMessage(
+                    content=f"Search Successful: {filtered_papers}",
+                    tool_call_id=tool_call_id
+                )
             ],
         }
     )

aiagents4pharma/talk2scholars/tools/s2/search.py

@@ -7,7 +7,6 @@ This tool is used to search for academic papers on Semantic Scholar.
 import logging
 from typing import Annotated, Any, Dict, Optional
 import hydra
-import pandas as pd
 import requests
 from langchain_core.messages import ToolMessage
 from langchain_core.tools import tool
@@ -16,6 +15,11 @@ from langgraph.types import Command
 from pydantic import BaseModel, Field
 
 
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
 class SearchInput(BaseModel):
     """Input schema for the search papers tool."""
 
@@ -65,8 +69,6 @@ def search_tool(
     params = {
         "query": query,
         "limit": min(limit, 100),
-        # "fields": "paperId,title,abstract,year,authors,
-        # citationCount,url,publicationTypes,openAccessPdf",
         "fields": ",".join(cfg.api_fields),
     }
 
@@ -77,7 +79,7 @@ def search_tool(
     response = requests.get(endpoint, params=params, timeout=10)
     data = response.json()
     papers = data.get("data", [])
-
+    logger.info("Received %d papers", len(papers))
     # Create a dictionary to store the papers
     filtered_papers = {
         paper["paperId"]: {
@@ -86,40 +88,19 @@ def search_tool(
             "Year": paper.get("year", "N/A"),
             "Citation Count": paper.get("citationCount", "N/A"),
             "URL": paper.get("url", "N/A"),
-            # "Publication Type": paper.get("publicationTypes", ["N/A"])[0]
-            # if paper.get("publicationTypes")
-            # else "N/A",
-            # "Open Access PDF": paper.get("openAccessPdf", {}).get("url", "N/A")
-            # if paper.get("openAccessPdf") is not None
-            # else "N/A",
         }
         for paper in papers
         if paper.get("title") and paper.get("authors")
     }
 
-    df = pd.DataFrame(filtered_papers)
-
-    # Format papers for state update
-    papers = [
-        f"Paper ID: {paper_id}\n"
-        f"Title: {paper_data['Title']}\n"
-        f"Abstract: {paper_data['Abstract']}\n"
-        f"Year: {paper_data['Year']}\n"
-        f"Citations: {paper_data['Citation Count']}\n"
-        f"URL: {paper_data['URL']}\n"
-        # f"Publication Type: {paper_data['Publication Type']}\n"
-        # f"Open Access PDF: {paper_data['Open Access PDF']}"
-        for paper_id, paper_data in filtered_papers.items()
-    ]
-
-    markdown_table = df.to_markdown(tablefmt="grid")
-    logging.info("Search results: %s", papers)
-
     return Command(
         update={
             "papers": filtered_papers,  # Now sending the dictionary directly
             "messages": [
-                ToolMessage(content=markdown_table, tool_call_id=tool_call_id)
+                ToolMessage(
+                    content=f"Search Successful: {filtered_papers}",
+                    tool_call_id=tool_call_id
+                )
             ],
         }
     )

aiagents4pharma/talk2scholars/tools/s2/single_paper_rec.py

@@ -7,7 +7,6 @@ This tool is used to return recommendations for a single paper.
 import logging
 from typing import Annotated, Any, Dict, Optional
 import hydra
-import pandas as pd
 import requests
 from langchain_core.messages import ToolMessage
 from langchain_core.tools import tool
@@ -84,7 +83,6 @@ def get_single_paper_recommendations(
 
     response = requests.get(endpoint, params=params, timeout=cfg.request_timeout)
     data = response.json()
-    papers = data.get("data", [])
     response = requests.get(endpoint, params=params, timeout=10)
     # print(f"API Response Status: {response.status_code}")
     logging.info(
@@ -106,42 +104,19 @@ def get_single_paper_recommendations(
             "Year": paper.get("year", "N/A"),
             "Citation Count": paper.get("citationCount", "N/A"),
             "URL": paper.get("url", "N/A"),
-            # "Publication Type": paper.get("publicationTypes", ["N/A"])[0]
-            # if paper.get("publicationTypes")
-            # else "N/A",
-            # "Open Access PDF": paper.get("openAccessPdf", {}).get("url", "N/A")
-            # if paper.get("openAccessPdf") is not None
-            # else "N/A",
         }
         for paper in recommendations
         if paper.get("title") and paper.get("authors")
     }
 
-    # Create a DataFrame for pretty printing
-    df = pd.DataFrame(filtered_papers)
-
-    # Format papers for state update
-    papers = [
-        f"Paper ID: {paper_id}\n"
-        f"Title: {paper_data['Title']}\n"
-        f"Abstract: {paper_data['Abstract']}\n"
-        f"Year: {paper_data['Year']}\n"
-        f"Citations: {paper_data['Citation Count']}\n"
-        f"URL: {paper_data['URL']}\n"
-        # f"Publication Type: {paper_data['Publication Type']}\n"
-        # f"Open Access PDF: {paper_data['Open Access PDF']}"
-        for paper_id, paper_data in filtered_papers.items()
-    ]
-
-    # Convert DataFrame to markdown table
-    markdown_table = df.to_markdown(tablefmt="grid")
-    logging.info("Search results: %s", papers)
-
     return Command(
         update={
             "papers": filtered_papers,  # Now sending the dictionary directly
             "messages": [
-                ToolMessage(content=markdown_table, tool_call_id=tool_call_id)
+                ToolMessage(
+                    content=f"Search Successful: {filtered_papers}",
+                    tool_call_id=tool_call_id
+                )
             ],
         }
     )

{aiagents4pharma-1.19.1.dist-info → aiagents4pharma-1.20.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: aiagents4pharma
-Version: 1.19.1
+Version: 1.20.1
 Summary: AI Agents for drug discovery, drug development, and other pharmaceutical R&D.
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
@@ -56,6 +56,7 @@ Requires-Dist: streamlit-feedback
 [![TESTS Talk2Scholars](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2scholars.yml/badge.svg)](https://github.com/VirtualPatientEngine/AIAgents4Pharma/actions/workflows/tests_talk2scholars.yml)
 ![GitHub Release](https://img.shields.io/github/v/release/VirtualPatientEngine/AIAgents4Pharma)
 ![Python Version from PEP 621 TOML](https://img.shields.io/python/required-version-toml?tomlFilePath=https%3A%2F%2Fraw.githubusercontent.com%2FVirtualPatientEngine%2FAIAgents4Pharma%2Frefs%2Fheads%2Fmain%2Fpyproject.toml)
+![Docker Pulls](https://img.shields.io/docker/pulls/virtualpatientengine/talk2biomodels?link=https%3A%2F%2Fhub.docker.com%2Frepository%2Fdocker%2Fvirtualpatientengine%2Ftalk2biomodels%2Fgeneral)
 
 
 ## Introduction
@@ -85,7 +86,22 @@ pip install aiagents4pharma
 
 Check out the tutorials on each agent for detailed instrcutions.
 
-#### Option 2: git
+#### Option 2: docker hub
+
+_Please note that this option is currently available only for Talk2Biomodels._
+
+1. **Pull the image**
+   ```
+   docker pull virtualpatientengine/talk2biomodels
+   ```
+2. **Run a container**
+   ```
+   docker run -e OPENAI_API_KEY=<openai_api_key> -e NVIDIA_API_KEY=<nvidia_api_key> -p 8501:8501 virtualpatientengine/talk2biomodels
+   ```
+_You can create a free account at NVIDIA and apply for their
+free credits [here](https://build.nvidia.com/explore/discover)._
+
+#### Option 3: git
 
 1. **Clone the repository:**
    ```bash
@@ -94,7 +110,7 @@ Check out the tutorials on each agent for detailed instrcutions.
    ```
 2. **Install dependencies:**
    ```bash
-   pip install .
+   pip install -r requirements.txt
    ```
 3. **Initialize OPENAI_API_KEY and NVIDIA_API_KEY**
    ```bash