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

aiagents4pharma/talk2scholars/tools/s2/search.py

@@ -5,7 +5,7 @@ This tool is used to search for academic papers on Semantic Scholar.
 """
 
 import logging
-from typing import Annotated, Any, Dict, Optional
+from typing import Annotated, Any, Optional
 import hydra
 import requests
 from langchain_core.messages import ToolMessage
@@ -14,7 +14,6 @@ from langchain_core.tools.base import InjectedToolCallId
 from langgraph.types import Command
 from pydantic import BaseModel, Field
 
-
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -28,7 +27,7 @@ class SearchInput(BaseModel):
         "Be specific and include relevant academic terms."
     )
     limit: int = Field(
-        default=2, description="Maximum number of results to return", ge=1, le=100
+        default=5, description="Maximum number of results to return", ge=1, le=100
     )
     year: Optional[str] = Field(
         default=None,
@@ -44,13 +43,13 @@ with hydra.initialize(version_base=None, config_path="../../configs"):
     cfg = cfg.tools.search
 
 
-@tool(args_schema=SearchInput)
+@tool("search_tool", args_schema=SearchInput, parse_docstring=True)
 def search_tool(
     query: str,
     tool_call_id: Annotated[str, InjectedToolCallId],
-    limit: int = 2,
+    limit: int = 5,
     year: Optional[str] = None,
-) -> Dict[str, Any]:
+) -> Command[Any]:
     """
     Search for academic papers on Semantic Scholar.
 
@@ -62,9 +61,9 @@ def search_tool(
         Supports formats like "2024-", "-2024", "2024:2025". Defaults to None.
 
     Returns:
-        Dict[str, Any]: The search results and related information.
+        The number of papers found on Semantic Scholar.
     """
-    print("Starting paper search...")
+    logger.info("Searching for papers on %s", query)
     endpoint = cfg.api_endpoint
     params = {
         "query": query,
@@ -80,26 +79,51 @@ def search_tool(
     data = response.json()
     papers = data.get("data", [])
     logger.info("Received %d papers", len(papers))
+    if not papers:
+        return Command(
+            update={  # Place 'messages' inside 'update'
+                "messages": [
+                    ToolMessage(
+                        content="No papers found. Please try a different search query.",
+                        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 papers
         if paper.get("title") and paper.get("authors")
     }
 
+    logger.info("Filtered %d papers", len(filtered_papers))
+
+    content = "Search was successful."
+    content += " Here is a summary of the search results:"
+    content += f"Number of papers found: {len(filtered_papers)}\n"
+    content += f"Query: {query}\n"
+    content += f"Year: {year}\n" if year else ""
+
     return Command(
         update={
             "papers": filtered_papers,  # Now sending the dictionary directly
+            "last_displayed_papers": "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/single_paper_rec.py

@@ -5,7 +5,7 @@ This tool is used to return recommendations for a single paper.
 """
 
 import logging
-from typing import Annotated, Any, Dict, Optional
+from typing import Annotated, Any, Optional
 import hydra
 import requests
 from langchain_core.messages import ToolMessage
@@ -26,7 +26,7 @@ class SinglePaperRecInput(BaseModel):
         description="Semantic Scholar Paper ID to get recommendations for (40-character string)"
     )
     limit: int = Field(
-        default=2,
+        default=5,
         description="Maximum number of recommendations to return",
         ge=1,
         le=500,
@@ -48,15 +48,16 @@ with hydra.initialize(version_base=None, config_path="../../configs"):
     cfg = cfg.tools.single_paper_recommendation
 
 
-@tool(args_schema=SinglePaperRecInput)
+@tool(args_schema=SinglePaperRecInput, parse_docstring=True)
 def get_single_paper_recommendations(
     paper_id: str,
     tool_call_id: Annotated[str, InjectedToolCallId],
-    limit: int = 2,
+    limit: int = 5,
     year: Optional[str] = None,
-) -> Dict[str, Any]:
+) -> Command[Any]:
     """
-    Get paper recommendations based on a single paper.
+    Get recommendations for on a single paper using its Semantic Scholar ID.
+    No other ID types are supported.
 
     Args:
         paper_id (str): The Semantic Scholar Paper ID to get recommendations for.
@@ -68,7 +69,9 @@ def get_single_paper_recommendations(
     Returns:
         Dict[str, Any]: The recommendations and related information.
     """
-    logger.info("Starting single paper recommendations search.")
+    logger.info(
+        "Starting single paper recommendations search with paper ID: %s", paper_id
+    )
 
     endpoint = f"{cfg.api_endpoint}/{paper_id}"
     params = {
@@ -90,32 +93,59 @@ def get_single_paper_recommendations(
         paper_id,
         response.status_code,
     )
+    if response.status_code != 200:
+        raise ValueError("Invalid paper ID or API error.")
     # print(f"Request params: {params}")
     logging.info("Request params: %s", params)
 
     data = response.json()
     recommendations = data.get("recommendedPapers", [])
 
+    if not recommendations:
+        return Command(
+            update={
+                "papers": {},
+                "messages": [
+                    ToolMessage(
+                        content=f"No recommendations found for {paper_id}.",
+                        tool_call_id=tool_call_id,
+                    )
+                ],
+            }
+        )
+
     # Extract paper ID and title from recommendations
     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("authors")
     }
 
+    content = "Recommendations based on a single paper were successful."
+    content += " Here is a summary of the recommendations:"
+    content += f"Number of papers found: {len(filtered_papers)}\n"
+    content += f"Query Paper ID: {paper_id}\n"
+    content += f"Year: {year}\n" if year else ""
+
     return Command(
         update={
             "papers": filtered_papers,  # Now sending the dictionary directly
+            "last_displayed_papers": "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-1.20.0.dist-info → aiagents4pharma-1.21.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: aiagents4pharma
-Version: 1.20.0
+Version: 1.21.0
 Summary: AI Agents for drug discovery, drug development, and other pharmaceutical R&D.
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
@@ -110,7 +110,7 @@ free credits [here](https://build.nvidia.com/explore/discover)._
    ```
 2. **Install dependencies:**
    ```bash
-   pip install .
+   pip install -r requirements.txt
    ```
 3. **Initialize OPENAI_API_KEY and NVIDIA_API_KEY**
    ```bash

{aiagents4pharma-1.20.0.dist-info → aiagents4pharma-1.21.0.dist-info}/RECORD

@@ -7,6 +7,7 @@ aiagents4pharma/talk2biomodels/api/kegg.py,sha256=QzYDAfJ16E7tbHGxP8ZNWRizMkMRS_
 aiagents4pharma/talk2biomodels/api/ols.py,sha256=qq0Qy-gJDxanQW-HfCChDsTQsY1M41ua8hMlTnfuzrA,2202
 aiagents4pharma/talk2biomodels/api/uniprot.py,sha256=aPUAVBR7UYXDuuhDpKezAK2aTMzo-NxFYFq6C0W5u6U,1175
 aiagents4pharma/talk2biomodels/configs/__init__.py,sha256=safyFKhkd5Wlirl9dMZIHWDLTpY2oLw9wjIM7ZtLIHk,88
+aiagents4pharma/talk2biomodels/configs/config.yaml,sha256=X0CMsnx6hHNvV04wsENQSGXadx0aKIy6mziSopVUdZI,116
 aiagents4pharma/talk2biomodels/configs/agents/__init__.py,sha256=_ZoG8snICK2bidWtc2KOGs738LWg9_r66V9mOMnEb-E,71
 aiagents4pharma/talk2biomodels/configs/agents/t2b_agent/__init__.py,sha256=-fAORvyFmG2iSvFOFDixmt9OTQRR58y89uhhu2EgbA8,46
 aiagents4pharma/talk2biomodels/configs/agents/t2b_agent/default.yaml,sha256=pSViMKwKyMQDm8LzbfIaGdxph73iHYaXMiv5YOuxM7k,536
@@ -114,42 +115,45 @@ aiagents4pharma/talk2knowledgegraphs/utils/extractions/__init__.py,sha256=7gwwtf
 aiagents4pharma/talk2knowledgegraphs/utils/extractions/pcst.py,sha256=m5p0yoJb7I19ua5yeQfXPf7c4r6S1XPwttsrM7Qoy94,9336
 aiagents4pharma/talk2scholars/__init__.py,sha256=gphERyVKZHvOnMQsml7TIHlaIshHJ75R1J3FKExkfuY,120
 aiagents4pharma/talk2scholars/agents/__init__.py,sha256=ykszlVGxz3egLHZAttlNoTPxIrnQJZYva_ssR8fwIFk,117
-aiagents4pharma/talk2scholars/agents/main_agent.py,sha256=_spqbC2lS-pJ8lq59S7ViVT_aQK3jlA5vqlOx1Q7JOc,7177
-aiagents4pharma/talk2scholars/agents/s2_agent.py,sha256=QBsar2QrTgi9H5DniJzvbEUH3vKb_dlYojEcfsBXhPM,2765
+aiagents4pharma/talk2scholars/agents/main_agent.py,sha256=UiiqCMUko4t39uR4NtTvG3kRMf89x67dU_3NSAM4DWw,7603
+aiagents4pharma/talk2scholars/agents/s2_agent.py,sha256=XilKQvlxEVNYK5cEc44x-62OZZ1qG77v1r0FkwLexcw,4581
 aiagents4pharma/talk2scholars/configs/__init__.py,sha256=tf2gz8n7M4ko6xLdX_C925ELVIxoP6SgkPcbeh59ad4,151
-aiagents4pharma/talk2scholars/configs/config.yaml,sha256=a3_jCFAmsVL6gZuvzoRe4jL94mQaSbp0CUXZDUtqhZA,254
+aiagents4pharma/talk2scholars/configs/config.yaml,sha256=XaUi1aP0VGiZFeW_ZP68bVSiCJezied5yLzcx3Uljhc,308
 aiagents4pharma/talk2scholars/configs/agents/__init__.py,sha256=yyh7PB2oY_JulnpSQCWS4wwCH_uzIdt47O2Ay48x_oU,75
 aiagents4pharma/talk2scholars/configs/agents/talk2scholars/__init__.py,sha256=Tj4urOkjpu2cTlpJl0Fmr_18RZCR88vns-Gt-XquDzs,95
 aiagents4pharma/talk2scholars/configs/agents/talk2scholars/main_agent/__init__.py,sha256=fqQQ-GlRcbzru2KmEk3oMma0R6_SzGM8dOXzYeU4oVA,46
-aiagents4pharma/talk2scholars/configs/agents/talk2scholars/main_agent/default.yaml,sha256=Z85fv6orEpvdn-JGEmx0-45fsTJrlTRnqUYAhp1gZpY,679
+aiagents4pharma/talk2scholars/configs/agents/talk2scholars/main_agent/default.yaml,sha256=kizyUjDasjv3zt2xHcGMqTxnpR_FHAcOs1vgVaH7tsY,1882
 aiagents4pharma/talk2scholars/configs/agents/talk2scholars/s2_agent/__init__.py,sha256=fqQQ-GlRcbzru2KmEk3oMma0R6_SzGM8dOXzYeU4oVA,46
-aiagents4pharma/talk2scholars/configs/agents/talk2scholars/s2_agent/default.yaml,sha256=OCQ6wEh48qX73m49WJj2RHaws4kkKRUsGt6T2thDAzo,766
+aiagents4pharma/talk2scholars/configs/agents/talk2scholars/s2_agent/default.yaml,sha256=lQUvQDLqqPIxn4TkL58atG2dQFhysNuNyAbd6P8km3g,585
 aiagents4pharma/talk2scholars/configs/app/__init__.py,sha256=JoSZV6N669kGMv5zLDszwf0ZjcRHx9TJfIqGhIIdPXE,70
 aiagents4pharma/talk2scholars/configs/app/frontend/__init__.py,sha256=fqQQ-GlRcbzru2KmEk3oMma0R6_SzGM8dOXzYeU4oVA,46
-aiagents4pharma/talk2scholars/configs/app/frontend/default.yaml,sha256=BX-J1zQyb0QJ7hcOFOnkJ8aWoWbjK4WE2VG7OZTOyKU,821
+aiagents4pharma/talk2scholars/configs/app/frontend/default.yaml,sha256=wsELBdRLv6UqZ9QZfwpS7K4xfMj5s-a99-aXqIs6WEI,868
 aiagents4pharma/talk2scholars/configs/tools/__init__.py,sha256=w0BJK0MR6Et8Pw1htP8JV0Lr9F_N68CqvbpV14KBy_8,151
 aiagents4pharma/talk2scholars/configs/tools/multi_paper_recommendation/__init__.py,sha256=fqQQ-GlRcbzru2KmEk3oMma0R6_SzGM8dOXzYeU4oVA,46
-aiagents4pharma/talk2scholars/configs/tools/multi_paper_recommendation/default.yaml,sha256=70HSJ8WbS2Qbhur5FpuOPBjrea9g3TioM0gjGn6U1bE,369
+aiagents4pharma/talk2scholars/configs/tools/multi_paper_recommendation/default.yaml,sha256=iEsEW89MlQwKsAW4ZAxLt4pDBwA1qxImYQ2dfONIf6c,442
+aiagents4pharma/talk2scholars/configs/tools/retrieve_semantic_scholar_paper_id/__init__.py,sha256=fqQQ-GlRcbzru2KmEk3oMma0R6_SzGM8dOXzYeU4oVA,46
 aiagents4pharma/talk2scholars/configs/tools/search/__init__.py,sha256=fqQQ-GlRcbzru2KmEk3oMma0R6_SzGM8dOXzYeU4oVA,46
-aiagents4pharma/talk2scholars/configs/tools/search/default.yaml,sha256=NznRVqB6EamMfsFc5hj5s9ygzl6rPuFPiy9ikcpqp68,486
+aiagents4pharma/talk2scholars/configs/tools/search/default.yaml,sha256=tw8N1Mms0qHQbIY3KGDNK1NuT19dQGPiagxzWDdOAJk,504
 aiagents4pharma/talk2scholars/configs/tools/single_paper_recommendation/__init__.py,sha256=fqQQ-GlRcbzru2KmEk3oMma0R6_SzGM8dOXzYeU4oVA,46
-aiagents4pharma/talk2scholars/configs/tools/single_paper_recommendation/default.yaml,sha256=4_gTvdVc-hf9GNxBKMGQd72s5h53Zy09j9qeZ9Fys04,578
+aiagents4pharma/talk2scholars/configs/tools/single_paper_recommendation/default.yaml,sha256=TILecrowsu5VGdJPeac6fl5AXSf3piSHN0oKdjY2q1o,596
 aiagents4pharma/talk2scholars/state/__init__.py,sha256=S6SxlszIMZSIMJehjevPF9sKyR-PAwWb5TEdo6xWXE8,103
-aiagents4pharma/talk2scholars/state/state_talk2scholars.py,sha256=6LwGCP1yluxcAf1PyWl7LaXEvxbaJ2-b2Nrxu4Z_KI0,978
+aiagents4pharma/talk2scholars/state/state_talk2scholars.py,sha256=Jbwckhhpv03KFQYkYkuPsYz5OxT0Si-mra7qrynH5Jo,2246
 aiagents4pharma/talk2scholars/tests/__init__.py,sha256=U3PsTiUZaUBD1IZanFGkDIOdFieDVJtGKQ5-woYUo8c,45
-aiagents4pharma/talk2scholars/tests/test_integration.py,sha256=JWRoXkYjB6mKjiGvzoqSjAFXX_L8s2dP67MGS9dk6Rs,7721
-aiagents4pharma/talk2scholars/tests/test_main_agent.py,sha256=_XPnyveZoBOgrrV_2ws8CLWAX338xBlyh1Zs6xoZKpk,6171
-aiagents4pharma/talk2scholars/tests/test_s2_agent.py,sha256=0VtCZ0C_4WyF5t7i8Cmv7P1FakmDTe6YfXNriFAiGEQ,4645
-aiagents4pharma/talk2scholars/tests/test_s2_tools.py,sha256=OU2ArjF07J1jKz4dLnhQ1RxNXgNBhgMqglxLIvKcXuk,7496
+aiagents4pharma/talk2scholars/tests/test_llm_main_integration.py,sha256=SAMG-Kb2S9sei8Us5vUWCUJikTKXPZVKQ6aJJPEhJsc,1880
+aiagents4pharma/talk2scholars/tests/test_main_agent.py,sha256=TTPfVGWWq6BXJVgfR958qttD6dGRnpJHZMqo86k4aMo,5562
+aiagents4pharma/talk2scholars/tests/test_s2_agent.py,sha256=-yEoG2v5SMkCLCrSA2DFcNE-xMOSn97N4UTomzCeW40,7559
+aiagents4pharma/talk2scholars/tests/test_s2_tools.py,sha256=QEwraJk9_Kp6ZSGYyYDXWH62wIjSwi1Pptwwbx1fuG0,13176
 aiagents4pharma/talk2scholars/tests/test_state.py,sha256=_iHXvoZnU_eruf8l1sQKBSCIVnxNkH_9VzkVtZZA6bY,384
 aiagents4pharma/talk2scholars/tools/__init__.py,sha256=YudBDRwaEzDnAcpxGZvEOfyh5-6xd51CTvTKTkywgXw,68
-aiagents4pharma/talk2scholars/tools/s2/__init__.py,sha256=9RQH3efTj6qkXk0ICKSc7Mzpkitt4gRGsQ1pGPrrREU,181
-aiagents4pharma/talk2scholars/tools/s2/display_results.py,sha256=Aap3P-9i80TDarkEZX0S62rzsVZw0ftzgUppK67rOhI,1446
-aiagents4pharma/talk2scholars/tools/s2/multi_paper_rec.py,sha256=7GrMaZ15zLSmBvmek9NCATM4-O6BXuxa5Pwd9qJ6ygU,3887
-aiagents4pharma/talk2scholars/tools/s2/search.py,sha256=fF8i2s8ikF0qdM_c7KSAQvguG9Hjq0s2bPvoyLdx5wA,3311
-aiagents4pharma/talk2scholars/tools/s2/single_paper_rec.py,sha256=Z3NQgf3hHcPSisYJnbEErsyGFZPgad8CdaWe4XxiOgs,3978
-aiagents4pharma-1.20.0.dist-info/LICENSE,sha256=IcIbyB1Hyk5ZDah03VNQvJkbNk2hkBCDqQ8qtnCvB4Q,1077
-aiagents4pharma-1.20.0.dist-info/METADATA,sha256=GZWGNk_tXc9mtrsvC5ubW6zbFUVPzIRDNNS2TPnMQBQ,7739
-aiagents4pharma-1.20.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-aiagents4pharma-1.20.0.dist-info/top_level.txt,sha256=-AH8rMmrSnJtq7HaAObS78UU-cTCwvX660dSxeM7a0A,16
-aiagents4pharma-1.20.0.dist-info/RECORD,,
+aiagents4pharma/talk2scholars/tools/s2/__init__.py,sha256=wytqCmGm8Fbl8y5qLdIkxhhG8VHLYMifCGjbH_LK2Fc,258
+aiagents4pharma/talk2scholars/tools/s2/display_results.py,sha256=UR0PtEHGDpOhPH0Di5HT8-Fip2RkEMTJgzROsChb1gc,2959
+aiagents4pharma/talk2scholars/tools/s2/multi_paper_rec.py,sha256=sYtoJWFmJuSTSHMrZmqiBfGS-mDKS1gAbKtNyjRnlwU,4979
+aiagents4pharma/talk2scholars/tools/s2/query_results.py,sha256=EUfzRh5Qc_tMl5fDIFb9PIsQkkrU4Xb5MR0sud_X5-c,2017
+aiagents4pharma/talk2scholars/tools/s2/retrieve_semantic_scholar_paper_id.py,sha256=Lg1L4HQCN2LaQEyWtLD73O67PMoXkPHi-Y8rCzHS0A4,2499
+aiagents4pharma/talk2scholars/tools/s2/search.py,sha256=mnBQWDuQ50UVw6B-bRuL8Ek1av-pEtdgzVMxpEA2BpI,4296
+aiagents4pharma/talk2scholars/tools/s2/single_paper_rec.py,sha256=xgnUj9W9JkeTvB2VJBJUAnia789GGNGqdqgJ_G16v2s,5120
+aiagents4pharma-1.21.0.dist-info/LICENSE,sha256=IcIbyB1Hyk5ZDah03VNQvJkbNk2hkBCDqQ8qtnCvB4Q,1077
+aiagents4pharma-1.21.0.dist-info/METADATA,sha256=YsjDHw3yfqfPClv0N3j35AObxHBhStDojFUUslyd_1Q,7757
+aiagents4pharma-1.21.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+aiagents4pharma-1.21.0.dist-info/top_level.txt,sha256=-AH8rMmrSnJtq7HaAObS78UU-cTCwvX660dSxeM7a0A,16
+aiagents4pharma-1.21.0.dist-info/RECORD,,

aiagents4pharma/talk2scholars/tests/test_integration.py

@@ -1,237 +0,0 @@
-"""
-Integration tests for talk2scholars system.
-
-These tests ensure that:
-1. The main agent and sub-agent work together.
-2. The agents correctly interact with tools (search, recommendations).
-3. The full pipeline processes queries and updates state correctly.
-"""
-
-# pylint: disable=redefined-outer-name
-from unittest.mock import patch, Mock
-import pytest
-from langchain_core.messages import HumanMessage
-from ..agents.main_agent import get_app as get_main_app
-from ..agents.s2_agent import get_app as get_s2_app
-from ..state.state_talk2scholars import Talk2Scholars
-
-
-@pytest.fixture(autouse=True)
-def mock_hydra():
-    """Mock Hydra configuration to prevent external dependencies."""
-    with patch("hydra.initialize"), patch("hydra.compose") as mock_compose:
-        cfg_mock = Mock()
-        cfg_mock.agents.talk2scholars.main_agent.temperature = 0
-        cfg_mock.agents.talk2scholars.main_agent.main_agent = "Test main agent prompt"
-        cfg_mock.agents.talk2scholars.s2_agent.temperature = 0
-        cfg_mock.agents.talk2scholars.s2_agent.s2_agent = "Test s2 agent prompt"
-        mock_compose.return_value = cfg_mock
-        yield mock_compose
-
-
-@pytest.fixture(autouse=True)
-def mock_tools():
-    """Mock tools to prevent execution of real API calls."""
-    with (
-        patch(
-            "aiagents4pharma.talk2scholars.tools.s2.search.search_tool"
-        ) as mock_s2_search,
-        patch(
-            "aiagents4pharma.talk2scholars.tools.s2.display_results.display_results"
-        ) as mock_s2_display,
-        patch(
-            "aiagents4pharma.talk2scholars.tools.s2.single_paper_rec."
-            "get_single_paper_recommendations"
-        ) as mock_s2_single_rec,
-        patch(
-            "aiagents4pharma.talk2scholars.tools.s2.multi_paper_rec."
-            "get_multi_paper_recommendations"
-        ) as mock_s2_multi_rec,
-    ):
-
-        mock_s2_search.return_value = {"papers": {"id123": "Mock Paper"}}
-        mock_s2_display.return_value = "Displaying Mock Results"
-        mock_s2_single_rec.return_value = {"recommendations": ["Paper A", "Paper B"]}
-        mock_s2_multi_rec.return_value = {
-            "multi_recommendations": ["Paper X", "Paper Y"]
-        }
-
-        yield {
-            "search_tool": mock_s2_search,
-            "display_results": mock_s2_display,
-            "single_paper_rec": mock_s2_single_rec,
-            "multi_paper_rec": mock_s2_multi_rec,
-        }
-
-
-def test_full_workflow():
-    """Test the full workflow from main agent to S2 agent."""
-    thread_id = "test_thread"
-    main_app = get_main_app(thread_id)
-
-    # Define expected mock response with the actual structure
-    expected_paper = {
-        "530a059cb48477ad1e3d4f8f4b153274c8997332": {
-            "Title": "Explainable Artificial Intelligence",
-            "Abstract": None,
-            "Citation Count": 5544,
-            "Year": "2024",
-            "URL": "https://example.com/paper",
-        }
-    }
-
-    # Mock the search tool instead of the app
-    with patch(
-        "aiagents4pharma.talk2scholars.tools.s2.search.search_tool",
-        return_value={"papers": expected_paper},
-    ):
-        state = Talk2Scholars(messages=[HumanMessage(content="Find AI papers")])
-        result = main_app.invoke(
-            state,
-            config={
-                "configurable": {
-                    "thread_id": thread_id,
-                    "checkpoint_ns": "test_ns",
-                    "checkpoint_id": "test_checkpoint",
-                }
-            },
-        )
-
-    # Check values
-    assert "papers" in result
-    assert "messages" in result
-    assert len(result["papers"]) > 0
-
-
-def test_s2_agent_execution():
-    """Test if the S2 agent processes requests correctly and updates state."""
-    thread_id = "test_thread"
-    s2_app = get_s2_app(thread_id)
-
-    state = Talk2Scholars(messages=[HumanMessage(content="Get recommendations")])
-
-    result = s2_app.invoke(
-        state,
-        config={
-            "configurable": {
-                "thread_id": thread_id,
-                "checkpoint_ns": "test_ns",
-                "checkpoint_id": "test_checkpoint",
-            }
-        },
-    )
-
-    assert "messages" in result
-    assert "multi_papers" in result
-    assert result["multi_papers"] is not None
-
-
-def test_tool_integration(mock_tools):
-    """Test if the tools interact correctly with the workflow."""
-    thread_id = "test_thread"
-    s2_app = get_s2_app(thread_id)
-
-    state = Talk2Scholars(
-        messages=[HumanMessage(content="Search for AI ethics papers")]
-    )
-
-    mock_paper_id = "11159bdb213aaa243916f42f576396d483ba474b"
-    mock_response = {
-        "papers": {
-            mock_paper_id: {
-                "Title": "Mock AI Ethics Paper",
-                "Abstract": "A study on AI ethics",
-                "Citation Count": 100,
-                "URL": "https://example.com/mock-paper",
-            }
-        }
-    }
-
-    # Update both the fixture mock and patch the actual tool
-    mock_tools["search_tool"].return_value = {"papers": mock_response["papers"]}
-
-    with patch(
-        "aiagents4pharma.talk2scholars.tools.s2.search.search_tool",
-        return_value={"papers": mock_response["papers"]},
-    ):
-        result = s2_app.invoke(
-            state,
-            config={
-                "configurable": {
-                    "thread_id": thread_id,
-                    "checkpoint_ns": "test_ns",
-                    "checkpoint_id": "test_checkpoint",
-                }
-            },
-        )
-
-    assert "papers" in result
-    assert len(result["papers"]) > 0  # Verify we have papers
-    assert isinstance(result["papers"], dict)  # Verify it's a dictionary
-
-
-def test_empty_query():
-    """Test how the system handles an empty query."""
-    thread_id = "test_thread"
-    main_app = get_main_app(thread_id)
-
-    state = Talk2Scholars(messages=[HumanMessage(content="")])
-
-    # Mock the s2_agent app
-    mock_s2_app = get_s2_app(thread_id)
-
-    with patch(
-        "aiagents4pharma.talk2scholars.agents.s2_agent.get_app",
-        return_value=mock_s2_app,
-    ):
-        result = main_app.invoke(
-            state,
-            config={
-                "configurable": {
-                    "thread_id": thread_id,
-                    "checkpoint_ns": "test_ns",
-                    "checkpoint_id": "test_checkpoint",
-                }
-            },
-        )
-
-    assert "messages" in result
-    last_message = result["messages"][-1].content.lower()
-    assert any(
-        phrase in last_message
-        for phrase in ["no valid input", "how can i assist", "please provide a query"]
-    )
-
-
-def test_api_failure_handling():
-    """Test if the system gracefully handles an API failure."""
-    thread_id = "test_thread"
-    s2_app = get_s2_app(thread_id)
-
-    expected_error = "API Timeout: Connection failed"
-    with patch("requests.get", side_effect=Exception(expected_error)):
-        state = Talk2Scholars(messages=[HumanMessage(content="Find latest NLP papers")])
-
-        result = s2_app.invoke(
-            state,
-            config={
-                "configurable": {
-                    "thread_id": thread_id,
-                    "checkpoint_ns": "test_ns",
-                    "checkpoint_id": "test_checkpoint",
-                }
-            },
-        )
-
-        assert "messages" in result
-        last_message = result["messages"][-1].content.lower()
-
-        # Update assertions to match actual error message
-        assert any(
-            [
-                "unable to retrieve" in last_message,
-                "connection issue" in last_message,
-                "please try again later" in last_message,
-            ]
-        )
-        assert "nlp papers" in last_message  # Verify context is maintained