From 2b5bd80a945a24072e578133eacc751feeec6188 Mon Sep 17 00:00:00 2001 From: Elias Thouant Date: Wed, 15 Jan 2025 21:28:43 +0100 Subject: [PATCH] fix: Schema parameter type --- scrapegraphai/graphs/abstract_graph.py | 4 ++-- scrapegraphai/graphs/code_generator_graph.py | 8 ++++++-- scrapegraphai/graphs/csv_scraper_graph.py | 10 +++++++--- scrapegraphai/graphs/csv_scraper_multi_graph.py | 4 ++-- scrapegraphai/graphs/depth_search_graph.py | 8 ++++++-- scrapegraphai/graphs/document_scraper_graph.py | 8 ++++++-- scrapegraphai/graphs/document_scraper_multi_graph.py | 4 ++-- scrapegraphai/graphs/json_scraper_graph.py | 8 ++++++-- scrapegraphai/graphs/json_scraper_multi_graph.py | 4 ++-- scrapegraphai/graphs/omni_scraper_graph.py | 8 ++++++-- scrapegraphai/graphs/omni_search_graph.py | 6 ++++-- scrapegraphai/graphs/screenshot_scraper_graph.py | 10 +++++++--- scrapegraphai/graphs/script_creator_graph.py | 8 ++++++-- scrapegraphai/graphs/script_creator_multi_graph.py | 4 ++-- scrapegraphai/graphs/search_graph.py | 6 ++++-- scrapegraphai/graphs/search_link_graph.py | 6 ++++-- scrapegraphai/graphs/smart_scraper_graph.py | 8 ++++++-- scrapegraphai/graphs/smart_scraper_lite_graph.py | 4 ++-- .../graphs/smart_scraper_multi_concat_graph.py | 4 ++-- scrapegraphai/graphs/smart_scraper_multi_graph.py | 4 ++-- scrapegraphai/graphs/smart_scraper_multi_lite_graph.py | 4 ++-- scrapegraphai/graphs/speech_graph.py | 8 ++++++-- scrapegraphai/graphs/xml_scraper_graph.py | 8 ++++++-- scrapegraphai/graphs/xml_scraper_multi_graph.py | 4 ++-- scrapegraphai/nodes/graph_iterator_node.py | 4 ++-- 25 files changed, 102 insertions(+), 52 deletions(-) diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py index a56c9954..86311cb6 100644 --- a/scrapegraphai/graphs/abstract_graph.py +++ b/scrapegraphai/graphs/abstract_graph.py @@ -6,7 +6,7 @@ import asyncio import uuid import warnings from abc import ABC, abstractmethod -from typing import Optional +from typing import Optional, Type from langchain.chat_models import init_chat_model from langchain_core.rate_limiters import InMemoryRateLimiter @@ -51,7 +51,7 @@ class AbstractGraph(ABC): prompt: str, config: dict, source: Optional[str] = None, - schema: Optional[BaseModel] = None, + schema: Optional[Type[BaseModel]] = None, ): if config.get("llm").get("temperature") is None: config["llm"]["temperature"] = 0 diff --git a/scrapegraphai/graphs/code_generator_graph.py b/scrapegraphai/graphs/code_generator_graph.py index 5b5b23d8..506b87a1 100644 --- a/scrapegraphai/graphs/code_generator_graph.py +++ b/scrapegraphai/graphs/code_generator_graph.py @@ -2,7 +2,7 @@ SmartScraperGraph Module """ -from typing import Optional +from typing import Optional, Type from pydantic import BaseModel @@ -56,7 +56,11 @@ class CodeGeneratorGraph(AbstractGraph): """ def __init__( - self, prompt: str, source: str, config: dict, schema: Optional[BaseModel] = None + self, + prompt: str, + source: str, + config: dict, + schema: Optional[Type[BaseModel]] = None, ): super().__init__(prompt, config, source, schema) diff --git a/scrapegraphai/graphs/csv_scraper_graph.py b/scrapegraphai/graphs/csv_scraper_graph.py index b2bcc712..11a18553 100644 --- a/scrapegraphai/graphs/csv_scraper_graph.py +++ b/scrapegraphai/graphs/csv_scraper_graph.py @@ -2,7 +2,7 @@ Module for creating the smart scraper """ -from typing import Optional +from typing import Optional, Type from pydantic import BaseModel @@ -22,7 +22,7 @@ class CSVScraperGraph(AbstractGraph): config (dict): Additional configuration parameters needed by some nodes in the graph. Methods: - __init__ (prompt: str, source: str, config: dict, schema: Optional[BaseModel] = None): + __init__ (prompt: str, source: str, config: dict, schema: Optional[Type[BaseModel]] = None): Initializes the CSVScraperGraph with a prompt, source, and configuration. __init__ initializes the CSVScraperGraph class. It requires the user's prompt as input, @@ -49,7 +49,11 @@ class CSVScraperGraph(AbstractGraph): """ def __init__( - self, prompt: str, source: str, config: dict, schema: Optional[BaseModel] = None + self, + prompt: str, + source: str, + config: dict, + schema: Optional[Type[BaseModel]] = None, ): """ Initializes the CSVScraperGraph with a prompt, source, and configuration. diff --git a/scrapegraphai/graphs/csv_scraper_multi_graph.py b/scrapegraphai/graphs/csv_scraper_multi_graph.py index b495c6d7..5e3a398d 100644 --- a/scrapegraphai/graphs/csv_scraper_multi_graph.py +++ b/scrapegraphai/graphs/csv_scraper_multi_graph.py @@ -3,7 +3,7 @@ CSVScraperMultiGraph Module """ from copy import deepcopy -from typing import List, Optional +from typing import List, Optional, Type from pydantic import BaseModel @@ -47,7 +47,7 @@ class CSVScraperMultiGraph(AbstractGraph): prompt: str, source: List[str], config: dict, - schema: Optional[BaseModel] = None, + schema: Optional[Type[BaseModel]] = None, ): self.copy_config = safe_deepcopy(config) diff --git a/scrapegraphai/graphs/depth_search_graph.py b/scrapegraphai/graphs/depth_search_graph.py index 4dd0e49d..b68aa21d 100644 --- a/scrapegraphai/graphs/depth_search_graph.py +++ b/scrapegraphai/graphs/depth_search_graph.py @@ -2,7 +2,7 @@ depth search graph Module """ -from typing import Optional +from typing import Optional, Type from pydantic import BaseModel @@ -54,7 +54,11 @@ class DepthSearchGraph(AbstractGraph): """ def __init__( - self, prompt: str, source: str, config: dict, schema: Optional[BaseModel] = None + self, + prompt: str, + source: str, + config: dict, + schema: Optional[Type[BaseModel]] = None, ): super().__init__(prompt, config, source, schema) diff --git a/scrapegraphai/graphs/document_scraper_graph.py b/scrapegraphai/graphs/document_scraper_graph.py index 92a0f3d1..012ab0b7 100644 --- a/scrapegraphai/graphs/document_scraper_graph.py +++ b/scrapegraphai/graphs/document_scraper_graph.py @@ -2,7 +2,7 @@ This module implements the Document Scraper Graph for the ScrapeGraphAI application. """ -from typing import Optional +from typing import Optional, Type from pydantic import BaseModel @@ -44,7 +44,11 @@ class DocumentScraperGraph(AbstractGraph): """ def __init__( - self, prompt: str, source: str, config: dict, schema: Optional[BaseModel] = None + self, + prompt: str, + source: str, + config: dict, + schema: Optional[Type[BaseModel]] = None, ): super().__init__(prompt, config, source, schema) diff --git a/scrapegraphai/graphs/document_scraper_multi_graph.py b/scrapegraphai/graphs/document_scraper_multi_graph.py index 555b3964..edd8c7ef 100644 --- a/scrapegraphai/graphs/document_scraper_multi_graph.py +++ b/scrapegraphai/graphs/document_scraper_multi_graph.py @@ -3,7 +3,7 @@ DocumentScraperMultiGraph Module """ from copy import deepcopy -from typing import List, Optional +from typing import List, Optional, Type from pydantic import BaseModel @@ -47,7 +47,7 @@ class DocumentScraperMultiGraph(AbstractGraph): prompt: str, source: List[str], config: dict, - schema: Optional[BaseModel] = None, + schema: Optional[Type[BaseModel]] = None, ): self.copy_config = safe_deepcopy(config) self.copy_schema = deepcopy(schema) diff --git a/scrapegraphai/graphs/json_scraper_graph.py b/scrapegraphai/graphs/json_scraper_graph.py index 29e96497..5d0dfbff 100644 --- a/scrapegraphai/graphs/json_scraper_graph.py +++ b/scrapegraphai/graphs/json_scraper_graph.py @@ -2,7 +2,7 @@ JSONScraperGraph Module """ -from typing import Optional +from typing import Optional, Type from pydantic import BaseModel @@ -42,7 +42,11 @@ class JSONScraperGraph(AbstractGraph): """ def __init__( - self, prompt: str, source: str, config: dict, schema: Optional[BaseModel] = None + self, + prompt: str, + source: str, + config: dict, + schema: Optional[Type[BaseModel]] = None, ): super().__init__(prompt, config, source, schema) diff --git a/scrapegraphai/graphs/json_scraper_multi_graph.py b/scrapegraphai/graphs/json_scraper_multi_graph.py index 7984c7b9..6623718b 100644 --- a/scrapegraphai/graphs/json_scraper_multi_graph.py +++ b/scrapegraphai/graphs/json_scraper_multi_graph.py @@ -3,7 +3,7 @@ JSONScraperMultiGraph Module """ from copy import deepcopy -from typing import List, Optional +from typing import List, Optional, Type from pydantic import BaseModel @@ -47,7 +47,7 @@ class JSONScraperMultiGraph(AbstractGraph): prompt: str, source: List[str], config: dict, - schema: Optional[BaseModel] = None, + schema: Optional[Type[BaseModel]] = None, ): self.copy_config = safe_deepcopy(config) diff --git a/scrapegraphai/graphs/omni_scraper_graph.py b/scrapegraphai/graphs/omni_scraper_graph.py index c2c13f88..ef55dc75 100644 --- a/scrapegraphai/graphs/omni_scraper_graph.py +++ b/scrapegraphai/graphs/omni_scraper_graph.py @@ -2,7 +2,7 @@ This module implements the Omni Scraper Graph for the ScrapeGraphAI application. """ -from typing import Optional +from typing import Optional, Type from pydantic import BaseModel @@ -47,7 +47,11 @@ class OmniScraperGraph(AbstractGraph): """ def __init__( - self, prompt: str, source: str, config: dict, schema: Optional[BaseModel] = None + self, + prompt: str, + source: str, + config: dict, + schema: Optional[Type[BaseModel]] = None, ): self.max_images = 5 if config is None else config.get("max_images", 5) diff --git a/scrapegraphai/graphs/omni_search_graph.py b/scrapegraphai/graphs/omni_search_graph.py index a02e31c6..c30033ab 100644 --- a/scrapegraphai/graphs/omni_search_graph.py +++ b/scrapegraphai/graphs/omni_search_graph.py @@ -3,7 +3,7 @@ OmniSearchGraph Module """ from copy import deepcopy -from typing import Optional +from typing import Optional, Type from pydantic import BaseModel @@ -41,7 +41,9 @@ class OmniSearchGraph(AbstractGraph): >>> result = search_graph.run() """ - def __init__(self, prompt: str, config: dict, schema: Optional[BaseModel] = None): + def __init__( + self, prompt: str, config: dict, schema: Optional[Type[BaseModel]] = None + ): self.max_results = config.get("max_results", 3) diff --git a/scrapegraphai/graphs/screenshot_scraper_graph.py b/scrapegraphai/graphs/screenshot_scraper_graph.py index c37e34f2..765bd428 100644 --- a/scrapegraphai/graphs/screenshot_scraper_graph.py +++ b/scrapegraphai/graphs/screenshot_scraper_graph.py @@ -2,7 +2,7 @@ ScreenshotScraperGraph Module """ -from typing import Optional +from typing import Optional, Type from pydantic import BaseModel @@ -21,7 +21,7 @@ class ScreenshotScraperGraph(AbstractGraph): source (str): The source URL or image link to scrape from. Methods: - __init__(prompt: str, source: str, config: dict, schema: Optional[BaseModel] = None) + __init__(prompt: str, source: str, config: dict, schema: Optional[Type[BaseModel]] = None) Initializes the ScreenshotScraperGraph instance with the given prompt, source, and configuration parameters. @@ -33,7 +33,11 @@ class ScreenshotScraperGraph(AbstractGraph): """ def __init__( - self, prompt: str, source: str, config: dict, schema: Optional[BaseModel] = None + self, + prompt: str, + source: str, + config: dict, + schema: Optional[Type[BaseModel]] = None, ): super().__init__(prompt, config, source, schema) diff --git a/scrapegraphai/graphs/script_creator_graph.py b/scrapegraphai/graphs/script_creator_graph.py index 98dd05e4..4d373a28 100644 --- a/scrapegraphai/graphs/script_creator_graph.py +++ b/scrapegraphai/graphs/script_creator_graph.py @@ -2,7 +2,7 @@ ScriptCreatorGraph Module """ -from typing import Optional +from typing import Optional, Type from pydantic import BaseModel @@ -44,7 +44,11 @@ class ScriptCreatorGraph(AbstractGraph): """ def __init__( - self, prompt: str, source: str, config: dict, schema: Optional[BaseModel] = None + self, + prompt: str, + source: str, + config: dict, + schema: Optional[Type[BaseModel]] = None, ): self.library = config["library"] diff --git a/scrapegraphai/graphs/script_creator_multi_graph.py b/scrapegraphai/graphs/script_creator_multi_graph.py index b3f21025..a8e01729 100644 --- a/scrapegraphai/graphs/script_creator_multi_graph.py +++ b/scrapegraphai/graphs/script_creator_multi_graph.py @@ -3,7 +3,7 @@ ScriptCreatorMultiGraph Module """ from copy import deepcopy -from typing import List, Optional +from typing import List, Optional, Type from pydantic import BaseModel @@ -46,7 +46,7 @@ class ScriptCreatorMultiGraph(AbstractGraph): prompt: str, source: List[str], config: dict, - schema: Optional[BaseModel] = None, + schema: Optional[Type[BaseModel]] = None, ): self.copy_config = safe_deepcopy(config) diff --git a/scrapegraphai/graphs/search_graph.py b/scrapegraphai/graphs/search_graph.py index 394afb2a..2458c1d8 100644 --- a/scrapegraphai/graphs/search_graph.py +++ b/scrapegraphai/graphs/search_graph.py @@ -3,7 +3,7 @@ SearchGraph Module """ from copy import deepcopy -from typing import List, Optional +from typing import List, Optional, Type from pydantic import BaseModel @@ -42,7 +42,9 @@ class SearchGraph(AbstractGraph): >>> print(search_graph.get_considered_urls()) """ - def __init__(self, prompt: str, config: dict, schema: Optional[BaseModel] = None): + def __init__( + self, prompt: str, config: dict, schema: Optional[Type[BaseModel]] = None + ): self.max_results = config.get("max_results", 3) self.copy_config = safe_deepcopy(config) diff --git a/scrapegraphai/graphs/search_link_graph.py b/scrapegraphai/graphs/search_link_graph.py index ba781363..c46fe9be 100644 --- a/scrapegraphai/graphs/search_link_graph.py +++ b/scrapegraphai/graphs/search_link_graph.py @@ -2,7 +2,7 @@ SearchLinkGraph Module """ -from typing import Optional +from typing import Optional, Type from pydantic import BaseModel @@ -36,7 +36,9 @@ class SearchLinkGraph(AbstractGraph): """ - def __init__(self, source: str, config: dict, schema: Optional[BaseModel] = None): + def __init__( + self, source: str, config: dict, schema: Optional[Type[BaseModel]] = None + ): super().__init__("", config, source, schema) self.input_key = "url" if source.startswith("http") else "local_dir" diff --git a/scrapegraphai/graphs/smart_scraper_graph.py b/scrapegraphai/graphs/smart_scraper_graph.py index 7719979d..16964d84 100644 --- a/scrapegraphai/graphs/smart_scraper_graph.py +++ b/scrapegraphai/graphs/smart_scraper_graph.py @@ -2,7 +2,7 @@ SmartScraperGraph Module """ -from typing import Optional +from typing import Optional, Type from pydantic import BaseModel @@ -52,7 +52,11 @@ class SmartScraperGraph(AbstractGraph): """ def __init__( - self, prompt: str, source: str, config: dict, schema: Optional[BaseModel] = None + self, + prompt: str, + source: str, + config: dict, + schema: Optional[Type[BaseModel]] = None, ): super().__init__(prompt, config, source, schema) diff --git a/scrapegraphai/graphs/smart_scraper_lite_graph.py b/scrapegraphai/graphs/smart_scraper_lite_graph.py index 7769e21b..a845dc7a 100644 --- a/scrapegraphai/graphs/smart_scraper_lite_graph.py +++ b/scrapegraphai/graphs/smart_scraper_lite_graph.py @@ -2,7 +2,7 @@ SmartScraperGraph Module """ -from typing import Optional +from typing import Optional, Type from pydantic import BaseModel @@ -44,7 +44,7 @@ class SmartScraperLiteGraph(AbstractGraph): source: str, config: dict, prompt: str = "", - schema: Optional[BaseModel] = None, + schema: Optional[Type[BaseModel]] = None, ): super().__init__(prompt, config, source, schema) diff --git a/scrapegraphai/graphs/smart_scraper_multi_concat_graph.py b/scrapegraphai/graphs/smart_scraper_multi_concat_graph.py index 8c856c01..ebd7b936 100644 --- a/scrapegraphai/graphs/smart_scraper_multi_concat_graph.py +++ b/scrapegraphai/graphs/smart_scraper_multi_concat_graph.py @@ -3,7 +3,7 @@ SmartScraperMultiCondGraph Module with ConditionalNode """ from copy import deepcopy -from typing import List, Optional +from typing import List, Optional, Type from pydantic import BaseModel @@ -51,7 +51,7 @@ class SmartScraperMultiConcatGraph(AbstractGraph): prompt: str, source: List[str], config: dict, - schema: Optional[BaseModel] = None, + schema: Optional[Type[BaseModel]] = None, ): self.copy_config = safe_deepcopy(config) diff --git a/scrapegraphai/graphs/smart_scraper_multi_graph.py b/scrapegraphai/graphs/smart_scraper_multi_graph.py index fa4bfd0f..a0518fb7 100644 --- a/scrapegraphai/graphs/smart_scraper_multi_graph.py +++ b/scrapegraphai/graphs/smart_scraper_multi_graph.py @@ -3,7 +3,7 @@ SmartScraperMultiGraph Module """ from copy import deepcopy -from typing import List, Optional +from typing import List, Optional, Type from pydantic import BaseModel @@ -53,7 +53,7 @@ class SmartScraperMultiGraph(AbstractGraph): prompt: str, source: List[str], config: dict, - schema: Optional[BaseModel] = None, + schema: Optional[Type[BaseModel]] = None, ): self.max_results = config.get("max_results", 3) diff --git a/scrapegraphai/graphs/smart_scraper_multi_lite_graph.py b/scrapegraphai/graphs/smart_scraper_multi_lite_graph.py index ea57bab0..d212b08a 100644 --- a/scrapegraphai/graphs/smart_scraper_multi_lite_graph.py +++ b/scrapegraphai/graphs/smart_scraper_multi_lite_graph.py @@ -3,7 +3,7 @@ SmartScraperMultiGraph Module """ from copy import deepcopy -from typing import List, Optional +from typing import List, Optional, Type from pydantic import BaseModel @@ -53,7 +53,7 @@ class SmartScraperMultiLiteGraph(AbstractGraph): prompt: str, source: List[str], config: dict, - schema: Optional[BaseModel] = None, + schema: Optional[Type[BaseModel]] = None, ): self.copy_config = safe_deepcopy(config) diff --git a/scrapegraphai/graphs/speech_graph.py b/scrapegraphai/graphs/speech_graph.py index 32d5be8c..11caea9b 100644 --- a/scrapegraphai/graphs/speech_graph.py +++ b/scrapegraphai/graphs/speech_graph.py @@ -2,7 +2,7 @@ SpeechGraph Module """ -from typing import Optional +from typing import Optional, Type from pydantic import BaseModel @@ -44,7 +44,11 @@ class SpeechGraph(AbstractGraph): """ def __init__( - self, prompt: str, source: str, config: dict, schema: Optional[BaseModel] = None + self, + prompt: str, + source: str, + config: dict, + schema: Optional[Type[BaseModel]] = None, ): super().__init__(prompt, config, source, schema) diff --git a/scrapegraphai/graphs/xml_scraper_graph.py b/scrapegraphai/graphs/xml_scraper_graph.py index c7dcd62e..162aa322 100644 --- a/scrapegraphai/graphs/xml_scraper_graph.py +++ b/scrapegraphai/graphs/xml_scraper_graph.py @@ -2,7 +2,7 @@ XMLScraperGraph Module """ -from typing import Optional +from typing import Optional, Type from pydantic import BaseModel @@ -44,7 +44,11 @@ class XMLScraperGraph(AbstractGraph): """ def __init__( - self, prompt: str, source: str, config: dict, schema: Optional[BaseModel] = None + self, + prompt: str, + source: str, + config: dict, + schema: Optional[Type[BaseModel]] = None, ): super().__init__(prompt, config, source, schema) diff --git a/scrapegraphai/graphs/xml_scraper_multi_graph.py b/scrapegraphai/graphs/xml_scraper_multi_graph.py index f44bf343..2a3848c9 100644 --- a/scrapegraphai/graphs/xml_scraper_multi_graph.py +++ b/scrapegraphai/graphs/xml_scraper_multi_graph.py @@ -3,7 +3,7 @@ XMLScraperMultiGraph Module """ from copy import deepcopy -from typing import List, Optional +from typing import List, Optional, Type from pydantic import BaseModel @@ -47,7 +47,7 @@ class XMLScraperMultiGraph(AbstractGraph): prompt: str, source: List[str], config: dict, - schema: Optional[BaseModel] = None, + schema: Optional[Type[BaseModel]] = None, ): self.copy_config = safe_deepcopy(config) diff --git a/scrapegraphai/nodes/graph_iterator_node.py b/scrapegraphai/nodes/graph_iterator_node.py index 15ae5524..82171258 100644 --- a/scrapegraphai/nodes/graph_iterator_node.py +++ b/scrapegraphai/nodes/graph_iterator_node.py @@ -3,7 +3,7 @@ GraphIterator Module """ import asyncio -from typing import List, Optional +from typing import List, Optional, Type from pydantic import BaseModel from tqdm.asyncio import tqdm @@ -34,7 +34,7 @@ class GraphIteratorNode(BaseNode): output: List[str], node_config: Optional[dict] = None, node_name: str = "GraphIterator", - schema: Optional[BaseModel] = None, + schema: Optional[Type[BaseModel]] = None, ): super().__init__(node_name, "node", input, output, 2, node_config)