From 6e003641bea035f0631e749b57cb1736e605740a Mon Sep 17 00:00:00 2001
From: Silvano Cerza <3314350+silvanocerza@users.noreply.github.com>
Date: Thu, 7 Mar 2024 17:38:33 +0100
Subject: [PATCH] Add missing import in NvidiaGenerator docstring (#559)

* Add missing import in NvidiaGenerator docstring

* Export NvidiaGeneratorModel in nvidia package
---
 .../components/generators/nvidia/__init__.py                   | 3 ++-
 .../components/generators/nvidia/generator.py                  | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/__init__.py b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/__init__.py
index 18354ea17..3a315843d 100644
--- a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/__init__.py
+++ b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/__init__.py
@@ -2,5 +2,6 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 from .generator import NvidiaGenerator
+from .models import NvidiaGeneratorModel
 
-__all__ = ["NvidiaGenerator"]
+__all__ = ["NvidiaGenerator", "NvidiaGeneratorModel"]
diff --git a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py
index cb24c7fe0..46550baab 100644
--- a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py
+++ b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py
@@ -19,7 +19,7 @@ class NvidiaGenerator:
 
     Usage example:
     ```python
-    from haystack_integrations.components.generators.nvidia import NvidiaGenerator
+    from haystack_integrations.components.generators.nvidia import NvidiaGenerator, NvidiaGeneratorModel
 
     generator = NvidiaGenerator(
         model=NvidiaGeneratorModel.NV_LLAMA2_RLHF_70B,