fix: update user agent (#36)

* fix: update user agent * move to common * fix import

Author: averikitsch

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