community[patch]: Fixing some caching issues for AzureCosmosDBSemanticCache (#18884)

Fixing some issues for AzureCosmosDBSemanticCache
- Added the entry for "AzureCosmosDBSemanticCache" which was missing in
langchain/cache.py
- Added application name when creating the MongoClient for the
AzureCosmosDBVectorSearch, for tracking purposes.

@baskaryan, can you please review this PR, we need this to go in asap.
These are just small fixes which we found today in our testing.
This commit is contained in:
Aayush Kataria 2024-03-25 19:06:17 -07:00 committed by GitHub
parent a6cbb755a7
commit 03c38005cb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 84 additions and 62 deletions

View File

@ -12,12 +12,12 @@
},
{
"cell_type": "code",
"execution_count": 15,
"execution_count": 3,
"id": "10ad9224",
"metadata": {
"ExecuteTime": {
"end_time": "2024-02-02T21:34:23.461332Z",
"start_time": "2024-02-02T21:34:23.394461Z"
"end_time": "2024-03-18T01:01:08.425930Z",
"start_time": "2024-03-18T01:01:08.327196Z"
}
},
"outputs": [],
@ -41,7 +41,7 @@
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": 11,
"id": "426ff912",
"metadata": {},
"outputs": [],
@ -1356,18 +1356,26 @@
},
{
"cell_type": "markdown",
"id": "40624c26e86b57a4",
"metadata": {
"collapsed": false
},
"source": [
"## Azure Cosmos DB Semantic Cache\n",
"\n",
"You can use this integrated [vector database](https://learn.microsoft.com/en-us/azure/cosmos-db/vector-database) for caching."
],
"metadata": {
"collapsed": false
},
"id": "40624c26e86b57a4"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "4a9d592db01b11b2",
"metadata": {
"ExecuteTime": {
"end_time": "2024-03-18T01:01:32.014750Z",
"start_time": "2024-03-18T01:01:31.955991Z"
}
},
"outputs": [],
"source": [
"from langchain.cache import AzureCosmosDBSemanticCache\n",
@ -1379,11 +1387,11 @@
"\n",
"# Read more about Azure CosmosDB Mongo vCore vector search here https://learn.microsoft.com/en-us/azure/cosmos-db/mongodb/vcore/vector-search\n",
"\n",
"INDEX_NAME = \"langchain-test-index\"\n",
"NAMESPACE = \"langchain_test_db.langchain_test_collection\"\n",
"CONNECTION_STRING = (\n",
" \"Please provide your azure cosmos mongo vCore vector db connection string\"\n",
")\n",
"\n",
"DB_NAME, COLLECTION_NAME = NAMESPACE.split(\".\")\n",
"\n",
"# Default value for these params\n",
@ -1394,7 +1402,9 @@
"m = 16\n",
"ef_construction = 64\n",
"ef_search = 40\n",
"score_threshold = 0.1\n",
"score_threshold = 0.9\n",
"application_name = \"LANGCHAIN_CACHING_PYTHON\"\n",
"\n",
"\n",
"set_llm_cache(\n",
" AzureCosmosDBSemanticCache(\n",
@ -1411,18 +1421,10 @@
" ef_construction=ef_construction,\n",
" ef_search=ef_search,\n",
" score_threshold=score_threshold,\n",
" application_name=application_name,\n",
" )\n",
")"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-02-02T21:34:49.457001Z",
"start_time": "2024-02-02T21:34:49.411293Z"
}
},
"id": "4a9d592db01b11b2",
"execution_count": 16
]
},
{
"cell_type": "code",
@ -1431,15 +1433,15 @@
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 43.4 ms, sys: 7.23 ms, total: 50.7 ms\n",
"Wall time: 1.61 s\n"
"CPU times: user 45.6 ms, sys: 19.7 ms, total: 65.3 ms\n",
"Wall time: 2.29 s\n"
]
},
{
"data": {
"text/plain": "\"\\n\\nWhy couldn't the bicycle stand up by itself?\\n\\nBecause it was two-tired!\""
"text/plain": "'\\n\\nWhy was the math book sad? Because it had too many problems.'"
},
"execution_count": 17,
"execution_count": 82,
"metadata": {},
"output_type": "execute_result"
}
@ -1452,47 +1454,46 @@
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-02-02T21:34:53.704234Z",
"start_time": "2024-02-02T21:34:52.091096Z"
"end_time": "2024-03-12T00:12:57.462226Z",
"start_time": "2024-03-12T00:12:55.166201Z"
}
},
"id": "8488cf9c97ec7ab",
"execution_count": 17
"id": "14ca942820e8140c",
"execution_count": 82
},
{
"cell_type": "code",
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 6.89 ms, sys: 2.24 ms, total: 9.13 ms\n",
"Wall time: 337 ms\n"
]
},
{
"data": {
"text/plain": "\"\\n\\nWhy couldn't the bicycle stand up by itself?\\n\\nBecause it was two-tired!\""
},
"execution_count": 18,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"%%time\n",
"# The first time, it is not yet in cache, so it should take longer\n",
"llm(\"Tell me a joke\")"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-02-02T21:34:56.004502Z",
"start_time": "2024-02-02T21:34:55.650136Z"
}
},
"execution_count": 83,
"id": "bc1570a2a77b58c8",
"execution_count": 18
"metadata": {
"ExecuteTime": {
"end_time": "2024-03-12T00:13:03.652755Z",
"start_time": "2024-03-12T00:13:03.159428Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 9.61 ms, sys: 3.42 ms, total: 13 ms\n",
"Wall time: 474 ms\n"
]
},
{
"data": {
"text/plain": "'\\n\\nWhy was the math book sad? Because it had too many problems.'"
},
"execution_count": 83,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"%%time\n",
"# The first time, it is not yet in cache, so it should take longer\n",
"llm(\"Tell me a joke\")"
]
},
{
"cell_type": "markdown",
@ -1743,7 +1744,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.17"
"version": "3.11.4"
}
},
"nbformat": 4,

View File

@ -1879,6 +1879,7 @@ class AzureCosmosDBSemanticCache(BaseCache):
ef_construction: int = 64,
ef_search: int = 40,
score_threshold: Optional[float] = None,
application_name: str = "LANGCHAIN_CACHING_PYTHON",
):
"""
Args:
@ -1920,6 +1921,7 @@ class AzureCosmosDBSemanticCache(BaseCache):
(40 by default). A higher value provides better
recall at the cost of speed.
score_threshold: Maximum score used to filter the vector search documents.
application_name: Application name for the client for tracking and logging
"""
self._validate_enum_value(similarity, CosmosDBSimilarityType)
@ -1942,6 +1944,7 @@ class AzureCosmosDBSemanticCache(BaseCache):
self.ef_search = ef_search
self.score_threshold = score_threshold
self._cache_dict: Dict[str, AzureCosmosDBVectorSearch] = {}
self.application_name = application_name
def _index_name(self, llm_string: str) -> str:
hashed_index = _hash(llm_string)
@ -1972,6 +1975,7 @@ class AzureCosmosDBSemanticCache(BaseCache):
namespace=namespace,
embedding=self.embedding,
index_name=index_name,
application_name=self.application_name,
)
# create index for the vectorstore

View File

@ -119,6 +119,7 @@ class AzureCosmosDBVectorSearch(VectorStore):
connection_string: str,
namespace: str,
embedding: Embeddings,
application_name: str = "LANGCHAIN_PYTHON",
**kwargs: Any,
) -> AzureCosmosDBVectorSearch:
"""Creates an Instance of AzureCosmosDBVectorSearch from a Connection String
@ -140,7 +141,8 @@ class AzureCosmosDBVectorSearch(VectorStore):
"Could not import pymongo, please install it with "
"`pip install pymongo`."
)
client: MongoClient = MongoClient(connection_string)
appname = application_name
client: MongoClient = MongoClient(connection_string, appname=appname)
db_name, collection_name = namespace.split(".")
collection = client[db_name][collection_name]
return cls(collection, embedding, **kwargs)

View File

@ -35,6 +35,7 @@ m = 16
ef_construction = 64
ef_search = 40
score_threshold = 0.1
application_name = "LANGCHAIN_PYTHON"
def prepare_collection() -> Any:
@ -108,6 +109,7 @@ class TestAzureCosmosDBVectorSearch:
azure_openai_embeddings,
collection=collection,
index_name=INDEX_NAME,
application_name=application_name,
)
sleep(1) # waits for Cosmos DB to save contents to the collection
@ -911,6 +913,7 @@ class TestAzureCosmosDBVectorSearch:
NAMESPACE,
azure_openai_embeddings,
index_name=INDEX_NAME,
application_name=application_name,
)
)
@ -926,6 +929,7 @@ class TestAzureCosmosDBVectorSearch:
NAMESPACE,
azure_openai_embeddings,
index_name=INDEX_NAME,
application_name=application_name,
)
)

View File

@ -1,6 +1,7 @@
from langchain_community.cache import (
AstraDBCache,
AstraDBSemanticCache,
AzureCosmosDBSemanticCache,
CassandraCache,
CassandraSemanticCache,
FullLLMCache,
@ -32,4 +33,5 @@ __all__ = [
"SQLAlchemyMd5Cache",
"AstraDBCache",
"AstraDBSemanticCache",
"AzureCosmosDBSemanticCache",
]

View File

@ -36,6 +36,7 @@ m = 16
ef_construction = 64
ef_search = 40
score_threshold = 0.1
application_name = "LANGCHAIN_CACHING_PYTHON"
def _has_env_vars() -> bool:
@ -66,6 +67,7 @@ def test_azure_cosmos_db_semantic_cache() -> None:
ef_construction=ef_construction,
ef_search=ef_search,
score_threshold=score_threshold,
application_name=application_name,
)
)
@ -103,6 +105,7 @@ def test_azure_cosmos_db_semantic_cache_inner_product() -> None:
ef_construction=ef_construction,
ef_search=ef_search,
score_threshold=score_threshold,
application_name=application_name,
)
)
@ -140,6 +143,7 @@ def test_azure_cosmos_db_semantic_cache_multi() -> None:
ef_construction=ef_construction,
ef_search=ef_search,
score_threshold=score_threshold,
application_name=application_name,
)
)
@ -179,6 +183,7 @@ def test_azure_cosmos_db_semantic_cache_multi_inner_product() -> None:
ef_construction=ef_construction,
ef_search=ef_search,
score_threshold=score_threshold,
application_name=application_name,
)
)
@ -218,6 +223,7 @@ def test_azure_cosmos_db_semantic_cache_hnsw() -> None:
ef_construction=ef_construction,
ef_search=ef_search,
score_threshold=score_threshold,
application_name=application_name,
)
)
@ -255,6 +261,7 @@ def test_azure_cosmos_db_semantic_cache_inner_product_hnsw() -> None:
ef_construction=ef_construction,
ef_search=ef_search,
score_threshold=score_threshold,
application_name=application_name,
)
)
@ -292,6 +299,7 @@ def test_azure_cosmos_db_semantic_cache_multi_hnsw() -> None:
ef_construction=ef_construction,
ef_search=ef_search,
score_threshold=score_threshold,
application_name=application_name,
)
)
@ -331,6 +339,7 @@ def test_azure_cosmos_db_semantic_cache_multi_inner_product_hnsw() -> None:
ef_construction=ef_construction,
ef_search=ef_search,
score_threshold=score_threshold,
application_name=application_name,
)
)