INTPYTHON-461 - Add additional client metadata on LangChain Integrations (#183)

Author: NoahStapp

libs/langchain-mongodb/langchain_mongodb/agent_toolkit/database.py

@@ -5,7 +5,6 @@
 import json
 import re
 from datetime import date, datetime
-from importlib.metadata import version
 from typing import Any, Dict, Iterable, List, Optional, Union
 
 from bson import ObjectId
@@ -14,9 +13,10 @@
 from bson.json_util import dumps
 from pymongo import MongoClient
 from pymongo.cursor import Cursor
-from pymongo.driver_info import DriverInfo
 from pymongo.errors import PyMongoError
 
+from langchain_mongodb.utils import DRIVER_METADATA, _append_client_metadata
+
 NUM_DOCUMENTS_TO_SAMPLE = 4
 MAX_STRING_LENGTH_OF_SAMPLE_DOCUMENT_VALUE = 20
 
@@ -62,6 +62,8 @@ def __init__(
         self._sample_docs_in_coll_info = sample_docs_in_collection_info
         self._indexes_in_coll_info = indexes_in_collection_info
 
+        _append_client_metadata(self._client)
+
     @classmethod
     def from_connection_string(
         cls,
@@ -72,7 +74,7 @@ def from_connection_string(
         """Construct a MongoDBDatabase from URI."""
         client: MongoClient[dict[str, Any]] = MongoClient(
             connection_string,
-            driver=DriverInfo(name="Langchain", version=version("langchain-mongodb")),
+            driver=DRIVER_METADATA,
         )
         database = database or client.get_default_database().name
         return cls(client, database, **kwargs)

libs/langchain-mongodb/langchain_mongodb/cache.py

@@ -3,7 +3,6 @@
 import json
 import logging
 import time
-from importlib.metadata import version
 from typing import Any, Callable, Dict, Optional, Union
 
 from langchain_core.caches import RETURN_VAL_TYPE, BaseCache
@@ -14,8 +13,8 @@
 from pymongo import MongoClient
 from pymongo.collection import Collection
 from pymongo.database import Database
-from pymongo.driver_info import DriverInfo
 
+from langchain_mongodb.utils import DRIVER_METADATA
 from langchain_mongodb.vectorstores import MongoDBAtlasVectorSearch
 
 logger = logging.getLogger(__file__)
@@ -211,7 +210,7 @@ def clear(self, **kwargs: Any) -> None:
 def _generate_mongo_client(connection_string: str) -> MongoClient:
     return MongoClient(
         connection_string,
-        driver=DriverInfo(name="Langchain", version=version("langchain-mongodb")),
+        driver=DRIVER_METADATA,
     )
 
 

libs/langchain-mongodb/langchain_mongodb/chat_message_histories.py

@@ -1,6 +1,5 @@
 import json
 import logging
-from importlib.metadata import version
 from typing import Dict, List, Optional
 
 from langchain_core.chat_history import BaseChatMessageHistory
@@ -10,7 +9,8 @@
     messages_from_dict,
 )
 from pymongo import MongoClient, errors
-from pymongo.driver_info import DriverInfo
+
+from langchain_mongodb.utils import DRIVER_METADATA, _append_client_metadata
 
 logger = logging.getLogger(__name__)
 
@@ -112,13 +112,12 @@ def __init__(
             if connection_string:
                 raise ValueError("Must provide connection_string or client, not both")
             self.client = client
+            _append_client_metadata(self.client)
         elif connection_string:
             try:
                 self.client = MongoClient(
                     connection_string,
-                    driver=DriverInfo(
-                        name="Langchain", version=version("langchain-mongodb")
-                    ),
+                    driver=DRIVER_METADATA,
                 )
             except errors.ConnectionFailure as error:
                 logger.error(error)

libs/langchain-mongodb/langchain_mongodb/docstores.py

@@ -1,15 +1,15 @@
 from __future__ import annotations
 
-from importlib.metadata import version
 from typing import Any, Generator, Iterable, Iterator, List, Optional, Sequence, Union
 
 from langchain_core.documents import Document
 from langchain_core.stores import BaseStore
 from pymongo import MongoClient
 from pymongo.collection import Collection
-from pymongo.driver_info import DriverInfo
 
 from langchain_mongodb.utils import (
+    DRIVER_METADATA,
+    _append_client_metadata,
     make_serializable,
 )
 
@@ -37,6 +37,8 @@ def __init__(self, collection: Collection, text_key: str = "page_content") -> No
         self.collection = collection
         self._text_key = text_key
 
+        _append_client_metadata(self.collection.database.client)
+
     @classmethod
     def from_connection_string(
         cls,
@@ -55,7 +57,7 @@ def from_connection_string(
         """
         client: MongoClient = MongoClient(
             connection_string,
-            driver=DriverInfo(name="Langchain", version=version("langchain-mongodb")),
+            driver=DRIVER_METADATA,
         )
         db_name, collection_name = namespace.split(".")
         collection = client[db_name][collection_name]

libs/langchain-mongodb/langchain_mongodb/graphrag/graph.py

@@ -3,7 +3,6 @@
 import json
 import logging
 from copy import deepcopy
-from importlib.metadata import version
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 from langchain_core.documents import Document
@@ -12,12 +11,12 @@
 from langchain_core.prompts.chat import ChatPromptTemplate
 from pymongo import MongoClient, UpdateOne
 from pymongo.collection import Collection
-from pymongo.driver_info import DriverInfo
 from pymongo.errors import OperationFailure
 from pymongo.results import BulkWriteResult
 
 from langchain_mongodb.graphrag import example_templates, prompts
 
+from ..utils import DRIVER_METADATA, _append_client_metadata
 from .prompts import rag_prompt
 from .schema import entity_schema
 
@@ -140,9 +139,7 @@ def __init__(
             assert database_name is not None
             client: MongoClient = MongoClient(
                 connection_string,
-                driver=DriverInfo(
-                    name="Langchain", version=version("langchain-mongodb")
-                ),
+                driver=DRIVER_METADATA,
             )
             db = client[database_name]
             if collection_name not in db.list_collection_names():
@@ -186,6 +183,9 @@ def __init__(
                     )
         self.collection = collection
 
+        # append_metadata was added in PyMongo 4.14.0, but is a valid database name on earlier versions
+        _append_client_metadata(collection.database.client)
+
         self.entity_extraction_model = entity_extraction_model
         self.entity_prompt = (
             prompts.entity_prompt if entity_prompt is None else entity_prompt
@@ -268,7 +268,7 @@ def from_connection_string(
         """
         client: MongoClient = MongoClient(
             connection_string,
-            driver=DriverInfo(name="Langchain", version=version("langchain-mongodb")),
+            driver=DRIVER_METADATA,
         )
         collection = client[database_name].create_collection(collection_name)
         return cls(

libs/langchain-mongodb/langchain_mongodb/indexes.py

@@ -3,16 +3,16 @@
 
 import functools
 import warnings
-from importlib.metadata import version
 from typing import Any, Dict, List, Optional, Sequence
 
 from langchain_core.indexing.base import RecordManager
 from langchain_core.runnables.config import run_in_executor
 from pymongo import MongoClient
 from pymongo.collection import Collection
-from pymongo.driver_info import DriverInfo
 from pymongo.errors import OperationFailure
 
+from langchain_mongodb.utils import DRIVER_METADATA, _append_client_metadata
+
 
 class MongoDBRecordManager(RecordManager):
     """A MongoDB-based implementation of the record manager."""
@@ -36,6 +36,8 @@ def __init__(self, collection: Collection) -> None:
         super().__init__(namespace=namespace)
         self._collection = collection
 
+        _append_client_metadata(self._collection.database.client)
+
     @classmethod
     def from_connection_string(
         cls, connection_string: str, namespace: str
@@ -51,7 +53,7 @@ def from_connection_string(
         """
         client: MongoClient = MongoClient(
             connection_string,
-            driver=DriverInfo(name="Langchain", version=version("langchain-mongodb")),
+            driver=DRIVER_METADATA,
         )
         db_name, collection_name = namespace.split(".")
         collection = client[db_name][collection_name]

libs/langchain-mongodb/langchain_mongodb/loaders.py

@@ -2,15 +2,15 @@
 from __future__ import annotations
 
 import logging
-from importlib.metadata import version
 from typing import Any, Dict, List, Optional, Sequence
 
 from langchain_community.document_loaders.base import BaseLoader
 from langchain_core.documents import Document
 from langchain_core.runnables.config import run_in_executor
 from pymongo import MongoClient
 from pymongo.collection import Collection
-from pymongo.driver_info import DriverInfo
+
+from langchain_mongodb.utils import DRIVER_METADATA, _append_client_metadata
 
 logger = logging.getLogger(__name__)
 
@@ -53,6 +53,9 @@ def __init__(
         self.metadata_names = metadata_names or []
         self.include_db_collection_in_metadata = include_db_collection_in_metadata
 
+        # append_metadata was added in PyMongo 4.14.0, but is a valid database name on earlier versions
+        _append_client_metadata(self.db.client)
+
     @classmethod
     def from_connection_string(
         cls,
@@ -84,7 +87,7 @@ def from_connection_string(
         """
         client: MongoClient[dict[str, Any]] = MongoClient(
             connection_string,
-            driver=DriverInfo(name="Langchain", version=version("langchain-mongodb")),
+            driver=DRIVER_METADATA,
         )
         collection = client[db_name][collection_name]
         return MongoDBLoader(

libs/langchain-mongodb/langchain_mongodb/retrievers/full_text_search.py

@@ -7,7 +7,7 @@
 from pymongo.collection import Collection
 
 from langchain_mongodb.pipelines import text_search_stage
-from langchain_mongodb.utils import make_serializable
+from langchain_mongodb.utils import _append_client_metadata, make_serializable
 
 
 class MongoDBAtlasFullTextSearchRetriever(BaseRetriever):
@@ -28,6 +28,7 @@ class MongoDBAtlasFullTextSearchRetriever(BaseRetriever):
     top_k: Annotated[
         Optional[int], Field(deprecated='top_k is deprecated, use "k" instead')
     ] = None
+    _added_metadata: bool = False
     """Number of documents to return. Default is no limit"""
 
     def close(self) -> None:
@@ -55,6 +56,10 @@ def _get_relevant_documents(
             include_scores=self.include_scores,
         )
 
+        if not self._added_metadata:
+            _append_client_metadata(self.collection.database.client)
+            self._added_metadata = True
+
         # Execution
         cursor = self.collection.aggregate(pipeline)  # type: ignore[arg-type]
 

libs/langchain-mongodb/langchain_mongodb/retrievers/parent_document.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from importlib.metadata import version
 from typing import Any, List, Optional
 
 import pymongo
@@ -14,12 +13,11 @@
 from langchain_core.runnables import run_in_executor
 from langchain_text_splitters import TextSplitter
 from pymongo import MongoClient
-from pymongo.driver_info import DriverInfo
 
 from langchain_mongodb import MongoDBAtlasVectorSearch
 from langchain_mongodb.docstores import MongoDBDocStore
 from langchain_mongodb.pipelines import vector_search_stage
-from langchain_mongodb.utils import make_serializable
+from langchain_mongodb.utils import DRIVER_METADATA, make_serializable
 
 
 class MongoDBAtlasParentDocumentRetriever(ParentDocumentRetriever):
@@ -168,7 +166,7 @@ def from_connection_string(
         """
         client: MongoClient = MongoClient(
             connection_string,
-            driver=DriverInfo(name="Langchain", version=version("langchain-mongodb")),
+            driver=DRIVER_METADATA,
         )
         collection = client[database_name][collection_name]
         vectorstore = MongoDBAtlasVectorSearch(

libs/langchain-mongodb/langchain_mongodb/utils.py

@@ -20,14 +20,25 @@
 
 import logging
 from datetime import date, datetime
+from importlib.metadata import version
 from typing import Any, Dict, List, Union
 
 import numpy as np
+from pymongo import MongoClient
+from pymongo.driver_info import DriverInfo
 
 logger = logging.getLogger(__name__)
 
 Matrix = Union[List[List[float]], List[np.ndarray], np.ndarray]
 
+DRIVER_METADATA = DriverInfo(name="Langchain", version=version("langchain-mongodb"))
+
+
+def _append_client_metadata(client: MongoClient) -> None:
+    # append_metadata was added in PyMongo 4.14.0, but is a valid database name on earlier versions
+    if callable(client.append_metadata):
+        client.append_metadata(DRIVER_METADATA)
+
 
 def cosine_similarity(X: Matrix, Y: Matrix) -> np.ndarray:
     """Row-wise cosine similarity between two equal-width matrices."""