docs: graphs and helpers docstrings

Author: PeriniM

scrapegraphai/builders/graph_builder.py

@@ -1,5 +1,5 @@
 """ 
-Module for making the graph building
+GraphBuilder Module
 """
 
 from langchain_core.prompts import ChatPromptTemplate

scrapegraphai/graphs/__init__.py

@@ -1,6 +1,7 @@
 """ 
 __init__.py file for graphs folder
 """
+
 from .base_graph import BaseGraph
 from .smart_scraper_graph import SmartScraperGraph
 from .speech_graph import SpeechGraph

scrapegraphai/graphs/abstract_graph.py

@@ -1,6 +1,7 @@
 """
-Module having abstract class for creating all the graphs
+AbstractGraph Module
 """
+
 from abc import ABC, abstractmethod
 from typing import Optional
 from ..models import OpenAI, Gemini, Ollama, AzureOpenAI, HuggingFace, Groq
@@ -9,13 +10,34 @@
 
 class AbstractGraph(ABC):
  """
- Abstract class representing a generic graph-based tool.
+ Scaffolding class for creating a graph representation and executing it.
+
+ Attributes:
+ prompt (str): The prompt for the graph.
+ source (str): The source of the graph.
+ config (dict): Configuration parameters for the graph.
+ llm_model: An instance of a language model client, configured for generating answers.
+ embedder_model: An instance of an embedding model client, configured for generating embeddings.
+ verbose (bool): A flag indicating whether to show print statements during execution.
+ headless (bool): A flag indicating whether to run the graph in headless mode.
+
+ Args:
+ prompt (str): The prompt for the graph.
+ config (dict): Configuration parameters for the graph.
+ source (str, optional): The source of the graph.
+
+ Example:
+ >>> class MyGraph(AbstractGraph):
+ ... def _create_graph(self):
+ ... # Implementation of graph creation here
+ ... return graph
+ ...
+ >>> my_graph = MyGraph("Example Graph", {"llm": {"model": "gpt-3.5-turbo"}}, "example_source")
+ >>> result = my_graph.run()
  """
 
  def __init__(self, prompt: str, config: dict, source: Optional[str] = None):
- """
- Initializes the AbstractGraph with a prompt, file source, and configuration.
- """
+
  self.prompt = prompt
  self.source = source
  self.config = config
@@ -32,10 +54,20 @@ def __init__(self, prompt: str, config: dict, source: Optional[str] = None):
  self.final_state = None
  self.execution_info = None
 
- def _create_llm(self, llm_config: dict):
+ def _create_llm(self, llm_config: dict) -> object:
  """
- Creates an instance of the language model (OpenAI or Gemini) based on configuration.
+ Create a large language model instance based on the configuration provided.
+
+ Args:
+ llm_config (dict): Configuration parameters for the language model.
+
+ Returns:
+ object: An instance of the language model client.
+
+ Raises:
+ KeyError: If the model is not supported.
  """
+
  llm_defaults = {
  "temperature": 0,
  "streaming": False
@@ -104,16 +136,27 @@ def _create_llm(self, llm_config: dict):
 
  def get_state(self, key=None) -> dict:
  """""
- Obtain the current state
+ Get the final state of the graph.
+
+ Args:
+ key (str, optional): The key of the final state to retrieve.
+
+ Returns:
+ dict: The final state of the graph.
  """
+
  if key is not None:
  return self.final_state[key]
  return self.final_state
 
  def get_execution_info(self):
  """
  Returns the execution information of the graph.
+
+ Returns:
+ dict: The execution information of the graph.
  """
+ 
  return self.execution_info
 
  @abstractmethod

scrapegraphai/graphs/base_graph.py

@@ -1,6 +1,7 @@
 """
-Module for creating the base graphs
- """
+BaseGraph Module
+"""
+
 import time
 import warnings
 from langchain_community.callbacks import get_openai_callback
@@ -16,21 +17,33 @@ class BaseGraph:
  key-value pair corresponds to the from-node and to-node relationship.
  entry_point (str): The name of the entry point node from which the graph execution begins.
 
- Methods:
- execute(initial_state): Executes the graph's nodes starting from the entry point and
- traverses the graph based on the provided initial state.
-
  Args:
  nodes (iterable): An iterable of node instances that will be part of the graph.
  edges (iterable): An iterable of tuples where each tuple represents a directed edge
  in the graph, defined by a pair of nodes (from_node, to_node).
  entry_point (BaseNode): The node instance that represents the entry point of the graph.
+
+ Raises:
+ Warning: If the entry point node is not the first node in the list.
+
+ Example:
+ >>> BaseGraph(
+ ... nodes=[
+ ... fetch_node,
+ ... parse_node,
+ ... rag_node,
+ ... generate_answer_node,
+ ... ],
+ ... edges=[
+ ... (fetch_node, parse_node),
+ ... (parse_node, rag_node),
+ ... (rag_node, generate_answer_node)
+ ... ],
+ ... entry_point=fetch_node
+ ... )
  """
 
  def __init__(self, nodes: list, edges: list, entry_point: str):
- """
- Initializes the graph with nodes, edges, and the entry point.
- """
 
  self.nodes = nodes
  self.edges = self._create_edges({e for e in edges})
@@ -51,6 +64,7 @@ def _create_edges(self, edges: list) -> dict:
  Returns:
  dict: A dictionary of edges with the from-node as keys and to-node as values.
  """
+
  edge_dict = {}
  for from_node, to_node in edges:
  edge_dict[from_node.node_name] = to_node.node_name
@@ -66,8 +80,10 @@ def execute(self, initial_state: dict) -> Tuple[dict, list]:
  initial_state (dict): The initial state to pass to the entry point node.
 
  Returns:
- dict: The state after execution has completed, which may have been altered by the nodes.
+ Tuple[dict, list]: A tuple containing the final state of the execution and a list
+ of execution information for each node.
  """
+ 
  current_node_name = self.nodes[0]
  state = initial_state
 

scrapegraphai/graphs/json_scraper_graph.py

@@ -1,6 +1,7 @@
 """
-Module for creating the smart scraper
+JSONScraperGraph Module
 """
+
 from .base_graph import BaseGraph
 from ..nodes import (
  FetchNode,
@@ -13,22 +14,44 @@
 
 class JSONScraperGraph(AbstractGraph):
  """
- SmartScraper is a comprehensive web scraping tool that automates the process of extracting
- information from web pages using a natural language model to interpret and answer prompts.
+ JSONScraperGraph defines a scraping pipeline for JSON files.
+
+ Attributes:
+ prompt (str): The prompt for the graph.
+ source (str): The source of the graph.
+ config (dict): Configuration parameters for the graph.
+ llm_model: An instance of a language model client, configured for generating answers.
+ embedder_model: An instance of an embedding model client, configured for generating embeddings.
+ verbose (bool): A flag indicating whether to show print statements during execution.
+ headless (bool): A flag indicating whether to run the graph in headless mode.
+
+ Args:
+ prompt (str): The prompt for the graph.
+ source (str): The source of the graph.
+ config (dict): Configuration parameters for the graph.
+
+ Example:
+ >>> json_scraper = JSONScraperGraph(
+ ... "List me all the attractions in Chioggia.",
+ ... "data/chioggia.json",
+ ... {"llm": {"model": "gpt-3.5-turbo"}}
+ ... )
+ >>> result = json_scraper.run()
  """
 
  def __init__(self, prompt: str, source: str, config: dict):
- """
- Initializes the JsonScraperGraph with a prompt, source, and configuration.
- """
  super().__init__(prompt, config, source)
 
  self.input_key = "json" if source.endswith("json") else "json_dir"
 
- def _create_graph(self):
+ def _create_graph(self) -> BaseGraph:
  """
  Creates the graph of nodes representing the workflow for web scraping.
+ 
+ Returns:
+ BaseGraph: A graph instance representing the web scraping workflow.
  """
+
  fetch_node = FetchNode(
  input="json_dir",
  output=["doc"],
@@ -81,7 +104,11 @@ def _create_graph(self):
  def run(self) -> str:
  """
  Executes the web scraping process and returns the answer to the prompt.
+
+ Returns:
+ str: The answer to the prompt.
  """
+
  inputs = {"user_prompt": self.prompt, self.input_key: self.source}
  self.final_state, self.execution_info = self.graph.execute(inputs)
 

scrapegraphai/graphs/script_creator_graph.py

@@ -1,6 +1,7 @@
 """
-Module for creating the smart scraper
+ScriptCreatorGraph Module
 """
+
 from .base_graph import BaseGraph
 from ..nodes import (
  FetchNode,
@@ -13,24 +14,47 @@
 
 class ScriptCreatorGraph(AbstractGraph):
  """
- SmartScraper is a comprehensive web scraping tool that automates the process of extracting
- information from web pages using a natural language model to interpret and answer prompts.
+ ScriptCreatorGraph defines a scraping pipeline for generating web scraping scripts.
+
+ Attributes:
+ prompt (str): The prompt for the graph.
+ source (str): The source of the graph.
+ config (dict): Configuration parameters for the graph.
+ llm_model: An instance of a language model client, configured for generating answers.
+ embedder_model: An instance of an embedding model client, configured for generating embeddings.
+ verbose (bool): A flag indicating whether to show print statements during execution.
+ headless (bool): A flag indicating whether to run the graph in headless mode.
+ model_token (int): The token limit for the language model.
+ library (str): The library used for web scraping.
+
+ Args:
+ prompt (str): The prompt for the graph.
+ source (str): The source of the graph.
+ config (dict): Configuration parameters for the graph.
+
+ Example:
+ >>> script_creator = ScriptCreatorGraph(
+ ... "List me all the attractions in Chioggia.",
+ ... "https://en.wikipedia.org/wiki/Chioggia",
+ ... {"llm": {"model": "gpt-3.5-turbo"}}
+ ... )
+ >>> result = script_creator.run()
  """
 
  def __init__(self, prompt: str, source: str, config: dict):
- """
- Initializes the ScriptCreatorGraph with a prompt, source, and configuration.
- """
- self.library = config['library']
-
  super().__init__(prompt, config, source)
 
  self.input_key = "url" if source.startswith("http") else "local_dir"
+ self.library = config['library']
 
- def _create_graph(self):
+ def _create_graph(self) -> BaseGraph:
  """
  Creates the graph of nodes representing the workflow for web scraping.
+
+ Returns:
+ BaseGraph: A graph instance representing the web scraping workflow.
  """
+
  fetch_node = FetchNode(
  input="url | local_dir",
  output=["doc"],
@@ -76,7 +100,11 @@ def _create_graph(self):
  def run(self) -> str:
  """
  Executes the web scraping process and returns the answer to the prompt.
+
+ Returns:
+ str: The answer to the prompt.
  """
+ 
  inputs = {"user_prompt": self.prompt, self.input_key: self.source}
  self.final_state, self.execution_info = self.graph.execute(inputs)
 

scrapegraphai/graphs/search_graph.py

@@ -1,6 +1,7 @@
 """ 
-Module for making the search on the intenet
+SearchGraph Module
 """
+
 from .base_graph import BaseGraph
 from ..nodes import (
  SearchInternetNode,
@@ -14,13 +15,37 @@
 
 class SearchGraph(AbstractGraph):
  """ 
- Module for searching info on the internet
+ SearchGraph is a scraping pipeline that searches the internet for answers to a given prompt.
+ It only requires a user prompt to search the internet and generate an answer.
+
+ Attributes:
+ prompt (str): The user prompt to search the internet.
+ llm_model (dict): The configuration for the language model.
+ embedder_model (dict): The configuration for the embedder model.
+ headless (bool): A flag to run the browser in headless mode.
+ verbose (bool): A flag to display the execution information.
+ model_token (int): The token limit for the language model.
+
+ Args:
+ prompt (str): The user prompt to search the internet.
+ config (dict): Configuration parameters for the graph.
+
+ Example:
+ >>> search_graph = SearchGraph(
+ ... "What is Chioggia famous for?",
+ ... {"llm": {"model": "gpt-3.5-turbo"}}
+ ... )
+ >>> result = search_graph.run()
  """
 
- def _create_graph(self):
+ def _create_graph(self) -> BaseGraph:
  """
  Creates the graph of nodes representing the workflow for web scraping and searching.
+
+ Returns:
+ BaseGraph: A graph instance representing the web scraping and searching workflow.
  """
+
  search_internet_node = SearchInternetNode(
  input="user_prompt",
  output=["url"],
@@ -83,7 +108,11 @@ def _create_graph(self):
  def run(self) -> str:
  """
  Executes the web scraping and searching process.
+ 
+ Returns:
+ str: The answer to the prompt.
  """
+ 
  inputs = {"user_prompt": self.prompt}
  self.final_state, self.execution_info = self.graph.execute(inputs)