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

aiagents4pharma/talk2scholars/tools/s2/multi_paper_rec.py

@@ -1,22 +1,18 @@
 #!/usr/bin/env python3
 
 """
-multi_paper_rec: Tool for getting recommendations
-                based on multiple papers
+This tool is used to return recommendations based on multiple papers
 """
 
-import json
 import logging
 from typing import Annotated, Any, List, Optional
-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 BaseModel, Field
+from .utils.multi_helper import MultiPaperRecData
 
-# pylint: disable=R0914,R0912,R0915
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -66,143 +62,21 @@ def get_multi_paper_recommendations(
     Returns:
         Dict[str, Any]: The recommendations and related information.
     """
-    # Load hydra configuration
-    with hydra.initialize(version_base=None, config_path="../../configs"):
-        cfg = hydra.compose(
-            config_name="config", overrides=["tools/multi_paper_recommendation=default"]
-        )
-        cfg = cfg.tools.multi_paper_recommendation
-        logger.info("Loaded configuration for multi-paper recommendation tool")
-    logger.info(
-        "Starting multi-paper recommendations search with paper IDs: %s", paper_ids
-    )
-
-    endpoint = cfg.api_endpoint
-    headers = cfg.headers
-    payload = {"positivePaperIds": paper_ids, "negativePaperIds": []}
-    params = {
-        "limit": min(limit, 500),
-        "fields": ",".join(cfg.api_fields),
-    }
-
-    # Add year parameter if provided
-    if year:
-        params["year"] = year
-
-    # Wrap API call in try/except to catch connectivity issues and validate response format
-    response = None
-    for attempt in range(10):
-        try:
-            response = requests.post(
-                endpoint,
-                headers=headers,
-                params=params,
-                data=json.dumps(payload),
-                timeout=cfg.request_timeout,
-            )
-            response.raise_for_status()  # Raises HTTPError for bad responses
-            break  # Exit loop if request is successful
-        except requests.exceptions.RequestException as e:
-            logger.error(
-                "Attempt %d: Failed to connect to Semantic Scholar API for "
-                "multi-paper recommendations: %s",
-                attempt + 1,
-                e,
-            )
-            if attempt == 9:  # Last attempt
-                raise RuntimeError(
-                    "Failed to connect to Semantic Scholar API after 10 attempts."
-                    "Please retry the same query."
-                ) from e
-
-    if response is None:
-        raise RuntimeError("Failed to obtain a response from the Semantic Scholar API.")
-
-    logger.info(
-        "API Response Status for multi-paper recommendations: %s", response.status_code
-    )
-    logger.info("Request params: %s", params)
-
-    data = response.json()
-
-    # Check for expected data format
-    if "recommendedPapers" not in data:
-        logger.error("Unexpected API response format: %s", data)
-        raise RuntimeError(
-            "Unexpected response from Semantic Scholar API. The results could not be "
-            "retrieved due to an unexpected format. "
-            "Please modify your search query and try again."
-        )
+    # Create recommendation data object to organize variables
+    rec_data = MultiPaperRecData(paper_ids, limit, year, tool_call_id)
 
-    recommendations = data.get("recommendedPapers", [])
-    if not recommendations:
-        logger.error(
-            "No recommendations returned from API for paper IDs: %s", paper_ids
-        )
-        raise RuntimeError(
-            "No recommendations were found for your query. Consider refining your search "
-            "by using more specific keywords or different terms."
-        )
-
-    # Create a dictionary to store the papers
-    filtered_papers = {
-        paper["paperId"]: {
-            "semantic_scholar_paper_id": paper["paperId"],
-            "Title": paper.get("title", "N/A"),
-            "Abstract": paper.get("abstract", "N/A"),
-            "Year": paper.get("year", "N/A"),
-            "Publication Date": paper.get("publicationDate", "N/A"),
-            "Venue": paper.get("venue", "N/A"),
-            # "Publication Venue": (paper.get("publicationVenue") or {}).get("name", "N/A"),
-            # "Venue Type": (paper.get("publicationVenue") or {}).get("name", "N/A"),
-            "Journal Name": (paper.get("journal") or {}).get("name", "N/A"),
-            # "Journal Volume": paper.get("journal", {}).get("volume", "N/A"),
-            # "Journal Pages": paper.get("journal", {}).get("pages", "N/A"),
-            "Citation Count": paper.get("citationCount", "N/A"),
-            "Authors": [
-                f"{author.get('name', 'N/A')} (ID: {author.get('authorId', 'N/A')})"
-                for author in paper.get("authors", [])
-            ],
-            "URL": paper.get("url", "N/A"),
-            "arxiv_id": paper.get("externalIds", {}).get("ArXiv", "N/A"),
-        }
-        for paper in recommendations
-        if paper.get("title") and paper.get("authors")
-    }
-
-    # Prepare content with top 3 paper titles and years
-    top_papers = list(filtered_papers.values())[:3]
-    top_papers_info = "\n".join(
-        [
-            # f"{i+1}. {paper['Title']} ({paper['Year']})"
-            f"{i+1}. {paper['Title']} ({paper['Year']}; "
-            f"semantic_scholar_paper_id: {paper['semantic_scholar_paper_id']}; "
-            f"arXiv ID: {paper['arxiv_id']})"
-            for i, paper in enumerate(top_papers)
-        ]
-    )
-
-    logger.info("Filtered %d papers", len(filtered_papers))
-
-    content = (
-        "Recommendations based on multiple papers were successful. "
-        "Papers are attached as an artifact."
-    )
-    content += " Here is a summary of the recommendations:\n"
-    content += f"Number of recommended papers found: {len(filtered_papers)}\n"
-    content += f"Query Paper IDs: {', '.join(paper_ids)}\n"
-    content += f"Year: {year}\n" if year else ""
-    content += "Here are a few of these papers:\n" + top_papers_info
+    # Process the recommendations
+    results = rec_data.process_recommendations()
 
     return Command(
         update={
-            "multi_papers": filtered_papers,  # Sending the dictionary directly
+            "multi_papers": results["papers"],
             "last_displayed_papers": "multi_papers",
             "messages": [
                 ToolMessage(
-                    content=content,
+                    content=results["content"],
                     tool_call_id=tool_call_id,
-                    artifact=filtered_papers,
+                    artifact=results["papers"],
                 )
             ],
         }

aiagents4pharma/talk2scholars/tools/s2/search.py

@@ -6,15 +6,13 @@ This tool is used to search for academic papers on Semantic Scholar.
 
 import logging
 from typing import Annotated, Any, Optional
-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 BaseModel, Field
+from .utils.search_helper import SearchData
 
-# pylint: disable=R0914,R0912,R0915
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -58,125 +56,21 @@ def search_tool(
     Returns:
         The number of papers found on Semantic Scholar.
     """
-    # Load hydra configuration
-    with hydra.initialize(version_base=None, config_path="../../configs"):
-        cfg = hydra.compose(config_name="config", overrides=["tools/search=default"])
-        cfg = cfg.tools.search
-        logger.info("Loaded configuration for search tool")
-    logger.info("Searching for papers on %s", query)
-    endpoint = cfg.api_endpoint
-    params = {
-        "query": query,
-        "limit": min(limit, 100),
-        "fields": ",".join(cfg.api_fields),
-    }
+    # Create search data object to organize variables
+    search_data = SearchData(query, limit, year, tool_call_id)
 
-    # Add year parameter if provided
-    if year:
-        params["year"] = year
-
-    # Wrap API call in try/except to catch connectivity issues
-    response = None
-    for attempt in range(10):
-        try:
-            response = requests.get(endpoint, params=params, timeout=10)
-            response.raise_for_status()  # Raises HTTPError for bad responses
-            break  # Exit loop if request is successful
-        except requests.exceptions.RequestException as e:
-            logger.error(
-                "Attempt %d: Failed to connect to Semantic Scholar API: %s",
-                attempt + 1,
-                e,
-            )
-            if attempt == 9:  # Last attempt
-                raise RuntimeError(
-                    "Failed to connect to Semantic Scholar API after 10 attempts."
-                    "Please retry the same query."
-                ) from e
-
-    if response is None:
-        raise RuntimeError("Failed to obtain a response from the Semantic Scholar API.")
-
-    data = response.json()
-
-    # Check for expected data format
-    if "data" not in data:
-        logger.error("Unexpected API response format: %s", data)
-        raise RuntimeError(
-            "Unexpected response from Semantic Scholar API. The results could not be "
-            "retrieved due to an unexpected format. "
-            "Please modify your search query and try again."
-        )
-
-    papers = data.get("data", [])
-    if not papers:
-        logger.error(
-            "No papers returned from Semantic Scholar API for query: %s", query
-        )
-        raise RuntimeError(
-            "No papers were found for your query. Consider refining your search "
-            "by using more specific keywords or different terms."
-        )
-
-    # Create a dictionary to store the papers
-    filtered_papers = {
-        paper["paperId"]: {
-            "semantic_scholar_paper_id": paper["paperId"],
-            "Title": paper.get("title", "N/A"),
-            "Abstract": paper.get("abstract", "N/A"),
-            "Year": paper.get("year", "N/A"),
-            "Publication Date": paper.get("publicationDate", "N/A"),
-            "Venue": paper.get("venue", "N/A"),
-            # "Publication Venue": (paper.get("publicationVenue") or {}).get("name", "N/A"),
-            # "Venue Type": (paper.get("publicationVenue") or {}).get("name", "N/A"),
-            "Journal Name": (paper.get("journal") or {}).get("name", "N/A"),
-            # "Journal Volume": paper.get("journal", {}).get("volume", "N/A"),
-            # "Journal Pages": paper.get("journal", {}).get("pages", "N/A"),
-            "Citation Count": paper.get("citationCount", "N/A"),
-            "Authors": [
-                f"{author.get('name', 'N/A')} (ID: {author.get('authorId', 'N/A')})"
-                for author in paper.get("authors", [])
-            ],
-            "URL": paper.get("url", "N/A"),
-            "arxiv_id": paper.get("externalIds", {}).get("ArXiv", "N/A"),
-        }
-        for paper in papers
-        if paper.get("title") and paper.get("authors")
-    }
-
-    logger.info("Filtered %d papers", len(filtered_papers))
-
-    # Prepare content with top 3 paper titles and years
-    top_papers = list(filtered_papers.values())[:3]
-    top_papers_info = "\n".join(
-        [
-            f"{i+1}. {paper['Title']} ({paper['Year']}; "
-            f"semantic_scholar_paper_id: {paper['semantic_scholar_paper_id']}; "
-            f"arXiv ID: {paper['arxiv_id']})"
-            for i, paper in enumerate(top_papers)
-        ]
-    )
-
-    logger.info("-----------Filtered %d papers", len(filtered_papers))
-
-    content = (
-        "Search was successful. Papers are attached as an artifact. "
-        "Here is a summary of the search results:\n"
-    )
-    content += f"Number of papers found: {len(filtered_papers)}\n"
-    content += f"Query: {query}\n"
-    content += f"Year: {year}\n" if year else ""
-    content += "Top 3 papers:\n" + top_papers_info
+    # Process the search
+    results = search_data.process_search()
 
     return Command(
         update={
-            "papers": filtered_papers,  # Sending the dictionary directly
+            "papers": results["papers"],
             "last_displayed_papers": "papers",
             "messages": [
                 ToolMessage(
-                    content=content,
+                    content=results["content"],
                     tool_call_id=tool_call_id,
-                    artifact=filtered_papers,
+                    artifact=results["papers"],
                 )
             ],
         }

aiagents4pharma/talk2scholars/tools/s2/single_paper_rec.py

@@ -6,15 +6,13 @@ This tool is used to return recommendations for a single paper.
 
 import logging
 from typing import Annotated, Any, Optional
-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 BaseModel, Field
+from .utils.single_helper import SinglePaperRecData
 
-# pylint: disable=R0914,R0912,R0915
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -62,137 +60,21 @@ def get_single_paper_recommendations(
     Returns:
         Dict[str, Any]: The recommendations and related information.
     """
-    # Load hydra configuration
-    with hydra.initialize(version_base=None, config_path="../../configs"):
-        cfg = hydra.compose(
-            config_name="config",
-            overrides=["tools/single_paper_recommendation=default"],
-        )
-        cfg = cfg.tools.single_paper_recommendation
-        logger.info("Loaded configuration for single paper recommendation tool")
-    logger.info(
-        "Starting single paper recommendations search with paper ID: %s", paper_id
-    )
-
-    endpoint = f"{cfg.api_endpoint}/{paper_id}"
-    params = {
-        "limit": min(limit, 500),  # Max 500 per API docs
-        "fields": ",".join(cfg.api_fields),
-        "from": cfg.recommendation_params.from_pool,
-    }
-
-    # Add year parameter if provided
-    if year:
-        params["year"] = year
-
-    # Wrap API call in try/except to catch connectivity issues and check response format
-    response = None
-    for attempt in range(10):
-        try:
-            response = requests.get(
-                endpoint, params=params, timeout=cfg.request_timeout
-            )
-            response.raise_for_status()  # Raises HTTPError for bad responses
-            break  # Exit loop if request is successful
-        except requests.exceptions.RequestException as e:
-            logger.error(
-                "Attempt %d: Failed to connect to Semantic Scholar API for recommendations: %s",
-                attempt + 1,
-                e,
-            )
-            if attempt == 9:  # Last attempt
-                raise RuntimeError(
-                    "Failed to connect to Semantic Scholar API after 10 attempts."
-                    "Please retry the same query."
-                ) from e
-
-    if response is None:
-        raise RuntimeError("Failed to obtain a response from the Semantic Scholar API.")
-
-    logger.info(
-        "API Response Status for recommendations of paper %s: %s",
-        paper_id,
-        response.status_code,
-    )
-    logger.info("Request params: %s", params)
-
-    data = response.json()
-
-    # Check for expected data format
-    if "recommendedPapers" not in data:
-        logger.error("Unexpected API response format: %s", data)
-        raise RuntimeError(
-            "Unexpected response from Semantic Scholar API. The results could not be "
-            "retrieved due to an unexpected format. "
-            "Please modify your search query and try again."
-        )
+    # Create recommendation data object to organize variables
+    rec_data = SinglePaperRecData(paper_id, limit, year, tool_call_id)
 
-    recommendations = data.get("recommendedPapers", [])
-    if not recommendations:
-        logger.error("No recommendations returned from API for paper: %s", paper_id)
-        raise RuntimeError(
-            "No recommendations were found for your query. Consider refining your search "
-            "by using more specific keywords or different terms."
-        )
-
-    # Extract paper ID and title from recommendations
-    filtered_papers = {
-        paper["paperId"]: {
-            "semantic_scholar_paper_id": paper["paperId"],
-            "Title": paper.get("title", "N/A"),
-            "Abstract": paper.get("abstract", "N/A"),
-            "Year": paper.get("year", "N/A"),
-            "Publication Date": paper.get("publicationDate", "N/A"),
-            "Venue": paper.get("venue", "N/A"),
-            # "Publication Venue": (paper.get("publicationVenue") or {}).get("name", "N/A"),
-            # "Venue Type": (paper.get("publicationVenue") or {}).get("name", "N/A"),
-            "Journal Name": (paper.get("journal") or {}).get("name", "N/A"),
-            # "Journal Volume": paper.get("journal", {}).get("volume", "N/A"),
-            # "Journal Pages": paper.get("journal", {}).get("pages", "N/A"),
-            "Citation Count": paper.get("citationCount", "N/A"),
-            "Authors": [
-                f"{author.get('name', 'N/A')} (ID: {author.get('authorId', 'N/A')})"
-                for author in paper.get("authors", [])
-            ],
-            "URL": paper.get("url", "N/A"),
-            "arxiv_id": paper.get("externalIds", {}).get("ArXiv", "N/A"),
-        }
-        for paper in recommendations
-        if paper.get("title") and paper.get("authors")
-    }
-
-    # Prepare content with top 3 paper titles and years
-    top_papers = list(filtered_papers.values())[:3]
-    top_papers_info = "\n".join(
-        [
-            # f"{i+1}. {paper['Title']} ({paper['Year']})"
-            f"{i+1}. {paper['Title']} ({paper['Year']}; "
-            f"semantic_scholar_paper_id: {paper['semantic_scholar_paper_id']}; "
-            f"arXiv ID: {paper['arxiv_id']})"
-            for i, paper in enumerate(top_papers)
-        ]
-    )
-
-    logger.info("Filtered %d papers", len(filtered_papers))
-
-    content = (
-        "Recommendations based on the single paper were successful. "
-        "Papers are attached as an artifact. "
-        "Here is a summary of the recommendations:\n"
-    )
-    content += f"Number of recommended papers found: {len(filtered_papers)}\n"
-    content += f"Query Paper ID: {paper_id}\n"
-    content += "Here are a few of these papers:\n" + top_papers_info
+    # Process the recommendations
+    results = rec_data.process_recommendations()
 
     return Command(
         update={
-            "papers": filtered_papers,  # Sending the dictionary directly
+            "papers": results["papers"],
             "last_displayed_papers": "papers",
             "messages": [
                 ToolMessage(
-                    content=content,
+                    content=results["content"],
                     tool_call_id=tool_call_id,
-                    artifact=filtered_papers,
+                    artifact=results["papers"],
                 )
             ],
         }

aiagents4pharma/talk2scholars/tools/s2/utils/__init__.py

@@ -0,0 +1,7 @@
+"""This module contains utility functions for the Semantic Scholar search tool."""
+
+from . import search_helper
+from . import single_helper
+from . import multi_helper
+
+__all__ = ["search_helper", "single_helper", "multi_helper"]