fix: Schema parameter type

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

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)
 

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.

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)

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)
 

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)
 

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)

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)
 

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)

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)
 

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)
 

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)
 

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"]
 

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)

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)

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"

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)
 

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)
 

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)

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)

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)

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)
 

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)
 

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)

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)