From 01a331afa5fc6f6d6aea4f1969cbf41f0b25f5e0 Mon Sep 17 00:00:00 2001 From: PeriniM Date: Mon, 6 Jan 2025 18:48:55 +0100 Subject: [PATCH] =?UTF-8?q?feat:=20=E2=8F=B0added=20graph=20timeout=20and?= =?UTF-8?q?=20fixed=20model=5Ftokens=20param=20(#810=20#856=20#853)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- examples/local_models/smart_scraper_ollama.py | 11 +++++++---- examples/openai/smart_scraper_openai.py | 2 +- pyproject.toml | 1 + scrapegraphai/graphs/abstract_graph.py | 4 +++- scrapegraphai/nodes/generate_answer_node.py | 2 +- uv.lock | 4 +++- 6 files changed, 16 insertions(+), 8 deletions(-) diff --git a/examples/local_models/smart_scraper_ollama.py b/examples/local_models/smart_scraper_ollama.py index d5585ff7..61294eaf 100644 --- a/examples/local_models/smart_scraper_ollama.py +++ b/examples/local_models/smart_scraper_ollama.py @@ -1,21 +1,24 @@ -""" +""" Basic example of scraping pipeline using SmartScraper """ + from scrapegraphai.graphs import SmartScraperGraph from scrapegraphai.utils import prettify_exec_info + # ************************************************ # Define the configuration for the graph # ************************************************ graph_config = { "llm": { - "model": "ollama/llama3.1", + "model": "ollama/llama3.2:3b", "temperature": 0, "format": "json", # Ollama needs the format to be specified explicitly # "base_url": "http://localhost:11434", # set ollama URL arbitrarily + "model_tokens": 1024, }, "verbose": True, - "headless": False + "headless": False, } # ************************************************ @@ -24,7 +27,7 @@ graph_config = { smart_scraper_graph = SmartScraperGraph( prompt="Find some information about what does the company do, the name and a contact email.", source="https://scrapegraphai.com/", - config=graph_config + config=graph_config, ) result = smart_scraper_graph.run() diff --git a/examples/openai/smart_scraper_openai.py b/examples/openai/smart_scraper_openai.py index e0d634e7..cbf3e21e 100644 --- a/examples/openai/smart_scraper_openai.py +++ b/examples/openai/smart_scraper_openai.py @@ -20,7 +20,7 @@ load_dotenv() graph_config = { "llm": { "api_key": os.getenv("OPENAI_API_KEY"), - "model": "openai/gpt-4o00", + "model": "openai/gpt-4o-mini", }, "verbose": True, "headless": False, diff --git a/pyproject.toml b/pyproject.toml index 66514ae1..033bacca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,7 @@ dependencies = [ "async-timeout>=4.0.3", "simpleeval>=1.0.0", "jsonschema>=4.23.0", + "transformers>=4.46.3", ] readme = "README.md" diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py index 3225a1df..812aaf80 100644 --- a/scrapegraphai/graphs/abstract_graph.py +++ b/scrapegraphai/graphs/abstract_graph.py @@ -68,6 +68,7 @@ class AbstractGraph(ABC): 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 @@ -86,6 +87,7 @@ class AbstractGraph(ABC): "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) @@ -194,7 +196,7 @@ class AbstractGraph(ABC): 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: self.model_token = models_tokens[llm_params["model_provider"]][ llm_params["model"] diff --git a/scrapegraphai/nodes/generate_answer_node.py b/scrapegraphai/nodes/generate_answer_node.py index 4e313352..688300cf 100644 --- a/scrapegraphai/nodes/generate_answer_node.py +++ b/scrapegraphai/nodes/generate_answer_node.py @@ -66,7 +66,7 @@ class GenerateAnswerNode(BaseNode): self.script_creator = node_config.get("script_creator", False) self.is_md_scraper = node_config.get("is_md_scraper", False) 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): """Helper method to invoke chain with timeout""" diff --git a/uv.lock b/uv.lock index a7e4f0c1..ff6f13f2 100644 --- a/uv.lock +++ b/uv.lock @@ -3429,7 +3429,7 @@ wheels = [ [[package]] name = "scrapegraphai" -version = "1.34.3b1" +version = "1.35.0b2" source = { editable = "." } dependencies = [ { 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 = "tiktoken" }, { name = "tqdm" }, + { name = "transformers" }, { name = "undetected-playwright" }, ] @@ -3515,6 +3516,7 @@ requires-dist = [ { name = "surya-ocr", marker = "extra == 'ocr'", specifier = ">=0.5.0" }, { name = "tiktoken", specifier = ">=0.7" }, { name = "tqdm", specifier = ">=4.66.4" }, + { name = "transformers", specifier = ">=4.46.3" }, { name = "undetected-playwright", specifier = ">=0.3.0" }, ]