[langchain-ibm/feat]: Added support for Model Gateway (Embeddings), Added async methods (#88)

* Added support for Model Gateway Embeddings, added async methods

* Added more informing docstring under model parameter

* Added warning when user get ApiRequestFailure

* Added warning for watsonx_embed_gateway

Author: MateuszOssGit

libs/ibm/langchain_ibm/chat_models.py

@@ -82,9 +82,11 @@
 from typing_extensions import Self
 
 from langchain_ibm.utils import (
+    async_gateway_error_handler,
     check_duplicate_chat_params,
     check_for_attribute,
     extract_chat_params,
+    gateway_error_handler,
 )
 
 logger = logging.getLogger(__name__)
@@ -430,7 +432,16 @@ class ChatWatsonx(BaseChatModel):
     """Type of model to use."""
 
     model: Optional[str] = None
-    """Name of model for given provider or alias."""
+    """
+    Name or alias of the foundation model to use.  
+    When using IBM’s watsonx.ai Model Gateway (public preview), you can specify any 
+    supported third-party model—OpenAI, Anthropic, NVIDIA, Cerebras, or IBM’s own 
+    Granite series—via a single, OpenAI-compatible interface. Models must be explicitly 
+    provisioned (opt-in) through the Gateway to ensure secure, vendor-agnostic access 
+    and easy switch-over without reconfiguration.
+
+    For more details on configuration and usage, see IBM watsonx Model Gateway docs: https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-model-gateway.html?context=wx&audience=wdp
+    """
 
     deployment_id: Optional[str] = None
     """Type of deployed model to use."""
@@ -753,6 +764,20 @@ def validate_environment(self) -> Self:
 
         return self
 
+    @gateway_error_handler
+    def _call_model_gateway(self, *, model: str, messages: list, **params: Any) -> Any:
+        return self.watsonx_model_gateway.chat.completions.create(
+            model=model, messages=messages, **params
+        )
+
+    @async_gateway_error_handler
+    async def _acall_model_gateway(
+        self, *, model: str, messages: list, **params: Any
+    ) -> Any:
+        return await self.watsonx_model_gateway.chat.completions.acreate(
+            model=model, messages=messages, **params
+        )
+
     def _generate(
         self,
         messages: List[BaseMessage],
@@ -769,8 +794,9 @@ def _generate(
         message_dicts, params = self._create_message_dicts(messages, stop, **kwargs)
         updated_params = self._merge_params(params, kwargs)
         if self.watsonx_model_gateway is not None:
-            response = self.watsonx_model_gateway.chat.completions.create(
-                model=self.model, messages=message_dicts, **(kwargs | updated_params)
+            call_kwargs = {**kwargs, **updated_params}
+            response = self._call_model_gateway(
+                model=self.model, messages=message_dicts, **call_kwargs
             )
         else:
             response = self.watsonx_model.chat(
@@ -794,8 +820,9 @@ async def _agenerate(
         message_dicts, params = self._create_message_dicts(messages, stop, **kwargs)
         updated_params = self._merge_params(params, kwargs)
         if self.watsonx_model_gateway is not None:
-            response = await self.watsonx_model_gateway.chat.completions.acreate(
-                model=self.model, messages=message_dicts, **(kwargs | updated_params)
+            call_kwargs = {**kwargs, **updated_params}
+            response = await self._acall_model_gateway(
+                model=self.model, messages=message_dicts, **call_kwargs
             )
         else:
             response = await self.watsonx_model.achat(
@@ -815,7 +842,7 @@ def _stream(
 
         if self.watsonx_model_gateway is not None:
             call_kwargs = {**kwargs, **updated_params, "stream": True}
-            chunk_iter = self.watsonx_model_gateway.chat.completions.create(
+            chunk_iter = self._call_model_gateway(
                 model=self.model, messages=message_dicts, **call_kwargs
             )
         else:
@@ -872,7 +899,7 @@ async def _astream(
 
         if self.watsonx_model_gateway is not None:
             call_kwargs = {**kwargs, **updated_params, "stream": True}
-            chunk_iter = await self.watsonx_model_gateway.chat.completions.acreate(
+            chunk_iter = await self._acall_model_gateway(
                 model=self.model, messages=message_dicts, **call_kwargs
             )
         else:

libs/ibm/langchain_ibm/embeddings.py

@@ -3,12 +3,18 @@
 
 from ibm_watsonx_ai import APIClient, Credentials  # type: ignore
 from ibm_watsonx_ai.foundation_models.embeddings import Embeddings  # type: ignore
+from ibm_watsonx_ai.gateway import Gateway  # type: ignore
 from langchain_core.embeddings import Embeddings as LangChainEmbeddings
 from langchain_core.utils.utils import secret_from_env
 from pydantic import BaseModel, ConfigDict, Field, SecretStr, model_validator
 from typing_extensions import Self
 
-from langchain_ibm.utils import check_for_attribute, extract_params
+from langchain_ibm.utils import (
+    async_gateway_error_handler,
+    check_for_attribute,
+    extract_params,
+    gateway_error_handler,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -19,6 +25,18 @@ class WatsonxEmbeddings(BaseModel, LangChainEmbeddings):
     model_id: Optional[str] = None
     """Type of model to use."""
 
+    model: Optional[str] = None
+    """
+    Name or alias of the foundation model to use.  
+    When using IBM’s watsonx.ai Model Gateway (public preview), you can specify any 
+    supported third-party model—OpenAI, Anthropic, NVIDIA, Cerebras, or IBM’s own 
+    Granite series—via a single, OpenAI-compatible interface. Models must be explicitly 
+    provisioned (opt-in) through the Gateway to ensure secure, vendor-agnostic access 
+    and easy switch-over without reconfiguration.
+    
+    For more details on configuration and usage, see IBM watsonx Model Gateway docs: https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-model-gateway.html?context=wx&audience=wdp
+    """
+
     project_id: Optional[str] = None
     """ID of the Watson Studio project."""
 
@@ -70,10 +88,15 @@ class WatsonxEmbeddings(BaseModel, LangChainEmbeddings):
         * the path to a CA_BUNDLE file
         * the path of directory with certificates of trusted CAs
         * True - default path to truststore will be taken
-        * False - no verification will be made"""
+        * False - no verification will be made
+    """
 
     watsonx_embed: Embeddings = Field(default=None)  #: :meta private:
 
+    watsonx_embed_gateway: Gateway = Field(
+        default=None, exclude=True
+    )  #: :meta private:
+
     watsonx_client: Optional[APIClient] = Field(default=None)  #: :meta private:
 
     model_config = ConfigDict(
@@ -85,6 +108,12 @@ class WatsonxEmbeddings(BaseModel, LangChainEmbeddings):
     @model_validator(mode="after")
     def validate_environment(self) -> Self:
         """Validate that credentials and python package exists in environment."""
+        if self.watsonx_embed_gateway is not None:
+            raise NotImplementedError(
+                "Passing the 'watsonx_embed_gateway' parameter to the "
+                "WatsonxEmbeddings constructor is not supported yet."
+            )
+
         if isinstance(self.watsonx_embed, Embeddings):
             self.model_id = getattr(self.watsonx_embed, "model_id")
             self.project_id = getattr(
@@ -98,17 +127,36 @@ def validate_environment(self) -> Self:
             self.params = getattr(self.watsonx_embed, "params")
 
         elif isinstance(self.watsonx_client, APIClient):
-            watsonx_embed = Embeddings(
-                model_id=self.model_id,
-                params=self.params,
-                api_client=self.watsonx_client,
-                project_id=self.project_id,
-                space_id=self.space_id,
-                verify=self.verify,
-            )
-            self.watsonx_embed = watsonx_embed
+            if sum(map(bool, (self.model, self.model_id))) != 1:
+                raise ValueError(
+                    "The parameters 'model' and 'model_id' are mutually exclusive. "
+                    "Please specify exactly one of these parameters when "
+                    "initializing WatsonxEmbeddings."
+                )
+            if self.model is not None:
+                watsonx_embed_gateway = Gateway(
+                    api_client=self.watsonx_client,
+                    verify=self.verify,
+                )
+                self.watsonx_embed_gateway = watsonx_embed_gateway
+            else:
+                watsonx_embed = Embeddings(
+                    model_id=self.model_id,
+                    params=self.params,
+                    api_client=self.watsonx_client,
+                    project_id=self.project_id,
+                    space_id=self.space_id,
+                    verify=self.verify,
+                )
+                self.watsonx_embed = watsonx_embed
 
         else:
+            if sum(map(bool, (self.model, self.model_id))) != 1:
+                raise ValueError(
+                    "The parameters 'model' and 'model_id' are mutually exclusive. "
+                    "Please specify exactly one of these parameters when "
+                    "initializing WatsonxEmbeddings."
+                )
             check_for_attribute(self.url, "url", "WATSONX_URL")
 
             if "cloud.ibm.com" in self.url.get_secret_value():
@@ -157,26 +205,77 @@ def validate_environment(self) -> Self:
                 version=self.version.get_secret_value() if self.version else None,
                 verify=self.verify,
             )
+            if self.model is not None:
+                watsonx_embed_gateway = Gateway(
+                    credentials=credentials,
+                    verify=self.verify,
+                )
+                self.watsonx_embed_gateway = watsonx_embed_gateway
 
-            watsonx_embed = Embeddings(
-                model_id=self.model_id,
-                params=self.params,
-                credentials=credentials,
-                project_id=self.project_id,
-                space_id=self.space_id,
-            )
+            else:
+                watsonx_embed = Embeddings(
+                    model_id=self.model_id,
+                    params=self.params,
+                    credentials=credentials,
+                    project_id=self.project_id,
+                    space_id=self.space_id,
+                )
 
-            self.watsonx_embed = watsonx_embed
+                self.watsonx_embed = watsonx_embed
 
         return self
 
+    @gateway_error_handler
+    def _call_model_gateway(
+        self, *, model: str, texts: List[str], **params: Any
+    ) -> Any:
+        return self.watsonx_embed_gateway.embeddings.create(
+            model=model, input=texts, **params
+        )
+
+    @async_gateway_error_handler
+    async def _acall_model_gateway(
+        self, *, model: str, texts: List[str], **params: Any
+    ) -> Any:
+        return await self.watsonx_embed_gateway.embeddings.acreate(
+            model=model, input=texts, **params
+        )
+
     def embed_documents(self, texts: List[str], **kwargs: Any) -> List[List[float]]:
         """Embed search docs."""
         params = extract_params(kwargs, self.params)
-        return self.watsonx_embed.embed_documents(
-            texts=texts, **(kwargs | {"params": params})
-        )
+        if self.watsonx_embed_gateway is not None:
+            call_kwargs = {**kwargs, **params}
+            embed_response = self._call_model_gateway(
+                model=self.model, texts=texts, **call_kwargs
+            )
+            return [embedding["embedding"] for embedding in embed_response["data"]]
+        else:
+            return self.watsonx_embed.embed_documents(
+                texts=texts, **(kwargs | {"params": params})
+            )
+
+    async def aembed_documents(
+        self, texts: List[str], **kwargs: Any
+    ) -> List[List[float]]:
+        """Asynchronous Embed search docs."""
+        params = extract_params(kwargs, self.params)
+        if self.watsonx_embed_gateway is not None:
+            call_kwargs = {**kwargs, **params}
+            embed_response = await self._acall_model_gateway(
+                model=self.model, texts=texts, **call_kwargs
+            )
+            return [embedding["embedding"] for embedding in embed_response["data"]]
+        else:
+            return await self.watsonx_embed.aembed_documents(
+                texts=texts, **(kwargs | {"params": params})
+            )
 
     def embed_query(self, text: str, **kwargs: Any) -> List[float]:
         """Embed query text."""
         return self.embed_documents([text], **kwargs)[0]
+
+    async def aembed_query(self, text: str, **kwargs: Any) -> List[float]:
+        """Asynchronous Embed query text."""
+        embeddings = await self.aembed_documents([text], **kwargs)
+        return embeddings[0]

libs/ibm/langchain_ibm/utils.py

@@ -1,12 +1,17 @@
+import functools
+import logging
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
 
 from ibm_watsonx_ai.foundation_models.schema import BaseSchema  # type: ignore
+from ibm_watsonx_ai.wml_client_error import ApiRequestFailure  # type: ignore
 from pydantic import SecretStr
 
 if TYPE_CHECKING:
     from langchain_ibm.toolkit import WatsonxTool
 
+logger = logging.getLogger(__name__)
+
 
 def check_for_attribute(value: SecretStr | None, key: str, env_key: str) -> None:
     if not value or not value.get_secret_value():
@@ -135,3 +140,48 @@ def parse_parameters(input_schema: dict | None) -> dict:
         },
     }
     return watsonx_tool
+
+
+def gateway_error_handler(func: Callable) -> Callable:
+    """Decorator to catch ApiRequestFailure on Model Gateway calls
+    and log a uniform warning."""
+
+    @functools.wraps(func)
+    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+        try:
+            return func(self, *args, **kwargs)
+        except ApiRequestFailure:
+            if any(
+                hasattr(self, attr)
+                for attr in ("watsonx_model_gateway", "watsonx_embed_gateway")
+            ):
+                logger.warning(
+                    "You are calling the Model Gateway endpoint using the 'model' "
+                    "parameter. Please ensure this model is registered with the "
+                    "Gateway. If you intend to use a watsonx.ai–hosted model, pass "
+                    "'model_id' instead of 'model'."
+                )
+            raise
+
+    return wrapper
+
+
+def async_gateway_error_handler(func: Callable) -> Callable:
+    """Async decorator to catch ApiRequestFailure on Model Gateway calls
+    and log a uniform warning."""
+
+    @functools.wraps(func)
+    async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+        try:
+            return await func(self, *args, **kwargs)
+        except ApiRequestFailure:
+            if getattr(self, "watsonx_model_gateway", None) is not None:
+                logger.warning(
+                    "You are calling the Model Gateway endpoint using the 'model' "
+                    "parameter. Please ensure this model is registered with the "
+                    "Gateway. If you intend to use a watsonx.ai–hosted model, pass "
+                    "'model_id' instead of 'model'."
+                )
+            raise
+
+    return wrapper