aiagents4pharma 1.27.2__py3-none-any.whl → 1.29.0__py3-none-any.whl

aiagents4pharma/talk2scholars/tools/pdf/question_and_answer.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python3
+"""
+question_and_answer: Tool for performing Q&A on PDF documents using retrieval augmented generation.
+
+This module provides functionality to extract text from PDF binary data, split it into 
+chunks, retrieve relevant segments via a vector store, and generate an answer to a 
+user-provided question using a language model chain.
+"""
+
+import io
+import logging
+from typing import Annotated, Dict, Any, List
+
+from PyPDF2 import PdfReader
+from pydantic import BaseModel, Field
+import hydra
+
+from langchain.chains.question_answering import load_qa_chain
+from langchain.docstore.document import Document
+from langchain.text_splitter import CharacterTextSplitter
+from langchain_community.vectorstores import Annoy
+from langchain_openai import OpenAIEmbeddings
+from langchain_core.language_models.chat_models import BaseChatModel
+
+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 langgraph.prebuilt import InjectedState
+
+# Set up logging.
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+# Load configuration using Hydra.
+with hydra.initialize(version_base=None, config_path="../../configs"):
+    cfg = hydra.compose(config_name="config", overrides=["tools/question_and_answer=default"])
+    cfg = cfg.tools.question_and_answer
+    logger.info("Loaded Question and Answer tool configuration.")
+
+class QuestionAndAnswerInput(BaseModel):
+    """
+    Input schema for the PDF Question and Answer tool.
+
+    Attributes:
+        question (str): The question to ask regarding the PDF content.
+        tool_call_id (str): Unique identifier for the tool call, injected automatically.
+    """
+    question: str = Field(
+        description="The question to ask regarding the PDF content."
+    )
+    tool_call_id: Annotated[str, InjectedToolCallId]
+    state: Annotated[dict, InjectedState]
+
+def extract_text_from_pdf_data(pdf_bytes: bytes) -> str:
+    """
+    Extract text content from PDF binary data.
+
+    This function uses PyPDF2 to read the provided PDF bytes and concatenates the text
+    extracted from each page.
+
+    Args:
+        pdf_bytes (bytes): The binary data of the PDF document.
+
+    Returns:
+        str: The complete text extracted from the PDF.
+    """
+    reader = PdfReader(io.BytesIO(pdf_bytes))
+    text = ""
+    for page in reader.pages:
+        page_text = page.extract_text() or ""
+        text += page_text
+    return text
+
+def generate_answer(question: str, pdf_bytes: bytes, llm_model: BaseChatModel) -> Dict[str, Any]:
+    """
+    Generate an answer for a question using retrieval augmented generation on PDF content.
+
+    This function extracts text from the PDF data, splits the text into manageable chunks,
+    performs a similarity search to retrieve the most relevant segments, and then uses a
+    question-answering chain (built using the provided llm_model) to generate an answer.
+
+    Args:
+        question (str): The question to be answered.
+        pdf_bytes (bytes): The binary content of the PDF document.
+        llm_model (BaseChatModel): The language model instance to use for answering.
+
+    Returns:
+        Dict[str, Any]: A dictionary containing the answer generated by the language model.
+    """
+    text = extract_text_from_pdf_data(pdf_bytes)
+    logger.info("Extracted text from PDF.")
+    text_splitter = CharacterTextSplitter(
+        separator="\n",
+        chunk_size=cfg.chunk_size,
+        chunk_overlap=cfg.chunk_overlap
+    )
+    chunks = text_splitter.split_text(text)
+    documents: List[Document] = [Document(page_content=chunk) for chunk in chunks]
+    logger.info("Split PDF text into %d chunks.", len(documents))
+
+    embeddings = OpenAIEmbeddings(openai_api_key=cfg.openai_api_key)
+    vector_store = Annoy.from_documents(documents, embeddings)
+    search_results = vector_store.similarity_search(
+        question,
+        k=cfg.num_retrievals
+    )
+    logger.info("Retrieved %d relevant document chunks.", len(search_results))
+    # Use the provided llm_model to build the QA chain.
+    qa_chain = load_qa_chain(llm_model, chain_type=cfg.qa_chain_type)
+    answer = qa_chain.invoke(
+        input={"input_documents": search_results, "question": question}
+    )
+    return answer
+
+@tool(args_schema=QuestionAndAnswerInput)
+def question_and_answer_tool(
+    question: str,
+    tool_call_id: Annotated[str, InjectedToolCallId],
+    state: Annotated[dict, InjectedState],
+) -> Dict[str, Any]:
+    """
+    Answer a question using PDF content stored in the state via retrieval augmented generation.
+
+    This tool retrieves the PDF binary data from the state (under the key "pdf_data"), extracts its
+    textual content, and generates an answer to the specified question. It also extracts the 
+    llm_model (of type BaseChatModel) from the state to use for answering.
+
+    Args:
+        question (str): The question regarding the PDF content.
+        tool_call_id (str): Unique identifier for the current tool call.
+        state (dict): A dictionary representing the current state, expected to contain PDF data
+                      under the key "pdf_data" with a sub-key "pdf_object" for the binary content,
+                      and a key "llm_model" holding the language model instance.
+
+    Returns:
+        Dict[str, Any]: A dictionary containing the generated answer or an error message.
+    """
+    logger.info("Starting PDF Question and Answer tool using PDF data from state.")
+    pdf_state = state.get("pdf_data")
+    if not pdf_state:
+        error_msg = "No pdf_data found in state."
+        logger.error(error_msg)
+        return Command(
+            update={
+                "messages": [
+                    ToolMessage(content=error_msg, tool_call_id=tool_call_id)
+                ]
+            }
+        )
+    pdf_bytes = pdf_state.get("pdf_object")
+    if not pdf_bytes:
+        error_msg = "PDF binary data is missing in the pdf_data from state."
+        logger.error(error_msg)
+        return Command(
+            update={
+                "messages": [
+                    ToolMessage(content=error_msg, tool_call_id=tool_call_id)
+                ]
+            }
+        )
+    # Retrieve llm_model from state; use a default if not provided.
+    llm_model = state.get("llm_model")
+    if not llm_model:
+        logger.error("Missing LLM model instance in state.")
+        return {"error": "No LLM model found in state."}
+    answer = generate_answer(question, pdf_bytes, llm_model)
+    logger.info("Generated answer: %s", answer)
+    return answer

aiagents4pharma/talk2scholars/tools/s2/multi_paper_rec.py

@@ -16,6 +16,7 @@ from langchain_core.tools.base import InjectedToolCallId
 from langgraph.types import Command
 from pydantic import BaseModel, Field
 
+# pylint: disable=R0914,R0912,R0915
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -26,7 +27,7 @@ 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")
+        description="List of Semantic Scholar Paper IDs to get recommendations for"
     )
     limit: int = Field(
         default=2,
@@ -44,14 +45,6 @@ class MultiPaperRecInput(BaseModel):
     model_config = {"arbitrary_types_allowed": True}
 
 
-# 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
-
-
 @tool(args_schema=MultiPaperRecInput, parse_docstring=True)
 def get_multi_paper_recommendations(
     paper_ids: List[str],
@@ -73,7 +66,14 @@ def get_multi_paper_recommendations(
     Returns:
         Dict[str, Any]: The recommendations and related information.
     """
-    logging.info(
+    # 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
     )
 
@@ -89,45 +89,61 @@ def get_multi_paper_recommendations(
     if year:
         params["year"] = year
 
-    # Getting recommendations
-    response = requests.post(
-        endpoint,
-        headers=headers,
-        params=params,
-        data=json.dumps(payload),
-        timeout=cfg.request_timeout,
-    )
-    logging.info(
+    # Wrap API call in try/except to catch connectivity issues and validate response format
+    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
+    except requests.exceptions.RequestException as e:
+        logger.error(
+            "Failed to connect to Semantic Scholar API for multi-paper recommendations: %s",
+            e,
+        )
+        raise RuntimeError(
+            "Failed to connect to Semantic Scholar API. Please retry the same query."
+        ) from e
+
+    logger.info(
         "API Response Status for multi-paper recommendations: %s", response.status_code
     )
+    logger.info("Request params: %s", params)
 
     data = response.json()
-    recommendations = data.get("recommendedPapers", [])
 
+    # 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."
+        )
+
+    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,
-                    )
-                ]
-            }
+        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_id": paper["paperId"],  # Store Semantic Scholar ID
+            "paper_id": 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"),
-            # "arXiv_ID": paper.get("externalIds", {}).get(
-            #     "ArXiv", "N/A"
-            # ),  # Extract arXiv ID
+            "arxiv_id": paper.get("externalIds", {}).get("ArXiv", "N/A"),
         }
         for paper in recommendations
         if paper.get("title") and paper.get("authors")
@@ -156,7 +172,7 @@ def get_multi_paper_recommendations(
 
     return Command(
         update={
-            "multi_papers": filtered_papers,  # Now sending the dictionary directly
+            "multi_papers": filtered_papers,  # Sending the dictionary directly
             "last_displayed_papers": "multi_papers",
             "messages": [
                 ToolMessage(

aiagents4pharma/talk2scholars/tools/s2/query_results.py

@@ -44,7 +44,7 @@ def query_results(question: str, state: Annotated[dict, InjectedState]) -> str:
         raise NoPapersFoundError(
             "No papers found. A search needs to be performed first."
         )
-    context_key = state.get("last_displayed_papers")
+    context_key = state.get("last_displayed_papers","pdf_data")
     dic_papers = state.get(context_key)
     df_papers = pd.DataFrame.from_dict(dic_papers, orient="index")
     df_agent = create_pandas_dataframe_agent(

aiagents4pharma/talk2scholars/tools/s2/retrieve_semantic_scholar_paper_id.py

@@ -19,14 +19,6 @@ from pydantic import Field
 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(
@@ -49,6 +41,14 @@ def retrieve_semantic_scholar_paper_id(
     Returns:
         ToolMessage: A message containing the paper ID.
     """
+    # 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
+        logger.info("Loaded configuration for Semantic Scholar paper ID retrieval tool")
     logger.info("Retrieving ID of paper with title: %s", paper_title)
     endpoint = cfg.api_endpoint
     params = {

aiagents4pharma/talk2scholars/tools/s2/search.py

@@ -37,12 +37,6 @@ class SearchInput(BaseModel):
     tool_call_id: Annotated[str, InjectedToolCallId]
 
 
-# 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
-
-
 @tool("search_tool", args_schema=SearchInput, parse_docstring=True)
 def search_tool(
     query: str,
@@ -56,13 +50,18 @@ def search_tool(
     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.
+        limit (int, optional): The maximum number of results to return. Defaults to 5.
         year (str, optional): Year range for papers.
         Supports formats like "2024-", "-2024", "2024:2025". Defaults to None.
 
     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 = {
@@ -75,33 +74,47 @@ def search_tool(
     if year:
         params["year"] = year
 
-    response = requests.get(endpoint, params=params, timeout=10)
+    # Wrap API call in try/except to catch connectivity issues
+    try:
+        response = requests.get(endpoint, params=params, timeout=10)
+        response.raise_for_status()  # Raises HTTPError for bad responses
+    except requests.exceptions.RequestException as e:
+        logger.error("Failed to connect to Semantic Scholar API: %s", e)
+        raise RuntimeError(
+            "Failed to connect to Semantic Scholar API. Please retry the same query."
+        ) from e
+
     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", [])
-    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,
-                    )
-                ]
-            }
+        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_id": paper["paperId"],  # Store Semantic Scholar ID
+            "paper_id": 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"),
-            # "arXiv_ID": paper.get("externalIds", {}).get(
-            #     "ArXiv", "N/A"
-            # ),  # Extract arXiv ID
+            "arxiv_id": paper.get("externalIds", {}).get("ArXiv", "N/A"),
         }
         for paper in papers
         if paper.get("title") and paper.get("authors")
@@ -129,7 +142,7 @@ def search_tool(
 
     return Command(
         update={
-            "papers": filtered_papers,  # Now sending the dictionary directly
+            "papers": filtered_papers,  # Sending the dictionary directly
             "last_displayed_papers": "papers",
             "messages": [
                 ToolMessage(

aiagents4pharma/talk2scholars/tools/s2/single_paper_rec.py

@@ -40,14 +40,6 @@ class SinglePaperRecInput(BaseModel):
     model_config = {"arbitrary_types_allowed": True}
 
 
-# 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
-
-
 @tool(args_schema=SinglePaperRecInput, parse_docstring=True)
 def get_single_paper_recommendations(
     paper_id: str,
@@ -56,19 +48,27 @@ def get_single_paper_recommendations(
     year: Optional[str] = None,
 ) -> Command[Any]:
     """
-    Get recommendations for on a single paper using its Semantic Scholar ID.
+    Get recommendations for 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.
         tool_call_id (Annotated[str, InjectedToolCallId]): The tool call ID.
-        limit (int, optional): The maximum number of recommendations to return. Defaults to 2.
+        limit (int, optional): The maximum number of recommendations to return. Defaults to 5.
         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.
     """
+    # 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
     )
@@ -84,48 +84,54 @@ def get_single_paper_recommendations(
     if year:
         params["year"] = year
 
-    response = requests.get(endpoint, params=params, timeout=cfg.request_timeout)
-    data = response.json()
-    response = requests.get(endpoint, params=params, timeout=10)
-    # print(f"API Response Status: {response.status_code}")
-    logging.info(
+    # Wrap API call in try/except to catch connectivity issues and check response format
+    try:
+        response = requests.get(endpoint, params=params, timeout=cfg.request_timeout)
+        response.raise_for_status()  # Raises HTTPError for bad responses
+    except requests.exceptions.RequestException as e:
+        logger.error(
+            "Failed to connect to Semantic Scholar API for recommendations: %s", e
+        )
+        raise RuntimeError(
+            "Failed to connect to Semantic Scholar API. Please retry the same query."
+        ) from e
+
+    logger.info(
         "API Response Status for recommendations of paper %s: %s",
         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)
+    logger.info("Request params: %s", params)
 
     data = response.json()
-    recommendations = data.get("recommendedPapers", [])
 
+    # 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."
+        )
+
+    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,
-                    )
-                ],
-            }
+        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_id": paper["paperId"],  # Store Semantic Scholar ID
+            "paper_id": 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"),
-            # "arXiv_ID": paper.get("externalIds", {}).get(
-            #     "ArXiv", "N/A"
-            # ),  # Extract arXiv ID
+            "arxiv_id": paper.get("externalIds", {}).get("ArXiv", "N/A"),
         }
         for paper in recommendations
         if paper.get("title") and paper.get("authors")
@@ -143,10 +149,10 @@ def get_single_paper_recommendations(
     logger.info("Filtered %d papers", len(filtered_papers))
 
     content = (
-        "Recommendations based on single paper were successful. "
-        "Papers are attached as an artifact."
+        "Recommendations based on the single paper were successful. "
+        "Papers are attached as an artifact. "
+        "Here is a summary of the recommendations:\n"
     )
-    content += " Here is a summary of the recommendations:\n"
     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 ""
@@ -154,7 +160,7 @@ def get_single_paper_recommendations(
 
     return Command(
         update={
-            "papers": filtered_papers,  # Now sending the dictionary directly
+            "papers": filtered_papers,  # Sending the dictionary directly
             "last_displayed_papers": "papers",
             "messages": [
                 ToolMessage(

aiagents4pharma/talk2scholars/tools/zotero/__init__.py

@@ -3,3 +3,5 @@ Import statements
 """
 
 from . import zotero_read
+from . import zotero_write
+from . import utils

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

@@ -0,0 +1,5 @@
+"""
+Import statements
+"""
+
+from . import zotero_path