From ad6b5f84e50d87317ed2444b5bdbeda2e73f8782 Mon Sep 17 00:00:00 2001
From: Eugene Yurtsev <eyurtsev@gmail.com>
Date: Tue, 23 Apr 2024 11:10:11 -0400
Subject: [PATCH] community[patch],core[minor]: Move in memory cache
 implementation to core (#20753)

This PR moves the InMemoryCache implementation from community to core.
---
 libs/core/langchain_core/caches.py            | 36 ++++++++++++++++++-
 libs/langchain/langchain/cache.py             |  2 +-
 .../tests/unit_tests/llms/test_base.py        |  3 +-
 libs/langchain/tests/unit_tests/test_cache.py |  3 +-
 4 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/libs/core/langchain_core/caches.py b/libs/core/langchain_core/caches.py
index 4c494c4fbc849..4e9e2993a0428 100644
--- a/libs/core/langchain_core/caches.py
+++ b/libs/core/langchain_core/caches.py
@@ -22,7 +22,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, Optional, Sequence
+from typing import Any, Dict, Optional, Sequence, Tuple
 
 from langchain_core.outputs import Generation
 from langchain_core.runnables import run_in_executor
@@ -105,3 +105,37 @@ async def aupdate(
     async def aclear(self, **kwargs: Any) -> None:
         """Clear cache that can take additional keyword arguments."""
         return await run_in_executor(None, self.clear, **kwargs)
+
+
+class InMemoryCache(BaseCache):
+    """Cache that stores things in memory."""
+
+    def __init__(self) -> None:
+        """Initialize with empty cache."""
+        self._cache: Dict[Tuple[str, str], RETURN_VAL_TYPE] = {}
+
+    def lookup(self, prompt: str, llm_string: str) -> Optional[RETURN_VAL_TYPE]:
+        """Look up based on prompt and llm_string."""
+        return self._cache.get((prompt, llm_string), None)
+
+    def update(self, prompt: str, llm_string: str, return_val: RETURN_VAL_TYPE) -> None:
+        """Update cache based on prompt and llm_string."""
+        self._cache[(prompt, llm_string)] = return_val
+
+    def clear(self, **kwargs: Any) -> None:
+        """Clear cache."""
+        self._cache = {}
+
+    async def alookup(self, prompt: str, llm_string: str) -> Optional[RETURN_VAL_TYPE]:
+        """Look up based on prompt and llm_string."""
+        return self.lookup(prompt, llm_string)
+
+    async def aupdate(
+        self, prompt: str, llm_string: str, return_val: RETURN_VAL_TYPE
+    ) -> None:
+        """Update cache based on prompt and llm_string."""
+        self.update(prompt, llm_string, return_val)
+
+    async def aclear(self, **kwargs: Any) -> None:
+        """Clear cache."""
+        self.clear()
diff --git a/libs/langchain/langchain/cache.py b/libs/langchain/langchain/cache.py
index d0badba70c79f..7d84971d16fa9 100644
--- a/libs/langchain/langchain/cache.py
+++ b/libs/langchain/langchain/cache.py
@@ -18,7 +18,6 @@
 )
 
 __all__ = [
-    "InMemoryCache",
     "FullLLMCache",
     "SQLAlchemyCache",
     "SQLiteCache",
@@ -27,6 +26,7 @@
     "RedisSemanticCache",
     "GPTCache",
     "MomentoCache",
+    "InMemoryCache",
     "CassandraCache",
     "CassandraSemanticCache",
     "FullMd5LLMCache",
diff --git a/libs/langchain/tests/unit_tests/llms/test_base.py b/libs/langchain/tests/unit_tests/llms/test_base.py
index 37d9b802ed275..1b19b88ee7209 100644
--- a/libs/langchain/tests/unit_tests/llms/test_base.py
+++ b/libs/langchain/tests/unit_tests/llms/test_base.py
@@ -6,9 +6,10 @@
 except ImportError:
     from sqlalchemy.ext.declarative import declarative_base
 
+from langchain_core.caches import InMemoryCache
 from langchain_core.outputs import Generation, LLMResult
 
-from langchain.cache import InMemoryCache, SQLAlchemyCache
+from langchain.cache import SQLAlchemyCache
 from langchain.globals import get_llm_cache, set_llm_cache
 from langchain.llms.base import __all__
 from tests.unit_tests.llms.fake_llm import FakeLLM
diff --git a/libs/langchain/tests/unit_tests/test_cache.py b/libs/langchain/tests/unit_tests/test_cache.py
index 42cc6c36b6847..13e27c8e024ad 100644
--- a/libs/langchain/tests/unit_tests/test_cache.py
+++ b/libs/langchain/tests/unit_tests/test_cache.py
@@ -6,6 +6,7 @@
 from _pytest.fixtures import FixtureRequest
 from langchain_community.chat_models import FakeListChatModel
 from langchain_community.llms import FakeListLLM
+from langchain_core.caches import InMemoryCache
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.language_models.llms import BaseLLM
 from langchain_core.load import dumps
@@ -14,7 +15,7 @@
 from sqlalchemy import create_engine
 from sqlalchemy.orm import Session
 
-from langchain.cache import InMemoryCache, SQLAlchemyCache
+from langchain.cache import SQLAlchemyCache
 from langchain.globals import get_llm_cache, set_llm_cache