test_embeddings_respects_token_limit() — langchain Function Reference

Source Code

libs/partners/openai/tests/unit_tests/embeddings/test_base.py lines 103–150

def test_embeddings_respects_token_limit() -> None:
    """Test that embeddings respect the 300k token per request limit."""
    # Create embeddings instance
    embeddings = OpenAIEmbeddings(
        model="text-embedding-ada-002", api_key=SecretStr("test-key")
    )

    call_counts = []

    def mock_create(**kwargs: Any) -> Mock:
        input_ = kwargs["input"]
        # Track how many tokens in this call
        if isinstance(input_, list):
            total_tokens = sum(
                len(t) if isinstance(t, list) else len(t.split()) for t in input_
            )
            call_counts.append(total_tokens)
            # Verify this call doesn't exceed limit
            assert total_tokens <= 300000, (
                f"Batch exceeded token limit: {total_tokens} tokens"
            )

        # Return mock response
        mock_response = Mock()
        mock_response.model_dump.return_value = {
            "data": [
                {"embedding": [0.1] * 1536}
                for _ in range(len(input_) if isinstance(input_, list) else 1)
            ]
        }
        return mock_response

    embeddings.client.create = mock_create

    # Create a scenario that would exceed 300k tokens in a single batch
    # with default chunk_size=1000
    # Simulate 500 texts with ~1000 tokens each = 500k tokens total
    large_texts = ["word " * 1000 for _ in range(500)]

    # This should not raise an error anymore
    embeddings.embed_documents(large_texts)

    # Verify we made multiple API calls to respect the limit
    assert len(call_counts) > 1, "Should have split into multiple batches"

    # Verify each call respected the limit
    for count in call_counts:
        assert count <= 300000, f"Batch exceeded limit: {count}"