aiagents4pharma 1.39.0__py3-none-any.whl → 1.39.2__py3-none-any.whl

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

@@ -0,0 +1,10 @@
+"""
+Utility modules for the PDF question_and_answer tool.
+"""
+
+from . import generate_answer
+from . import nvidia_nim_reranker
+from . import retrieve_chunks
+from . import vector_store
+
+__all__ = ["generate_answer", "nvidia_nim_reranker", "retrieve_chunks", "vector_store"]

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

@@ -0,0 +1,97 @@
+"""
+Generate an answer for a question using retrieved chunks of documents.
+"""
+
+import logging
+import os
+from typing import Any, Dict, List
+
+import hydra
+from langchain_core.documents import Document
+from langchain_core.language_models.chat_models import BaseChatModel
+
+# Set up logging with configurable level
+log_level = os.environ.get("LOG_LEVEL", "INFO")
+logging.basicConfig(level=getattr(logging, log_level))
+logger = logging.getLogger(__name__)
+logger.setLevel(getattr(logging, log_level))
+
+
+def _build_context_and_sources(
+    retrieved_chunks: List[Document],
+) -> tuple[str, set[str]]:
+    """
+    Build the combined context string and set of paper_ids from retrieved chunks.
+    """
+    papers = {}
+    for doc in retrieved_chunks:
+        pid = doc.metadata.get("paper_id", "unknown")
+        papers.setdefault(pid, []).append(doc)
+    formatted = []
+    idx = 1
+    for pid, chunks in papers.items():
+        title = chunks[0].metadata.get("title", "Unknown")
+        formatted.append(f"[Document {idx}] From: '{title}' (ID: {pid})")
+        for chunk in chunks:
+            page = chunk.metadata.get("page", "unknown")
+            formatted.append(f"Page {page}: {chunk.page_content}")
+        idx += 1
+    context = "\n\n".join(formatted)
+    sources: set[str] = set()
+    for doc in retrieved_chunks:
+        pid = doc.metadata.get("paper_id")
+        if isinstance(pid, str):
+            sources.add(pid)
+    return context, sources
+
+
+def load_hydra_config() -> Any:
+    """
+    Load the configuration using Hydra and return the configuration for the Q&A tool.
+    """
+    with hydra.initialize(version_base=None, config_path="../../../configs"):
+        cfg = hydra.compose(
+            config_name="config",
+            overrides=["tools/question_and_answer=default"],
+        )
+        config = cfg.tools.question_and_answer
+        logger.debug("Loaded Question and Answer tool configuration.")
+        return config
+
+
+def generate_answer(
+    question: str,
+    retrieved_chunks: List[Document],
+    llm_model: BaseChatModel,
+    config: Any,
+) -> Dict[str, Any]:
+    """
+    Generate an answer for a question using retrieved chunks.
+
+    Args:
+        question (str): The question to answer
+        retrieved_chunks (List[Document]): List of relevant document chunks
+        llm_model (BaseChatModel): Language model for generating answers
+        config (Any): Configuration for answer generation
+
+    Returns:
+        Dict[str, Any]: Dictionary with the answer and metadata
+    """
+    # Ensure the configuration is provided and has the prompt_template.
+    if config is None:
+        raise ValueError("Configuration for generate_answer is required.")
+    if "prompt_template" not in config:
+        raise ValueError("The prompt_template is missing from the configuration.")
+
+    # Build context and sources, then invoke LLM
+    context, paper_sources = _build_context_and_sources(retrieved_chunks)
+    prompt = config["prompt_template"].format(context=context, question=question)
+    response = llm_model.invoke(prompt)
+
+    # Return the response with metadata
+    return {
+        "output_text": response.content,
+        "sources": [doc.metadata for doc in retrieved_chunks],
+        "num_sources": len(retrieved_chunks),
+        "papers_used": list(paper_sources),
+    }

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

@@ -0,0 +1,77 @@
+"""
+NVIDIA NIM Reranker Utility
+"""
+
+import logging
+import os
+
+
+from typing import Any, List
+
+from langchain_core.documents import Document
+from langchain_nvidia_ai_endpoints import NVIDIARerank
+
+# Set up logging with configurable level
+log_level = os.environ.get("LOG_LEVEL", "INFO")
+logging.basicConfig(level=getattr(logging, log_level))
+logger = logging.getLogger(__name__)
+logger.setLevel(getattr(logging, log_level))
+
+
+def rank_papers_by_query(self, query: str, config: Any, top_k: int = 5) -> List[str]:
+    """
+    Rank papers by relevance to the query using NVIDIA's off-the-shelf re-ranker.
+
+    This function aggregates all chunks per paper, ranks them using the NVIDIA model,
+    and returns the top-k papers.
+
+    Args:
+        query (str): The query string.
+        config (Any): Configuration containing reranker settings (model, api_key).
+        top_k (int): Number of top papers to return.
+
+    Returns:
+        List of tuples (paper_id, dummy_score) sorted by relevance.
+    """
+
+    logger.info("Starting NVIDIA re-ranker for query: '%s' with top_k=%d", query, top_k)
+    # Aggregate all document chunks for each paper
+    paper_texts = {}
+    for doc in self.documents.values():
+        paper_id = doc.metadata["paper_id"]
+        paper_texts.setdefault(paper_id, []).append(doc.page_content)
+
+    aggregated_documents = []
+    for paper_id, texts in paper_texts.items():
+        aggregated_text = " ".join(texts)
+        aggregated_documents.append(
+            Document(page_content=aggregated_text, metadata={"paper_id": paper_id})
+        )
+
+    logger.info(
+        "Aggregated %d papers into %d documents for reranking",
+        len(paper_texts),
+        len(aggregated_documents),
+    )
+    # Instantiate the NVIDIA re-ranker client using provided config
+    # Use NVIDIA API key from Hydra configuration (expected to be resolved via oc.env)
+    api_key = config.reranker.api_key
+    if not api_key:
+        logger.error("No NVIDIA API key found in configuration for reranking")
+        raise ValueError("Configuration 'reranker.api_key' must be set for reranking")
+    logger.info("Using NVIDIA API key from configuration for reranking")
+    # Truncate long inputs at the model-end to avoid exceeding max token size
+    logger.info("Setting NVIDIA reranker truncate mode to END to limit input length")
+    reranker = NVIDIARerank(
+        model=config.reranker.model,
+        api_key=api_key,
+        truncate="END",
+    )
+
+    # Get the ranked list of documents based on the query
+    response = reranker.compress_documents(query=query, documents=aggregated_documents)
+    logger.info("Received %d documents from NVIDIA reranker", len(response))
+
+    ranked_papers = [doc.metadata["paper_id"] for doc in response[:top_k]]
+    logger.info("Top %d papers after reranking: %s", top_k, ranked_papers)
+    return ranked_papers

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

@@ -0,0 +1,83 @@
+"""
+Retrieve relevant chunks from a vector store using MMR (Maximal Marginal Relevance).
+"""
+
+import logging
+import os
+from typing import List, Optional
+
+import numpy as np
+from langchain_core.documents import Document
+from langchain_core.vectorstores.utils import maximal_marginal_relevance
+
+
+# Set up logging with configurable level
+log_level = os.environ.get("LOG_LEVEL", "INFO")
+logging.basicConfig(level=getattr(logging, log_level))
+logger = logging.getLogger(__name__)
+logger.setLevel(getattr(logging, log_level))
+
+
+def retrieve_relevant_chunks(
+    self,
+    query: str,
+    paper_ids: Optional[List[str]] = None,
+    top_k: int = 25,
+    mmr_diversity: float = 1.00,
+) -> List[Document]:
+    """
+    Retrieve the most relevant chunks for a query using maximal marginal relevance.
+
+    Args:
+        query: Query string
+        paper_ids: Optional list of paper IDs to filter by
+        top_k: Number of chunks to retrieve
+        mmr_diversity: Diversity parameter for MMR (higher = more diverse)
+
+    Returns:
+        List of document chunks
+    """
+    if not self.vector_store:
+        logger.error("Failed to build vector store")
+        return []
+
+    if paper_ids:
+        logger.info("Filtering retrieval to papers: %s", paper_ids)
+
+    # Step 1: Embed the query
+    logger.info("Embedding query using model: %s", type(self.embedding_model).__name__)
+    query_embedding = np.array(self.embedding_model.embed_query(query))
+
+    # Step 2: Filter relevant documents
+    all_docs = [
+        doc
+        for doc in self.documents.values()
+        if not paper_ids or doc.metadata["paper_id"] in paper_ids
+    ]
+
+    if not all_docs:
+        logger.warning("No documents found after filtering by paper_ids.")
+        return []
+
+    # Step 3: Retrieve or compute embeddings for all documents using cache
+    logger.info("Retrieving embeddings for %d chunks...", len(all_docs))
+    all_embeddings = []
+    for doc in all_docs:
+        doc_id = f"{doc.metadata['paper_id']}_{doc.metadata['chunk_id']}"
+        if doc_id not in self.embeddings:
+            logger.info("Embedding missing chunk %s", doc_id)
+            emb = self.embedding_model.embed_documents([doc.page_content])[0]
+            self.embeddings[doc_id] = emb
+        all_embeddings.append(self.embeddings[doc_id])
+
+    # Step 4: Apply MMR
+    mmr_indices = maximal_marginal_relevance(
+        query_embedding,
+        all_embeddings,
+        k=top_k,
+        lambda_mult=mmr_diversity,
+    )
+
+    results = [all_docs[i] for i in mmr_indices]
+    logger.info("Retrieved %d chunks using MMR", len(results))
+    return results

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

@@ -0,0 +1,125 @@
+"""
+Helper class for PDF Q&A tool orchestration: state validation, vectorstore init,
+paper loading, reranking, and answer formatting.
+"""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from .generate_answer import generate_answer
+from .nvidia_nim_reranker import rank_papers_by_query
+from .vector_store import Vectorstore
+
+logger = logging.getLogger(__name__)
+
+
+class QAToolHelper:
+    """Encapsulates helper routines for the PDF Question & Answer tool."""
+
+    def __init__(self) -> None:
+        self.prebuilt_vector_store: Optional[Vectorstore] = None
+        self.config: Any = None
+        self.call_id: str = ""
+        logger.debug("Initialized QAToolHelper")
+
+    def start_call(self, config: Any, call_id: str) -> None:
+        """Initialize helper with current config and call identifier."""
+        self.config = config
+        self.call_id = call_id
+        logger.debug("QAToolHelper started call %s", call_id)
+
+    def get_state_models_and_data(self, state: dict) -> tuple[Any, Any, Dict[str, Any]]:
+        """Retrieve embedding model, LLM, and article data from agent state."""
+        text_emb = state.get("text_embedding_model")
+        if not text_emb:
+            msg = "No text embedding model found in state."
+            logger.error("%s: %s", self.call_id, msg)
+            raise ValueError(msg)
+        llm = state.get("llm_model")
+        if not llm:
+            msg = "No LLM model found in state."
+            logger.error("%s: %s", self.call_id, msg)
+            raise ValueError(msg)
+        articles = state.get("article_data", {})
+        if not articles:
+            msg = "No article_data found in state."
+            logger.error("%s: %s", self.call_id, msg)
+            raise ValueError(msg)
+        return text_emb, llm, articles
+
+    def init_vector_store(self, emb_model: Any) -> Vectorstore:
+        """Return shared or new Vectorstore instance."""
+        if self.prebuilt_vector_store is not None:
+            logger.info("Using shared pre-built vector store from memory")
+            return self.prebuilt_vector_store
+        vs = Vectorstore(embedding_model=emb_model, config=self.config)
+        logger.info("Initialized new vector store with provided configuration")
+        self.prebuilt_vector_store = vs
+        return vs
+
+    def load_candidate_papers(
+        self,
+        vs: Vectorstore,
+        articles: Dict[str, Any],
+        candidates: List[str],
+    ) -> None:
+        """Ensure each candidate paper is loaded into the vector store."""
+        for pid in candidates:
+            if pid not in vs.loaded_papers:
+                pdf_url = articles.get(pid, {}).get("pdf_url")
+                if not pdf_url:
+                    continue
+                try:
+                    vs.add_paper(pid, pdf_url, articles[pid])
+                except (IOError, ValueError) as exc:
+                    logger.warning(
+                        "%s: Error loading paper %s: %s", self.call_id, pid, exc
+                    )
+
+    def run_reranker(
+        self,
+        vs: Vectorstore,
+        query: str,
+        candidates: List[str],
+    ) -> List[str]:
+        """Rank papers by relevance and return filtered paper IDs."""
+        try:
+            ranked = rank_papers_by_query(
+                vs, query, self.config, top_k=self.config.top_k_papers
+            )
+            logger.info("%s: Papers after NVIDIA reranking: %s", self.call_id, ranked)
+            return [pid for pid in ranked if pid in candidates]
+        except (ValueError, RuntimeError) as exc:
+            logger.error("%s: NVIDIA reranker failed: %s", self.call_id, exc)
+            logger.info(
+                "%s: Falling back to all %d candidate papers",
+                self.call_id,
+                len(candidates),
+            )
+            return candidates
+
+    def format_answer(
+        self,
+        question: str,
+        chunks: List[Any],
+        llm: Any,
+        articles: Dict[str, Any],
+    ) -> str:
+        """Generate the final answer text with source attributions."""
+        result = generate_answer(question, chunks, llm, self.config)
+        answer = result.get("output_text", "No answer generated.")
+        titles: Dict[str, str] = {}
+        for pid in result.get("papers_used", []):
+            if pid in articles:
+                titles[pid] = articles[pid].get("Title", "Unknown paper")
+        if titles:
+            srcs = "\n\nSources:\n" + "\n".join(f"- {t}" for t in titles.values())
+        else:
+            srcs = ""
+        logger.info(
+            "%s: Generated answer using %d chunks from %d papers",
+            self.call_id,
+            len(chunks),
+            len(titles),
+        )
+        return f"{answer}{srcs}"

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

@@ -0,0 +1,162 @@
+"""
+Vectorstore class for managing document embeddings and retrieval.
+"""
+
+import logging
+import os
+import time
+from typing import Any, Dict, List, Optional
+
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_community.vectorstores import FAISS
+from langchain_core.documents import Document
+from langchain_core.embeddings import Embeddings
+from langchain_core.vectorstores import VectorStore
+
+
+# Set up logging with configurable level
+log_level = os.environ.get("LOG_LEVEL", "INFO")
+logging.basicConfig(level=getattr(logging, log_level))
+logger = logging.getLogger(__name__)
+logger.setLevel(getattr(logging, log_level))
+
+
+class Vectorstore:
+    """
+    A class for managing document embeddings and retrieval.
+    Provides unified access to documents across multiple papers.
+    """
+
+    def __init__(
+        self,
+        embedding_model: Embeddings,
+        metadata_fields: Optional[List[str]] = None,
+        config: Any = None,
+    ):
+        """
+        Initialize the document store.
+
+        Args:
+            embedding_model: The embedding model to use
+            metadata_fields: Fields to include in document metadata for filtering/retrieval
+        """
+        self.embedding_model = embedding_model
+        self.config = config
+        self.metadata_fields = metadata_fields or [
+            "title",
+            "paper_id",
+            "page",
+            "chunk_id",
+        ]
+        self.initialization_time = time.time()
+        logger.info("Vectorstore initialized at: %s", self.initialization_time)
+
+        # Track loaded papers to prevent duplicate loading
+        self.loaded_papers = set()
+        self.vector_store_class = FAISS
+        logger.info("Using FAISS vector store")
+
+        # Store for initialized documents
+        self.documents: Dict[str, Document] = {}
+        self.vector_store: Optional[VectorStore] = None
+        self.paper_metadata: Dict[str, Dict[str, Any]] = {}
+        # Cache for document chunk embeddings to avoid recomputation
+        self.embeddings: Dict[str, Any] = {}
+
+    def add_paper(
+        self,
+        paper_id: str,
+        pdf_url: str,
+        paper_metadata: Dict[str, Any],
+    ) -> None:
+        """
+        Add a paper to the document store.
+
+        Args:
+            paper_id: Unique identifier for the paper
+            pdf_url: URL to the PDF
+            paper_metadata: Metadata about the paper
+        """
+        # Skip if already loaded
+        if paper_id in self.loaded_papers:
+            logger.info("Paper %s already loaded, skipping", paper_id)
+            return
+
+        logger.info("Loading paper %s from %s", paper_id, pdf_url)
+
+        # Store paper metadata
+        self.paper_metadata[paper_id] = paper_metadata
+
+        # Load the PDF and split into chunks according to Hydra config
+        loader = PyPDFLoader(pdf_url)
+        documents = loader.load()
+        logger.info("Loaded %d pages from %s", len(documents), paper_id)
+
+        # Create text splitter according to provided configuration
+        if self.config is None:
+            raise ValueError(
+                "Configuration is required for text splitting in Vectorstore."
+            )
+        splitter = RecursiveCharacterTextSplitter(
+            chunk_size=self.config.chunk_size,
+            chunk_overlap=self.config.chunk_overlap,
+            separators=["\n\n", "\n", ". ", " ", ""],
+        )
+
+        # Split documents and add metadata for each chunk
+        chunks = splitter.split_documents(documents)
+        logger.info("Split %s into %d chunks", paper_id, len(chunks))
+        # Embed and cache chunk embeddings
+        chunk_texts = [chunk.page_content for chunk in chunks]
+        chunk_embeddings = self.embedding_model.embed_documents(chunk_texts)
+        logger.info("Embedded %d chunks for paper %s", len(chunk_embeddings), paper_id)
+
+        # Enhance document metadata
+        for i, chunk in enumerate(chunks):
+            # Add paper metadata to each chunk
+            chunk.metadata.update(
+                {
+                    "paper_id": paper_id,
+                    "title": paper_metadata.get("Title", "Unknown"),
+                    "chunk_id": i,
+                    # Keep existing page number if available
+                    "page": chunk.metadata.get("page", 0),
+                }
+            )
+
+            # Add any additional metadata fields
+            for field in self.metadata_fields:
+                if field in paper_metadata and field not in chunk.metadata:
+                    chunk.metadata[field] = paper_metadata[field]
+
+            # Store chunk
+            doc_id = f"{paper_id}_{i}"
+            self.documents[doc_id] = chunk
+            # Cache embedding if available
+            if chunk_embeddings[i] is not None:
+                self.embeddings[doc_id] = chunk_embeddings[i]
+
+        # Mark as loaded to prevent duplicate loading
+        self.loaded_papers.add(paper_id)
+        logger.info("Added %d chunks from paper %s", len(chunks), paper_id)
+
+    def build_vector_store(self) -> None:
+        """
+        Build the vector store from all loaded documents.
+        Should be called after all papers are added.
+        """
+        if not self.documents:
+            logger.warning("No documents added to build vector store")
+            return
+
+        if self.vector_store is not None:
+            logger.info("Vector store already built, skipping")
+            return
+
+        # Create vector store from documents
+        documents_list = list(self.documents.values())
+        self.vector_store = self.vector_store_class.from_documents(
+            documents=documents_list, embedding=self.embedding_model
+        )
+        logger.info("Built vector store with %d documents", len(documents_list))

aiagents4pharma/talk2scholars/tools/s2/display_dataframe.py

@@ -4,17 +4,19 @@
 """
 Tool for rendering the most recently displayed papers as a DataFrame artifact for the front-end.
 
-This module defines a tool that retrieves the paper metadata stored under the state key
-'last_displayed_papers' and returns it as an artifact (dictionary of papers). The front-end
-can then render this artifact as a pandas DataFrame for display. If no papers are found,
-a NoPapersFoundError is raised to indicate that a search or recommendation should be
-performed first.
+Call this tool when you need to present the current set of retrieved papers to the user
+(e.g., "show me the papers", "display results"). It reads the 'last_displayed_papers'
+dictionary from the agent state and returns it as an artifact that the UI will render
+as a pandas DataFrame. This tool does not perform any new searches or filtering; it
+only displays the existing list. If no papers are available, it raises NoPapersFoundError
+to signal that a search or recommendation must be executed first.
 """
 
 
 import logging
 
 from typing import Annotated
+from pydantic import BaseModel, Field
 from langchain_core.messages import ToolMessage
 from langchain_core.tools import tool
 from langchain_core.tools.base import InjectedToolCallId
@@ -40,10 +42,31 @@ class NoPapersFoundError(Exception):
     """
 
 
-@tool("display_dataframe", parse_docstring=True)
+class DisplayDataFrameInput(BaseModel):
+    """
+    Pydantic schema for displaying the last set of papers as a DataFrame artifact.
+
+    Fields:
+      state: Agent state dict containing the 'last_displayed_papers' key.
+      tool_call_id: LangGraph-injected identifier for this tool invocation.
+    """
+
+    state: Annotated[dict, InjectedState] = Field(
+        ..., description="Agent state containing the 'last_displayed_papers' reference."
+    )
+    tool_call_id: Annotated[str, InjectedToolCallId] = Field(
+        ..., description="LangGraph-injected identifier for this tool call."
+    )
+
+
+@tool(
+    "display_dataframe",
+    args_schema=DisplayDataFrameInput,
+    parse_docstring=True,
+)
 def display_dataframe(
-    tool_call_id: Annotated[str, InjectedToolCallId],
-    state: Annotated[dict, InjectedState],
+    tool_call_id: str,
+    state: dict,
 ) -> Command:
     """
     Render the last set of retrieved papers as a DataFrame in the front-end.
@@ -55,7 +78,7 @@ def display_dataframe(
     that a search or recommendation must be performed first.
 
     Args:
-        tool_call_id (InjectedToolCallId): Unique ID of this tool invocation.
+        tool_call_id (str): LangGraph-injected unique ID for this tool call.
         state (dict): The agent's state containing the 'last_displayed_papers' reference.
 
     Returns:
@@ -65,7 +88,7 @@ def display_dataframe(
     Raises:
         NoPapersFoundError: If no entries exist under 'last_displayed_papers' in state.
     """
-    logger.info("Displaying papers")
+    logger.info("Displaying papers from 'last_displayed_papers'")
     context_val = state.get("last_displayed_papers")
     # Support both key reference (str) and direct mapping
     if isinstance(context_val, dict):