test_rate_limiting.py — langchain Source File

Source Code

import time

import pytest
from blockbuster import BlockBuster

from langchain_core.caches import InMemoryCache
from langchain_core.language_models import GenericFakeChatModel
from langchain_core.load import dumps
from langchain_core.rate_limiters import InMemoryRateLimiter


@pytest.fixture(autouse=True)
def deactivate_blockbuster(blockbuster: BlockBuster) -> None:
    # Deactivate BlockBuster to not disturb the rate limiter timings
    blockbuster.deactivate()


def test_rate_limit_invoke() -> None:
    """Add rate limiter."""
    model = GenericFakeChatModel(
        messages=iter(["hello", "world"]),
        rate_limiter=InMemoryRateLimiter(
            requests_per_second=20,
            check_every_n_seconds=0.1,
            max_bucket_size=10,
            # At 20 requests per second we see a refresh every 0.05 seconds
        ),
    )
    tic = time.time()
    model.invoke("foo")
    toc = time.time()
    # Should be larger than check every n seconds since the token bucket starts
    # with 0 tokens.
    assert 0.10 < toc - tic < 0.15

    tic = time.time()
    model.invoke("foo")
    toc = time.time()
    # Second time we check the model, we should have 1 extra token
    # since the sleep time is 0.1 seconds
    assert 0.00 < toc - tic < 0.10


async def test_rate_limit_ainvoke() -> None:
    """Add rate limiter."""
    model = GenericFakeChatModel(
        messages=iter(["hello", "world", "!"]),
        rate_limiter=InMemoryRateLimiter(
            requests_per_second=20,
            check_every_n_seconds=0.1,
            max_bucket_size=10,
            # At 20 requests per second we see a refresh every 0.05 seconds
        ),
    )
    tic = time.time()
    await model.ainvoke("foo")
    toc = time.time()
    # Should be larger than check every n seconds since the token bucket starts
    # with 0 tokens.
    assert 0.1 < toc - tic < 0.2
// ... (211 more lines)