mirror of
https://github.com/VinciGit00/Scrapegraph-ai.git
synced 2026-07-12 21:01:56 +08:00
252 lines
9.1 KiB
Python
252 lines
9.1 KiB
Python
"""
|
|
AbstractGraph Module
|
|
"""
|
|
from abc import ABC, abstractmethod
|
|
from typing import Optional
|
|
import uuid
|
|
import warnings
|
|
from pydantic import BaseModel
|
|
from langchain.chat_models import init_chat_model
|
|
from langchain_core.rate_limiters import InMemoryRateLimiter
|
|
from ..helpers import models_tokens
|
|
from ..models import (
|
|
OneApi,
|
|
DeepSeek
|
|
)
|
|
from ..utils.logging import set_verbosity_warning, set_verbosity_info
|
|
|
|
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[BaseModel] = None):
|
|
|
|
if config.get("llm").get("temperature") is None:
|
|
config["llm"]["temperature"] = 0
|
|
|
|
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.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,
|
|
}
|
|
|
|
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 = {"temperature": 0, "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", "togetherai"}
|
|
|
|
split_model_provider = llm_params["model"].split("/", 1)
|
|
llm_params["model_provider"] = split_model_provider[0]
|
|
llm_params["model"] = split_model_provider[1]
|
|
|
|
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 "model_tokens" not in llm_params:
|
|
try:
|
|
self.model_token = models_tokens[llm_params["model_provider"]][llm_params["model"]]
|
|
except KeyError:
|
|
print(f"""Model {llm_params['model_provider']}/{llm_params['model']} not found,
|
|
using default token size (8192)""")
|
|
self.model_token = 8192
|
|
else:
|
|
self.model_token = llm_params["model_tokens"]
|
|
|
|
try:
|
|
if llm_params["model_provider"] not in \
|
|
{"oneapi","nvidia","ernie","deepseek","togetherai"}:
|
|
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 == "deepseek":
|
|
return DeepSeek(**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 == "togetherai":
|
|
try:
|
|
from langchain_together import ChatTogether
|
|
except ImportError:
|
|
raise ImportError("""The langchain_together module is not installed.
|
|
Please install it using `pip install scrapegraphai[other-language-models]`.""")
|
|
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 scrapegraphai[other-language-models]`.""")
|
|
return ChatNVIDIA(**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.
|
|
"""
|