Skip to content

Commit 09cb6e9

Browse files
committed
refactor: add missing schemas and renamed files
1 parent d8a99e9 commit 09cb6e9

File tree

6 files changed

+18
-7
lines changed

6 files changed

+18
-7
lines changed

scrapegraphai/graphs/__init__.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@
1616
from .omni_scraper_graph import OmniScraperGraph
1717
from .omni_search_graph import OmniSearchGraph
1818
from .smart_scraper_multi_graph import SmartScraperMultiGraph
19-
from .pdf_scraper_multi import PdfScraperMultiGraph
20-
from .json_scraper_multi import JSONScraperMultiGraph
21-
from .csv_scraper_graph_multi import CSVScraperMultiGraph
22-
from .xml_scraper_graph_multi import XMLScraperMultiGraph
19+
from .pdf_scraper_multi_graph import PdfScraperMultiGraph
20+
from .json_scraper_multi_graph import JSONScraperMultiGraph
21+
from .csv_scraper_multi_graph import CSVScraperMultiGraph
22+
from .xml_scraper_multi_graph import XMLScraperMultiGraph
2323
from .script_creator_multi_graph import ScriptCreatorMultiGraph

scrapegraphai/graphs/json_scraper_multi.py renamed to scrapegraphai/graphs/json_scraper_multi_graph.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
from copy import copy, deepcopy
66
from typing import List, Optional
7+
from pydantic import BaseModel
78

89
from .base_graph import BaseGraph
910
from .abstract_graph import AbstractGraph
@@ -42,7 +43,7 @@ class JSONScraperMultiGraph(AbstractGraph):
4243
>>> result = search_graph.run()
4344
"""
4445

45-
def __init__(self, prompt: str, source: List[str], config: dict, schema: Optional[str] = None):
46+
def __init__(self, prompt: str, source: List[str], config: dict, schema: Optional[BaseModel] = None):
4647

4748
self.max_results = config.get("max_results", 3)
4849

@@ -51,6 +52,8 @@ def __init__(self, prompt: str, source: List[str], config: dict, schema: Optiona
5152
else:
5253
self.copy_config = deepcopy(config)
5354

55+
self.copy_schema = deepcopy(schema)
56+
5457
super().__init__(prompt, config, source, schema)
5558

5659
def _create_graph(self) -> BaseGraph:
@@ -69,6 +72,7 @@ def _create_graph(self) -> BaseGraph:
6972
prompt="",
7073
source="",
7174
config=self.copy_config,
75+
schema=self.copy_schema
7276
)
7377

7478
# ************************************************

scrapegraphai/graphs/omni_search_graph.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,8 @@ def __init__(self, prompt: str, config: dict, schema: Optional[BaseModel] = None
5353
else:
5454
self.copy_config = deepcopy(config)
5555

56+
self.copy_schema = deepcopy(schema)
57+
5658
super().__init__(prompt, config, schema)
5759

5860
def _create_graph(self) -> BaseGraph:
@@ -70,7 +72,8 @@ def _create_graph(self) -> BaseGraph:
7072
omni_scraper_instance = OmniScraperGraph(
7173
prompt="",
7274
source="",
73-
config=self.copy_config
75+
config=self.copy_config,
76+
schema=self.copy_schema
7477
)
7578

7679
# ************************************************

scrapegraphai/graphs/xml_scraper_graph_multi.py renamed to scrapegraphai/graphs/xml_scraper_multi_graph.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
from copy import copy, deepcopy
66
from typing import List, Optional
7+
from pydantic import BaseModel
78

89
from .base_graph import BaseGraph
910
from .abstract_graph import AbstractGraph
@@ -43,7 +44,7 @@ class XMLScraperMultiGraph(AbstractGraph):
4344
>>> result = search_graph.run()
4445
"""
4546

46-
def __init__(self, prompt: str, source: List[str], config: dict, schema: Optional[str] = None):
47+
def __init__(self, prompt: str, source: List[str], config: dict, schema: Optional[BaseModel] = None):
4748

4849
self.max_results = config.get("max_results", 3)
4950

@@ -52,6 +53,8 @@ def __init__(self, prompt: str, source: List[str], config: dict, schema: Optiona
5253
else:
5354
self.copy_config = deepcopy(config)
5455

56+
self.copy_schema = deepcopy(schema)
57+
5558
super().__init__(prompt, config, source, schema)
5659

5760
def _create_graph(self) -> BaseGraph:
@@ -70,6 +73,7 @@ def _create_graph(self) -> BaseGraph:
7073
prompt="",
7174
source="",
7275
config=self.copy_config,
76+
schema=self.copy_schema
7377
)
7478

7579
# ************************************************

0 commit comments

Comments
 (0)