feat(drivers-rerank-nvidia-nim): add NvidiaNimRerankDriver (#1926)

* feat(drivers-rerank-nvidia-nim): add `NvidiaNimRerankDriver`

* skipped test

* clean up logic

Author: vachillo

docs/griptape-framework/drivers/rerank-drivers.md

@@ -44,3 +44,13 @@ The [CohereRerankDriver](../../reference/griptape/drivers/rerank/cohere_rerank_d
     ```text
     --8<-- "docs/griptape-framework/drivers/logs/cohere_rerank_driver.txt"
     ```
+
+### Nvidia NIM
+
+The [NvidiaNimRerankDriver](../../reference/griptape/drivers/rerank/nvidia_nim_rerank_driver.md) uses the [Nvidia NIM Reranking API](https://docs.nvidia.com/nim/nemo-retriever/text-reranking/latest/index.html).
+
+=== "Code"
+
+    ```python
+    --8<-- "docs/griptape-framework/drivers/src/nvidia_nim_rerank_driver.py"
+    ```

docs/griptape-framework/drivers/src/nvidia_nim_rerank_driver.py

@@ -0,0 +1,19 @@
+from griptape.artifacts import TextArtifact
+from griptape.drivers.rerank.nvidia_nim import NvidiaNimRerankDriver
+
+rerank_driver = NvidiaNimRerankDriver(
+    model="nvidia/bert-base-uncased",
+    base_url="http://localhost:8000",
+)
+
+artifacts = rerank_driver.run(
+    "Where is NYC located?",
+    [
+        TextArtifact("NYC Media"),
+        TextArtifact("New York City Police Department"),
+        TextArtifact("New York City"),
+        TextArtifact("New York City Subway"),
+    ],
+)
+for artifact in artifacts:
+    print(artifact.value)

griptape/drivers/rerank/nvidia_nim/__init__.py

@@ -0,0 +1,3 @@
+from griptape.drivers.rerank.nvidia_nim_rerank_driver import NvidiaNimRerankDriver
+
+__all__ = ["NvidiaNimRerankDriver"]

griptape/drivers/rerank/nvidia_nim_rerank_driver.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+import requests
+from attrs import define, field
+
+from griptape.drivers.rerank.base_rerank_driver import BaseRerankDriver
+
+if TYPE_CHECKING:
+    from griptape.artifacts import TextArtifact
+
+
+@define(kw_only=True)
+class NvidiaNimRerankDriver(BaseRerankDriver):
+    """Nvidia Rerank Driver."""
+
+    model: str = field()
+    base_url: str = field()
+    truncate: Literal["NONE", "END"] = field(default="NONE")
+    headers: dict = field(factory=dict)
+
+    def run(self, query: str, artifacts: list[TextArtifact]) -> list[TextArtifact]:
+        if not artifacts:
+            return []
+
+        response = requests.post(
+            url=f"{self.base_url.rstrip('/')}/v1/ranking",
+            json=self._get_body(query, artifacts),
+            headers=self.headers,
+        )
+
+        response.raise_for_status()
+
+        ranked_artifacts = []
+        for ranking in response.json()["rankings"]:
+            artifact = artifacts[ranking["index"]]
+            artifact.meta.update({"logit": ranking["logit"], "usage": ranking.get("usage")})
+            ranked_artifacts.append(artifact)
+
+        return ranked_artifacts
+
+    def _get_body(self, query: str, artifacts: list[TextArtifact]) -> dict:
+        return {
+            "model": self.model,
+            "query": {"text": query},
+            "passages": [{"text": artifact.value} for artifact in artifacts],
+            "truncate": self.truncate,
+        }

tests/integration/test_code_blocks.py

@@ -19,6 +19,7 @@
     "docs/griptape-framework/drivers/src/prompt_drivers_14.py",
     "docs/griptape-framework/drivers/src/observability_drivers_1.py",
     "docs/griptape-framework/drivers/src/observability_drivers_2.py",
+    "docs/griptape-framework/drivers/src/nvidia_nim_rerank_driver.py",
     "docs/griptape-framework/structures/src/observability_1.py",
     "docs/griptape-framework/structures/src/observability_2.py",
     "docs/griptape-framework/data/src/loaders_9.py",

tests/unit/drivers/rerank/test_nvidia_nim_rerank_driver.py

@@ -0,0 +1,62 @@
+import pytest
+import requests
+
+from griptape.artifacts import TextArtifact
+from griptape.drivers.rerank.nvidia_nim import NvidiaNimRerankDriver
+
+
+class TestNvidiaNimRerankDriver:
+    @pytest.fixture()
+    def mock_client(self, mocker):
+        def mock_post(*args, **kwargs):
+            return mocker.Mock(
+                status_code=200,
+                json=lambda: {
+                    "rankings": [
+                        {"index": 0, "logit": 0.1, "usage": {"prompt_tokens": 10, "total_tokens": 20}},
+                        {"index": 1, "logit": 0.2, "usage": {"prompt_tokens": 10, "total_tokens": 20}},
+                    ]
+                },
+            )
+
+        mocker.patch("griptape.drivers.rerank.nvidia_nim_rerank_driver.requests.post", side_effect=mock_post)
+
+    @pytest.fixture()
+    def mock_empty_client(self, mocker):
+        def mock_post(*args, **kwargs):
+            return mocker.Mock(
+                status_code=200,
+                json=lambda: {"rankings": []},
+            )
+
+        mocker.patch("griptape.drivers.rerank.nvidia_nim_rerank_driver.requests.post", side_effect=mock_post)
+
+    def test_run(self, mock_client):
+        driver = NvidiaNimRerankDriver(model="model-name", base_url="http://localhost:8000")
+        result = driver.run("hello", artifacts=[TextArtifact("foo"), TextArtifact("bar")])
+
+        assert len(result) == 2
+
+    def test_run_empty_artifacts(self, mock_empty_client):
+        driver = NvidiaNimRerankDriver(model="model-name", base_url="http://localhost:8000")
+        result = driver.run("hello", artifacts=[TextArtifact(""), TextArtifact("  ")])
+
+        assert len(result) == 0
+
+        result = driver.run("hello", artifacts=[])
+        assert len(result) == 0
+
+    def test_run_error(self, mocker):
+        mocker.patch(
+            "griptape.drivers.rerank.nvidia_nim_rerank_driver.requests.post",
+            return_value=mocker.Mock(
+                status_code=500,
+                text="Internal Server Error",
+                raise_for_status=lambda: (_ for _ in ()).throw(requests.exceptions.HTTPError()),
+            ),
+        )
+
+        driver = NvidiaNimRerankDriver(model="model-name", base_url="http://localhost:8000")
+
+        with pytest.raises(requests.exceptions.HTTPError):
+            driver.run("hello", artifacts=[TextArtifact("foo"), TextArtifact("bar")])