feat: added graph timeout and fixed model_tokens param (#810 #856 #853)

This commit is contained in:
PeriniM 2025-01-06 18:48:55 +01:00
parent 5f2df70040
commit 01a331afa5
6 changed files with 16 additions and 8 deletions

View File

@ -1,21 +1,24 @@
""" """
Basic example of scraping pipeline using SmartScraper Basic example of scraping pipeline using SmartScraper
""" """
from scrapegraphai.graphs import SmartScraperGraph from scrapegraphai.graphs import SmartScraperGraph
from scrapegraphai.utils import prettify_exec_info from scrapegraphai.utils import prettify_exec_info
# ************************************************ # ************************************************
# Define the configuration for the graph # Define the configuration for the graph
# ************************************************ # ************************************************
graph_config = { graph_config = {
"llm": { "llm": {
"model": "ollama/llama3.1", "model": "ollama/llama3.2:3b",
"temperature": 0, "temperature": 0,
"format": "json", # Ollama needs the format to be specified explicitly "format": "json", # Ollama needs the format to be specified explicitly
# "base_url": "http://localhost:11434", # set ollama URL arbitrarily # "base_url": "http://localhost:11434", # set ollama URL arbitrarily
"model_tokens": 1024,
}, },
"verbose": True, "verbose": True,
"headless": False "headless": False,
} }
# ************************************************ # ************************************************
@ -24,7 +27,7 @@ graph_config = {
smart_scraper_graph = SmartScraperGraph( smart_scraper_graph = SmartScraperGraph(
prompt="Find some information about what does the company do, the name and a contact email.", prompt="Find some information about what does the company do, the name and a contact email.",
source="https://scrapegraphai.com/", source="https://scrapegraphai.com/",
config=graph_config config=graph_config,
) )
result = smart_scraper_graph.run() result = smart_scraper_graph.run()

View File

@ -20,7 +20,7 @@ load_dotenv()
graph_config = { graph_config = {
"llm": { "llm": {
"api_key": os.getenv("OPENAI_API_KEY"), "api_key": os.getenv("OPENAI_API_KEY"),
"model": "openai/gpt-4o00", "model": "openai/gpt-4o-mini",
}, },
"verbose": True, "verbose": True,
"headless": False, "headless": False,

View File

@ -31,6 +31,7 @@ dependencies = [
"async-timeout>=4.0.3", "async-timeout>=4.0.3",
"simpleeval>=1.0.0", "simpleeval>=1.0.0",
"jsonschema>=4.23.0", "jsonschema>=4.23.0",
"transformers>=4.46.3",
] ]
readme = "README.md" readme = "README.md"

View File

@ -68,6 +68,7 @@ class AbstractGraph(ABC):
self.browser_base = self.config.get("browser_base") self.browser_base = self.config.get("browser_base")
self.scrape_do = self.config.get("scrape_do") self.scrape_do = self.config.get("scrape_do")
self.storage_state = self.config.get("storage_state") self.storage_state = self.config.get("storage_state")
self.timeout = self.config.get("timeout", 480)
self.graph = self._create_graph() self.graph = self._create_graph()
self.final_state = None self.final_state = None
@ -86,6 +87,7 @@ class AbstractGraph(ABC):
"loader_kwargs": self.loader_kwargs, "loader_kwargs": self.loader_kwargs,
"llm_model": self.llm_model, "llm_model": self.llm_model,
"cache_path": self.cache_path, "cache_path": self.cache_path,
"timeout": self.timeout,
} }
self.set_common_params(common_params, overwrite=True) self.set_common_params(common_params, overwrite=True)
@ -194,7 +196,7 @@ class AbstractGraph(ABC):
If possible, try to use a model instance instead.""" If possible, try to use a model instance instead."""
) )
if "model_tokens" not in llm_params: if llm_params.get("model_tokens", None) is None:
try: try:
self.model_token = models_tokens[llm_params["model_provider"]][ self.model_token = models_tokens[llm_params["model_provider"]][
llm_params["model"] llm_params["model"]

View File

@ -66,7 +66,7 @@ class GenerateAnswerNode(BaseNode):
self.script_creator = node_config.get("script_creator", False) self.script_creator = node_config.get("script_creator", False)
self.is_md_scraper = node_config.get("is_md_scraper", False) self.is_md_scraper = node_config.get("is_md_scraper", False)
self.additional_info = node_config.get("additional_info") self.additional_info = node_config.get("additional_info")
self.timeout = node_config.get("timeout", 120) self.timeout = node_config.get("timeout", 480)
def invoke_with_timeout(self, chain, inputs, timeout): def invoke_with_timeout(self, chain, inputs, timeout):
"""Helper method to invoke chain with timeout""" """Helper method to invoke chain with timeout"""

View File

@ -3429,7 +3429,7 @@ wheels = [
[[package]] [[package]]
name = "scrapegraphai" name = "scrapegraphai"
version = "1.34.3b1" version = "1.35.0b2"
source = { editable = "." } source = { editable = "." }
dependencies = [ dependencies = [
{ name = "async-timeout", version = "4.0.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "async-timeout", version = "4.0.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
@ -3452,6 +3452,7 @@ dependencies = [
{ name = "simpleeval" }, { name = "simpleeval" },
{ name = "tiktoken" }, { name = "tiktoken" },
{ name = "tqdm" }, { name = "tqdm" },
{ name = "transformers" },
{ name = "undetected-playwright" }, { name = "undetected-playwright" },
] ]
@ -3515,6 +3516,7 @@ requires-dist = [
{ name = "surya-ocr", marker = "extra == 'ocr'", specifier = ">=0.5.0" }, { name = "surya-ocr", marker = "extra == 'ocr'", specifier = ">=0.5.0" },
{ name = "tiktoken", specifier = ">=0.7" }, { name = "tiktoken", specifier = ">=0.7" },
{ name = "tqdm", specifier = ">=4.66.4" }, { name = "tqdm", specifier = ">=4.66.4" },
{ name = "transformers", specifier = ">=4.46.3" },
{ name = "undetected-playwright", specifier = ">=0.3.0" }, { name = "undetected-playwright", specifier = ">=0.3.0" },
] ]