import logging
from typing import Any, Dict, List, Optional, Union
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 (
async_gateway_error_handler,
check_for_attribute,
extract_params,
gateway_error_handler,
)
logger = logging.getLogger(__name__)
[docs]
class WatsonxEmbeddings(BaseModel, LangChainEmbeddings):
"""IBM watsonx.ai embedding models."""
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."""
space_id: Optional[str] = None
"""ID of the Watson Studio space."""
url: SecretStr = Field(
alias="url",
default_factory=secret_from_env("WATSONX_URL", default=None), # type: ignore[assignment]
)
"""URL to the Watson Machine Learning or CPD instance."""
apikey: Optional[SecretStr] = Field(
alias="apikey", default_factory=secret_from_env("WATSONX_APIKEY", default=None)
)
"""API key to the Watson Machine Learning or CPD instance."""
token: Optional[SecretStr] = Field(
alias="token", default_factory=secret_from_env("WATSONX_TOKEN", default=None)
)
"""Token to the CPD instance."""
password: Optional[SecretStr] = Field(
alias="password",
default_factory=secret_from_env("WATSONX_PASSWORD", default=None),
)
"""Password to the CPD instance."""
username: Optional[SecretStr] = Field(
alias="username",
default_factory=secret_from_env("WATSONX_USERNAME", default=None),
)
"""Username to the CPD instance."""
instance_id: Optional[SecretStr] = Field(
alias="instance_id",
default_factory=secret_from_env("WATSONX_INSTANCE_ID", default=None),
)
"""Instance_id of the CPD instance."""
version: Optional[SecretStr] = None
"""Version of the CPD instance."""
params: Optional[Dict] = None
"""Model parameters to use during request generation."""
verify: Union[str, bool, None] = None
"""You can pass one of following as verify:
* 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
"""
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(
extra="forbid",
arbitrary_types_allowed=True,
protected_namespaces=(),
)
@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(
getattr(self.watsonx_embed, "_client"),
"default_project_id",
)
self.space_id = getattr(
getattr(self.watsonx_embed, "_client"), "default_space_id"
)
self.params = getattr(self.watsonx_embed, "params")
elif isinstance(self.watsonx_client, APIClient):
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():
if not self.token and not self.apikey:
raise ValueError(
"Did not find 'apikey' or 'token',"
" please add an environment variable"
" `WATSONX_APIKEY` or 'WATSONX_TOKEN' "
"which contains it,"
" or pass 'apikey' or 'token'"
" as a named parameter."
)
else:
if not self.token and not self.password and not self.apikey:
raise ValueError(
"Did not find 'token', 'password' or 'apikey',"
" please add an environment variable"
" `WATSONX_TOKEN`, 'WATSONX_PASSWORD' or 'WATSONX_APIKEY' "
"which contains it,"
" or pass 'token', 'password' or 'apikey'"
" as a named parameter."
)
elif self.token:
check_for_attribute(self.token, "token", "WATSONX_TOKEN")
elif self.password:
check_for_attribute(self.password, "password", "WATSONX_PASSWORD")
check_for_attribute(self.username, "username", "WATSONX_USERNAME")
elif self.apikey:
check_for_attribute(self.apikey, "apikey", "WATSONX_APIKEY")
check_for_attribute(self.username, "username", "WATSONX_USERNAME")
if not self.instance_id:
check_for_attribute(
self.instance_id, "instance_id", "WATSONX_INSTANCE_ID"
)
credentials = Credentials(
url=self.url.get_secret_value() if self.url else None,
api_key=self.apikey.get_secret_value() if self.apikey else None,
token=self.token.get_secret_value() if self.token else None,
password=self.password.get_secret_value() if self.password else None,
username=self.username.get_secret_value() if self.username else None,
instance_id=self.instance_id.get_secret_value()
if self.instance_id
else None,
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
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
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
)
[docs]
def embed_documents(self, texts: List[str], **kwargs: Any) -> List[List[float]]:
"""Embed search docs."""
params = extract_params(kwargs, self.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})
)
[docs]
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})
)
[docs]
def embed_query(self, text: str, **kwargs: Any) -> List[float]:
"""Embed query text."""
return self.embed_documents([text], **kwargs)[0]
[docs]
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]
