mirror of https://github.com/hwchase17/langchain
Upstash redis integration (#10871)
- **Description:** Introduced Upstash provider with following wrappers: UpstashRedisCache, UpstashRedisEntityStore, UpstashRedisChatMessageHistory, UpstashRedisStore - **Issue:** -, - **Dependencies:** upstash-redis python package is needed, - **Tag maintainer:** @baskaryan - **Twitter handle:** @BurakY744 --------- Co-authored-by: Burak Yılmaz <burakyilmaz@Buraks-MacBook-Pro.local> Co-authored-by: Bagatur <baskaryan@gmail.com>pull/11748/head
Burak Yılmaz
12 months ago
committed by
GitHub
parent
a9db2b0b92
commit
63e516c2b0
@ -0,0 +1,61 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Upstash Redis Chat Message History\n",
|
||||
"\n",
|
||||
"This notebook goes over how to use Upstash Redis to store chat message history."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from langchain.memory.chat_message_histories.upstash_redis import UpstashRedisChatMessageHistory\n",
|
||||
"\n",
|
||||
"URL = \"<UPSTASH_REDIS_REST_URL>\"\n",
|
||||
"TOKEN = \"<UPSTASH_REDIS_REST_TOKEN>\"\n",
|
||||
"\n",
|
||||
"history = UpstashRedisChatMessageHistory(url=URL, token=TOKEN, ttl=10, session_id=\"my-test-session\")\n",
|
||||
"\n",
|
||||
"history.add_user_message(\"hello llm!\")\n",
|
||||
"history.add_ai_message(\"hello user!\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"history.messages"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": ".venv",
|
||||
"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"
|
||||
},
|
||||
"orig_nbformat": 4
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
@ -0,0 +1,42 @@
|
||||
# Upstash Redis
|
||||
|
||||
Upstash offers developers serverless databases and messaging platforms to build powerful applications without having to worry about the operational complexity of running databases at scale.
|
||||
|
||||
This page covers how to use [Upstash Redis](https://upstash.com/redis) with LangChain.
|
||||
|
||||
## Installation and Setup
|
||||
- Upstash Redis Python SDK can be installed with `pip install upstash-redis`
|
||||
- A globally distributed, low-latency and highly available database can be created at the [Upstash Console](https://console.upstash.com)
|
||||
|
||||
|
||||
## Integrations
|
||||
All of Upstash-LangChain integrations are based on `upstash-redis` Python SDK being utilized as wrappers for LangChain.
|
||||
This SDK utilizes Upstash Redis DB by giving UPSTASH_REDIS_REST_URL and UPSTASH_REDIS_REST_TOKEN parameters from the console.
|
||||
One significant advantage of this is that, this SDK uses a REST API. This means, you can run this in serverless platforms, edge or any platform that does not support TCP connections.
|
||||
|
||||
|
||||
### Cache
|
||||
|
||||
[Upstash Redis](https://upstash.com/redis) can be used as a cache for LLM prompts and responses.
|
||||
|
||||
To import this cache:
|
||||
```python
|
||||
from langchain.cache import UpstashRedisCache
|
||||
```
|
||||
|
||||
To use with your LLMs:
|
||||
```python
|
||||
import langchain
|
||||
from upstash_redis import Redis
|
||||
|
||||
URL = "<UPSTASH_REDIS_REST_URL>"
|
||||
TOKEN = "<UPSTASH_REDIS_REST_TOKEN>"
|
||||
|
||||
langchain.llm_cache = UpstashRedisCache(redis_=Redis(url=URL, token=TOKEN))
|
||||
```
|
||||
|
||||
### Memory
|
||||
Upstash Redis can be used to persist LLM conversations.
|
||||
|
||||
#### Chat Message History Memory
|
||||
An example of Upstash Redis for caching conversation message history can be seen in [this notebook](/docs/integrations/memory/upstash_redis_chat_message_history.html).
|
||||
@ -0,0 +1,67 @@
|
||||
import json
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
from langchain.schema import (
|
||||
BaseChatMessageHistory,
|
||||
)
|
||||
from langchain.schema.messages import BaseMessage, _message_to_dict, messages_from_dict
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class UpstashRedisChatMessageHistory(BaseChatMessageHistory):
|
||||
"""Chat message history stored in an Upstash Redis database."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
session_id: str,
|
||||
url: str = "",
|
||||
token: str = "",
|
||||
key_prefix: str = "message_store:",
|
||||
ttl: Optional[int] = None,
|
||||
):
|
||||
try:
|
||||
from upstash_redis import Redis
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Could not import upstash redis python package. "
|
||||
"Please install it with `pip install upstash_redis`."
|
||||
)
|
||||
|
||||
if url == "" or token == "":
|
||||
raise ValueError(
|
||||
"UPSTASH_REDIS_REST_URL and UPSTASH_REDIS_REST_TOKEN are needed."
|
||||
)
|
||||
|
||||
try:
|
||||
self.redis_client = Redis(url=url, token=token)
|
||||
except Exception:
|
||||
logger.error("Upstash Redis instance could not be initiated.")
|
||||
|
||||
self.session_id = session_id
|
||||
self.key_prefix = key_prefix
|
||||
self.ttl = ttl
|
||||
|
||||
@property
|
||||
def key(self) -> str:
|
||||
"""Construct the record key to use"""
|
||||
return self.key_prefix + self.session_id
|
||||
|
||||
@property
|
||||
def messages(self) -> List[BaseMessage]: # type: ignore
|
||||
"""Retrieve the messages from Upstash Redis"""
|
||||
_items = self.redis_client.lrange(self.key, 0, -1)
|
||||
items = [json.loads(m) for m in _items[::-1]]
|
||||
messages = messages_from_dict(items)
|
||||
return messages
|
||||
|
||||
def add_message(self, message: BaseMessage) -> None:
|
||||
"""Append the message to the record in Upstash Redis"""
|
||||
self.redis_client.lpush(self.key, json.dumps(_message_to_dict(message)))
|
||||
if self.ttl:
|
||||
self.redis_client.expire(self.key, self.ttl)
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear session memory from Upstash Redis"""
|
||||
self.redis_client.delete(self.key)
|
||||
@ -0,0 +1,119 @@
|
||||
from typing import Any, Iterator, List, Optional, Sequence, Tuple, cast
|
||||
|
||||
from langchain.schema import BaseStore
|
||||
|
||||
|
||||
class UpstashRedisStore(BaseStore[str, str]):
|
||||
"""BaseStore implementation using Upstash Redis as the underlying store."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
client: Any = None,
|
||||
url: Optional[str] = None,
|
||||
token: Optional[str] = None,
|
||||
ttl: Optional[int] = None,
|
||||
namespace: Optional[str] = None,
|
||||
) -> None:
|
||||
"""Initialize the UpstashRedisStore with HTTP API.
|
||||
|
||||
Must provide either an Upstash Redis client or a url.
|
||||
|
||||
Args:
|
||||
client: An Upstash Redis instance
|
||||
url: UPSTASH_REDIS_REST_URL
|
||||
token: UPSTASH_REDIS_REST_TOKEN
|
||||
ttl: time to expire keys in seconds if provided,
|
||||
if None keys will never expire
|
||||
namespace: if provided, all keys will be prefixed with this namespace
|
||||
"""
|
||||
try:
|
||||
from upstash_redis import Redis
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"UpstashRedisStore requires the upstash_redis library to be installed. "
|
||||
"pip install upstash_redis"
|
||||
) from e
|
||||
|
||||
if client and url:
|
||||
raise ValueError(
|
||||
"Either an Upstash Redis client or a url must be provided, not both."
|
||||
)
|
||||
|
||||
if client:
|
||||
if not isinstance(client, Redis):
|
||||
raise TypeError(
|
||||
f"Expected Upstash Redis client, got {type(client).__name__}."
|
||||
)
|
||||
_client = client
|
||||
else:
|
||||
if not url or not token:
|
||||
raise ValueError(
|
||||
"Either an Upstash Redis client or url and token must be provided."
|
||||
)
|
||||
_client = Redis(url=url, token=token)
|
||||
|
||||
self.client = _client
|
||||
|
||||
if not isinstance(ttl, int) and ttl is not None:
|
||||
raise TypeError(f"Expected int or None, got {type(ttl)} instead.")
|
||||
|
||||
self.ttl = ttl
|
||||
self.namespace = namespace
|
||||
|
||||
def _get_prefixed_key(self, key: str) -> str:
|
||||
"""Get the key with the namespace prefix.
|
||||
|
||||
Args:
|
||||
key (str): The original key.
|
||||
|
||||
Returns:
|
||||
str: The key with the namespace prefix.
|
||||
"""
|
||||
delimiter = "/"
|
||||
if self.namespace:
|
||||
return f"{self.namespace}{delimiter}{key}"
|
||||
return key
|
||||
|
||||
def mget(self, keys: Sequence[str]) -> List[Optional[str]]:
|
||||
"""Get the values associated with the given keys."""
|
||||
|
||||
keys = [self._get_prefixed_key(key) for key in keys]
|
||||
return cast(
|
||||
List[Optional[str]],
|
||||
self.client.mget(*keys),
|
||||
)
|
||||
|
||||
def mset(self, key_value_pairs: Sequence[Tuple[str, str]]) -> None:
|
||||
"""Set the given key-value pairs."""
|
||||
for key, value in key_value_pairs:
|
||||
self.client.set(self._get_prefixed_key(key), value, ex=self.ttl)
|
||||
|
||||
def mdelete(self, keys: Sequence[str]) -> None:
|
||||
"""Delete the given keys."""
|
||||
_keys = [self._get_prefixed_key(key) for key in keys]
|
||||
self.client.delete(*_keys)
|
||||
|
||||
def yield_keys(self, *, prefix: Optional[str] = None) -> Iterator[str]:
|
||||
"""Yield keys in the store."""
|
||||
if prefix:
|
||||
pattern = self._get_prefixed_key(prefix)
|
||||
else:
|
||||
pattern = self._get_prefixed_key("*")
|
||||
|
||||
cursor, keys = self.client.scan(0, match=pattern)
|
||||
for key in keys:
|
||||
if self.namespace:
|
||||
relative_key = key[len(self.namespace) + 1 :]
|
||||
yield relative_key
|
||||
else:
|
||||
yield key
|
||||
|
||||
while cursor != 0:
|
||||
cursor, keys = self.client.scan(cursor, match=pattern)
|
||||
for key in keys:
|
||||
if self.namespace:
|
||||
relative_key = key[len(self.namespace) + 1 :]
|
||||
yield relative_key
|
||||
else:
|
||||
yield key
|
||||
@ -0,0 +1,91 @@
|
||||
"""Test Upstash Redis cache functionality."""
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
import langchain
|
||||
from langchain.cache import UpstashRedisCache
|
||||
from langchain.schema import Generation, LLMResult
|
||||
from tests.unit_tests.llms.fake_chat_model import FakeChatModel
|
||||
from tests.unit_tests.llms.fake_llm import FakeLLM
|
||||
|
||||
URL = "<UPSTASH_REDIS_REST_URL>"
|
||||
TOKEN = "<UPSTASH_REDIS_REST_TOKEN>"
|
||||
|
||||
|
||||
def random_string() -> str:
|
||||
return str(uuid.uuid4())
|
||||
|
||||
|
||||
@pytest.mark.requires("upstash_redis")
|
||||
def test_redis_cache_ttl() -> None:
|
||||
from upstash_redis import Redis
|
||||
|
||||
langchain.llm_cache = UpstashRedisCache(redis_=Redis(url=URL, token=TOKEN), ttl=1)
|
||||
langchain.llm_cache.update("foo", "bar", [Generation(text="fizz")])
|
||||
key = langchain.llm_cache._key("foo", "bar")
|
||||
assert langchain.llm_cache.redis.pttl(key) > 0
|
||||
|
||||
|
||||
@pytest.mark.requires("upstash_redis")
|
||||
def test_redis_cache() -> None:
|
||||
from upstash_redis import Redis
|
||||
|
||||
langchain.llm_cache = UpstashRedisCache(redis_=Redis(url=URL, token=TOKEN), ttl=1)
|
||||
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"])
|
||||
expected_output = LLMResult(
|
||||
generations=[[Generation(text="fizz")]],
|
||||
llm_output={},
|
||||
)
|
||||
assert output == expected_output
|
||||
|
||||
lookup_output = langchain.llm_cache.lookup("foo", llm_string)
|
||||
if lookup_output and len(lookup_output) > 0:
|
||||
assert lookup_output == expected_output.generations[0]
|
||||
|
||||
langchain.llm_cache.clear()
|
||||
output = llm.generate(["foo"])
|
||||
|
||||
assert output != expected_output
|
||||
langchain.llm_cache.redis.flushall()
|
||||
|
||||
|
||||
def test_redis_cache_multi() -> None:
|
||||
from upstash_redis import Redis
|
||||
|
||||
langchain.llm_cache = UpstashRedisCache(redis_=Redis(url=URL, token=TOKEN), ttl=1)
|
||||
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"), Generation(text="Buzz")]
|
||||
)
|
||||
output = llm.generate(
|
||||
["foo"]
|
||||
) # foo and bar will have the same embedding produced by FakeEmbeddings
|
||||
expected_output = LLMResult(
|
||||
generations=[[Generation(text="fizz"), Generation(text="Buzz")]],
|
||||
llm_output={},
|
||||
)
|
||||
assert output == expected_output
|
||||
# clear the cache
|
||||
langchain.llm_cache.clear()
|
||||
|
||||
|
||||
@pytest.mark.requires("upstash_redis")
|
||||
def test_redis_cache_chat() -> None:
|
||||
from upstash_redis import Redis
|
||||
|
||||
langchain.llm_cache = UpstashRedisCache(redis_=Redis(url=URL, token=TOKEN), ttl=1)
|
||||
llm = FakeChatModel()
|
||||
params = llm.dict()
|
||||
params["stop"] = None
|
||||
with pytest.warns():
|
||||
llm.predict("foo")
|
||||
langchain.llm_cache.redis.flushall()
|
||||
@ -0,0 +1,38 @@
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from langchain.memory import ConversationBufferMemory
|
||||
from langchain.memory.chat_message_histories.upstash_redis import (
|
||||
UpstashRedisChatMessageHistory,
|
||||
)
|
||||
from langchain.schema.messages import _message_to_dict
|
||||
|
||||
URL = "<UPSTASH_REDIS_REST_URL>"
|
||||
TOKEN = "<UPSTASH_REDIS_REST_TOKEN>"
|
||||
|
||||
|
||||
@pytest.mark.requires("upstash_redis")
|
||||
def test_memory_with_message_store() -> None:
|
||||
"""Test the memory with a message store."""
|
||||
# setup Upstash Redis as a message store
|
||||
message_history = UpstashRedisChatMessageHistory(
|
||||
url=URL, token=TOKEN, ttl=10, session_id="my-test-session"
|
||||
)
|
||||
memory = ConversationBufferMemory(
|
||||
memory_key="baz", chat_memory=message_history, return_messages=True
|
||||
)
|
||||
|
||||
# add some messages
|
||||
memory.chat_memory.add_ai_message("This is me, the AI")
|
||||
memory.chat_memory.add_user_message("This is me, the human")
|
||||
|
||||
# get the message history from the memory store and turn it into a json
|
||||
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
|
||||
|
||||
# remove the record from Redis, so the next test run won't pick it up
|
||||
memory.chat_memory.clear()
|
||||
@ -0,0 +1,95 @@
|
||||
"""Implement integration tests for Redis storage."""
|
||||
|
||||
import pytest
|
||||
from upstash_redis import Redis
|
||||
|
||||
from langchain.storage.upstash_redis import UpstashRedisStore
|
||||
|
||||
pytest.importorskip("upstash_redis")
|
||||
|
||||
URL = "<UPSTASH_REDIS_REST_URL>"
|
||||
TOKEN = "<UPSTASH_REDIS_REST_TOKEN>"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def redis_client() -> Redis:
|
||||
"""Yield redis client."""
|
||||
from upstash_redis import Redis
|
||||
|
||||
# This fixture flushes the database!
|
||||
|
||||
client = Redis(url=URL, token=TOKEN)
|
||||
try:
|
||||
client.ping()
|
||||
except Exception:
|
||||
pytest.skip("Ping request failed. Verify that credentials are correct.")
|
||||
|
||||
client.flushdb()
|
||||
return client
|
||||
|
||||
|
||||
def test_mget(redis_client: Redis) -> None:
|
||||
store = UpstashRedisStore(client=redis_client, ttl=None)
|
||||
keys = ["key1", "key2"]
|
||||
redis_client.mset({"key1": "value1", "key2": "value2"})
|
||||
result = store.mget(keys)
|
||||
assert result == ["value1", "value2"]
|
||||
|
||||
|
||||
def test_mset(redis_client: Redis) -> None:
|
||||
store = UpstashRedisStore(client=redis_client, ttl=None)
|
||||
key_value_pairs = [("key1", "value1"), ("key2", "value2")]
|
||||
store.mset(key_value_pairs)
|
||||
result = redis_client.mget("key1", "key2")
|
||||
assert result == ["value1", "value2"]
|
||||
|
||||
|
||||
def test_mdelete(redis_client: Redis) -> None:
|
||||
"""Test that deletion works as expected."""
|
||||
store = UpstashRedisStore(client=redis_client, ttl=None)
|
||||
keys = ["key1", "key2"]
|
||||
redis_client.mset({"key1": "value1", "key2": "value2"})
|
||||
store.mdelete(keys)
|
||||
result = redis_client.mget(*keys)
|
||||
assert result == [None, None]
|
||||
|
||||
|
||||
def test_yield_keys(redis_client: Redis) -> None:
|
||||
store = UpstashRedisStore(client=redis_client, ttl=None)
|
||||
redis_client.mset({"key1": "value2", "key2": "value2"})
|
||||
assert sorted(store.yield_keys()) == ["key1", "key2"]
|
||||
assert sorted(store.yield_keys(prefix="key*")) == ["key1", "key2"]
|
||||
assert sorted(store.yield_keys(prefix="lang*")) == []
|
||||
|
||||
|
||||
def test_namespace(redis_client: Redis) -> None:
|
||||
store = UpstashRedisStore(client=redis_client, ttl=None, namespace="meow")
|
||||
key_value_pairs = [("key1", "value1"), ("key2", "value2")]
|
||||
store.mset(key_value_pairs)
|
||||
|
||||
cursor, all_keys = redis_client.scan(0)
|
||||
while cursor != 0:
|
||||
cursor, keys = redis_client.scan(cursor)
|
||||
if len(keys) != 0:
|
||||
all_keys.extend(keys)
|
||||
|
||||
assert sorted(all_keys) == [
|
||||
"meow/key1",
|
||||
"meow/key2",
|
||||
]
|
||||
|
||||
store.mdelete(["key1"])
|
||||
|
||||
cursor, all_keys = redis_client.scan(0, match="*")
|
||||
while cursor != 0:
|
||||
cursor, keys = redis_client.scan(cursor, match="*")
|
||||
if len(keys) != 0:
|
||||
all_keys.extend(keys)
|
||||
|
||||
assert sorted(all_keys) == [
|
||||
"meow/key2",
|
||||
]
|
||||
|
||||
assert list(store.yield_keys()) == ["key2"]
|
||||
assert list(store.yield_keys(prefix="key*")) == ["key2"]
|
||||
assert list(store.yield_keys(prefix="key1")) == []
|
||||
@ -0,0 +1,8 @@
|
||||
"""Light weight unit test that attempts to import UpstashRedisStore.
|
||||
"""
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.requires("upstash_redis")
|
||||
def test_import_storage() -> None:
|
||||
from langchain.storage.upstash_redis import UpstashRedisStore # noqa
|
||||
Loading…
Reference in New Issue