fix: bug with fetch node

This commit is contained in:
VinciGit00 2024-04-27 21:23:35 +02:00
parent 3c77acbb1d
commit 9cd516507c
5 changed files with 76 additions and 14 deletions

View File

@ -0,0 +1,27 @@
"""
Example of custom graph using existing nodes
"""
from scrapegraphai.nodes import FetchNode
# ************************************************
# Define the node
# ************************************************
robots_node = FetchNode(
input="url | local_dir",
output=["doc"],
)
# ************************************************
# Test the node
# ************************************************
state = {
"url": "https://twitter.com/home"
}
result = robots_node.execute(state)
print(result)

View File

@ -59,7 +59,7 @@ class FetchNode(BaseNode):
Raises: Raises:
KeyError: If the 'url' key is not found in the state, indicating that the KeyError: If the 'url' key is not found in the state, indicating that the
necessary information to perform the operation is missing. necessary information to perform the operation is missing.
""" """
print(f"--- Executing {self.node_name} Node ---") print(f"--- Executing {self.node_name} Node ---")
@ -78,7 +78,7 @@ class FetchNode(BaseNode):
})] })]
else: else:
if self.node_config.get("endpoint") is not None: if self.node_config is not None and self.node_config.get("endpoint") is not None:
loader = AsyncHtmlLoader( loader = AsyncHtmlLoader(
source, proxies={"http": self.node_config["endpoint"]}) source, proxies={"http": self.node_config["endpoint"]})
else: else:

View File

@ -1 +0,0 @@
OPENAI_APIKEY="your openai api key"

View File

@ -0,0 +1,44 @@
"""
Module for testinh robot_node
"""
import pytest
from scrapegraphai.nodes import FetchNode
@pytest.fixture
def setup():
"""
setup
"""
# ************************************************
# Define the node
# ************************************************
robots_node = FetchNode(
input="url | local_dir",
output=["doc"],
)
return robots_node
# ************************************************
# Test the node
# ************************************************
def test_robots_node(setup):
"""
Run the tests
"""
state = {
"url": "https://twitter.com/home"
}
result = setup.execute(state)
assert result is not None
# If you need to run this script directly
if __name__ == "__main__":
pytest.main()

View File

@ -1,15 +1,10 @@
""" """
Module for testinh robot_node Module for testinh robot_node
""" """
import os
from dotenv import load_dotenv
import pytest import pytest
from scrapegraphai.models import OpenAI from scrapegraphai.models import Ollama
from scrapegraphai.nodes import RobotsNode from scrapegraphai.nodes import RobotsNode
# Load environment variables from .env file
load_dotenv()
@pytest.fixture @pytest.fixture
def setup(): def setup():
@ -20,12 +15,9 @@ def setup():
# Define the configuration for the graph # Define the configuration for the graph
# ************************************************ # ************************************************
openai_key = os.getenv("OPENAI_APIKEY")
graph_config = { graph_config = {
"llm": { "llm": {
"api_key": openai_key, "model": "ollama/llama3",
"model": "gpt-3.5-turbo",
"temperature": 0, "temperature": 0,
"streaming": True "streaming": True
}, },
@ -35,7 +27,7 @@ def setup():
# Define the node # Define the node
# ************************************************ # ************************************************
llm_model = OpenAI(graph_config["llm"]) llm_model = Ollama(graph_config["llm"])
robots_node = RobotsNode( robots_node = RobotsNode(
input="url", input="url",