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

aiagents4pharma/talk2scholars/tests/test_pdf_singleton_manager.py

@@ -0,0 +1,156 @@
+"""Tests for singleton_manager: manages vector store connections and event loops."""
+
+from unittest.mock import MagicMock, patch
+import pytest
+from pymilvus.exceptions import MilvusException
+
+from aiagents4pharma.talk2scholars.tools.pdf.utils.get_vectorstore import (
+    get_vectorstore,
+)
+from aiagents4pharma.talk2scholars.tools.pdf.utils.singleton_manager import (
+    VectorstoreSingleton,
+)
+
+
+def test_singleton_instance_identity():
+    """Singleton should return the same instance."""
+    a = VectorstoreSingleton()
+    b = VectorstoreSingleton()
+    assert a is b
+
+
+@patch(
+    "aiagents4pharma.talk2scholars.tools.pdf.utils.singleton_manager.detect_nvidia_gpu"
+)
+def test_detect_gpu_once(mock_detect):
+    """Ensure GPU detection is cached."""
+    mock_detect.return_value = True
+    singleton = VectorstoreSingleton()
+    setattr(singleton, "_gpu_detected", None)
+
+    result = singleton.detect_gpu_once()
+    assert result is True
+
+    result2 = singleton.detect_gpu_once()
+    assert result2 is True
+
+    mock_detect.assert_called_once()
+
+
+def test_get_event_loop_reuses_existing():
+    """get_event_loop should return the same loop if it exists."""
+    singleton = VectorstoreSingleton()
+    loop1 = singleton.get_event_loop()
+    loop2 = singleton.get_event_loop()
+    assert loop1 is loop2
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.singleton_manager.connections")
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.singleton_manager.db")
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.singleton_manager.utility")
+def test_get_connection_creates_connection(_, mock_db, mock_conns):
+    """get_connection should create a new connection if none exists."""
+    singleton = VectorstoreSingleton()
+    mock_conns.has_connection.return_value = True
+    mock_db.list_database.return_value = []
+
+    conn_key = singleton.get_connection("localhost", 19530, "test_db")
+    assert conn_key == "default"
+    mock_conns.remove_connection.assert_called_once()
+    mock_conns.connect.assert_called_once()
+    mock_db.create_database.assert_called_once_with("test_db")
+    mock_db.using_database.assert_called_once_with("test_db")
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.singleton_manager.Milvus")
+def test_get_vector_store_creates_if_missing(mock_milvus):
+    """get_vector_store should create a new vector store if missing."""
+    singleton = VectorstoreSingleton()
+    setattr(singleton, "_vector_stores", {})
+    setattr(singleton, "_event_loops", {})
+
+    mock_embed = MagicMock()
+    connection_args = {"host": "localhost", "port": 19530}
+
+    vs = singleton.get_vector_store("collection1", mock_embed, connection_args)
+
+    vector_stores = getattr(singleton, "_vector_stores")
+    assert vs is vector_stores["collection1"]
+    assert "collection1" in vector_stores
+    mock_milvus.assert_called_once()
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.get_vectorstore.Vectorstore")
+def test_get_vectorstore_factory(mock_vectorstore_cls):
+    """get_vectorstore should reuse or create Vectorstore."""
+    mock_config = MagicMock()
+    mock_config.milvus.collection_name = "demo"
+    mock_config.milvus.embedding_dim = 768
+    mock_embed = MagicMock()
+
+    result1 = get_vectorstore(mock_embed, mock_config, force_new=True)
+    assert result1 == mock_vectorstore_cls.return_value
+
+    result2 = get_vectorstore(mock_embed, mock_config)
+    assert result2 == result1
+
+
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.get_vectorstore.Vectorstore")
+def test_get_vectorstore_force_new(mock_vectorstore_cls):
+    """get_vectorstore should return a new instance if force_new=True."""
+    mock_vs1 = MagicMock(name="Vectorstore1")
+    mock_vs2 = MagicMock(name="Vectorstore2")
+    mock_vectorstore_cls.side_effect = [mock_vs1, mock_vs2]
+
+    dummy_config = MagicMock()
+    dummy_config.milvus.collection_name = "my_test_collection"
+    dummy_config.milvus.embedding_dim = 768
+
+    vs1 = get_vectorstore(mock_vs1, dummy_config)
+    vs2 = get_vectorstore(mock_vs2, dummy_config, force_new=True)
+
+    assert vs1 is mock_vs1
+    assert vs2 is mock_vs2
+    assert vs1 != vs2
+
+
+@patch(
+    "aiagents4pharma.talk2scholars.tools.pdf.utils.singleton_manager.connections.connect"
+)
+@patch(
+    "aiagents4pharma.talk2scholars.tools.pdf.utils.singleton_manager.connections.has_connection"
+)
+@patch("aiagents4pharma.talk2scholars.tools.pdf.utils.singleton_manager.db")
+def test_get_connection_milvus_error(_, mock_has_connection, mock_connect):
+    """get_connection should raise MilvusException on connection failure."""
+    manager = VectorstoreSingleton()
+    setattr(manager, "_connections", {})
+
+    mock_has_connection.return_value = False
+    mock_connect.side_effect = MilvusException("Connection failed")
+
+    with pytest.raises(MilvusException, match="Connection failed"):
+        manager.get_connection("localhost", 19530, "test_db")
+
+
+def test_get_event_loop_creates_new_loop_on_closed():
+    """Ensure get_event_loop creates a new loop if current one is closed."""
+    manager = VectorstoreSingleton()
+    setattr(manager, "_event_loops", {})
+
+    mock_loop = MagicMock()
+    mock_loop.is_closed.return_value = True
+
+    with (
+        patch("asyncio.get_event_loop", return_value=mock_loop),
+        patch("asyncio.new_event_loop") as mock_new_loop,
+        patch("asyncio.set_event_loop") as mock_set_loop,
+    ):
+        new_loop = MagicMock()
+        mock_new_loop.return_value = new_loop
+
+        result_loop = manager.get_event_loop()
+
+        mock_new_loop.assert_called_once()
+        mock_set_loop.assert_called_once_with(new_loop)
+        assert result_loop == new_loop

aiagents4pharma/talk2scholars/tests/test_pdf_vector_normalization.py

@@ -0,0 +1,121 @@
+"""Unit tests for vector normalization utilities for GPU COSINE support."""
+
+import logging
+
+import pytest
+from langchain_core.embeddings import Embeddings
+
+from aiagents4pharma.talk2scholars.tools.pdf.utils import vector_normalization as vn
+
+
+def test_normalize_vector_nonzero():
+    """Test normalizing a non-zero vector."""
+    vec = [3.0, 4.0]
+    result = vn.normalize_vector(vec)
+    expected = [0.6, 0.8]
+    assert pytest.approx(result) == expected
+
+
+def test_normalize_vector_zero_logs_warning(caplog):
+    """Test normalizing a zero vector logs a warning."""
+    with caplog.at_level(logging.WARNING):
+        result = vn.normalize_vector([0.0, 0.0])
+        assert result == [0.0, 0.0]
+        assert "Zero vector encountered" in caplog.text
+
+
+def test_normalize_vectors_batch_empty():
+    """Test that an empty batch returns unchanged."""
+    result = vn.normalize_vectors_batch([])
+    assert result == []
+
+
+def test_normalize_vectors_batch_normal_case():
+    """Test batch normalization of valid vectors with equal dimensions."""
+    vectors = [[3, 4], [6, 8]]
+    result = vn.normalize_vectors_batch(vectors)
+    expected = [
+        [0.6, 0.8],
+        [0.6, 0.8],
+    ]
+    for r, e in zip(result, expected):
+        assert pytest.approx(r) == e
+
+
+def test_normalize_vectors_batch_with_zero_vector(caplog):
+    """Test that zero vectors are handled and logged."""
+    vectors = [[0.0, 0.0], [1.0, 0.0]]
+    with caplog.at_level(logging.WARNING):
+        result = vn.normalize_vectors_batch(vectors)
+        assert len(result) == 2
+        assert "zero vectors during batch normalization" in caplog.text
+        assert pytest.approx(result[1]) == [1.0, 0.0]
+
+
+class DummyEmbedding(Embeddings):
+    """A dummy embedding class for testing normalization wrapper."""
+
+    def __init__(self):
+        self.test_attr = "test"
+
+    def embed_documents(self, texts):
+        return [[3.0, 4.0] for _ in texts]
+
+    def embed_query(self, text):
+        return [3.0, 4.0]
+
+
+def test_normalizing_embeddings_embed_documents():
+    """Test that document embeddings are normalized."""
+    model = vn.NormalizingEmbeddings(DummyEmbedding())
+    result = model.embed_documents(["doc1", "doc2"])
+    assert len(result) == 2
+    assert pytest.approx(result[0]) == [0.6, 0.8]
+
+
+def test_normalizing_embeddings_embed_query():
+    """Test that query embeddings are normalized."""
+    model = vn.NormalizingEmbeddings(DummyEmbedding())
+    result = model.embed_query("query")
+    assert pytest.approx(result) == [0.6, 0.8]
+
+
+def test_normalizing_embeddings_passthrough():
+    """Test attribute delegation to base embedding model."""
+    dummy = DummyEmbedding()
+    model = vn.NormalizingEmbeddings(dummy)
+    assert model.test_attr == "test"
+
+
+@pytest.mark.parametrize(
+    "has_gpu,use_cosine,expected_log",
+    [
+        (True, True, "ENABLED"),
+        (False, True, "DISABLED"),
+        (True, False, "DISABLED"),
+        (False, False, "DISABLED"),
+    ],
+)
+def test_should_normalize_vectors_logging(has_gpu, use_cosine, expected_log, caplog):
+    """Test should_normalize_vectors decision logic and logging."""
+    with caplog.at_level(logging.INFO):
+        result = vn.should_normalize_vectors(has_gpu, use_cosine)
+        if has_gpu and use_cosine:
+            assert result is True
+        else:
+            assert result is False
+        assert expected_log in caplog.text
+
+
+def test_wrap_embedding_model_if_needed_enabled():
+    """Test that wrapping is applied when needed."""
+    base = DummyEmbedding()
+    wrapped = vn.wrap_embedding_model_if_needed(base, has_gpu=True, use_cosine=True)
+    assert isinstance(wrapped, vn.NormalizingEmbeddings)
+
+
+def test_wrap_embedding_model_if_needed_disabled():
+    """Test that original model is returned when normalization not needed."""
+    base = DummyEmbedding()
+    wrapped = vn.wrap_embedding_model_if_needed(base, has_gpu=False, use_cosine=True)
+    assert wrapped is base