diff --git a/scrapegraphai/graphs/__init__.py b/scrapegraphai/graphs/__init__.py index 5a38574b..8819811c 100644 --- a/scrapegraphai/graphs/__init__.py +++ b/scrapegraphai/graphs/__init__.py @@ -16,8 +16,8 @@ from .pdf_scraper_graph import PDFScraperGraph from .omni_scraper_graph import OmniScraperGraph from .omni_search_graph import OmniSearchGraph from .smart_scraper_multi_graph import SmartScraperMultiGraph -from .pdf_scraper_multi import PdfScraperMultiGraph -from .json_scraper_multi import JSONScraperMultiGraph -from .csv_scraper_graph_multi import CSVScraperMultiGraph -from .xml_scraper_graph_multi import XMLScraperMultiGraph +from .pdf_scraper_multi_graph import PdfScraperMultiGraph +from .json_scraper_multi_graph import JSONScraperMultiGraph +from .csv_scraper_multi_graph import CSVScraperMultiGraph +from .xml_scraper_multi_graph import XMLScraperMultiGraph from .script_creator_multi_graph import ScriptCreatorMultiGraph diff --git a/scrapegraphai/graphs/csv_scraper_graph_multi.py b/scrapegraphai/graphs/csv_scraper_multi_graph.py similarity index 100% rename from scrapegraphai/graphs/csv_scraper_graph_multi.py rename to scrapegraphai/graphs/csv_scraper_multi_graph.py diff --git a/scrapegraphai/graphs/json_scraper_multi.py b/scrapegraphai/graphs/json_scraper_multi_graph.py similarity index 95% rename from scrapegraphai/graphs/json_scraper_multi.py rename to scrapegraphai/graphs/json_scraper_multi_graph.py index 2010c856..f86fdc67 100644 --- a/scrapegraphai/graphs/json_scraper_multi.py +++ b/scrapegraphai/graphs/json_scraper_multi_graph.py @@ -4,6 +4,7 @@ JSONScraperMultiGraph Module from copy import copy, deepcopy from typing import List, Optional +from pydantic import BaseModel from .base_graph import BaseGraph from .abstract_graph import AbstractGraph @@ -42,7 +43,7 @@ class JSONScraperMultiGraph(AbstractGraph): >>> result = search_graph.run() """ - def __init__(self, prompt: str, source: List[str], config: dict, schema: Optional[str] = None): + def __init__(self, prompt: str, source: List[str], config: dict, schema: Optional[BaseModel] = None): self.max_results = config.get("max_results", 3) @@ -51,6 +52,8 @@ class JSONScraperMultiGraph(AbstractGraph): else: self.copy_config = deepcopy(config) + self.copy_schema = deepcopy(schema) + super().__init__(prompt, config, source, schema) def _create_graph(self) -> BaseGraph: @@ -69,6 +72,7 @@ class JSONScraperMultiGraph(AbstractGraph): prompt="", source="", config=self.copy_config, + schema=self.copy_schema ) # ************************************************ diff --git a/scrapegraphai/graphs/omni_search_graph.py b/scrapegraphai/graphs/omni_search_graph.py index 2185dd09..d5783729 100644 --- a/scrapegraphai/graphs/omni_search_graph.py +++ b/scrapegraphai/graphs/omni_search_graph.py @@ -53,6 +53,8 @@ class OmniSearchGraph(AbstractGraph): else: self.copy_config = deepcopy(config) + self.copy_schema = deepcopy(schema) + super().__init__(prompt, config, schema) def _create_graph(self) -> BaseGraph: @@ -70,7 +72,8 @@ class OmniSearchGraph(AbstractGraph): omni_scraper_instance = OmniScraperGraph( prompt="", source="", - config=self.copy_config + config=self.copy_config, + schema=self.copy_schema ) # ************************************************ diff --git a/scrapegraphai/graphs/pdf_scraper_multi.py b/scrapegraphai/graphs/pdf_scraper_multi_graph.py similarity index 100% rename from scrapegraphai/graphs/pdf_scraper_multi.py rename to scrapegraphai/graphs/pdf_scraper_multi_graph.py diff --git a/scrapegraphai/graphs/xml_scraper_graph_multi.py b/scrapegraphai/graphs/xml_scraper_multi_graph.py similarity index 95% rename from scrapegraphai/graphs/xml_scraper_graph_multi.py rename to scrapegraphai/graphs/xml_scraper_multi_graph.py index 1198f580..a9127d5b 100644 --- a/scrapegraphai/graphs/xml_scraper_graph_multi.py +++ b/scrapegraphai/graphs/xml_scraper_multi_graph.py @@ -4,6 +4,7 @@ XMLScraperMultiGraph Module from copy import copy, deepcopy from typing import List, Optional +from pydantic import BaseModel from .base_graph import BaseGraph from .abstract_graph import AbstractGraph @@ -43,7 +44,7 @@ class XMLScraperMultiGraph(AbstractGraph): >>> result = search_graph.run() """ - def __init__(self, prompt: str, source: List[str], config: dict, schema: Optional[str] = None): + def __init__(self, prompt: str, source: List[str], config: dict, schema: Optional[BaseModel] = None): self.max_results = config.get("max_results", 3) @@ -52,6 +53,8 @@ class XMLScraperMultiGraph(AbstractGraph): else: self.copy_config = deepcopy(config) + self.copy_schema = deepcopy(schema) + super().__init__(prompt, config, source, schema) def _create_graph(self) -> BaseGraph: @@ -70,6 +73,7 @@ class XMLScraperMultiGraph(AbstractGraph): prompt="", source="", config=self.copy_config, + schema=self.copy_schema ) # ************************************************