Raise warning for base_url ../embeddings .../completions .../rankings (…

…#922) * add validation for base url routes * move url validation to utils * update docstring for url validation * add typing for arg type * return docstring update Co-authored-by: Madeesh Kannan <[email protected]> * fix typo error Co-authored-by: Madeesh Kannan <[email protected]> --------- Co-authored-by: Madeesh Kannan <[email protected]>
deepset-ai · Jul 29, 2024 · 2f6f134 · 2f6f134
1 parent 0cdda5c
commit 2f6f134
Show file tree

Hide file tree

Showing 9 changed files with 160 additions and 43 deletions.
diff --git a/...rations/nvidia/src/haystack_integrations/components/embedders/nvidia/document_embedder.py b/...rations/nvidia/src/haystack_integrations/components/embedders/nvidia/document_embedder.py
@@ -2,12 +2,15 @@
 
 from haystack import Document, component, default_from_dict, default_to_dict
 from haystack.utils import Secret, deserialize_secrets_inplace
+from haystack_integrations.utils.nvidia import url_validation
 from tqdm import tqdm
 
 from ._nim_backend import NimBackend
 from .backend import EmbedderBackend
 from .truncate import EmbeddingTruncateMode
 
+_DEFAULT_API_URL = "https://ai.api.nvidia.com/v1/retrieval/nvidia"
+
 
 @component
 class NvidiaDocumentEmbedder:
@@ -33,7 +36,7 @@ def __init__(
         self,
         model: str = "NV-Embed-QA",
         api_key: Optional[Secret] = Secret.from_env_var("NVIDIA_API_KEY"),
-        api_url: str = "https://ai.api.nvidia.com/v1/retrieval/nvidia",
+        api_url: str = _DEFAULT_API_URL,
         prefix: str = "",
         suffix: str = "",
         batch_size: int = 32,
@@ -51,6 +54,7 @@ def __init__(
             API key for the NVIDIA NIM.
         :param api_url:
             Custom API URL for the NVIDIA NIM.
+            Format for API URL is http://host:port
         :param prefix:
             A string to add to the beginning of each text.
         :param suffix:
@@ -71,7 +75,7 @@ def __init__(
 
         self.api_key = api_key
         self.model = model
-        self.api_url = api_url
+        self.api_url = url_validation(api_url, _DEFAULT_API_URL, ["v1/embeddings"])
         self.prefix = prefix
         self.suffix = suffix
         self.batch_size = batch_size

diff --git a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/text_embedder.py b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/text_embedder.py
@@ -2,11 +2,14 @@
 
 from haystack import component, default_from_dict, default_to_dict
 from haystack.utils import Secret, deserialize_secrets_inplace
+from haystack_integrations.utils.nvidia import url_validation
 
 from ._nim_backend import NimBackend
 from .backend import EmbedderBackend
 from .truncate import EmbeddingTruncateMode
 
+_DEFAULT_API_URL = "https://ai.api.nvidia.com/v1/retrieval/nvidia"
+
 
 @component
 class NvidiaTextEmbedder:
@@ -34,7 +37,7 @@ def __init__(
         self,
         model: str = "NV-Embed-QA",
         api_key: Optional[Secret] = Secret.from_env_var("NVIDIA_API_KEY"),
-        api_url: str = "https://ai.api.nvidia.com/v1/retrieval/nvidia",
+        api_url: str = _DEFAULT_API_URL,
         prefix: str = "",
         suffix: str = "",
         truncate: Optional[Union[EmbeddingTruncateMode, str]] = None,
@@ -48,6 +51,7 @@ def __init__(
             API key for the NVIDIA NIM.
         :param api_url:
             Custom API URL for the NVIDIA NIM.
+            Format for API URL is http://host:port
         :param prefix:
             A string to add to the beginning of each text.
         :param suffix:
@@ -59,7 +63,7 @@ def __init__(
 
         self.api_key = api_key
         self.model = model
-        self.api_url = api_url
+        self.api_url = url_validation(api_url, _DEFAULT_API_URL, ["v1/embeddings"])
         self.prefix = prefix
         self.suffix = suffix
 

diff --git a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py
@@ -5,6 +5,7 @@
 
 from haystack import component, default_from_dict, default_to_dict
 from haystack.utils.auth import Secret, deserialize_secrets_inplace
+from haystack_integrations.utils.nvidia import url_validation
 
 from ._nim_backend import NimBackend
 from .backend import GeneratorBackend
@@ -63,7 +64,7 @@ def __init__(
             to know the supported arguments.
         """
         self._model = model
-        self._api_url = api_url
+        self._api_url = url_validation(api_url, _DEFAULT_API_URL, ["v1/chat/completions"])
         self._api_key = api_key
         self._model_arguments = model_arguments or {}
 

diff --git a/integrations/nvidia/src/haystack_integrations/utils/nvidia/__init__.py b/integrations/nvidia/src/haystack_integrations/utils/nvidia/__init__.py
@@ -0,0 +1,3 @@
+from .utils import url_validation
+
+__all__ = ["url_validation"]
diff --git a/integrations/nvidia/src/haystack_integrations/utils/nvidia/utils.py b/integrations/nvidia/src/haystack_integrations/utils/nvidia/utils.py
@@ -0,0 +1,39 @@
+import warnings
+from typing import List
+from urllib.parse import urlparse, urlunparse
+
+
+def url_validation(api_url: str, default_api_url: str, allowed_paths: List[str]) -> str:
+    """
+    Validate and normalize an API URL.
+
+    :param api_url:
+        The API URL to validate and normalize.
+    :param default_api_url:
+        The default API URL for comparison.
+    :param allowed_paths:
+        A list of allowed base paths that are valid if present in the URL.
+    :returns:
+        A normalized version of the API URL with '/v1' path appended, if needed.
+    :raises ValueError:
+        If the base URL path is not recognized or does not match expected format.
+    """
+    ## Making sure /v1 in added to the url, followed by infer_path
+    result = urlparse(api_url)
+    expected_format = "Expected format is 'http://host:port'."
+
+    if api_url == default_api_url:
+        return api_url
+    if result.path:
+        normalized_path = result.path.strip("/")
+        if normalized_path == "v1":
+            pass
+        elif normalized_path in allowed_paths:
+            warn_msg = f"{expected_format} Rest is ignored."
+            warnings.warn(warn_msg, stacklevel=2)
+        else:
+            err_msg = f"Base URL path is not recognized. {expected_format}"
+            raise ValueError(err_msg)
+
+    base_url = urlunparse((result.scheme, result.netloc, "v1", "", "", ""))
+    return base_url
diff --git a/integrations/nvidia/tests/test_base_url.py b/integrations/nvidia/tests/test_base_url.py
@@ -0,0 +1,64 @@
+import pytest
+from haystack_integrations.components.embedders.nvidia import NvidiaDocumentEmbedder, NvidiaTextEmbedder
+from haystack_integrations.components.generators.nvidia import NvidiaGenerator
+
+
+@pytest.mark.parametrize(
+    "base_url",
+    [
+        "http://localhost:8888/embeddings",
+        "http://0.0.0.0:8888/rankings",
+        "http://0.0.0.0:8888/v1/rankings",
+        "http://localhost:8888/chat/completions",
+        "http://localhost:8888/v1/chat/completions",
+    ],
+)
+@pytest.mark.parametrize(
+    "embedder",
+    [NvidiaDocumentEmbedder, NvidiaTextEmbedder],
+)
+def test_base_url_invalid_not_hosted(base_url: str, embedder) -> None:
+    with pytest.raises(ValueError):
+        embedder(api_url=base_url, model="x")
+
+
+@pytest.mark.parametrize(
+    "base_url",
+    ["http://localhost:8080/v1/embeddings", "http://0.0.0.0:8888/v1/embeddings"],
+)
+@pytest.mark.parametrize(
+    "embedder",
+    [NvidiaDocumentEmbedder, NvidiaTextEmbedder],
+)
+def test_base_url_valid_embedder(base_url: str, embedder) -> None:
+    with pytest.warns(UserWarning):
+        embedder(api_url=base_url)
+
+
+@pytest.mark.parametrize(
+    "base_url",
+    [
+        "http://localhost:8080/v1/chat/completions",
+        "http://0.0.0.0:8888/v1/chat/completions",
+    ],
+)
+def test_base_url_valid_generator(base_url: str) -> None:
+    with pytest.warns(UserWarning):
+        NvidiaGenerator(
+            api_url=base_url,
+            model="mistralai/mixtral-8x7b-instruct-v0.1",
+        )
+
+
+@pytest.mark.parametrize(
+    "base_url",
+    [
+        "http://localhost:8888/embeddings",
+        "http://0.0.0.0:8888/rankings",
+        "http://0.0.0.0:8888/v1/rankings",
+        "http://localhost:8888/chat/completions",
+    ],
+)
+def test_base_url_invalid_generator(base_url: str) -> None:
+    with pytest.raises(ValueError):
+        NvidiaGenerator(api_url=base_url, model="x")
diff --git a/integrations/nvidia/tests/test_document_embedder.py b/integrations/nvidia/tests/test_document_embedder.py
@@ -33,27 +33,28 @@ def test_init_default(self, monkeypatch):
         assert embedder.embedding_separator == "\n"
 
     def test_init_with_parameters(self):
-        embedder = NvidiaDocumentEmbedder(
-            api_key=Secret.from_token("fake-api-key"),
-            model="nvolveqa_40k",
-            api_url="https://ai.api.nvidia.com/v1/retrieval/nvidia/test",
-            prefix="prefix",
-            suffix="suffix",
-            batch_size=30,
-            progress_bar=False,
-            meta_fields_to_embed=["test_field"],
-            embedding_separator=" | ",
-        )
-
-        assert embedder.api_key == Secret.from_token("fake-api-key")
-        assert embedder.model == "nvolveqa_40k"
-        assert embedder.api_url == "https://ai.api.nvidia.com/v1/retrieval/nvidia/test"
-        assert embedder.prefix == "prefix"
-        assert embedder.suffix == "suffix"
-        assert embedder.batch_size == 30
-        assert embedder.progress_bar is False
-        assert embedder.meta_fields_to_embed == ["test_field"]
-        assert embedder.embedding_separator == " | "
+        with pytest.raises(ValueError):
+            embedder = NvidiaDocumentEmbedder(
+                api_key=Secret.from_token("fake-api-key"),
+                model="nvolveqa_40k",
+                api_url="https://ai.api.nvidia.com/v1/retrieval/nvidia/test",
+                prefix="prefix",
+                suffix="suffix",
+                batch_size=30,
+                progress_bar=False,
+                meta_fields_to_embed=["test_field"],
+                embedding_separator=" | ",
+            )
+
+            assert embedder.api_key == Secret.from_token("fake-api-key")
+            assert embedder.model == "nvolveqa_40k"
+            assert embedder.api_url == "https://ai.api.nvidia.com/v1/retrieval/nvidia/test"
+            assert embedder.prefix == "prefix"
+            assert embedder.suffix == "suffix"
+            assert embedder.batch_size == 30
+            assert embedder.progress_bar is False
+            assert embedder.meta_fields_to_embed == ["test_field"]
+            assert embedder.embedding_separator == " | "
 
     def test_init_fail_wo_api_key(self, monkeypatch):
         monkeypatch.delenv("NVIDIA_API_KEY", raising=False)
@@ -99,7 +100,7 @@ def test_to_dict_with_custom_init_parameters(self, monkeypatch):
             "type": "haystack_integrations.components.embedders.nvidia.document_embedder.NvidiaDocumentEmbedder",
             "init_parameters": {
                 "api_key": {"env_vars": ["NVIDIA_API_KEY"], "strict": True, "type": "env_var"},
-                "api_url": "https://example.com",
+                "api_url": "https://example.com/v1",
                 "model": "playground_nvolveqa_40k",
                 "prefix": "prefix",
                 "suffix": "suffix",
@@ -130,7 +131,7 @@ def from_dict(self, monkeypatch):
         }
         component = NvidiaDocumentEmbedder.from_dict(data)
         assert component.model == "nvolveqa_40k"
-        assert component.api_url == "https://example.com"
+        assert component.api_url == "https://example.com/v1"
         assert component.prefix == "prefix"
         assert component.suffix == "suffix"
         assert component.batch_size == 32

diff --git a/integrations/nvidia/tests/test_generator.py b/integrations/nvidia/tests/test_generator.py
@@ -80,7 +80,7 @@ def test_to_dict_with_custom_init_parameters(self, monkeypatch):
             "type": "haystack_integrations.components.generators.nvidia.generator.NvidiaGenerator",
             "init_parameters": {
                 "api_key": {"env_vars": ["NVIDIA_API_KEY"], "strict": True, "type": "env_var"},
-                "api_url": "https://my.url.com",
+                "api_url": "https://my.url.com/v1",
                 "model": "playground_nemotron_steerlm_8b",
                 "model_arguments": {
                     "temperature": 0.2,

diff --git a/integrations/nvidia/tests/test_text_embedder.py b/integrations/nvidia/tests/test_text_embedder.py
@@ -28,18 +28,19 @@ def test_init_default(self, monkeypatch):
         assert embedder.suffix == ""
 
     def test_init_with_parameters(self):
-        embedder = NvidiaTextEmbedder(
-            api_key=Secret.from_token("fake-api-key"),
-            model="nvolveqa_40k",
-            api_url="https://ai.api.nvidia.com/v1/retrieval/nvidia/test",
-            prefix="prefix",
-            suffix="suffix",
-        )
-        assert embedder.api_key == Secret.from_token("fake-api-key")
-        assert embedder.model == "nvolveqa_40k"
-        assert embedder.api_url == "https://ai.api.nvidia.com/v1/retrieval/nvidia/test"
-        assert embedder.prefix == "prefix"
-        assert embedder.suffix == "suffix"
+        with pytest.raises(ValueError):
+            embedder = NvidiaTextEmbedder(
+                api_key=Secret.from_token("fake-api-key"),
+                model="nvolveqa_40k",
+                api_url="https://ai.api.nvidia.com/v1/retrieval/nvidia/test",
+                prefix="prefix",
+                suffix="suffix",
+            )
+            assert embedder.api_key == Secret.from_token("fake-api-key")
+            assert embedder.model == "nvolveqa_40k"
+            assert embedder.api_url == "https://ai.api.nvidia.com/v1/retrieval/nvidia/test"
+            assert embedder.prefix == "prefix"
+            assert embedder.suffix == "suffix"
 
     def test_init_fail_wo_api_key(self, monkeypatch):
         monkeypatch.delenv("NVIDIA_API_KEY", raising=False)
@@ -77,7 +78,7 @@ def test_to_dict_with_custom_init_parameters(self, monkeypatch):
             "type": "haystack_integrations.components.embedders.nvidia.text_embedder.NvidiaTextEmbedder",
             "init_parameters": {
                 "api_key": {"env_vars": ["NVIDIA_API_KEY"], "strict": True, "type": "env_var"},
-                "api_url": "https://example.com",
+                "api_url": "https://example.com/v1",
                 "model": "nvolveqa_40k",
                 "prefix": "prefix",
                 "suffix": "suffix",
@@ -100,7 +101,7 @@ def from_dict(self, monkeypatch):
         }
         component = NvidiaTextEmbedder.from_dict(data)
         assert component.model == "nvolveqa_40k"
-        assert component.api_url == "https://example.com"
+        assert component.api_url == "https://example.com/v1"
         assert component.prefix == "prefix"
         assert component.suffix == "suffix"
         assert component.truncate == "START"
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		from .utils import url_validation

		__all__ = ["url_validation"]