mirror of
https://github.com/hwchase17/langchain.git
synced 2025-08-02 01:23:07 +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:
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,
|
||||
|
@ -1,325 +1,354 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "a283d2fd-e26e-4811-a486-d3cf0ecf6749",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Couchbase\n",
|
||||
"> Couchbase is an award-winning distributed NoSQL cloud database that delivers unmatched versatility, performance, scalability, and financial value for all of your cloud, mobile, AI, and edge computing applications. Couchbase embraces AI with coding assistance for developers and vector search for their applications.\n",
|
||||
"\n",
|
||||
"This notebook goes over how to use the `CouchbaseChatMessageHistory` class to store the chat message history in a Couchbase cluster\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ff868a6c-3e17-4c3d-8d32-67b01f4d7bcc",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Set Up Couchbase Cluster\n",
|
||||
"To run this demo, you need a Couchbase Cluster. \n",
|
||||
"\n",
|
||||
"You can work with both [Couchbase Capella](https://www.couchbase.com/products/capella/) and your self-managed Couchbase Server."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "41fa85e7-6968-45e4-a445-de305d80f332",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Install Dependencies\n",
|
||||
"`CouchbaseChatMessageHistory` lives inside the `langchain-couchbase` package. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"id": "b744ca05-b8c6-458c-91df-f50ca2c20b3c",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Note: you may need to restart the kernel to use updated packages.\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"%pip install --upgrade --quiet langchain-couchbase"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "41f29205-6452-493b-ba18-8a3b006bcca4",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Create Couchbase Connection Object\n",
|
||||
"We create a connection to the Couchbase cluster initially and then pass the cluster object to the Vector Store. \n",
|
||||
"\n",
|
||||
"Here, we are connecting using the username and password. You can also connect using any other supported way to your cluster. \n",
|
||||
"\n",
|
||||
"For more information on connecting to the Couchbase cluster, please check the [Python SDK documentation](https://docs.couchbase.com/python-sdk/current/hello-world/start-using-sdk.html#connect)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"id": "f394908e-f5fe-408a-84d7-b97fdebcfa26",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"COUCHBASE_CONNECTION_STRING = (\n",
|
||||
" \"couchbase://localhost\" # or \"couchbases://localhost\" if using TLS\n",
|
||||
")\n",
|
||||
"DB_USERNAME = \"Administrator\"\n",
|
||||
"DB_PASSWORD = \"Password\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"id": "ad4dce21-d80c-465a-b709-fd366ba5ce35",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from datetime import timedelta\n",
|
||||
"\n",
|
||||
"from couchbase.auth import PasswordAuthenticator\n",
|
||||
"from couchbase.cluster import Cluster\n",
|
||||
"from couchbase.options import ClusterOptions\n",
|
||||
"\n",
|
||||
"auth = PasswordAuthenticator(DB_USERNAME, DB_PASSWORD)\n",
|
||||
"options = ClusterOptions(auth)\n",
|
||||
"cluster = Cluster(COUCHBASE_CONNECTION_STRING, options)\n",
|
||||
"\n",
|
||||
"# Wait until the cluster is ready for use.\n",
|
||||
"cluster.wait_until_ready(timedelta(seconds=5))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "e3d0210c-e2e6-437a-86f3-7397a1899fef",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"We will now set the bucket, scope, and collection names in the Couchbase cluster that we want to use for storing the message history.\n",
|
||||
"\n",
|
||||
"Note that the bucket, scope, and collection need to exist before using them to store the message history."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"id": "e8c7f846-a5c4-4465-a40e-4a9a23ac71bd",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"BUCKET_NAME = \"langchain-testing\"\n",
|
||||
"SCOPE_NAME = \"_default\"\n",
|
||||
"COLLECTION_NAME = \"conversational_cache\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "283959e1-6af7-4768-9211-5b0facc6ef65",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Usage\n",
|
||||
"In order to store the messages, you need the following:\n",
|
||||
"- Couchbase Cluster object: Valid connection to the Couchbase cluster\n",
|
||||
"- bucket_name: Bucket in cluster to store the chat message history\n",
|
||||
"- scope_name: Scope in bucket to store the message history\n",
|
||||
"- collection_name: Collection in scope to store the message history\n",
|
||||
"- session_id: Unique identifier for the session\n",
|
||||
"\n",
|
||||
"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"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"id": "43c3b2d5-aae2-44a9-9e9f-f10adf054cfa",
|
||||
"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",
|
||||
")\n",
|
||||
"\n",
|
||||
"message_history.add_user_message(\"hi!\")\n",
|
||||
"\n",
|
||||
"message_history.add_ai_message(\"how are you doing?\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"id": "e7e348ef-79e9-481c-aeef-969ae03dea6a",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[HumanMessage(content='hi!'), AIMessage(content='how are you doing?')]"
|
||||
]
|
||||
},
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"message_history.messages"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "c8b942a7-93fa-4cd9-8414-d047135c2733",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Chaining\n",
|
||||
"The chat message history class can be used with [LCEL Runnables](https://python.langchain.com/docs/how_to/message_history/)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "8a9f0d91-d1d6-481d-8137-ea11229f485a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import getpass\n",
|
||||
"import os\n",
|
||||
"\n",
|
||||
"from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
|
||||
"from langchain_core.runnables.history import RunnableWithMessageHistory\n",
|
||||
"from langchain_openai import ChatOpenAI\n",
|
||||
"\n",
|
||||
"os.environ[\"OPENAI_API_KEY\"] = getpass.getpass()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"id": "946d45aa-5a61-49ae-816b-1c3949c56d9a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"prompt = ChatPromptTemplate.from_messages(\n",
|
||||
" [\n",
|
||||
" (\"system\", \"You are a helpful assistant.\"),\n",
|
||||
" MessagesPlaceholder(variable_name=\"history\"),\n",
|
||||
" (\"human\", \"{question}\"),\n",
|
||||
" ]\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Create the LCEL runnable\n",
|
||||
"chain = prompt | ChatOpenAI()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"id": "20dfd838-b549-42ed-b3ba-ac005f7e024c",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"chain_with_history = RunnableWithMessageHistory(\n",
|
||||
" chain,\n",
|
||||
" lambda session_id: CouchbaseChatMessageHistory(\n",
|
||||
" cluster=cluster,\n",
|
||||
" bucket_name=BUCKET_NAME,\n",
|
||||
" scope_name=SCOPE_NAME,\n",
|
||||
" collection_name=COLLECTION_NAME,\n",
|
||||
" session_id=session_id,\n",
|
||||
" ),\n",
|
||||
" input_messages_key=\"question\",\n",
|
||||
" history_messages_key=\"history\",\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"id": "17bd09f4-896d-433d-bb9a-369a06e7aa8a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# This is where we configure the session id\n",
|
||||
"config = {\"configurable\": {\"session_id\": \"testing\"}}"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 11,
|
||||
"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})"
|
||||
]
|
||||
},
|
||||
"execution_count": 11,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"chain_with_history.invoke({\"question\": \"Hi! I'm bob\"}, config=config)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 12,
|
||||
"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})"
|
||||
]
|
||||
},
|
||||
"execution_count": 12,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"chain_with_history.invoke({\"question\": \"Whats my name\"}, config=config)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"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.10.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "a283d2fd-e26e-4811-a486-d3cf0ecf6749",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Couchbase\n",
|
||||
"> Couchbase is an award-winning distributed NoSQL cloud database that delivers unmatched versatility, performance, scalability, and financial value for all of your cloud, mobile, AI, and edge computing applications. Couchbase embraces AI with coding assistance for developers and vector search for their applications.\n",
|
||||
"\n",
|
||||
"This notebook goes over how to use the `CouchbaseChatMessageHistory` class to store the chat message history in a Couchbase cluster\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ff868a6c-3e17-4c3d-8d32-67b01f4d7bcc",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Set Up Couchbase Cluster\n",
|
||||
"To run this demo, you need a Couchbase Cluster. \n",
|
||||
"\n",
|
||||
"You can work with both [Couchbase Capella](https://www.couchbase.com/products/capella/) and your self-managed Couchbase Server."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "41fa85e7-6968-45e4-a445-de305d80f332",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Install Dependencies\n",
|
||||
"`CouchbaseChatMessageHistory` lives inside the `langchain-couchbase` package. "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"id": "b744ca05-b8c6-458c-91df-f50ca2c20b3c",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Note: you may need to restart the kernel to use updated packages.\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"%pip install --upgrade --quiet langchain-couchbase"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "41f29205-6452-493b-ba18-8a3b006bcca4",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Create Couchbase Connection Object\n",
|
||||
"We create a connection to the Couchbase cluster initially and then pass the cluster object to the Vector Store. \n",
|
||||
"\n",
|
||||
"Here, we are connecting using the username and password. You can also connect using any other supported way to your cluster. \n",
|
||||
"\n",
|
||||
"For more information on connecting to the Couchbase cluster, please check the [Python SDK documentation](https://docs.couchbase.com/python-sdk/current/hello-world/start-using-sdk.html#connect)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"id": "f394908e-f5fe-408a-84d7-b97fdebcfa26",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"COUCHBASE_CONNECTION_STRING = (\n",
|
||||
" \"couchbase://localhost\" # or \"couchbases://localhost\" if using TLS\n",
|
||||
")\n",
|
||||
"DB_USERNAME = \"Administrator\"\n",
|
||||
"DB_PASSWORD = \"Password\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"id": "ad4dce21-d80c-465a-b709-fd366ba5ce35",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from datetime import timedelta\n",
|
||||
"\n",
|
||||
"from couchbase.auth import PasswordAuthenticator\n",
|
||||
"from couchbase.cluster import Cluster\n",
|
||||
"from couchbase.options import ClusterOptions\n",
|
||||
"\n",
|
||||
"auth = PasswordAuthenticator(DB_USERNAME, DB_PASSWORD)\n",
|
||||
"options = ClusterOptions(auth)\n",
|
||||
"cluster = Cluster(COUCHBASE_CONNECTION_STRING, options)\n",
|
||||
"\n",
|
||||
"# Wait until the cluster is ready for use.\n",
|
||||
"cluster.wait_until_ready(timedelta(seconds=5))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "e3d0210c-e2e6-437a-86f3-7397a1899fef",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"We will now set the bucket, scope, and collection names in the Couchbase cluster that we want to use for storing the message history.\n",
|
||||
"\n",
|
||||
"Note that the bucket, scope, and collection need to exist before using them to store the message history."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"id": "e8c7f846-a5c4-4465-a40e-4a9a23ac71bd",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"BUCKET_NAME = \"langchain-testing\"\n",
|
||||
"SCOPE_NAME = \"_default\"\n",
|
||||
"COLLECTION_NAME = \"conversational_cache\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "283959e1-6af7-4768-9211-5b0facc6ef65",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Usage\n",
|
||||
"In order to store the messages, you need the following:\n",
|
||||
"- Couchbase Cluster object: Valid connection to the Couchbase cluster\n",
|
||||
"- bucket_name: Bucket in cluster to store the chat message history\n",
|
||||
"- scope_name: Scope in bucket to store the message history\n",
|
||||
"- collection_name: Collection in scope to store the message history\n",
|
||||
"- session_id: Unique identifier for the session\n",
|
||||
"\n",
|
||||
"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\n",
|
||||
"- ttl: Used to specify a time in `timedelta` to live for the documents after which they will get deleted automatically from the storage."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"id": "43c3b2d5-aae2-44a9-9e9f-f10adf054cfa",
|
||||
"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",
|
||||
")\n",
|
||||
"\n",
|
||||
"message_history.add_user_message(\"hi!\")\n",
|
||||
"\n",
|
||||
"message_history.add_ai_message(\"how are you doing?\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"id": "e7e348ef-79e9-481c-aeef-969ae03dea6a",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[HumanMessage(content='hi!'), AIMessage(content='how are you doing?')]"
|
||||
]
|
||||
},
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"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",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Chaining\n",
|
||||
"The chat message history class can be used with [LCEL Runnables](https://python.langchain.com/docs/how_to/message_history/)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"id": "8a9f0d91-d1d6-481d-8137-ea11229f485a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import getpass\n",
|
||||
"import os\n",
|
||||
"\n",
|
||||
"from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
|
||||
"from langchain_core.runnables.history import RunnableWithMessageHistory\n",
|
||||
"from langchain_openai import ChatOpenAI\n",
|
||||
"\n",
|
||||
"os.environ[\"OPENAI_API_KEY\"] = getpass.getpass()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"id": "946d45aa-5a61-49ae-816b-1c3949c56d9a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"prompt = ChatPromptTemplate.from_messages(\n",
|
||||
" [\n",
|
||||
" (\"system\", \"You are a helpful assistant.\"),\n",
|
||||
" MessagesPlaceholder(variable_name=\"history\"),\n",
|
||||
" (\"human\", \"{question}\"),\n",
|
||||
" ]\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Create the LCEL runnable\n",
|
||||
"chain = prompt | ChatOpenAI()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"id": "20dfd838-b549-42ed-b3ba-ac005f7e024c",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"chain_with_history = RunnableWithMessageHistory(\n",
|
||||
" chain,\n",
|
||||
" lambda session_id: CouchbaseChatMessageHistory(\n",
|
||||
" cluster=cluster,\n",
|
||||
" bucket_name=BUCKET_NAME,\n",
|
||||
" scope_name=SCOPE_NAME,\n",
|
||||
" collection_name=COLLECTION_NAME,\n",
|
||||
" session_id=session_id,\n",
|
||||
" ),\n",
|
||||
" input_messages_key=\"question\",\n",
|
||||
" history_messages_key=\"history\",\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 11,
|
||||
"id": "17bd09f4-896d-433d-bb9a-369a06e7aa8a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# This is where we configure the session id\n",
|
||||
"config = {\"configurable\": {\"session_id\": \"testing\"}}"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 12,
|
||||
"id": "4bda1096-2fc2-40d7-a046-0d5d8e3a8f75",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"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": 12,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"chain_with_history.invoke({\"question\": \"Hi! I'm bob\"}, config=config)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 13,
|
||||
"id": "1cfb31da-51bb-4c5f-909a-b7118b0ae08d",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"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": 13,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"chain_with_history.invoke({\"question\": \"Whats my name\"}, config=config)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"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.10.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
@ -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:
|
||||
self._collection.upsert(
|
||||
key=self._generate_key(prompt, llm_string), value=doc
|
||||
)
|
||||
if self._ttl:
|
||||
self._collection.upsert(
|
||||
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,15 +175,27 @@ class CouchbaseChatMessageHistory(BaseChatMessageHistory):
|
||||
# get utc timestamp for ordering the messages
|
||||
timestamp = time.time()
|
||||
message_content = message_to_dict(message)
|
||||
|
||||
try:
|
||||
self._collection.insert(
|
||||
document_key,
|
||||
value={
|
||||
self._message_key: message_content,
|
||||
self._session_id_key: self._session_id,
|
||||
self._ts_key: timestamp,
|
||||
},
|
||||
)
|
||||
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={
|
||||
self._message_key: message_content,
|
||||
self._session_id_key: self._session_id,
|
||||
self._ts_key: timestamp,
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("Error adding message: ", e)
|
||||
|
||||
@ -192,7 +223,10 @@ 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}
|
||||
self._collection.insert_multi(insert_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,7 +399,11 @@ class CouchbaseVectorStore(VectorStore):
|
||||
for i in range(0, len(documents_to_insert), batch_size):
|
||||
batch = documents_to_insert[i : i + batch_size]
|
||||
try:
|
||||
result = self._collection.upsert_multi(batch[0])
|
||||
# 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())
|
||||
except DocumentExistsException as e:
|
||||
|
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