"""Test Upstash Redis cache functionality."""
import uuid

import pytest

import langchain
from langchain.cache import UpstashRedisCache
from langchain.schema import Generation, LLMResult
from tests.unit_tests.llms.fake_chat_model import FakeChatModel
from tests.unit_tests.llms.fake_llm import FakeLLM

URL = "<UPSTASH_REDIS_REST_URL>"
TOKEN = "<UPSTASH_REDIS_REST_TOKEN>"


def random_string() -> str:
    return str(uuid.uuid4())


@pytest.mark.requires("upstash_redis")
def test_redis_cache_ttl() -> None:
    from upstash_redis import Redis

    langchain.llm_cache = UpstashRedisCache(redis_=Redis(url=URL, token=TOKEN), ttl=1)
    langchain.llm_cache.update("foo", "bar", [Generation(text="fizz")])
    key = langchain.llm_cache._key("foo", "bar")
    assert langchain.llm_cache.redis.pttl(key) > 0


@pytest.mark.requires("upstash_redis")
def test_redis_cache() -> None:
    from upstash_redis import Redis

    langchain.llm_cache = UpstashRedisCache(redis_=Redis(url=URL, token=TOKEN), ttl=1)
    llm = FakeLLM()
    params = llm.dict()
    params["stop"] = None
    llm_string = str(sorted([(k, v) for k, v in params.items()]))
    langchain.llm_cache.update("foo", llm_string, [Generation(text="fizz")])
    output = llm.generate(["foo"])
    expected_output = LLMResult(
        generations=[[Generation(text="fizz")]],
        llm_output={},
    )
    assert output == expected_output

    lookup_output = langchain.llm_cache.lookup("foo", llm_string)
    if lookup_output and len(lookup_output) > 0:
        assert lookup_output == expected_output.generations[0]

    langchain.llm_cache.clear()
    output = llm.generate(["foo"])

    assert output != expected_output
    langchain.llm_cache.redis.flushall()


def test_redis_cache_multi() -> None:
    from upstash_redis import Redis

    langchain.llm_cache = UpstashRedisCache(redis_=Redis(url=URL, token=TOKEN), ttl=1)
    llm = FakeLLM()
    params = llm.dict()
    params["stop"] = None
    llm_string = str(sorted([(k, v) for k, v in params.items()]))
    langchain.llm_cache.update(
        "foo", llm_string, [Generation(text="fizz"), Generation(text="Buzz")]
    )
    output = llm.generate(
        ["foo"]
    )  # foo and bar will have the same embedding produced by FakeEmbeddings
    expected_output = LLMResult(
        generations=[[Generation(text="fizz"), Generation(text="Buzz")]],
        llm_output={},
    )
    assert output == expected_output
    # clear the cache
    langchain.llm_cache.clear()


@pytest.mark.requires("upstash_redis")
def test_redis_cache_chat() -> None:
    from upstash_redis import Redis

    langchain.llm_cache = UpstashRedisCache(redis_=Redis(url=URL, token=TOKEN), ttl=1)
    llm = FakeChatModel()
    params = llm.dict()
    params["stop"] = None
    with pytest.warns():
        llm.predict("foo")
    langchain.llm_cache.redis.flushall()

langchain.llm_cache.clear langchain.schema.Generation langchain.llm_cache.redis.flushall langchain.llm_cache.redis.pttl langchain.llm_cache._key langchain.llm_cache.lookup
[((436, 473), 'pytest.mark.requires', 'pytest.mark.requires', (['"""upstash_redis"""'], {}), "('upstash_redis')\n", (456, 473), False, 'import pytest\n'), ((809, 846), 'pytest.mark.requires', 'pytest.mark.requires', (['"""upstash_redis"""'], {}), "('upstash_redis')\n", (829, 846), False, 'import pytest\n'), ((2491, 2528), 'pytest.mark.requires', 'pytest.mark.requires', (['"""upstash_redis"""'], {}), "('upstash_redis')\n", (2511, 2528), False, 'import pytest\n'), ((716, 754), 'langchain.llm_cache._key', 'langchain.llm_cache._key', (['"""foo"""', '"""bar"""'], {}), "('foo', 'bar')\n", (740, 754), False, 'import langchain\n'), ((1013, 1022), 'tests.unit_tests.llms.fake_llm.FakeLLM', 'FakeLLM', ([], {}), '()\n', (1020, 1022), False, 'from tests.unit_tests.llms.fake_llm import FakeLLM\n'), ((1420, 1465), 'langchain.llm_cache.lookup', 'langchain.llm_cache.lookup', (['"""foo"""', 'llm_string'], {}), "('foo', llm_string)\n", (1446, 1465), False, 'import langchain\n'), ((1583, 1610), 'langchain.llm_cache.clear', 'langchain.llm_cache.clear', ([], {}), '()\n', (1608, 1610), False, 'import langchain\n'), ((1688, 1724), 'langchain.llm_cache.redis.flushall', 'langchain.llm_cache.redis.flushall', ([], {}), '()\n', (1722, 1724), False, 'import langchain\n'), ((1899, 1908), 'tests.unit_tests.llms.fake_llm.FakeLLM', 'FakeLLM', ([], {}), '()\n', (1906, 1908), False, 'from tests.unit_tests.llms.fake_llm import FakeLLM\n'), ((2460, 2487), 'langchain.llm_cache.clear', 'langchain.llm_cache.clear', ([], {}), '()\n', (2485, 2487), False, 'import langchain\n'), ((2700, 2715), 'tests.unit_tests.llms.fake_chat_model.FakeChatModel', 'FakeChatModel', ([], {}), '()\n', (2713, 2715), False, 'from tests.unit_tests.llms.fake_chat_model import FakeChatModel\n'), ((2822, 2858), 'langchain.llm_cache.redis.flushall', 'langchain.llm_cache.redis.flushall', ([], {}), '()\n', (2856, 2858), False, 'import langchain\n'), ((419, 431), 'uuid.uuid4', 'uuid.uuid4', ([], {}), '()\n', (429, 431), False, 'import uuid\n'), ((766, 801), 'langchain.llm_cache.redis.pttl', 'langchain.llm_cache.redis.pttl', (['key'], {}), '(key)\n', (796, 801), False, 'import langchain\n'), ((2775, 2789), 'pytest.warns', 'pytest.warns', ([], {}), '()\n', (2787, 2789), False, 'import pytest\n'), ((598, 625), 'upstash_redis.Redis', 'Redis', ([], {'url': 'URL', 'token': 'TOKEN'}), '(url=URL, token=TOKEN)\n', (603, 625), False, 'from upstash_redis import Redis\n'), ((680, 703), 'langchain.schema.Generation', 'Generation', ([], {'text': '"""fizz"""'}), "(text='fizz')\n", (690, 703), False, 'from langchain.schema import Generation, LLMResult\n'), ((967, 994), 'upstash_redis.Redis', 'Redis', ([], {'url': 'URL', 'token': 'TOKEN'}), '(url=URL, token=TOKEN)\n', (972, 994), False, 'from upstash_redis import Redis\n'), ((1190, 1213), 'langchain.schema.Generation', 'Generation', ([], {'text': '"""fizz"""'}), "(text='fizz')\n", (1200, 1213), False, 'from langchain.schema import Generation, LLMResult\n'), ((1853, 1880), 'upstash_redis.Redis', 'Redis', ([], {'url': 'URL', 'token': 'TOKEN'}), '(url=URL, token=TOKEN)\n', (1858, 1880), False, 'from upstash_redis import Redis\n'), ((2085, 2108), 'langchain.schema.Generation', 'Generation', ([], {'text': '"""fizz"""'}), "(text='fizz')\n", (2095, 2108), False, 'from langchain.schema import Generation, LLMResult\n'), ((2110, 2133), 'langchain.schema.Generation', 'Generation', ([], {'text': '"""Buzz"""'}), "(text='Buzz')\n", (2120, 2133), False, 'from langchain.schema import Generation, LLMResult\n'), ((2654, 2681), 'upstash_redis.Redis', 'Redis', ([], {'url': 'URL', 'token': 'TOKEN'}), '(url=URL, token=TOKEN)\n', (2659, 2681), False, 'from upstash_redis import Redis\n'), ((1306, 1329), 'langchain.schema.Generation', 'Generation', ([], {'text': '"""fizz"""'}), "(text='fizz')\n", (1316, 1329), False, 'from langchain.schema import Generation, LLMResult\n'), ((2316, 2339), 'langchain.schema.Generation', 'Generation', ([], {'text': '"""fizz"""'}), "(text='fizz')\n", (2326, 2339), False, 'from langchain.schema import Generation, LLMResult\n'), ((2341, 2364), 'langchain.schema.Generation', 'Generation', ([], {'text': '"""Buzz"""'}), "(text='Buzz')\n", (2351, 2364), False, 'from langchain.schema import Generation, LLMResult\n')]