"""Test caching for LLMs and ChatModels."""
from typing import Dict, Generator, List, Union

import pytest
from _pytest.fixtures import FixtureRequest
from sqlalchemy import create_engine
from sqlalchemy.orm import Session

import langchain
from langchain.cache import (
    InMemoryCache,
    SQLAlchemyCache,
)
from langchain.chat_models import FakeListChatModel
from langchain.chat_models.base import BaseChatModel, dumps
from langchain.llms import FakeListLLM
from langchain.llms.base import BaseLLM
from langchain.schema import (
    AIMessage,
    BaseMessage,
    ChatGeneration,
    Generation,
    HumanMessage,
)


def get_sqlite_cache() -> SQLAlchemyCache:
    return SQLAlchemyCache(engine=create_engine("sqlite://"))


CACHE_OPTIONS = [
    InMemoryCache,
    get_sqlite_cache,
]


@pytest.fixture(autouse=True, params=CACHE_OPTIONS)
def set_cache_and_teardown(request: FixtureRequest) -> Generator[None, None, None]:
    # Will be run before each test
    cache_instance = request.param
    langchain.llm_cache = cache_instance()
    if langchain.llm_cache:
        langchain.llm_cache.clear()
    else:
        raise ValueError("Cache not set. This should never happen.")

    yield

    # Will be run after each test
    if langchain.llm_cache:
        langchain.llm_cache.clear()
    else:
        raise ValueError("Cache not set. This should never happen.")


def test_llm_caching() -> None:
    prompt = "How are you?"
    response = "Test response"
    cached_response = "Cached test response"
    llm = FakeListLLM(responses=[response])
    if langchain.llm_cache:
        langchain.llm_cache.update(
            prompt=prompt,
            llm_string=create_llm_string(llm),
            return_val=[Generation(text=cached_response)],
        )
        assert llm(prompt) == cached_response
    else:
        raise ValueError(
            "The cache not set. This should never happen, as the pytest fixture "
            "`set_cache_and_teardown` always sets the cache."
        )


def test_old_sqlite_llm_caching() -> None:
    if isinstance(langchain.llm_cache, SQLAlchemyCache):
        prompt = "How are you?"
        response = "Test response"
        cached_response = "Cached test response"
        llm = FakeListLLM(responses=[response])
        items = [
            langchain.llm_cache.cache_schema(
                prompt=prompt,
                llm=create_llm_string(llm),
                response=cached_response,
                idx=0,
            )
        ]
        with Session(langchain.llm_cache.engine) as session, session.begin():
            for item in items:
                session.merge(item)
        assert llm(prompt) == cached_response


def test_chat_model_caching() -> None:
    prompt: List[BaseMessage] = [HumanMessage(content="How are you?")]
    response = "Test response"
    cached_response = "Cached test response"
    cached_message = AIMessage(content=cached_response)
    llm = FakeListChatModel(responses=[response])
    if langchain.llm_cache:
        langchain.llm_cache.update(
            prompt=dumps(prompt),
            llm_string=llm._get_llm_string(),
            return_val=[ChatGeneration(message=cached_message)],
        )
        result = llm(prompt)
        assert isinstance(result, AIMessage)
        assert result.content == cached_response
    else:
        raise ValueError(
            "The cache not set. This should never happen, as the pytest fixture "
            "`set_cache_and_teardown` always sets the cache."
        )


def test_chat_model_caching_params() -> None:
    prompt: List[BaseMessage] = [HumanMessage(content="How are you?")]
    response = "Test response"
    cached_response = "Cached test response"
    cached_message = AIMessage(content=cached_response)
    llm = FakeListChatModel(responses=[response])
    if langchain.llm_cache:
        langchain.llm_cache.update(
            prompt=dumps(prompt),
            llm_string=llm._get_llm_string(functions=[]),
            return_val=[ChatGeneration(message=cached_message)],
        )
        result = llm(prompt, functions=[])
        assert isinstance(result, AIMessage)
        assert result.content == cached_response
        result_no_params = llm(prompt)
        assert isinstance(result_no_params, AIMessage)
        assert result_no_params.content == response

    else:
        raise ValueError(
            "The cache not set. This should never happen, as the pytest fixture "
            "`set_cache_and_teardown` always sets the cache."
        )


def create_llm_string(llm: Union[BaseLLM, BaseChatModel]) -> str:
    _dict: Dict = llm.dict()
    _dict["stop"] = None
    return str(sorted([(k, v) for k, v in _dict.items()]))

langchain.llm_cache.clear langchain.schema.Generation langchain.chat_models.FakeListChatModel langchain.llms.FakeListLLM langchain.schema.ChatGeneration langchain.schema.HumanMessage langchain.schema.AIMessage langchain.chat_models.base.dumps
[((796, 846), 'pytest.fixture', 'pytest.fixture', ([], {'autouse': '(True)', 'params': 'CACHE_OPTIONS'}), '(autouse=True, params=CACHE_OPTIONS)\n', (810, 846), False, 'import pytest\n'), ((1524, 1557), 'langchain.llms.FakeListLLM', 'FakeListLLM', ([], {'responses': '[response]'}), '(responses=[response])\n', (1535, 1557), False, 'from langchain.llms import FakeListLLM\n'), ((2895, 2929), 'langchain.schema.AIMessage', 'AIMessage', ([], {'content': 'cached_response'}), '(content=cached_response)\n', (2904, 2929), False, 'from langchain.schema import AIMessage, BaseMessage, ChatGeneration, Generation, HumanMessage\n'), ((2940, 2979), 'langchain.chat_models.FakeListChatModel', 'FakeListChatModel', ([], {'responses': '[response]'}), '(responses=[response])\n', (2957, 2979), False, 'from langchain.chat_models import FakeListChatModel\n'), ((3728, 3762), 'langchain.schema.AIMessage', 'AIMessage', ([], {'content': 'cached_response'}), '(content=cached_response)\n', (3737, 3762), False, 'from langchain.schema import AIMessage, BaseMessage, ChatGeneration, Generation, HumanMessage\n'), ((3773, 3812), 'langchain.chat_models.FakeListChatModel', 'FakeListChatModel', ([], {'responses': '[response]'}), '(responses=[response])\n', (3790, 3812), False, 'from langchain.chat_models import FakeListChatModel\n'), ((1080, 1107), 'langchain.llm_cache.clear', 'langchain.llm_cache.clear', ([], {}), '()\n', (1105, 1107), False, 'import langchain\n'), ((1269, 1296), 'langchain.llm_cache.clear', 'langchain.llm_cache.clear', ([], {}), '()\n', (1294, 1296), False, 'import langchain\n'), ((2233, 2266), 'langchain.llms.FakeListLLM', 'FakeListLLM', ([], {'responses': '[response]'}), '(responses=[response])\n', (2244, 2266), False, 'from langchain.llms import FakeListLLM\n'), ((2760, 2796), 'langchain.schema.HumanMessage', 'HumanMessage', ([], {'content': '"""How are you?"""'}), "(content='How are you?')\n", (2772, 2796), False, 'from langchain.schema import AIMessage, BaseMessage, ChatGeneration, Generation, HumanMessage\n'), ((3593, 3629), 'langchain.schema.HumanMessage', 'HumanMessage', ([], {'content': '"""How are you?"""'}), "(content='How are you?')\n", (3605, 3629), False, 'from langchain.schema import AIMessage, BaseMessage, ChatGeneration, Generation, HumanMessage\n'), ((702, 728), 'sqlalchemy.create_engine', 'create_engine', (['"""sqlite://"""'], {}), "('sqlite://')\n", (715, 728), False, 'from sqlalchemy import create_engine\n'), ((2508, 2543), 'sqlalchemy.orm.Session', 'Session', (['langchain.llm_cache.engine'], {}), '(langchain.llm_cache.engine)\n', (2515, 2543), False, 'from sqlalchemy.orm import Session\n'), ((3063, 3076), 'langchain.chat_models.base.dumps', 'dumps', (['prompt'], {}), '(prompt)\n', (3068, 3076), False, 'from langchain.chat_models.base import BaseChatModel, dumps\n'), ((3896, 3909), 'langchain.chat_models.base.dumps', 'dumps', (['prompt'], {}), '(prompt)\n', (3901, 3909), False, 'from langchain.chat_models.base import BaseChatModel, dumps\n'), ((1720, 1752), 'langchain.schema.Generation', 'Generation', ([], {'text': 'cached_response'}), '(text=cached_response)\n', (1730, 1752), False, 'from langchain.schema import AIMessage, BaseMessage, ChatGeneration, Generation, HumanMessage\n'), ((3148, 3186), 'langchain.schema.ChatGeneration', 'ChatGeneration', ([], {'message': 'cached_message'}), '(message=cached_message)\n', (3162, 3186), False, 'from langchain.schema import AIMessage, BaseMessage, ChatGeneration, Generation, HumanMessage\n'), ((3993, 4031), 'langchain.schema.ChatGeneration', 'ChatGeneration', ([], {'message': 'cached_message'}), '(message=cached_message)\n', (4007, 4031), False, 'from langchain.schema import AIMessage, BaseMessage, ChatGeneration, Generation, HumanMessage\n')]