mirror of https://github.com/hwchase17/langchain
Cassandra support for LLM cache (exact-match and semantic) (#9772)
This PR implements two new classes in the cache module: `CassandraCache` and `CassandraSemanticCache`, similar in structure and functionality to their Redis counterpart: providing a cache for the response to a (prompt, llm) pair. Integration tests are included. Moreover, linting and type checks are all passing on my machine. Dependencies: the `pyproject.toml` and `poetry.lock` have the newest version of cassIO (the very same as in the Cassandra vector store metadata PR, submitted as #9280). If I may suggest, this issue and #9280 might be reviewed together (as they bring the same poetry changes along), so I'm tagging @baskaryan who already helped out a little with poetry-related conflicts there. (Thank you!) I'd be happy to add a short notebook if this is deemed necessary (but it seems to me that, contrary e.g. to vector stores, caches are not covered in specific notebooks). Thank you! --------- Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>pull/10160/head
parent
16945c9922
commit
c9ff0ab2e9
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,100 @@
|
||||
"""Test Cassandra caches. Requires a running vector-capable Cassandra cluster."""
|
||||
import os
|
||||
import time
|
||||
from typing import Any, Iterator, Tuple
|
||||
|
||||
import pytest
|
||||
|
||||
import langchain
|
||||
from langchain.cache import CassandraCache, CassandraSemanticCache
|
||||
from langchain.schema import Generation, LLMResult
|
||||
from tests.integration_tests.vectorstores.fake_embeddings import FakeEmbeddings
|
||||
from tests.unit_tests.llms.fake_llm import FakeLLM
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def cassandra_connection() -> Iterator[Tuple[Any, str]]:
|
||||
from cassandra.cluster import Cluster
|
||||
|
||||
keyspace = "langchain_cache_test_keyspace"
|
||||
# get db connection
|
||||
if "CASSANDRA_CONTACT_POINTS" in os.environ:
|
||||
contact_points = os.environ["CONTACT_POINTS"].split(",")
|
||||
cluster = Cluster(contact_points)
|
||||
else:
|
||||
cluster = Cluster()
|
||||
#
|
||||
session = cluster.connect()
|
||||
# ensure keyspace exists
|
||||
session.execute(
|
||||
(
|
||||
f"CREATE KEYSPACE IF NOT EXISTS {keyspace} "
|
||||
f"WITH replication = {{'class': 'SimpleStrategy', 'replication_factor': 1}}"
|
||||
)
|
||||
)
|
||||
|
||||
yield (session, keyspace)
|
||||
|
||||
|
||||
def test_cassandra_cache(cassandra_connection: Tuple[Any, str]) -> None:
|
||||
session, keyspace = cassandra_connection
|
||||
cache = CassandraCache(session, keyspace)
|
||||
langchain.llm_cache = cache
|
||||
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"])
|
||||
print(output)
|
||||
expected_output = LLMResult(
|
||||
generations=[[Generation(text="fizz")]],
|
||||
llm_output={},
|
||||
)
|
||||
print(expected_output)
|
||||
assert output == expected_output
|
||||
cache.clear()
|
||||
|
||||
|
||||
def test_cassandra_cache_ttl(cassandra_connection: Tuple[Any, str]) -> None:
|
||||
session, keyspace = cassandra_connection
|
||||
cache = CassandraCache(session, keyspace, ttl_seconds=2)
|
||||
langchain.llm_cache = cache
|
||||
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")])
|
||||
expected_output = LLMResult(
|
||||
generations=[[Generation(text="fizz")]],
|
||||
llm_output={},
|
||||
)
|
||||
output = llm.generate(["foo"])
|
||||
assert output == expected_output
|
||||
time.sleep(2.5)
|
||||
# entry has expired away.
|
||||
output = llm.generate(["foo"])
|
||||
assert output != expected_output
|
||||
cache.clear()
|
||||
|
||||
|
||||
def test_cassandra_semantic_cache(cassandra_connection: Tuple[Any, str]) -> None:
|
||||
session, keyspace = cassandra_connection
|
||||
sem_cache = CassandraSemanticCache(session, keyspace, embedding=FakeEmbeddings())
|
||||
langchain.llm_cache = sem_cache
|
||||
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(["bar"]) # same embedding as 'foo'
|
||||
expected_output = LLMResult(
|
||||
generations=[[Generation(text="fizz")]],
|
||||
llm_output={},
|
||||
)
|
||||
assert output == expected_output
|
||||
# clear the cache
|
||||
sem_cache.clear()
|
||||
output = llm.generate(["bar"]) # 'fizz' is erased away now
|
||||
assert output != expected_output
|
||||
sem_cache.clear()
|
Loading…
Reference in New Issue