unstructured-io--unstructured/test_unstructured/embed/test_voyageai.py

from unittest.mock import Mock

from unstructured.documents.elements import Text
from unstructured.embed.voyageai import VoyageAIEmbeddingConfig, VoyageAIEmbeddingEncoder


def test_embed_documents_does_not_break_element_to_dict(mocker):
    # Mocked client with the desired behavior for embed_documents
    embed_response = Mock()
    embed_response.embeddings = [[1], [2]]
    mock_client = mocker.MagicMock()
    mock_client.embed.return_value = embed_response
    mock_client.tokenize.return_value = [[1], [1]]  # Mock token counts

    # Mock get_client to return our mock_client
    mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mock_client)

    encoder = VoyageAIEmbeddingEncoder(
        config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3-large")
    )
    elements = encoder.embed_documents(
        elements=[Text("This is sentence 1"), Text("This is sentence 2")],
    )
    assert len(elements) == 2
    assert elements[0].to_dict()["text"] == "This is sentence 1"
    assert elements[1].to_dict()["text"] == "This is sentence 2"


def test_embed_documents_voyage_3_5(mocker):
    """Test embedding with voyage-3.5 model."""
    embed_response = Mock()
    embed_response.embeddings = [[1.0] * 1024, [2.0] * 1024]
    mock_client = mocker.MagicMock()
    mock_client.embed.return_value = embed_response
    mock_client.tokenize.return_value = [[1, 2, 3], [1, 2]]  # Mock token counts

    mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mock_client)

    encoder = VoyageAIEmbeddingEncoder(
        config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5")
    )
    elements = encoder.embed_documents(
        elements=[Text("Test document 1"), Text("Test document 2")],
    )
    assert len(elements) == 2
    assert len(elements[0].embeddings) == 1024
    assert len(elements[1].embeddings) == 1024


def test_embed_documents_voyage_3_5_lite(mocker):
    """Test embedding with voyage-3.5-lite model."""
    embed_response = Mock()
    embed_response.embeddings = [[1.0] * 512, [2.0] * 512, [3.0] * 512]
    mock_client = mocker.MagicMock()
    mock_client.embed.return_value = embed_response
    mock_client.tokenize.return_value = [[1], [1], [1]]  # Mock token counts

    mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mock_client)

    encoder = VoyageAIEmbeddingEncoder(
        config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5-lite")
    )
    elements = encoder.embed_documents(
        elements=[Text("Test 1"), Text("Test 2"), Text("Test 3")],
    )
    assert len(elements) == 3
    assert all(len(e.embeddings) == 512 for e in elements)


def test_embed_documents_contextual_model(mocker):
    """Test embedding with voyage-context-3 model."""
    # Mock contextualized_embed response
    contextualized_response = Mock()
    result_item = Mock()
    result_item.embeddings = [[1.0] * 1024, [2.0] * 1024]
    contextualized_response.results = [result_item]

    mock_client = mocker.MagicMock()
    mock_client.contextualized_embed.return_value = contextualized_response
    mock_client.tokenize.return_value = [[1, 2], [1, 2, 3]]  # Mock token counts

    mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mock_client)

    encoder = VoyageAIEmbeddingEncoder(
        config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-context-3")
    )
    elements = encoder.embed_documents(
        elements=[Text("Context document 1"), Text("Context document 2")],
    )
    assert len(elements) == 2
    assert len(elements[0].embeddings) == 1024
    assert len(elements[1].embeddings) == 1024
    # Verify contextualized_embed was called
    mock_client.contextualized_embed.assert_called_once()


def test_count_tokens(mocker):
    """Test token counting functionality."""
    mock_client = mocker.MagicMock()
    mock_client.tokenize.return_value = [[1, 2], [1, 2, 3, 4, 5]]  # Different token counts

    mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mock_client)

    encoder = VoyageAIEmbeddingEncoder(
        config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5")
    )
    texts = ["short text", "this is a longer text with more tokens"]
    token_counts = encoder.count_tokens(texts)

    assert len(token_counts) == 2
    assert token_counts[0] == 2
    assert token_counts[1] == 5


def test_count_tokens_empty_list(mocker):
    """Test token counting with empty list."""
    mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mocker.MagicMock())

    encoder = VoyageAIEmbeddingEncoder(
        config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5")
    )
    token_counts = encoder.count_tokens([])
    assert token_counts == []


def test_get_token_limit(mocker):
    """Test getting token limit for different models."""
    mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mocker.MagicMock())

    # Test voyage-3.5 model
    config = VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5")
    assert config.get_token_limit() == 320_000

    # Test voyage-3.5-lite model
    config_lite = VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5-lite")
    assert config_lite.get_token_limit() == 1_000_000

    # Test context model
    config_context = VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-context-3")
    assert config_context.get_token_limit() == 32_000

    # Test voyage-2 model
    config_v2 = VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-2")
    assert config_v2.get_token_limit() == 320_000

    # Test unknown model (should use default)
    config_unknown = VoyageAIEmbeddingConfig(api_key="api_key", model_name="unknown-model")
    assert config_unknown.get_token_limit() == 120_000


def test_is_context_model(mocker):
    """Test the _is_context_model helper method."""
    mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mocker.MagicMock())

    # Test with context model
    encoder_context = VoyageAIEmbeddingEncoder(
        config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-context-3")
    )
    assert encoder_context._is_context_model() is True

    # Test with regular model
    encoder_regular = VoyageAIEmbeddingEncoder(
        config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5")
    )
    assert encoder_regular._is_context_model() is False


def test_build_batches_with_token_limits(mocker):
    """Test that batching respects token limits."""
    mock_client = mocker.MagicMock()
    # Simulate different token counts for each text
    mock_client.tokenize.return_value = [[1] * 10, [1] * 20, [1] * 15, [1] * 25]

    mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mock_client)

    encoder = VoyageAIEmbeddingEncoder(
        config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-2")
    )
    texts = ["text1", "text2", "text3", "text4"]
    batches = list(encoder._build_batches(texts, mock_client))

    # Should create at least one batch
    assert len(batches) >= 1
    # Total texts should be preserved
    total_texts = sum(len(batch) for batch in batches)
    assert total_texts == len(texts)


def test_embed_query(mocker):
    """Test embedding a single query."""
    embed_response = Mock()
    embed_response.embeddings = [[1.0] * 1024]
    mock_client = mocker.MagicMock()
    mock_client.embed.return_value = embed_response

    mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mock_client)

    encoder = VoyageAIEmbeddingEncoder(
        config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5")
    )
    embedding = encoder.embed_query("test query")

    assert len(embedding) == 1024
    # Verify embed was called with input_type="query"
    mock_client.embed.assert_called_once()
    call_kwargs = mock_client.embed.call_args[1]
    assert call_kwargs["input_type"] == "query"


def test_embed_documents_with_output_dimension(mocker):
    """Test embedding with custom output dimension."""
    embed_response = Mock()
    embed_response.embeddings = [[1.0] * 512, [2.0] * 512]
    mock_client = mocker.MagicMock()
    mock_client.embed.return_value = embed_response
    mock_client.tokenize.return_value = [[1], [1]]

    mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mock_client)

    encoder = VoyageAIEmbeddingEncoder(
        config=VoyageAIEmbeddingConfig(
            api_key="api_key", model_name="voyage-3.5", output_dimension=512
        )
    )
    elements = encoder.embed_documents(
        elements=[Text("Test 1"), Text("Test 2")],
    )
    assert len(elements) == 2
    # Verify output_dimension was passed
    call_kwargs = mock_client.embed.call_args[1]
    assert call_kwargs["output_dimension"] == 512


def test_embed_documents_empty_list(mocker):
    """Test embedding empty list of documents."""
    mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mocker.MagicMock())

    encoder = VoyageAIEmbeddingEncoder(
        config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5")
    )
    elements = encoder.embed_documents(elements=[])
    assert elements == []