From 2c00a1dfc71ec513353aa10159d980b68c3c7e70 Mon Sep 17 00:00:00 2001
From: Karol Zmorski <46227818+KarolZm@users.noreply.github.com>
Date: Thu, 27 Feb 2025 22:04:56 +0100
Subject: [PATCH] feat: Add `ibm-watsonx-ai` rerank integration (#17900)
---
.../node_postprocessor/ibm_watsonx.ipynb | 603 ++++++++++++++++++
.../llama-index-postprocessor-ibm/.gitignore | 153 +++++
.../llama-index-postprocessor-ibm/BUILD | 3 +
.../llama-index-postprocessor-ibm/Makefile | 17 +
.../llama-index-postprocessor-ibm/README.md | 224 +++++++
.../llama_index/postprocessor/ibm/BUILD | 1 +
.../llama_index/postprocessor/ibm/__init__.py | 4 +
.../llama_index/postprocessor/ibm/base.py | 265 ++++++++
.../llama_index/postprocessor/ibm/utils.py | 111 ++++
.../pyproject.toml | 60 ++
.../llama-index-postprocessor-ibm/tests/BUILD | 6 +
.../tests/__init__.py | 0
.../tests/test_ibm.py | 34 +
.../tests/test_postprocessor_ibm_rerank.py | 7 +
14 files changed, 1488 insertions(+)
create mode 100644 docs/docs/examples/node_postprocessor/ibm_watsonx.ipynb
create mode 100644 llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/.gitignore
create mode 100644 llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/BUILD
create mode 100644 llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/Makefile
create mode 100644 llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/README.md
create mode 100644 llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/llama_index/postprocessor/ibm/BUILD
create mode 100644 llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/llama_index/postprocessor/ibm/__init__.py
create mode 100644 llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/llama_index/postprocessor/ibm/base.py
create mode 100644 llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/llama_index/postprocessor/ibm/utils.py
create mode 100644 llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/pyproject.toml
create mode 100644 llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/tests/BUILD
create mode 100644 llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/tests/__init__.py
create mode 100644 llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/tests/test_ibm.py
create mode 100644 llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/tests/test_postprocessor_ibm_rerank.py
diff --git a/docs/docs/examples/node_postprocessor/ibm_watsonx.ipynb b/docs/docs/examples/node_postprocessor/ibm_watsonx.ipynb
new file mode 100644
index 0000000000..2761e18cf6
--- /dev/null
+++ b/docs/docs/examples/node_postprocessor/ibm_watsonx.ipynb
@@ -0,0 +1,603 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "id": "c9d89b5c",
+ "metadata": {},
+ "source": [
+ "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/docs/examples/node_postprocessor/ibm_watsonx.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "70996d8a",
+ "metadata": {},
+ "source": [
+ "# IBM watsonx.ai\n",
+ "\n",
+ ">WatsonxRerank is a wrapper for IBM [watsonx.ai](https://www.ibm.com/products/watsonx-ai) Rerank.\n",
+ "\n",
+ "The aim of these examples is to show how to take advantage of `watsonx.ai` Rerank, Embeddings and LLMs using the `LlamaIndex` postprocessor API."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "ea35b2b7",
+ "metadata": {},
+ "source": [
+ "## Setting up\n",
+ "\n",
+ "Install required packages:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "2f1fff4e",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%pip install -qU llama-index\n",
+ "%pip install -qU llama-index-llms-ibm\n",
+ "%pip install -qU llama-index-postprocessor-ibm\n",
+ "%pip install -qU llama-index-embeddings-ibm"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "f406e092",
+ "metadata": {},
+ "source": [
+ "The cell below defines the credentials required to work with watsonx Foundation Models, Embeddings and Rerank.\n",
+ "\n",
+ "**Action:** Provide the IBM Cloud user API key. For details, see\n",
+ "[Managing user API keys](https://cloud.ibm.com/docs/account?topic=account-userapikey&interface=ui)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "11d572a1",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "from getpass import getpass\n",
+ "\n",
+ "watsonx_api_key = getpass()\n",
+ "os.environ[\"WATSONX_APIKEY\"] = watsonx_api_key"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "c59782a7",
+ "metadata": {},
+ "source": [
+ "Additionally, you can pass additional secrets as an environment variable:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "f98c573c",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "\n",
+ "os.environ[\"WATSONX_URL\"] = \"your service instance url\"\n",
+ "os.environ[\"WATSONX_TOKEN\"] = \"your token for accessing the CPD cluster\"\n",
+ "os.environ[\"WATSONX_PASSWORD\"] = \"your password for accessing the CPD cluster\"\n",
+ "os.environ[\"WATSONX_USERNAME\"] = \"your username for accessing the CPD cluster\"\n",
+ "os.environ[\n",
+ " \"WATSONX_INSTANCE_ID\"\n",
+ "] = \"your instance_id for accessing the CPD cluster\""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "b0de9336",
+ "metadata": {},
+ "source": [
+ "**Note**: \n",
+ "\n",
+ "- To provide context for the API call, you must pass the `project_id` or `space_id`. To get your project or space ID, open your project or space, go to the **Manage** tab, and click **General**. For more information see: [Project documentation](https://www.ibm.com/docs/en/watsonx-as-a-service?topic=projects) or [Deployment space documentation](https://www.ibm.com/docs/en/watsonx/saas?topic=spaces-creating-deployment).\n",
+ "- Depending on the region of your provisioned service instance, use one of the urls listed in [watsonx.ai API Authentication](https://ibm.github.io/watsonx-ai-python-sdk/setup_cloud.html#authentication).\n",
+ "\n",
+ "In this example, we’ll use the `project_id` and Dallas URL."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "a3ef7659",
+ "metadata": {},
+ "source": [
+ "Provide `PROJECT_ID` that will be used for initialize each watsonx integration instance."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "b933bb7f",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "PROJECT_ID = \"PASTE YOUR PROJECT_ID HERE\"\n",
+ "URL = \"https://us-south.ml.cloud.ibm.com\""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "f0308874",
+ "metadata": {},
+ "source": [
+ "## Download data and load documents\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "3d67b5f7",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "--2025-02-24 10:46:16-- https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
+ "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8000::154, 2606:50c0:8001::154, 2606:50c0:8002::154, ...\n",
+ "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8000::154|:443... connected.\n",
+ "HTTP request sent, awaiting response... 200 OK\n",
+ "Length: 75042 (73K) [text/plain]\n",
+ "Saving to: ‘data/paul_graham/paul_graham_essay.txt’\n",
+ "\n",
+ "data/paul_graham/pa 100%[===================>] 73,28K --.-KB/s in 0,06s \n",
+ "\n",
+ "2025-02-24 10:46:17 (1,30 MB/s) - ‘data/paul_graham/paul_graham_essay.txt’ saved [75042/75042]\n",
+ "\n"
+ ]
+ }
+ ],
+ "source": [
+ "!mkdir -p 'data/paul_graham/'\n",
+ "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "a65dc946",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.core import SimpleDirectoryReader\n",
+ "\n",
+ "documents = SimpleDirectoryReader(\"./data/paul_graham/\").load_data()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "cd7ac083",
+ "metadata": {},
+ "source": [
+ "## Load the Rerank"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "166b40ca",
+ "metadata": {},
+ "source": [
+ "You might need to adjust rerank parameters for different tasks:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "15647312",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "truncate_input_tokens = 512"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "a03571f0",
+ "metadata": {},
+ "source": [
+ "#### Initialize `WatsonxRerank` instance."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "79dcb01b",
+ "metadata": {},
+ "source": [
+ "You need to specify the `model_id` that will be used for rerank. You can find the list of all the available models in [Supported reranker models](https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-models-embed.html?context=wx#rerank)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "33efcbc5",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.postprocessor.ibm import WatsonxRerank\n",
+ "\n",
+ "watsonx_rerank = WatsonxRerank(\n",
+ " model_id=\"cross-encoder/ms-marco-minilm-l-12-v2\",\n",
+ " top_n=2,\n",
+ " url=URL,\n",
+ " project_id=PROJECT_ID,\n",
+ " truncate_input_tokens=truncate_input_tokens,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "7e79e20d",
+ "metadata": {},
+ "source": [
+ "Alternatively, you can use Cloud Pak for Data credentials. For details, see [watsonx.ai software setup](https://ibm.github.io/watsonx-ai-python-sdk/setup_cpd.html). "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "c08ee5a1",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.postprocessor.ibm import WatsonxRerank\n",
+ "\n",
+ "watsonx_rerank = WatsonxRerank(\n",
+ " model_id=\"cross-encoder/ms-marco-minilm-l-12-v2\",\n",
+ " url=URL,\n",
+ " username=\"PASTE YOUR USERNAME HERE\",\n",
+ " password=\"PASTE YOUR PASSWORD HERE\",\n",
+ " instance_id=\"openshift\",\n",
+ " version=\"5.1\",\n",
+ " project_id=PROJECT_ID,\n",
+ " truncate_input_tokens=truncate_input_tokens,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "106479c9",
+ "metadata": {},
+ "source": [
+ "## Load the embedding model"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "10b86b6b",
+ "metadata": {},
+ "source": [
+ "#### Initialize the `WatsonxEmbeddings` instance.\n",
+ "\n",
+ ">For more information about `WatsonxEmbeddings` please refer to the sample notebook: <a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/docs/examples/embeddings/ibm_watsonx.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "e5bf6aaf",
+ "metadata": {},
+ "source": [
+ "You might need to adjust embedding parameters for different tasks:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "66c3226c",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "truncate_input_tokens = 512"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "cf2a9d4f",
+ "metadata": {},
+ "source": [
+ "You need to specify the `model_id` that will be used for embedding. You can find the list of all the available models in [Supported embedding models](https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-models-embed.html?context=wx#embed)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "c81d51e1",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.embeddings.ibm import WatsonxEmbeddings\n",
+ "\n",
+ "watsonx_embedding = WatsonxEmbeddings(\n",
+ " model_id=\"ibm/slate-30m-english-rtrvr\",\n",
+ " url=URL,\n",
+ " project_id=PROJECT_ID,\n",
+ " truncate_input_tokens=truncate_input_tokens,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "8647a00c",
+ "metadata": {},
+ "source": [
+ "Change default settings"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "12645356",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.core import Settings\n",
+ "\n",
+ "Settings.chunk_size = 512"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "b1a5570f",
+ "metadata": {},
+ "source": [
+ "#### Build index"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "eb518f53",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.core import VectorStoreIndex\n",
+ "\n",
+ "index = VectorStoreIndex.from_documents(\n",
+ " documents=documents, embed_model=watsonx_embedding\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "e36acbef",
+ "metadata": {},
+ "source": [
+ "## Load the LLM"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "e2924c37",
+ "metadata": {},
+ "source": [
+ "#### Initialize the `WatsonxLLM` instance.\n",
+ "\n",
+ ">For more information about `WatsonxLLM` please refer to the sample notebook: <a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/docs/examples/llm/ibm_watsonx.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "ae227410",
+ "metadata": {},
+ "source": [
+ "You need to specify the `model_id` that will be used for inferencing. You can find the list of all the available models in [Supported foundation models](https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-models.html?context=wx)."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "ba54c1de",
+ "metadata": {},
+ "source": [
+ "You might need to adjust model `parameters` for different models or tasks. For details, refer to [Available MetaNames](https://ibm.github.io/watsonx-ai-python-sdk/fm_model.html#metanames.GenTextParamsMetaNames)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "54e38eb1",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "max_new_tokens = 128"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "359898de",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.llms.ibm import WatsonxLLM\n",
+ "\n",
+ "watsonx_llm = WatsonxLLM(\n",
+ " model_id=\"meta-llama/llama-3-3-70b-instruct\",\n",
+ " url=URL,\n",
+ " project_id=PROJECT_ID,\n",
+ " max_new_tokens=max_new_tokens,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "55aa40af",
+ "metadata": {},
+ "source": [
+ "## Send a query"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "9be7fc6c",
+ "metadata": {},
+ "source": [
+ "#### Retrieve top 10 most relevant nodes, then filter with `WatsonxRerank`"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "b84f0158",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "query_engine = index.as_query_engine(\n",
+ " llm=watsonx_llm,\n",
+ " similarity_top_k=10,\n",
+ " node_postprocessors=[watsonx_rerank],\n",
+ ")\n",
+ "response = query_engine.query(\n",
+ " \"What did Sam Altman do in this essay?\",\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "3fdc1eb4",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Final Response: In this essay, Sam Altman was recruited to be the\n",
+ "president of Y Combinator (YC), and he agreed to take over the role\n",
+ "starting with the winter 2014 batch. He initially declined the offer,\n",
+ "wanting to start a startup to make nuclear reactors, but eventually\n",
+ "agreed after being persuaded. He began learning the job and taking\n",
+ "over responsibilities from the author in the latter part of 2013, and\n",
+ "officially took over as president in 2014.\n",
+ "______________________________________________________________________\n",
+ "Source Node 1/2\n",
+ "Node ID: 2ed5d8e7-2681-49b0-a112-ea35cc9a8b9e\n",
+ "Similarity: 3.2075154781341553\n",
+ "Text: \"You know,\" he said, \"you should make sure Y Combinator isn't\n",
+ "the last cool thing you do.\" At the time I didn't understand what he\n",
+ "meant, but gradually it dawned on me that he was saying I should quit.\n",
+ "This seemed strange advice, because YC was doing great. But if there\n",
+ "was one thing rarer than Rtm offering advice, it was Rtm being wrong.\n",
+ "So th...\n",
+ "______________________________________________________________________\n",
+ "Source Node 2/2\n",
+ "Node ID: 6ae17865-aaa7-46a5-bc49-f38abf4a825e\n",
+ "Similarity: -1.3127477169036865\n",
+ "Text: I asked Jessica if she wanted to be president, but she didn't,\n",
+ "so we decided we'd try to recruit Sam Altman. We talked to Robert and\n",
+ "Trevor and we agreed to make it a complete changing of the guard. Up\n",
+ "till that point YC had been controlled by the original LLC we four had\n",
+ "started. But we wanted YC to last for a long time, and to do that it\n",
+ "could...\n"
+ ]
+ }
+ ],
+ "source": [
+ "from llama_index.core.response.pprint_utils import pprint_response\n",
+ "\n",
+ "pprint_response(response, show_source=True)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "baf5df0b",
+ "metadata": {},
+ "source": [
+ "#### Directly retrieve top 2 most similar nodes"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "1f784254",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "query_engine = index.as_query_engine(\n",
+ " llm=watsonx_llm,\n",
+ " similarity_top_k=2,\n",
+ ")\n",
+ "response = query_engine.query(\n",
+ " \"What did Sam Altman do in this essay?\",\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "8465e7a9",
+ "metadata": {},
+ "source": [
+ "Retrieved context is irrelevant and response is hallucinated."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "4dd6af0d",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Final Response: Sam Altman was one of the founders of the first batch\n",
+ "of startups funded by the Summer Founders Program, and he later became\n",
+ "the second president of YC.\n",
+ "______________________________________________________________________\n",
+ "Source Node 1/2\n",
+ "Node ID: ba52769a-7342-4e6c-af02-4159216a79a8\n",
+ "Similarity: 0.6396056863136902\n",
+ "Text: We knew undergrads were deciding then about summer jobs, so in a\n",
+ "matter of days we cooked up something we called the Summer Founders\n",
+ "Program, and I posted an announcement on my site, inviting undergrads\n",
+ "to apply. I had never imagined that writing essays would be a way to\n",
+ "get \"deal flow,\" as investors call it, but it turned out to be the\n",
+ "perfect ...\n",
+ "______________________________________________________________________\n",
+ "Source Node 2/2\n",
+ "Node ID: 43a6cf9f-8284-45db-bbbd-44109fcb9373\n",
+ "Similarity: 0.6334836031239921\n",
+ "Text: I wrote this new Lisp, called Bel, in itself in Arc. That may\n",
+ "sound like a contradiction, but it's an indication of the sort of\n",
+ "trickery I had to engage in to make this work. By means of an\n",
+ "egregious collection of hacks I managed to make something close enough\n",
+ "to an interpreter written in itself that could actually run. Not fast,\n",
+ "but fast enough...\n"
+ ]
+ }
+ ],
+ "source": [
+ "pprint_response(response, show_source=True)"
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3 (ipykernel)",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/.gitignore b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/.gitignore
new file mode 100644
index 0000000000..990c18de22
--- /dev/null
+++ b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/.gitignore
@@ -0,0 +1,153 @@
+llama_index/_static
+.DS_Store
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+bin/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+etc/
+include/
+lib/
+lib64/
+parts/
+sdist/
+share/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+# Usually these files are written by a python script from a template
+# before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+.ruff_cache
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+notebooks/
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+# However, in case of collaboration, if having platform-specific dependencies or dependencies
+# having no cross-platform support, pipenv may install dependencies that don't work, or not
+# install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+pyvenv.cfg
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# Jetbrains
+.idea
+modules/
+*.swp
+
+# VsCode
+.vscode
+
+# pipenv
+Pipfile
+Pipfile.lock
+
+# pyright
+pyrightconfig.json
diff --git a/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/BUILD b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/BUILD
new file mode 100644
index 0000000000..0896ca890d
--- /dev/null
+++ b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/BUILD
@@ -0,0 +1,3 @@
+poetry_requirements(
+ name="poetry",
+)
diff --git a/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/Makefile b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/Makefile
new file mode 100644
index 0000000000..b9eab05aa3
--- /dev/null
+++ b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/Makefile
@@ -0,0 +1,17 @@
+GIT_ROOT ?= $(shell git rev-parse --show-toplevel)
+
+help: ## Show all Makefile targets.
+ @grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[33m%-30s\033[0m %s\n", $$1, $$2}'
+
+format: ## Run code autoformatters (black).
+ pre-commit install
+ git ls-files | xargs pre-commit run black --files
+
+lint: ## Run linters: pre-commit (black, ruff, codespell) and mypy
+ pre-commit install && git ls-files | xargs pre-commit run --show-diff-on-failure --files
+
+test: ## Run tests via pytest.
+ pytest tests
+
+watch-docs: ## Build and watch documentation.
+ sphinx-autobuild docs/ docs/_build/html --open-browser --watch $(GIT_ROOT)/llama_index/
diff --git a/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/README.md b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/README.md
new file mode 100644
index 0000000000..0091701ce5
--- /dev/null
+++ b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/README.md
@@ -0,0 +1,224 @@
+# LlamaIndex Postprocessor Integration: IBM
+
+This package integrates the LlamaIndex Postprocessor API with the IBM watsonx.ai Rerank API by leveraging `ibm-watsonx-ai` [SDK](https://ibm.github.io/watsonx-ai-python-sdk/index.html).
+
+## Installation
+
+```bash
+pip install llama-index-postprocessor-ibm
+```
+
+## Usage
+
+### Setting up
+
+#### Install other required packages:
+
+```bash
+pip install -qU llama-index
+pip install -qU llama-index-llms-ibm
+pip install -qU llama-index-embeddings-ibm
+```
+
+To use IBM's Foundation Models, Embeddings and Rerank, you must have an IBM Cloud user API key. Here's how to obtain and set up your API key:
+
+1. **Obtain an API Key:** For more details on how to create and manage an API key, refer to [Managing user API keys](https://cloud.ibm.com/docs/account?topic=account-userapikey&interface=ui).
+2. **Set the API Key as an Environment Variable:** For security reasons, it's recommended to not hard-code your API key directly in your scripts. Instead, set it up as an environment variable. You can use the following code to prompt for the API key and set it as an environment variable:
+
+```python
+import os
+from getpass import getpass
+
+watsonx_api_key = getpass()
+os.environ["WATSONX_APIKEY"] = watsonx_api_key
+```
+
+Alternatively, you can set the environment variable in your terminal.
+
+- **Linux/macOS:** Open your terminal and execute the following command:
+
+ ```bash
+ export WATSONX_APIKEY='your_ibm_api_key'
+ ```
+
+ To make this environment variable persistent across terminal sessions, add the above line to your `~/.bashrc`, `~/.bash_profile`, or `~/.zshrc` file.
+
+- **Windows:** For Command Prompt, use:
+ ```cmd
+ set WATSONX_APIKEY=your_ibm_api_key
+ ```
+
+**Note**:
+
+- To provide context for the API call, you must pass the `project_id` or `space_id`. To get your project or space ID, open your project or space, go to the **Manage** tab, and click **General**. For more information see: [Project documentation](https://www.ibm.com/docs/en/watsonx-as-a-service?topic=projects) or [Deployment space documentation](https://www.ibm.com/docs/en/watsonx/saas?topic=spaces-creating-deployment).
+- Depending on the region of your provisioned service instance, use one of the urls listed in [watsonx.ai API Authentication](https://ibm.github.io/watsonx-ai-python-sdk/setup_cloud.html#authentication).
+
+In this example, we’ll use the `project_id` and Dallas URL.
+
+Provide `PROJECT_ID` that will be used for initialize each watsonx integration instance.
+
+```python
+PROJECT_ID = "PASTE YOUR PROJECT_ID HERE"
+URL = "https://us-south.ml.cloud.ibm.com"
+```
+
+### Download data and load documents
+
+```bash
+!mkdir -p 'data/paul_graham/'
+!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'
+```
+
+```python
+from llama_index.core import SimpleDirectoryReader
+
+documents = SimpleDirectoryReader("./data/paul_graham/").load_data()
+```
+
+### Load the Rerank
+
+You might need to adjust rerank parameters for different tasks:
+
+```python
+truncate_input_tokens = 512
+```
+
+#### Initialize `WatsonxRerank` instance.
+
+You need to specify the `model_id` that will be used for rerank. You can find the list of all the available models in [Supported reranker models](https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-models-embed.html?context=wx#rerank).
+
+```python
+from llama_index.postprocessor.ibm import WatsonxRerank
+
+watsonx_rerank = WatsonxRerank(
+ model_id="cross-encoder/ms-marco-minilm-l-12-v2",
+ top_n=2,
+ url=URL,
+ project_id=PROJECT_ID,
+ truncate_input_tokens=truncate_input_tokens,
+)
+```
+
+Alternatively, you can use Cloud Pak for Data credentials. For details, see [watsonx.ai software setup](https://ibm.github.io/watsonx-ai-python-sdk/setup_cpd.html).
+
+```python
+from llama_index.postprocessor.ibm import WatsonxRerank
+
+watsonx_rerank = WatsonxRerank(
+ model_id="cross-encoder/ms-marco-minilm-l-12-v2",
+ url=URL,
+ username="PASTE YOUR USERNAME HERE",
+ password="PASTE YOUR PASSWORD HERE",
+ instance_id="openshift",
+ version="5.1",
+ project_id=PROJECT_ID,
+ truncate_input_tokens=truncate_input_tokens,
+)
+```
+
+### Load the embedding model
+
+#### Initialize the `WatsonxEmbeddings` instance.
+
+> For more information about `WatsonxEmbeddings` please refer to the `llama-index-embeddings-ibm` package description.
+
+You might need to adjust embedding parameters for different tasks:
+
+```python
+truncate_input_tokens = 512
+```
+
+You need to specify the `model_id` that will be used for embedding. You can find the list of all the available models in [Supported embedding models](https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-models-embed.html?context=wx#embed).
+
+```python
+from llama_index.embeddings.ibm import WatsonxEmbeddings
+
+watsonx_embedding = WatsonxEmbeddings(
+ model_id="ibm/slate-30m-english-rtrvr",
+ url=URL,
+ project_id=PROJECT_ID,
+ truncate_input_tokens=truncate_input_tokens,
+)
+```
+
+Change default settings
+
+```python
+from llama_index.core import Settings
+
+Settings.chunk_size = 512
+```
+
+#### Build index
+
+```python
+from llama_index.core import VectorStoreIndex
+
+index = VectorStoreIndex.from_documents(
+ documents=documents, embed_model=watsonx_embedding
+)
+```
+
+### Load the LLM
+
+#### Initialize the `WatsonxLLM` instance.
+
+> For more information about `WatsonxLLM` please refer to the `llama-index-llms-ibm` package description.
+
+You need to specify the `model_id` that will be used for inferencing. You can find the list of all the available models in [Supported foundation models](https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-models.html?context=wx).
+
+You might need to adjust model `parameters` for different models or tasks. For details, refer to [Available MetaNames](https://ibm.github.io/watsonx-ai-python-sdk/fm_model.html#metanames.GenTextParamsMetaNames).
+
+```python
+max_new_tokens = 128
+```
+
+```python
+from llama_index.llms.ibm import WatsonxLLM
+
+watsonx_llm = WatsonxLLM(
+ model_id="meta-llama/llama-3-3-70b-instruct",
+ url=URL,
+ project_id=PROJECT_ID,
+ max_new_tokens=max_new_tokens,
+)
+```
+
+### Send a query
+
+#### Retrieve top 10 most relevant nodes, then filter with `WatsonxRerank`
+
+```python
+query_engine = index.as_query_engine(
+ llm=watsonx_llm,
+ similarity_top_k=10,
+ node_postprocessors=[watsonx_rerank],
+)
+response = query_engine.query(
+ "What did Sam Altman do in this essay?",
+)
+```
+
+```python
+from llama_index.core.response.pprint_utils import pprint_response
+
+pprint_response(response, show_source=True)
+```
+
+#### Directly retrieve top 2 most similar nodes
+
+```python
+query_engine = index.as_query_engine(
+ llm=watsonx_llm,
+ similarity_top_k=2,
+)
+response = query_engine.query(
+ "What did Sam Altman do in this essay?",
+)
+```
+
+Retrieved context is irrelevant and response is hallucinated.
+
+```python
+pprint_response(response, show_source=True)
+```
diff --git a/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/llama_index/postprocessor/ibm/BUILD b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/llama_index/postprocessor/ibm/BUILD
new file mode 100644
index 0000000000..db46e8d6c9
--- /dev/null
+++ b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/llama_index/postprocessor/ibm/BUILD
@@ -0,0 +1 @@
+python_sources()
diff --git a/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/llama_index/postprocessor/ibm/__init__.py b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/llama_index/postprocessor/ibm/__init__.py
new file mode 100644
index 0000000000..2908136e1d
--- /dev/null
+++ b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/llama_index/postprocessor/ibm/__init__.py
@@ -0,0 +1,4 @@
+from llama_index.postprocessor.ibm.base import WatsonxRerank
+
+
+__all__ = ["WatsonxRerank"]
diff --git a/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/llama_index/postprocessor/ibm/base.py b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/llama_index/postprocessor/ibm/base.py
new file mode 100644
index 0000000000..e77319a5cc
--- /dev/null
+++ b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/llama_index/postprocessor/ibm/base.py
@@ -0,0 +1,265 @@
+from typing import Any, Dict, Optional, Union, List
+
+from ibm_watsonx_ai import Credentials, APIClient
+
+from ibm_watsonx_ai.foundation_models import Rerank
+from ibm_watsonx_ai.foundation_models.schema import (
+ RerankParameters,
+ RerankReturnOptions,
+)
+
+from llama_index.core.bridge.pydantic import (
+ Field,
+ PrivateAttr,
+)
+
+# Import SecretStr directly from pydantic
+# since there is not one in llama_index.core.bridge.pydantic
+from pydantic import SecretStr
+
+from llama_index.core.postprocessor.types import BaseNodePostprocessor
+from llama_index.core.schema import NodeWithScore, QueryBundle, MetadataMode
+from llama_index.core.instrumentation import get_dispatcher
+from llama_index.core.instrumentation.events.rerank import (
+ ReRankEndEvent,
+ ReRankStartEvent,
+)
+
+from llama_index.core.callbacks import CallbackManager
+
+from llama_index.postprocessor.ibm.utils import resolve_watsonx_credentials
+
+dispatcher = get_dispatcher(__name__)
+
+
+class WatsonxRerank(BaseNodePostprocessor):
+ """
+ IBM watsonx.ai Rerank.
+
+ Example:
+ `pip install llama-index-postprocessor-ibm`
+
+ ```python
+
+ from llama_index.postprocessor.ibm import WatsonxRerank
+ watsonx_llm = WatsonxRerank(
+ model_id="<RERANK MODEL>",
+ url="https://us-south.ml.cloud.ibm.com",
+ apikey="*****",
+ project_id="*****",
+ )
+ ```
+ """
+
+ model_id: str = Field(description="Type of model to use.")
+
+ top_n: Optional[int] = Field(
+ default=None,
+ description="Number of top results to return.",
+ )
+
+ truncate_input_tokens: Optional[int] = Field(
+ default=None,
+ description="""Represents the maximum number of input tokens accepted.""",
+ )
+
+ project_id: Optional[str] = Field(
+ default=None,
+ description="ID of the Watson Studio project.",
+ frozen=True,
+ )
+
+ space_id: Optional[str] = Field(
+ default=None, description="ID of the Watson Studio space.", frozen=True
+ )
+
+ url: Optional[SecretStr] = Field(
+ default=None,
+ description="Url to Watson Machine Learning or CPD instance",
+ frozen=True,
+ )
+
+ apikey: Optional[SecretStr] = Field(
+ default=None,
+ description="Apikey to Watson Machine Learning or CPD instance",
+ frozen=True,
+ )
+
+ token: Optional[SecretStr] = Field(
+ default=None, description="Token to CPD instance", frozen=True
+ )
+
+ password: Optional[SecretStr] = Field(
+ default=None, description="Password to CPD instance", frozen=True
+ )
+
+ username: Optional[SecretStr] = Field(
+ default=None, description="Username to CPD instance", frozen=True
+ )
+
+ instance_id: Optional[SecretStr] = Field(
+ default=None, description="Instance_id of CPD instance", frozen=True
+ )
+
+ version: Optional[SecretStr] = Field(
+ default=None, description="Version of CPD instance", frozen=True
+ )
+
+ verify: Union[str, bool, None] = Field(
+ default=None,
+ description="""
+ User can pass as verify one of following:
+ 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
+ """,
+ frozen=True,
+ )
+
+ _client: Optional[APIClient] = PrivateAttr()
+ _watsonx_rerank: Rerank = PrivateAttr()
+
+ def __init__(
+ self,
+ model_id: Optional[str] = None,
+ top_n: Optional[int] = None,
+ truncate_input_tokens: Optional[int] = None,
+ project_id: Optional[str] = None,
+ space_id: Optional[str] = None,
+ url: Optional[str] = None,
+ apikey: Optional[str] = None,
+ token: Optional[str] = None,
+ password: Optional[str] = None,
+ username: Optional[str] = None,
+ instance_id: Optional[str] = None,
+ version: Optional[str] = None,
+ verify: Union[str, bool, None] = None,
+ api_client: Optional[APIClient] = None,
+ callback_manager: Optional[CallbackManager] = None,
+ **kwargs: Any,
+ ) -> None:
+ """
+ Initialize watsonx.ai Rerank.
+ """
+ callback_manager = callback_manager or CallbackManager([])
+
+ creds = (
+ resolve_watsonx_credentials(
+ url=url,
+ apikey=apikey,
+ token=token,
+ username=username,
+ password=password,
+ instance_id=instance_id,
+ )
+ if not isinstance(api_client, APIClient)
+ else {}
+ )
+
+ super().__init__(
+ model_id=model_id,
+ top_n=top_n,
+ truncate_input_tokens=truncate_input_tokens,
+ project_id=project_id,
+ space_id=space_id,
+ url=creds.get("url"),
+ apikey=creds.get("apikey"),
+ token=creds.get("token"),
+ password=creds.get("password"),
+ username=creds.get("username"),
+ instance_id=creds.get("instance_id"),
+ version=version,
+ verify=verify,
+ _client=api_client,
+ callback_manager=callback_manager,
+ **kwargs,
+ )
+
+ self._client = api_client
+ self._watsonx_rerank = Rerank(
+ model_id=model_id,
+ params=self.params,
+ credentials=(
+ Credentials.from_dict(
+ {
+ key: value.get_secret_value() if value else None
+ for key, value in self._get_credential_kwargs().items()
+ },
+ _verify=self.verify,
+ )
+ if creds
+ else None
+ ),
+ project_id=self.project_id,
+ space_id=self.space_id,
+ verify=verify,
+ api_client=api_client,
+ )
+
+ @classmethod
+ def class_name(cls) -> str:
+ """Get Class Name."""
+ return "WatsonxRerank"
+
+ def _get_credential_kwargs(self) -> Dict[str, SecretStr | None]:
+ return {
+ "url": self.url,
+ "apikey": self.apikey,
+ "token": self.token,
+ "password": self.password,
+ "username": self.username,
+ "instance_id": self.instance_id,
+ "version": self.version,
+ }
+
+ @property
+ def params(self) -> RerankParameters:
+ rerank_return_options: RerankReturnOptions = RerankReturnOptions(
+ top_n=self.top_n,
+ inputs=False,
+ query=False,
+ )
+ return RerankParameters(
+ truncate_input_tokens=self.truncate_input_tokens,
+ return_options=rerank_return_options,
+ )
+
+ def _postprocess_nodes(
+ self,
+ nodes: List[NodeWithScore],
+ query_bundle: Optional[QueryBundle] = None,
+ ) -> List[NodeWithScore]:
+ dispatcher.event(
+ ReRankStartEvent(
+ query=query_bundle,
+ nodes=nodes,
+ top_n=self.top_n,
+ model_name=self.model_id,
+ )
+ )
+
+ if query_bundle is None:
+ raise ValueError("Query bundle must be provided.")
+ if len(nodes) == 0:
+ return []
+
+ texts = [
+ node.node.get_content(metadata_mode=MetadataMode.EMBED) for node in nodes
+ ]
+ results = self._watsonx_rerank.generate(
+ query=query_bundle.query_str,
+ inputs=texts,
+ params=self.params,
+ )
+
+ new_nodes = []
+ for result in results.get("results", []):
+ new_node_with_score = NodeWithScore(
+ node=nodes[result["index"]].node,
+ score=result["score"],
+ )
+ new_nodes.append(new_node_with_score)
+
+ dispatcher.event(ReRankEndEvent(nodes=new_nodes))
+ return new_nodes
diff --git a/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/llama_index/postprocessor/ibm/utils.py b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/llama_index/postprocessor/ibm/utils.py
new file mode 100644
index 0000000000..86227cfc22
--- /dev/null
+++ b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/llama_index/postprocessor/ibm/utils.py
@@ -0,0 +1,111 @@
+import os
+import urllib.parse
+from typing import Dict, Union, Optional
+
+
+from llama_index.core.base.llms.generic_utils import (
+ get_from_param_or_env,
+)
+
+# Import SecretStr directly from pydantic
+# since there is not one in llama_index.core.bridge.pydantic
+from pydantic import SecretStr
+
+
+def resolve_watsonx_credentials(
+ *,
+ url: Optional[str] = None,
+ apikey: Optional[str] = None,
+ token: Optional[str] = None,
+ username: Optional[str] = None,
+ password: Optional[str] = None,
+ instance_id: Optional[str] = None
+) -> Dict[str, SecretStr]:
+ """
+ Resolve watsonx.ai credentials. If the value of given param is None
+ then tries to find corresponding environment variable.
+
+
+ :raises ValueError: raises when value of required attribute is not found
+ :return: Dictionary with resolved credentials items
+ :rtype: Dict[str, SecretStr]
+ """
+ creds = {}
+ creds["url"] = convert_to_secret_str(
+ get_from_param_or_env("url", url, "WATSONX_URL")
+ )
+
+ parsed_url = urllib.parse.urlparse(creds["url"].get_secret_value())
+ if parsed_url.netloc.endswith(".cloud.ibm.com"):
+ if not (apikey or "WATSONX_APIKEY" in os.environ) and not (
+ token or "WATSONX_TOKEN" in os.environ
+ ):
+ 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."
+ )
+ elif apikey or "WATSONX_APIKEY" in os.environ:
+ creds["apikey"] = convert_to_secret_str(
+ get_from_param_or_env("apikey", apikey, "WATSONX_APIKEY")
+ )
+ else:
+ creds["token"] = convert_to_secret_str(
+ get_from_param_or_env("token", token, "WATSONX_TOKEN")
+ )
+ else:
+ if (
+ not token
+ and "WATSONX_TOKEN" not in os.environ
+ and not password
+ and "WATSONX_PASSWORD" not in os.environ
+ and not apikey
+ and "WATSONX_APIKEY" not in os.environ
+ ):
+ 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 token or "WATSONX_TOKEN" in os.environ:
+ creds["token"] = convert_to_secret_str(
+ get_from_param_or_env("token", token, "WATSONX_TOKEN")
+ )
+
+ elif password or "WATSONX_PASSWORD" in os.environ:
+ creds["password"] = convert_to_secret_str(
+ get_from_param_or_env("password", password, "WATSONX_PASSWORD")
+ )
+
+ creds["username"] = convert_to_secret_str(
+ get_from_param_or_env("username", username, "WATSONX_USERNAME")
+ )
+
+ elif apikey or "WATSONX_APIKEY" in os.environ:
+ creds["apikey"] = convert_to_secret_str(
+ get_from_param_or_env("apikey", apikey, "WATSONX_APIKEY")
+ )
+
+ creds["username"] = convert_to_secret_str(
+ get_from_param_or_env("username", username, "WATSONX_USERNAME")
+ )
+
+ if not instance_id or "WATSONX_INSTANCE_ID" not in os.environ:
+ creds["instance_id"] = convert_to_secret_str(
+ get_from_param_or_env("instance_id", instance_id, "WATSONX_INSTANCE_ID")
+ )
+
+ return creds
+
+
+def convert_to_secret_str(value: Union[SecretStr, str]) -> SecretStr:
+ """Convert a string to a SecretStr."""
+ if isinstance(value, SecretStr):
+ return value
+ return SecretStr(value)
diff --git a/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/pyproject.toml b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/pyproject.toml
new file mode 100644
index 0000000000..a487bf838b
--- /dev/null
+++ b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/pyproject.toml
@@ -0,0 +1,60 @@
+[build-system]
+build-backend = "poetry.core.masonry.api"
+requires = ["poetry-core"]
+
+[tool.codespell]
+check-filenames = true
+check-hidden = true
+# Feel free to un-skip examples, and experimental, you will just need to
+# work through many typos (--write-changes and --interactive will help)
+skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb"
+
+[tool.llamahub]
+contains_example = false
+import_path = "llama_index.postprocessor.ibm"
+
+[tool.llamahub.class_authors]
+WatsonxRerank = "llama-index"
+
+[tool.mypy]
+disallow_untyped_defs = true
+# Remove venv skip when integrated with pre-commit
+exclude = ["_static", "build", "examples", "notebooks", "venv"]
+ignore_missing_imports = true
+python_version = "3.10"
+
+[tool.poetry]
+authors = ["IBM"]
+description = "llama-index postprocessor IBM watsonx.ai integration"
+exclude = ["**/BUILD"]
+license = "MIT"
+name = "llama-index-postprocessor-ibm"
+packages = [{include = "llama_index/"}]
+readme = "README.md"
+version = "0.1.0"
+
+[tool.poetry.dependencies]
+python = ">=3.10,<3.13"
+ibm-watsonx-ai = ">=1.1.16"
+pyarrow = "*"
+llama-index-core = "^0.12.0"
+
+[tool.poetry.group.dev.dependencies]
+black = {extras = ["jupyter"], version = "<=23.9.1,>=23.7.0"}
+codespell = {extras = ["toml"], version = ">=v2.2.6"}
+ipython = "8.10.0"
+jupyter = "^1.0.0"
+mypy = "0.991"
+pre-commit = "3.2.0"
+pylint = "2.15.10"
+pytest = "7.2.1"
+pytest-asyncio = "0.23.7"
+pytest-mock = "3.11.1"
+ruff = "0.0.292"
+tree-sitter-languages = "^1.8.0"
+types-Deprecated = ">=0.1.0"
+types-PyYAML = "^6.0.12.12"
+types-protobuf = "^4.24.0.4"
+types-redis = "4.5.5.0"
+types-requests = "2.28.11.8" # TODO: unpin when mypy>0.991
+types-setuptools = "67.1.0.0"
diff --git a/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/tests/BUILD b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/tests/BUILD
new file mode 100644
index 0000000000..47ca122a7b
--- /dev/null
+++ b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/tests/BUILD
@@ -0,0 +1,6 @@
+python_tests(
+ interpreter_constraints=[">=3.10"],
+ dependencies=[
+ "llama-index-integrations/postprocessor/llama-index-postprocessor-ibm:poetry#pyarrow"
+ ],
+)
diff --git a/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/tests/__init__.py b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/tests/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/tests/test_ibm.py b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/tests/test_ibm.py
new file mode 100644
index 0000000000..c4e8a0242b
--- /dev/null
+++ b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/tests/test_ibm.py
@@ -0,0 +1,34 @@
+import pytest
+
+from llama_index.postprocessor.ibm import WatsonxRerank
+
+
+class TestWasonxRerank:
+ TEST_URL = "https://us-south.ml.cloud.ibm.com"
+ TEST_APIKEY = "test_apikey"
+ TEST_PROJECT_ID = "test_project_id"
+
+ TEST_MODEL = "test_rerank_model"
+
+ def test_initialization(self) -> None:
+ with pytest.raises(ValueError, match=r"^Did not find") as e_info:
+ _ = WatsonxRerank(model_id=self.TEST_MODEL, project_id=self.TEST_PROJECT_ID)
+
+ # Cloud scenario
+ with pytest.raises(
+ ValueError, match=r"^Did not find 'apikey' or 'token',"
+ ) as e_info:
+ _ = WatsonxRerank(
+ model_id=self.TEST_MODEL,
+ url=self.TEST_URL,
+ project_id=self.TEST_PROJECT_ID,
+ )
+
+ # CPD scenario
+ with pytest.raises(ValueError, match=r"^Did not find instance_id") as e_info:
+ _ = WatsonxRerank(
+ model_id=self.TEST_MODEL,
+ token="test-token",
+ url="test-cpd-instance",
+ project_id=self.TEST_PROJECT_ID,
+ )
diff --git a/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/tests/test_postprocessor_ibm_rerank.py b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/tests/test_postprocessor_ibm_rerank.py
new file mode 100644
index 0000000000..8c204078ab
--- /dev/null
+++ b/llama-index-integrations/postprocessor/llama-index-postprocessor-ibm/tests/test_postprocessor_ibm_rerank.py
@@ -0,0 +1,7 @@
+from llama_index.core.postprocessor.types import BaseNodePostprocessor
+from llama_index.postprocessor.ibm import WatsonxRerank
+
+
+def test_class():
+ names_of_base_classes = [b.__name__ for b in WatsonxRerank.__mro__]
+ assert BaseNodePostprocessor.__name__ in names_of_base_classes
--
GitLab