Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
13 changes: 7 additions & 6 deletions src/langchain_google_bigtable/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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