partner[minor]: Astra DB clients identify themselves as coming through LangChain package (#18131)

**Description**

This PR sets the "caller identity" of the Astra DB clients used by the
integration plugins (`AstraDBChatMessageHistory`, `AstraDBStore`,
`AstraDBByteStore` and, pending #17767 , `AstraDBVectorStore`). In this
way, the requests to the Astra DB Data API coming from within LangChain
are identified as such (the purpose is anonymous usage stats to best
improve the Astra DB service).
This commit is contained in:
Stefano Lottini 2024-02-28 23:13:22 +01:00 committed by GitHub
parent 4899a72b56
commit 6d863bed51
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -6,6 +6,7 @@ from asyncio import InvalidStateError, Task
from enum import Enum from enum import Enum
from typing import Awaitable, Optional, Union from typing import Awaitable, Optional, Union
import langchain_core
from astrapy.db import AstraDB, AsyncAstraDB from astrapy.db import AstraDB, AsyncAstraDB
@ -51,13 +52,13 @@ class _AstraDBEnvironment:
) )
if astra_db: if astra_db:
self.astra_db = astra_db self.astra_db = astra_db.copy()
if async_astra_db: if async_astra_db:
self.async_astra_db = async_astra_db self.async_astra_db = async_astra_db.copy()
else: else:
self.async_astra_db = self.astra_db.to_async() self.async_astra_db = self.astra_db.to_async()
elif async_astra_db: elif async_astra_db:
self.async_astra_db = async_astra_db self.async_astra_db = async_astra_db.copy()
self.astra_db = self.async_astra_db.to_sync() self.astra_db = self.async_astra_db.to_sync()
else: else:
raise ValueError( raise ValueError(
@ -65,6 +66,15 @@ class _AstraDBEnvironment:
"'token' and 'api_endpoint'" "'token' and 'api_endpoint'"
) )
self.astra_db.set_caller(
caller_name="langchain",
caller_version=getattr(langchain_core, "__version__", None),
)
self.async_astra_db.set_caller(
caller_name="langchain",
caller_version=getattr(langchain_core, "__version__", None),
)
class _AstraDBCollectionEnvironment(_AstraDBEnvironment): class _AstraDBCollectionEnvironment(_AstraDBEnvironment):
def __init__( def __init__(