diff --git a/examples/smart_scraper_graph/nvidia/smart_scraper_nvidia.py b/examples/smart_scraper_graph/nvidia/smart_scraper_nvidia.py new file mode 100644 index 00000000..9a65e9c2 --- /dev/null +++ b/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)) diff --git a/pyproject.toml b/pyproject.toml index 30278838..b1dcde3c 100644 --- a/pyproject.toml +++ b/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", diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py index c980acd7..30dd9a05 100644 --- a/scrapegraphai/graphs/abstract_graph.py +++ b/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}") diff --git a/scrapegraphai/models/__init__.py b/scrapegraphai/models/__init__.py index 0031f9a1..8ec6e490 100644 --- a/scrapegraphai/models/__init__.py +++ b/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"] diff --git a/scrapegraphai/models/nvidia.py b/scrapegraphai/models/nvidia.py new file mode 100644 index 00000000..7fcb0d32 --- /dev/null +++ b/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)