mirror of
https://github.com/hwchase17/langchain.git
synced 2025-08-03 18:24:10 +00:00
couchbase: Add ttl support to caches & chat_message_history (#26214)
**Description:** Add support to delete documents automatically from the caches & chat message history by adding a new optional parameter, `ttl`. - [x] **Add tests and docs**: If you're adding a new integration, please include 1. a test for the integration, preferably unit tests that do not rely on network access, 2. an example notebook showing its use. It lives in `docs/docs/integrations` directory. - [x] **Lint and test**: Run `make format`, `make lint` and `make test` from the root of the package(s) you've modified. See contribution guidelines for more: https://python.langchain.com/docs/contributing/ --------- Co-authored-by: Nithish Raghunandanan <nithishr@users.noreply.github.com> Co-authored-by: Erick Friis <erick@langchain.dev>
This commit is contained in:
Nithish Raghunandanan
GitHub
parent
c6c508ee96
commit
2d21274bf6
@ -14,7 +14,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"execution_count": 2,
|
||||
"id": "88486f6f",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
@ -30,7 +30,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"execution_count": 3,
|
||||
"id": "10ad9224",
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
@ -2400,7 +2400,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"execution_count": 4,
|
||||
"id": "9b4764e4-c75f-4185-b326-524287a826be",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
@ -2430,7 +2430,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"execution_count": 5,
|
||||
"id": "4b5e73c5-92c1-4eab-84e2-77924ea9c123",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
@ -2452,7 +2452,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"execution_count": 6,
|
||||
"id": "db8d28cc-8d93-47b4-8326-57a29a06fb3c",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
@ -2483,7 +2483,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"execution_count": 7,
|
||||
"id": "b470dc81-2e7f-4743-9435-ce9071394eea",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
@ -2491,17 +2491,17 @@
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"CPU times: user 53 ms, sys: 29 ms, total: 82 ms\n",
|
||||
"Wall time: 84.2 ms\n"
|
||||
"CPU times: user 25.9 ms, sys: 15.3 ms, total: 41.3 ms\n",
|
||||
"Wall time: 144 ms\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"\"\\n\\nWhy couldn't the bicycle stand up by itself? Because it was two-tired!\""
|
||||
"\"\\n\\nWhy don't scientists trust atoms?\\n\\nBecause they make up everything.\""
|
||||
]
|
||||
},
|
||||
"execution_count": 5,
|
||||
"execution_count": 7,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
@ -2512,6 +2512,35 @@
|
||||
"llm.invoke(\"Tell me a joke\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "1dca39d8-233a-45ba-ad7d-0920dfbc4a50",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Specifying a Time to Live (TTL) for the Cached entries\n",
|
||||
"The Cached documents can be deleted after a specified time automatically by specifying a `ttl` parameter along with the initialization of the Cache."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"id": "a3edcc6f-2ccd-45ba-8ca4-f65b5e51b461",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from datetime import timedelta\n",
|
||||
"\n",
|
||||
"set_llm_cache(\n",
|
||||
" CouchbaseCache(\n",
|
||||
" cluster=cluster,\n",
|
||||
" bucket_name=BUCKET_NAME,\n",
|
||||
" scope_name=SCOPE_NAME,\n",
|
||||
" collection_name=COLLECTION_NAME,\n",
|
||||
" ttl=timedelta(minutes=5),\n",
|
||||
" )\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "43626f33-d184-4260-b641-c9341cef5842",
|
||||
@ -2523,7 +2552,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"execution_count": 9,
|
||||
"id": "6b470c03-d7fe-4270-89e1-638251619a53",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
@ -2653,7 +2682,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"execution_count": 10,
|
||||
"id": "ae0766c8-ea34-4604-b0dc-cf2bbe8077f4",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
@ -2679,7 +2708,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"execution_count": 11,
|
||||
"id": "a2e82743-10ea-4319-b43e-193475ae5449",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
@ -2703,7 +2732,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"execution_count": 12,
|
||||
"id": "c36f4e29-d872-4334-a1f1-0e6d10c5d9f2",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
@ -2725,6 +2754,38 @@
|
||||
"print(llm.invoke(\"What is the expected lifespan of a dog?\"))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f6f674fa-70b5-4cf9-a208-992aad2c3c89",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Specifying a Time to Live (TTL) for the Cached entries\n",
|
||||
"The Cached documents can be deleted after a specified time automatically by specifying a `ttl` parameter along with the initialization of the Cache."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 13,
|
||||
"id": "7e127d0a-5049-47e0-abf4-6424ad7d9fec",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from datetime import timedelta\n",
|
||||
"\n",
|
||||
"set_llm_cache(\n",
|
||||
" CouchbaseSemanticCache(\n",
|
||||
" cluster=cluster,\n",
|
||||
" embedding=embeddings,\n",
|
||||
" bucket_name=BUCKET_NAME,\n",
|
||||
" scope_name=SCOPE_NAME,\n",
|
||||
" collection_name=COLLECTION_NAME,\n",
|
||||
" index_name=INDEX_NAME,\n",
|
||||
" score_threshold=0.8,\n",
|
||||
" ttl=timedelta(minutes=5),\n",
|
||||
" )\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ae1f5e1c-085e-4998-9f2d-b5867d2c3d5b",
|
||||
@ -2802,7 +2863,7 @@
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.12.3"
|
||||
"version": "3.10.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
|
||||
@ -135,7 +135,8 @@
|
||||
"Optionally you can configure the following:\n",
|
||||
"- session_id_key: Field in the chat message documents to store the `session_id`\n",
|
||||
"- message_key: Field in the chat message documents to store the message content\n",
|
||||
"- create_index: Used to specify if the index needs to be created on the collection. By default, an index is created on the `message_key` and the `session_id_key` of the documents"
|
||||
"- create_index: Used to specify if the index needs to be created on the collection. By default, an index is created on the `message_key` and the `session_id_key` of the documents\n",
|
||||
"- ttl: Used to specify a time in `timedelta` to live for the documents after which they will get deleted automatically from the storage."
|
||||
]
|
||||
},
|
||||
{
|
||||
@ -181,6 +182,34 @@
|
||||
"message_history.messages"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "b993fe69-4462-4cb4-ad0a-eb31a1a4d7d9",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Specifying a Time to Live (TTL) for the Chat Messages\n",
|
||||
"The stored messages can be deleted after a specified time automatically by specifying a `ttl` parameter along with the initialization of the chat message history store."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"id": "d32d9302-de97-4319-a484-c83530bab508",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from langchain_couchbase.chat_message_histories import CouchbaseChatMessageHistory\n",
|
||||
"\n",
|
||||
"message_history = CouchbaseChatMessageHistory(\n",
|
||||
" cluster=cluster,\n",
|
||||
" bucket_name=BUCKET_NAME,\n",
|
||||
" scope_name=SCOPE_NAME,\n",
|
||||
" collection_name=COLLECTION_NAME,\n",
|
||||
" session_id=\"test-session\",\n",
|
||||
" ttl=timedelta(hours=24),\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "c8b942a7-93fa-4cd9-8414-d047135c2733",
|
||||
@ -192,7 +221,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"execution_count": 8,
|
||||
"id": "8a9f0d91-d1d6-481d-8137-ea11229f485a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
@ -209,7 +238,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"execution_count": 9,
|
||||
"id": "946d45aa-5a61-49ae-816b-1c3949c56d9a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
@ -228,7 +257,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"execution_count": 10,
|
||||
"id": "20dfd838-b549-42ed-b3ba-ac005f7e024c",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
@ -249,7 +278,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"execution_count": 11,
|
||||
"id": "17bd09f4-896d-433d-bb9a-369a06e7aa8a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
@ -260,17 +289,17 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 11,
|
||||
"execution_count": 12,
|
||||
"id": "4bda1096-2fc2-40d7-a046-0d5d8e3a8f75",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"AIMessage(content='Hello Bob! How can I assist you today?', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 22, 'total_tokens': 32}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-a0f8a29e-ddf4-4e06-a1fe-cf8c325a2b72-0', usage_metadata={'input_tokens': 22, 'output_tokens': 10, 'total_tokens': 32})"
|
||||
"AIMessage(content='Hello, Bob! How can I assist you today?', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 11, 'prompt_tokens': 22, 'total_tokens': 33}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-62e54e3d-db70-429d-9ee0-e5e8eb2489a1-0', usage_metadata={'input_tokens': 22, 'output_tokens': 11, 'total_tokens': 33})"
|
||||
]
|
||||
},
|
||||
"execution_count": 11,
|
||||
"execution_count": 12,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
@ -281,17 +310,17 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 12,
|
||||
"execution_count": 13,
|
||||
"id": "1cfb31da-51bb-4c5f-909a-b7118b0ae08d",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"AIMessage(content='Your name is Bob.', response_metadata={'token_usage': {'completion_tokens': 5, 'prompt_tokens': 43, 'total_tokens': 48}, 'model_name': 'gpt-4o-mini', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-f764a9eb-999e-4042-96b6-fe47b7ae4779-0', usage_metadata={'input_tokens': 43, 'output_tokens': 5, 'total_tokens': 48})"
|
||||
"AIMessage(content='Your name is Bob.', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 5, 'prompt_tokens': 44, 'total_tokens': 49}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-d84a570a-45f3-4931-814a-078761170bca-0', usage_metadata={'input_tokens': 44, 'output_tokens': 5, 'total_tokens': 49})"
|
||||
]
|
||||
},
|
||||
"execution_count": 12,
|
||||
"execution_count": 13,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
|
||||
@ -9,6 +9,7 @@ are duplicated in this utility from modules:
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
from datetime import timedelta
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
from couchbase.cluster import Cluster
|
||||
@ -87,6 +88,16 @@ def _loads_generations(generations_str: str) -> Union[RETURN_VAL_TYPE, None]:
|
||||
return None
|
||||
|
||||
|
||||
def _validate_ttl(ttl: Optional[timedelta]) -> None:
|
||||
"""Validate the time to live"""
|
||||
if not isinstance(ttl, timedelta):
|
||||
raise ValueError(f"ttl should be of type timedelta but was {type(ttl)}.")
|
||||
if ttl <= timedelta(seconds=0):
|
||||
raise ValueError(
|
||||
f"ttl must be greater than 0 but was {ttl.total_seconds()} seconds."
|
||||
)
|
||||
|
||||
|
||||
class CouchbaseCache(BaseCache):
|
||||
"""Couchbase LLM Cache
|
||||
LLM Cache that uses Couchbase as the backend
|
||||
@ -140,6 +151,7 @@ class CouchbaseCache(BaseCache):
|
||||
bucket_name: str,
|
||||
scope_name: str,
|
||||
collection_name: str,
|
||||
ttl: Optional[timedelta] = None,
|
||||
**kwargs: Dict[str, Any],
|
||||
) -> None:
|
||||
"""Initialize the Couchbase LLM Cache
|
||||
@ -149,6 +161,8 @@ class CouchbaseCache(BaseCache):
|
||||
scope_name (str): name of the scope in bucket to store documents in.
|
||||
collection_name (str): name of the collection in the scope to store
|
||||
documents in.
|
||||
ttl (Optional[timedelta]): TTL or time for the document to live in the cache
|
||||
After this time, the document will get deleted from the cache.
|
||||
"""
|
||||
if not isinstance(cluster, Cluster):
|
||||
raise ValueError(
|
||||
@ -162,6 +176,8 @@ class CouchbaseCache(BaseCache):
|
||||
self._scope_name = scope_name
|
||||
self._collection_name = collection_name
|
||||
|
||||
self._ttl = None
|
||||
|
||||
# Check if the bucket exists
|
||||
if not self._check_bucket_exists():
|
||||
raise ValueError(
|
||||
@ -185,6 +201,11 @@ class CouchbaseCache(BaseCache):
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
# Check if the time to live is provided and valid
|
||||
if ttl is not None:
|
||||
_validate_ttl(ttl)
|
||||
self._ttl = ttl
|
||||
|
||||
def lookup(self, prompt: str, llm_string: str) -> Optional[RETURN_VAL_TYPE]:
|
||||
"""Look up from cache based on prompt and llm_string."""
|
||||
try:
|
||||
@ -206,10 +227,16 @@ class CouchbaseCache(BaseCache):
|
||||
self.LLM: llm_string,
|
||||
self.RETURN_VAL: _dumps_generations(return_val),
|
||||
}
|
||||
document_key = self._generate_key(prompt, llm_string)
|
||||
try:
|
||||
if self._ttl:
|
||||
self._collection.upsert(
|
||||
key=self._generate_key(prompt, llm_string), value=doc
|
||||
key=document_key,
|
||||
value=doc,
|
||||
expiry=self._ttl,
|
||||
)
|
||||
else:
|
||||
self._collection.upsert(key=document_key, value=doc)
|
||||
except Exception:
|
||||
logger.error("Error updating cache")
|
||||
|
||||
@ -242,6 +269,7 @@ class CouchbaseSemanticCache(BaseCache, CouchbaseVectorStore):
|
||||
collection_name: str,
|
||||
index_name: str,
|
||||
score_threshold: Optional[float] = None,
|
||||
ttl: Optional[timedelta] = None,
|
||||
) -> None:
|
||||
"""Initialize the Couchbase LLM Cache
|
||||
Args:
|
||||
@ -253,6 +281,8 @@ class CouchbaseSemanticCache(BaseCache, CouchbaseVectorStore):
|
||||
documents in.
|
||||
index_name (str): name of the Search index to use.
|
||||
score_threshold (float): score threshold to use for filtering results.
|
||||
ttl (Optional[timedelta]): TTL or time for the document to live in the cache
|
||||
After this time, the document will get deleted from the cache.
|
||||
"""
|
||||
if not isinstance(cluster, Cluster):
|
||||
raise ValueError(
|
||||
@ -265,6 +295,7 @@ class CouchbaseSemanticCache(BaseCache, CouchbaseVectorStore):
|
||||
self._bucket_name = bucket_name
|
||||
self._scope_name = scope_name
|
||||
self._collection_name = collection_name
|
||||
self._ttl = None
|
||||
|
||||
# Check if the bucket exists
|
||||
if not self._check_bucket_exists():
|
||||
@ -291,6 +322,10 @@ class CouchbaseSemanticCache(BaseCache, CouchbaseVectorStore):
|
||||
|
||||
self.score_threshold = score_threshold
|
||||
|
||||
if ttl is not None:
|
||||
_validate_ttl(ttl)
|
||||
self._ttl = ttl
|
||||
|
||||
# Initialize the vector store
|
||||
super().__init__(
|
||||
cluster=cluster,
|
||||
@ -334,6 +369,7 @@ class CouchbaseSemanticCache(BaseCache, CouchbaseVectorStore):
|
||||
self.RETURN_VAL: _dumps_generations(return_val),
|
||||
}
|
||||
],
|
||||
ttl=self._ttl,
|
||||
)
|
||||
except Exception:
|
||||
logger.error("Error updating cache")
|
||||
|
||||
@ -3,7 +3,8 @@
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Sequence
|
||||
from datetime import timedelta
|
||||
from typing import Any, Dict, List, Optional, Sequence
|
||||
|
||||
from couchbase.cluster import Cluster
|
||||
from langchain_core.chat_history import BaseChatMessageHistory
|
||||
@ -22,6 +23,16 @@ DEFAULT_INDEX_NAME = "LANGCHAIN_CHAT_HISTORY"
|
||||
DEFAULT_BATCH_SIZE = 100
|
||||
|
||||
|
||||
def _validate_ttl(ttl: Optional[timedelta]) -> None:
|
||||
"""Validate the time to live"""
|
||||
if not isinstance(ttl, timedelta):
|
||||
raise ValueError(f"ttl should be of type timedelta but was {type(ttl)}.")
|
||||
if ttl <= timedelta(seconds=0):
|
||||
raise ValueError(
|
||||
f"ttl must be greater than 0 but was {ttl.total_seconds()} seconds."
|
||||
)
|
||||
|
||||
|
||||
class CouchbaseChatMessageHistory(BaseChatMessageHistory):
|
||||
"""Couchbase Chat Message History
|
||||
Chat message history that uses Couchbase as the storage
|
||||
@ -77,6 +88,7 @@ class CouchbaseChatMessageHistory(BaseChatMessageHistory):
|
||||
session_id_key: str = DEFAULT_SESSION_ID_KEY,
|
||||
message_key: str = DEFAULT_MESSAGE_KEY,
|
||||
create_index: bool = True,
|
||||
ttl: Optional[timedelta] = None,
|
||||
) -> None:
|
||||
"""Initialize the Couchbase Chat Message History
|
||||
Args:
|
||||
@ -92,6 +104,8 @@ class CouchbaseChatMessageHistory(BaseChatMessageHistory):
|
||||
message_key (str): name of the field to use for the messages
|
||||
Set to "message" by default.
|
||||
create_index (bool): create an index if True. Set to True by default.
|
||||
ttl (timedelta): time to live for the documents in the collection.
|
||||
When set, the documents are automatically deleted after the ttl expires.
|
||||
"""
|
||||
if not isinstance(cluster, Cluster):
|
||||
raise ValueError(
|
||||
@ -104,6 +118,7 @@ class CouchbaseChatMessageHistory(BaseChatMessageHistory):
|
||||
self._bucket_name = bucket_name
|
||||
self._scope_name = scope_name
|
||||
self._collection_name = collection_name
|
||||
self._ttl = None
|
||||
|
||||
# Check if the bucket exists
|
||||
if not self._check_bucket_exists():
|
||||
@ -134,6 +149,10 @@ class CouchbaseChatMessageHistory(BaseChatMessageHistory):
|
||||
self._session_id = session_id
|
||||
self._ts_key = DEFAULT_TS_KEY
|
||||
|
||||
if ttl is not None:
|
||||
_validate_ttl(ttl)
|
||||
self._ttl = ttl
|
||||
|
||||
# Create an index if it does not exist if requested
|
||||
if create_index:
|
||||
index_fields = (
|
||||
@ -156,7 +175,19 @@ class CouchbaseChatMessageHistory(BaseChatMessageHistory):
|
||||
# get utc timestamp for ordering the messages
|
||||
timestamp = time.time()
|
||||
message_content = message_to_dict(message)
|
||||
|
||||
try:
|
||||
if self._ttl:
|
||||
self._collection.insert(
|
||||
document_key,
|
||||
value={
|
||||
self._message_key: message_content,
|
||||
self._session_id_key: self._session_id,
|
||||
self._ts_key: timestamp,
|
||||
},
|
||||
expiry=self._ttl,
|
||||
)
|
||||
else:
|
||||
self._collection.insert(
|
||||
document_key,
|
||||
value={
|
||||
@ -192,6 +223,9 @@ class CouchbaseChatMessageHistory(BaseChatMessageHistory):
|
||||
batch = messages_to_insert[i : i + batch_size]
|
||||
# Convert list of dictionaries to a single dictionary to insert
|
||||
insert_batch = {list(d.keys())[0]: list(d.values())[0] for d in batch}
|
||||
if self._ttl:
|
||||
self._collection.insert_multi(insert_batch, expiry=self._ttl)
|
||||
else:
|
||||
self._collection.insert_multi(insert_batch)
|
||||
except Exception as e:
|
||||
logger.error("Error adding messages: ", e)
|
||||
|
||||
@ -377,6 +377,9 @@ class CouchbaseVectorStore(VectorStore):
|
||||
if metadatas is None:
|
||||
metadatas = [{} for _ in texts]
|
||||
|
||||
# Check if TTL is provided
|
||||
ttl = kwargs.get("ttl", None)
|
||||
|
||||
embedded_texts = self._embedding_function.embed_documents(list(texts))
|
||||
|
||||
documents_to_insert = [
|
||||
@ -396,6 +399,10 @@ class CouchbaseVectorStore(VectorStore):
|
||||
for i in range(0, len(documents_to_insert), batch_size):
|
||||
batch = documents_to_insert[i : i + batch_size]
|
||||
try:
|
||||
# Insert with TTL if provided
|
||||
if ttl:
|
||||
result = self._collection.upsert_multi(batch[0], expiry=ttl)
|
||||
else:
|
||||
result = self._collection.upsert_multi(batch[0])
|
||||
if result.all_ok:
|
||||
doc_ids.extend(batch[0].keys())
|
||||
|
||||
170
libs/partners/couchbase/poetry.lock
generated
170
libs/partners/couchbase/poetry.lock
generated
@ -121,8 +121,27 @@ files = [
|
||||
{file = "annotated_types-0.7.0.tar.gz", hash = "sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89"},
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "anyio"
|
||||
version = "4.4.0"
|
||||
description = "High level compatibility layer for multiple asynchronous event loop implementations"
|
||||
optional = false
|
||||
python-versions = ">=3.8"
|
||||
files = [
|
||||
{file = "anyio-4.4.0-py3-none-any.whl", hash = "sha256:c1b2d8f46a8a812513012e1107cb0e68c17159a7a594208005a57dc776e1bdc7"},
|
||||
{file = "anyio-4.4.0.tar.gz", hash = "sha256:5aadc6a1bbb7cdb0bede386cac5e2940f5e2ff3aa20277e991cf028e0585ce94"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
typing-extensions = {version = ">=4.0.0", markers = "python_version < \"3.9\""}
|
||||
exceptiongroup = {version = ">=1.0.2", markers = "python_version < \"3.11\""}
|
||||
idna = ">=2.8"
|
||||
sniffio = ">=1.1"
|
||||
typing-extensions = {version = ">=4.1", markers = "python_version < \"3.11\""}
|
||||
|
||||
[package.extras]
|
||||
doc = ["Sphinx (>=7)", "packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx-rtd-theme"]
|
||||
test = ["anyio[trio]", "coverage[toml] (>=7)", "exceptiongroup (>=1.2.0)", "hypothesis (>=4.0)", "psutil (>=5.9)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "uvloop (>=0.17)"]
|
||||
trio = ["trio (>=0.23)"]
|
||||
|
||||
[[package]]
|
||||
name = "async-timeout"
|
||||
@ -497,6 +516,63 @@ files = [
|
||||
docs = ["Sphinx", "furo"]
|
||||
test = ["objgraph", "psutil"]
|
||||
|
||||
[[package]]
|
||||
name = "h11"
|
||||
version = "0.14.0"
|
||||
description = "A pure-Python, bring-your-own-I/O implementation of HTTP/1.1"
|
||||
optional = false
|
||||
python-versions = ">=3.7"
|
||||
files = [
|
||||
{file = "h11-0.14.0-py3-none-any.whl", hash = "sha256:e3fe4ac4b851c468cc8363d500db52c2ead036020723024a109d37346efaa761"},
|
||||
{file = "h11-0.14.0.tar.gz", hash = "sha256:8f19fbbe99e72420ff35c00b27a34cb9937e902a8b810e2c88300c6f0a3b699d"},
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "httpcore"
|
||||
version = "1.0.5"
|
||||
description = "A minimal low-level HTTP client."
|
||||
optional = false
|
||||
python-versions = ">=3.8"
|
||||
files = [
|
||||
{file = "httpcore-1.0.5-py3-none-any.whl", hash = "sha256:421f18bac248b25d310f3cacd198d55b8e6125c107797b609ff9b7a6ba7991b5"},
|
||||
{file = "httpcore-1.0.5.tar.gz", hash = "sha256:34a38e2f9291467ee3b44e89dd52615370e152954ba21721378a87b2960f7a61"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
certifi = "*"
|
||||
h11 = ">=0.13,<0.15"
|
||||
|
||||
[package.extras]
|
||||
asyncio = ["anyio (>=4.0,<5.0)"]
|
||||
http2 = ["h2 (>=3,<5)"]
|
||||
socks = ["socksio (==1.*)"]
|
||||
trio = ["trio (>=0.22.0,<0.26.0)"]
|
||||
|
||||
[[package]]
|
||||
name = "httpx"
|
||||
version = "0.27.2"
|
||||
description = "The next generation HTTP client."
|
||||
optional = false
|
||||
python-versions = ">=3.8"
|
||||
files = [
|
||||
{file = "httpx-0.27.2-py3-none-any.whl", hash = "sha256:7bb2708e112d8fdd7829cd4243970f0c223274051cb35ee80c03301ee29a3df0"},
|
||||
{file = "httpx-0.27.2.tar.gz", hash = "sha256:f7c2be1d2f3c3c3160d441802406b206c2b76f5947b11115e6df10c6c65e66c2"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
anyio = "*"
|
||||
certifi = "*"
|
||||
httpcore = "==1.*"
|
||||
idna = "*"
|
||||
sniffio = "*"
|
||||
|
||||
[package.extras]
|
||||
brotli = ["brotli", "brotlicffi"]
|
||||
cli = ["click (==8.*)", "pygments (==2.*)", "rich (>=10,<14)"]
|
||||
http2 = ["h2 (>=3,<5)"]
|
||||
socks = ["socksio (==1.*)"]
|
||||
zstd = ["zstandard (>=0.18.0)"]
|
||||
|
||||
[[package]]
|
||||
name = "idna"
|
||||
version = "3.7"
|
||||
@ -546,26 +622,26 @@ files = [
|
||||
|
||||
[[package]]
|
||||
name = "langchain"
|
||||
version = "0.2.9"
|
||||
version = "0.3.0"
|
||||
description = "Building applications with LLMs through composability"
|
||||
optional = false
|
||||
python-versions = "<4.0,>=3.8.1"
|
||||
python-versions = "<4.0,>=3.9"
|
||||
files = [
|
||||
{file = "langchain-0.2.9-py3-none-any.whl", hash = "sha256:be23fcb29adbd5059944f1fed08fa575f0739d420b1c4127531e0fbf5663fcca"},
|
||||
{file = "langchain-0.2.9.tar.gz", hash = "sha256:cc326a7f6347787a19882928c324433b1a79df629bba45604b2d26495ee5d69c"},
|
||||
{file = "langchain-0.3.0-py3-none-any.whl", hash = "sha256:59a75a6a1eb7bfd2a6bf0c7a5816409a8fdc9046187b07af287b23b9899617af"},
|
||||
{file = "langchain-0.3.0.tar.gz", hash = "sha256:a7c23892440bd1f5b9e029ff0dd709dd881ae927c4c0a3210ac64dba9bbf3f7f"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
aiohttp = ">=3.8.3,<4.0.0"
|
||||
async-timeout = {version = ">=4.0.0,<5.0.0", markers = "python_version < \"3.11\""}
|
||||
langchain-core = ">=0.2.20,<0.3.0"
|
||||
langchain-text-splitters = ">=0.2.0,<0.3.0"
|
||||
langchain-core = ">=0.3.0,<0.4.0"
|
||||
langchain-text-splitters = ">=0.3.0,<0.4.0"
|
||||
langsmith = ">=0.1.17,<0.2.0"
|
||||
numpy = [
|
||||
{version = ">=1,<2", markers = "python_version < \"3.12\""},
|
||||
{version = ">=1.26.0,<2.0.0", markers = "python_version >= \"3.12\""},
|
||||
]
|
||||
pydantic = ">=1,<3"
|
||||
pydantic = ">=2.7.4,<3.0.0"
|
||||
PyYAML = ">=5.3"
|
||||
requests = ">=2,<3"
|
||||
SQLAlchemy = ">=1.4,<3"
|
||||
@ -573,23 +649,24 @@ tenacity = ">=8.1.0,<8.4.0 || >8.4.0,<9.0.0"
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "0.2.20"
|
||||
version = "0.3.0"
|
||||
description = "Building applications with LLMs through composability"
|
||||
optional = false
|
||||
python-versions = ">=3.8.1,<4.0"
|
||||
python-versions = ">=3.9,<4.0"
|
||||
files = []
|
||||
develop = true
|
||||
|
||||
[package.dependencies]
|
||||
jsonpatch = "^1.33"
|
||||
langsmith = "^0.1.75"
|
||||
langsmith = "^0.1.117"
|
||||
packaging = ">=23.2,<25"
|
||||
pydantic = [
|
||||
{version = ">=1,<3", markers = "python_full_version < \"3.12.4\""},
|
||||
{version = ">=2.5.2,<3.0.0", markers = "python_full_version < \"3.12.4\""},
|
||||
{version = ">=2.7.4,<3.0.0", markers = "python_full_version >= \"3.12.4\""},
|
||||
]
|
||||
PyYAML = ">=5.3"
|
||||
tenacity = "^8.1.0,!=8.4.0"
|
||||
typing-extensions = ">=4.7"
|
||||
|
||||
[package.source]
|
||||
type = "directory"
|
||||
@ -597,30 +674,31 @@ url = "../../core"
|
||||
|
||||
[[package]]
|
||||
name = "langchain-text-splitters"
|
||||
version = "0.2.2"
|
||||
version = "0.3.0"
|
||||
description = "LangChain text splitting utilities"
|
||||
optional = false
|
||||
python-versions = "<4.0,>=3.8.1"
|
||||
python-versions = "<4.0,>=3.9"
|
||||
files = [
|
||||
{file = "langchain_text_splitters-0.2.2-py3-none-any.whl", hash = "sha256:1c80d4b11b55e2995f02d2a326c0323ee1eeff24507329bb22924e420c782dff"},
|
||||
{file = "langchain_text_splitters-0.2.2.tar.gz", hash = "sha256:a1e45de10919fa6fb080ef0525deab56557e9552083600455cb9fa4238076140"},
|
||||
{file = "langchain_text_splitters-0.3.0-py3-none-any.whl", hash = "sha256:e84243e45eaff16e5b776cd9c81b6d07c55c010ebcb1965deb3d1792b7358e83"},
|
||||
{file = "langchain_text_splitters-0.3.0.tar.gz", hash = "sha256:f9fe0b4d244db1d6de211e7343d4abc4aa90295aa22e1f0c89e51f33c55cd7ce"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
langchain-core = ">=0.2.10,<0.3.0"
|
||||
langchain-core = ">=0.3.0,<0.4.0"
|
||||
|
||||
[[package]]
|
||||
name = "langsmith"
|
||||
version = "0.1.87"
|
||||
version = "0.1.120"
|
||||
description = "Client library to connect to the LangSmith LLM Tracing and Evaluation Platform."
|
||||
optional = false
|
||||
python-versions = "<4.0,>=3.8.1"
|
||||
files = [
|
||||
{file = "langsmith-0.1.87-py3-none-any.whl", hash = "sha256:4cd19539c29367f812667e43c0fb0b6d304a078246508df85c38c4ea3df2d0cf"},
|
||||
{file = "langsmith-0.1.87.tar.gz", hash = "sha256:d2422099708af5717d01559731c1c62d49aebf05a420015e30f6dca5ed44227c"},
|
||||
{file = "langsmith-0.1.120-py3-none-any.whl", hash = "sha256:54d2785e301646c0988e0a69ebe4d976488c87b41928b358cb153b6ddd8db62b"},
|
||||
{file = "langsmith-0.1.120.tar.gz", hash = "sha256:25499ca187b41bd89d784b272b97a8d76f60e0e21bdf20336e8a2aa6a9b23ac9"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
httpx = ">=0.23.0,<1"
|
||||
orjson = ">=3.9.14,<4.0.0"
|
||||
pydantic = [
|
||||
{version = ">=1,<3", markers = "python_full_version < \"3.12.4\""},
|
||||
@ -785,43 +863,6 @@ files = [
|
||||
{file = "mypy_extensions-1.0.0.tar.gz", hash = "sha256:75dbf8955dc00442a438fc4d0666508a9a97b6bd41aa2f0ffe9d2f2725af0782"},
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "numpy"
|
||||
version = "1.24.4"
|
||||
description = "Fundamental package for array computing in Python"
|
||||
optional = false
|
||||
python-versions = ">=3.8"
|
||||
files = [
|
||||
{file = "numpy-1.24.4-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:c0bfb52d2169d58c1cdb8cc1f16989101639b34c7d3ce60ed70b19c63eba0b64"},
|
||||
{file = "numpy-1.24.4-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:ed094d4f0c177b1b8e7aa9cba7d6ceed51c0e569a5318ac0ca9a090680a6a1b1"},
|
||||
{file = "numpy-1.24.4-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:79fc682a374c4a8ed08b331bef9c5f582585d1048fa6d80bc6c35bc384eee9b4"},
|
||||
{file = "numpy-1.24.4-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7ffe43c74893dbf38c2b0a1f5428760a1a9c98285553c89e12d70a96a7f3a4d6"},
|
||||
{file = "numpy-1.24.4-cp310-cp310-win32.whl", hash = "sha256:4c21decb6ea94057331e111a5bed9a79d335658c27ce2adb580fb4d54f2ad9bc"},
|
||||
{file = "numpy-1.24.4-cp310-cp310-win_amd64.whl", hash = "sha256:b4bea75e47d9586d31e892a7401f76e909712a0fd510f58f5337bea9572c571e"},
|
||||
{file = "numpy-1.24.4-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:f136bab9c2cfd8da131132c2cf6cc27331dd6fae65f95f69dcd4ae3c3639c810"},
|
||||
{file = "numpy-1.24.4-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:e2926dac25b313635e4d6cf4dc4e51c8c0ebfed60b801c799ffc4c32bf3d1254"},
|
||||
{file = "numpy-1.24.4-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:222e40d0e2548690405b0b3c7b21d1169117391c2e82c378467ef9ab4c8f0da7"},
|
||||
{file = "numpy-1.24.4-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7215847ce88a85ce39baf9e89070cb860c98fdddacbaa6c0da3ffb31b3350bd5"},
|
||||
{file = "numpy-1.24.4-cp311-cp311-win32.whl", hash = "sha256:4979217d7de511a8d57f4b4b5b2b965f707768440c17cb70fbf254c4b225238d"},
|
||||
{file = "numpy-1.24.4-cp311-cp311-win_amd64.whl", hash = "sha256:b7b1fc9864d7d39e28f41d089bfd6353cb5f27ecd9905348c24187a768c79694"},
|
||||
{file = "numpy-1.24.4-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:1452241c290f3e2a312c137a9999cdbf63f78864d63c79039bda65ee86943f61"},
|
||||
{file = "numpy-1.24.4-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:04640dab83f7c6c85abf9cd729c5b65f1ebd0ccf9de90b270cd61935eef0197f"},
|
||||
{file = "numpy-1.24.4-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a5425b114831d1e77e4b5d812b69d11d962e104095a5b9c3b641a218abcc050e"},
|
||||
{file = "numpy-1.24.4-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dd80e219fd4c71fc3699fc1dadac5dcf4fd882bfc6f7ec53d30fa197b8ee22dc"},
|
||||
{file = "numpy-1.24.4-cp38-cp38-win32.whl", hash = "sha256:4602244f345453db537be5314d3983dbf5834a9701b7723ec28923e2889e0bb2"},
|
||||
{file = "numpy-1.24.4-cp38-cp38-win_amd64.whl", hash = "sha256:692f2e0f55794943c5bfff12b3f56f99af76f902fc47487bdfe97856de51a706"},
|
||||
{file = "numpy-1.24.4-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:2541312fbf09977f3b3ad449c4e5f4bb55d0dbf79226d7724211acc905049400"},
|
||||
{file = "numpy-1.24.4-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:9667575fb6d13c95f1b36aca12c5ee3356bf001b714fc354eb5465ce1609e62f"},
|
||||
{file = "numpy-1.24.4-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f3a86ed21e4f87050382c7bc96571755193c4c1392490744ac73d660e8f564a9"},
|
||||
{file = "numpy-1.24.4-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d11efb4dbecbdf22508d55e48d9c8384db795e1b7b51ea735289ff96613ff74d"},
|
||||
{file = "numpy-1.24.4-cp39-cp39-win32.whl", hash = "sha256:6620c0acd41dbcb368610bb2f4d83145674040025e5536954782467100aa8835"},
|
||||
{file = "numpy-1.24.4-cp39-cp39-win_amd64.whl", hash = "sha256:befe2bf740fd8373cf56149a5c23a0f601e82869598d41f8e188a0e9869926f8"},
|
||||
{file = "numpy-1.24.4-pp38-pypy38_pp73-macosx_10_9_x86_64.whl", hash = "sha256:31f13e25b4e304632a4619d0e0777662c2ffea99fcae2029556b17d8ff958aef"},
|
||||
{file = "numpy-1.24.4-pp38-pypy38_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:95f7ac6540e95bc440ad77f56e520da5bf877f87dca58bd095288dce8940532a"},
|
||||
{file = "numpy-1.24.4-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:e98f220aa76ca2a977fe435f5b04d7b3470c0a2e6312907b37ba6068f26787f2"},
|
||||
{file = "numpy-1.24.4.tar.gz", hash = "sha256:80f5e3a4e498641401868df4208b74581206afbee7cf7b8329daae82676d9463"},
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "numpy"
|
||||
version = "1.26.4"
|
||||
@ -1237,6 +1278,17 @@ files = [
|
||||
{file = "ruff-0.5.2.tar.gz", hash = "sha256:2c0df2d2de685433794a14d8d2e240df619b748fbe3367346baa519d8e6f1ca2"},
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sniffio"
|
||||
version = "1.3.1"
|
||||
description = "Sniff out which async library your code is running under"
|
||||
optional = false
|
||||
python-versions = ">=3.7"
|
||||
files = [
|
||||
{file = "sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2"},
|
||||
{file = "sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc"},
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sqlalchemy"
|
||||
version = "2.0.31"
|
||||
@ -1497,5 +1549,5 @@ multidict = ">=4.0"
|
||||
|
||||
[metadata]
|
||||
lock-version = "2.0"
|
||||
python-versions = ">=3.8.1,<4.0"
|
||||
content-hash = "0a2c4601a26691fd3bf061961c95a19a3869c491d4670dc38b8a4c46d2c01a6e"
|
||||
python-versions = ">=3.9,<4.0"
|
||||
content-hash = "0834872527b82aef5921c6b869a70c49ce36507d858e1371c05a079a77cfad86"
|
||||
|
||||
@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
|
||||
|
||||
[tool.poetry]
|
||||
name = "langchain-couchbase"
|
||||
version = "0.1.1"
|
||||
version = "0.2.0"
|
||||
description = "An integration package connecting Couchbase and LangChain"
|
||||
authors = []
|
||||
readme = "README.md"
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
"""Test Couchbase Cache functionality"""
|
||||
|
||||
import os
|
||||
from datetime import timedelta
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
@ -12,7 +12,13 @@ from langchain_core.globals import get_llm_cache, set_llm_cache
|
||||
from langchain_core.outputs import Generation
|
||||
|
||||
from langchain_couchbase.cache import CouchbaseCache, CouchbaseSemanticCache
|
||||
from tests.utils import FakeEmbeddings, FakeLLM
|
||||
from tests.utils import (
|
||||
FakeEmbeddings,
|
||||
FakeLLM,
|
||||
cache_key_hash_function,
|
||||
fetch_document_expiry_time,
|
||||
get_document_keys,
|
||||
)
|
||||
|
||||
CONNECTION_STRING = os.getenv("COUCHBASE_CONNECTION_STRING", "")
|
||||
BUCKET_NAME = os.getenv("COUCHBASE_BUCKET_NAME", "")
|
||||
@ -88,6 +94,39 @@ class TestCouchbaseCache:
|
||||
output = get_llm_cache().lookup("bar", llm_string)
|
||||
assert output != [Generation(text="fizz")]
|
||||
|
||||
def test_cache_with_ttl(self, cluster: Any) -> None:
|
||||
"""Test standard LLM cache functionality with TTL"""
|
||||
ttl = timedelta(minutes=10)
|
||||
set_llm_cache(
|
||||
CouchbaseCache(
|
||||
cluster=cluster,
|
||||
bucket_name=BUCKET_NAME,
|
||||
scope_name=SCOPE_NAME,
|
||||
collection_name=CACHE_COLLECTION_NAME,
|
||||
ttl=ttl,
|
||||
)
|
||||
)
|
||||
|
||||
llm = FakeLLM()
|
||||
|
||||
params = llm.dict()
|
||||
params["stop"] = None
|
||||
llm_string = str(sorted([(k, v) for k, v in params.items()]))
|
||||
get_llm_cache().update("foo", llm_string, [Generation(text="fizz")])
|
||||
cache_output = get_llm_cache().lookup("foo", llm_string)
|
||||
assert cache_output == [Generation(text="fizz")]
|
||||
|
||||
# Check the document's expiry time by fetching it from the database
|
||||
document_key = cache_key_hash_function("foo" + llm_string)
|
||||
document_expiry_time = fetch_document_expiry_time(
|
||||
cluster, BUCKET_NAME, SCOPE_NAME, CACHE_COLLECTION_NAME, document_key
|
||||
)
|
||||
current_time = datetime.now()
|
||||
|
||||
time_to_expiry = document_expiry_time - current_time
|
||||
|
||||
assert time_to_expiry < ttl
|
||||
|
||||
def test_semantic_cache(self, cluster: Any) -> None:
|
||||
"""Test semantic LLM cache functionality"""
|
||||
set_llm_cache(
|
||||
@ -118,3 +157,62 @@ class TestCouchbaseCache:
|
||||
get_llm_cache().clear()
|
||||
output = get_llm_cache().lookup("bar", llm_string)
|
||||
assert output != [Generation(text="fizz"), Generation(text="Buzz")]
|
||||
|
||||
def test_semantic_cache_with_ttl(self, cluster: Any) -> None:
|
||||
"""Test semantic LLM cache functionality with TTL"""
|
||||
ttl = timedelta(minutes=10)
|
||||
|
||||
set_llm_cache(
|
||||
CouchbaseSemanticCache(
|
||||
cluster=cluster,
|
||||
embedding=FakeEmbeddings(),
|
||||
index_name=INDEX_NAME,
|
||||
bucket_name=BUCKET_NAME,
|
||||
scope_name=SCOPE_NAME,
|
||||
collection_name=SEMANTIC_CACHE_COLLECTION_NAME,
|
||||
ttl=ttl,
|
||||
)
|
||||
)
|
||||
|
||||
llm = FakeLLM()
|
||||
|
||||
params = llm.dict()
|
||||
params["stop"] = None
|
||||
llm_string = str(sorted([(k, v) for k, v in params.items()]))
|
||||
|
||||
# Add a document to the cache
|
||||
seed_prompt = "foo"
|
||||
get_llm_cache().update(
|
||||
seed_prompt, llm_string, [Generation(text="fizz"), Generation(text="Buzz")]
|
||||
)
|
||||
|
||||
# foo and bar will have the same embedding produced by FakeEmbeddings
|
||||
cache_output = get_llm_cache().lookup("bar", llm_string)
|
||||
assert cache_output == [Generation(text="fizz"), Generation(text="Buzz")]
|
||||
|
||||
# Check the document's expiry time by fetching it from the database
|
||||
fetch_document_query = (
|
||||
f"SELECT meta().id, * from `{SEMANTIC_CACHE_COLLECTION_NAME}` doc "
|
||||
f"WHERE doc.text = '{seed_prompt}'"
|
||||
)
|
||||
|
||||
document_keys = get_document_keys(
|
||||
cluster=cluster,
|
||||
bucket_name=BUCKET_NAME,
|
||||
scope_name=SCOPE_NAME,
|
||||
query=fetch_document_query,
|
||||
)
|
||||
assert len(document_keys) == 1
|
||||
|
||||
document_expiry_time = fetch_document_expiry_time(
|
||||
cluster=cluster,
|
||||
bucket_name=BUCKET_NAME,
|
||||
scope_name=SCOPE_NAME,
|
||||
collection_name=SEMANTIC_CACHE_COLLECTION_NAME,
|
||||
document_key=document_keys[0],
|
||||
)
|
||||
current_time = datetime.now()
|
||||
|
||||
time_to_expiry = document_expiry_time - current_time
|
||||
|
||||
assert time_to_expiry < ttl
|
||||
|
||||
@ -2,7 +2,7 @@
|
||||
|
||||
import os
|
||||
import time
|
||||
from datetime import timedelta
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
@ -13,6 +13,7 @@ from langchain.memory import ConversationBufferMemory
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
|
||||
from langchain_couchbase.chat_message_histories import CouchbaseChatMessageHistory
|
||||
from tests.utils import fetch_document_expiry_time, get_document_keys
|
||||
|
||||
CONNECTION_STRING = os.getenv("COUCHBASE_CONNECTION_STRING", "")
|
||||
BUCKET_NAME = os.getenv("COUCHBASE_BUCKET_NAME", "")
|
||||
@ -162,3 +163,132 @@ class TestCouchbaseCache:
|
||||
memory_b.chat_memory.clear()
|
||||
time.sleep(SLEEP_DURATION)
|
||||
assert memory_b.chat_memory.messages == []
|
||||
|
||||
def test_memory_message_with_ttl(self, cluster: Any) -> None:
|
||||
"""Test chat message history with a message being saved with a TTL"""
|
||||
ttl = timedelta(minutes=5)
|
||||
session_id = "test-session-ttl"
|
||||
message_history = CouchbaseChatMessageHistory(
|
||||
cluster=cluster,
|
||||
bucket_name=BUCKET_NAME,
|
||||
scope_name=SCOPE_NAME,
|
||||
collection_name=MESSAGE_HISTORY_COLLECTION_NAME,
|
||||
session_id=session_id,
|
||||
ttl=ttl,
|
||||
)
|
||||
|
||||
memory = ConversationBufferMemory(
|
||||
memory_key="baz", chat_memory=message_history, return_messages=True
|
||||
)
|
||||
|
||||
# clear the memory
|
||||
memory.chat_memory.clear()
|
||||
|
||||
# wait for the messages to be cleared
|
||||
time.sleep(SLEEP_DURATION)
|
||||
assert memory.chat_memory.messages == []
|
||||
|
||||
# add some messages
|
||||
ai_message = AIMessage(content="Hello, how are you doing ?")
|
||||
memory.chat_memory.add_ai_message(ai_message)
|
||||
|
||||
# wait until the messages can be retrieved
|
||||
time.sleep(SLEEP_DURATION)
|
||||
|
||||
# check that the messages are in the memory
|
||||
messages = memory.chat_memory.messages
|
||||
assert len(messages) == 1
|
||||
|
||||
# check that the messages are in the order of creation
|
||||
assert messages == [ai_message]
|
||||
|
||||
# Check the document's expiry time by fetching it from the database
|
||||
fetch_documents_query = (
|
||||
f"SELECT meta().id, * from `{MESSAGE_HISTORY_COLLECTION_NAME}` doc"
|
||||
f" WHERE doc.session_id = '{session_id}'"
|
||||
)
|
||||
|
||||
document_keys = get_document_keys(
|
||||
cluster=cluster,
|
||||
bucket_name=BUCKET_NAME,
|
||||
scope_name=SCOPE_NAME,
|
||||
query=fetch_documents_query,
|
||||
)
|
||||
assert len(document_keys) == 1
|
||||
|
||||
# Ensure that the document will expire within the TTL
|
||||
|
||||
document_expiry_time = fetch_document_expiry_time(
|
||||
cluster=cluster,
|
||||
bucket_name=BUCKET_NAME,
|
||||
scope_name=SCOPE_NAME,
|
||||
collection_name=MESSAGE_HISTORY_COLLECTION_NAME,
|
||||
document_key=document_keys[0],
|
||||
)
|
||||
current_time = datetime.now()
|
||||
assert document_expiry_time - current_time < ttl
|
||||
|
||||
def test_memory_messages_with_ttl(self, cluster: Any) -> None:
|
||||
"""Test chat message history with messages being stored with a TTL"""
|
||||
ttl = timedelta(minutes=5)
|
||||
session_id = "test-session-ttl"
|
||||
message_history = CouchbaseChatMessageHistory(
|
||||
cluster=cluster,
|
||||
bucket_name=BUCKET_NAME,
|
||||
scope_name=SCOPE_NAME,
|
||||
collection_name=MESSAGE_HISTORY_COLLECTION_NAME,
|
||||
session_id=session_id,
|
||||
ttl=ttl,
|
||||
)
|
||||
|
||||
memory = ConversationBufferMemory(
|
||||
memory_key="baz", chat_memory=message_history, return_messages=True
|
||||
)
|
||||
|
||||
# clear the memory
|
||||
memory.chat_memory.clear()
|
||||
|
||||
# wait for the messages to be cleared
|
||||
time.sleep(SLEEP_DURATION)
|
||||
assert memory.chat_memory.messages == []
|
||||
|
||||
# add some messages
|
||||
ai_message = AIMessage(content="Hello, how are you doing ?")
|
||||
user_message = HumanMessage(content="I'm good, how are you?")
|
||||
memory.chat_memory.add_messages([ai_message, user_message])
|
||||
|
||||
# wait until the messages can be retrieved
|
||||
time.sleep(SLEEP_DURATION)
|
||||
|
||||
# check that the messages are in the memory
|
||||
messages = memory.chat_memory.messages
|
||||
assert len(messages) == 2
|
||||
|
||||
# check that the messages are in the order of creation
|
||||
assert messages == [ai_message, user_message]
|
||||
|
||||
# Check the documents' expiry time by fetching the documents from the database
|
||||
fetch_documents_query = (
|
||||
f"SELECT meta().id, * from `{MESSAGE_HISTORY_COLLECTION_NAME}` doc"
|
||||
f" WHERE doc.session_id = '{session_id}'"
|
||||
)
|
||||
|
||||
document_keys = get_document_keys(
|
||||
cluster=cluster,
|
||||
bucket_name=BUCKET_NAME,
|
||||
scope_name=SCOPE_NAME,
|
||||
query=fetch_documents_query,
|
||||
)
|
||||
assert len(document_keys) == 2
|
||||
|
||||
# Ensure that each document will expire within the TTL
|
||||
for document_key in document_keys:
|
||||
document_expiry_time = fetch_document_expiry_time(
|
||||
cluster=cluster,
|
||||
bucket_name=BUCKET_NAME,
|
||||
scope_name=SCOPE_NAME,
|
||||
collection_name=MESSAGE_HISTORY_COLLECTION_NAME,
|
||||
document_key=document_key,
|
||||
)
|
||||
current_time = datetime.now()
|
||||
assert document_expiry_time - current_time < ttl
|
||||
|
||||
@ -1,7 +1,11 @@
|
||||
"""Fake Embedding class for testing purposes."""
|
||||
"""Utilities for testing purposes."""
|
||||
|
||||
import hashlib
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, List, Mapping, Optional, cast
|
||||
|
||||
from couchbase.cluster import Cluster
|
||||
from couchbase.options import GetOptions
|
||||
from langchain_core.callbacks import CallbackManagerForLLMRun
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from langchain_core.language_models.llms import LLM
|
||||
@ -97,3 +101,36 @@ class FakeLLM(LLM):
|
||||
response = queries[list(queries.keys())[self.response_index]]
|
||||
self.response_index = self.response_index + 1
|
||||
return response
|
||||
|
||||
|
||||
def cache_key_hash_function(_input: str) -> str:
|
||||
"""Use a deterministic hashing approach."""
|
||||
return hashlib.md5(_input.encode()).hexdigest()
|
||||
|
||||
|
||||
def fetch_document_expiry_time(
|
||||
cluster: Cluster,
|
||||
bucket_name: str,
|
||||
scope_name: str,
|
||||
collection_name: str,
|
||||
document_key: str,
|
||||
) -> datetime:
|
||||
"""Fetch the document's expiry time from the database."""
|
||||
collection = (
|
||||
cluster.bucket(bucket_name).scope(scope_name).collection(collection_name)
|
||||
)
|
||||
result = collection.get(document_key, GetOptions(with_expiry=True))
|
||||
|
||||
return result.expiryTime
|
||||
|
||||
|
||||
def get_document_keys(
|
||||
cluster: Cluster, bucket_name: str, scope_name: str, query: str
|
||||
) -> List[str]:
|
||||
"""Get the document key from the database based on the query using meta().id."""
|
||||
scope = cluster.bucket(bucket_name).scope(scope_name)
|
||||
|
||||
result = scope.query(query).execute()
|
||||
|
||||
document_keys = [row["id"] for row in result]
|
||||
return document_keys
|
||||
|
||||
Loading…
Reference in New Issue
Block a user