aiagents4pharma 1.40.0__py3-none-any.whl → 1.41.0__py3-none-any.whl

aiagents4pharma/talk2scholars/tests/test_tool_helper_utils.py

@@ -10,28 +10,34 @@ from aiagents4pharma.talk2scholars.tools.pdf.utils.tool_helper import QAToolHelp
 
 
 class TestQAToolHelper(unittest.TestCase):
-    """tests for QAToolHelper routines in tool_helper.py"""
+    """tests for QAToolHelper routines"""
 
     def setUp(self):
-        """set up test case"""
+        """setup for each test"""
         self.helper = QAToolHelper()
 
     def test_start_call_sets_config_and_call_id(self):
-        """test start_call sets config and call_id"""
+        """start_call should set config and call_id"""
         cfg = SimpleNamespace(foo="bar")
         self.helper.start_call(cfg, "call123")
         self.assertIs(self.helper.config, cfg)
         self.assertEqual(self.helper.call_id, "call123")
 
-    def test_init_vector_store_reuse(self):
-        """test init_vector_store reuses existing instance"""
+    @patch("aiagents4pharma.talk2scholars.tools.pdf.utils.tool_helper.get_vectorstore")
+    def test_init_vector_store_reuse(self, mock_get_vectorstore):
+        """Mock vector store reuse test"""
         emb_model = MagicMock()
+        mock_instance = MagicMock()
+        mock_get_vectorstore.return_value = mock_instance
+
         first = self.helper.init_vector_store(emb_model)
         second = self.helper.init_vector_store(emb_model)
-        self.assertIs(second, first)
+
+        self.assertIs(first, second)
+        self.assertIs(second, mock_instance)
 
     def test_get_state_models_and_data_success(self):
-        """test get_state_models_and_data returns models and data"""
+        """get_state_models_and_data should return models and data from state"""
         emb = MagicMock()
         llm = MagicMock()
         articles = {"p": {}}
@@ -46,95 +52,34 @@ class TestQAToolHelper(unittest.TestCase):
         self.assertIs(ret_articles, articles)
 
     def test_get_state_models_and_data_missing_text_embedding(self):
-        """test get_state_models_and_data raises ValueError if missing text embedding"""
+        """get_state_models_and_data should raise ValueError if text_embedding_model is missing"""
         state = {"llm_model": MagicMock(), "article_data": {"p": {}}}
-        with self.assertRaises(ValueError) as cm:
+        with self.assertRaises(ValueError):
             self.helper.get_state_models_and_data(state)
-        self.assertEqual(str(cm.exception), "No text embedding model found in state.")
 
     def test_get_state_models_and_data_missing_llm(self):
-        """test get_state_models_and_data raises ValueError if missing LLM"""
+        """should raise ValueError if llm_model is missing"""
         state = {"text_embedding_model": MagicMock(), "article_data": {"p": {}}}
-        with self.assertRaises(ValueError) as cm:
+        with self.assertRaises(ValueError):
             self.helper.get_state_models_and_data(state)
-        self.assertEqual(str(cm.exception), "No LLM model found in state.")
 
     def test_get_state_models_and_data_missing_article_data(self):
-        """test get_state_models_and_data raises ValueError if missing article data"""
+        """get_state_models_and_data should raise ValueError if article_data is missing"""
         state = {"text_embedding_model": MagicMock(), "llm_model": MagicMock()}
-        with self.assertRaises(ValueError) as cm:
+        with self.assertRaises(ValueError):
             self.helper.get_state_models_and_data(state)
-        self.assertEqual(str(cm.exception), "No article_data found in state.")
-
-    def test_load_candidate_papers_calls_add_paper_only_for_valid(self):
-        """test load_candidate_papers calls add_paper only for valid candidates"""
-        vs = SimpleNamespace(loaded_papers=set(), add_paper=MagicMock())
-        articles = {"p1": {"pdf_url": "url1"}, "p2": {}, "p3": {"pdf_url": None}}
-        candidates = ["p1", "p2", "p3"]
-        self.helper.load_candidate_papers(vs, articles, candidates)
-        vs.add_paper.assert_called_once_with("p1", "url1", articles["p1"])
-
-    def test_load_candidate_papers_handles_add_paper_exception(self):
-        """test load_candidate_papers handles add_paper exception"""
-        # If add_paper raises, it should be caught and not propagate
-        vs = SimpleNamespace(
-            loaded_papers=set(), add_paper=MagicMock(side_effect=ValueError("oops"))
-        )
-        articles = {"p1": {"pdf_url": "url1"}}
-        # Start call to set call_id (used in logging)
-        self.helper.start_call(SimpleNamespace(), "call001")
-        # Should not raise despite exception
-        self.helper.load_candidate_papers(vs, articles, ["p1"])
-        vs.add_paper.assert_called_once_with("p1", "url1", articles["p1"])
-
-    def test_run_reranker_success_and_filtering(self):
-        """test run_reranker success and filtering"""
-        # Successful rerank returns filtered candidates
-        cfg = SimpleNamespace(top_k_papers=2)
-        self.helper.config = cfg
-        vs = MagicMock()
-        with patch(
-            "aiagents4pharma.talk2scholars.tools.pdf.utils.tool_helper.rank_papers_by_query",
-            return_value=["a", "c"],
-        ):
-            out = self.helper.run_reranker(vs, "q", ["a", "b"])
-        self.assertEqual(out, ["a"])
-
-    def test_run_reranker_exception_fallback(self):
-        """test run_reranker exception fallback"""
-        # On reranker failure, should return original candidates
-        cfg = SimpleNamespace(top_k_papers=5)
-        self.helper.config = cfg
-        vs = MagicMock()
-
-        def fail(*args, **kwargs):
-            raise RuntimeError("fail")
-
-        with patch(
-            "aiagents4pharma.talk2scholars.tools.pdf.utils.tool_helper.rank_papers_by_query",
-            side_effect=fail,
-        ):
-            candidates = ["x", "y"]
-            out = self.helper.run_reranker(vs, "q", candidates)
-        self.assertEqual(out, candidates)
-
-    def test_format_answer_with_and_without_sources(self):
-        """test format_answer with and without sources"""
-        articles = {"p1": {"Title": "T1"}, "p2": {"Title": "T2"}}
-        # With sources
-        with patch(
-            "aiagents4pharma.talk2scholars.tools.pdf.utils.tool_helper.generate_answer",
-            return_value={"output_text": "ans", "papers_used": ["p1", "p2"]},
-        ):
-            res = self.helper.format_answer("q", [], MagicMock(), articles)
-            self.assertIn("ans", res)
-            self.assertIn("Sources:", res)
-            self.assertIn("- T1", res)
-            self.assertIn("- T2", res)
-        # Without sources
-        with patch(
-            "aiagents4pharma.talk2scholars.tools.pdf.utils.tool_helper.generate_answer",
-            return_value={"output_text": "ans", "papers_used": []},
-        ):
-            res2 = self.helper.format_answer("q", [], MagicMock(), {})
-            self.assertEqual(res2, "ans")
+
+    def test_get_hardware_stats(self):
+        """get_hardware_stats should return correct GPU and hardware mode"""
+        helper = QAToolHelper()
+        helper.call_id = "test_call"
+
+        helper.has_gpu = False
+        stats = helper.get_hardware_stats()
+        self.assertEqual(stats["gpu_available"], False)
+        self.assertEqual(stats["hardware_mode"], "CPU-only")
+
+        helper.has_gpu = True
+        stats = helper.get_hardware_stats()
+        self.assertEqual(stats["gpu_available"], True)
+        self.assertEqual(stats["hardware_mode"], "GPU-accelerated")

aiagents4pharma/talk2scholars/tools/paper_download/download_biorxiv_input.py

@@ -22,20 +22,20 @@ logger = logging.getLogger(__name__)
 class DownloadBiorxivPaperInput(BaseModel):
     """Input schema for the bioRxiv paper download tool."""
 
-    doi: str = Field(description=
-    """The bioRxiv DOI, from search_helper or multi_helper or single_helper, 
+    doi: str = Field(
+        description="""The bioRxiv DOI, from search_helper or multi_helper or single_helper,
     used to retrieve the paper details and PDF URL."""
     )
-    logger.info("DOI Received: %s", doi)
     tool_call_id: Annotated[str, InjectedToolCallId]
 
+
 def fetch_biorxiv_metadata(doi: str, api_url: str, request_timeout: int) -> dict:
     """
     Fetch metadata for a bioRxiv paper using its DOI and extract relevant fields.
-    
+
     Parameters:
         doi (str): The DOI of the bioRxiv paper.
-    
+
     Returns:
         dict: A dictionary containing the title, authors, abstract, publication date, and URLs.
     """
@@ -55,6 +55,7 @@ def fetch_biorxiv_metadata(doi: str, api_url: str, request_timeout: int) -> dict
 
     return data["collection"][0]
 
+
 def extract_metadata(paper: dict, doi: str) -> dict:
     """
     Extract relevant metadata fields from a bioRxiv paper entry.
@@ -75,9 +76,10 @@ def extract_metadata(paper: dict, doi: str) -> dict:
         "pdf_url": pdf_url,
         "filename": f"{doi_suffix}.pdf",
         "source": "biorxiv",
-        "biorxiv_id": doi
+        "biorxiv_id": doi,
     }
 
+
 @tool(args_schema=DownloadBiorxivPaperInput, parse_docstring=True)
 def download_biorxiv_paper(
     doi: str,

aiagents4pharma/talk2scholars/tools/paper_download/download_medrxiv_input.py

@@ -22,13 +22,13 @@ logger = logging.getLogger(__name__)
 class DownloadMedrxivPaperInput(BaseModel):
     """Input schema for the medRxiv paper download tool."""
 
-    doi: str = Field(description=
-    """The medRxiv DOI, from search_helper or multi_helper or single_helper, 
+    doi: str = Field(
+        description="""The medRxiv DOI, from search_helper or multi_helper or single_helper,
     used to retrieve the paper details and PDF URL."""
     )
-    logger.info("DOI Received: %s", doi)
     tool_call_id: Annotated[str, InjectedToolCallId]
 
+
 # Fetching raw metadata from medRxiv API for a given DOI
 def fetch_medrxiv_metadata(doi: str, api_url: str, request_timeout: int) -> dict:
     """
@@ -54,6 +54,7 @@ def fetch_medrxiv_metadata(doi: str, api_url: str, request_timeout: int) -> dict
 
     return data["collection"][0]
 
+
 # Extracting relevant metadata fields from the raw data
 def extract_metadata(paper: dict, doi: str) -> dict:
     """
@@ -75,9 +76,10 @@ def extract_metadata(paper: dict, doi: str) -> dict:
         "pdf_url": pdf_url,
         "filename": f"{doi_suffix}.pdf",
         "source": "medrxiv",
-        "medrxiv_id": doi
+        "medrxiv_id": doi,
     }
 
+
 # Tool to download medRxiv paper metadata and PDF URL
 @tool(args_schema=DownloadMedrxivPaperInput, parse_docstring=True)
 def download_medrxiv_paper(

aiagents4pharma/talk2scholars/tools/pdf/question_and_answer.py

@@ -1,17 +1,16 @@
 """
 LangGraph PDF Retrieval-Augmented Generation (RAG) Tool
 
-This tool answers user questions by retrieving and ranking relevant text chunks from PDFs
-and invoking an LLM to generate a concise, source-attributed response. It supports
-single or multiple PDF sources—such as Zotero libraries, arXiv papers, or direct uploads.
-
-Workflow:
-  1. (Optional) Load PDFs from diverse sources into a FAISS vector store of embeddings.
-  2. Rerank candidate papers using NVIDIA NIM semantic re-ranker.
-  3. Retrieve top-K diverse text chunks via Maximal Marginal Relevance (MMR).
-  4. Build a context-rich prompt combining retrieved chunks and the user question.
-  5. Invoke the LLM to craft a clear answer with source citations.
-  6. Return the answer in a ToolMessage for LangGraph to dispatch.
+This tool answers user questions using the traditional RAG pipeline:
+1. Retrieve relevant chunks from ALL papers in the vector store
+2. Rerank chunks using NVIDIA NIM reranker to find the most relevant ones
+3. Generate answer using the top reranked chunks
+
+Traditional RAG Pipeline Flow:
+  Query → Retrieve chunks from ALL papers → Rerank chunks → Generate answer
+
+This ensures the best possible chunks are selected across all available papers,
+not just from pre-selected papers.
 """
 
 import logging
@@ -27,8 +26,10 @@ from langgraph.types import Command
 from pydantic import BaseModel, Field
 
 from .utils.generate_answer import load_hydra_config
-from .utils.retrieve_chunks import retrieve_relevant_chunks
 from .utils.tool_helper import QAToolHelper
+from .utils.paper_loader import load_all_papers
+from .utils.rag_pipeline import retrieve_and_rerank_chunks
+from .utils.answer_formatter import format_answer
 
 # Helper for managing state, vectorstore, reranking, and formatting
 helper = QAToolHelper()
@@ -53,7 +54,6 @@ class QuestionAndAnswerInput(BaseModel):
         - article_data: metadata mapping of paper IDs to info (e.g., 'pdf_url', title).
         - text_embedding_model: embedding model instance for chunk indexing.
         - llm_model: chat/LLM instance for answer generation.
-        - vector_store: optional pre-built Vectorstore for retrieval.
     """
 
     question: str = Field(
@@ -70,17 +70,16 @@ def question_and_answer(
     tool_call_id: Annotated[str, InjectedToolCallId],
 ) -> Command[Any]:
     """
-    LangGraph tool for Retrieval-Augmented Generation over PDFs.
-
-    Given a user question, this tool applies the following pipeline:
-      1. Validates that embedding and LLM models, plus article metadata, are in state.
-      2. Initializes or reuses a FAISS-based Vectorstore for PDF embeddings.
-      3. Loads one or more PDFs (from Zotero, arXiv, uploads) as text chunks into the store.
-      4. Uses NVIDIA NIM semantic re-ranker to select top candidate papers.
-      5. Retrieves the most relevant and diverse text chunks via Maximal Marginal Relevance.
-      6. Constructs an LLM prompt combining contextual chunks and the query.
-      7. Invokes the LLM to generate an answer, appending source attributions.
-      8. Returns a LangGraph Command with a ToolMessage containing the answer.
+    LangGraph tool for Retrieval-Augmented Generation over PDFs using traditional RAG pipeline.
+
+    Traditional RAG Pipeline Implementation:
+      1. Load ALL available PDFs into Milvus vector store (if not already loaded)
+      2. Retrieve relevant chunks from ALL papers using vector similarity search
+      3. Rerank retrieved chunks using NVIDIA NIM semantic reranker
+      4. Generate answer using top reranked chunks with source attribution
+
+    This approach ensures the best chunks are selected across all available papers,
+    rather than pre-selecting papers and potentially missing relevant information.
 
     Args:
       question (str): The free-text question to answer.
@@ -99,35 +98,70 @@ def question_and_answer(
     """
     call_id = f"qa_call_{time.time()}"
     logger.info(
-        "Starting PDF Question and Answer tool call %s for question: %s",
+        "Starting PDF Question and Answer tool (Traditional RAG Pipeline) - Call %s",
         call_id,
-        question,
     )
+    logger.info("%s: Question: '%s'", call_id, question)
+
     helper.start_call(config, call_id)
 
     # Extract models and article metadata
     text_emb, llm_model, article_data = helper.get_state_models_and_data(state)
 
-    # Initialize or reuse vector store, then load candidate papers
+    # Initialize or reuse Milvus vector store
+    logger.info("%s: Initializing vector store", call_id)
     vs = helper.init_vector_store(text_emb)
-    candidate_ids = list(article_data.keys())
-    logger.info("%s: Candidate paper IDs for reranking: %s", call_id, candidate_ids)
-    helper.load_candidate_papers(vs, article_data, candidate_ids)
-
-    # Rerank papers and retrieve top chunks
-    selected_ids = helper.run_reranker(vs, question, candidate_ids)
-    relevant_chunks = retrieve_relevant_chunks(
-        vs, query=question, paper_ids=selected_ids, top_k=config.top_k_chunks
+
+    # Load ALL papers (traditional RAG approach)
+    logger.info(
+        "%s: Loading all %d papers into vector store (traditional RAG approach)",
+        call_id,
+        len(article_data),
+    )
+    load_all_papers(
+        vector_store=vs,
+        articles=article_data,
+        call_id=call_id,
+        config=config,
+        has_gpu=helper.has_gpu,
+    )
+
+    # Traditional RAG Pipeline: Retrieve from ALL papers, then rerank
+    logger.info(
+        "%s: Starting traditional RAG pipeline: retrieve → rerank → generate",
+        call_id,
+    )
+
+    # Retrieve and rerank chunks in one step
+    reranked_chunks = retrieve_and_rerank_chunks(
+        vs, question, config, call_id, helper.has_gpu
     )
-    if not relevant_chunks:
+
+    if not reranked_chunks:
         msg = f"No relevant chunks found for question: '{question}'"
         logger.warning("%s: %s", call_id, msg)
-        raise RuntimeError(msg)
 
-    # Generate answer and format with sources
-    response_text = helper.format_answer(
-        question, relevant_chunks, llm_model, article_data
+    # Generate answer using reranked chunks
+    logger.info(
+        "%s: Generating answer using %d reranked chunks",
+        call_id,
+        len(reranked_chunks),
     )
+    response_text = format_answer(
+        question,
+        reranked_chunks,
+        llm_model,
+        article_data,
+        config,
+        call_id=call_id,
+        has_gpu=helper.has_gpu,
+    )
+
+    logger.info(
+        "%s: Successfully traditional completed RAG pipeline",
+        call_id,
+    )
+
     return Command(
         update={
             "messages": [

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

@@ -2,9 +2,34 @@
 Utility modules for the PDF question_and_answer tool.
 """
 
+from . import answer_formatter
+from . import batch_processor
+from . import collection_manager
 from . import generate_answer
+from . import get_vectorstore
+from . import gpu_detection
 from . import nvidia_nim_reranker
+from . import paper_loader
+from . import rag_pipeline
 from . import retrieve_chunks
+from . import singleton_manager
+from . import tool_helper
+from . import vector_normalization
 from . import vector_store
 
-__all__ = ["generate_answer", "nvidia_nim_reranker", "retrieve_chunks", "vector_store"]
+__all__ = [
+    "answer_formatter",
+    "batch_processor",
+    "collection_manager",
+    "generate_answer",
+    "get_vectorstore",
+    "gpu_detection",
+    "nvidia_nim_reranker",
+    "paper_loader",
+    "rag_pipeline",
+    "retrieve_chunks",
+    "singleton_manager",
+    "tool_helper",
+    "vector_normalization",
+    "vector_store",
+]

aiagents4pharma/talk2scholars/tools/pdf/utils/answer_formatter.py

@@ -0,0 +1,62 @@
+"""
+Format the final answer text with source attributions and hardware info.
+"""
+
+import logging
+from typing import Any, Dict, List
+
+from .generate_answer import generate_answer
+
+logger = logging.getLogger(__name__)
+
+
+def format_answer(
+    question: str,
+    chunks: List[Any],
+    llm: Any,
+    articles: Dict[str, Any],
+    config: Any,
+    **kwargs: Any,
+) -> str:
+    """
+    Generate the final answer text with source attributions and hardware info.
+
+    Expects `call_id` and `has_gpu` in kwargs.
+    """
+    result = generate_answer(question, chunks, llm, config)
+    answer = result.get("output_text", "No answer generated.")
+
+    # Get unique paper titles for source attribution
+    titles: Dict[str, str] = {}
+    for pid in result.get("papers_used", []):
+        if pid in articles:
+            titles[pid] = articles[pid].get("Title", "Unknown paper")
+
+    # Format sources
+    if titles:
+        srcs = "\n\nSources:\n" + "\n".join(f"- {t}" for t in titles.values())
+    else:
+        srcs = ""
+
+    # Extract logging metadata
+    call_id = kwargs.get("call_id", "<no-call-id>")
+    has_gpu = kwargs.get("has_gpu", False)
+    hardware_info = "GPU-accelerated" if has_gpu else "CPU-processed"
+
+    # Log final statistics with hardware info
+    logger.info(
+        "%s: Generated answer using %d chunks from %d papers (%s)",
+        call_id,
+        len(chunks),
+        len(titles),
+        hardware_info,
+    )
+
+    # Add subtle hardware info to logs but not to user output
+    logger.debug(
+        "%s: Answer generation completed with %s processing",
+        call_id,
+        hardware_info,
+    )
+
+    return f"{answer}{srcs}"