From 1db164e9e682eefbc1414989a043fefa2e9009c2 Mon Sep 17 00:00:00 2001 From: Marco Vinciguerra Date: Tue, 30 Jul 2024 16:12:31 +0200 Subject: [PATCH] feat: fix tests --- examples/single_node/robot_node.py | 4 ++-- tests/nodes/robot_node_test.py | 3 +-- tests/nodes/search_internet_node_test.py | 4 ++-- tests/nodes/search_link_node_test.py | 4 ++-- 4 files changed, 7 insertions(+), 8 deletions(-) diff --git a/examples/single_node/robot_node.py b/examples/single_node/robot_node.py index f51f8649..c2bcbbd1 100644 --- a/examples/single_node/robot_node.py +++ b/examples/single_node/robot_node.py @@ -2,7 +2,7 @@ Example of custom graph using existing nodes """ -from scrapegraphai.models import Ollama +from langchain_community.chat_models import ChatOllama from scrapegraphai.nodes import RobotsNode # ************************************************ @@ -26,7 +26,7 @@ graph_config = { # Define the node # ************************************************ -llm_model = Ollama(graph_config["llm"]) +llm_model = ChatOllama(graph_config["llm"]) robots_node = RobotsNode( input="url", diff --git a/tests/nodes/robot_node_test.py b/tests/nodes/robot_node_test.py index 00a45b05..62527dda 100644 --- a/tests/nodes/robot_node_test.py +++ b/tests/nodes/robot_node_test.py @@ -1,7 +1,6 @@ import pytest from unittest.mock import MagicMock - -from scrapegraphai.models import Ollama +from langchain_community.chat_models import ChatOllama from scrapegraphai.nodes import RobotsNode @pytest.fixture diff --git a/tests/nodes/search_internet_node_test.py b/tests/nodes/search_internet_node_test.py index db2cbdee..8e198448 100644 --- a/tests/nodes/search_internet_node_test.py +++ b/tests/nodes/search_internet_node_test.py @@ -1,5 +1,5 @@ import unittest -from scrapegraphai.models import Ollama +from langchain_community.chat_models import ChatOllama from scrapegraphai.nodes import SearchInternetNode class TestSearchInternetNode(unittest.TestCase): @@ -18,7 +18,7 @@ class TestSearchInternetNode(unittest.TestCase): } # Define the model - self.llm_model = Ollama(self.graph_config["llm"]) + self.llm_model = ChatOllama(self.graph_config["llm"]) # Initialize the SearchInternetNode self.search_node = SearchInternetNode( diff --git a/tests/nodes/search_link_node_test.py b/tests/nodes/search_link_node_test.py index 648db4ee..1f8c5a58 100644 --- a/tests/nodes/search_link_node_test.py +++ b/tests/nodes/search_link_node_test.py @@ -1,5 +1,5 @@ import pytest -from scrapegraphai.models import Ollama +from langchain_community.chat_models import ChatOllama from scrapegraphai.nodes import SearchLinkNode from unittest.mock import patch, MagicMock @@ -18,7 +18,7 @@ def setup(): } # 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 search_link_node = SearchLinkNode(