""" AbstractGraph Module """ import asyncio import uuid import warnings from abc import ABC, abstractmethod from typing import Optional, Type from langchain.chat_models import init_chat_model from langchain_core.rate_limiters import InMemoryRateLimiter from pydantic import BaseModel from ..helpers import models_tokens from ..models import XAI, CLoD, DeepSeek, MiniMax, Nvidia, OneApi from ..utils.logging import get_logger, set_verbosity_info, set_verbosity_warning logger = get_logger(__name__) # ANSI escape sequence for hyperlink CLICKABLE_URL = ( "\033]8;;https://scrapegraphai.com\033\\https://scrapegraphai.com\033]8;;\033\\" ) class AbstractGraph(ABC): """ Scaffolding class for creating a graph representation and executing it. prompt (str): The prompt for the graph. source (str): The source of the graph. config (dict): Configuration parameters for the graph. schema (BaseModel): The schema for the graph output. llm_model: An instance of a language model client, configured for generating answers. verbose (bool): A flag indicating whether to show print statements during execution. headless (bool): A flag indicating whether to run the graph in headless mode. Args: prompt (str): The prompt for the graph. config (dict): Configuration parameters for the graph. source (str, optional): The source of the graph. schema (str, optional): The schema for the graph output. Example: >>> class MyGraph(AbstractGraph): ... def _create_graph(self): ... # Implementation of graph creation here ... return graph ... >>> my_graph = MyGraph("Example Graph", {"llm": {"model": "gpt-3.5-turbo"}}, "example_source") >>> result = my_graph.run() """ def __init__( self, prompt: str, config: dict, source: Optional[str] = None, schema: Optional[Type[BaseModel]] = None, ): self.prompt = prompt self.source = source self.config = config self.schema = schema self.llm_model = self._create_llm(config["llm"]) self.verbose = False if config is None else config.get("verbose", False) self.headless = True if self.config is None else config.get("headless", True) self.loader_kwargs = self.config.get("loader_kwargs", {}) self.cache_path = self.config.get("cache_path", False) self.browser_base = self.config.get("browser_base") self.scrape_do = self.config.get("scrape_do") self.storage_state = self.config.get("storage_state") self.timeout = self.config.get("timeout", 480) self.graph = self._create_graph() self.final_state = None self.execution_info = None verbose = bool(config and config.get("verbose")) if verbose: set_verbosity_info() else: set_verbosity_warning() common_params = { "headless": self.headless, "verbose": self.verbose, "loader_kwargs": self.loader_kwargs, "llm_model": self.llm_model, "cache_path": self.cache_path, "timeout": self.timeout, } self.set_common_params(common_params, overwrite=True) self.burr_kwargs = config.get("burr_kwargs", None) if self.burr_kwargs is not None: self.graph.use_burr = True if "app_instance_id" not in self.burr_kwargs: self.burr_kwargs["app_instance_id"] = str(uuid.uuid4()) self.graph.burr_config = self.burr_kwargs def set_common_params(self, params: dict, overwrite=False): """ Pass parameters to every node in the graph unless otherwise defined in the graph. Args: params (dict): Common parameters and their values. """ for node in self.graph.nodes: node.update_config(params, overwrite) def _create_llm(self, llm_config: dict) -> object: """ Create a large language model instance based on the configuration provided. Args: llm_config (dict): Configuration parameters for the language model. Returns: object: An instance of the language model client. Raises: KeyError: If the model is not supported. """ llm_defaults = {"streaming": False} llm_params = {**llm_defaults, **llm_config} rate_limit_params = llm_params.pop("rate_limit", {}) if rate_limit_params: requests_per_second = rate_limit_params.get("requests_per_second") max_retries = rate_limit_params.get("max_retries") if requests_per_second is not None: with warnings.catch_warnings(): warnings.simplefilter("ignore") llm_params["rate_limiter"] = InMemoryRateLimiter( requests_per_second=requests_per_second ) if max_retries is not None: llm_params["max_retries"] = max_retries if "model_instance" in llm_params: try: self.model_token = llm_params["model_tokens"] except KeyError as exc: raise KeyError("model_tokens not specified") from exc return llm_params["model_instance"] known_providers = { "openai", "azure_openai", "google_genai", "google_vertexai", "ollama", "oneapi", "nvidia", "groq", "anthropic", "bedrock", "mistralai", "hugging_face", "deepseek", "ernie", "fireworks", "clod", "togetherai", "xai", "minimax", } if "/" in llm_params["model"]: split_model_provider = llm_params["model"].split("/", 1) llm_params["model_provider"] = split_model_provider[0] llm_params["model"] = split_model_provider[1] else: possible_providers = [ provider for provider, models_d in models_tokens.items() if llm_params["model"] in models_d ] if len(possible_providers) <= 0: raise ValueError( f"""Provider {llm_params["model_provider"]} is not supported. If possible, try to use a model instance instead.""" ) llm_params["model_provider"] = possible_providers[0] logger.info( "Found providers %s for model %s, using %s. " "If it was not intended please specify the model provider in the graph configuration", possible_providers, llm_params["model"], llm_params["model_provider"], ) if llm_params["model_provider"] not in known_providers: raise ValueError( f"""Provider {llm_params["model_provider"]} is not supported. If possible, try to use a model instance instead.""" ) if llm_params.get("model_tokens", None) is None: try: self.model_token = models_tokens[llm_params["model_provider"]][ llm_params["model"] ] except KeyError: logger.warning( "Max input tokens for model %s/%s not found, " "please specify the model_tokens parameter in the llm section of the graph configuration. " "Using default token size: 8192", llm_params["model_provider"], llm_params["model"], ) self.model_token = 8192 else: self.model_token = llm_params["model_tokens"] # Consumed by ScrapeGraphAI; must not be forwarded to the model client. llm_params.pop("model_tokens", None) try: if llm_params["model_provider"] not in { "oneapi", "nvidia", "ernie", "deepseek", "togetherai", "clod", "xai", "minimax", }: if llm_params["model_provider"] == "bedrock": llm_params["model_kwargs"] = { "temperature": llm_params.pop("temperature") } with warnings.catch_warnings(): warnings.simplefilter("ignore") return init_chat_model(**llm_params) else: model_provider = llm_params.pop("model_provider") if model_provider == "clod": return CLoD(**llm_params) if model_provider == "deepseek": return DeepSeek(**llm_params) if model_provider == "minimax": return MiniMax(**llm_params) if model_provider == "ernie": from langchain_community.chat_models import ErnieBotChat return ErnieBotChat(**llm_params) elif model_provider == "oneapi": return OneApi(**llm_params) elif model_provider == "xai": return XAI(**llm_params) elif model_provider == "togetherai": try: from langchain_together import ChatTogether except ImportError: raise ImportError( """The langchain_together module is not installed. Please install it using `pip install langchain-together`.""" ) return ChatTogether(**llm_params) elif model_provider == "nvidia": return Nvidia(**llm_params) except Exception as e: raise Exception(f"Error instancing model: {e}") def get_state(self, key=None) -> dict: """ "" Get the final state of the graph. Args: key (str, optional): The key of the final state to retrieve. Returns: dict: The final state of the graph. """ if key is not None: return self.final_state[key] return self.final_state def append_node(self, node): """ Add a node to the graph. Args: node (BaseNode): The node to add to the graph. """ self.graph.append_node(node) def get_execution_info(self): """ Returns the execution information of the graph. Returns: dict: The execution information of the graph. """ return self.execution_info @abstractmethod def _create_graph(self): """ Abstract method to create a graph representation. """ @abstractmethod def run(self) -> str: """ Abstract method to execute the graph and return the result. """ inputs = {"user_prompt": self.prompt, self.input_key: self.source} self.final_state, self.execution_info = self.graph.execute(inputs) result = self.final_state.get("answer", "No answer found.") return result async def run_safe_async(self) -> str: """ Executes the run process asynchronously safety. Returns: str: The answer to the prompt. """ loop = asyncio.get_event_loop() return await loop.run_in_executor(None, self.run)