Skip to content

Commit

Permalink
feat: refactoring of ScrapeGraph to SmartScraperLiteGraph
Browse files Browse the repository at this point in the history
  • Loading branch information
VinciGit00 committed Oct 21, 2024
1 parent b84883b commit 52b6bf5
Show file tree
Hide file tree
Showing 3 changed files with 8 additions and 7 deletions.
2 changes: 1 addition & 1 deletion scrapegraphai/graphs/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,4 +26,4 @@
from .code_generator_graph import CodeGeneratorGraph
from .depth_search_graph import DepthSearchGraph
from .smart_scraper_multi_lite_graph import SmartScraperMultiLiteGraph
from .scrape_graph import ScrapeGraph
from .smart_scraper_lite_graph import SmartScraperLiteGraph
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,9 @@
ParseNode,
)

class ScrapeGraph(AbstractGraph):
class SmartScraperLiteGraph(AbstractGraph):
"""
ScrapeGraph is a scraping pipeline that automates the process of
SmartScraperLiteGraph is a scraping pipeline that automates the process of
extracting information from web pages.
Attributes:
Expand All @@ -30,15 +30,16 @@ class ScrapeGraph(AbstractGraph):
schema (BaseModel): The schema for the graph output.
Example:
>>> scraper = ScraperGraph(
>>> scraper = SmartScraperLiteGraph(
... "https://en.wikipedia.org/wiki/Chioggia",
... {"llm": {"model": "openai/gpt-3.5-turbo"}}
... )
>>> result = smart_scraper.run()
)
"""

def __init__(self, source: str, config: dict, prompt: str = "", schema: Optional[BaseModel] = None):
def __init__(self, source: str, config: dict, prompt: str = "",
schema: Optional[BaseModel] = None):
super().__init__(prompt, config, source, schema)

self.input_key = "url" if source.startswith("http") else "local_dir"
Expand Down
4 changes: 2 additions & 2 deletions scrapegraphai/graphs/smart_scraper_multi_lite_graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from pydantic import BaseModel
from .base_graph import BaseGraph
from .abstract_graph import AbstractGraph
from .scrape_graph import ScrapeGraph
from .smart_scraper_lite_graph import SmartScraperLiteGraph
from ..nodes import (
GraphIteratorNode,
MergeAnswersNode,
Expand Down Expand Up @@ -63,7 +63,7 @@ def _create_graph(self) -> BaseGraph:
input="user_prompt & urls",
output=["parsed_doc"],
node_config={
"graph_instance": ScrapeGraph,
"graph_instance": SmartScraperLiteGraph,
"scraper_config": self.copy_config,
},
schema=self.copy_schema
Expand Down

0 comments on commit 52b6bf5

Please sign in to comment.