fest: add structured output support in NvidiaChatGenerator (#2405)

* Add structured output

* Update ruff version

* restore python 3.9

* Fix linting

* Fix linting

* Restore

* Restore

* Restore

* Add a new test

* Update nvidia.yml

* Update tools

* remove openrouter

Author: Amnah199

integrations/nvidia/examples/chat_generator_with_structured_outputs.py

@@ -0,0 +1,34 @@
+# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+# This example demonstrates how to use the NvidiaChatGenerator component
+# with structured outputs.
+# To run this example, you will need to
+# set `NVIDIA_API_KEY` environment variable
+
+from haystack.dataclasses import ChatMessage
+
+from haystack_integrations.components.generators.nvidia import NvidiaChatGenerator
+
+json_schema = {
+    "type": "object",
+    "properties": {"title": {"type": "string"}, "rating": {"type": "number"}},
+    "required": ["title", "rating"],
+}
+chat_messages = [
+    ChatMessage.from_user(
+        """
+    Return the title and the rating based on the following movie review according
+    to the provided json schema.
+    Review: Inception is a really well made film. I rate it four stars out of five."""
+    )
+]
+
+component = NvidiaChatGenerator(
+    model="meta/llama-3.1-70b-instruct",
+    generation_kwargs={"extra_body": {"nvext": {"guided_json": json_schema}}},
+)
+results = component.run(chat_messages)
+# print(results)

integrations/nvidia/pyproject.toml

@@ -23,7 +23,7 @@ classifiers = [
   "Programming Language :: Python :: Implementation :: CPython",
   "Programming Language :: Python :: Implementation :: PyPy",
 ]
-dependencies = ["haystack-ai>=2.13.0", "requests>=2.25.0", "tqdm>=4.21.0"]
+dependencies = ["haystack-ai>=2.19.0", "requests>=2.25.0", "tqdm>=4.21.0"]
 
 [project.urls]
 Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/nvidia#readme"
@@ -160,4 +160,4 @@ addopts = "--strict-markers"
 markers = [
   "integration: integration tests",
 ]
-log_cli = true
+log_cli = true

integrations/nvidia/src/haystack_integrations/components/generators/nvidia/chat/chat_generator.py

@@ -3,12 +3,12 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import os
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, Optional
 
 from haystack import component, default_to_dict, logging
 from haystack.components.generators.chat import OpenAIChatGenerator
 from haystack.dataclasses import StreamingCallbackT
-from haystack.tools import Tool, Toolset, serialize_tools_or_toolset
+from haystack.tools import ToolsType, serialize_tools_or_toolset
 from haystack.utils import serialize_callable
 from haystack.utils.auth import Secret
 
@@ -56,7 +56,7 @@ def __init__(
         streaming_callback: Optional[StreamingCallbackT] = None,
         api_base_url: Optional[str] = os.getenv("NVIDIA_API_URL", DEFAULT_API_URL),
         generation_kwargs: Optional[Dict[str, Any]] = None,
-        tools: Optional[Union[List[Tool], Toolset]] = None,
+        tools: Optional[ToolsType] = None,
         timeout: Optional[float] = None,
         max_retries: Optional[int] = None,
         http_client_kwargs: Optional[Dict[str, Any]] = None,
@@ -86,6 +86,22 @@ def __init__(
                 comprising the top 10% probability mass are considered.
             - `stream`: Whether to stream back partial progress. If set, tokens will be sent as data-only server-sent
                 events as they become available, with the stream terminated by a data: [DONE] message.
+            - `response_format`: For NVIDIA NIM servers, this parameter has limited support.
+                - The basic JSON mode with `{"type": "json_object"}` is supported by compatible models, to produce
+                valid JSON output.
+                To pass the JSON schema to the model, use the `guided_json` parameter in `extra_body`.
+                For example:
+                ```python
+                generation_kwargs={
+                    "extra_body": {
+                        "nvext": {
+                            "guided_json": {
+                                json_schema
+                        }
+                    }
+                }
+                ```
+                For more details, see the [NVIDIA NIM documentation](https://docs.nvidia.com/nim/large-language-models/latest/structured-generation.html).
         :param tools:
             A list of tools or a Toolset for which the model can prepare calls. This parameter can accept either a
             list of `Tool` objects or a `Toolset` instance.

integrations/nvidia/tests/test_nvidia_chat_generator.py

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import json
 import os
 from datetime import datetime
 from unittest.mock import AsyncMock, patch
@@ -267,6 +268,67 @@ def __call__(self, chunk: StreamingChunk) -> None:
         assert callback.counter > 1
         assert "Paris" in callback.responses
 
+    @pytest.mark.skipif(
+        not os.environ.get("NVIDIA_API_KEY", None),
+        reason="Export an env var called NVIDIA_API_KEY containing the NVIDIA API key to run this test.",
+    )
+    @pytest.mark.integration
+    def test_live_run_with_guided_json_schema(self):
+        json_schema = {
+            "type": "object",
+            "properties": {"title": {"type": "string"}, "rating": {"type": "number"}},
+            "required": ["title", "rating"],
+        }
+        chat_messages = [
+            ChatMessage.from_user(
+                """
+            Return the title and the rating based on the following movie review according
+            to the provided json schema.
+            Review: Inception is a really well made film. I rate it four stars out of five."""
+            )
+        ]
+
+        component = NvidiaChatGenerator(
+            model="meta/llama-3.1-70b-instruct",
+            generation_kwargs={"extra_body": {"nvext": {"guided_json": json_schema}}},
+        )
+
+        results = component.run(chat_messages)
+        assert len(results["replies"]) == 1
+        message = results["replies"][0].text
+        output = json.loads(message)
+        assert output["title"] == "Inception"
+        assert "rating" in output
+
+    @pytest.mark.skipif(
+        not os.environ.get("NVIDIA_API_KEY", None),
+        reason="Export an env var called NVIDIA_API_KEY containing the NVIDIA API key to run this test.",
+    )
+    @pytest.mark.integration
+    def test_live_run_with_json_object(self):
+        chat_messages = [
+            ChatMessage.from_user(
+                """
+            Return the title and the rating based on the following movie review according
+            to the provided json schema.
+            Review: Inception is a really well made film. I rate it four stars out of five."""
+            )
+        ]
+
+        component = NvidiaChatGenerator(
+            model="meta/llama-3.1-70b-instruct",
+            generation_kwargs={"response_format": {"type": "json_object"}},
+        )
+
+        results = component.run(chat_messages)
+        assert len(results["replies"]) == 1
+        message = results["replies"][0].text
+        output = json.loads(message)
+        assert "title" in output
+        assert "rating" in output
+        assert isinstance(output["rating"], int)
+        assert "Inception" in output["title"]
+
 
 class TestNvidiaChatGeneratorAsync:
     def test_init_default_async(self, monkeypatch):