fix: update user agent (#36)

* fix: update user agent

* move to common

* fix import

src/langchain_google_datastore/chat_message_history.py

@@ -15,12 +15,12 @@
 from __future__ import annotations
 
 import json
-from typing import TYPE_CHECKING, Any, Iterator, List, Optional
+from typing import TYPE_CHECKING, List, Optional
 
-from google.cloud import datastore
 from langchain_core.chat_history import BaseChatMessageHistory
 from langchain_core.messages import BaseMessage, messages_from_dict
 
+from .common import client_with_user_agent
 from .version import __version__
 
 USER_AGENT = "langchain-google-datastore-python:chat_history/" + __version__
@@ -45,12 +45,7 @@ def __init__(
                 and by default it will use `ChatHistory` as the kind.
             client: Client for interacting with the Google Cloud Firestore API.
         """
-        self.client = client or datastore.Client()
-        client_agent = self.client._client_info.user_agent
-        if not client_agent:
-            self.client._client_info.user_agent = USER_AGENT
-        elif USER_AGENT not in client_agent:
-            self.client._client_info.user_agent = " ".join([client_agent, USER_AGENT])
+        self.client = client_with_user_agent(client, USER_AGENT)
         self.session_id = session_id
         self.key = self.client.key(kind, session_id)
         self.messages: List[BaseMessage] = []

src/langchain_google_datastore/common.py

@@ -0,0 +1,33 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.cloud import datastore  # type: ignore
+from google.cloud.datastore_v1.services.datastore.transports.base import (
+    DEFAULT_CLIENT_INFO,  # type: ignore
+)
+
+
+def client_with_user_agent(
+    client: datastore.Client | None, user_agent: str
+) -> datastore.Client:
+    client_info = DEFAULT_CLIENT_INFO
+    client_info.user_agent = user_agent
+    if not client:
+        client = datastore.Client(client_info=client_info)
+    client_agent = client._client_info.user_agent
+    if not client_agent:
+        client._client_info.user_agent = user_agent
+    elif user_agent not in client_agent:
+        client._client_info.user_agent = " ".join([user_agent, client_agent])
+    return client

src/langchain_google_datastore/document_loader.py

@@ -15,13 +15,13 @@
 from __future__ import annotations
 
 import itertools
-from typing import TYPE_CHECKING, Any, Iterator, List, Optional
+from typing import TYPE_CHECKING, Iterator, List, Optional
 
 import more_itertools
-from google.cloud import datastore
 from langchain_community.document_loaders.base import BaseLoader
 from langchain_core.documents import Document
 
+from .common import client_with_user_agent
 from .document_converter import (
     DATASTORE_TYPE,
     KEY,
@@ -155,14 +155,3 @@ def delete_documents(
                     )
                 db_batch.delete(key)
             db_batch.commit()
-
-
-def client_with_user_agent(client: Client | None, user_agent: str) -> Client:
-    if not client:
-        client = datastore.Client()
-    client_agent = client._client_info.user_agent
-    if not client_agent:
-        client._client_info.user_agent = user_agent
-    elif user_agent not in client_agent:
-        client._client_info.user_agent = " ".join([client_agent, user_agent])
-    return client