feat: fix tests

This commit is contained in:
Marco Vinciguerra 2024-07-30 16:12:31 +02:00
parent 88710f1a7c
commit 1db164e9e6
4 changed files with 7 additions and 8 deletions

View File

@ -2,7 +2,7 @@
Example of custom graph using existing nodes Example of custom graph using existing nodes
""" """
from scrapegraphai.models import Ollama from langchain_community.chat_models import ChatOllama
from scrapegraphai.nodes import RobotsNode from scrapegraphai.nodes import RobotsNode
# ************************************************ # ************************************************
@ -26,7 +26,7 @@ graph_config = {
# Define the node # Define the node
# ************************************************ # ************************************************
llm_model = Ollama(graph_config["llm"]) llm_model = ChatOllama(graph_config["llm"])
robots_node = RobotsNode( robots_node = RobotsNode(
input="url", input="url",

View File

@ -1,7 +1,6 @@
import pytest import pytest
from unittest.mock import MagicMock from unittest.mock import MagicMock
from langchain_community.chat_models import ChatOllama
from scrapegraphai.models import Ollama
from scrapegraphai.nodes import RobotsNode from scrapegraphai.nodes import RobotsNode
@pytest.fixture @pytest.fixture

View File

@ -1,5 +1,5 @@
import unittest import unittest
from scrapegraphai.models import Ollama from langchain_community.chat_models import ChatOllama
from scrapegraphai.nodes import SearchInternetNode from scrapegraphai.nodes import SearchInternetNode
class TestSearchInternetNode(unittest.TestCase): class TestSearchInternetNode(unittest.TestCase):
@ -18,7 +18,7 @@ class TestSearchInternetNode(unittest.TestCase):
} }
# Define the model # Define the model
self.llm_model = Ollama(self.graph_config["llm"]) self.llm_model = ChatOllama(self.graph_config["llm"])
# Initialize the SearchInternetNode # Initialize the SearchInternetNode
self.search_node = SearchInternetNode( self.search_node = SearchInternetNode(

View File

@ -1,5 +1,5 @@
import pytest import pytest
from scrapegraphai.models import Ollama from langchain_community.chat_models import ChatOllama
from scrapegraphai.nodes import SearchLinkNode from scrapegraphai.nodes import SearchLinkNode
from unittest.mock import patch, MagicMock from unittest.mock import patch, MagicMock
@ -18,7 +18,7 @@ def setup():
} }
# Instantiate the LLM model with the configuration # Instantiate the LLM model with the configuration
llm_model = Ollama(graph_config["llm"]) llm_model = ChatOllama(graph_config["llm"])
# Define the SearchLinkNode with necessary configurations # Define the SearchLinkNode with necessary configurations
search_link_node = SearchLinkNode( search_link_node = SearchLinkNode(