diff --git a/pyproject.toml b/pyproject.toml index b60a96c..b646a46 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,3 +42,7 @@ profile = "black" [tool.mypy] python_version = "3.11" warn_unused_configs = true + +[[tool.mypy.overrides]] +module="google.cloud.*" +ignore_missing_imports = true \ No newline at end of file diff --git a/src/langchain_google_bigtable/common.py b/src/langchain_google_bigtable/common.py index f87add8..4d11568 100644 --- a/src/langchain_google_bigtable/common.py +++ b/src/langchain_google_bigtable/common.py @@ -12,29 +12,30 @@ # See the License for the specific language governing permissions and # limitations under the License. -from importlib.metadata import version from typing import Dict, Optional from google.cloud import bigtable # type: ignore +from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin.transports.base import ( + DEFAULT_CLIENT_INFO, # type: ignore +) from .version import __version__ PACKAGE_PREFIX = "langchain-google-bigtable-python:" -clients: Dict[str, bigtable.Client] = {} def use_client_or_default( client: Optional[bigtable.Client], client_name: str ) -> bigtable.Client: user_agent = PACKAGE_PREFIX + client_name + "/" + __version__ - global clients + client_info = DEFAULT_CLIENT_INFO + client_info.user_agent = user_agent if not client: - client = clients.get(client_name, bigtable.Client(admin=True)) - clients[client_name] = client + client = bigtable.Client(admin=True, 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([client_agent, user_agent]) + client._client_info.user_agent = " ".join([user_agent, client_agent]) return client