Add caching to BaseChatModel (issue #1644) (#5089)

#  Add caching to BaseChatModel
Fixes #1644

(Sidenote: While testing, I noticed we have multiple implementations of
Fake LLMs, used for testing. I consolidated them.)

## Who can review?
Community members can review the PR once tests pass. Tag
maintainers/contributors who might be interested:
Models
- @hwchase17
- @agola11

Twitter: [@UmerHAdil](https://twitter.com/@UmerHAdil) | Discord:
RicChilligerDude#7589

---------

Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>

tests/unit_tests/agents/test_react.py

@@ -1,13 +1,12 @@
 """Unit tests for ReAct."""
 
-from typing import Any, List, Mapping, Optional, Union
+from typing import Union
 
 from langchain.agents.react.base import ReActChain, ReActDocstoreAgent
 from langchain.agents.tools import Tool
-from langchain.callbacks.manager import CallbackManagerForLLMRun
 from langchain.docstore.base import Docstore
 from langchain.docstore.document import Document
-from langchain.llms.base import LLM
+from langchain.llms.fake import FakeListLLM
 from langchain.prompts.prompt import PromptTemplate
 from langchain.schema import AgentAction
 
@@ -22,33 +21,6 @@ Made in 2022."""
 _FAKE_PROMPT = PromptTemplate(input_variables=["input"], template="{input}")
 
 
-class FakeListLLM(LLM):
-    """Fake LLM for testing that outputs elements of a list."""
-
-    responses: List[str]
-    i: int = -1
-
-    @property
-    def _llm_type(self) -> str:
-        """Return type of llm."""
-        return "fake_list"
-
-    def _call(
-        self,
-        prompt: str,
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> str:
-        """Increment counter, and then return response in that index."""
-        self.i += 1
-        return self.responses[self.i]
-
-    @property
-    def _identifying_params(self) -> Mapping[str, Any]:
-        return {}
-
-
 class FakeDocstore(Docstore):
     """Fake docstore for testing purposes."""
 

tests/unit_tests/llms/test_callbacks.py

@@ -1,17 +1,17 @@
 """Test LLM callbacks."""
+from langchain.chat_models.fake import FakeListChatModel
+from langchain.llms.fake import FakeListLLM
 from langchain.schema import HumanMessage
 from tests.unit_tests.callbacks.fake_callback_handler import (
     FakeCallbackHandler,
     FakeCallbackHandlerWithChatStart,
 )
-from tests.unit_tests.llms.fake_chat_model import FakeChatModel
-from tests.unit_tests.llms.fake_llm import FakeLLM
 
 
 def test_llm_with_callbacks() -> None:
     """Test LLM callbacks."""
     handler = FakeCallbackHandler()
-    llm = FakeLLM(callbacks=[handler], verbose=True)
+    llm = FakeListLLM(callbacks=[handler], verbose=True, responses=["foo"])
     output = llm("foo")
     assert output == "foo"
     assert handler.starts == 1
@@ -22,7 +22,9 @@ def test_llm_with_callbacks() -> None:
 def test_chat_model_with_v1_callbacks() -> None:
     """Test chat model callbacks fall back to on_llm_start."""
     handler = FakeCallbackHandler()
-    llm = FakeChatModel(callbacks=[handler], verbose=True)
+    llm = FakeListChatModel(
+        callbacks=[handler], verbose=True, responses=["fake response"]
+    )
     output = llm([HumanMessage(content="foo")])
     assert output.content == "fake response"
     assert handler.starts == 1
@@ -35,7 +37,9 @@ def test_chat_model_with_v1_callbacks() -> None:
 def test_chat_model_with_v2_callbacks() -> None:
     """Test chat model callbacks fall back to on_llm_start."""
     handler = FakeCallbackHandlerWithChatStart()
-    llm = FakeChatModel(callbacks=[handler], verbose=True)
+    llm = FakeListChatModel(
+        callbacks=[handler], verbose=True, responses=["fake response"]
+    )
     output = llm([HumanMessage(content="foo")])
     assert output.content == "fake response"
     assert handler.starts == 1

tests/unit_tests/test_cache.py

@@ -0,0 +1,146 @@
+"""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()]))