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

aiagents4pharma/talk2scholars/tests/test_pdf_gpu_detection.py

@@ -0,0 +1,140 @@
+"""gpu detection and index configuration tests."""
+
+import subprocess
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+
+from aiagents4pharma.talk2scholars.tools.pdf.utils.gpu_detection import (
+    detect_nvidia_gpu,
+    get_optimal_index_config,
+    log_index_configuration,
+)
+
+# === detect_nvidia_gpu ===
+
+
+def test_detect_nvidia_gpu_force_cpu_from_config():
+    """detect_nvidia_gpu should return False if force_cpu_mode is set."""
+    config = SimpleNamespace(gpu_detection=SimpleNamespace(force_cpu_mode=True))
+    assert detect_nvidia_gpu(config) is False
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.gpu_detection.subprocess.run")
+def test_detect_nvidia_gpu_success(mock_run):
+    """detect_nvidia_gpu should return True if NVIDIA GPUs are detected."""
+    mock_run.return_value = MagicMock(
+        returncode=0, stdout="NVIDIA A100\nNVIDIA RTX 3090"
+    )
+
+    assert detect_nvidia_gpu() is True
+    mock_run.assert_called_once()
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.gpu_detection.subprocess.run")
+def test_detect_nvidia_gpu_no_output(mock_run):
+    """detect_nvidia_gpu should return False if no GPUs are detected."""
+    mock_run.return_value = MagicMock(returncode=0, stdout="")
+
+    assert detect_nvidia_gpu() is False
+
+
+# === get_optimal_index_config ===
+
+
+def test_get_optimal_index_config_gpu():
+    """get_optimal_index_config should return GPU_CAGRA for GPU setup."""
+    index_params, search_params = get_optimal_index_config(
+        has_gpu=True, embedding_dim=768
+    )
+
+    assert index_params["index_type"] == "GPU_CAGRA"
+    assert "cache_dataset_on_device" in index_params["params"]
+    assert search_params["params"]["search_width"] == 16
+
+
+def test_get_optimal_index_config_cpu():
+    """get_optimal_index_config should return IVF_FLAT for CPU setup."""
+    index_params, search_params = get_optimal_index_config(
+        has_gpu=False, embedding_dim=768
+    )
+
+    assert index_params["index_type"] == "IVF_FLAT"
+    assert index_params["params"]["nlist"] == 96  # 768 / 8 = 96
+    assert search_params["params"]["nprobe"] == 16
+
+
+# === log_index_configuration ===
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.gpu_detection.logger")
+def test_log_index_configuration_logs_all(mock_logger):
+    """log_index_configuration should log all parameters correctly."""
+    index_params = {
+        "index_type": "IVF_FLAT",
+        "metric_type": "COSINE",
+        "params": {"nlist": 128},
+    }
+    search_params = {"metric_type": "COSINE", "params": {"nprobe": 16}}
+
+    log_index_configuration(index_params, search_params)
+
+    assert mock_logger.info.call_count >= 5
+
+
+def test_get_optimal_index_config_gpu_without_cosine():
+    """Ensure GPU config defaults to IP when use_cosine is False."""
+    index_params, search_params = get_optimal_index_config(
+        has_gpu=True, embedding_dim=768, use_cosine=False
+    )
+
+    assert index_params["index_type"] == "GPU_CAGRA"
+    assert index_params["metric_type"] == "IP"
+    assert search_params["metric_type"] == "IP"
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.gpu_detection.logger")
+def test_log_index_configuration_logs_cosine_simulation_note(mock_logger):
+    """Test GPU_CAGRA COSINE -> IP note is logged properly."""
+    index_params = {
+        "index_type": "GPU_CAGRA",
+        "metric_type": "IP",
+        "params": {"itopk_size": 128},
+    }
+    search_params = {
+        "metric_type": "IP",
+        "params": {"search_width": 16},
+    }
+
+    log_index_configuration(index_params, search_params, use_cosine=True)
+
+    log_messages = [str(call.args[0]) for call in mock_logger.info.call_args_list]
+    assert any("simulate COSINE for GPU" in msg for msg in log_messages)
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.gpu_detection.logger")
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.gpu_detection.subprocess.run")
+def test_detect_nvidia_gpu_timeout_raises_false(mock_run, mock_logger):
+    """detect_nvidia_gpu should return False and log info on subprocess.TimeoutExpired."""
+    # Simulate a timeout
+    mock_run.side_effect = subprocess.TimeoutExpired(cmd="nvidia-smi", timeout=10)
+
+    result = detect_nvidia_gpu()
+    assert result is False
+    mock_logger.info.assert_called_with(
+        "NVIDIA GPU detection failed: %s", mock_run.side_effect
+    )
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.gpu_detection.logger")
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.gpu_detection.subprocess.run")
+def test_detect_nvidia_gpu_file_not_found_raises_false(mock_run, mock_logger):
+    """detect_nvidia_gpu should return False and log info on FileNotFoundError."""
+    # Simulate nvidia-smi not installed
+    mock_run.side_effect = FileNotFoundError("nvidia-smi not found")
+
+    result = detect_nvidia_gpu()
+    assert result is False
+    mock_logger.info.assert_called_with(
+        "NVIDIA GPU detection failed: %s", mock_run.side_effect
+    )

aiagents4pharma/talk2scholars/tests/test_pdf_paper_loader.py

@@ -0,0 +1,116 @@
+"""paper_loader tests for the load_all_papers function."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from aiagents4pharma.talk2scholars.tools.pdf.utils.paper_loader import (
+    load_all_papers,
+)
+
+
+@pytest.fixture
+def articles():
+    """A fixture to provide a sample articles dictionary."""
+    return {
+        "p1": {"pdf_url": "http://example.com/p1.pdf", "title": "Paper 1"},
+        "p2": {"pdf_url": "http://example.com/p2.pdf", "title": "Paper 2"},
+        "p3": {"title": "No PDF paper"},
+    }
+
+
+@pytest.fixture
+def mock_vector_store():
+    """Mock vector store fixture."""
+    return MagicMock(
+        loaded_papers={"p1"},
+        paper_metadata={},
+        documents={},
+        metadata_fields=["title"],
+        config={"embedding_batch_size": 1234},
+        has_gpu=False,
+        vector_store=MagicMock(),
+    )
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.paper_loader.add_papers_batch")
+def test_all_papers_loaded_returns_early(mock_batch, request):
+    """Test early return when all papers are already loaded."""
+    article_data = request.getfixturevalue("articles")
+    vector_store = request.getfixturevalue("mock_vector_store")
+    vector_store.loaded_papers = set(article_data.keys())
+
+    load_all_papers(
+        vector_store=vector_store,
+        articles=article_data,
+        call_id="test_call",
+        config={"embedding_batch_size": 1000},
+        has_gpu=False,
+    )
+
+    mock_batch.assert_not_called()
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.paper_loader.add_papers_batch")
+def test_skips_papers_without_pdf(mock_batch, request):
+    """Test that papers without PDF URLs are skipped."""
+    article_data = request.getfixturevalue("articles")
+    vector_store = request.getfixturevalue("mock_vector_store")
+    vector_store.loaded_papers = {"p2"}  # p1 not loaded, p3 has no pdf
+
+    load_all_papers(
+        vector_store=vector_store,
+        articles=article_data,
+        call_id="test_call",
+        config={"embedding_batch_size": 1000},
+        has_gpu=False,
+    )
+
+    assert mock_batch.call_count == 1
+    call_args = mock_batch.call_args[1]["papers_to_add"]
+    assert len(call_args) == 1
+    assert call_args[0][0] == "p1"
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.paper_loader.add_papers_batch")
+def test_gpu_parameters_used(mock_batch, request):
+    """Test GPU-based parameters are used if has_gpu is True."""
+    article_data = request.getfixturevalue("articles")
+    vector_store = request.getfixturevalue("mock_vector_store")
+    vector_store.loaded_papers = set()
+    vector_store.has_gpu = True
+
+    load_all_papers(
+        vector_store=vector_store,
+        articles=article_data,
+        call_id="gpu_call",
+        config={"embedding_batch_size": 2048},
+        has_gpu=True,
+    )
+
+    args = mock_batch.call_args[1]
+    assert args["has_gpu"] is True
+    assert args["batch_size"] == 2048
+    assert args["max_workers"] >= 4
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.paper_loader.add_papers_batch")
+def test_cpu_parameters_used(mock_batch, request):
+    """Test CPU-based parameters are used if has_gpu is False."""
+    article_data = request.getfixturevalue("articles")
+    vector_store = request.getfixturevalue("mock_vector_store")
+    vector_store.loaded_papers = set()
+    vector_store.has_gpu = False
+
+    load_all_papers(
+        vector_store=vector_store,
+        articles=article_data,
+        call_id="cpu_call",
+        config={"embedding_batch_size": 512},
+        has_gpu=False,
+    )
+
+    args = mock_batch.call_args[1]
+    assert args["has_gpu"] is False
+    assert args["batch_size"] == 512
+    assert args["max_workers"] >= 3

aiagents4pharma/talk2scholars/tests/test_pdf_rag_pipeline.py

@@ -0,0 +1,98 @@
+"""pdf rag pipeline tests."""
+
+from unittest.mock import MagicMock, patch
+import pytest
+from langchain_core.documents import Document
+
+from aiagents4pharma.talk2scholars.tools.pdf.utils.rag_pipeline import (
+    retrieve_and_rerank_chunks,
+)
+
+
+@pytest.fixture(name="base_config_fixture")
+def _base_config_fixture():
+    """Provides a config-like object for testing."""
+    config = MagicMock()
+    config.get.side_effect = lambda key, default=None: {
+        "initial_retrieval_k": 120,
+        "mmr_diversity": 0.7,
+    }.get(key, default)
+    config.top_k_chunks = 5
+    return config
+
+
+@pytest.fixture(name="mock_docs_fixture")
+def _mock_docs_fixture():
+    """Simulates PDF document chunks."""
+    return [
+        Document(page_content=f"chunk {i}", metadata={"paper_id": f"P{i % 2}"})
+        for i in range(10)
+    ]
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.rag_pipeline.rerank_chunks")
+@patch(
+    "aiagents4pharma.talk2scholars.tools.pdf.utils.rag_pipeline.retrieve_relevant_chunks"
+)
+def test_rag_pipeline_gpu_path(
+    mock_retrieve, mock_rerank, base_config_fixture, mock_docs_fixture
+):
+    """test RAG pipeline with GPU path."""
+    mock_retrieve.return_value = mock_docs_fixture
+    mock_rerank.return_value = mock_docs_fixture[:5]
+
+    result = retrieve_and_rerank_chunks(
+        vector_store=MagicMock(),
+        query="Explain AI.",
+        config=base_config_fixture,
+        call_id="gpu_test",
+        has_gpu=True,
+    )
+
+    assert result == mock_docs_fixture[:5]
+    mock_retrieve.assert_called_once()
+    mock_rerank.assert_called_once()
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.rag_pipeline.rerank_chunks")
+@patch(
+    "aiagents4pharma.talk2scholars.tools.pdf.utils.rag_pipeline.retrieve_relevant_chunks"
+)
+def test_rag_pipeline_cpu_path(
+    mock_retrieve, mock_rerank, base_config_fixture, mock_docs_fixture
+):
+    """rag pipeline with CPU path."""
+    mock_retrieve.return_value = mock_docs_fixture
+    mock_rerank.return_value = mock_docs_fixture[:5]
+
+    result = retrieve_and_rerank_chunks(
+        vector_store=MagicMock(),
+        query="Explain quantum physics.",
+        config=base_config_fixture,
+        call_id="cpu_test",
+        has_gpu=False,
+    )
+
+    assert result == mock_docs_fixture[:5]
+    mock_retrieve.assert_called_once()
+    mock_rerank.assert_called_once()
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.rag_pipeline.rerank_chunks")
+@patch(
+    "aiagents4pharma.talk2scholars.tools.pdf.utils.rag_pipeline.retrieve_relevant_chunks"
+)
+def test_rag_pipeline_empty_results(mock_retrieve, mock_rerank, base_config_fixture):
+    """rag pipeline with no results."""
+    mock_retrieve.return_value = []
+
+    result = retrieve_and_rerank_chunks(
+        vector_store=MagicMock(),
+        query="No match?",
+        config=base_config_fixture,
+        call_id="empty_test",
+        has_gpu=False,
+    )
+
+    assert result == []
+    mock_rerank.assert_not_called()

aiagents4pharma/talk2scholars/tests/test_pdf_retrieve_chunks.py

@@ -0,0 +1,197 @@
+"""retrieve_chunks for PDF tool tests"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+from langchain_core.documents import Document
+
+from aiagents4pharma.talk2scholars.tools.pdf.utils.retrieve_chunks import (
+    retrieve_relevant_chunks,
+    retrieve_relevant_chunks_with_scores,
+)
+
+
+@pytest.fixture
+def mock_vector_store():
+    """Fixture to simulate a vector store."""
+    return MagicMock()
+
+
+@pytest.fixture
+def mock_chunks():
+    """Fixture to simulate PDF chunks."""
+    return [
+        Document(page_content=f"chunk {i}", metadata={"paper_id": f"P{i%2}"})
+        for i in range(5)
+    ]
+
+
+@pytest.fixture
+def mock_scored_chunks():
+    """Fixture to simulate scored PDF chunks."""
+    return [
+        (Document(page_content=f"chunk {i}", metadata={}), score)
+        for i, score in enumerate([0.9, 0.8, 0.4, 0.95])
+    ]
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.retrieve_chunks.logger")
+def test_retrieve_chunks_cpu_success(mock_logger, request):
+    """Test retrieve_relevant_chunks with CPU path."""
+    vector_store = request.getfixturevalue("mock_vector_store")
+    chunks = request.getfixturevalue("mock_chunks")
+    vector_store.has_gpu = False
+    mock_logger.debug = MagicMock()
+    vector_store.max_marginal_relevance_search.return_value = chunks
+
+    results = retrieve_relevant_chunks(vector_store, query="AI", top_k=5)
+
+    assert results == chunks
+    vector_store.max_marginal_relevance_search.assert_called_once()
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.retrieve_chunks.logger")
+def test_retrieve_chunks_gpu_success(mock_logger, request):
+    """Test retrieve_relevant_chunks with GPU path."""
+    vector_store = request.getfixturevalue("mock_vector_store")
+    chunks = request.getfixturevalue("mock_chunks")
+    vector_store.has_gpu = True
+    mock_logger.debug = MagicMock()
+    vector_store.max_marginal_relevance_search.return_value = chunks
+
+    results = retrieve_relevant_chunks(vector_store, query="AI", top_k=5)
+
+    assert results == chunks
+    vector_store.max_marginal_relevance_search.assert_called_once()
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.retrieve_chunks.logger")
+def test_retrieve_chunks_with_filter(mock_logger, request):
+    """Test retrieve_relevant_chunks with paper_id filter."""
+    vector_store = request.getfixturevalue("mock_vector_store")
+    chunks = request.getfixturevalue("mock_chunks")
+    vector_store.has_gpu = False
+    mock_logger.debug = MagicMock()
+    vector_store.max_marginal_relevance_search.return_value = chunks
+
+    results = retrieve_relevant_chunks(
+        vector_store, query="filter test", paper_ids=["P1"], top_k=3
+    )
+    assert results == chunks
+    args, kwargs = vector_store.max_marginal_relevance_search.call_args
+    assert len(args) == 0
+    assert kwargs["filter"] == {"paper_id": ["P1"]}
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.retrieve_chunks.logger")
+def test_retrieve_chunks_no_vector_store(mock_logger):
+    """Test when vector store is None."""
+    result = retrieve_relevant_chunks(vector_store=None, query="irrelevant")
+    assert result == []
+    mock_logger.error.assert_called_with("Vector store is not initialized")
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.retrieve_chunks.logger")
+def test_retrieve_chunks_with_scores_no_vector_store(mock_logger):
+    """Test retrieve_relevant_chunks_with_scores when vector store is None."""
+    result = retrieve_relevant_chunks_with_scores(vector_store=None, query="none")
+    assert result == []
+    mock_logger.error.assert_called_with("Vector store is not initialized")
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.retrieve_chunks.logger")
+def test_retrieve_chunks_default_search_params(mock_logger, request):
+    """Test default search params used when not defined."""
+    vector_store = request.getfixturevalue("mock_vector_store")
+    chunks = request.getfixturevalue("mock_chunks")
+    vector_store.has_gpu = False
+    delattr(vector_store, "search_params")
+    vector_store.max_marginal_relevance_search.return_value = chunks
+
+    results = retrieve_relevant_chunks(
+        vector_store,
+        query="default search param test",
+        top_k=5,
+    )
+
+    assert results == chunks
+    mock_logger.debug.assert_any_call(
+        "Using default search parameters (no hardware optimization)"
+    )
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.retrieve_chunks.logger")
+def test_retrieve_chunks_with_scores_paper_filter(mock_logger, request):
+    """Test retrieve_relevant_chunks_with_scores applies paper_id filter."""
+    vector_store = request.getfixturevalue("mock_vector_store")
+    scored_chunks = request.getfixturevalue("mock_scored_chunks")
+    vector_store.similarity_search_with_score.return_value = scored_chunks
+    mock_logger.debug = MagicMock()
+
+    results = retrieve_relevant_chunks_with_scores(
+        vector_store=vector_store,
+        query="filtered score",
+        paper_ids=["P123"],
+        top_k=5,
+        score_threshold=0.0,
+    )
+
+    assert isinstance(results, list)
+    assert vector_store.similarity_search_with_score.call_args[1]["filter"] == {
+        "paper_id": ["P123"]
+    }
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.retrieve_chunks.logger")
+def test_retrieve_chunks_with_scores_gpu_debug(mock_logger, request):
+    """Test GPU debug log and correct return in retrieve_relevant_chunks_with_scores."""
+    vector_store = request.getfixturevalue("mock_vector_store")
+    scored_chunks = request.getfixturevalue("mock_scored_chunks")
+    vector_store.has_gpu = True
+    vector_store.similarity_search_with_score.return_value = scored_chunks
+    mock_logger.debug = MagicMock()
+
+    results = retrieve_relevant_chunks_with_scores(
+        vector_store=vector_store, query="gpu test", top_k=4, score_threshold=0.0
+    )
+
+    # Should return all scored_chunks since threshold=0.0
+    assert results == scored_chunks
+    mock_logger.debug.assert_called_with("GPU-accelerated similarity search enabled")
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.retrieve_chunks.logger")
+def test_retrieve_chunks_with_scores_cpu_debug(mock_logger, request):
+    """Test CPU debug log and correct return in retrieve_relevant_chunks_with_scores."""
+    vector_store = request.getfixturevalue("mock_vector_store")
+    scored_chunks = request.getfixturevalue("mock_scored_chunks")
+    vector_store.has_gpu = False
+    vector_store.similarity_search_with_score.return_value = scored_chunks
+    mock_logger.debug = MagicMock()
+
+    results = retrieve_relevant_chunks_with_scores(
+        vector_store=vector_store, query="cpu test", top_k=2, score_threshold=0.0
+    )
+
+    assert results == scored_chunks
+    mock_logger.debug.assert_called_with("Standard CPU similarity search")
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.retrieve_chunks.logger")
+def test_retrieve_chunks_with_scores_not_implemented(mock_logger, request):
+    """Test NotImplementedError path when similarity_search_with_score is missing."""
+    vector_store = request.getfixturevalue("mock_vector_store")
+    vector_store.has_gpu = True
+    # Remove the method to trigger NotImplementedError
+    if hasattr(vector_store, "similarity_search_with_score"):
+        delattr(vector_store, "similarity_search_with_score")
+    mock_logger.debug = MagicMock()
+
+    with pytest.raises(NotImplementedError) as excinfo:
+        retrieve_relevant_chunks_with_scores(
+            vector_store=vector_store, query="fail test", top_k=1, score_threshold=0.0
+        )
+    assert "Vector store does not support similarity_search_with_score" in str(
+        excinfo.value
+    )
+    mock_logger.debug.assert_called_with("GPU-accelerated similarity search enabled")