Skip to content

fix: Add IBM watsonx.ai support to EmbeddingModel #10677

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
edwinjosechittilappilly merged 13 commits into from
Nov 21, 2025
Merged

fix: Add IBM watsonx.ai support to EmbeddingModel #10677

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
55cf5d9
Add IBM watsonx.ai support to EmbeddingModel
edwinjosechittilappilly Nov 20, 2025
d927cc6
update watsonx default models
edwinjosechittilappilly Nov 21, 2025
ce8aabd
Merge branch 'main' into fix-emebddings-component
edwinjosechittilappilly Nov 21, 2025
ec14524
update index
edwinjosechittilappilly Nov 21, 2025
9fad275
[autofix.ci] apply automated fixes
autofix-ci[bot] Nov 21, 2025
7079625
[autofix.ci] apply automated fixes (attempt 2/3)
autofix-ci[bot] Nov 21, 2025
4dc4ec9
[autofix.ci] apply automated fixes (attempt 3/3)
autofix-ci[bot] Nov 21, 2025
5444d54
Update test_embedding_model_component.py
edwinjosechittilappilly Nov 21, 2025
04f4a78
Merge branch 'main' into fix-emebddings-component
edwinjosechittilappilly Nov 21, 2025
9f3541f
Update component_index.json
edwinjosechittilappilly Nov 21, 2025
968a006
[autofix.ci] apply automated fixes
autofix-ci[bot] Nov 21, 2025
b2ee642
[autofix.ci] apply automated fixes (attempt 2/3)
autofix-ci[bot] Nov 21, 2025
3f93924
[autofix.ci] apply automated fixes (attempt 3/3)
autofix-ci[bot] Nov 21, 2025
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
52 changes: 49 additions & 3 deletions src/backend/base/langflow/initial_setup/starter_projects/Nvidia Remix.json
View file
Open in desktop

Large diffs are not rendered by default.

73 changes: 57 additions & 16 deletions src/backend/tests/unit/components/models_and_agents/test_embedding_model_component.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,9 @@
from unittest.mock import MagicMock, patch

import pytest
from lfx.components.models_and_agents.embedding_model import (
OPENAI_EMBEDDING_MODEL_NAMES,
WATSONX_EMBEDDING_MODEL_NAMES,
EmbeddingModelComponent,
)
from lfx.base.models.openai_constants import OPENAI_EMBEDDING_MODEL_NAMES
from lfx.base.models.watsonx_constants import WATSONX_EMBEDDING_MODEL_NAMES
from lfx.components.models_and_agents.embedding_model import EmbeddingModelComponent

from tests.base import ComponentTestBaseWithClient

Expand Down Expand Up @@ -36,10 +34,12 @@ async def test_update_build_config_openai(self, component_class, default_kwargs)
build_config = {
"model": {"options": [], "value": ""},
"api_key": {"display_name": "API Key", "required": True, "show": True},
"api_base": {"display_name": "API Base URL", "value": ""},
"api_base": {"display_name": "API Base URL", "value": "", "advanced": False, "show": False},
"base_url_ibm_watsonx": {"show": False},
"project_id": {"show": False},
"ollama_base_url": {"show": False},
"truncate_input_tokens": {"show": False},
"input_text": {"show": False},
}
updated_config = await component.update_build_config(build_config, "OpenAI", "provider")
assert updated_config["model"]["options"] == OPENAI_EMBEDDING_MODEL_NAMES
Expand All @@ -48,12 +48,16 @@ async def test_update_build_config_openai(self, component_class, default_kwargs)
assert updated_config["api_key"]["required"] is True
assert updated_config["api_key"]["show"] is True
assert updated_config["api_base"]["display_name"] == "OpenAI API Base URL"
assert updated_config["api_base"]["advanced"] is True
assert updated_config["api_base"]["show"] is True
assert updated_config["project_id"]["show"] is False
assert updated_config["base_url_ibm_watsonx"]["show"] is False
assert updated_config["ollama_base_url"]["show"] is False
assert updated_config["truncate_input_tokens"]["show"] is False
assert updated_config["input_text"]["show"] is False

@patch("lfx.components.models.embedding_model.get_ollama_models")
@patch("lfx.components.models.embedding_model.is_valid_ollama_url")
@patch("lfx.components.models_and_agents.embedding_model.get_ollama_models")
@patch("lfx.components.models_and_agents.embedding_model.is_valid_ollama_url")
async def test_update_build_config_ollama(
self, mock_is_valid_url, mock_get_ollama_models, component_class, default_kwargs
):
Expand All @@ -66,10 +70,12 @@ async def test_update_build_config_ollama(
build_config = {
"model": {"options": [], "value": ""},
"api_key": {"display_name": "API Key", "required": True, "show": True},
"api_base": {"display_name": "API Base URL", "value": ""},
"api_base": {"display_name": "API Base URL", "value": "", "show": False},
"project_id": {"show": False},
"base_url_ibm_watsonx": {"show": False},
"ollama_base_url": {"show": False},
"truncate_input_tokens": {"show": False},
"input_text": {"show": False},
}
updated_config = await component.update_build_config(build_config, "Ollama", "provider")
assert updated_config["model"]["options"] == ["nomic-embed-text", "mxbai-embed-large"]
Expand All @@ -81,16 +87,23 @@ async def test_update_build_config_ollama(
assert updated_config["project_id"]["show"] is False
assert updated_config["base_url_ibm_watsonx"]["show"] is False
assert updated_config["ollama_base_url"]["show"] is True
assert updated_config["truncate_input_tokens"]["show"] is False
assert updated_config["input_text"]["show"] is False

@patch.object(EmbeddingModelComponent, "fetch_ibm_models")
async def test_update_build_config_watsonx(self, mock_fetch_ibm_models, component_class, default_kwargs):
mock_fetch_ibm_models.return_value = WATSONX_EMBEDDING_MODEL_NAMES

async def test_update_build_config_watsonx(self, component_class, default_kwargs):
component = component_class(**default_kwargs)
build_config = {
"model": {"options": [], "value": ""},
"api_key": {"display_name": "API Key", "required": True, "show": True},
"api_base": {"display_name": "API Base URL", "value": ""},
"api_base": {"display_name": "API Base URL", "value": "", "show": False},
"project_id": {"show": False},
"base_url_ibm_watsonx": {"show": False},
"ollama_base_url": {"show": False},
"truncate_input_tokens": {"show": False},
"input_text": {"show": False},
}
updated_config = await component.update_build_config(build_config, "IBM watsonx.ai", "provider")
assert updated_config["model"]["options"] == WATSONX_EMBEDDING_MODEL_NAMES
Expand All @@ -102,6 +115,8 @@ async def test_update_build_config_watsonx(self, component_class, default_kwargs
assert updated_config["base_url_ibm_watsonx"]["show"] is True
assert updated_config["project_id"]["show"] is True
assert updated_config["ollama_base_url"]["show"] is False
assert updated_config["truncate_input_tokens"]["show"] is True
assert updated_config["input_text"]["show"] is True

@patch("lfx.components.models_and_agents.embedding_model.OpenAIEmbeddings")
async def test_build_embeddings_openai(self, mock_openai_embeddings, component_class, default_kwargs):
Expand Down Expand Up @@ -145,8 +160,9 @@ async def test_build_embeddings_ollama(self, mock_ollama_embeddings, component_c
kwargs = default_kwargs.copy()
kwargs["provider"] = "Ollama"
kwargs["model"] = "nomic-embed-text"
kwargs["model_kwargs"] = {}
component = component_class(**kwargs)
component.api_base = "http://localhost:11434"
component.ollama_base_url = "http://localhost:11434"

# Build the embeddings
embeddings = component.build_embeddings()
Expand All @@ -158,27 +174,52 @@ async def test_build_embeddings_ollama(self, mock_ollama_embeddings, component_c
)
assert embeddings == mock_instance

@patch("ibm_watsonx_ai.APIClient")
@patch("ibm_watsonx_ai.Credentials")
@patch("langchain_ibm.WatsonxEmbeddings")
async def test_build_embeddings_watsonx(self, mock_watsonx_embeddings, component_class, default_kwargs):
# Setup mock
async def test_build_embeddings_watsonx(
self, mock_watsonx_embeddings, mock_credentials, mock_api_client, component_class, default_kwargs
):
# Setup mocks
mock_instance = MagicMock()
mock_watsonx_embeddings.return_value = mock_instance
mock_cred_instance = MagicMock()
mock_credentials.return_value = mock_cred_instance
mock_client_instance = MagicMock()
mock_api_client.return_value = mock_client_instance

# Create and configure the component
kwargs = default_kwargs.copy()
kwargs["provider"] = "IBM watsonx.ai"
kwargs["model"] = "ibm/granite-embedding-125m-english"
component = component_class(**kwargs)
component.project_id = "test-project-id"
component.truncate_input_tokens = 200
component.input_text = True

# Build the embeddings
embeddings = component.build_embeddings()

# Verify Credentials was created correctly
mock_credentials.assert_called_once_with(
api_key="test-api-key", # pragma:allowlist secret
url="https://us-south.ml.cloud.ibm.com",
)

# Verify APIClient was created with credentials
mock_api_client.assert_called_once_with(mock_cred_instance)

# Verify the WatsonxEmbeddings was called with the correct parameters
from ibm_watsonx_ai.metanames import EmbedTextParamsMetaNames

expected_params = {
EmbedTextParamsMetaNames.TRUNCATE_INPUT_TOKENS: 200,
EmbedTextParamsMetaNames.RETURN_OPTIONS: {"input_text": True},
}
mock_watsonx_embeddings.assert_called_once_with(
model_id="ibm/granite-embedding-125m-english",
url="https://us-south.ml.cloud.ibm.com",
apikey="test-api-key", # pragma:allowlist secret
params=expected_params,
watsonx_client=mock_client_instance,
project_id="test-project-id",
)
assert embeddings == mock_instance
Expand Down
2 changes: 1 addition & 1 deletion src/lfx/src/lfx/_assets/component_index.json
View file
Open in desktop

Large diffs are not rendered by default.

31 changes: 23 additions & 8 deletions src/lfx/src/lfx/base/models/watsonx_constants.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,15 +1,30 @@
from .model_metadata import create_model_metadata

# Granite Embedding models
WATSONX_EMBEDDING_MODELS_DETAILED = [
create_model_metadata(provider="IBM Watsonx", name="ibm/granite-embedding-125m-english", icon="IBMWatsonx"),
create_model_metadata(provider="IBM Watsonx", name="ibm/granite-embedding-278m-multilingual", icon="IBMWatsonx"),
create_model_metadata(provider="IBM Watsonx", name="ibm/granite-embedding-30m-english", icon="IBMWatsonx"),
create_model_metadata(provider="IBM Watsonx", name="ibm/granite-embedding-107m-multilingual", icon="IBMWatsonx"),
create_model_metadata(provider="IBM Watsonx", name="ibm/granite-embedding-30m-sparse", icon="IBMWatsonx"),
WATSONX_DEFAULT_EMBEDDING_MODELS = [
create_model_metadata(
provider="IBM Watsonx",
name="sentence-transformers/all-minilm-l12-v2",
icon="WatsonxAI",
),
create_model_metadata(
provider="IBM Watsonx",
name="ibm/slate-125m-english-rtrvr-v2",
icon="WatsonxAI",
),
create_model_metadata(
provider="IBM Watsonx",
name="ibm/slate-30m-english-rtrvr-v2",
icon="WatsonxAI",
),
create_model_metadata(
provider="IBM Watsonx",
name="intfloat/multilingual-e5-large",
icon="WatsonxAI",
),
]

WATSONX_EMBEDDING_MODEL_NAMES = [metadata["name"] for metadata in WATSONX_EMBEDDING_MODELS_DETAILED]

WATSONX_EMBEDDING_MODEL_NAMES = [metadata["name"] for metadata in WATSONX_DEFAULT_EMBEDDING_MODELS]

IBM_WATSONX_URLS = [
"https://us-south.ml.cloud.ibm.com",
Expand Down
80 changes: 68 additions & 12 deletions src/lfx/src/lfx/components/models_and_agents/embedding_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,16 @@
from typing import Any

import requests
from ibm_watsonx_ai.metanames import EmbedTextParamsMetaNames
from langchain_openai import OpenAIEmbeddings

from lfx.base.embeddings.model import LCEmbeddingsModel
from lfx.base.models.model_utils import get_ollama_models, is_valid_ollama_url
from lfx.base.models.openai_constants import OPENAI_EMBEDDING_MODEL_NAMES
from lfx.base.models.watsonx_constants import IBM_WATSONX_URLS, WATSONX_EMBEDDING_MODEL_NAMES
from lfx.base.models.watsonx_constants import (
IBM_WATSONX_URLS,
WATSONX_EMBEDDING_MODEL_NAMES,
)
from lfx.field_typing import Embeddings
from lfx.io import (
BoolInput,
Expand Down Expand Up @@ -77,6 +82,8 @@ class EmbeddingModelComponent(LCEmbeddingsModel):
options=OPENAI_EMBEDDING_MODEL_NAMES,
value=OPENAI_EMBEDDING_MODEL_NAMES[0],
info="Select the embedding model to use",
real_time_refresh=True,
refresh_button=True,
),
SecretStrInput(
name="api_key",
Expand Down Expand Up @@ -110,8 +117,40 @@ class EmbeddingModelComponent(LCEmbeddingsModel):
advanced=True,
info="Additional keyword arguments to pass to the model.",
),
IntInput(
name="truncate_input_tokens",
display_name="Truncate Input Tokens",
advanced=True,
value=200,
show=False,
),
BoolInput(
name="input_text",
display_name="Include the original text in the output",
value=True,
advanced=True,
show=False,
),
]

@staticmethod
def fetch_ibm_models(base_url: str) -> list[str]:
"""Fetch available models from the watsonx.ai API."""
try:
endpoint = f"{base_url}/ml/v1/foundation_model_specs"
params = {
"version": "2024-09-16",
"filters": "function_embedding,!lifecycle_withdrawn:and",
Copy link

Copilot AI Nov 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The filter syntax "function_embedding,!lifecycle_withdrawn:and" has an unusual :and suffix at the end. Comparing with the similar implementation in language_model.py (line 57), which uses "function_text_chat,!lifecycle_withdrawn" without the :and suffix, this appears to be inconsistent. Consider removing the :and suffix or verifying the correct filter syntax with the IBM watsonx.ai API documentation.

Copilot uses AI. Check for mistakes.
}
response = requests.get(endpoint, params=params, timeout=10)
Copy link

Copilot AI Nov 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The API request to fetch IBM models is made without authentication. This endpoint likely requires authentication but the request doesn't include any API key or credentials. Consider adding authentication headers or verifying if this endpoint is intended to be publicly accessible.

Copilot uses AI. Check for mistakes.
response.raise_for_status()
data = response.json()
models = [model["model_id"] for model in data.get("resources", [])]
return sorted(models)
except Exception: # noqa: BLE001
logger.exception("Error fetching models")
return WATSONX_EMBEDDING_MODEL_NAMES

def build_embeddings(self) -> Embeddings:
provider = self.provider
model = self.model
Expand Down Expand Up @@ -188,15 +227,26 @@ def build_embeddings(self) -> Embeddings:
msg = "Project ID is required for IBM watsonx.ai provider"
raise ValueError(msg)

from ibm_watsonx_ai import APIClient, Credentials

credentials = Credentials(
api_key=self.api_key,
url=base_url_ibm_watsonx or "https://us-south.ml.cloud.ibm.com",
)

api_client = APIClient(credentials)

params = {
"model_id": model,
"url": base_url_ibm_watsonx or "https://us-south.ml.cloud.ibm.com",
"apikey": api_key,
EmbedTextParamsMetaNames.TRUNCATE_INPUT_TOKENS: self.truncate_input_tokens,
EmbedTextParamsMetaNames.RETURN_OPTIONS: {"input_text": self.input_text},
}

params["project_id"] = project_id

return WatsonxEmbeddings(**params)
return WatsonxEmbeddings(
model_id=model,
params=params,
watsonx_client=api_client,
project_id=project_id,
)

msg = f"Unknown provider: {provider}"
raise ValueError(msg)
Expand All @@ -217,7 +267,8 @@ async def update_build_config(
build_config["ollama_base_url"]["show"] = False
build_config["project_id"]["show"] = False
build_config["base_url_ibm_watsonx"]["show"] = False

build_config["truncate_input_tokens"]["show"] = False
build_config["input_text"]["show"] = False
elif field_value == "Ollama":
build_config["ollama_base_url"]["show"] = True

Expand All @@ -238,7 +289,8 @@ async def update_build_config(
else:
build_config["model"]["options"] = []
build_config["model"]["value"] = ""

build_config["truncate_input_tokens"]["show"] = False
build_config["input_text"]["show"] = False
build_config["api_key"]["display_name"] = "API Key (Optional)"
build_config["api_key"]["required"] = False
build_config["api_key"]["show"] = False
Expand All @@ -247,16 +299,20 @@ async def update_build_config(
build_config["base_url_ibm_watsonx"]["show"] = False

elif field_value == "IBM watsonx.ai":
build_config["model"]["options"] = WATSONX_EMBEDDING_MODEL_NAMES
build_config["model"]["value"] = WATSONX_EMBEDDING_MODEL_NAMES[0]
build_config["model"]["options"] = self.fetch_ibm_models(base_url=self.base_url_ibm_watsonx)
build_config["model"]["value"] = self.fetch_ibm_models(base_url=self.base_url_ibm_watsonx)[0]
Comment on lines +302 to +303
Copy link

Copilot AI Nov 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The fetch_ibm_models method is called twice with the same base_url parameter on consecutive lines. This results in duplicate API requests. Consider storing the result in a variable and reusing it:

ibm_models = self.fetch_ibm_models(base_url=self.base_url_ibm_watsonx)
build_config["model"]["options"] = ibm_models
build_config["model"]["value"] = ibm_models[0]

Copilot uses AI. Check for mistakes.
Copy link

Copilot AI Nov 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Potential IndexError if fetch_ibm_models returns an empty list. The code accesses [0] without checking if the list is non-empty. Consider adding a check or providing a fallback value:

ibm_models = self.fetch_ibm_models(base_url=self.base_url_ibm_watsonx)
build_config["model"]["options"] = ibm_models
build_config["model"]["value"] = ibm_models[0] if ibm_models else WATSONX_EMBEDDING_MODEL_NAMES[0]

Copilot uses AI. Check for mistakes.
build_config["api_key"]["display_name"] = "IBM watsonx.ai API Key"
build_config["api_key"]["required"] = True
build_config["api_key"]["show"] = True
build_config["api_base"]["show"] = False
build_config["ollama_base_url"]["show"] = False
build_config["base_url_ibm_watsonx"]["show"] = True
build_config["project_id"]["show"] = True

build_config["truncate_input_tokens"]["show"] = True
build_config["input_text"]["show"] = True
elif field_name == "base_url_ibm_watsonx":
build_config["model"]["options"] = self.fetch_ibm_models(base_url=field_value)
build_config["model"]["value"] = self.fetch_ibm_models(base_url=field_value)[0]
Comment on lines +314 to +315
Copy link

Copilot AI Nov 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The fetch_ibm_models method is called twice with the same field_value parameter on consecutive lines. This results in duplicate API requests. Consider storing the result in a variable and reusing it:

ibm_models = self.fetch_ibm_models(base_url=field_value)
build_config["model"]["options"] = ibm_models
build_config["model"]["value"] = ibm_models[0]

Copilot uses AI. Check for mistakes.
Copy link

Copilot AI Nov 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Potential IndexError if fetch_ibm_models returns an empty list. The code accesses [0] without checking if the list is non-empty. Consider adding a check or providing a fallback value:

ibm_models = self.fetch_ibm_models(base_url=field_value)
build_config["model"]["options"] = ibm_models
build_config["model"]["value"] = ibm_models[0] if ibm_models else WATSONX_EMBEDDING_MODEL_NAMES[0]

Copilot uses AI. Check for mistakes.
elif field_name == "ollama_base_url":
# # Refresh Ollama models when base URL changes
# if hasattr(self, "provider") and self.provider == "Ollama":
Expand Down
Loading