Merge pull request #895 from Telsho/main

fix: Schema parameter type
ScrapeGraphAI · Jan 21, 2025 · fe9570b · fe9570b
2 parents 67de52a + 2b5bd80
commit fe9570b
Show file tree

Hide file tree

Showing 25 changed files with 102 additions and 52 deletions.
diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py
@@ -6,7 +6,7 @@
 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 @@ def __init__(
         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
@@ -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
@@ -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
@@ -3,7 +3,7 @@
 """
 
 from copy import deepcopy
-from typing import List, Optional
+from typing import List, Optional, Type
 
 from pydantic import BaseModel
 
@@ -47,7 +47,7 @@ def __init__(
         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
@@ -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
@@ -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
@@ -3,7 +3,7 @@
 """
 
 from copy import deepcopy
-from typing import List, Optional
+from typing import List, Optional, Type
 
 from pydantic import BaseModel
 
@@ -47,7 +47,7 @@ def __init__(
         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
@@ -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
@@ -3,7 +3,7 @@
 """
 
 from copy import deepcopy
-from typing import List, Optional
+from typing import List, Optional, Type
 
 from pydantic import BaseModel
 
@@ -47,7 +47,7 @@ def __init__(
         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
@@ -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
@@ -3,7 +3,7 @@
 """
 
 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
@@ -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
@@ -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
@@ -3,7 +3,7 @@
 """
 
 from copy import deepcopy
-from typing import List, Optional
+from typing import List, Optional, Type
 
 from pydantic import BaseModel
 
@@ -46,7 +46,7 @@ def __init__(
         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
@@ -3,7 +3,7 @@
 """
 
 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
@@ -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
@@ -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
@@ -2,7 +2,7 @@
 SmartScraperGraph Module
 """
 
-from typing import Optional
+from typing import Optional, Type
 
 from pydantic import BaseModel
 
@@ -44,7 +44,7 @@ def __init__(
         source: str,
         config: dict,
         prompt: str = "",
-        schema: Optional[BaseModel] = None,
+        schema: Optional[Type[BaseModel]] = None,
     ):
         super().__init__(prompt, config, source, schema)