fix: examples and graphs

This commit is contained in:
VinciGit00 2024-05-02 09:20:46 +02:00
parent ba2b24b4cd
commit 5cf4e4f92f
12 changed files with 95 additions and 21 deletions

View File

@ -51,6 +51,7 @@ Please make sure to format your code accordingly before submitting a pull reques
- [Style Guide for Python Code](https://www.python.org/dev/peps/pep-0008/)
- [Google Python Style Guide](https://google.github.io/styleguide/pyguide.html)
- [The Hitchhiker's Guide to Python](https://docs.python-guide.org/writing/style/)
- [Pylint style of code for the documentation](https://pylint.pycqa.org/en/1.6.0/tutorial.html)
## Submitting a Pull Request

View File

@ -12,6 +12,9 @@ from scrapegraphai.nodes import FetchNode
robots_node = FetchNode(
input="url | local_dir",
output=["doc"],
node_config={
"headless": False
}
)
# ************************************************

View File

@ -26,7 +26,9 @@ llm_model = Ollama(graph_config["llm"])
robots_node = RobotsNode(
input="url",
output=["is_scrapable"],
node_config={"llm": llm_model}
node_config={"llm": llm_model,
"headless": False
}
)
# ************************************************

View File

@ -21,7 +21,8 @@ class JSONScraperGraph(AbstractGraph):
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.
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.
@ -47,7 +48,7 @@ class JSONScraperGraph(AbstractGraph):
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.
"""

View File

@ -21,7 +21,8 @@ class ScriptCreatorGraph(AbstractGraph):
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.
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.
@ -44,7 +45,7 @@ class ScriptCreatorGraph(AbstractGraph):
def __init__(self, prompt: str, source: str, config: dict):
self.library = config['library']
super().__init__(prompt, config, source)
self.input_key = "url" if source.startswith("http") else "local_dir"
@ -61,25 +62,29 @@ class ScriptCreatorGraph(AbstractGraph):
input="url | local_dir",
output=["doc"],
node_config={
"headless": True if self.config is None else self.config.get("headless", True)}
"headless": True if self.config is None else self.config.get("headless", True),
"verbose": self.verbose}
)
parse_node = ParseNode(
input="doc",
output=["parsed_doc"],
node_config={"chunk_size": self.model_token}
node_config={"chunk_size": self.model_token,
"verbose": self.verbose}
)
rag_node = RAGNode(
input="user_prompt & (parsed_doc | doc)",
output=["relevant_chunks"],
node_config={
"llm": self.llm_model,
"embedder_model": self.embedder_model
"embedder_model": self.embedder_model,
"verbose": self.verbose
}
)
generate_scraper_node = GenerateScraperNode(
input="user_prompt & (relevant_chunks | parsed_doc | doc)",
output=["answer"],
node_config={"llm": self.llm_model},
node_config={"llm": self.llm_model,
"verbose": self.verbose},
library=self.library,
website=self.source
)
@ -106,7 +111,7 @@ class ScriptCreatorGraph(AbstractGraph):
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)

View File

@ -14,7 +14,8 @@ from .abstract_graph import AbstractGraph
class SmartScraperGraph(AbstractGraph):
"""
SmartScraper is a scraping pipeline that automates the process of extracting information from web pages
SmartScraper is a scraping pipeline that automates the process of
extracting information from web pages
using a natural language model to interpret and answer prompts.
Attributes:
@ -22,7 +23,8 @@ class SmartScraperGraph(AbstractGraph):
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.
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.
@ -45,7 +47,7 @@ class SmartScraperGraph(AbstractGraph):
super().__init__(prompt, config, source)
self.input_key = "url" if source.startswith("http") else "local_dir"
def _create_graph(self) -> BaseGraph:
"""
Creates the graph of nodes representing the workflow for web scraping.

View File

@ -22,7 +22,8 @@ class XMLScraperGraph(AbstractGraph):
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.
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.
@ -49,7 +50,7 @@ class XMLScraperGraph(AbstractGraph):
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.
"""
@ -110,7 +111,7 @@ class XMLScraperGraph(AbstractGraph):
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)

View File

@ -0,0 +1,56 @@
"""
Module for scraping json documents
"""
import os
import pytest
from scrapegraphai.graphs import JSONScraperGraph
@pytest.fixture
def sample_json():
"""
Example of text
"""
file_name = "inputs/example.json"
curr_dir = os.path.dirname(os.path.realpath(__file__))
file_path = os.path.join(curr_dir, file_name)
with open(file_path, 'r', encoding="utf-8") as file:
text = file.read()
return text
@pytest.fixture
def graph_config():
"""
Configuration of the graph
"""
return {
"llm": {
"model": "ollama/mistral",
"temperature": 0,
"format": "json",
"base_url": "http://localhost:11434",
},
"embeddings": {
"model": "ollama/nomic-embed-text",
"temperature": 0,
"base_url": "http://localhost:11434",
}
}
def test_scraping_pipeline(sample_json: str, graph_config: dict):
"""
Start of the scraping pipeline
"""
smart_scraper_graph = JSONScraperGraph(
prompt="List me all the titles",
source=sample_json,
config=graph_config
)
result = smart_scraper_graph.run()
assert result is not None

View File

@ -3,7 +3,7 @@ Module for scraping XML documents
"""
import os
import pytest
from scrapegraphai.graphs import SmartScraperGraph
from scrapegraphai.graphs import XMLScraperGraph
@pytest.fixture
@ -45,7 +45,7 @@ def test_scraping_pipeline(sample_xml: str, graph_config: dict):
"""
Start of the scraping pipeline
"""
smart_scraper_graph = SmartScraperGraph(
smart_scraper_graph = XMLScraperGraph(
prompt="List me all the authors, title and genres of the books",
source=sample_xml,
config=graph_config

View File

@ -46,6 +46,4 @@ def test_script_creator_graph(graph_config: dict):
assert graph_exec_info is not None
assert isinstance(graph_exec_info, dict)
print(prettify_exec_info(graph_exec_info))

View File

@ -17,6 +17,9 @@ def setup():
robots_node = FetchNode(
input="url | local_dir",
output=["doc"],
node_config={
"headless": False
}
)
return robots_node

View File

@ -32,7 +32,9 @@ def setup():
robots_node = RobotsNode(
input="url",
output=["is_scrapable"],
node_config={"llm": llm_model}
node_config={"llm": llm_model,
"headless": False
}
)
return robots_node