forked from Archives/langchain
feat: add Momento as a standard cache and chat message history provider (#5221)
# Add Momento as a standard cache and chat message history provider This PR adds Momento as a standard caching provider. Implements the interface, adds integration tests, and documentation. We also add Momento as a chat history message provider along with integration tests, and documentation. [Momento](https://www.gomomento.com/) is a fully serverless cache. Similar to S3 or DynamoDB, it requires zero configuration, infrastructure management, and is instantly available. Users sign up for free and get 50GB of data in/out for free every month. ## Before submitting ✅ We have added documentation, notebooks, and integration tests demonstrating usage. Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>searx_updates
Michael Landis
12 months ago
committed by
GitHub
parent
56ad56c812
commit
7047a2c1af
@ -0,0 +1,53 @@
|
||||
# Momento
|
||||
|
||||
This page covers how to use the [Momento](https://gomomento.com) ecosystem within LangChain.
|
||||
It is broken into two parts: installation and setup, and then references to specific Momento wrappers.
|
||||
|
||||
## Installation and Setup
|
||||
|
||||
- Sign up for a free account [here](https://docs.momentohq.com/getting-started) and get an auth token
|
||||
- Install the Momento Python SDK with `pip install momento`
|
||||
|
||||
## Wrappers
|
||||
|
||||
### Cache
|
||||
|
||||
The Cache wrapper allows for [Momento](https://gomomento.com) to be used as a serverless, distributed, low-latency cache for LLM prompts and responses.
|
||||
|
||||
#### Standard Cache
|
||||
|
||||
The standard cache is the go-to use case for [Momento](https://gomomento.com) users in any environment.
|
||||
|
||||
Import the cache as follows:
|
||||
|
||||
```python
|
||||
from langchain.cache import MomentoCache
|
||||
```
|
||||
|
||||
And set up like so:
|
||||
|
||||
```python
|
||||
from datetime import timedelta
|
||||
from momento import CacheClient, Configurations, CredentialProvider
|
||||
import langchain
|
||||
|
||||
# Instantiate the Momento client
|
||||
cache_client = CacheClient(
|
||||
Configurations.Laptop.v1(),
|
||||
CredentialProvider.from_environment_variable("MOMENTO_AUTH_TOKEN"),
|
||||
default_ttl=timedelta(days=1))
|
||||
|
||||
# Choose a Momento cache name of your choice
|
||||
cache_name = "langchain"
|
||||
|
||||
# Instantiate the LLM cache
|
||||
langchain.llm_cache = MomentoCache(cache_client, cache_name)
|
||||
```
|
||||
|
||||
### Memory
|
||||
|
||||
Momento can be used as a distributed memory store for LLMs.
|
||||
|
||||
#### Chat Message History Memory
|
||||
|
||||
See [this notebook](../modules/memory/examples/momento_chat_message_history.ipynb) for a walkthrough of how to use Momento as a memory store for chat message history.
|
||||
@ -0,0 +1,86 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "91c6a7ef",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Momento\n",
|
||||
"\n",
|
||||
"This notebook goes over how to use [Momento Cache](https://gomomento.com) to store chat message history using the `MomentoChatMessageHistory` class. See the Momento [docs](https://docs.momentohq.com/getting-started) for more detail on how to get set up with Momento.\n",
|
||||
"\n",
|
||||
"Note that, by default we will create a cache if one with the given name doesn't already exist.\n",
|
||||
"\n",
|
||||
"You'll need to get a Momento auth token to use this class. This can either be passed in to a momento.CacheClient if you'd like to instantiate that directly, as a named parameter `auth_token` to `MomentoChatMessageHistory.from_client_params`, or can just be set as an environment variable `MOMENTO_AUTH_TOKEN`."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"id": "d15e3302",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from datetime import timedelta\n",
|
||||
"\n",
|
||||
"from langchain.memory import MomentoChatMessageHistory\n",
|
||||
"\n",
|
||||
"session_id = \"foo\"\n",
|
||||
"cache_name = \"langchain\"\n",
|
||||
"ttl = timedelta(days=1),\n",
|
||||
"history = MomentoChatMessageHistory.from_client_params(\n",
|
||||
" session_id, \n",
|
||||
" cache_name,\n",
|
||||
" ttl,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"history.add_user_message(\"hi!\")\n",
|
||||
"\n",
|
||||
"history.add_ai_message(\"whats up?\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"id": "64fc465e",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[HumanMessage(content='hi!', additional_kwargs={}, example=False),\n",
|
||||
" AIMessage(content='whats up?', additional_kwargs={}, example=False)]"
|
||||
]
|
||||
},
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"history.messages"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.3"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@ -0,0 +1,200 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from datetime import timedelta
|
||||
from typing import TYPE_CHECKING, Any, Optional
|
||||
|
||||
from langchain.schema import (
|
||||
AIMessage,
|
||||
BaseChatMessageHistory,
|
||||
BaseMessage,
|
||||
HumanMessage,
|
||||
_message_to_dict,
|
||||
messages_from_dict,
|
||||
)
|
||||
from langchain.utils import get_from_env
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import momento
|
||||
|
||||
|
||||
def _ensure_cache_exists(cache_client: momento.CacheClient, cache_name: str) -> None:
|
||||
"""Create cache if it doesn't exist.
|
||||
|
||||
Raises:
|
||||
SdkException: Momento service or network error
|
||||
Exception: Unexpected response
|
||||
"""
|
||||
from momento.responses import CreateCache
|
||||
|
||||
create_cache_response = cache_client.create_cache(cache_name)
|
||||
if isinstance(create_cache_response, CreateCache.Success) or isinstance(
|
||||
create_cache_response, CreateCache.CacheAlreadyExists
|
||||
):
|
||||
return None
|
||||
elif isinstance(create_cache_response, CreateCache.Error):
|
||||
raise create_cache_response.inner_exception
|
||||
else:
|
||||
raise Exception(f"Unexpected response cache creation: {create_cache_response}")
|
||||
|
||||
|
||||
class MomentoChatMessageHistory(BaseChatMessageHistory):
|
||||
"""Chat message history cache that uses Momento as a backend.
|
||||
See https://gomomento.com/"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
session_id: str,
|
||||
cache_client: momento.CacheClient,
|
||||
cache_name: str,
|
||||
*,
|
||||
key_prefix: str = "message_store:",
|
||||
ttl: Optional[timedelta] = None,
|
||||
ensure_cache_exists: bool = True,
|
||||
):
|
||||
"""Instantiate a chat message history cache that uses Momento as a backend.
|
||||
|
||||
Note: to instantiate the cache client passed to MomentoChatMessageHistory,
|
||||
you must have a Momento account at https://gomomento.com/.
|
||||
|
||||
Args:
|
||||
session_id (str): The session ID to use for this chat session.
|
||||
cache_client (CacheClient): The Momento cache client.
|
||||
cache_name (str): The name of the cache to use to store the messages.
|
||||
key_prefix (str, optional): The prefix to apply to the cache key.
|
||||
Defaults to "message_store:".
|
||||
ttl (Optional[timedelta], optional): The TTL to use for the messages.
|
||||
Defaults to None, ie the default TTL of the cache will be used.
|
||||
ensure_cache_exists (bool, optional): Create the cache if it doesn't exist.
|
||||
Defaults to True.
|
||||
|
||||
Raises:
|
||||
ImportError: Momento python package is not installed.
|
||||
TypeError: cache_client is not of type momento.CacheClientObject
|
||||
"""
|
||||
try:
|
||||
from momento import CacheClient
|
||||
from momento.requests import CollectionTtl
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Could not import momento python package. "
|
||||
"Please install it with `pip install momento`."
|
||||
)
|
||||
if not isinstance(cache_client, CacheClient):
|
||||
raise TypeError("cache_client must be a momento.CacheClient object.")
|
||||
if ensure_cache_exists:
|
||||
_ensure_cache_exists(cache_client, cache_name)
|
||||
self.key = key_prefix + session_id
|
||||
self.cache_client = cache_client
|
||||
self.cache_name = cache_name
|
||||
if ttl is not None:
|
||||
self.ttl = CollectionTtl.of(ttl)
|
||||
else:
|
||||
self.ttl = CollectionTtl.from_cache_ttl()
|
||||
|
||||
@classmethod
|
||||
def from_client_params(
|
||||
cls,
|
||||
session_id: str,
|
||||
cache_name: str,
|
||||
ttl: timedelta,
|
||||
*,
|
||||
configuration: Optional[momento.config.Configuration] = None,
|
||||
auth_token: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> MomentoChatMessageHistory:
|
||||
"""Construct cache from CacheClient parameters."""
|
||||
try:
|
||||
from momento import CacheClient, Configurations, CredentialProvider
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Could not import momento python package. "
|
||||
"Please install it with `pip install momento`."
|
||||
)
|
||||
if configuration is None:
|
||||
configuration = Configurations.Laptop.v1()
|
||||
auth_token = auth_token or get_from_env("auth_token", "MOMENTO_AUTH_TOKEN")
|
||||
credentials = CredentialProvider.from_string(auth_token)
|
||||
cache_client = CacheClient(configuration, credentials, default_ttl=ttl)
|
||||
return cls(session_id, cache_client, cache_name, ttl=ttl, **kwargs)
|
||||
|
||||
@property
|
||||
def messages(self) -> list[BaseMessage]: # type: ignore[override]
|
||||
"""Retrieve the messages from Momento.
|
||||
|
||||
Raises:
|
||||
SdkException: Momento service or network error
|
||||
Exception: Unexpected response
|
||||
|
||||
Returns:
|
||||
list[BaseMessage]: List of cached messages
|
||||
"""
|
||||
from momento.responses import CacheListFetch
|
||||
|
||||
fetch_response = self.cache_client.list_fetch(self.cache_name, self.key)
|
||||
|
||||
if isinstance(fetch_response, CacheListFetch.Hit):
|
||||
items = [json.loads(m) for m in fetch_response.value_list_string]
|
||||
return messages_from_dict(items)
|
||||
elif isinstance(fetch_response, CacheListFetch.Miss):
|
||||
return []
|
||||
elif isinstance(fetch_response, CacheListFetch.Error):
|
||||
raise fetch_response.inner_exception
|
||||
else:
|
||||
raise Exception(f"Unexpected response: {fetch_response}")
|
||||
|
||||
def add_user_message(self, message: str) -> None:
|
||||
"""Store a user message in the cache.
|
||||
|
||||
Args:
|
||||
message (str): The message to store.
|
||||
"""
|
||||
self.__add_message(HumanMessage(content=message))
|
||||
|
||||
def add_ai_message(self, message: str) -> None:
|
||||
"""Store an AI message in the cache.
|
||||
|
||||
Args:
|
||||
message (str): The message to store.
|
||||
"""
|
||||
self.__add_message(AIMessage(content=message))
|
||||
|
||||
def __add_message(self, message: BaseMessage) -> None:
|
||||
"""Store a message in the cache.
|
||||
|
||||
Args:
|
||||
message (BaseMessage): The message object to store.
|
||||
|
||||
Raises:
|
||||
SdkException: Momento service or network error.
|
||||
Exception: Unexpected response.
|
||||
"""
|
||||
from momento.responses import CacheListPushBack
|
||||
|
||||
item = json.dumps(_message_to_dict(message))
|
||||
push_response = self.cache_client.list_push_back(
|
||||
self.cache_name, self.key, item, ttl=self.ttl
|
||||
)
|
||||
if isinstance(push_response, CacheListPushBack.Success):
|
||||
return None
|
||||
elif isinstance(push_response, CacheListPushBack.Error):
|
||||
raise push_response.inner_exception
|
||||
else:
|
||||
raise Exception(f"Unexpected response: {push_response}")
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Remove the session's messages from the cache.
|
||||
|
||||
Raises:
|
||||
SdkException: Momento service or network error.
|
||||
Exception: Unexpected response.
|
||||
"""
|
||||
from momento.responses import CacheDelete
|
||||
|
||||
delete_response = self.cache_client.delete(self.cache_name, self.key)
|
||||
if isinstance(delete_response, CacheDelete.Success):
|
||||
return None
|
||||
elif isinstance(delete_response, CacheDelete.Error):
|
||||
raise delete_response.inner_exception
|
||||
else:
|
||||
raise Exception(f"Unexpected response: {delete_response}")
|
||||
@ -0,0 +1,94 @@
|
||||
"""Test Momento cache functionality.
|
||||
|
||||
To run tests, set the environment variable MOMENTO_AUTH_TOKEN to a valid
|
||||
Momento auth token. This can be obtained by signing up for a free
|
||||
Momento account at https://gomomento.com/.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from datetime import timedelta
|
||||
from typing import Iterator
|
||||
|
||||
import pytest
|
||||
from momento import CacheClient, Configurations, CredentialProvider
|
||||
|
||||
import langchain
|
||||
from langchain.cache import MomentoCache
|
||||
from langchain.schema import Generation, LLMResult
|
||||
from tests.unit_tests.llms.fake_llm import FakeLLM
|
||||
|
||||
|
||||
def random_string() -> str:
|
||||
return str(uuid.uuid4())
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def momento_cache() -> Iterator[MomentoCache]:
|
||||
cache_name = f"langchain-test-cache-{random_string()}"
|
||||
client = CacheClient(
|
||||
Configurations.Laptop.v1(),
|
||||
CredentialProvider.from_environment_variable("MOMENTO_AUTH_TOKEN"),
|
||||
default_ttl=timedelta(seconds=30),
|
||||
)
|
||||
try:
|
||||
llm_cache = MomentoCache(client, cache_name)
|
||||
langchain.llm_cache = llm_cache
|
||||
yield llm_cache
|
||||
finally:
|
||||
client.delete_cache(cache_name)
|
||||
|
||||
|
||||
def test_invalid_ttl() -> None:
|
||||
client = CacheClient(
|
||||
Configurations.Laptop.v1(),
|
||||
CredentialProvider.from_environment_variable("MOMENTO_AUTH_TOKEN"),
|
||||
default_ttl=timedelta(seconds=30),
|
||||
)
|
||||
with pytest.raises(ValueError):
|
||||
MomentoCache(client, cache_name=random_string(), ttl=timedelta(seconds=-1))
|
||||
|
||||
|
||||
def test_momento_cache_miss(momento_cache: MomentoCache) -> None:
|
||||
llm = FakeLLM()
|
||||
stub_llm_output = LLMResult(generations=[[Generation(text="foo")]])
|
||||
assert llm.generate([random_string()]) == stub_llm_output
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"prompts, generations",
|
||||
[
|
||||
# Single prompt, single generation
|
||||
([random_string()], [[random_string()]]),
|
||||
# Single prompt, multiple generations
|
||||
([random_string()], [[random_string(), random_string()]]),
|
||||
# Single prompt, multiple generations
|
||||
([random_string()], [[random_string(), random_string(), random_string()]]),
|
||||
# Multiple prompts, multiple generations
|
||||
(
|
||||
[random_string(), random_string()],
|
||||
[[random_string()], [random_string(), random_string()]],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_momento_cache_hit(
|
||||
momento_cache: MomentoCache, prompts: list[str], generations: list[list[str]]
|
||||
) -> None:
|
||||
llm = FakeLLM()
|
||||
params = llm.dict()
|
||||
params["stop"] = None
|
||||
llm_string = str(sorted([(k, v) for k, v in params.items()]))
|
||||
|
||||
llm_generations = [
|
||||
[
|
||||
Generation(text=generation, generation_info=params)
|
||||
for generation in prompt_i_generations
|
||||
]
|
||||
for prompt_i_generations in generations
|
||||
]
|
||||
for prompt_i, llm_generations_i in zip(prompts, llm_generations):
|
||||
momento_cache.update(prompt_i, llm_string, llm_generations_i)
|
||||
|
||||
assert llm.generate(prompts) == LLMResult(
|
||||
generations=llm_generations, llm_output={}
|
||||
)
|
||||
@ -0,0 +1,70 @@
|
||||
"""Test Momento chat message history functionality.
|
||||
|
||||
To run tests, set the environment variable MOMENTO_AUTH_TOKEN to a valid
|
||||
Momento auth token. This can be obtained by signing up for a free
|
||||
Momento account at https://gomomento.com/.
|
||||
"""
|
||||
import json
|
||||
import uuid
|
||||
from datetime import timedelta
|
||||
from typing import Iterator
|
||||
|
||||
import pytest
|
||||
from momento import CacheClient, Configurations, CredentialProvider
|
||||
|
||||
from langchain.memory import ConversationBufferMemory
|
||||
from langchain.memory.chat_message_histories import MomentoChatMessageHistory
|
||||
from langchain.schema import _message_to_dict
|
||||
|
||||
|
||||
def random_string() -> str:
|
||||
return str(uuid.uuid4())
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def message_history() -> Iterator[MomentoChatMessageHistory]:
|
||||
cache_name = f"langchain-test-cache-{random_string()}"
|
||||
client = CacheClient(
|
||||
Configurations.Laptop.v1(),
|
||||
CredentialProvider.from_environment_variable("MOMENTO_AUTH_TOKEN"),
|
||||
default_ttl=timedelta(seconds=30),
|
||||
)
|
||||
try:
|
||||
chat_message_history = MomentoChatMessageHistory(
|
||||
session_id="my-test-session",
|
||||
cache_client=client,
|
||||
cache_name=cache_name,
|
||||
)
|
||||
yield chat_message_history
|
||||
finally:
|
||||
client.delete_cache(cache_name)
|
||||
|
||||
|
||||
def test_memory_empty_on_new_session(
|
||||
message_history: MomentoChatMessageHistory,
|
||||
) -> None:
|
||||
memory = ConversationBufferMemory(
|
||||
memory_key="foo", chat_memory=message_history, return_messages=True
|
||||
)
|
||||
assert memory.chat_memory.messages == []
|
||||
|
||||
|
||||
def test_memory_with_message_store(message_history: MomentoChatMessageHistory) -> None:
|
||||
memory = ConversationBufferMemory(
|
||||
memory_key="baz", chat_memory=message_history, return_messages=True
|
||||
)
|
||||
|
||||
# Add some messages to the memory store
|
||||
memory.chat_memory.add_ai_message("This is me, the AI")
|
||||
memory.chat_memory.add_user_message("This is me, the human")
|
||||
|
||||
# Verify that the messages are in the store
|
||||
messages = memory.chat_memory.messages
|
||||
messages_json = json.dumps([_message_to_dict(msg) for msg in messages])
|
||||
|
||||
assert "This is me, the AI" in messages_json
|
||||
assert "This is me, the human" in messages_json
|
||||
|
||||
# Verify clearing the store
|
||||
memory.chat_memory.clear()
|
||||
assert memory.chat_memory.messages == []
|
||||
Loading…
Reference in New Issue