partner(google-vertex-ai): Add Custom User Agent to Vertex AI SDK

libs/community/langchain_community/chat_models/vertexai.py

@@ -228,6 +228,7 @@ def get_lc_namespace(cls) -> List[str]:
     def validate_environment(cls, values: Dict) -> Dict:
         """Validate that the python package exists in environment."""
         is_gemini = is_gemini_model(values["model_name"])
+        values["module"] = "vertexai-chat"
         cls._try_init_vertexai(values)
         try:
             from vertexai.language_models import ChatModel, CodeChatModel

libs/community/langchain_community/embeddings/vertexai.py

@@ -36,6 +36,7 @@ class VertexAIEmbeddings(_VertexAICommon, Embeddings):
     @root_validator()
     def validate_environment(cls, values: Dict) -> Dict:
         """Validates that the python package exists in environment."""
+        values["module"] = "vertexai-embeddings"
         cls._try_init_vertexai(values)
         if values["model_name"] == "textembedding-gecko-default":
             logger.warning(

libs/community/langchain_community/llms/vertexai.py

@@ -181,7 +181,7 @@ def _default_params(self) -> Dict[str, Any]:
 
     @classmethod
     def _try_init_vertexai(cls, values: Dict) -> None:
-        allowed_params = ["project", "location", "credentials"]
+        allowed_params = ["project", "location", "credentials", "module"]
         params = {k: v for k, v in values.items() if k in allowed_params}
         init_vertexai(**params)
         return None
@@ -229,6 +229,7 @@ def validate_environment(cls, values: Dict) -> Dict:
         tuned_model_name = values.get("tuned_model_name")
         model_name = values["model_name"]
         is_gemini = is_gemini_model(values["model_name"])
+        values["module"] = f"vertexai-llm-{model_name}"
         cls._try_init_vertexai(values)
         try:
             from vertexai.language_models import (

libs/community/langchain_community/utilities/vertexai.py

@@ -1,6 +1,6 @@
 """Utilities to init Vertex AI."""
 from importlib import metadata
-from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple, Union
 
 from langchain_core.callbacks import (
     AsyncCallbackManagerForLLMRun,
@@ -52,10 +52,22 @@ def raise_vertex_import_error(minimum_expected_version: str = "1.38.0") -> None:
     )
 
 
+def _get_user_agent(module: Optional[str] = None) -> Tuple[str, str]:
+    try:
+        langchain_version = metadata.version("langchain")
+    except metadata.PackageNotFoundError:
+        langchain_version = "0.0.0"
+    client_library_version = (
+        f"{langchain_version}-{module}" if module else langchain_version
+    )
+    return client_library_version, f"langchain/{client_library_version}"
+
+
 def init_vertexai(
     project: Optional[str] = None,
     location: Optional[str] = None,
     credentials: Optional["Credentials"] = None,
+    module: Optional[str] = None,
 ) -> None:
     """Init Vertex AI.
 
@@ -65,12 +77,14 @@ def init_vertexai(
         credentials: The default custom
             credentials to use when making API calls. If not provided credentials
             will be ascertained from the environment.
+        module: The module for a custom user agent header.
 
     Raises:
         ImportError: If importing vertexai SDK did not succeed.
     """
     try:
         import vertexai
+        from google.cloud.aiplatform import initializer
     except ImportError:
         raise_vertex_import_error()
 
@@ -80,6 +94,9 @@ def init_vertexai(
         credentials=credentials,
     )
 
+    _, user_agent = _get_user_agent(module)
+    initializer.global_config.append_user_agent(user_agent)
+
 
 def get_client_info(module: Optional[str] = None) -> "ClientInfo":
     r"""Returns a custom user agent header.
@@ -98,13 +115,10 @@ def get_client_info(module: Optional[str] = None) -> "ClientInfo":
             "pip install google-api-core"
         ) from exc
 
-    langchain_version = metadata.version("langchain")
-    client_library_version = (
-        f"{langchain_version}-{module}" if module else langchain_version
-    )
+    client_library_version, user_agent = _get_user_agent(module)
     return ClientInfo(
         client_library_version=client_library_version,
-        user_agent=f"langchain/{client_library_version}",
+        user_agent=user_agent,
     )
 
 

libs/partners/google-vertexai/langchain_google_vertexai/_utils.py

@@ -2,7 +2,7 @@
 
 import dataclasses
 from importlib import metadata
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Any, Callable, Dict, Optional, Tuple, Union
 
 import google.api_core
 import proto  # type: ignore[import-untyped]
@@ -58,22 +58,38 @@ def raise_vertex_import_error(minimum_expected_version: str = "1.38.0") -> None:
     )
 
 
-def get_client_info(module: Optional[str] = None) -> "ClientInfo":
+def get_user_agent(module: Optional[str] = None) -> Tuple[str, str]:
     r"""Returns a custom user agent header.
 
     Args:
         module (Optional[str]):
             Optional. The module for a custom user agent header.
     Returns:
-        google.api_core.gapic_v1.client_info.ClientInfo
+        Tuple[str, str]: The client library version and user agent.
     """
-    langchain_version = metadata.version("langchain")
+    try:
+        langchain_version = metadata.version("langchain")
+    except metadata.PackageNotFoundError:
+        langchain_version = "0.0.0"
     client_library_version = (
         f"{langchain_version}-{module}" if module else langchain_version
     )
+    return client_library_version, f"langchain/{client_library_version}"
+
+
+def get_client_info(module: Optional[str] = None) -> ClientInfo:
+    r"""Returns ClientInfo with a custom user agent header.
+
+    Args:
+        module (Optional[str]):
+            Optional. The module for a custom user agent header.
+    Returns:
+        google.api_core.gapic_v1.client_info.ClientInfo
+    """
+    client_library_version, user_agent = get_user_agent(module)
     return ClientInfo(
         client_library_version=client_library_version,
-        user_agent=f"langchain/{client_library_version}",
+        user_agent=user_agent,
     )
 
 

libs/partners/google-vertexai/langchain_google_vertexai/chat_models.py

@@ -319,7 +319,7 @@ def validate_environment(cls, values: Dict) -> Dict:
 
         if safety_settings and not is_gemini:
             raise ValueError("Safety settings are only supported for Gemini models")
-
+        values["module"] = "vertexai-chat"
         cls._init_vertexai(values)
         if is_gemini:
             values["client"] = GenerativeModel(

libs/partners/google-vertexai/langchain_google_vertexai/llms.py

@@ -5,6 +5,7 @@
 
 import vertexai  # type: ignore[import-untyped]
 from google.api_core.client_options import ClientOptions
+from google.cloud.aiplatform import initializer
 from google.cloud.aiplatform.gapic import (
     PredictionServiceAsyncClient,
     PredictionServiceClient,
@@ -48,6 +49,7 @@
     create_retry_decorator,
     get_client_info,
     get_generation_info,
+    get_user_agent,
     is_codey_model,
     is_gemini_model,
 )
@@ -229,6 +231,8 @@ def _init_vertexai(cls, values: Dict) -> None:
             location=values.get("location"),
             credentials=values.get("credentials"),
         )
+        _, user_agent = get_user_agent(values.get("module"))
+        initializer.global_config.append_user_agent(user_agent)
         return None
 
     def _prepare_params(
@@ -291,6 +295,7 @@ def validate_environment(cls, values: Dict) -> Dict:
         model_name = values["model_name"]
         safety_settings = values["safety_settings"]
         is_gemini = is_gemini_model(values["model_name"])
+        values["module"] = f"vertexai-llm-{model_name}"
         cls._init_vertexai(values)
 
         if safety_settings and (not is_gemini or tuned_model_name):

libs/partners/google-vertexai/pyproject.toml

@@ -13,7 +13,7 @@ license = "MIT"
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"
 langchain-core = "^0.1.7"
-google-cloud-aiplatform = "^1.39.0"
+google-cloud-aiplatform = "^1.41.0"
 google-cloud-storage = "^2.14.0"
 
 [tool.poetry.group.test]