Add NVIDIA LLM integration with wrapper class pattern

examples/smart_scraper_graph/nvidia/smart_scraper_nvidia.py

@@ -0,0 +1,48 @@
+"""
+Basic example of scraping pipeline using SmartScraper with NVIDIA
+"""
+
+import json
+import os
+
+from dotenv import load_dotenv
+
+from scrapegraphai.graphs import SmartScraperGraph
+from scrapegraphai.utils import prettify_exec_info
+
+load_dotenv()
+
+# ************************************************
+# Define the configuration for the graph
+# ************************************************
+
+
+graph_config = {
+    "llm": {
+        "api_key": os.getenv("NVIDIA_API_KEY"),
+        "model": "nvidia/meta/llama3-70b-instruct",
+        "model_provider": "nvidia",
+    },
+    "verbose": True,
+    "headless": False,
+}
+
+# ************************************************
+# Create the SmartScraperGraph instance and run it
+# ************************************************
+
+smart_scraper_graph = SmartScraperGraph(
+    prompt="Extract me the first article",
+    source="https://www.wired.com",
+    config=graph_config,
+)
+
+result = smart_scraper_graph.run()
+print(json.dumps(result, indent=4))
+
+# ************************************************
+# Get graph execution info
+# ************************************************
+
+graph_exec_info = smart_scraper_graph.get_execution_info()
+print(prettify_exec_info(graph_exec_info))

pyproject.toml

@@ -69,6 +69,7 @@ requires-python = ">=3.10,<4.0"
 [project.optional-dependencies]
 burr = ["burr[start]==0.22.1"]
 docs = ["sphinx==6.0", "furo==2024.5.6"]
+nvidia = ["langchain-nvidia-ai-endpoints>=0.1.0"]
 ocr = [
     "surya-ocr>=0.5.0",
     "matplotlib>=3.7.2",

scrapegraphai/graphs/abstract_graph.py

@@ -5,7 +5,6 @@
 import asyncio
 import uuid
 import warnings
-import time
 from abc import ABC, abstractmethod
 from typing import Optional, Type
 
@@ -14,9 +13,8 @@
 from pydantic import BaseModel
 
 from ..helpers import models_tokens
-from ..models import CLoD, DeepSeek, OneApi, XAI
+from ..models import CLoD, DeepSeek, Nvidia, OneApi, XAI
 from ..utils.logging import set_verbosity_info, set_verbosity_warning, get_logger
-from ..telemetry import log_graph_execution
 
 logger = get_logger(__name__)
 
@@ -264,14 +262,7 @@ def _create_llm(self, llm_config: dict) -> object:
                     return ChatTogether(**llm_params)
 
                 elif model_provider == "nvidia":
-                    try:
-                        from langchain_nvidia_ai_endpoints import ChatNVIDIA
-                    except ImportError:
-                        raise ImportError(
-                            """The langchain_nvidia_ai_endpoints module is not installed.
-                                          Please install it using `pip install langchain-nvidia-ai-endpoints`."""
-                        )
-                    return ChatNVIDIA(**llm_params)
+                    return Nvidia(**llm_params)
 
         except Exception as e:
             raise Exception(f"Error instancing model: {e}")

scrapegraphai/models/__init__.py

@@ -4,9 +4,10 @@
 
 from .clod import CLoD
 from .deepseek import DeepSeek
+from .nvidia import Nvidia
 from .oneapi import OneApi
 from .openai_itt import OpenAIImageToText
 from .openai_tts import OpenAITextToSpeech
 from .xai import XAI
 
-__all__ = ["DeepSeek", "OneApi", "OpenAIImageToText", "OpenAITextToSpeech", "CLoD", "XAI"]
+__all__ = ["DeepSeek", "OneApi", "OpenAIImageToText", "OpenAITextToSpeech", "CLoD", "XAI", "Nvidia"]

scrapegraphai/models/nvidia.py

@@ -0,0 +1,32 @@
+"""
+NVIDIA Module
+"""
+
+
+class Nvidia:
+    """
+    A wrapper for the ChatNVIDIA class that provides default configuration
+    and could be extended with additional methods if needed.
+
+    Note: This class uses __new__ instead of __init__ because langchain_nvidia_ai_endpoints
+    is an optional dependency. We cannot inherit from ChatNVIDIA at class definition time
+    since the module may not be installed. The __new__ method allows us to lazily import
+    and return a ChatNVIDIA instance only when Nvidia() is instantiated.
+
+    Args:
+        llm_config (dict): Configuration parameters for the language model.
+    """
+
+    def __new__(cls, **llm_config):
+        try:
+            from langchain_nvidia_ai_endpoints import ChatNVIDIA
+        except ImportError:
+            raise ImportError(
+                """The langchain_nvidia_ai_endpoints module is not installed.
+                              Please install it using `pip install langchain-nvidia-ai-endpoints`."""
+            )
+
+        if "api_key" in llm_config:
+            llm_config["nvidia_api_key"] = llm_config.pop("api_key")
+
+        return ChatNVIDIA(**llm_config)