aiagents4pharma 1.9.0__py3-none-any.whl → 1.11.0__py3-none-any.whl

aiagents4pharma/talk2competitors/tests/test_langgraph.py

@@ -0,0 +1,274 @@
+"""
+Unit and integration tests for Talk2Competitors system.
+Each test focuses on a single, specific functionality.
+Tests are deterministic and independent of each other.
+"""
+
+from unittest.mock import Mock, patch
+
+import pytest
+from langchain_core.messages import AIMessage, HumanMessage
+
+from ..agents.main_agent import get_app, make_supervisor_node
+from ..state.state_talk2competitors import replace_dict
+from ..tools.s2.display_results import display_results
+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
+
+# Fixed test data for deterministic results
+MOCK_SEARCH_RESPONSE = {
+    "data": [
+        {
+            "paperId": "123",
+            "title": "Machine Learning Basics",
+            "abstract": "An introduction to ML",
+            "year": 2023,
+            "citationCount": 100,
+            "url": "https://example.com/paper1",
+            "authors": [{"name": "Test Author"}],
+        }
+    ]
+}
+
+MOCK_STATE_PAPER = {
+    "123": {
+        "Title": "Machine Learning Basics",
+        "Abstract": "An introduction to ML",
+        "Year": 2023,
+        "Citation Count": 100,
+        "URL": "https://example.com/paper1",
+    }
+}
+
+
+@pytest.fixture
+def initial_state():
+    """Create a base state for tests"""
+    return {
+        "messages": [],
+        "papers": {},
+        "is_last_step": False,
+        "current_agent": None,
+        "llm_model": "gpt-4o-mini",
+    }
+
+
+class TestMainAgent:
+    """Unit tests for main agent functionality"""
+
+    def test_supervisor_routes_search_to_s2(self, initial_state):
+        """Verifies that search-related queries are routed to S2 agent"""
+        llm_mock = Mock()
+        llm_mock.invoke.return_value = AIMessage(content="Search initiated")
+
+        supervisor = make_supervisor_node(llm_mock)
+        state = initial_state.copy()
+        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):
+        """Verifies that non-search queries end the conversation"""
+        llm_mock = Mock()
+        llm_mock.invoke.return_value = AIMessage(content="General response")
+
+        supervisor = make_supervisor_node(llm_mock)
+        state = initial_state.copy()
+        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):
+        """Verifies display_results tool correctly returns papers from state"""
+        state = initial_state.copy()
+        state["papers"] = MOCK_STATE_PAPER
+        result = display_results.invoke(input={"state": state})
+        assert result == MOCK_STATE_PAPER
+        assert isinstance(result, dict)
+
+    @patch("requests.get")
+    def test_search_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",
+                "id": "test123",
+            }
+        )
+
+        assert "papers" in result.update
+        assert "messages" in result.update
+        papers = result.update["papers"]
+        assert isinstance(papers, dict)
+        assert len(papers) > 0
+        paper = next(iter(papers.values()))
+        assert paper["Title"] == "Machine Learning Basics"
+        assert paper["Year"] == 2023
+
+    @patch("requests.get")
+    def test_single_paper_rec_basic(self, mock_get):
+        """Tests basic single paper recommendation functionality"""
+        mock_get.return_value.json.return_value = {
+            "recommendedPapers": [MOCK_SEARCH_RESPONSE["data"][0]]
+        }
+        mock_get.return_value.status_code = 200
+
+        result = get_single_paper_recommendations.invoke(
+            input={
+                "paper_id": "123",
+                "limit": 1,
+                "tool_call_id": "test123",
+                "id": "test123",
+            }
+        )
+        assert "papers" in result.update
+        assert len(result.update["messages"]) == 1
+
+    @patch("requests.get")
+    def test_single_paper_rec_with_optional_params(self, mock_get):
+        """Tests single paper recommendations with year parameter"""
+        mock_get.return_value.json.return_value = {
+            "recommendedPapers": [MOCK_SEARCH_RESPONSE["data"][0]]
+        }
+        mock_get.return_value.status_code = 200
+
+        result = get_single_paper_recommendations.invoke(
+            input={
+                "paper_id": "123",
+                "limit": 1,
+                "year": "2023-",
+                "tool_call_id": "test123",
+                "id": "test123",
+            }
+        )
+        assert "papers" in result.update
+
+    @patch("requests.post")
+    def test_multi_paper_rec_basic(self, mock_post):
+        """Tests basic multi-paper recommendation functionality"""
+        mock_post.return_value.json.return_value = {
+            "recommendedPapers": [MOCK_SEARCH_RESPONSE["data"][0]]
+        }
+        mock_post.return_value.status_code = 200
+
+        result = get_multi_paper_recommendations.invoke(
+            input={
+                "paper_ids": ["123", "456"],
+                "limit": 1,
+                "tool_call_id": "test123",
+                "id": "test123",
+            }
+        )
+        assert "papers" in result.update
+        assert len(result.update["messages"]) == 1
+
+    @patch("requests.post")
+    def test_multi_paper_rec_with_optional_params(self, mock_post):
+        """Tests multi-paper recommendations with all optional parameters"""
+        mock_post.return_value.json.return_value = {
+            "recommendedPapers": [MOCK_SEARCH_RESPONSE["data"][0]]
+        }
+        mock_post.return_value.status_code = 200
+
+        result = get_multi_paper_recommendations.invoke(
+            input={
+                "paper_ids": ["123", "456"],
+                "limit": 1,
+                "year": "2023-",
+                "tool_call_id": "test123",
+                "id": "test123",
+            }
+        )
+        assert "papers" in result.update
+        assert len(result.update["messages"]) == 1
+
+    @patch("requests.get")
+    def test_single_paper_rec_empty_response(self, mock_get):
+        """Tests single paper recommendations with empty response"""
+        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",
+                "id": "test123",
+            }
+        )
+        assert "papers" in result.update
+        assert len(result.update["papers"]) == 0
+
+    @patch("requests.post")
+    def test_multi_paper_rec_empty_response(self, mock_post):
+        """Tests multi-paper recommendations with empty response"""
+        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",
+                "id": "test123",
+            }
+        )
+        assert "papers" in result.update
+        assert len(result.update["papers"]) == 0
+
+
+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):
+    """Integration test: Complete search workflow"""
+    with (
+        patch("requests.get") as mock_get,
+        patch("langchain_openai.ChatOpenAI") as mock_llm,
+    ):
+        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.copy()
+        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/talk2competitors/tools/__init__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+
+'''
+Import statements
+'''
+
+from . import s2

aiagents4pharma/talk2competitors/tools/s2/__init__.py

@@ -0,0 +1,8 @@
+'''
+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

aiagents4pharma/talk2competitors/tools/s2/display_results.py

@@ -0,0 +1,25 @@
+#!/usr/bin/env python3
+
+'''
+This tool is used to display the table of studies.
+'''
+
+import logging
+from typing import Annotated
+from langchain_core.tools import tool
+from langgraph.prebuilt import InjectedState
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+@tool('display_results')
+def display_results(state: Annotated[dict, InjectedState]):
+    """
+    Display the papers in the state.
+
+    Args:
+        state (dict): The state of the agent.
+    """
+    logger.info("Displaying papers from the state")
+    return state["papers"]

aiagents4pharma/talk2competitors/tools/s2/multi_paper_rec.py

@@ -0,0 +1,132 @@
+#!/usr/bin/env python3
+
+"""
+multi_paper_rec: Tool for getting recommendations
+                based on multiple papers
+"""
+
+import json
+import logging
+from typing import Annotated, Any, Dict, List, Optional
+
+import pandas as pd
+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 BaseModel, Field
+
+
+class MultiPaperRecInput(BaseModel):
+    """Input schema for multiple paper recommendations tool."""
+
+    paper_ids: List[str] = Field(
+        description=("List of Semantic Scholar Paper IDs to get recommendations for")
+    )
+    limit: int = Field(
+        default=2,
+        description="Maximum total number of recommendations to return",
+        ge=1,
+        le=500,
+    )
+    year: Optional[str] = Field(
+        default=None,
+        description="Year range in format: YYYY for specific year, "
+        "YYYY- for papers after year, -YYYY for papers before year, or YYYY:YYYY for range",
+    )
+    tool_call_id: Annotated[str, InjectedToolCallId]
+
+    model_config = {"arbitrary_types_allowed": True}
+
+
+@tool(args_schema=MultiPaperRecInput)
+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]:
+    """
+    Get paper recommendations based on multiple papers.
+
+    Args:
+        paper_ids (List[str]): The list of paper IDs to base recommendations on.
+        tool_call_id (Annotated[str, InjectedToolCallId]): The tool call ID.
+        limit (int, optional): The maximum number of recommendations to return. Defaults to 2.
+        year (str, optional): Year range for papers.
+        Supports formats like "2024-", "-2024", "2024:2025". Defaults to None.
+
+    Returns:
+        Dict[str, Any]: The recommendations and related information.
+    """
+    logging.info("Starting multi-paper recommendations search.")
+
+    endpoint = "https://api.semanticscholar.org/recommendations/v1/papers"
+    headers = {"Content-Type": "application/json"}
+    payload = {"positivePaperIds": paper_ids, "negativePaperIds": []}
+    params = {
+        "limit": min(limit, 500),
+        "fields": "paperId,title,abstract,year,authors,citationCount,url",
+    }
+
+    # Add year parameter if provided
+    if year:
+        params["year"] = year
+
+    # Getting recommendations
+    response = requests.post(
+        endpoint,
+        headers=headers,
+        params=params,
+        data=json.dumps(payload),
+        timeout=10,
+    )
+    logging.info(
+        "API Response Status for multi-paper recommendations: %s", response.status_code
+    )
+
+    data = response.json()
+    recommendations = data.get("recommendedPapers", [])
+
+    # Create a dictionary to store the papers
+    filtered_papers = {
+        paper["paperId"]: {
+            "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"),
+        }
+        for paper in 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
+            "messages": [
+                ToolMessage(content=markdown_table, tool_call_id=tool_call_id)
+            ],
+        }
+    )

aiagents4pharma/talk2competitors/tools/s2/search.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python3
+
+"""
+This tool is used to search for academic papers on Semantic Scholar.
+"""
+
+import logging
+from typing import Annotated, Any, Dict, Optional
+
+import pandas as pd
+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 BaseModel, Field
+
+
+class SearchInput(BaseModel):
+    """Input schema for the search papers tool."""
+
+    query: str = Field(
+        description="Search query string to find academic papers."
+        "Be specific and include relevant academic terms."
+    )
+    limit: int = Field(
+        default=2, description="Maximum number of results to return", ge=1, le=100
+    )
+    year: Optional[str] = Field(
+        default=None,
+        description="Year range in format: YYYY for specific year, "
+        "YYYY- for papers after year, -YYYY for papers before year, or YYYY:YYYY for range",
+    )
+    tool_call_id: Annotated[str, InjectedToolCallId]
+
+
+@tool(args_schema=SearchInput)
+def search_tool(
+    query: str,
+    tool_call_id: Annotated[str, InjectedToolCallId],
+    limit: int = 2,
+    year: Optional[str] = None,
+) -> Dict[str, Any]:
+    """
+    Search for academic papers on Semantic Scholar.
+
+    Args:
+        query (str): The search query string to find academic papers.
+        tool_call_id (Annotated[str, InjectedToolCallId]): The tool call ID.
+        limit (int, optional): The maximum number of results to return. Defaults to 2.
+        year (str, optional): Year range for papers.
+        Supports formats like "2024-", "-2024", "2024:2025". Defaults to None.
+
+    Returns:
+        Dict[str, Any]: The search results and related information.
+    """
+    print("Starting paper search...")
+    endpoint = "https://api.semanticscholar.org/graph/v1/paper/search"
+    params = {
+        "query": query,
+        "limit": min(limit, 100),
+        # "fields": "paperId,title,abstract,year,authors,
+        # citationCount,url,publicationTypes,openAccessPdf",
+        "fields": "paperId,title,abstract,year,authors,citationCount,url",
+    }
+
+    # Add year parameter if provided
+    if year:
+        params["year"] = year
+
+    response = requests.get(endpoint, params=params, timeout=10)
+    data = response.json()
+    papers = data.get("data", [])
+
+    # Create a dictionary to store the papers
+    filtered_papers = {
+        paper["paperId"]: {
+            "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"),
+            # "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)
+            ],
+        }
+    )